Skip to content

[Bugfix][Attention] Keep DSA shared-indexer top-k state per DBO microbatch - #60734

Open
aniravich wants to merge 1 commit into
vllm-project:mainfrom
aniravich:fix/dsa-shared-indexer-dbo
Open

aniravich wants to merge 1 commit into
vllm-project:mainfrom
aniravich:fix/dsa-shared-indexer-dbo

Conversation

@aniravich

@aniravich aniravich commented Oct 8, 2026 •

Copy link
Copy Markdown
Contributor

Overview

With --enable-dbo, GLM-5.x (DSA with shared top-k layers, index_topk_freq > 1) loses about a third of its GSM8K accuracy: 0.93 with DBO off, 0.61 with DBO on. The layers that reuse an earlier layer's sparse-attention top-k (57 of GLM-5.3's 78 layers) read it from buffers that have no microbatch dimension. DBO runs microbatch 1's index-producing layer between microbatch 0's producer and its consumers, so microbatch 0's consumers attend over microbatch 1's token selection. This PR gives every DBO microbatch after the first its own copy of that state. With it, DBO-on accuracy matches DBO off.

Companion to #60732 (DSA indexer crash under DBO, which hid this bug). The two fixes are independent; DBO with DeepEP high-throughput needs both. No existing issue or PR covers this (searched dbo topk, index_topk_freq dbo, shared indexer microbatch, GLM-5 dbo accuracy: 0 results). Related open PR: #56476 fixes the same class of bug for DeepSeek-V4.1 (its top-k and candidate buffers, via a new dbo_select_buffer helper). It does not cover the DeepSeek-V3.2 / GLM-5.x IndexCache path changed here; if it lands first, this PR can switch to its helper.

Claims

  • Fixes GLM-5.x accuracy under DBO on Model Runner V1 and V2: GSM8K 0.61 → 0.92 (DBO off: 0.92–0.93). Also with DBO microbatches replayed from FULL decode CUDA graphs (DeepEP low-latency): 0.44 → 0.85 (DBO off: 0.84–0.86).
  • No change without DBO, and none for DSA models without shared top-k layers (DeepSeek-V3.2, GLM-5.1: every layer computes its own top-k, so no microbatch state is allocated). Measured: both score the same with DBO on as off, with and without this change, and report identical KV cache memory with and without it.
  • Memory under DBO, per GPU and per extra microbatch: (max_num_batched_tokens + G * (max_decode_rows + 1)) * index_topk * 4 bytes, where G is the number of index groups with followers (19 for GLM-5.3). GLM-5.3 at max_num_batched_tokens=8192, max_num_seqs=1024: 216 MiB, measured as 0.22 GiB less KV cache memory.
  • No host syncs, no extra kernels. The per-microbatch lookup is one dict read per access, the same one DBO already uses everywhere (dbo_current_ubatch_id()).
  • HiSparse with DBO and shared top-k layers now raises at startup (it previously ran with the same corruption); supporting it needs per-microbatch HiSparse staging, out of scope here.

Validation

Unit test (new), test_sparse_mla_index_group_keeps_topk_per_dbo_ubatch in tests/v1/attention/test_sparse_mla_backends.py: one leader group without followers and one with a follower, built with num_ubatches=2. It checks that the leader-only group allocates nothing, that leader and follower impls resolve to the same per-microbatch buffer, and that after both microbatches' leaders write (in DBO order: leader 0, leader 1, then followers) each follower's logical and converted physical top-k equal its own leader's bit for bit. Runs in the existing "V1 Attention" CI shards.

Build Result
this branch passed
base (without fix) fails (the builder has no num_ubatches)
mutation: index group always returns microbatch 0's state fails
mutation: impl not wired to the per-microbatch logical buffers fails

The two mutations each break one half of the fix and confirm the assertions catch it. Existing tests on this branch (H200): tests/v1/attention/test_sparse_mla_backends.py 293 passed (1,188 skipped by hardware), tests/v1/worker/test_workspace.py + test_gpu_ubatch_slicing.py 40 passed. tests/v1/distributed/test_dbo.py (DeepSeek-V2-Lite, DP2, DBO): 2 passed (deepep_low_latency, deepep_high_throughput). pre-commit's ruff check and format pass; mypy on the four changed source files reports no new errors (two fewer than base).

