[Feat] Added JAX-Triton bridge for ROCm - #649
Conversation
AllenFarcas
added
ci-level 3
Claude WalkthroughIntent. Extend TE's JAX/Triton custom-call bridge in Key changes.
Walkthrough.
Testing. New Notes for reviewers.
Generated by Claude. To request a code review, comment |
Claude reviewReviewed the PR-specific diff (
Verdict: looks clean. The ROCm/HSACO path, Gluon Copyright headers: OK — both modified files carry an AMD 2026 line above the preserved NVIDIA No new inline findings. |
|
@AllenFarcas Please resolve conflict and test against the latest CI image. |
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
Claude review (re-review)Re-reviewed the full PR diff against the current base ( 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 Note: the PR is currently in a conflicting state against |
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) <noreply@anthropic.com>
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) <noreply@anthropic.com>
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_<rand>/<rand>.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) <noreply@anthropic.com>
The prewarm guarded itself with a single _PREWARM_ACTIVE flag and fired only when the kernel cache had grown during this lowering. Both fail once more than one thread lowers at a time: the flag is a reentrancy guard rather than mutual exclusion, so a thread lowering one kernel suppressed the load for another, and a thread that lost the race to populate the cache saw no growth and skipped the load for the kernel it was about to launch. Key the load on a digest of the serialized call proto, record loaded keys in a process-wide set, and hold a per-key lock across the launch so a second lowering of the same kernel waits for it rather than racing the plugin. The key is recorded in a finally: a persistent failure would never succeed on retry and would repeat its warning once per lowering. _hsaco_dir and _prewarm_primitive had the same unguarded lazy-init shape; two threads could build two temp dirs, whichever lost taking the winner's blobs with it at finalization. Both now initialize under a lock, and the locks are replaced in a forked child, where one held at fork would never be released. The child detaches the inherited TemporaryDirectory before dropping it, so the parent's blobs survive, and clears the loaded set, since the plugin's module cache does not cross the fork. Rebuild the FFI lowering inside the rule. ffi_lowering captured an ir.StringAttr belonging to whichever MLIR context was current when the rule was built, and the prewarm replays that rule in a context of its own. Call _require_triton_permutation from the token_dispatch, token_combine and sort_chunks_by_index forward rules. The guard existed but nothing invoked it, so a Triton-less install reached the kernel wrappers with the names bound to None. Widen both import guards to RuntimeError and ValueError, which a mismatched Triton raises at import, and chain the original cause onto the message. Import triton inside the ROCm branch of compile_triton, its only use, which also clears the two pylint messages the module-scope import was adding. Measured with the prewarm on: the permutation suite passes 24 of 24 at L1 in each of three runs and 72 of 72 at L2; with it off, all three L1 runs fail. A four-thread concurrent lowering is clean in five runs. Eight threads on one kernel produce a single load and no early return; on eight distinct kernels they do not serialize. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
wangye805
left a comment
wangye805
left a comment
There was a problem hiding this comment.
LGTM with minor comments
Import the NVIDIA Triton backend at module scope again, under `if not is_hip_extension()`. Moving it into compile_triton kept a Triton built for AMD only importable, but it also deferred the failure a CUDA build gets from a broken backend to the first lowering; the guard keeps the import-time error while leaving the AMD case alone. pylint cannot see that the name is bound wherever it is read, hence the suppression at the use site. Register the fork handler only on ROCm, so a CUDA build no longer installs a callback for state it never creates. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
4 participants
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
Changes
Please list the changes introduced in this PR:
GluonASTSourcewith a full constexpr-marked signature.nanobind std::string), not raw bytes.num_warps/num_stagesfor non-autotuned Gluon layout matching.Checklist: