Skip to content

[Bugfix] Avoid unused sparse decode split-KV workspace - #19

Merged
LucasWilkinson merged 1 commit into
vllm-project:mainfrom
MatthewBonanni:fix/skip-unused-splitkv-workspace
Aug 26, 2026
Merged

LucasWilkinson merged 1 commit into
vllm-project:mainfrom
MatthewBonanni:fix/skip-unused-splitkv-workspace

Conversation

@MatthewBonanni

@MatthewBonanni MatthewBonanni commented Aug 25, 2026 •

Copy link
Copy Markdown
Member

Summary

  • avoid allocating FP32 split-KV accumulators when sparse decode has only one SM partition
  • return the kernel output directly instead of launching the unnecessary combine kernel
  • make the optional accumulator handling safe for the SM100 head64x2/head128 paths
  • add a peak-allocation regression test

Companion vLLM integration PR: vllm-project/vllm#53755.

AI assistance

OpenAI Codex assisted with analysis, implementation, and validation. This PR is intentionally draft pending the human submitter's line-by-line review; the submitter will own and defend the change before marking it ready.

Co-authored-by: OpenAI Codex <codex@openai.com>
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
@LucasWilkinson
LucasWilkinson merged commit 0397728 into vllm-project:main Aug 26, 2026
1 check passed
drakosha added a commit to drakosha/FlashMLA that referenced this pull request Nov 3, 2026
vllm-project#19 allocated the split-KV accumulators only when the scheduler has more
than one partition. The V4.1 sync (vllm-project#20) brought back the upstream
interface, which allocates them whenever split-KV scheduling is enabled,
and on SM90 that is always. With a single partition no request is split:
the SM90 kernel writes `out` / `lse` directly and the combine has nothing
to merge, so the (b + num_sm_parts) x s_q x h_q x d_v fp32 buffers are
never read. In vLLM's mixed-batch layout (b=1, s_q=num_tokens) they grow
to 8 GiB at 32k tokens and reintroduce the runtime OOM of
vllm-project/vllm#53413 (vllm-project/vllm#59305).

Skip the accumulators and the combine launch on SM90 when there is a
single partition, and bring back the workspace test from vllm-project#19 for SM90.
SM100 is left to vllm-project#42: its kernels in the new tree build the o_accum tensor
map and offset its pointers unconditionally, so there split-KV has to be
turned off before scheduling instead.

Signed-off-by: Mikhail Kostryukov <mike@triptrack.net>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants