Repository navigation
fix(rocm): reuse staged moe alignment on Triton 3.8 - #691
Merged
Merged
Conversation
jason-fxz
reviewed
Oct 15, 2026
Comment on lines
+93
to
+95
| # Avoid the ROCm Triton 3.8 histogram-mask miscompile until 3.9. | ||
| # https://github.com/triton-lang/triton/pull/11246 | ||
| if Version(triton.__version__) < Version("3.9"): |
Collaborator
There was a problem hiding this comment.
The miscompile is 3.8-only: tensor cmpi was not folded before triton#9805, and triton#11577 confirms 3.7.1 is unaffected. pyproject.toml still allows triton>=3.6, so < 3.9 also sends 3.6/3.7 to the slower staged path. Could you gate on 3.8.x only and mark the branch as temporary?
Suggested change
| # Avoid the ROCm Triton 3.8 histogram-mask miscompile until 3.9. | |
| # https://github.com/triton-lang/triton/pull/11246 | |
| if Version(triton.__version__) < Version("3.9"): | |
| # FIXME: remove once Triton 3.8 is unsupported. Its AMD backend folds compares on tl.histogram results to false. | |
| # https://github.com/triton-lang/triton/pull/11246 | |
| if Version(triton.__version__).release[:2] == (3, 8): |
LZ-QWQ
changed the title
zhangnju
added a commit
to zhangnju/FreeToken
that referenced
this pull request
Oct 15, 2026
…Cm compile) Slimmed from the original torch2.14/triton3.8 enablement now that upstream main landed its own ROCm support (setup.py IS_ROCM, native gpu_sync_hip.cpp, torch 2.14, FlashML-org#691 moe_align). These three are the RDNA fixes main still lacks: - qsa/attend.py: query the device's shared_memory_per_block and shrink the KV tile when it won't fit, instead of assuming the GB300's 228KB LDS -> RDNA (64KB/CU) no longer overflows LDS. ROCm-gated; NVIDIA tiles byte-identical. - gguf/dispatch.h: ROCm's __shfl_xor_sync needs a 64-bit warp mask (amd_warp_sync_functions static_assert); the vLLM/sgl donor passes uint32. Widen to uint64_t on HIP. - gguf.py: --expt-relaxed-constexpr and -ccbin are nvcc-only; ROCm compiles hipified sources with clang++ which rejects them.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
On ROCm with Triton 3.8, the existing histogram align can be miscompiled, producing invalid routing buffers that may cause incorrect MoE outputs or illegal GPU memory accesses. The existing grouped-MoE test in
test_fused_moe.pyreproduces an illegal memory access before this fix. The compiler fix is Triton #11246.This PR proposes a simple workaround: reuse the existing moe_align_block_size_triton on ROCm with Triton 3.8.x. Other platform/version dispatch remains unchanged. This temporary workaround can be removed once Triton 3.8 is no longer supported.
Related: #541, #587, #687. Both existing proposals add Torch fallbacks; this alternative reuses the staged Triton implementation without adding another implementation or configuration option.
Performance
Radeon AI PRO R9700 (
gfx1201), ROCm 10.1, Torch 2.14. ActualQwen/Qwen3.8-Flash-Next-FP8prefill routes from all 48 layers,K=10,E=512,block=64.Cells report mean ± SD (us/call) across 48 layers. Each layer averages all its chunks, using the median of three independent run medians per chunk (10 warmups, 30 samples/run, 5 graph replays/sample).
±describes variation across the measured layer values, not a confidence interval.M128 has one chunk per layer. The 8192-token requests provide all eight M1024 or both M4096 chunks per layer. The 3.9 preview uses a ROCm 10.2 wheel in the same Torch/ROCm environment and is an experimental reference, not an end-to-end speedup or a supported framework upgrade.
The staged fallback has mixed performance against Torch on 3.8. The purpose is a minimal correctness workaround while preserving the substantially faster histogram path on fixed compilers. Full Triton 3.9 migration remains separate because other kernels use its removed block-pointer API.
Validation
Tests are unchanged. The existing grouped-MoE test fails before the fix and passes afterward on ROCm with Triton 3.8.
cc @zihaomu