[AMD][DSv4] Fuse the DSv4 FP4 indexer prefill-schedule preamble into one kernel - #37764
Merged
HaiShaw merged 3 commits intoSep 4, 2026
Merged
Conversation
Between the C4 K-cache write and the indexer logits kernel, every C4 layer dispatched a 33-kernel preamble to build the persistent-grid schedule: the page-table pad (new_zeros + masked copy_), a row_to_batch arange, a local_starts zero-fill, AITER's ~27-op compute_prefill_schedule, its cta_info kernel, and the logits -inf pre-fill. Each of those touches a few hundred elements, so the block is pure launch latency -- on a 128k/1k conc-64 DP8TP8 MTP trace it measured 146us per layer per step, 11.5% of all GPU kernel time, to schedule an 8.6us logits kernel. Everything the preamble computes is a reduction or a scan over the row count, so it collapses into one Triton kernel; the only real dependency is that cta_info needs the finished prefix sums. The one step that needed rethinking is AITER's search for the smallest split factor s with sum_i ceil(chunks_i / s) <= P, which materialises an [s_max, T] tensor: that sum is monotonically non-increasing in s, so a binary search over s returns the identical answer in registers. Pinning cta_info also lets the logits kernel skip its -inf pre-fill, the same way the workspace path already does -- the length-aware top-k only reads [0, c4_seq_len), which is always covered by a CTA. Decode / MTP target-verify now costs 2 dispatches and 14.2us per layer instead of 33 and 145.8us. Rows past MAX_FUSED_ROWS (one row per lane) keep AITER's preamble: a prefill that wide is compute-bound and does not care about ~30 extra launches. Verified bit-identical against AITER's compute_prefill_schedule over 144 shape/length cases (cta_info, safe, padded page table) and end-to-end over the logits every row's top-k reads. GSM8k 1319q DP8TP8+MTP+FP4 indexer: 0.940/0.920 before, 0.948/0.931 after. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
heachary
marked this pull request as ready for review
September 3, 2026 12:14
heachary
requested review from
BBuf,
DarkSharpness,
HaiShaw,
HydraQYH,
celve and
yuan-luo
as code owners
September 3, 2026 12:14
One conflict, in `_guard_page_table`: main extracted the padding rule into `_guarded_pages()`, this branch replaced the body with the fused `pad_page_table` dispatch. Kept the fused dispatch and pointed `padded_page_table_shape()` at main's `_guarded_pages()` so the rule has a single definition. The two sides compose rather than collide. Main's pooled prefill logits block is safe precisely because a pinned `cta_info` makes the kernel skip its -inf pre-fill and the length-aware top-k reads only `[0, c4_seq_len)`; this branch extends that same pinning to the no-workspace fallback, so the pooled block stays sound there too. Main's new row chunking (`logits_rows_per_chunk`) calls the indexer with row subsets, which makes the prefill workspace stale by row count and routes every chunk through that fallback -- now two dispatches instead of 33. Re-verified on gfx950 against AITER's `compute_prefill_schedule`: 144 shape/length cases bit-identical (`cta_info`, `safe`, padded page table), and 5 end-to-end cases with bit-identical logits over every row's `[0, c4_seq_len)`. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
This was referenced Sep 3, 2026
StevenChenSE
pushed a commit
to StevenChenSE/sglang
that referenced
this pull request
Sep 6, 2026
This was referenced Sep 6, 2026
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
3 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.
Motivation
The FP4 indexer (#37353) builds its persistent-grid schedule with a torch-op
preamble that runs once per C4 layer. Every op touches only a few hundred
elements, so the whole block is launch latency bound.
On a DeepSeek-V4-Pro 128k/1k conc-64 DP8TP8 MTP trace (MI355X) that is
33 dispatches and 146 us per layer per step.
These 33 kernels can be fused into a single kernel.
Trace results
Before
After
The trace shows all 33 kernels being fused into a single kernel.
Summary of changes
Adds
python/sglang/kernels/ops/attention/dsv4/fp4_indexer_schedule_hip.py,which collapses the preamble into a single Triton kernel. Everything it
computes is a reduction or a scan over the row count; the only real dependency
is that AITER's
_prefill_cta_info_kernelneeds the finished prefix sums, sothe schedule now costs two dispatches instead of 33.
swithsum_i ceil(chunks_i / s) <= Pby materialising an[s_max, T]tensor andextracting the first feasible index arithmetically (~18 of the 27 ops). That
sum is monotonically non-increasing in
s, so a binary search oversreturns the identical answer in registers, with no intermediate tensor.
destination no longer needs pre-zeroing.
row_to_batch/local_startsare emitted in-kernel � sglang alwaysschedules one row per query token over its whole window.
cta_infolets the logits kernel skip its-infpre-fill, thesame way the existing workspace path already does: the length-aware top-k
only reads
[0, c4_seq_len), which is always covered by a CTA.MAX_FUSED_ROWS(one row per lane) keep AITER's preamble � aprefill that wide is compute-bound and does not care about ~30 extra launches.
No new flags; active whenever
--enable-deepseek-v4-fp4-indexeris on (HIP).Accuracy
GSM8k (1319 questions, DP8TP8 + MTP + FP4 indexer): 0.940 / 0.920 before, 0.948 / 0.931 after.
Performance
Fixed length
isl/osl = 8k1k, conc=4/64, mtp=on, image = lmsysorg/sglang-rocm:v0.5.18-rocm720-mi35x-20260902
The launch latency bound kernels are a fixed cost and the benefits of fusing them are more pronounced at lower latencies.
Agent-X
conc64, mtp=on, image = lmsysorg/sglang-rocm:v0.5.18-rocm720-mi35x-20260902
~3% improvement in decode step time
CI States
Latest PR Test (Base): ⏳ Run #33844863732
Latest PR Test (Extra): ❌ Run #33844863368
Latest PR Test (AMD ROCm 7.2): ⏳ Run #33844863663