diff --git a/tests/jax/test_triton_custom_calls.py b/tests/jax/test_triton_custom_calls.py index 846d26a41..7d207fb1e 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,176 @@ 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) + + @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) + + @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 + 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 2a86321c3..3629ec735 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. @@ -37,8 +39,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 @@ -53,6 +56,7 @@ is_triton_autotuned_alias_safe, is_triton_extension_supported, ) +from ..util import is_hip_extension # Placeholder package version on PyPI that should never be used @@ -184,6 +188,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 @@ -201,6 +206,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. @@ -278,23 +294,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( @@ -307,6 +329,8 @@ def compile_triton( num_ctas, enable_fp_fusion, compute_capability, + is_hip, + is_gluon, ) ).encode() ).hexdigest() @@ -314,36 +338,84 @@ 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 + its native warp size); HSACO binary. + 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" - compiled = tc.compile( - src, - target=tc.GPUTarget("cuda", compute_capability, 32), - options=options.__dict__, - ) + # 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=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. + 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 # Create kernel object for JAX # From jax/jaxlib/gpu/triton_kernels.cc: @@ -353,7 +425,7 @@ def compile_triton( 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) + binary, # arg4: ptx (str) "", # arg5: ttir (str) - empty compute_capability, # arg6: compute_capability (int) ) @@ -362,7 +434,7 @@ def compile_triton( compiled.name, num_warps, compiled.metadata.shared, - compiled.asm["ptx"], + binary, "", # ttir compute_capability, 1, @@ -381,6 +453,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. @@ -388,13 +462,18 @@ 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 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 @@ -441,12 +520,16 @@ 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: + # 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 num_ctas = 1 kernel_constexprs = constexprs if constexprs is not None else {}