From ae8aeb5393702c553c4e51f03770de120073e51c Mon Sep 17 00:00:00 2001 From: Kirthi Shankar Sivamani Date: Fri, 23 Oct 2026 11:49:44 -0700 Subject: [PATCH] Use device CC check from TE; fix GH doc CI Signed-off-by: Kirthi Shankar Sivamani --- transformer_engine/jax/triton_extensions/mhc.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/transformer_engine/jax/triton_extensions/mhc.py b/transformer_engine/jax/triton_extensions/mhc.py index 8913dd6906..f36a00b004 100644 --- a/transformer_engine/jax/triton_extensions/mhc.py +++ b/transformer_engine/jax/triton_extensions/mhc.py @@ -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 ( @@ -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", @@ -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)