Repository navigation
[Perf] [ROCM] diffusion: Bit-exact fast QK RMSNorm and fused RoPE for MiniMax-H3 - #47569
Merged
Merged
Conversation
ROCm keeps H3's per-head QK norm on aten rms_norm, whose vectorized_layer_norm_kernel runs one 256-thread block per 128-wide row (~0.6 TB/s) after a .contiguous() copy of each strided q/k view. Add a Triton kernel that replays that kernel's arithmetic (4 contiguous squares per lane, shfl-down fold 16/8/4/2/1, rsqrtf, w * (rstd * x)) and reads the fused-qkv views directly, behind a first-sight BitExactFusionGate. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…MSNorm Apply the bf16 NeoX rotary_embedding (each product and the sum rounded to bf16, dims past rot_dim passed through) in the QK norm's pass instead of a second read-modify-write of q and k, behind its own first-sight gate. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
There was a problem hiding this comment.
Claude Code Review
This repository is configured for manual code reviews. Comment @claude review for a one-time review, or @claude review always to subscribe this PR to a review on every future push.
Tip: disable this comment in your organization's Code Review settings.
mickqian
requested review from
AgainstEntropy,
BBuf,
HaiShaw,
kevin-mii,
niehen6174,
ping1jing2,
yichiche and
yingluosanqian
as code owners
October 31, 2026 02:16
github-actions
Bot
added
documentation
This comment has been minimized.
This comment has been minimized.
Collaborator
Author
|
/tag-and-rerun-ci |
mickqian
changed the title
niehen6174
added a commit
to niehen6174/sglang
that referenced
this pull request
Nov 5, 2026
Reviewed the six upstream minimax_h3.py changes since the last pin (sgl-project#46047, sgl-project#46709, sgl-project#47567, sgl-project#47569, sgl-project#47566 and the BCG output buffers). None touches what h3_comfyui_v1 relies on in its single-GPU eager torch_sdpa profile: DiTBlock.forward is unchanged, _attend still calls fused_inplace_qknorm_rope on the _use_fused_qknorm_rope branch with the same arguments (the new pipelined branch needs Ulysses, which the profile rejects), and _embed still returns (embeddings, t_emb) in eager mode.
Yuxingwang-intel
pushed a commit
to Yuxingwang-intel/sglang
that referenced
this pull request
Nov 6, 2026
… MiniMax-H3 (sgl-project#47569) Co-authored-by: Mick Qian <mickqian@users.noreply.github.com> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
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
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.
Motivation
On ROCm, H3 keeps the per-head QK RMSNorm on aten
rms_norm. Itsvectorized_layer_norm_kernelruns one 256-thread block per 128-wide row (~0.6 TB/s) after a.contiguous()copy of each strided q/k view of the fused QKV: QK norm plus that copy is 5.1% of a denoise step (43.7 ms/step) at 1280x704 on 4xMI355X.Modifications
aten_head_rmsnorm_triton.pythat replays aten's arithmetic (4 contiguous squares per lane, shfl-down fold 16/8/4/2/1, rsqrtf,w * (rstd * x)) and reads the strided q/k views directly. It is used only after a first-sightBitExactFusionGatecheck against the eager path passes.test/registered/kernels/ops/diffusion/test_norm.py(fail on main, pass here).Accuracy Tests
Bitwise: MP4 md5 identical to main for every prompt checked (22/22 in the final 20-prompt run, which includes the other ROCm branches), plus the kernel unit test comparing against aten bit for bit.
Speed Tests and Profiling
Setup for all numbers below: 4x MI355X, MiniMax-H3 T2VA, 1280x704 (SemiAnalysis "720p"), 124 frames, 49 forwards, Ulysses 4 + aiter, single stream. Each arm is a fresh server (1 warmup + 3 timed requests), ABAB across two rounds, base = main
75a5da7f.QK norm per step: 43.7 ms -> 7.1 ms.
Checklist
🤖 Generated with Claude Code
CI States
Latest PR Test (Base): ✅ Run #41159226768
Latest PR Test (Extra): ❌ Run #41159226703
Latest PR Test (AMD ROCm 10): ⏳ Run #41159226745