[Feat] Added JAX-Triton bridge for ROCm

[AllenFarcas]

Description

Extend TE's JAX Triton custom-call bridge to compile and dispatch AMD ROCm (HSACO) and Gluon kernels. This PR enables AMD's layout-explicit Gluon kernels to be called from JAX, mirroring NVIDIA's existing support.

Fixes https://github.com/ROCm/frameworks-internal/issues/16044

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Changes

Please list the changes introduced in this PR:

• Add ROCm/HIP backend path emitting HSACO alongside the existing CUDA/PTX path.
• Support @gluon.jit kernels via GluonASTSource with a full constexpr-marked signature.
• Pass HSACO as a temp-file path (nanobind std::string), not raw bytes.
• Add optional num_warps/num_stages for non-autotuned Gluon layout matching.
• Add version-guarded Gluon binding test to test_triton_custom_calls.py (requires Triton base > 3.4.0).

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

[ipanfilo]

ROCm arch does not necessary have 64-lane wavefront

[ipanfilo]

Keep upstream comments as-is

[ipanfilo]

Why 4?

[AllenFarcas]

I picked 4 because that's Triton's own default num_warps, it is what triton.Config() and unspecified @triton.jit launch use.

[ipanfilo]

Add it to comment

[github-actions]

Claude Walkthrough

Intent. Extend TE's JAX/Triton custom-call bridge in transformer_engine/jax/triton_extensions/utils.py so it also compiles and launches ROCm kernels (emitting HSACO instead of PTX) and accepts Gluon (@gluon.jit) sources alongside plain @triton.jit. This unlocks AMD's layout-explicit Gluon kernels through the same JAX lowering path NVIDIA already uses.

Key changes.

• Dual-backend compile in compile_triton at transformer_engine/jax/triton_extensions/utils.py:294 — ROCm branch uses triton.runtime.driver.active.get_current_target() + make_backend(...).parse_options(...) and the CUDA branch keeps the existing CUDAOptions/GPUTarget("cuda", ...) path.
• Gluon source path: when kernel_fn.is_gluon() is true, swaps tc.ASTSource for triton.experimental.gluon._runtime.GluonASTSource and, because Gluon requires every constexpr in the signature, back-fills any missing constants (utils.py:369).
• HSACO delivery via temp file: because TritonKernel's binary field is a std::string and nanobind won't coerce raw bytes, HSACO is written to a process-scoped TemporaryDirectory (_hsaco_dir, utils.py:216) and the path is passed instead of the blob (utils.py:410).
• triton_call_lowering now takes optional num_warps / num_stages (utils.py:453) so callers matching a Gluon BlockedLayout can pin the launch geometry; default num_warps becomes 4 on ROCm (wavefronts are 32 or 64 lanes, so the CUDA-tuned default of 32 would exceed 1024 threads/block).
• Cache key now folds in is_hip and is_gluon so CUDA/ROCm and Triton/Gluon variants don't collide in _TRITON_KERNEL_CACHE.
• Adds two Gluon test classes to tests/jax/test_triton_custom_calls.py behind a HAS_GLUON guard keyed on Triton > 3.4.0.

Walkthrough.

transformer_engine/jax/triton_extensions/utils.py — the compile function was previously CUDA-only (cb.CUDAOptions → tc.compile(..., target=GPUTarget("cuda", ...)) → compiled.asm["ptx"]). It's now split by is_hip_extension(). On ROCm the target is discovered from the active driver so warp_size (32 or 64 depending on gfx arch) can be threaded into parse_options; binary_key becomes backend.binary_ext (i.e. "hsaco"). Source construction is orthogonal to backend: Gluon detection (hasattr(kernel_fn, "is_gluon") and kernel_fn.is_gluon()) selects GluonASTSource, which unlike ASTSource requires every constexpr — including those already handled as constants — to appear in the signature, hence the explicit back-fill loop over kernel_fn.arg_names. The GluonASTSource import is wrapped in a try/except ImportError that raises a targeted message pointing users at a Triton upgrade. After tc.compile, the ROCm branch persists HSACO bytes to a temp file (_HSACO_TMPDIR, created lazily and torn down at interpreter exit) and swaps the path in where PTX text used to go, so both TritonKernel(...) signatures downstream (the pre- and post-3.6 shapes) work unchanged. triton_call_lowering grows two optional launch knobs; if unset, num_warps defaults to 4 on ROCm and 32 on CUDA, matching each backend's "warps × warp_size ≤ 1024 threads" ceiling.

tests/jax/test_triton_custom_calls.py — adds a Gluon _double_kernel that doubles its input using an explicit BlockedLayout whose warps_per_cta must match the launch num_warps. Two test classes exercise the bridge: TestGluonBinding for the JIT path (passes num_warps=NUM_WARPS) and TestGluonAutotunedBinding for the autotune path (each triton.Config carries BLOCK_SIZE/NUM_WARPS as constexpr kwargs plus a matching num_warps=, so the layout stays consistent across configs). Both are gated by HAS_GLUON, computed by parsing triton.__version__ with packaging.Version and comparing .release against (3, 4, 0) — using .release rather than raw string comparison so a +rocm.git local version segment can't spuriously satisfy the check. WARP_SIZE is read from the active Triton target, keeping the test agnostic to gfx942 vs. gfx950 vs. NVIDIA.

Testing. New TestGluonBinding (JIT) and TestGluonAutotunedBinding (autotune) classes cover the Gluon end-to-end path; each parametrizes on shapes (1024, 1024) and (1000, 1000) (the latter exercises the trailing masked block) and checks bitwise-close equality against x * 2.0. Both are silently skipped when Gluon is missing or the Triton base version is <= 3.4.0. Existing TestTritonBinding cases are unchanged, so the CUDA/PTX path continues to be exercised.

Notes for reviewers.

• The HSACO temp file is not deleted eagerly — _HSACO_TMPDIR lives for the process lifetime. Deliberate (the plugin loads the file at launch time, potentially long after compile_triton returns), but long-running processes accumulate one file per unique cache key.
• Gluon detection uses duck typing (hasattr(kernel_fn, "is_gluon")); autotuned kernels forward is_gluon to the underlying JITFunction, so autotuned Gluon kernels are correctly identified too.
• The ROCm num_warps default changed from 32 to 4. Existing ROCm callers that relied on the implicit 32 will now launch with fewer warps unless they pass num_warps explicitly — worth flagging in the CHANGELOG if any downstream users depend on the old value.
• compute_capability is still passed through to TritonKernel(...) on ROCm even though it's meaningless there; the JAX-side plugin ignores it for HIP targets, but future cleanup could carry the ROCm arch string instead.
• Gluon on ROCm requires an upstream Triton make_ir fix that landed after 3.4.0 — the HAS_GLUON guard encodes that requirement.

---

Generated by Claude. To request a code review, comment /claude review.

[github-actions]

Claude review

Reviewed the PR-specific diff (git diff <merge-base>..HEAD, two files):

• transformer_engine/jax/triton_extensions/utils.py
• tests/jax/test_triton_custom_calls.py

Verdict: looks clean. The ROCm/HSACO path, Gluon GluonASTSource handling, and the temp-file HSACO plumbing all read correctly, and the PR-level changes address the earlier reviewer feedback in-place (warp-size comment reworded, num_warps=4 default justified in a code comment).

Copyright headers: OK — both modified files carry an AMD 2026 line above the preserved NVIDIA 2022-2026 line, matching the existing repo convention (This file was modified for portability to AMDGPU prefix + AMD copyright + untouched NVIDIA copyright).

No new inline findings.

[wenchenvincent]

@AllenFarcas Please resolve conflict and test against the latest CI image.

[github-actions]

The HSACO path ends up inside the serialized custom call, which makes the compiled executable non-portable across processes.

