Repository navigation
[Perf][DCP] Sparse indexer: key only significant bits in the DCP top-k merge - #64129
Open
NickLucche wants to merge 2 commits into
Open
NickLucche wants to merge 2 commits into
NickLucche wants to merge 2 commits into
Conversation
…k merge Since the DeepGEMM bump in vllm-project#61149 the SM100 MQA logits are BF16 and are upcast to FP32 for the DCP top-k merge, so the low 16 bits of every candidate score are zero. The all-gather selector radix-sorts the 64-bit key score << 32 | ~token_id in 11-bit passes. BF16 scores almost always tie at the top-k boundary, so the radix then walks the zero score bits and the constant high token-id bits, and each of those passes funnels every remaining candidate into one histogram bin. Pack only the significant bits at the top of the key: the high 16 score bits when the logits were BF16, and the id bits that positions below max_model_len (max_pool_len for kpool indexers) need. The selected set is unchanged; FP32 logits keep 32 score bits. Final radix passes can now be narrower than the 512-thread block, so the bin scan covers whole block slices. Signed-off-by: NickLucche <nicolo.lucchesi@mistral.ai>
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
Member
Author
Member
Author
|
/ci run |
|
✅ Triggered Buildkite CI #97780 for commit ℹ️ This PR is 39 commits behind upstream |
Contributor
|
This pull request has merge conflicts that must be resolved before it can be |
2 of 6 tasks
This branch has not been deployed
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.
Overview
Since #61149 (DeepGEMM 2026-09-30), the SM100 MQA logits of the sparse indexer are BF16. Under DCP they are upcast to FP32 before the top-k, so the low 16 bits of every candidate score that reaches the DCP top-k merge are zero. The merge selector,
StableTopKFromGatheredCandidatesKernel, radix-sorts the 64-bit keyscore << 32 | ~token_idin 11-bit passes. BF16 scores almost always tie at the top-k boundary. The radix then also walks the 16 zero score bits and the constant high bits of the token id, and each of those passes funnels every remaining candidate into a single histogram bin.This PR packs only the bits that carry information into the key: the 16 high score bits when the logits were BF16, and the token-id bits that positions below
max_model_lenneed. The selected set is unchanged.Claims
Validation
Serving: GLM-5.3 (
zai-org/GLM-5.3, FP8)Setup:
a2a);--attention-backend FLASHMLA_SPARSE --kv-cache-dtype fp8_ds_mla;--max-model-len 65536 --max-num-batched-tokens 8192 --max-num-seqs 64 --no-enable-prefix-caching;-cc.cudagraph_mode=FULL_DECODE_ONLY: decode steps replay full CUDA graphs. With the defaultFULL_AND_PIECEWISE,mainitself did not start on this setup: the ranks diverged during PIECEWISE capture. That is unrelated to this PR.vllm bench serve --dataset-name random --random-input-len 60000 --random-output-len 1024 --ignore-eos, with 6 prompts at c1, 12 at c2 and 24 at c8.Method:
main, PR, main, PRandPR, main, PR, main.mainis 803bc13.mainand 13.43-13.46 ms with this PR.mainand 19.04-19.21 ms with this PR.GLM-5.3 has 21 sparse-indexer layers per step (
index_topk_freq=4). Models with an indexer in every layer, like DeepSeek-V3.2, or DCP8 with 2x the candidates, save proportionally more per step.Kernel time in the model (torch profiler)
Same setup, 60k-token prompts, 8192 candidates per row (4 ranks × top-2048). The table gives the median per call on ranks 0 and 2 (they agree within 0.1 us). Each decode trace covers 29-40 steps (609-840 selector calls).
StableTopKFromGatheredCandidatesKernel), decode c1 (1 row)cooperative_topk(control), c1 / c8The per-step saving (119 us at c1, 143 us at c8) matches the measured TPOT/ITL gain. The prefill-chunk saving is where the ~1% lower c1 TTFT comes from: every prefill chunk runs the merge in each of the 21 indexer layers.
Microbenchmark
All-gather selector alone on one GB300, top-2048, ids drawn from a 64k context (16 id bits). us per call, CUDA graph of 100 launches. "Wide" scores are N(0, 3²). "Narrow" scores are 60 + N(0, 4²): they are more concentrated, so more of them tie at BF16 precision. "set == ref" checks the selection against an FP64 reference.
FP32 logits keep 32 score bits and gain a little from the shorter id field.
Correctness
main's selector (64-bit key) on the same gathered candidates after every DCP merge and compared the selected sets. It ran in eager mode, so it could sync.main's selection.zai-org/GLM-5.3-Flash, kpool indexer, top-512 pools out of 2048 candidates), TP4 DCP4 with--cp-kv-cache-interleave-size 64, covered the long greedy prompts and 4 random 60k-token requests at c8.tests/evals/gsm8k/gsm8k_eval.py, 1319 questions, 5-shot, completions, max concurrency 64), same setup as above:Greedy (
temperature=0), sequential at c1, 38 prompts: 34 short ones (128 tokens each) and 4 real-text prompts of 8k, 20k, 38k and 51k tokens (256 tokens each). This setup is not deterministic run to run, even onmain:main, the same server queried twice: 11/38 and 16/38 outputs differ;mainacross server restarts: 14/38 and 15/38;main: 11/38, 13/38 and 16/38.Most of the differing prompts are short ones. There every valid candidate is selected, so this PR cannot change their result. The outputs diverge after a common prefix, as numerical noise does. So greedy identity cannot separate the two here. The check on real model data above can.
Tests
tests/v1/attention/test_indexer_dcp_localize.py(GB300)pre-commit(ruff-check, ruff-format, typos, mypy, ...)New test
test_cutedsl_gathered_select_compact_keys_match_reference[16384|131072|1048576]:num_positions - 1;score_bits=16,id_bits=dcp_topk_id_bits(num_positions)) both select the FP64 reference set.The server logs show no CuTeDSL JIT compilation during inference with this PR (
jit_monitor_mode=warn), on GLM-5.3 and on GLM-5.3-Flash (kpool,max_model_len // 4id bits). So the new warm-up keys match the runtime keys.Details
Root cause.
fp8_fp4_mqa_logits,fp8_fp4_paged_mqa_logits) BF16.sparse_attn_indexerupcasts them (as_fp32_logits) whenever DCP is on.StableTopKFromGatheredCandidatesKernelbuildsscore_key << 32 | ~token_idand runs up to 6 radix passes (11 bits each, the last one 9). It stops at the first pass whose threshold bin fits.max_model_len(16 bits for a 64k context);Fix.
CompileKeygainsscore_bitsandid_bits.score_key >> (32 - score_bits)followed by the lowid_bitsof~token_id, packed at the top of the 64-bit word. With BF16 logits and a 64k context, that is 32 bits and 3 passes.~id & maskpreserves the order of ids below2**id_bits.sparse_attn_indexer(DSA) and the GLM-5.3-Flash kpool indexer pass the following:score_bits=16when the raw logits are BF16;id_bitsfrommax_model_len, ormax_pool_lenfor kpool pool ids.score_bitsfollows the platform. That keeps it on the warmed-up key. 32 bits would also select correctly.max_model_len // index_kpool.num_binsstay zero.Scope.
score_bitscomes from the logits dtype and the platform, andid_bitsfrommax_model_len. All three are host-side metadata.pytest v1/attention), which has CuTeDSL.Why this is not a duplicate
gh pr list --repo vllm-project/vllm --state open --searchwith "dcp topk", "top-k merge", "indexer dcp", "radix select bf16", "61149", "StableTopKFromGatheredCandidates", "dcp_indexer_cutedsl", "dcp top-k" and "sparse indexer bf16 logits" found no PR that changes the merge selector's key:_merge_dcp_topk_global. The changes are independent, and a rebase conflict would be trivial.AI assistance
AI assistance was used to investigate the regression, write the change and its test, and run the benchmarks and evals above. The submitter reviewed every changed line and the measurements.
Pull Request Checklist
I used vLLM's
/pr-checklistskill. (Mandatory for agents, optional for humans).AI assistance was used during the creation of this PR.
Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.
Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI.
Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.
Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.