[GDN] [Backport] optimize: avoid fla l2-norm recompilation by token count (#31558) - #27
Open
sammysun0711 wants to merge 2 commits into
Conversation
…#31558) Signed-off-by: Xiake Sun <xiake.sun@amd.com>
sammysun0711
force-pushed
the
cherry-pick/gdn-l2norm-dynamic-token-count
branch
from
August 29, 2026 14:44
e0e9343 to
bf765a0
Compare
Reduce GDN prefill launches for benchmarked BF16/FP32 GPU shapes while preserving the original path on unsupported platforms, dtypes, layouts, and launch-bound wide shapes. Signed-off-by: Xiake Sun <xiake.sun@amd.com>
There was a problem hiding this comment.
🟡 Changes recommended
The benchmark underreports memory traffic, and the fused implementation retains an unreachable kernel path.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Pull request overview
Backports runtime token-count handling and fused AMD Q/K L2 normalization to reduce Triton recompilation and launch overhead.
Changes:
- Makes token counts runtime Triton arguments.
- Adds fused HIP Q/K normalization with fallback dispatch.
- Adds correctness, dispatch, and benchmark coverage.
File summaries
| File | Description |
|---|---|
python/sglang/srt/layers/attention/fla/l2norm.py |
Updates and fuses L2Norm kernels. |
python/sglang/srt/layers/attention/fla/chunk.py |
Adds HIP fused dispatch. |
test/registered/jit/test_fused_gdn_l2norm.py |
Tests correctness and dispatch. |
test/registered/jit/benchmark/bench_gdn_l2norm.py |
Benchmarks fused versus separate normalization. |
Review details
- Files reviewed: 4/4 changed files
- Comments generated: 3
- Review effort level: Balanced
💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
| replay_iters=1200, | ||
| graph_clone_args=(0, 1), | ||
| memory_args=(q, k), | ||
| memory_output=None, |
Comment on lines
145
to
+146
| else: | ||
| fused_l2norm_qk_kernel1[(tokens,)]( | ||
| fused_l2norm_qk_kernel1[(max(TQ, TK),)]( |
Comment on lines
+41
to
+45
| @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32]) | ||
| def test_fused_l2norm_qk_generic_d256(dtype): | ||
| torch.manual_seed(99) | ||
| q = torch.randn(257, 8, 256, dtype=dtype, device="cuda") | ||
| k = torch.randn(257, 8, 256, dtype=dtype, device="cuda") |
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
Backports upstream SGLang PR31558, PR33290 to the
Qwen3.5 v0.5.15branch.Previous
Qwen3.5_v0.5.9branch used different, fused, already-warmed kernel fused_l2norm_qkCurrent
qwen3_5_v0.5.15_dflashbranch call two l2norm_fwd for q and k.Modifications
The change makes the L2Norm token count a runtime argument, preventing Triton from compiling a separate kernel for every unseen prefill length.
Accuracy Tests
Speed Tests and Profiling
Qwen3.5 397B PTPC FP8 TP8 Prefill 8K in/500 out, con=1 on MI308X, PTL:
Checklist
Review and Merge Process
/tag-and-rerun-ci,/tag-run-ci-label,/rerun-failed-ci