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)