[Perf][GLM] Fuse the kpool tail slot mapping into one Triton kernel - #57534
Merged
ZJY0516 merged 3 commits intoSep 20, 2026
Merged
Conversation
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
ZJY0516
reviewed
Sep 18, 2026
| tl.store(out_ptr + offs, tl.load(slot_mapping_ptr + offs, mask=mask), mask=mask) | ||
|
|
||
|
|
||
| def compute_kpool_tail_slot_mapping( |
Member
There was a problem hiding this comment.
why we don't fuse the entire operation?
Semantic rebase of PR #11 onto current upstream main: the sparse-indexer decode expansion and token_to_req_indices parts have already landed upstream (13e221f, _indexer_decode_metadata_kernel/_token_request_mapping_kernel). The remaining delta is the kpool tail slot mapping, which still ran an arange/searchsorted/index_select/remainder chain per step. Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: Jared Wen <jaredwen@inferact.ai>
JaredforReal
force-pushed
the
perf/sparse-indexer-decode-metadata
branch
from
September 18, 2026 10:15
c610955 to
731c67f
Compare
ZJY0516
approved these changes
Sep 19, 2026
|
✅ @JaredforReal, CI is now available for this PR.
|
Contributor
Author
|
/ci run |
|
❌ This PR is 18 commits behind upstream |
Contributor
Author
|
/ci run |
|
✅ Triggered Buildkite CI #90030 for commit |
Morrowmake
pushed a commit
to Morrowmake/vllm-cmp170hx
that referenced
this pull request
Sep 22, 2026
Upstream vllm-project#57534 (1b9fa3e) fused the kpool tail slot mapping into a single Triton kernel inside compute_kpool_tail_slot_mapping(), with a stronger contract than ours: it takes request query-start/end boundaries (extending the last interval to the actual token count), shares the launch with the optional tail slot-mapping copy, and bounds work by the actual token count. Our VLLM_GLM5_PROLOGUE_FUSE_KPOOL path predates that and is now strictly weaker, so drop it: the early-return hook in mla/indexer.py, the kernel, kpool_tail_slot_mapping_ref/kpool_tail_slot_mapping, the `kpool` settings field and its _KPOOL parser entry, and the tests of the removed private implementation. Upstream's kpool tests are untouched. The Mamba block-table and GDN speculative-metadata families keep working exactly as before behind VLLM_GLM5_PROLOGUE_FUSE{,_MAMBA_BT,_GDN,_DEBUG}.
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.
Purpose
GLM-5.3-Flash sparse-MLA models keep a one-block-per-request circular tail cache (kpool). Every decode step — twice per step with MTP, once for the target model and once for the drafter —
KpoolTailMetadataBuilder.build()maps each token to its request's tail ring with a chain of small torch ops (arange+searchsorted+clamp+index_select+remainder+copy_): 12 kernel launches and ~250 us of CPU enqueue time per build.This PR replaces the chain with a single Triton kernel driven directly by
query_start_loc, including the CUDA-graph padding rows. The CPU-tensor path keeps the torch implementation, and the emitted values are bit-identical to it.With #57317 the generic
_compute_slot_mappings_kernelno longer runs forKpoolTailSpec, socompute_kpool_tail_slot_mappingis now the only slot-mapping path for the tail group — this fusion covers all of it. It composes with #57477 (the seed-kernel stride fix): this PR only produces the logical slot ids (block * kpool + pos % kpool) that consumers decompose with// kpool/% kpool; physical addressing of the padded tail view is untouched, and the kernel reads the block table through its realstride(0).Not a duplicate: #56562 covers the sparse-indexer decode expansion and the token→request mapping (an earlier revision of this PR had similar pieces and was rescoped); #57161 reworks the kpool compress kernel and #57458 the NoPE query packing — neither touches the tail slot mapping.
Test Plan
E2E setup: 4x GB300 (SM103), TP4,
zai-org/GLM-5.3-Flash(FP8,flashinfer_trtllmMoE),--attention-backend FLASHINFER_MLA_SPARSE --max-model-len 69632 --max-num-seqs 256 --max-num-batched-tokens 16384 --no-enable-prefix-caching --speculative-config '{"method":"mtp","num_speculative_tokens":1}', V2 model runner;vllm bench serve --dataset-name random --random-input-len 8192 --random-output-len 1024 --random-range-ratio 0 --ignore-eos, two runs per build, base =42a33039eb.Test Result
test_kpool_tail_slot_mapping.py: 17 passed — 4 new CUDA cases check the Triton path against the torch reference, including tokens past the last request boundary, CUDA-graph padding rows, and a non-contiguous (strided-view) block table.test_sparse_indexer_decode_seq_lens.py+test_indexer_native_next_n.py: 16 passed.Micro (
compute_kpool_tail_slot_mapping, kpool=4, GB300; CPU = enqueue time without sync, GPU = profiler device time):E2E (8k in / 1k out, output tok/s, two runs per build):
Median TPOT deltas are within noise at every point (-0.6% to +2.7%). The c=1 throughput ordering was consistent across both pairs of runs (and c=1 mean TPOT was 4.19 ms on both PR runs vs 4.37/4.47 ms on main), but with n=2 per build this sits at the edge of the noise floor. The durable claim is the removed per-step CPU work (~0.44 ms/step with MTP: two builds x ~220 us) and 11 fewer launches per build, which async scheduling largely overlaps on this hardware.
AI assistance (Claude) was used for the profiling, the implementation and the benchmark harness; every changed line was reviewed and the tests/benchmarks above were run by the submitter.