[aniravich]

Overview

With --enable-dbo and an all2all backend that reserves communication SMs (e.g. DeepEP high-throughput), DSA models (DeepSeek-V3.2, GLM-5.x) crash in the sparse-attention indexer as soon as a microbatched step contains a decode-branch request:

RuntimeError: Assertion error (csrc/apis/attention.hpp:346): _schedule_meta_size == num_sms + 1 and _meta_info_size == 2

The indexer builder plans DeepGEMM's paged-MQA-logits schedule for the whole GPU, but DBO lowers DeepGEMM's SM count (set_num_sms(total - VLLM_DBO_COMM_SMS)) while the microbatches run. This PR re-plans the schedule at the call site when the two counts differ. No existing issue or PR covers this (searched schedule_meta_size, dbo indexer num_sms, paged_mqa_logits set_num_sms: 0 results). Rechecked 2026-10-08 against open PRs: #56476 (DeepSeek-V4.1 per-microbatch index state) also edits sparse_attn_indexer.py, in different lines, and does not touch the schedule metadata.

Claims

• Fixes the DSA indexer decode crash under DBO (Model Runner V1 and V2, CUDA + DeepGEMM).
• The re-planned schedule gives bit-identical indexer logits, so the fix does not change results: with it, GLM-5.3 accuracy under DBO matches DBO with SM partitioning disabled (VLLM_DBO_COMM_SMS=0), where this code path never runs.
• Separately, GLM-5.3 accuracy with DBO on is well below DBO off (GSM8K ~0.6 vs 0.93) with or without this change. That is a separate, pre-existing DBO bug in GLM-5.x's shared top-k layers that this crash was hiding, fixed in [Bugfix][Attention] Keep DSA shared-indexer top-k state per DBO microbatch #60734 (see Validation).
• No cost when the SM count is unchanged (one host-side get_num_sms() per indexer layer). Under DBO, one ~9 µs metadata kernel per microbatch per forward, reused by all indexer layers. No host syncs; CUDA-graph safe.
• XPU / ROCm paths unchanged.

Validation

This branch is rebased onto main @ 240785b; the rebase did not change the patch (same git patch-id). The unit and related tests below were re-run on the rebased commit (H200) with the same results. The GSM8K runs predate the rebase and used main @ 4a705c7 with and without this patch.

Unit test (new), tests/v1/attention/test_indexer_schedule_metadata.py: plans a schedule for the full GPU, lowers DeepGEMM's SM count by 20 (what create_sm_control_context does), and checks that get_schedule_metadata() re-plans once, the kernel runs, and logits match the full-SM result on all valid columns; the original schedule comes back after the SM count is restored. Runs in the existing H100/B200 "V1 Attention" CI shards; skipped without DeepGEMM or below SM90.

Build	GPU	Result
this branch (main @ 4a705c7, precompiled from 2e06395)	H200	next_n 1, 2, 4: 3 passed
vLLM 0.29.0 + backport	H200	next_n 1, 2: passed (next_n=4 needs a newer DeepGEMM than 0.29 ships)

Related existing tests on this branch (H200): tests/kernels/attention/test_deepgemm_attention.py 5 passed, test_indexer_native_next_n.py 9 passed, test_sparse_indexer_decode_seq_lens.py 7 passed, test_indexer_dcp_localize.py 114 passed, test_sparse_mla_backends.py 292 passed (1,188 skipped by hardware). test_indexer_deepseek_v4_slot_mapping.py: 30 passed, 3 could not run locally (they load the gated meta-llama/Meta-Llama-3-8B config); they fail identically on main and do not touch the changed code.

Kernel-level check (H200, DeepGEMM fp8_fp4_paged_mqa_logits, GLM-5.3 indexer shapes: 32 heads x 128, block 64, 13 requests with contexts 1 to 32k):