binary flows into TritonKernel → TritonKernelCall.to_proto() → zlib.compress(call_proto) → the backend_config of the triton_kernel_call custom call (utils.py:665-679). On CUDA that blob contains the PTX text, so the executable is self-contained. On ROCm it now contains /tmp/te_jax_hsaco_<rand>/tmp<rand>.hsaco, and _HSACO_TMPDIR is deleted when the interpreter exits.

So if JAX's persistent compilation cache is enabled (JAX_COMPILATION_CACHE_DIR — ci/jax.sh:73 sets it for the multi-GPU run), a cache hit in a later process replays an executable pointing at a temp dir that no longer exists, and the kernel fails to load at launch. Same issue for anything else that outlives the process that compiled it.

Writing to a stable, content-addressed location instead of mkstemp would fix it and would also avoid re-emitting an identical blob on every process start — cache_key is already a hash of everything that determines the binary, so e.g. <cache_root>/<cache_key>.hsaco under a persistent dir works, with the write done atomically (temp file + os.replace) to stay safe under concurrent processes.

[github-actions]

cb is now only used on the CUDA branch (utils.py:368), but this import is still unconditional at module scope. A Triton built with the AMD backend only (TRITON_CODEGEN_BACKENDS=amd) does not ship triton/backends/nvidia/, so this raises ImportError and gets rewritten by the handler below into "Triton is required … Install with: pip install triton" — a misleading message, and the ROCm path this PR adds never gets a chance to run.

Since make_backend is already imported lazily inside the is_hip branch, doing the same for cb in the else branch would make the module import cleanly on ROCm-only Triton builds regardless of how the container's Triton was configured.

[github-actions]

Two consecutive blank lines here, but this is inside the if HAS_GLUON: block (depth 1). Black only permits two blank lines at module level — below that the limit is one — so black --line-length=100 --preview (.pre-commit-config.yaml:31) will rewrite this and the lint gate will fail as-is.

Suggested change

[github-actions]

This combination makes a Gluon failure indistinguishable from a Gluon pass in CI.

except Exception swallows everything — a genuine breakage in triton.experimental.gluon, a driver that fails to initialise, a warp_size attribute rename — and sets HAS_GLUON = False. Because the test classes are then defined inside if HAS_GLUON: (:152), they are never collected, so the report shows no failures and no skipped entries either. The suite goes green having exercised nothing of what this PR adds.

Defining the classes unconditionally and gating with @pytest.mark.skipif(not HAS_GLUON, reason=...) would keep the same behaviour while making the non-execution visible in the junit XML. Narrowing the catch to (ImportError, AttributeError, RuntimeError) would also stop it from masking unexpected failures.

Related: tests/jax/test_triton_custom_calls.py is not in ci/jax.sh::run_test_config, so neither the existing Triton test nor these new Gluon tests run in ROCm CI at all. Given that ROCm/Gluon enablement is the whole point of the PR, adding run_default_fa 1 test_triton_custom_calls.py there is what would actually demonstrate it works — and it lines up with the "test against the latest CI image" ask.

[github-actions]

Claude review (re-review)

Re-reviewed the full PR diff against the current base (HEAD^1...HEAD^2, two files: transformer_engine/jax/triton_extensions/utils.py, tests/jax/test_triton_custom_calls.py). The PR head is unchanged since the last review, but the base has moved, so this pass re-read the whole diff rather than an incremental slice — it surfaced four findings the earlier pass missed.

Verdict: the ROCm/HSACO and Gluon compile paths are structurally sound and the CUDA path stays guarded and unchanged. Four issues worth addressing before merge, posted inline: one ROCm correctness risk around how the HSACO binary is delivered, one import-time fragility on ROCm-only Triton builds, one Black violation that will fail the lint gate, and one test-visibility gap that lets the new Gluon coverage silently not run.

Copyright headers: OK — both modified files carry an AMD 2026 line above the preserved NVIDIA 2022-2026 line, with the AMDGPU portability marker.

Note: the PR is currently in a conflicting state against dev.