Skip to content

Fix sparse MQA merge partitions with repeated block indices - #16

Closed
WoosukKwon wants to merge 1 commit into
devfrom
fix/sparse-mqa-duplicate-merge
Closed

WoosukKwon wants to merge 1 commit into
devfrom
fix/sparse-mqa-duplicate-merge

Conversation

@WoosukKwon

Copy link
Copy Markdown

Sparse MQA metadata can assert when sorted candidate rows contain repeated block IDs. vLLM emits repeated padding despite the previous native API requirement for unique indices. In a real DeepSeek-V4.1 TP4 startup, two rows containing 389 copies of block 388 (context lengths 3105/3106) trigger a sparse-slot assertion: the merge consumes Q1 entries too early, then the next merge-path partition resumes at an earlier slot, underflowing the unsigned offset.

Pair only the last equal Q0 slot with the first Q1 slot, and skip only that carried match at a partition boundary. This preserves output slots and supports repeated indices without removing bounds checks. Update the API contract and add regression coverage to the existing attention suite.

Validation on GB200, commit 2a95b271c7decb99a4773d17299f393167ab0a01:

  • Captured two-row input fails before the fix and passes afterward; the complete 8192-row metadata capture also passes.
  • 12 regression combinations (MXFP4/MXFP8 × aligned/unaligned/paged × repeated-only/mixed-prefix) match dense logits bitwise.
  • Four existing sparse cases pass, including their 30 self-consistency repetitions; coverage includes 32K–256K KV and 512–2048 selected blocks.
  • 2000 randomized CPU merge cases preserve every sparse slot exactly once with monotonic split offsets.
  • Original real-weight TP4 MegaAttention serving starts with sparse logits, autotuning, DSpark5 and FULL decode graphs enabled. GSM8K smoke: 31/32 flexible extraction, 26/32 strict. This is not a full accuracy evaluation. A 16,504-token prefill executes without CUDA errors but immediately stops generation; this establishes execution only.

Test commands executed from the vLLM worktree, using its virtualenv and an isolated copy of the pinned DeepGEMM runtime with the changed JIT header (unchanged native extension):

srun --jobid=29623 --overlap -N1 -n1 /usr/bin/env \
  PYTHONPATH=/home/woosuk/.cache/v41-dspark-opt/indexer-assert/patched-runtime:/home/woosuk/workspace/deepgemm-sparse-merge-fix/tests \
  OMP_NUM_THREADS=1 .venv/bin/python -c \
  'from test_attention import test_sparse_mqa_logits_repeated_blocks; test_sparse_mqa_logits_repeated_blocks()'

srun --jobid=29623 --overlap -N1 -n1 /usr/bin/env \
  PYTHONPATH=/home/woosuk/.cache/v41-dspark-opt/indexer-assert/patched-runtime:/home/woosuk/workspace/deepgemm-sparse-merge-fix/tests \
  OMP_NUM_THREADS=1 DG_MQA_NUM_CASES=4 .venv/bin/python -c \
  'from test_attention import test_sparse_mqa_logits; test_sparse_mqa_logits()'

Local reproduction artifacts, serving argv and logs are recorded in /home/woosuk/.cache/v41-dspark-opt/indexer-assert/REPORT.md and fixed-command.json. Default vendored-import GPU replay also passes after integrating the fixed header.

Duplicate-work check: reviewed all four open PRs in vllm-project/DeepGEMM and searched sparse metadata / duplicate merge; none addresses this fix. No linked issue. This is separate from vllm-project/vllm#57432, which fixes FlashInfer DSpark attention metadata.

AI assistance: implemented and validated with OpenAI Codex.

Pair only the last Q0 duplicate with the first Q1 duplicate so stable merge partitions preserve sparse slot order and coverage.

Co-authored-by: Codex <noreply@openai.com>
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
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