SM count	Schedule planned for	Result
132	132	reference
112 (inside vLLM's SMControlContextManager)	132	the production assertion
112	112 (this PR)	runs; logits bit-identical to the reference

End-to-end reproducer, before/after. A 4-layer GLM-5.3 (GlmMoeDsaForCausalLM, production attention/indexer geometry, dummy weights) on 2 x H200, DP2/EP2, --all2all-backend deepep_high_throughput --enable-dbo --dbo-prefill-token-threshold 32 --dbo-decode-token-threshold 32, cudagraph_mode NONE, fp8_per_block + fp8 KV, prefix caching. Workload: random-token prompts of 64*k+1 tokens, then exact repeats (a full prefix-cache hit leaves a 1-token short extend at the head of the batch) mixed with fresh long prompts, pinned per DP rank. vLLM 0.29.0 with and without the backport of this patch (same two files):

Run	DBO steps	Assertion	Requests
DBO off	0	0	6 rounds, 0 errors
DBO on, without fix	4	yes, 6 s after the first DBO step	29/32 failed in round 2, engine dead
DBO on, with fix	28	0	6 rounds, 0 errors
DBO on, with fix, 30 rounds incl. decodes	158	0	944 requests, 0 errors

With the fix, DBO-on generations equal DBO-off: 944/944 identical sequences (7,440/7,440 tokens), first-token logprob difference at most 1.9e-6.

Model evaluation, GLM-5.3 FP8 (1 x 8 H200, DP4 x TP2, EP8, DeepEP high-throughput, cudagraph_mode NONE, fp8 KV, prefix caching, VLLM_USE_DEEP_GEMM=0 VLLM_MOE_USE_DEEP_GEMM=0 as on our production prefill, FlashInfer prebuilt kernels; lm_eval GSM8K 5-shot, 1,319 questions, num_concurrent=256):

vLLM	Config	GSM8K exact_match (flexible / strict)
this branch	DBO off	0.927 / 0.928
main (without fix)	DBO on	crash: the assertion on 8 workers, 4 engines dead, eval aborted
this branch	DBO on	0.592 / 0.582 (all 1,319 answered)
main (without fix)	DBO on, VLLM_DBO_COMM_SMS=0 (fix path never runs)	0.610 / 0.608
this branch	DBO on, VLLM_DBO_COMM_SMS=0	0.623 / 0.606
0.29.0	DBO on	crash: the assertion on 8 workers, eval aborted
0.29.0 + backport	DBO on	0.669 / 0.681 (all 1,319 answered)

The fix turns the crash into a completed eval. The drop from 0.93 to ~0.6 with DBO on is not caused by it: the VLLM_DBO_COMM_SMS=0 rows never re-plan (and never crash) and land in the same range. DBO with GLM-5.x has an accuracy bug of its own that this crash has been hiding: layers that reuse another layer's top-k read it from state shared across microbatches. #60734 fixes it; with both PRs, GLM-5.3 with DBO on scores 0.922 / 0.923 (DBO off: 0.927 / 0.928).

Reproduction commands

# unit test
pytest -v tests/v1/attention/test_indexer_schedule_metadata.py

# GSM8K
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):

• create_sm_control_context (vllm/v1/worker/ubatch_utils.py:83-126) uses deep_gemm.set_num_sms as the compute-SM setter whenever has_deep_gemm() and comm_sms > 0, so inside the microbatches DeepGEMM sees total - VLLM_DBO_COMM_SMS SMs (132 - 20 = 112 on H100/H200). It keys on DeepGEMM being importable, not on VLLM_USE_DEEP_GEMM. comm_sms is 0 only for DeepEP low-latency and NIXL-EP.
• The context wraps the microbatched forward in both runners (gpu_ubatch_wrapper.py:426,453; MRV2 gpu/ubatch_utils.py:496), after attention metadata has been built.
• The indexer builder plans the decode schedule for num_compute_units SMs (v1/attention/backends/mla/indexer.py:967-968, :1688-1697, MTP drafts :1782-1789), and sparse_attn_indexer.py:729 / glm5next/nvidia/sparse_indexer.py:582 pass it to fp8_fp4_paged_mqa_logits, which asserts it against the current SM count.
• The decode branch takes every request with query_len <= 1 + num_speculative_tokens, including 1-token short extends of prefix-cached prompts, which the reorder puts first. A step is microbatched only when every DP rank votes for it.

How it was found: GLM-5.3 RL with P/D disaggregation; prefill DP32/EP32 with DeepEP high-throughput, DBO, cudagraph_mode NONE, prefix caching, vLLM 0.29.0 on H200. One prefill rank hit the assertion after ~28 h; the other 31 then timed out in DeepEP dispatch and failed with torch.cat(): expected a non-empty list of Tensors.

Fix: DeepSeekV32IndexerDecodeMetadata records the SM count and block size the schedule was planned with. get_schedule_metadata() returns the planned schedule when DeepGEMM's current SM count matches; otherwise it re-plans with the same inputs and caches the result per SM count. The planned count is stored rather than derived from the tensor shape because SM90 next_n=4 plans num_sms // 2 + 1 schedule rows. update_draft_decode_metadata rewrites seq_lens in place, so it clears the cache. In FULL-graph DBO capture the re-plan kernel is captured and reads the persistent seq_lens/indices buffers; eager warmups under the same context JIT-compile it before capture.

Alternatives considered

1. Plan for the reduced SM count in the builder: the builder would need to replicate create_sm_control_context (all2all max_sms_used(), VLLM_DBO_COMM_SMS, the ROCm override) for both runners, and anything else that changes DeepGEMM's SM count would break it again.
2. Skip set_num_sms in SM control: removes DBO's compute/communication SM partitioning for every DeepGEMM GEMM.
3. Re-plan inside the fp8_fp4_paged_mqa_logits wrapper: no per-forward object to cache on, so every indexer layer would re-plan.
4. Build both schedules up front: an extra kernel every step even when DBO does not split.

Not affected: the sparse-MQA schedule of DeepSeek-V4.1 (planned inside the forward), FlashMLA scheduler metadata, and prefill fp8_fp4_mqa_logits.

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

[github-actions]

👋 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.

🚀

[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.

[zyongye]

why can not we build metadata based on total - VLLM_DBO_COMM_SMS in the first place? Instead of changing it at forward time?

[aniravich]

why can not we build metadata based on total - VLLM_DBO_COMM_SMS in the first place? Instead of changing it at forward time?

@zyongye I considered that, so:

• Whether a step is microbatched is only known per step, and the builder doesn't know. With --enable-dbo, DP ranks decide each step whether to split, and steps that don’t split run with all SMs. The microbatch-0 builder (get_metadata_builder(ubid or 0)) builds metadata for both split and unsplit steps, and build() gets no microbatch flag. A schedule planned for the reduced count up front would hit the same _schedule_meta_size == num_sms + 1 assertion on every unsplit step. Doing it at build time would mean passing a “microbatched / compute SM count” argument from the runner into the builder, in both MRV1 and MRV2.

• The compute SM count actually isn’t just total - VLLM_DBO_COMM_SMS. create_sm_control_context uses min(VLLM_DBO_COMM_SMS, all2all_manager.max_sms_used()), forces 0 on ROCm with DeepEP HT, and only calls deep_gemm.set_num_sms when DeepGEMM is available and the comm SM count is > 0. The builder would need to duplicate that and stay in sync with it.

Re-planning at the call site uses the SM count DeepGEMM is actually running with (get_num_sms()), so it stays correct whichever code path changes it (The cost ofc is one small metadata kernel per microbatch, on the first indexer layer.)

[mergify]

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

[zyongye]

Thanks for the reply.

First of all, the metadata build is done at every engine step. The metadata don't reuse unless we explicitly cache them.

The microbatch-0 builder (get_metadata_builder(ubid or 0)) builds metadata for both split and unsplit steps, and build() gets no microbatch flag.

My sense is that for we can't get the exact SM count at metadata build time because the model runner is not exposed such value?

If that's the case, I think this approach make sense to me. But I would double check if get_paged_mqa_logits_metadata implement caching already so we can save a lot of code in the engine.