[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

[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]

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.