Skip to content

[Bugfix][DS4] Keep sparse decode workspace valid across CUDA Graph replay - #29

Draft
aoshen02 wants to merge 1 commit into
ISEEKYAN:ds4-v9-rc1from
aoshen02:codex/bugfix-sparse-decode-capacity
Draft

aoshen02 wants to merge 1 commit into
ISEEKYAN:ds4-v9-rc1from
aoshen02:codex/bugfix-sparse-decode-capacity

Conversation

@aoshen02

Copy link
Copy Markdown

Problem

In this fork, VLLM_DS4_DECODE_KERNEL=sparse can run decode through the sparse prefill kernel. With BI disabled, its temporary KV buffer and SWA offset are sized from CPU sequence lengths. CUDA Graph capture freezes those sizes, but sequence lengths grow during replay. Later compressed keys can overlap the SWA region or exceed the captured workspace.

This is a fork-specific correctness fix, not a claim that upstream main's paged decode has this bug.

Minimal fix

  • Use model capacity for compressed KV workspace in both modes, as BI already does.
  • Bound the SWA workspace by the static query tensor size rather than current sequence lengths.
  • Remove obsolete CPU length arithmetic.
  • Add a real-kernel CUDA Graph regression for C1/C4/C128 that grows sequences after capture.

No default backend switch, recipe change, native rebuild, or new kernel. BI's sizing policy is unchanged. The non-BI sparse path may reserve more workspace; that is required to cover future replays. The metadata assertion only makes an existing invariant explicit for mypy.

Validation

On GB200, old-wheel dev9 runtime with FlashInfer 0.6.18:

/opt/ds4-venv/bin/python -m pytest --import-mode=importlib \
  --confcutdir=<worktree>/tests/models \
  <worktree>/tests/models/test_deepseek_v4_decode_sparse.py -v

Candidate source was loaded into the installed module before pytest to retain the image's compiled extensions. Eight tests passed. Unpatched installed source fails the new C4 replay/eager comparison at lengths [15,17] after capture at [3,5]. The test uses actual quantized caches, gather, index combine and FlashMLA, not mocked attention math.

Full 43-layer rollout, 8 nodes / 32 GB200, 128 prompts x 8 samples, natural EOS, seqs64, batched tokens8192: 1,024 requests, 5,156,038 output tokens in156.085s (33.034 token/ms), driver exit0. The unpatched sparse run had abnormal short output. This is a rollout smoke/performance result, NOT an accuracy evaluation or training-alignment proof. The paged-vs-sparse speed difference must not be attributed solely to this fix. Cleanup emitted a DataLoader-worker killed warning after the completed result.

The GPU tests and rollout preceded final formatting, explicit metadata assertion, and the lint-required replacement of test synchronization calls with torch.accelerator.synchronize(). Final GPU rerun and human review are pending; hence Draft.

Duplicate check / review

Searched open PRs in vllm-project/vllm and ISEEKYAN/vllm for sparse decode, workspace/capture, _forward_decode_sparse, and VLLM_DS4_DECODE_KERNEL. No dedicated fix for this capacity issue was found. Upstream vllm-project#54955 is the broader fork integration; fork #19 optimizes gather scheduling, not this allocation contract.

AI assistance: OpenAI Codex investigated and prepared this patch and tests. Human submitter review of every changed line and validation is required before marking ready/merging. No GSM8K/AIME evaluation was run for this change.

Use static capacity rather than capture-time CPU sequence lengths when BI is disabled. Cover growing sequences with real-kernel CUDA graph replay tests.

Co-authored-by: OpenAI Codex <noreply@openai.com>
Signed-off-by: aoshen02 <aoshen02@users.noreply.github.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.

1 participant