[Feat] Added JAX-Triton bridge for ROCm

ci/jax.sh

@@ -65,8 +65,10 @@ run_test_config() {
     # group-mode per-segment dq_acc layout matters (see the equal-dim-128 RAGGED_SELF config).
     NVTE_CK_IS_V3_ATOMIC_FP32=0 run_default_fa_lbl "atomic16" 3 test_fused_attn.py -k "test_backward and RAGGED"
     run_default_fa 1 test_layer.py # it effectively always uses unfused attention
+    run_default_fa 1 test_permutation.py
     run_default_fa 1 test_sanity_import.py
     run_default_fa 1 test_softmax.py
+    run_default_fa 1 test_triton_custom_calls.py
 }
 
 run_test_config_mgpu() {
@@ -101,6 +103,7 @@ run_test_config_mgpu() {
     export XLA_FLAGS="${XLA_FLAGS} --xla_gpu_enable_triton_gemm=false"
     run_default_fa 2 test_distributed_layernorm_mlp.py -k "${_layernorm_mlp_jax_gemm_k}"
     export XLA_FLAGS="$_saved_xla_flags"
+    run_default_fa 3 test_distributed_permutation.py
     run_default_fa 3 test_distributed_softmax.py
 
     run_default_fa 3 test_sanity_import.py

tests/jax/test_distributed_permutation.py

@@ -1,3 +1,5 @@
+# This file was modified for portability to AMDGPU
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
 # Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 #
 # See LICENSE for license information.
@@ -75,6 +77,7 @@ def _inject_permutation(request):
 ]
 DISPATCH_COMBINE_CASES = {
     "L0": ALL_DISPATCH_COMBINE_CASES[0:1],
+    "L1": ALL_DISPATCH_COMBINE_CASES[0:2],
     "L2": ALL_DISPATCH_COMBINE_CASES,
 }
 
@@ -86,13 +89,15 @@ def _inject_permutation(request):
 ]
 DISPATCH_COMBINE_PADDING_CASES = {
     "L0": ALL_DISPATCH_COMBINE_PADDING_CASES[0:1],
+    "L1": ALL_DISPATCH_COMBINE_PADDING_CASES[0:2],
     "L2": ALL_DISPATCH_COMBINE_PADDING_CASES,
 }
 
 # Dtypes for testing
 ALL_DTYPES = [jnp.float32, jnp.bfloat16]
 DTYPES = {
     "L0": [jnp.float32],
+    "L1": [jnp.float32],
     "L2": ALL_DTYPES,
 }
 

tests/jax/test_permutation.py

@@ -1,3 +1,5 @@
+# This file was modified for portability to AMDGPU
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
 # Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 #
 # See LICENSE for license information.
@@ -48,6 +50,7 @@ def _inject_permutation(request):
 ]
 DISPATCH_COMBINE_CASES = {
     "L0": ALL_DISPATCH_COMBINE_CASES[0:2],
+    "L1": ALL_DISPATCH_COMBINE_CASES[0:3],
     "L2": ALL_DISPATCH_COMBINE_CASES,
 }
 
@@ -58,6 +61,7 @@ def _inject_permutation(request):
 ]
 SORT_CHUNKS_CASES = {
     "L0": ALL_SORT_CHUNKS_CASES[0:2],
+    "L1": ALL_SORT_CHUNKS_CASES[0:2],
     "L2": ALL_SORT_CHUNKS_CASES,
 }
 
@@ -69,18 +73,21 @@ def _inject_permutation(request):
 ]
 DISPATCH_COMBINE_PADDING_CASES = {
     "L0": ALL_DISPATCH_COMBINE_PADDING_CASES[0:2],
+    "L1": ALL_DISPATCH_COMBINE_PADDING_CASES[0:3],
     "L2": ALL_DISPATCH_COMBINE_PADDING_CASES,
 }
 
 ALL_DTYPES = [jnp.float32, jnp.bfloat16]
 DTYPES = {
     "L0": ALL_DTYPES,
+    "L1": ALL_DTYPES,
     "L2": ALL_DTYPES,
 }
 
 ALL_WITH_PROBS = [True, False]
 WITH_PROBS = {
     "L0": [True],
+    "L1": [True],
     "L2": ALL_WITH_PROBS,
 }
 

tests/jax/test_triton_custom_calls.py

@@ -1,3 +1,5 @@
+# This file was modified for portability to AMDGPU
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
 # Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 #
 # See LICENSE for license information.
@@ -11,12 +13,50 @@
 
 require_triton_or_skip_test_file()
 
+# require_triton_or_skip_test_file only checks the JAX version, so guard triton itself.
+pytest.importorskip("triton")
+
 import triton
 import triton.language as tl
 
 from transformer_engine.jax.cpp_extensions.base import BasePrimitive, register_primitive
 from transformer_engine.jax.triton_extensions import triton_call_lowering
 
