From d85d1468065e1276f87ee1e390213b8fef3b5e68 Mon Sep 17 00:00:00 2001 From: AllenFarcas Date: Wed, 24 Jun 2026 21:56:25 +0000 Subject: [PATCH 1/7] [Feat] Added JAX-Triton bridge for ROCm --- tests/jax/test_triton_custom_calls.py | 128 +++++++++++++++ .../jax/triton_extensions/utils.py | 151 +++++++++++++----- 2 files changed, 237 insertions(+), 42 deletions(-) diff --git a/tests/jax/test_triton_custom_calls.py b/tests/jax/test_triton_custom_calls.py index 846d26a417..55ad229a7b 100644 --- a/tests/jax/test_triton_custom_calls.py +++ b/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. @@ -17,6 +19,33 @@ 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 + +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) + except Exception: # pragma: no cover - Gluon or active GPU target unavailable + gluon = gl = None + WARP_SIZE = None + HAS_GLUON = False + @pytest.fixture(autouse=True, scope="module") def init(): @@ -116,3 +145,102 @@ def test_triton_amax(self, shape, dtype): result = jitted_amax(x) assert_allclose(result, expected, dtype=jnp.float32) + + +# Gluon binding test. Drives 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 + + @pytest.mark.triton + class TestGluonBinding: + """Test Gluon binding primitive through the Triton custom-call bridge.""" + + # Elementwise out = x * 2. Gluon is layout-explicit: `arange` needs a + # BlockedLayout whose warps_per_cta matches the launch num_warps. + @staticmethod + @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) + + # 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, + TestGluonBinding.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)]) + @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) diff --git a/transformer_engine/jax/triton_extensions/utils.py b/transformer_engine/jax/triton_extensions/utils.py index 28e3f08e18..4d987e5f7a 100644 --- a/transformer_engine/jax/triton_extensions/utils.py +++ b/transformer_engine/jax/triton_extensions/utils.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. @@ -32,8 +34,9 @@ import hashlib import os +import tempfile import warnings -from typing import Any, Callable, Mapping +from typing import Any, Callable, Mapping, Optional import zlib from packaging import version @@ -46,6 +49,7 @@ TRITON_EXTENSION_MIN_JAX_VERSION, is_triton_extension_supported, ) +from ..util import is_hip_extension # Placeholder package version on PyPI that should never be used @@ -171,6 +175,7 @@ def _check_triton_compatibility(): ) try: + import triton from jax._src.lib import gpu_triton from triton.compiler import compiler as tc from triton.backends.nvidia import compiler as cb @@ -259,23 +264,29 @@ def compile_triton( compute_capability: int, enable_fp_fusion: bool = False, ): - """Compile a Triton kernel to PTX. + """Compile a Triton or Gluon kernel to a GPU binary (PTX on CUDA, HSACO on ROCm). Kernels are cached to avoid recompilation. Args: - kernel_fn: Triton kernel function (decorated with @triton.jit) + kernel_fn: Triton (@triton.jit) or Gluon (@gluon.jit) kernel function signature: Dict mapping arg names to types (e.g., {"x_ptr": "*fp32", "n": "i32"}) constants: Dict of compile-time constants num_warps: Number of warps per block num_stages: Number of pipeline stages num_ctas: Number of CTAs (cooperative thread arrays) - compute_capability: CUDA compute capability + compute_capability: CUDA compute capability (CUDA only; ignored on ROCm, whose + target is auto-detected from the active GPU) enable_fp_fusion: Enable FP fusion optimizations (default False for accuracy) Returns: TritonKernel object for JAX """ + # Backend (CUDA/ROCm) and kernel flavor (Triton/Gluon); only the source and + # compile options differ, the resulting binary is handled the same way. + is_hip = is_hip_extension() + is_gluon = hasattr(kernel_fn, "is_gluon") and kernel_fn.is_gluon() + # Create cache key cache_key = hashlib.md5( str( @@ -288,6 +299,8 @@ def compile_triton( num_ctas, enable_fp_fusion, compute_capability, + is_hip, + is_gluon, ) ).encode() ).hexdigest() @@ -295,55 +308,102 @@ def compile_triton( if cache_key in _TRITON_KERNEL_CACHE: return _TRITON_KERNEL_CACHE[cache_key] - # Compile kernel - cuda_option_kwargs = {} - if version.parse(_TRITON_VERSION) < version.parse("3.6.0"): - cuda_option_kwargs["cluster_dims"] = (1, 1, 1) - options = cb.CUDAOptions( - num_warps=num_warps, - num_stages=num_stages, - num_ctas=num_ctas, - debug=False, - enable_fp_fusion=enable_fp_fusion, - **cuda_option_kwargs, - ) - # Mark constants as constexpr in signature signature_with_constexpr = dict(signature) for const_name in constants.keys(): if const_name in signature_with_constexpr: signature_with_constexpr[const_name] = "constexpr" - src = tc.ASTSource( - fn=kernel_fn, - constexprs=constants, - signature=signature_with_constexpr, - ) + if is_hip: + # ROCm: active GPU target (gfx arch + 64-lane warp); binary is HSACO. + from triton.compiler import make_backend + + target = triton.runtime.driver.active.get_current_target() + backend = make_backend(target) + options = backend.parse_options( + { + "num_warps": num_warps, + "num_ctas": num_ctas, + "num_stages": num_stages, + "warp_size": target.warp_size, + "enable_fp_fusion": enable_fp_fusion, + } + ) + binary_key = backend.binary_ext + else: + # CUDA path (unchanged). + cuda_option_kwargs = {} + if version.parse(_TRITON_VERSION) < version.parse("3.6.0"): + cuda_option_kwargs["cluster_dims"] = (1, 1, 1) + options = cb.CUDAOptions( + num_warps=num_warps, + num_stages=num_stages, + num_ctas=num_ctas, + debug=False, + enable_fp_fusion=enable_fp_fusion, + **cuda_option_kwargs, + ) + target = tc.GPUTarget("cuda", compute_capability, 32) + binary_key = "ptx" + + # Gluon uses GluonASTSource, which (unlike ASTSource) requires every constexpr + # parameter to be listed in the signature. + if is_gluon: + try: + from triton.experimental.gluon._runtime import GluonASTSource + except ImportError as exc: + raise ImportError( + "A Gluon kernel was passed but GluonASTSource is unavailable in this " + "Triton build (triton.experimental.gluon._runtime). Upgrade to a Triton " + "version that ships Gluon, or pass a @triton.jit kernel instead." + ) from exc + + gluon_signature = dict(signature_with_constexpr) + for name in kernel_fn.arg_names: + if name not in gluon_signature and name in constants: + gluon_signature[name] = "constexpr" + + src = GluonASTSource( + fn=kernel_fn, + constexprs=constants, + signature=gluon_signature, + ) + else: + src = tc.ASTSource( + fn=kernel_fn, + constexprs=constants, + signature=signature_with_constexpr, + ) - compiled = tc.compile( - src, - target=tc.GPUTarget("cuda", compute_capability, 32), - options=options.__dict__, - ) + compiled = tc.compile(src, target=target, options=options.__dict__) + + # TritonKernel's binary arg is a std::string: PTX is text, but HSACO is bytes + # (nanobind won't coerce), so pass HSACO as a temp-file path. The file is left + # in place because the plugin loads it at launch. + binary = compiled.asm[binary_key] + if is_hip: + fd, hsaco_path = tempfile.mkstemp(suffix=".hsaco") + with os.fdopen(fd, "wb") as f: + f.write(binary) + binary = hsaco_path - # Create kernel object for JAX - # From jax/jaxlib/gpu/triton_kernels.cc: + # TritonKernel signature per jax/jaxlib/gpu/triton_kernels.cc (changed in 0.8.2). if version.parse(jax.__version__) >= version.parse("0.8.2"): kernel = gpu_triton.TritonKernel( - compiled.name, # arg0: kernel_name (str) - num_warps, # arg1: num_warps (int) - num_ctas, # arg2: num_ctas (int) - compiled.metadata.shared, # arg3: shared_mem_bytes (int) - compiled.asm["ptx"], # arg4: ptx (str) - "", # arg5: ttir (str) - empty - compute_capability, # arg6: compute_capability (int) + compiled.name, + num_warps, + num_ctas, + compiled.metadata.shared, + binary, + "", # ttir + compute_capability, ) else: kernel = gpu_triton.TritonKernel( compiled.name, num_warps, compiled.metadata.shared, - compiled.asm["ptx"], + binary, "", # ttir compute_capability, 1, @@ -362,6 +422,8 @@ def triton_call_lowering( grid, input_output_aliases: Mapping[int, int] = None, constexprs: Mapping[str, Any] = None, + num_warps: Optional[int] = None, + num_stages: Optional[int] = None, ): """Helper for MLIR lowering that calls a Triton kernel. @@ -376,6 +438,10 @@ def triton_call_lowering( constexprs: Compile-time constants for the kernel. This includes both tl.constexpr arguments AND scalar runtime arguments (like num_tokens, strides) that are known at JAX trace time. + num_warps: Warps per block for non-autotuned kernels (required for Gluon + layouts; default 4 on ROCm, 32 on CUDA). Ignored when autotuned. + num_stages: Pipeline stages for non-autotuned kernels (default 1). Ignored + when autotuned. Returns: MLIR lowering result @@ -422,12 +488,13 @@ def lowering(ctx, x, *, block_size): else: grid_tuple = grid[:3] - # Default values for the kernel + # Resolve launch defaults (Gluon layouts need a matching num_warps). actual_kernel_fn = kernel_fn - num_warps = 32 - num_stages = ( - 1 # TODO(Phuong): consider if it is beneficial to expose num_warps, num_stages, num_ctas - ) + if num_warps is None: + # 32 warps would exceed the 1024-thread block limit on AMD's 64-lane warp. + num_warps = 4 if is_hip_extension() else 32 + if num_stages is None: + num_stages = 1 num_ctas = 1 kernel_constexprs = constexprs if constexprs is not None else {} From ad6e24599950710c06efa6139ed7622e82b3222e Mon Sep 17 00:00:00 2001 From: AllenFarcas Date: Wed, 24 Jun 2026 22:29:13 +0000 Subject: [PATCH 2/7] [Fix] Refactored and added Gluon autotuned tests --- tests/jax/test_triton_custom_calls.py | 136 ++++++++++++++---- .../jax/triton_extensions/utils.py | 20 ++- 2 files changed, 121 insertions(+), 35 deletions(-) diff --git a/tests/jax/test_triton_custom_calls.py b/tests/jax/test_triton_custom_calls.py index 55ad229a7b..7d207fb1e4 100644 --- a/tests/jax/test_triton_custom_calls.py +++ b/tests/jax/test_triton_custom_calls.py @@ -147,44 +147,44 @@ def test_triton_amax(self, shape, dtype): assert_allclose(result, expected, dtype=jnp.float32) -# Gluon binding test. Drives a @gluon.jit kernel through -# the same triton_call_lowering bridge as the Triton test above. +# 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) + @pytest.mark.triton class TestGluonBinding: """Test Gluon binding primitive through the Triton custom-call bridge.""" - # Elementwise out = x * 2. Gluon is layout-explicit: `arange` needs a - # BlockedLayout whose warps_per_cta matches the launch num_warps. - @staticmethod - @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) - # Define test primitive class DoubleGluonPrimitive(BasePrimitive): """Test primitive using a Gluon kernel.""" @@ -213,7 +213,7 @@ def lowering(ctx, x): return triton_call_lowering( ctx, - TestGluonBinding.double_kernel, + _double_kernel, x, grid=grid, constexprs={ @@ -232,7 +232,7 @@ 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)]) + @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.""" @@ -244,3 +244,77 @@ def test_gluon_double(self, shape, dtype): result = jitted_double(x) assert_allclose(result, expected, dtype=dtype) + + + @pytest.mark.triton + 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) + + 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) diff --git a/transformer_engine/jax/triton_extensions/utils.py b/transformer_engine/jax/triton_extensions/utils.py index 4d987e5f7a..99c2211f45 100644 --- a/transformer_engine/jax/triton_extensions/utils.py +++ b/transformer_engine/jax/triton_extensions/utils.py @@ -193,6 +193,17 @@ def _check_triton_compatibility(): # Triton kernel cache (module-level, shared across all kernels) _TRITON_KERNEL_CACHE = {} +# Process-scoped temp dir for ROCm HSACO blobs (removed at interpreter exit). +_HSACO_TMPDIR = None + + +def _hsaco_dir(): + """Return a process-scoped temp directory for HSACO blobs (created lazily).""" + global _HSACO_TMPDIR + if _HSACO_TMPDIR is None: + _HSACO_TMPDIR = tempfile.TemporaryDirectory(prefix="te_jax_hsaco_") + return _HSACO_TMPDIR.name + def get_triton_info(): """Get information about the installed Triton package. @@ -378,11 +389,11 @@ def compile_triton( compiled = tc.compile(src, target=target, options=options.__dict__) # TritonKernel's binary arg is a std::string: PTX is text, but HSACO is bytes - # (nanobind won't coerce), so pass HSACO as a temp-file path. The file is left - # in place because the plugin loads it at launch. + # (nanobind won't coerce), so write HSACO to a process-scoped temp dir and + # pass the path. The plugin loads it at launch; the dir is removed at exit. binary = compiled.asm[binary_key] if is_hip: - fd, hsaco_path = tempfile.mkstemp(suffix=".hsaco") + fd, hsaco_path = tempfile.mkstemp(suffix=".hsaco", dir=_hsaco_dir()) with os.fdopen(fd, "wb") as f: f.write(binary) binary = hsaco_path @@ -431,7 +442,8 @@ def triton_call_lowering( Args: ctx: MLIR lowering context - kernel_fn: Triton kernel function + kernel_fn: Triton (@triton.jit) or Gluon (@gluon.jit) kernel, optionally + wrapped with @triton.autotune *array_args: Input arrays (from ctx) grid: Grid dimensions (int or tuple) input_output_aliases: Mapping of input to output aliases From b944c51553f75a2fde407bb52f608fafdcccd4d4 Mon Sep 17 00:00:00 2001 From: AllenFarcas Date: Mon, 6 Jul 2026 15:36:45 +0000 Subject: [PATCH 3/7] [Fix] Addressed PR comments --- .../jax/triton_extensions/utils.py | 24 +++++++++++-------- 1 file changed, 14 insertions(+), 10 deletions(-) diff --git a/transformer_engine/jax/triton_extensions/utils.py b/transformer_engine/jax/triton_extensions/utils.py index 09516abdc4..3629ec735b 100644 --- a/transformer_engine/jax/triton_extensions/utils.py +++ b/transformer_engine/jax/triton_extensions/utils.py @@ -345,7 +345,7 @@ def compile_triton( signature_with_constexpr[const_name] = "constexpr" if is_hip: - # ROCm: active GPU target (gfx arch + 64-lane warp); binary is HSACO. + # ROCm: active GPU target (gfx arch + its native warp size); HSACO binary. from triton.compiler import make_backend target = triton.runtime.driver.active.get_current_target() @@ -417,16 +417,17 @@ def compile_triton( f.write(binary) binary = hsaco_path - # TritonKernel signature per jax/jaxlib/gpu/triton_kernels.cc (changed in 0.8.2). + # Create kernel object for JAX + # From jax/jaxlib/gpu/triton_kernels.cc: if version.parse(jax.__version__) >= version.parse("0.8.2"): kernel = gpu_triton.TritonKernel( - compiled.name, - num_warps, - num_ctas, - compiled.metadata.shared, - binary, - "", # ttir - compute_capability, + compiled.name, # arg0: kernel_name (str) + num_warps, # arg1: num_warps (int) + num_ctas, # arg2: num_ctas (int) + compiled.metadata.shared, # arg3: shared_mem_bytes (int) + binary, # arg4: ptx (str) + "", # arg5: ttir (str) - empty + compute_capability, # arg6: compute_capability (int) ) else: kernel = gpu_triton.TritonKernel( @@ -522,7 +523,10 @@ def lowering(ctx, x, *, block_size): # Resolve launch defaults (Gluon layouts need a matching num_warps). actual_kernel_fn = kernel_fn if num_warps is None: - # 32 warps would exceed the 1024-thread block limit on AMD's 64-lane warp. + # The upstream default of 32 assumes 32-lane warps (32*32 = 1024, the max + # threads per block). AMD wavefronts are 32 or 64 lanes, so use a smaller + # default that stays within 1024 threads either way; callers needing a + # specific count (e.g. Gluon layouts) pass num_warps explicitly. num_warps = 4 if is_hip_extension() else 32 if num_stages is None: num_stages = 1 From 58b3e3dd344e755abd9ceed65b13851b4e4b0f21 Mon Sep 17 00:00:00 2001 From: AllenFarcas Date: Wed, 12 Aug 2026 17:19:09 +0000 Subject: [PATCH 4/7] Address PR review on the ROCm JAX-Triton bridge Write HSACO to a persistent, content-addressed path. The blob path is serialized into the custom call, so a JAX persistent-cache hit in a later process replayed a deleted temp dir. Naming the file by its own digest also avoids the kernel cache key, which covers neither kernel source nor Triton version and would serve a stale binary after an edit. Import the NVIDIA Triton backend lazily. A Triton built for AMD only ships no triton/backends/nvidia, so the module-scope import failed and the ROCm path never ran. Collect the Gluon tests always and skip them with a reason. Defined inside "if HAS_GLUON" they were never collected, so the suite passed having run nothing. Narrow the probe's except for the same reason. Guard the triton import: require_triton_or_skip_test_file only checks the JAX version. Run test_triton_custom_calls.py in CI; it was in no leg. Co-Authored-By: Claude Opus 5 (1M context) --- ci/jax.sh | 1 + tests/jax/test_triton_custom_calls.py | 255 ++++++++++-------- .../jax/triton_extensions/utils.py | 55 ++-- 3 files changed, 175 insertions(+), 136 deletions(-) diff --git a/ci/jax.sh b/ci/jax.sh index dbca249c9b..6d17551d5f 100755 --- a/ci/jax.sh +++ b/ci/jax.sh @@ -67,6 +67,7 @@ run_test_config() { run_default_fa 1 test_layer.py # it effectively always uses unfused attention 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() { diff --git a/tests/jax/test_triton_custom_calls.py b/tests/jax/test_triton_custom_calls.py index 7d207fb1e4..8dbad028e1 100644 --- a/tests/jax/test_triton_custom_calls.py +++ b/tests/jax/test_triton_custom_calls.py @@ -13,6 +13,9 @@ 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 @@ -23,6 +26,8 @@ # 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 @@ -41,10 +46,16 @@ # 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) - except Exception: # pragma: no cover - Gluon or active GPU target unavailable + 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") @@ -181,140 +192,150 @@ def _double_kernel( x = gl.load(x_ptr + offsets, mask) gl.store(out_ptr + offsets, x * 2.0, mask) - @pytest.mark.triton - 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) +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 _gluon_double(x: jnp.ndarray) -> jnp.ndarray: - """Double each element using the Gluon kernel.""" - return TestGluonBinding.DoubleGluonPrimitive.outer_primitive.bind(x) + def lowering(ctx, x): + """MLIR lowering using the Gluon kernel.""" + n_elements = 1 + for dim in ctx.avals_in[0].shape: + n_elements *= dim - @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) + grid = (triton.cdiv(n_elements, BLOCK_SIZE),) - expected = (x * 2.0).astype(dtype) - jitted_double = jax.jit(self._gluon_double) - result = jitted_double(x) + 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 + ) - assert_allclose(result, expected, dtype=dtype) + 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) - @pytest.mark.triton - class TestGluonAutotunedBinding: - """Autotuned Gluon binding; exercises the TritonAutotunedKernelCall path.""" + expected = (x * 2.0).astype(dtype) + jitted_double = jax.jit(self._gluon_double) + result = jitted_double(x) - # 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( + 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) + class DoubleGluonAutotunedPrimitive(BasePrimitive): + """Test primitive using an autotuned Gluon kernel.""" + + name = "te_double_gluon_autotuned_test" + multiple_results = False + impl_static_args = () @staticmethod - def _gluon_double_autotuned(x: jnp.ndarray) -> jnp.ndarray: - """Double each element using the autotuned Gluon kernel.""" + def abstract(x_aval): + return jax.core.ShapedArray(x_aval.shape, x_aval.dtype) + + @staticmethod + def impl(x): prim = TestGluonAutotunedBinding.DoubleGluonAutotunedPrimitive - return prim.outer_primitive.bind(x) + assert prim.inner_primitive is not None + return prim.inner_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) + @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) + expected = (x * 2.0).astype(dtype) + result = jax.jit(self._gluon_double_autotuned)(x) - assert_allclose(result, expected, dtype=dtype) + assert_allclose(result, expected, dtype=dtype) diff --git a/transformer_engine/jax/triton_extensions/utils.py b/transformer_engine/jax/triton_extensions/utils.py index beb8b445dd..6813c95afc 100644 --- a/transformer_engine/jax/triton_extensions/utils.py +++ b/transformer_engine/jax/triton_extensions/utils.py @@ -194,7 +194,6 @@ def _check_triton_compatibility(): import triton from jax._src.lib import gpu_triton from triton.compiler import compiler as tc - from triton.backends.nvidia import compiler as cb from triton.runtime import autotuner except ImportError as e: raise ImportError( @@ -209,16 +208,37 @@ def _check_triton_compatibility(): # Triton kernel cache (module-level, shared across all kernels) _TRITON_KERNEL_CACHE = {} -# Process-scoped temp dir for ROCm HSACO blobs (removed at interpreter exit). -_HSACO_TMPDIR = None - def _hsaco_dir(): - """Return a process-scoped temp directory for HSACO blobs (created lazily).""" - global _HSACO_TMPDIR - if _HSACO_TMPDIR is None: - _HSACO_TMPDIR = tempfile.TemporaryDirectory(prefix="te_jax_hsaco_") - return _HSACO_TMPDIR.name + """Return the HSACO blob dir; a JAX persistent-cache hit replays the path in a later process.""" + jax_cache = os.environ.get("JAX_COMPILATION_CACHE_DIR") + if jax_cache: + root = os.path.join(jax_cache, "te_jax_hsaco") + else: + home = os.path.expanduser("~") + if not home or home == "/" or home.startswith("~"): + home = tempfile.gettempdir() + xdg = os.environ.get("XDG_CACHE_HOME") or os.path.join(home, ".cache") + root = os.path.join(xdg, "te_jax_hsaco") + root = os.path.abspath(root) + os.makedirs(root, exist_ok=True) + return root + + +def _write_hsaco(binary): + """Write HSACO to a content-addressed path; the kernel cache key can go stale.""" + path = os.path.join(_hsaco_dir(), f"{hashlib.sha256(binary).hexdigest()}.hsaco") + if os.path.exists(path): + return path + fd, tmp_path = tempfile.mkstemp(suffix=".hsaco.tmp", dir=os.path.dirname(path)) + try: + with os.fdopen(fd, "wb") as f: + f.write(binary) + os.replace(tmp_path, path) + finally: + if os.path.exists(tmp_path): + os.unlink(tmp_path) + return path def get_triton_info(): @@ -364,7 +384,9 @@ def compile_triton( ) binary_key = backend.binary_ext else: - # CUDA path (unchanged). + # Triton built for AMD only does not ship triton/backends/nvidia. + from triton.backends.nvidia import compiler as cb + cuda_option_kwargs = {} if version.parse(_TRITON_VERSION) < version.parse("3.6.0"): cuda_option_kwargs["cluster_dims"] = (1, 1, 1) @@ -410,15 +432,10 @@ def compile_triton( compiled = tc.compile(src, target=target, options=options.__dict__) - # TritonKernel's binary arg is a std::string: PTX is text, but HSACO is bytes - # (nanobind won't coerce), so write HSACO to a process-scoped temp dir and - # pass the path. The plugin loads it at launch; the dir is removed at exit. + # HSACO is bytes and nanobind won't coerce it to the std::string binary arg, so pass a path. binary = compiled.asm[binary_key] if is_hip: - fd, hsaco_path = tempfile.mkstemp(suffix=".hsaco", dir=_hsaco_dir()) - with os.fdopen(fd, "wb") as f: - f.write(binary) - binary = hsaco_path + binary = _write_hsaco(binary) # Create kernel object for JAX # From jax/jaxlib/gpu/triton_kernels.cc: @@ -483,8 +500,8 @@ def triton_call_lowering( tl.constexpr arguments AND scalar runtime arguments (like num_tokens, strides) that are known at JAX trace time. num_warps: Warps per block for non-autotuned kernels (required for Gluon - layouts; default 4 on ROCm, 32 on CUDA). Ignored when autotuned. - num_stages: Pipeline stages for non-autotuned kernels (default 1). Ignored + layouts; default 4). Ignored when autotuned. + num_stages: Pipeline stages for non-autotuned kernels (default 3). Ignored when autotuned. Returns: From 8cf848984e081f929dd7d583e599615a0d1af06c Mon Sep 17 00:00:00 2001 From: AllenFarcas Date: Wed, 12 Aug 2026 18:48:53 +0000 Subject: [PATCH 5/7] Keep HSACO blobs in the process-scoped temp dir The persistent content-addressed store was written for a failure that does not occur. Measured on gfx950: the ROCm plugin unlinks the blob once it has loaded it, and JAX derives its persistent-cache key from the lowered HLO, so lowering reruns in every process and rewrites the file before launch. A cached executable therefore never replays a stale path. Persisting the blobs only leaked the ones no kernel loaded, which the temp dir had been reclaiming. Co-Authored-By: Claude Opus 5 (1M context) --- .../jax/triton_extensions/utils.py | 43 ++++++------------- 1 file changed, 13 insertions(+), 30 deletions(-) diff --git a/transformer_engine/jax/triton_extensions/utils.py b/transformer_engine/jax/triton_extensions/utils.py index 6813c95afc..6ffcb1a3c5 100644 --- a/transformer_engine/jax/triton_extensions/utils.py +++ b/transformer_engine/jax/triton_extensions/utils.py @@ -208,37 +208,16 @@ def _check_triton_compatibility(): # Triton kernel cache (module-level, shared across all kernels) _TRITON_KERNEL_CACHE = {} +# Process-scoped temp dir for ROCm HSACO blobs (removed at interpreter exit). +_HSACO_TMPDIR = None + def _hsaco_dir(): - """Return the HSACO blob dir; a JAX persistent-cache hit replays the path in a later process.""" - jax_cache = os.environ.get("JAX_COMPILATION_CACHE_DIR") - if jax_cache: - root = os.path.join(jax_cache, "te_jax_hsaco") - else: - home = os.path.expanduser("~") - if not home or home == "/" or home.startswith("~"): - home = tempfile.gettempdir() - xdg = os.environ.get("XDG_CACHE_HOME") or os.path.join(home, ".cache") - root = os.path.join(xdg, "te_jax_hsaco") - root = os.path.abspath(root) - os.makedirs(root, exist_ok=True) - return root - - -def _write_hsaco(binary): - """Write HSACO to a content-addressed path; the kernel cache key can go stale.""" - path = os.path.join(_hsaco_dir(), f"{hashlib.sha256(binary).hexdigest()}.hsaco") - if os.path.exists(path): - return path - fd, tmp_path = tempfile.mkstemp(suffix=".hsaco.tmp", dir=os.path.dirname(path)) - try: - with os.fdopen(fd, "wb") as f: - f.write(binary) - os.replace(tmp_path, path) - finally: - if os.path.exists(tmp_path): - os.unlink(tmp_path) - return path + """Return a process-scoped temp directory for HSACO blobs (created lazily).""" + global _HSACO_TMPDIR + if _HSACO_TMPDIR is None: + _HSACO_TMPDIR = tempfile.TemporaryDirectory(prefix="te_jax_hsaco_") + return _HSACO_TMPDIR.name def get_triton_info(): @@ -433,9 +412,13 @@ def compile_triton( compiled = tc.compile(src, target=target, options=options.__dict__) # HSACO is bytes and nanobind won't coerce it to the std::string binary arg, so pass a path. + # The plugin unlinks the file once loaded; lowering reruns per process, so it is rewritten. binary = compiled.asm[binary_key] if is_hip: - binary = _write_hsaco(binary) + fd, hsaco_path = tempfile.mkstemp(suffix=".hsaco", dir=_hsaco_dir()) + with os.fdopen(fd, "wb") as f: + f.write(binary) + binary = hsaco_path # Create kernel object for JAX # From jax/jaxlib/gpu/triton_kernels.cc: From 727f47be56580866c7e98ac00c0c92892e94a662 Mon Sep 17 00:00:00 2001 From: AllenFarcas Date: Mon, 17 Aug 2026 14:57:20 +0000 Subject: [PATCH 6/7] Adjusted comment --- transformer_engine/jax/triton_extensions/utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/transformer_engine/jax/triton_extensions/utils.py b/transformer_engine/jax/triton_extensions/utils.py index 6ffcb1a3c5..515e7d4b75 100644 --- a/transformer_engine/jax/triton_extensions/utils.py +++ b/transformer_engine/jax/triton_extensions/utils.py @@ -412,7 +412,7 @@ def compile_triton( compiled = tc.compile(src, target=target, options=options.__dict__) # HSACO is bytes and nanobind won't coerce it to the std::string binary arg, so pass a path. - # The plugin unlinks the file once loaded; lowering reruns per process, so it is rewritten. + # The plugin reads the file once and unlinks it, then serves the blob from its own cache. binary = compiled.asm[binary_key] if is_hip: fd, hsaco_path = tempfile.mkstemp(suffix=".hsaco", dir=_hsaco_dir()) From 16cfe56a85c79ec9ce23b79454c4030094d62f18 Mon Sep 17 00:00:00 2001 From: AllenFarcas Date: Mon, 17 Aug 2026 21:22:14 +0000 Subject: [PATCH 7/7] Run the permutation tests in CI and prewarm ROCm Triton kernels The bridge's own tests covered a synthetic kernel; the code that ships through it, transformer_engine.jax.permutation, ran in no CI leg. Wiring those two files in exposed a launch failure with multiple local devices: INTERNAL: /tmp/te_jax_hsaco_/.hsaco; No such file or directory The plugin's HIP branch reads the blob from its path and unlinks it before taking the lock guarding its module cache, so two device threads reaching an unseen kernel together race and the loser opens a file the winner deleted. Launching each freshly compiled kernel once from a single device inserts the cache entry before any sharded executable can race for it. Measured on the distributed test at L1: 3 failures in 3 runs without it, none in 3 runs with it. NVTE_JAX_TRITON_PREWARM=0 disables it; the whole block goes away once the plugin does its read and unlink under the writer lock. The prewarm reuses the lowering closure rather than rebuilding the call, so the kernel cache key cannot drift, and it zero-fills operands because the gather kernels index with them. Adds the L1 parameter sets both files lacked, without which the multi-GPU leg fails collection. Co-Authored-By: Claude Opus 5 (1M context) --- ci/jax.sh | 2 + tests/jax/test_distributed_permutation.py | 5 ++ tests/jax/test_permutation.py | 7 ++ .../jax/triton_extensions/utils.py | 68 ++++++++++++++++++- 4 files changed, 80 insertions(+), 2 deletions(-) diff --git a/ci/jax.sh b/ci/jax.sh index 6d17551d5f..a67765f000 100755 --- a/ci/jax.sh +++ b/ci/jax.sh @@ -65,6 +65,7 @@ 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 @@ -102,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 diff --git a/tests/jax/test_distributed_permutation.py b/tests/jax/test_distributed_permutation.py index ee7a56a7ec..ad31b8a290 100644 --- a/tests/jax/test_distributed_permutation.py +++ b/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,6 +89,7 @@ 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, } @@ -93,6 +97,7 @@ def _inject_permutation(request): ALL_DTYPES = [jnp.float32, jnp.bfloat16] DTYPES = { "L0": [jnp.float32], + "L1": [jnp.float32], "L2": ALL_DTYPES, } diff --git a/tests/jax/test_permutation.py b/tests/jax/test_permutation.py index 38fbee18e3..c77b8e6c15 100644 --- a/tests/jax/test_permutation.py +++ b/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, } diff --git a/transformer_engine/jax/triton_extensions/utils.py b/transformer_engine/jax/triton_extensions/utils.py index 515e7d4b75..daa2f4d8fd 100644 --- a/transformer_engine/jax/triton_extensions/utils.py +++ b/transformer_engine/jax/triton_extensions/utils.py @@ -35,6 +35,13 @@ (jax-ml/jax#35218) instead of silently falling back to non-autotuned dispatch. Useful for CI or debugging to ensure autotuning is active. Default is "0" (silent compatibility fallback). + NVTE_JAX_TRITON_PREWARM: ROCm only. If set to "0", do not launch a freshly + compiled kernel once on a single device before returning from lowering. + The prewarm works around a race in the ROCm plugin, which unlinks the + HSACO it was handed before locking its module cache, so two device + threads reaching an unseen kernel together can make one of them fail + with ENOENT. Default is "1" (prewarm runs when more than one local + device is visible). """ import hashlib @@ -211,6 +218,10 @@ def _check_triton_compatibility(): # Process-scoped temp dir for ROCm HSACO blobs (removed at interpreter exit). _HSACO_TMPDIR = None +# Single-device prewarm state (ROCm only; see _prewarm_kernel_call). +_PREWARM_PRIM = None +_PREWARM_ACTIVE = False + def _hsaco_dir(): """Return a process-scoped temp directory for HSACO blobs (created lazily).""" @@ -220,6 +231,50 @@ def _hsaco_dir(): return _HSACO_TMPDIR.name +def _prewarm_primitive(): + """Primitive that re-emits an already-built Triton custom call for prewarming.""" + global _PREWARM_PRIM + if _PREWARM_PRIM is None: + from jax.extend.core import Primitive + from jax.interpreters import mlir + + prim = Primitive("te_triton_prewarm") + prim.multiple_results = True + prim.def_abstract_eval(lambda *_, out_avals, rule: out_avals) + mlir.register_lowering(prim, lambda ctx, *args, out_avals, rule: rule(ctx, *args)) + _PREWARM_PRIM = prim + return _PREWARM_PRIM + + +def _prewarm_kernel_call(ctx, rule): + """Load a freshly compiled HSACO into the ROCm plugin from a single device. + + The plugin reads the blob from its path and unlinks it before taking the lock that + guards its module cache, so two device threads reaching an unseen kernel together + race: the winner deletes the file, the loser fails with ENOENT. Launching once here, + on one device, inserts the cache entry before any sharded executable can race for it. + """ + global _PREWARM_ACTIVE + if _PREWARM_ACTIVE: + return + _PREWARM_ACTIVE = True + try: + prim = _prewarm_primitive() + device = jax.local_devices()[0] + operands = [ + jax.device_put(jnp.zeros(aval.shape, aval.dtype), device) for aval in ctx.avals_in + ] + out_avals = tuple(core.ShapedArray(a.shape, a.dtype) for a in ctx.avals_out) + jax.block_until_ready( + jax.jit(lambda *xs: prim.bind(*xs, out_avals=out_avals, rule=rule))(*operands) + ) + except Exception as e: # pylint: disable=broad-except + # Prewarming is an optimization; a failure here must not break the real compile. + warnings.warn(f"Triton kernel prewarm failed ({type(e).__name__}: {e})", RuntimeWarning) + finally: + _PREWARM_ACTIVE = False + + def get_triton_info(): """Get information about the installed Triton package. @@ -411,8 +466,8 @@ def compile_triton( compiled = tc.compile(src, target=target, options=options.__dict__) - # HSACO is bytes and nanobind won't coerce it to the std::string binary arg, so pass a path. - # The plugin reads the file once and unlinks it, then serves the blob from its own cache. + # The plugin's HIP branch takes a filename, not the blob; it reads the file once, unlinks it, + # and serves every later launch from its own cache. binary = compiled.asm[binary_key] if is_hip: fd, hsaco_path = tempfile.mkstemp(suffix=".hsaco", dir=_hsaco_dir()) @@ -567,6 +622,7 @@ def _normalize_grid(grid_tuple): # single non-autotuned dispatch for compatibility. Set # NVTE_JAX_ENFORCE_TRITON_AUTOTUNING=1 to raise an error instead, prompting the # user to upgrade JAX for improved performance. + kernels_before = len(_TRITON_KERNEL_CACHE) is_autotuned = isinstance(kernel_fn, autotuner.Autotuner) used_autotuned_launch = False if is_autotuned and not is_triton_autotuned_alias_safe(): @@ -731,4 +787,12 @@ def _normalize_grid(grid_tuple): operand_output_aliases=ffi_operand_output_aliases, ) + if ( + is_hip_extension() + and len(_TRITON_KERNEL_CACHE) > kernels_before + and len(jax.local_devices()) > 1 + and os.environ.get("NVTE_JAX_TRITON_PREWARM", "1") != "0" + ): + _prewarm_kernel_call(ctx, rule) + return rule(ctx, *array_args)