Skip to content

Use get_device_compute_capability from te jax extension; Fix GH docs CI #3735

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
ksivaman merged 1 commit into from
Oct 23, 2026
+3 −2
Merged

Use get_device_compute_capability from te jax extension; Fix GH docs CI #3735

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 3 additions & 2 deletions transformer_engine/jax/triton_extensions/mhc.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@
from jax.experimental.custom_partitioning import SdyShardingRule
import triton

from transformer_engine_jax import get_device_compute_capability
from transformer_engine.jax.cpp_extensions.base import BasePrimitive, register_primitive
from transformer_engine.jax.cpp_extensions.misc import get_padded_spec, NamedSharding
from transformer_engine.common.triton.mhc import (
Expand All @@ -35,7 +36,7 @@
TRITON_EXTENSION_GLOBAL_SCRATCH_MIN_JAX_VERSION,
jax_version_meet_requirement,
)
from .utils import triton_call_lowering, gpu_triton
from .utils import triton_call_lowering

__all__ = [
"mhc_projection_fwd",
Expand Down Expand Up @@ -232,7 +233,7 @@ def _projection_fwd_use_tma(x_aval, phi_aval) -> bool:
return mode == "1"
if not jax_version_meet_requirement(TRITON_EXTENSION_GLOBAL_SCRATCH_MIN_JAX_VERSION):
return False
if gpu_triton.get_compute_capability(0) < 90:
if get_device_compute_capability(0) < 90:
return False
return all(
(aval.shape[-1] * jnp.dtype(aval.dtype).itemsize) % 16 == 0 for aval in (x_aval, phi_aval)
Expand Down
Loading