+# Gluon support is optional and warp-size explicit. HAS_GLUON gates the Gluon
+# test below so the Triton test still runs when Gluon is unavailable.
+from typing import TYPE_CHECKING
+
+GLUON_SKIP_REASON = ""
+
+if TYPE_CHECKING:
+    # Resolve Gluon symbols for the type checker; the runtime guard is below.
+    from triton.experimental import gluon
+    from triton.experimental.gluon import language as gl
+
+    HAS_GLUON = True
+    WARP_SIZE = 64
+else:
+    try:
+        from packaging.version import Version
+
+        from triton.experimental import gluon
+        from triton.experimental.gluon import language as gl
+
+        WARP_SIZE = triton.runtime.driver.active.get_current_target().warp_size
+        # Gluon-on-ROCm needs the make_ir fix shipped after Triton 3.4.0. Compare
+        # the base release so a "+rocm.git" local segment can't slip past.
+        HAS_GLUON = Version(triton.__version__).release > (3, 4, 0)
+        if not HAS_GLUON:
+            GLUON_SKIP_REASON = f"Gluon needs Triton > 3.4.0, found {triton.__version__}"
+    except (ImportError, AttributeError, RuntimeError) as e:
+        # Anything else is a real breakage, not a missing Gluon.
+        gluon = gl = None
+        WARP_SIZE = None
+        HAS_GLUON = False
+        GLUON_SKIP_REASON = f"Gluon unavailable: {type(e).__name__}: {e}"
+
+requires_gluon = pytest.mark.skipif(not HAS_GLUON, reason=GLUON_SKIP_REASON or "Gluon unavailable")
+
 
 @pytest.fixture(autouse=True, scope="module")
 def init():
@@ -116,3 +156,186 @@ def test_triton_amax(self, shape, dtype):
         result = jitted_amax(x)
 
         assert_allclose(result, expected, dtype=jnp.float32)
