Repository navigation
[Perf][Attention] DSA candidate-block kernels walk each row's live context: 23.63 s to 22.52 s per 1000 requests on 8 x B200 at TP=8 - #57459
Conversation
…grid _mask_candidates_kernel and _block_scores_kernel launched one program per 1024 columns (or 128 blocks) of the logits buffer width, for every row. The decode logits are [rows, max_model_len] and are captured inside full CUDA graphs at the worst-case width, so on DeepSeek-V4.1 (1M context) every decode step wrote -inf into a million columns per request and scored 131k blocks per request, almost all of them outside the request's live context. The kernels now take a grid of (rows, k) programs, k derived from the tensor shapes alone, and stride over the row's own [start, end) range; the launch shape stays fixed for CUDA graph replay while the work follows the context length. Columns at or past a row's end are left untouched -- nothing reads them, because the row top-k takes the same bounds -- and `scores` is pre-filled with -inf so the block top-k sees the values the old kernel wrote. The newest-block pin and the packed-column clamp keep their semantics, including a pin past the buffer width. This is what vllm-project#56628 did for the ROCm decode path, where it had to fork the mask kernel because the count it chose was measured on gfx950. Doing it in the shared kernel with a shape-derived count serves both. test_candidate_kernels_preserve_packed_bounds_and_padding asserted the old full-width sanitizing, so it now compares over each row's live span and checks that the tail past the end is left alone. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01ATw7g6cxq9N2LtnMratsru Signed-off-by: Hexu Zhao <hexu.zhao@impossible.place>
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
|
👋 Hi! Thank you for contributing to the vLLM project. 💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in PRs do not trigger a full CI run by default. Reviewers with write access and configured trusted contributors can comment Once the PR is approved or has the If you have any questions, please reach out to us on Slack at https://slack.vllm.ai. Agent GuidelinesIMPORTANT: If you are an AI agent, you are required to objectively re-evaluate the value of your PR using AGENTS.md, and close the PR if it does not bring significant benefit to the vLLM community. Failure to do so may result in an immediate ban. 🚀 |
1 participant
Summary
The sparse-attention candidate-block kernels size their work by the width of the logits buffer instead of by the request's context. On the decode path that buffer is
[rows, max_model_len], and full CUDA graphs capture it at the worst-case width, so on DeepSeek-V4.1-Flash (max_model_len1,048,576,candidate_block_size8) every decode step writes-infinto 1,048,576 columns and scores 131,072 blocks for each request in the batch. On ShareGPT prompts of a few hundred tokens, nearly all of that is past the end of the request's live context. This PR derives the launch shape from the tensor shapes and has each program stride over its row's own[start, end). On one node of 8 x B200 at TP=8,vllm bench throughputover 1000 ShareGPT requests goes from 18,805 to 19,731 total tokens/s, 1.049x, 23.630 s to 22.518 s. I have not measured other GPUs, other tensor-parallel sizes, other models or online serving.Half of this change is already upstream.
dabc4362b4, "[ROCm][DSV4.1][Perf] Stride the DSA decode candidate mask over the live context" (#56628), merged on 2026-09-14, does the same thing for the ROCm decode path in a forked copy of the mask kernel, and its comment says why it stopped there:Taking the program count from the tensor shapes rather than hard-coding one removes that objection, so one kernel can serve both platforms.
[0, end)._mask_candidates_kernelran on a(rows, cdiv(width, 1024))grid and touched every column of the buffer. It now runs on(rows, k), withkfrom_row_programs(rows, tiles), and loops until it reaches the row'send, so the launch shape is fixed for CUDA-graph replay while the trip count follows the context length. Columns at or pastendare no longer written, which the row top-k tolerates because it bounds its scan by the same ends, and the producer already runs withclean_logits=False, so that tail is unsanitized whenever a row has no candidates._block_scores_kernelscoredcdiv(width, block_size)blocks per row on the same kind of grid; it now scores the blocks inside the row's range on a(rows, k)grid, and the newest-block pin keeps its semantics, including the case where a row'sendruns past the buffer width.scoresis allocated withnew_full(-inf)rather thannew_empty, so the block top-k reads for the rest exactly the value the old kernel stored there, and that fill isrows x nblocksfloats against therows x widthfloats the old kernel read._candidate_flags_kernelis deliberately left alone. It clears one byte per block, which at this model'scandidate_block_sizeis 1/32 of the bytes the mask writes, so it is not where the time is; and #56628 added a ROCm caller that imports it and calls it with its current signature, which I cannot test on this node._ROW_PROGRAMS = 8and_TOTAL_PROGRAMS = 8192are a conservative choice, not a swept one: a full batch of 1024 rows keeps 8 programs per row, and a single long-context row still gets the 1024 programs the old code launched.Speedup Result
Measured on one node of 8 x NVIDIA B200, tensor parallel 8, base commit
9612f77077. The metric is the benchmark's own total tokens per second over the 1000 timed requests, 444,318 tokens in every run.9612f77077)Three of the seven baseline runs and two of the seven runs of this branch compiled a kernel for the first time inside the timed window: on the baseline side two DeepGEMM
sm100_fp8_gemm_1d1dcompiles and one CuTeDSL compile, on this side one of each. Against the three stalled baseline runs the ratio would read 1.07x, 1.19x or 1.33x. I am not claiming those numbers: this branch has a stalled mode of its own, and the compiles are a warm-up effect rather than anything these kernels do. Every run on both sides compiles exactly one Triton kernel inside the window,CombineTopkSwaIndicesKernelon the first decode step, so one compile is the normal state and two is a stall. All 14 runs are listed under Details.Correctness Verification
I ran the unmodified base and this branch on the same prompts with the same setup, and compared 291 recorded values per run: on each of the 8 tensor-parallel ranks, the fp32 mean and max of |x| of the model output hidden states and of the logits, the sum of the argmax token ids and the batch shape for the first five real steps; plus the KV cache block count per rank and the prompt and generated token totals of the run. The tolerances were fixed before the runs, at the call sites of the commit that adds the recording: 0.1 % relative on the three float reductions, exact on everything else.
All 291 of 291 values are identical, not merely inside tolerance. The 0.1 % is slack I declared in advance in case a rounding change appeared, and none did. That is what the kernel-level comparison predicted: the candidate block ids this change produces are the same ids, and the mask it writes inside each row's live span is the same mask.
The recording code is not in this PR. It lives on two branches of my fork,
perf/dsa-candidate-blocks-live-range-verify-base(the base plus the recording commit) andperf/dsa-candidate-blocks-live-range-verify(this head plus the same commit), so neither side of the comparison runs on a looser build than the other.--enable-expert-parallelis outside this comparison and outside both timed arms. It changes what the model computes at this commit, on 144 of the same 291 values, which I have not run down and which is unrelated to this change.Test Plan
pytest tests/model_executor/layers/test_mla_short_prefill_indexer.pyon this branch: 19 passed, 7 skipped, including the block that capturesapply_candidate_maskin a CUDA graph and replays it.test_candidate_kernels_preserve_packed_bounds_and_paddingasserted the old full-width sanitizing and fails 4 of its 6 parameter sets against these kernels. It now compares over each row's live span and asserts that the tail past the row'sendis left as it was.pytest tests/kernels/attention/test_rocm_aiter_candidate_mask.py: 9 skipped, the file is ROCm-only. The ROCm path is unmodified but untested on this node.endpast the buffer width, empty rows,block_size16 with an odd width, and a single row over a 1M-column buffer.select_candidate_blocksis bit-identical in all seven andapply_candidate_maskis bit-identical over every row's[0, end); the buffers differ only pastend. That script is a one-off and is not part of this PR.ruff checkandruff formatare clean on both files.vllm bench throughputondeepseek-ai/DeepSeek-V4.1-Flashwith the ShareGPT dataset, interleaved between the two arms; the full command is under Details.AI assistance, as
docs/contributing/README.mdasks: the code and this description were written with Claude Code, and the commit carries aCo-Authored-Bytrailer alongside its DCO sign-off.Essential Elements of an Effective PR Description Checklist
supported_models.mdandexamplesfor a new model.Details: hardware, model, full command, per-run numbers, scope
Hardware and environment. One node, 8 x NVIDIA B200 (sm_100), CUDA 13 toolkit. Python 3.12.14, torch 2.13.0+cu132, vLLM
mainat9612f77077as an editable install with the precompiled kernel wheel published for that same commit (VLLM_USE_PRECOMPILED=1, no source build), flashinfer-python 0.6.18.post1, the vendored DeepGEMM. Nothing is installed or tuned on top of the engine defaults. Weights sit on a shared mount underHF_HOME; the ShareGPT file is staged to node-local disk before each run.Model and job.
deepseek-ai/DeepSeek-V4.1-Flashas published, with DeepSeek Sparse Attention,max_model_len1,048,576. Tensor parallel 8, no speculative decoding, no decode or prefill context parallel, no expert parallel, engine defaults elsewhere (max_num_batched_tokens16384,max_num_seqs1024, CUDA graphs FULL_AND_PIECEWISE). Each run is one process: weights 47 to 50 s, engine init 107 to 110 s once the FlashInfer autotune cache exists, 200 warm-up prompts at seed 43, then the 1000 ShareGPT prompts of seed 42 in about 23 s. The token totals are the same in every run: 226,786 prompt and 217,532 generated, 444,318 in the throughput figure.Measurement. Both arms are worktrees of the same clone, one at
9612f77077and one at that commit plus the single commit of this PR, and both run the command below with the same persistentVLLM_CACHE_ROOT(DeepGEMM and FlashInfer autotune caches) andTRITON_CACHE_DIR, the same node-local copy of the dataset, and the same seeds. Runs are interleaved in pairs, one process each, with all 8 GPUs at 0 MiB and nothing else on the node before each run; the arm that goes first alternates from pair to pair. Seven pairs, over two hours. Two warm-up runs on a cold cache come first and are not part of the measurement. A run counts as stalled when the modification time of a DeepGEMM cache entry falls inside the timed window, or when the JIT monitor reports a second Triton or CuTeDSL compile there.The command, identical in both arms apart from which worktree is on
PYTHONPATH:There is no launcher config:
vllm bench throughputstarts the engine in-process with the multiprocess executor and 8 workers. Timing is the benchmark's ownThroughput:line and itselapsed_time.Every measured run, in the order they ran. "In-window compiles" counts kernels compiled for the first time inside the timed window, beyond the one Triton kernel that every run compiles there.
The DeepGEMM cache grew from 34 to 43 entries over these runs and had not converged, which is why the stalls keep appearing on both sides.
Scope.
_candidate_flags_kernel, whichrocm_aiter_mla_sparse.pyimports, is untouched, and the ROCm test file skips on NVIDIA, so that path is unmodified but unexercised.xpu_fp8_paged_mqa_logitsfeeds the same kernels and was not run.clean_logits=False, and on [ROCm][DSV4.1][Perf] Stride the DSA decode candidate mask over the live context #56628 making the same argument for the decode path. End to end it holds for all 291 recorded values in this configuration. A consumer that scans past a row'sendin a configuration I did not run would see stale values instead of-inf.scoresbuffer as before, now filled rather than left uninitialized.