[Bugfix][DSA] Reserve the decode indexer logits in the memory profile by gcanlin · Pull Request #62860 · vllm-project/vllm

gcanlin · 2026-10-21T08:54:29Z

Overview

The DSA indexer's decode logits are not counted by the memory profiler, so eager and piecewise steps can OOM at runtime. This PR reserves them in the profiling run, the way DeepInfer's profile run does.

OOM

DeepSeek-V4.1-Flash, 4×B200 DEP4 (InferenceX recipe, --max-num-seqs 256, --max-model-len 1048576, 5 spec tokens, --gpu-memory-utilization 0.97), AgentX at concurrency 512. Rank 0 died 5 minutes into the load:

File "vllm/models/deepseek_v41/nvidia/model.py", line 1086, in forward
    hidden_states, pre_mix, *late_aux = self.decoder_replay_layers(
...
File "vllm/model_executor/layers/sparse_attn_indexer.py", line 893, in sparse_attn_indexer
    logits = as_fp32_logits(logits, seq_lens.view(-1))
File "vllm/model_executor/layers/sparse_attn_indexer.py", line 312, in as_fp32_logits
    out = torch.empty_strided(
torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 3.07 GiB. GPU 0 has a total capacity of 178.35 GiB of which 3.04 GiB is free.

Why

The indexer's decode path allocates the paged MQA logits on every call. DeepGEMM returns [decode rows, max_model_len], in BF16 on SM100.
Every top-k backend except deep_select then reads an FP32 copy made by as_fp32_logits. Here that copy alone is 786 rows (131 decodes × 6) × 1,049,088 columns × 4 B = 3.07 GiB.
The profiling run reserves only VLLM_SPARSE_INDEXER_MAX_LOGITS_MB (512 MiB, the prefill chunk budget), so the KV cache takes this memory.
FULL decode graphs keep the logits in their graph pool. Mixed batches, batches above the capture sizes, and eager mode allocate them from whatever the KV cache leaves.

How DeepInfer handles it

DeepInfer@774ce24 sizes the decode logits the same way: max_kv_len = runtime_max_model_len // kv_compress_ratio (index_attn.py#L95, #L434-L441). It keeps them in BF16 (#L87).

On decode engines, its profile_run runs max_num_seqs sequences of 1 + nextn query tokens (model_runner.py#L444-L511). So the [max_num_seqs × (1 + nextn), max_kv_len] logits are part of the measured peak before the KV cache is sized.

Fix

vLLM's shared indexer now does the same accounting. Its profiling (dummy) path reserves the larger of the prefill budget and the decode worst case:

Rows: min(max_num_seqs, max_num_batched_tokens) × (1 + num_speculative_tokens). The profiling branch reads it from the current vLLM config.
Width: max_model_len, padded as DeepGEMM pads the rows: whole KV splits (384 columns on SM100, 256 elsewhere) at a 1024-byte stride, i.e. 1536 BF16 columns on SM100 (1,049,088 for 1M, as in the OOM above) and 256 FP32 columns elsewhere.
Bytes per element: 6 on SM100 (BF16 + FP32 copy), 2 with deep_select and no DCP, 4 on other platforms (FP32).

The runtime upcast and the reservation use the same predicate, so they cannot drift apart.

Memory

Same configuration, KV cache per rank at startup:

	top-k backend	decode logits reserved	CUDA graphs	KV cache / rank
main	`auto`	none (512 MiB budget)	18.40 GiB	41,050,274
this PR	`auto`	9.0 GiB (1536 rows × 1M × 6 B)	18.40 GiB	36,636,372 (−10.8%)
main	`deep_select`	none (512 MiB budget)	12.39 GiB	48,253,887
this PR	`deep_select`	3.0 GiB (1536 rows × 1M × 2 B)	12.39 GiB	48,163,989 (−0.2%)

Measured before the width was aligned to 1536 columns, which adds 4.5 MiB (auto) or 1.5 MiB (deep_select) to the reservation.

The KV cache shrinks by less than the reservation, because the reservation stacks on top of what is live inside the indexer, not on top of the run's global peak. With deep_select the reservation costs 0.2% of the KV cache. The decode graphs also drop the FP32 copy, so deep_select leaves 17.5% more KV cache than auto on main.

The DSA indexer's decode path allocates its paged MQA logits per call as [decode rows, max_model_len] (BF16 on SM100, FP32 elsewhere) and, on SM100, upcasts them to FP32 for every top-k backend but DeepSelect. The profiling run reserved only VLLM_SPARSE_INDEXER_MAX_LOGITS_MB for logits, the prefill chunks' budget, so the KV cache took the memory the decode logits need. CUDA graphs keep their own copy in the graph pool; eager and piecewise steps allocate from the remainder and, with many decodes at a long max_model_len, run out of memory. Reserve the larger of the prefill budget and the decode worst case: min(max_num_seqs, max_num_batched_tokens) * (1 + speculative tokens) rows at the padded max_model_len width, counting the FP32 upcast where it runs. The layer computes the row bound at construction and passes it to the op; the runtime path and the reservation share the upcast decision. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Canlin <canlinguosdu@gmail.com>

claude

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

gcanlin · 2026-10-21T09:23:19Z

cc @Juntian777 @ivanium @zyongye

gau-nernst · 2026-10-27T06:50:43Z

I think you can drop the test and make the logic a bit more concise i.e. inline functions appropriately

gcanlin · 2026-10-27T09:09:50Z

I think you can drop the test and make the logic a bit more concise i.e. inline functions appropriately

Thanks! I will do some clean up for it.

Per review: the profiling-path reservation is covered by the memory numbers in the PR description; the unit test mostly re-stated the formula. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Canlin <canlinguosdu@gmail.com>

Per review: compute the reserved bytes in the profiling branch and keep the FP32-copy predicate in one local shared with the runtime upcast, instead of two module-level helpers. Same reservation as before. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Canlin <canlinguosdu@gmail.com>

On SM100 DeepGEMM pads each paged MQA logits row to whole 384-column KV splits and 1024-byte strides, lcm 1536 BF16 columns (1,049,088 for a 1M max_model_len, as in the OOM report), and the FP32 copy keeps that stride. Rounding to 1024 columns reserved up to 1535 columns per row too little. Other platforms keep the 1024-column bound, above their 256-column padding. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Canlin <canlinguosdu@gmail.com>

…yout Spell the reservation as DeepGEMM's layout instead of bare byte counts and alignments: the logits dtype (BF16 on SM100, FP32 elsewhere) and its KV split width (384 / 256 columns), rows aligned to lcm(split, 1024-byte stride), plus the FP32 copy on SM100 when the top-k reads FP32. Other platforms now pad to DeepGEMM's 256 columns instead of an upper bound of 1024. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Canlin <canlinguosdu@gmail.com>

gcanlin · 2026-10-27T11:25:04Z

@gau-nernst READY now. Could you take a look again?

…rofiling run Per review: the bound only sizes the profiling reservation, so compute it there from get_current_vllm_config_or_none() instead of passing it through the custom op. The profiling run sets the config; other dummy runs (e.g. DP idle-rank batches) may not, and keep the prefill budget. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Canlin <canlinguosdu@gmail.com>

gcanlin · 2026-10-27T16:25:02Z

I think you can drop the test and make the logic a bit more concise i.e. inline functions appropriately

BTW, I also feel troubled by the generated AI code style(and the PR description), like overly defensive code, excessive checks, hard-to-read or messy abstractions, leading underscores and unnecessary UT. Is it possible to add a code style readme to avoid these in vLLM repo? Just some random thoughts.

It looks like it has been added in pr-checklist skill.

gau-nernst · 2026-10-27T23:46:48Z

/ci run

github-actions · 2026-10-27T23:48:01Z

✅ Triggered Buildkite CI #97788 for commit a914b1a02bc9.

github-actions · 2026-10-27T23:48:30Z

✅ @gcanlin, CI is now available for this PR.

/ci run starts upstream CI; /amd-ci run starts AMD CI only.
Run commands test your branch with every commit currently on its upstream target branch. If it is behind, the bot merges the target branch into it first (pull before pushing again). Append --allow-stale to a run command to test an outdated branch as is, at your own risk.
/ci retry retries failed jobs in the CI build for the current PR head. If the current head has no CI build, it starts a new CI build for the current head containing only jobs that failed in the latest earlier CI build for this PR.
/amd-ci retry retries failed jobs in AMD CI for the current PR head. Use /amd-ci run when the current head has no AMD CI build.
/ci cancel cancels scheduled or running CI builds for this PR branch; /amd-ci cancel does the same for AMD CI only.

gau-nernst · 2026-10-28T03:04:20Z

@gcanlin help me fix the CI. This is another reason I hate mocking via Namespace object 😂

The profiling branch now reads max_num_seqs, max_num_batched_tokens and num_speculative_tokens from the current vLLM config, so the test's config stub needs them. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Canlin <canlinguosdu@gmail.com>

gcanlin · 2026-10-28T04:44:24Z

/ci retry

github-actions · 2026-10-28T04:44:44Z

✅ Triggered Buildkite CI #97875 for commit b53b231a994a, running 3 failed step(s) from Buildkite CI #97788.

gcanlin · 2026-10-28T04:50:25Z

@gcanlin help me fix the CI. This is another reason I hate mocking via Namespace object 😂

Fixed. I can open a new PR to refactor the UT.

gcanlin · 2026-10-28T07:02:00Z

/ci retry

github-actions · 2026-10-28T07:02:20Z

✅ Triggered Buildkite CI #97897 for commit c21aa36affbf, running 1 failed step(s) from Buildkite CI #97875.

gau-nernst · 2026-10-28T07:27:15Z

I can open a new PR to refactor the UT.

probably not necessary. it's funny to have a PR just to fix it. we just have to live with it for now.

gcanlin · 2026-10-28T08:23:13Z

@gau-nernst CI pass now. ready to merge

…vllm-project#62860) Signed-off-by: Canlin <canlinguosdu@gmail.com> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> (cherry picked from commit 317081f)

Conflict in vllm/model_executor/layers/sparse_attn_indexer.py: the dense prefill path keeps the fused_mqa_topk row-window loop and takes upstream's DeepSelect BF16 ranking (vllm-project#62240) per window; the profiling branch takes upstream's logits reservation (vllm-project#62860, vllm-project#63047) next to the fused scratch. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: LopezCastroRoberto <rocastro@redhat.com>

mergify Bot added the bug Something isn't working label Oct 21, 2026

gcanlin marked this pull request as ready for review October 21, 2026 09:22

gcanlin requested a review from zyongye as a code owner October 21, 2026 09:22

claude Bot reviewed Oct 21, 2026

View reviewed changes

gcanlin mentioned this pull request Oct 22, 2026

[Bugfix][DSv4.1] Budget prefill memory the startup profile misses #63047

Merged

5 of 6 tasks

This was referenced Oct 23, 2026

[Core][CUDA Graph] Separate prefill (token) and decode (request) CUDA graph capture sizes Pb314314/vllm#2

Open

[Core][CUDA Graph] Separate prefill (token) and decode (request) CUDA graph capture sizes #63406

Open

gcanlin and others added 4 commits October 27, 2026 17:48

gau-nernst reviewed Oct 27, 2026

View reviewed changes

Comment thread vllm/model_executor/layers/sparse_attn_indexer.py Resolved

Pb314314 mentioned this pull request Oct 27, 2026

[Perf][DSA] Read BF16 logits directly in persistent top-k #64138

Merged

6 tasks done

gau-nernst approved these changes Oct 27, 2026

View reviewed changes

gau-nernst added the ready ONLY add when PR is ready to merge/full CI is needed label Oct 27, 2026

Merge branch 'main' into fix/dsa-decode-logits-profile

a914b1a

Merge branch 'main' into fix/dsa-decode-logits-profile

c21aa36

gau-nernst merged commit 317081f into vllm-project:main Oct 28, 2026
15 checks passed

GuoxiangZu mentioned this pull request Oct 31, 2026

[Bug]: [ROCm][GLM-5.3-Flash]: "Workspace is locked" AssertionError in rocm_fp8_paged_mqa_logits under high decode concurrency (glm5next AMD indexer reserves decode logits via torch.empty, not the workspace) #64874

Open

1 task done

Conversation

gcanlin commented Oct 21, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Overview

OOM

Why

How DeepInfer handles it

Fix

Memory

Uh oh!

claude Bot left a comment

Choose a reason for hiding this comment

Claude Code Review

Uh oh!

gcanlin commented Oct 21, 2026

Uh oh!

gau-nernst commented Oct 27, 2026

Uh oh!

gcanlin commented Oct 27, 2026

Uh oh!

gcanlin commented Oct 27, 2026

Uh oh!

Uh oh!

gcanlin commented Oct 27, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

gau-nernst commented Oct 27, 2026

Uh oh!

github-actions Bot commented Oct 27, 2026

Uh oh!

github-actions Bot commented Oct 27, 2026

Uh oh!

gau-nernst commented Oct 28, 2026

Uh oh!

gcanlin commented Oct 28, 2026

Uh oh!

github-actions Bot commented Oct 28, 2026

Uh oh!

gcanlin commented Oct 28, 2026

Uh oh!

gcanlin commented Oct 28, 2026

Uh oh!

github-actions Bot commented Oct 28, 2026

Uh oh!

gau-nernst commented Oct 28, 2026

Uh oh!

gcanlin commented Oct 28, 2026

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

gcanlin commented Oct 21, 2026 •

edited

Loading

gcanlin commented Oct 27, 2026 •

edited

Loading