+
+
+# Gluon binding tests. Drive a @gluon.jit kernel through the same
+# triton_call_lowering bridge as the Triton test above.
+if HAS_GLUON:
+    # BLOCK_SIZE must divide NUM_WARPS * WARP_SIZE so the 1-D layout tiles exactly.
+    BLOCK_SIZE = 1024
+    NUM_WARPS = 4
+
+    # Elementwise out = x * 2, shared by the jit and autotuned tests. Gluon is
+    # layout-explicit: `arange` needs a BlockedLayout whose warps_per_cta matches
+    # the launch num_warps.
+    @gluon.jit
+    def _double_kernel(
+        x_ptr,
+        out_ptr,
+        n_elements: gl.constexpr,
+        BLOCK_SIZE: gl.constexpr,
+        NUM_WARPS: gl.constexpr,
+        WARP_SIZE: gl.constexpr,
+    ):
+        """Multiply each element by 2 using Gluon."""
+        SIZE_PER_THREAD: gl.constexpr = BLOCK_SIZE // (NUM_WARPS * WARP_SIZE)
+        layout: gl.constexpr = gl.BlockedLayout(
+            size_per_thread=[SIZE_PER_THREAD],
+            threads_per_warp=[WARP_SIZE],
+            warps_per_cta=[NUM_WARPS],
+            order=[0],
+        )
+
+        pid = gl.program_id(0)
+        offsets = pid * BLOCK_SIZE + gl.arange(0, BLOCK_SIZE, layout=layout)
+        mask = offsets < n_elements
+        x = gl.load(x_ptr + offsets, mask)
+        gl.store(out_ptr + offsets, x * 2.0, mask)
+
+else:
+    _double_kernel = None
+
+
+@pytest.mark.triton
+@requires_gluon
+class TestGluonBinding:
+    """Test Gluon binding primitive through the Triton custom-call bridge."""
+
+    # Define test primitive
+    class DoubleGluonPrimitive(BasePrimitive):
+        """Test primitive using a Gluon kernel."""
+
+        name = "te_double_gluon_test"
+        multiple_results = False
+        impl_static_args = ()
+
+        @staticmethod
+        def abstract(x_aval):
+            return jax.core.ShapedArray(x_aval.shape, x_aval.dtype)
+
+        @staticmethod
+        def impl(x):
+            assert TestGluonBinding.DoubleGluonPrimitive.inner_primitive is not None
+            return TestGluonBinding.DoubleGluonPrimitive.inner_primitive.bind(x)
+
+        @staticmethod
+        def lowering(ctx, x):
+            """MLIR lowering using the Gluon kernel."""
+            n_elements = 1
+            for dim in ctx.avals_in[0].shape:
+                n_elements *= dim
+
+            grid = (triton.cdiv(n_elements, BLOCK_SIZE),)
+
+            return triton_call_lowering(
+                ctx,
+                _double_kernel,
+                x,
+                grid=grid,
+                constexprs={
+                    "n_elements": n_elements,
+                    "BLOCK_SIZE": BLOCK_SIZE,
+                    "NUM_WARPS": NUM_WARPS,
+                    "WARP_SIZE": WARP_SIZE,
+                },
+                num_warps=NUM_WARPS,  # must match warps_per_cta in the layout
+            )
+
+    register_primitive(DoubleGluonPrimitive)
+
+    @staticmethod
+    def _gluon_double(x: jnp.ndarray) -> jnp.ndarray:
+        """Double each element using the Gluon kernel."""
+        return TestGluonBinding.DoubleGluonPrimitive.outer_primitive.bind(x)
+
+    @pytest_parametrize_wrapper("shape", [(1024, 1024), (1000, 1000)])
+    @pytest_parametrize_wrapper("dtype", [jnp.float32])
+    def test_gluon_double(self, shape, dtype):
+        """Test the Gluon double kernel with JIT."""
+        key = jax.random.PRNGKey(0)
+        x = jax.random.uniform(key, shape, dtype)
+
+        expected = (x * 2.0).astype(dtype)
+        jitted_double = jax.jit(self._gluon_double)
+        result = jitted_double(x)
+
+        assert_allclose(result, expected, dtype=dtype)
+
+
+@pytest.mark.triton
+@requires_gluon
+class TestGluonAutotunedBinding:
+    """Autotuned Gluon binding; exercises the TritonAutotunedKernelCall path."""
+
+    # Reuse the shared kernel; each config carries BLOCK_SIZE and NUM_WARPS as
+    # constexprs (the layout needs them) with a matching launch num_warps.
+    double_kernel_autotuned = (
+        triton.autotune(
+            configs=[
+                triton.Config({"BLOCK_SIZE": 256, "NUM_WARPS": 4}, num_warps=4),
+                triton.Config({"BLOCK_SIZE": 1024, "NUM_WARPS": 8}, num_warps=8),
+            ],
+            key=["n_elements"],
+        )(_double_kernel)
+        if HAS_GLUON
+        else None
+    )
+
+    class DoubleGluonAutotunedPrimitive(BasePrimitive):
+        """Test primitive using an autotuned Gluon kernel."""
+
+        name = "te_double_gluon_autotuned_test"
+        multiple_results = False
+        impl_static_args = ()
+
+        @staticmethod
+        def abstract(x_aval):
+            return jax.core.ShapedArray(x_aval.shape, x_aval.dtype)
+
+        @staticmethod
+        def impl(x):
+            prim = TestGluonAutotunedBinding.DoubleGluonAutotunedPrimitive
+            assert prim.inner_primitive is not None
+            return prim.inner_primitive.bind(x)
+
+        @staticmethod
+        def lowering(ctx, x):
+            """MLIR lowering using the autotuned Gluon kernel."""
+            n_elements = 1
+            for dim in ctx.avals_in[0].shape:
+                n_elements *= dim
+
+            kernel = TestGluonAutotunedBinding.double_kernel_autotuned
+            # Smallest BLOCK_SIZE so every config's grid covers all elements.
+            block_size = min(c.kwargs.get("BLOCK_SIZE") for c in kernel.configs)
+            grid = (triton.cdiv(n_elements, block_size),)
+
+            return triton_call_lowering(
+                ctx,
+                kernel,
+                x,
+                grid=grid,
+                # BLOCK_SIZE/NUM_WARPS come from each autotune config; only
+                # n_elements and WARP_SIZE are passed here.
+                constexprs={"n_elements": n_elements, "WARP_SIZE": WARP_SIZE},
+            )
+
+    register_primitive(DoubleGluonAutotunedPrimitive)
+
+    @staticmethod
+    def _gluon_double_autotuned(x: jnp.ndarray) -> jnp.ndarray:
+        """Double each element using the autotuned Gluon kernel."""
+        prim = TestGluonAutotunedBinding.DoubleGluonAutotunedPrimitive
+        return prim.outer_primitive.bind(x)
+
+    @pytest_parametrize_wrapper("shape", [(1024, 1024), (1000, 1000)])
+    @pytest_parametrize_wrapper("dtype", [jnp.float32])
+    def test_gluon_double_autotuned(self, shape, dtype):
+        """Test the autotuned Gluon double kernel with JIT."""
+        key = jax.random.PRNGKey(0)
+        x = jax.random.uniform(key, shape, dtype)
+
+        expected = (x * 2.0).astype(dtype)
+        result = jax.jit(self._gluon_double_autotuned)(x)
+
+        assert_allclose(result, expected, dtype=dtype)