All tests and GPU runs in this description used main @ 4a705c7. This branch is rebased onto main @ 240785b, where the five files it changes differ only in HiSparse code it does not touch (HiSparseMLAIndexGroup routing, #60682; HiSparse tests, #60429 and #60682), and ubatching.py / ubatch_utils.py are unchanged. Re-run on the rebased commit (H200): the new test passes (and still fails on main without the fix), test_sparse_mla_backends.py 293 passed (1,188 skipped), tests/v1/worker/test_workspace.py and test_gpu_ubatch_slicing.py 40 passed, and pytest -v -s tests/v1/distributed/test_dbo.py 2 passed (deepep_low_latency, deepep_high_throughput). The GSM8K and long-prompt runs were not repeated on the new base.

Model evaluation, GLM-5.3 FP8 (1 x 8 H200, DP4 x TP2, EP8, DeepEP high-throughput unless noted, cudagraph_mode NONE, fp8 KV, prefix caching, VLLM_USE_DEEP_GEMM=0 VLLM_MOE_USE_DEEP_GEMM=0; lm_eval GSM8K 5-shot, 1,319 questions, num_concurrent=256; standard error ±0.007 at 0.92, ±0.013 at 0.61). Every DBO-on run microbatched 100+ steps per worker.

vLLM Runner DBO GSM8K flexible / strict
main + #60732 V2 off 0.927 / 0.928
main + #60732 V1 off 0.924 / 0.924
main + #60732 V1 on 0.613 / 0.607
main + #60732 V2 (VLLM_USE_V2_MODEL_RUNNER=1) on 0.614 / 0.610
main + #60732 + this PR V1 on 0.922 / 0.923
main + #60732 + this PR (repeat) V1 on 0.923 / 0.921
main + #60732 + this PR V2 on 0.919 / 0.919
main + this PR only, VLLM_DBO_COMM_SMS=0 V1 on 0.923 / 0.923

Note that --enable-dbo alone falls back to Model Runner V1; the V2 rows set VLLM_USE_V2_MODEL_RUNNER=1 (logs confirm "Using V2 Model Runner" and "Dual batch overlap is enabled").

CUDA graphs (DBO microbatches replayed from FULL decode graphs). This needs DeepEP low-latency (high-throughput cannot be captured) and pure decode steps on every DP rank, which lm_eval's streaming load never produces at low-latency's 256-token step budget (0 microbatched steps in a full GSM8K run). Workload instead: the first 1,024 GSM8K test questions zero-shot, sent at once (256 per DP rank), 192 forced greedy tokens each, scored for the final answer. Same GLM-5.3 setup with --all2all-backend deepep_low_latency, cudagraph_mode FULL_DECODE_ONLY, Model Runner V1 in all runs. These runs also carry a one-line fix for a separate bug: the batched Triton FP8 MoE that low-latency mode uses (with DeepGEMM off) does not clamp BLOCK_SIZE_K to the 128-wide FP8 block (fused_batched_moe.py:452, unlike fused_moe.py:844), which gives GSM8K 0.02 even with DBO off; with the clamp, 0.93. That fix is its own PR (#60733).

With --max-num-batched-tokens 512 (NVSHMEM_QP_DEPTH=4096, --gpu-memory-utilization 0.93, so the prompts prefill in about 50 steps and most answer tokens are decoded in microbatched steps), each DBO-on worker replayed 180+ microbatched decode graphs and ran 50+ eager microbatched prefill steps:

Run (1,024 questions, low-latency, FULL_DECODE_ONLY) Accuracy vs DBO off #1: lost / gained (McNemar z)
DBO off #1 0.843
DBO off #2 0.857 11 / 26 (-2.5)
DBO on, main + #60732 0.440 442 / 30 (19.0)
DBO on, main + #60732 + this PR 0.845 40 / 42 (-0.2)
DBO on, main + #60732 + this PR (repeat) 0.850 31 / 38 (-0.8)

At the default 256-token budget (fewer microbatched steps) the picture is the same: DBO off 0.849 / 0.848, without this PR 0.775 (z 6.4 vs off), with it 0.838 (z 1.4). Exact token sequences are not a usable signal in this mode: two DBO-off runs already differ at the first token for 100 of 1,024 prompts.

Negative control, GLM-5.1 FP8 (DSA, index_topk_freq=1: no shared top-k layers), same setup (main + #60732): DBO off 0.946 / 0.948; DBO on without this PR 0.952 / 0.950; DBO on with this PR 0.947 / 0.950 (100+ microbatched steps per worker in both DBO runs). With and without this PR the DBO-on runs report the same model memory (100.48 GiB) and KV cache memory (15.16 GiB): nothing is allocated when every layer computes its own top-k. DSA plus DBO is fine when no layer reuses another's top-k.

Second negative control, DeepSeek-V3.2 (61 layers, every one computes its own top-k), same setup. Note: in this configuration vLLM main crashes with an illegal memory access in the first microbatched profiling step for DeepSeek-V3.2 with DBO on, with or without either fix (also without #60732, where it shows up as a DeepEP timeout instead). It disappears with CUDA_LAUNCH_BLOCKING=1, so it is a separate stream race, not investigated here. The DBO-on runs below therefore use CUDA_LAUNCH_BLOCKING=1, which keeps DBO's microbatch order: GLM-5.3 without this PR still shows this bug under it (0.607 / 0.596, vs 0.613 / 0.607 without blocking).

DeepSeek-V3.2 GSM8K flexible / strict Long prompts vs DBO off: differ at token 0, median |Δ logprob| KV cache memory
DBO off (second DBO-off run = noise floor) 0.952 / 0.952 0, 0.0001 24.46 GiB
DBO on, without this PR 0.954 / 0.952 0, 0.0000 23.22 GiB
DBO on, with this PR 0.952 / 0.951 0, 0.0000 23.22 GiB

Isolating the cause (GLM-5.3, main + #60732, DBO on):

Experiment GSM8K flexible Reading
DBO only on decode steps 0.601 decode steps alone reproduce the full drop
DBO only on prefill steps (3 microbatched steps) 0.901 few prefill splits, small drop
prefix caching off 0.619 not prefix caching
VLLM_DBO_COMM_SMS=0 0.610 not the DeepGEMM SM partition
prototype: per-microbatch converted (physical) top-k only, logical buffer still shared 0.920 short steps only read the converted copy

Long prompts, greedy (GLM-5.3, same setup, Model Runner V1): 128 GSM8K-train prompts of 4.5k–7.7k tokens sent at once (prefill steps microbatch as 2 x 4096), 32 greedy tokens each with logprobs, compared against a DBO-off run. A second DBO-off run gives the noise floor (batching is not invariant, so long greedy continuations drift even without DBO).

Run vs. DBO off Identical 32-token outputs Mean first differing token Differ at token 0 Median |Δ logprob| of token 0
DBO off (second run) 35 / 128 18.2 0 0.0001
DBO on, main + #60732 0 / 128 3.0 7 0.0135
DBO on, prototype with only the converted (physical) top-k per microbatch 2 / 128 3.4 10 0.0081
DBO on, main + #60732 + this PR 24 / 128 17.1 0 0.0001

A step with more tokens than the conversion workspace (max_decode_rows + 1, 1,025 here) skips the cached conversion, and every follower converts from the logical buffer again. That is why the prototype fixes GSM8K but not long prefills, and why this PR keeps both the logical buffer and the converted state per microbatch.

Reproduction commands
# unit test
pytest -v tests/v1/attention/test_sparse_mla_backends.py -k dbo_ubatch

# GSM8K (add VLLM_USE_V2_MODEL_RUNNER=1 for Model Runner V2)
VLLM_USE_DEEP_GEMM=0 VLLM_MOE_USE_DEEP_GEMM=0 vllm serve zai-org/GLM-5.3 --tensor-parallel-size 2 --data-parallel-size 4 --enable-expert-parallel \
  --all2all-backend deepep_high_throughput --enable-dbo --compilation-config '{"cudagraph_mode": "NONE"}' \
  --kv-cache-dtype fp8 --max-model-len 8192 --gpu-memory-utilization 0.90 --enable-prefix-caching
lm_eval --model local-completions \
  --model_args model=zai-org/GLM-5.3,base_url=http://127.0.0.1:8000/v1/completions,num_concurrent=256,tokenized_requests=False \
  --tasks gsm8k --num_fewshot 5

Details

Root cause (main @ 240785b). GLM-5.x configs set index_topk_freq=4, index_skip_topk_offset=3: only 21 of 78 layers run the sparse indexer ("leaders": 0, 1, 2, 6, 10, …, 74); each of the other 57 reuses the top-k of the leader before it. That hand-off goes through state with no microbatch dimension:

  • the model-wide logical buffer DeepseekV32Model.topk_indices_buffer, which the leader writes (fused_norm_rope, sparse_attn_indexer) and every follower's sparse MLA impl reads;
  • per index group, SparseMLAIndexGroup.physical_topk_indices / valid_topk_counts (logical → physical conversion, done once by the leader and reused by its followers) and the logical_topk_ready / physical_topk_ready events.

DBO yields inside MoE dispatch/combine, so between microbatch 0's leader L and its follower L+1, microbatch 1 runs its own layer L and overwrites all of the above. Every follower of microbatch 0 except layer 3 (layers 0–2 are dense, so there is no yield between leader 2 and follower 3) then attends over microbatch 1's selected tokens, with block-table conversion done for microbatch 1's requests. Microbatch 1's followers read their own data and are unaffected. Leaders without followers (DeepSeek-V3.2, GLM-5.1, layers 0 and 1 here) are unaffected because they write and read within one layer.

How it was found: while validating #60732, GSM8K with DBO on came out at ~0.6 against 0.93 off, with or without that fix and with SM partitioning disabled. The table above narrowed it to decode-step microbatching independent of prefix caching and SM count; GLM-5.1 (same architecture, no shared top-k) was unaffected, which pointed at the shared top-k hand-off. A prototype that gave each microbatch its own converted top-k restored 0.920; the long-prompt test below shows it is not enough on its own.

Fix:

  • SparseMLAIndexGroupBuilder takes num_ubatches. When a group gets its first follower, it allocates one UbatchTopkState (physical indices, valid counts, two events) per extra microbatch for that group, and, once per builder, one logical buffer per extra microbatch. Groups without followers allocate nothing.
  • SparseMLAIndexGroup._topk_state() and SharedTopkIndicesBuffer.topk_indices_buffer return microbatch 0's existing state, or the extra microbatch's copy when dbo_current_ubatch_id() > 0.
  • The leader's writers (fused_norm_rope, sparse_attn_indexer in DeepseekV32Attention) now take the buffer from self.impl.topk_indices_buffer, so writer and readers resolve the same per-microbatch buffer.
  • DeepseekV32Model passes num_ubatches=get_num_ubatches(parallel_config).

The extra buffers are allocated at model construction, before the KV-cache profile, so they come out of KV cache memory. In FULL-graph DBO capture, each microbatch thread resolves its own buffers at capture time.

Alternatives considered
  1. Add a microbatch dimension to topk_indices_buffer itself: touches the MTP proposer's buffer sharing and every reader; the per-impl property was already the single resolution point.
  2. Make followers recompute top-k under DBO: correct but runs the indexer on all 78 layers, which is what index_topk_freq exists to avoid.
  3. Disable DBO for models with shared top-k layers: loses the overlap for exactly the models (GLM-5.x) that use DBO in production.
  4. Allocate per-microbatch state in every index group: wastes memory for DeepSeek-V3.2 and GLM-5.1, which have no followers.

MTP: the draft's indexers are re-pointed to the target buffer through the module setter, which clears any microbatch buffers. The MTP layer is a leader-only group and the draft runs outside the microbatch threads, so it is unchanged.

Not covered: the ROCm DSA model (vllm/models/deepseek_v32/amd) and the non-CUDA fallback in deepseek_v2.py/glm4_moe_lite.py still build their groups with one microbatch; they keep the old behavior if they run DBO with shared top-k layers. HiSparse raises instead (see Claims).

AI assistance: Claude Code was used to trace the root cause, write the patch and test, run the validations above, and draft this description. I reviewed every changed line, can explain and defend the change, and checked the test results above.


Pull Request Checklist
  • I used vLLM's /pr-checklist skill. (Mandatory for agents, optional for humans).

  • AI assistance was used during the creation of this PR.

  • Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.

  • Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI.

  • Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.

  • Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.

🤖 Generated with Claude Code

…batch

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Anirudh Ravichandran <ani@primeintellect.ai>
@github-actions

github-actions Bot commented Oct 8, 2026

Copy link
Copy Markdown

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

PRs do not trigger a full CI run by default. Reviewers with write access and configured trusted contributors can comment /ci run for upstream CI or /amd-ci run for AMD CI only whenever CI signals are needed.

Once the PR is approved or has the ready label, the PR author can also use the corresponding /ci run, /ci retry, and /ci cancel commands, or their /amd-ci variants. New commits do not start upstream CI automatically.

If you have any questions, please reach out to us on Slack at https://slack.vllm.ai.

Agent Guidelines

IMPORTANT: 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.

🚀

@aniravich
aniravich marked this pull request as ready for review October 8, 2026 23:37

@claude claude Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@aniravich

Copy link
Copy Markdown
Contributor Author

@MatthewBonanni could you take a look when you get a chance? It fixes IndexCache top-k state being shared between DBO microbatches on GLM-5.x (related to vllm-project/vllm#60732).

@mergify

mergify Bot commented Oct 20, 2026

Copy link
Copy Markdown
Contributor

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @aniravich.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

@mergify mergify Bot added the needs-rebase label Oct 20, 2026

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

bug Something isn't working deepseek Related to DeepSeek models needs-rebase

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant