Fix GLM DCP global top-k MTP

[voipmonitor]

Summary

This replaces the previous PR contents with one clean commit for the validated fast GLM-5.2 TP8/DCP4/MTP3 path on current dev/dark-devotion.

The fix keeps the B12X DCP global-topk path rank-consistent during graph warmup and makes the MTP draft path provide the same top-k score buffer that the target GLM path already provides.

Changes included:

• Keep empty DCP prefill chunks for the B12X sparse-indexer global-topk path instead of dropping them on ranks with no local KV.
• Prewarm the actual B12X DCP global-topk merge path, not a shape-only placeholder, so graph/JIT ordering is stable before MLA attention collectives begin.
• Keep no-local-KV B12X top-k chunks in the same global merge path with -inf scores and mask results back to -1.
• Prewarm B12X sparse MLA extend kernels using caller-owned scratch via plan.bind(scratch=...), matching the vLLM eager binding rule.
• Allocate and pass topk_scores_buffer from the MTP draft layer when DCP + B12X sparse indexer is active. Without this, the latest B12X global-topk stack failed or ran in non-comparable slow paths for MTP.

No B12X source patch is included in this PR. The validated run uses B12X 5af873a7b6c81fbf533ef96bede13fbf4744ad2a unchanged.

Validation

Clean image built through the standard blackwell-llm-docker pipeline:

voipmonitor/vllm:dark-devotion-df8ad3b-b12x5af873a-mtptopkscores-cu132-20260621

Build inputs:

• vLLM: this PR, 000807e2b0e33277ac6b3ae51ae2e52d8472c9ab
• Base: dev/dark-devotion at 4e4a0b91a73d474374e8e5da528a24bb6a16b0eb
• B12X: 5af873a7b6c81fbf533ef96bede13fbf4744ad2a
• FlashInfer: 9c5ed7c194e7412780862491742fc655daaad6ac
• torch: 2.12.0+cu132
• local NCCL: 2.30.4
• cublas runtime symlinked to 13.4.1.2

Validated runtime matching the fast reference:

• GLM-5.2 NVFP4
• TP8 / DCP4 / MTP3
• --dcp-comm-backend ag_rs
• --attention-backend B12X_MLA_SPARSE
• --moe-backend b12x
• B12X_MOE_FORCE_A16=1
• VLLM_USE_B12X_SPARSE_INDEXER=1
• VLLM_DCP_GLOBAL_TOPK=1
• VLLM_DCP_SHARD_DRAFT=1
• FP8 KV
• --max-model-len 256000
• --max-num-seqs 4
• --max-num-batched-tokens 8192
• --max-cudagraph-capture-size 24

Measured with /mnt/test.py --port 5543 -L:

• 4 samples
• mean 128.18 tok/s
• median 127.91 tok/s
• min 126.17 tok/s
• max 130.73 tok/s
• CJK 0

Reference PR30 image on the same host measured /mnt/test.py --port 5544 -L:

• 25 samples
• mean 119.96 tok/s
• median 120.27 tok/s
• max 128.47 tok/s
• CJK 0

Engine logs for the fixed latest stack showed MTP acceptance around 0.95 / 0.83 / 0.64-0.70, average draft acceptance about 78-82%.

Notes

The B12X vLLM bindings remain eager and caller-scratch-owned. This PR does not introduce workspace/arena ownership or cached workspace binding in the vLLM path.

Summary by CodeRabbit

• Performance Improvements

  • Enhanced sparse attention synchronization and warmup for distributed computing setups.
  • Optimized kernel initialization procedures for multi-GPU environments.

• Bug Fixes

  • Improved handling of distributed prefill chunk metadata to prevent early returns in specific multi-GPU scenarios.
  • Enhanced buffer allocation logic for multi-GPU sparse attention operations.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: 920ce89f-7bbd-494f-b322-368170cb51e2

📥 Commits

Reviewing files that changed from the base of the PR and between 000807e and 79f154c.

📒 Files selected for processing (4)

• vllm/model_executor/layers/sparse_attn_indexer.py
• vllm/model_executor/models/deepseek_mtp.py
• vllm/v1/attention/backends/mla/b12x_mla_sparse.py
• vllm/v1/attention/backends/mla/indexer.py

🚧 Files skipped from review as they are similar to previous changes (2)

• vllm/v1/attention/backends/mla/indexer.py
• vllm/model_executor/models/deepseek_mtp.py

---

📝 Walkthrough

Walkthrough

Adds _sync_dcp_warmup() and one-time prewarm helpers for B12X DCP top-k merge and extend kernels in sparse_attn_indexer.py and b12x_mla_sparse.py. Introduces _convert_b12x_dcp_local_topk_to_global() and re-routes prefill and decode top-k handling based on the dcp_global_topk flag. Allocates topk_scores_buffer when distributed context parallelism is enabled, and allows empty prefill chunks to propagate when DCP global top-k is active.

Changes

B12X DCP top-k warmup and routing

Layer / File(s)	Summary
DCP sync helper and prewarm kernel infrastructure vllm/model_executor/layers/sparse_attn_indexer.py, vllm/v1/attention/backends/mla/b12x_mla_sparse.py	Adds _sync_dcp_warmup() (CUDA sync + optional dist.barrier on DCP device group). Reworks _prewarm_b12x_dcp_topk_merge() to loop over q_rows sizes calling _merge_b12x_dcp_topk() then _sync_dcp_warmup(). Adds _EXTEND_PREWARM_DONE deduplication set and _prewarm_extend_kernels_once() in b12x_mla_sparse.py, which allocates dummy tensors, binds the extend plan, calls _sparse_mla_extend_forward, and syncs DCP ranks.
Local-to-global conversion helper and prefill/decode routing vllm/model_executor/layers/sparse_attn_indexer.py	Adds _convert_b12x_dcp_local_topk_to_global() using triton_convert_dcp_local_topk_to_global. Re-routes prefill dcp_global_topk branch: B12X calls _merge_b12x_dcp_topk() and masks invalid scores into topk_indices; non-B12X calls _dcp_global_topk_remap(). B12X prefill and decode paths now dispatch conditionally: merge when dcp_global_topk is true, otherwise convert local-to-global.
topk_scores_buffer allocation and empty chunk propagation vllm/model_executor/models/deepseek_mtp.py, vllm/v1/attention/backends/mla/indexer.py	Imports use_b12x_sparse_indexer and conditionally allocates topk_scores_buffer in DeepSeekMultiTokenPredictorLayer when DCP is enabled and B12X sparse indexer is active; wires the buffer into DeepseekV2DecoderLayer. Replaces the unconditional total_seq_lens == 0 early-return in build_prefill_chunk_metadata with a keep_empty_dcp_chunk flag (gated on VLLM_USE_B12X_SPARSE_INDEXER, dcp_world_size > 1, and VLLM_DCP_GLOBAL_TOPK). Returns None only when neither condition holds.

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~45 minutes

🚥 Pre-merge checks | ✅ 3 | ❌ 2

❌ Failed checks (1 warning, 1 inconclusive)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 13.33% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.
Title check	❓ Inconclusive	The PR title 'Fix GLM DCP global top-k MTP' is too vague and lacks specificity about the actual changes. While it mentions three related components (GLM, DCP, MTP), it doesn't clearly communicate the core improvements: empty chunk preservation, global-topk prewarming, sparse kernel binding, and score buffer allocation.	Consider a more descriptive title such as 'Fix DCP global top-k prewarming and score buffer allocation for B12X sparse indexer' to better summarize the primary changes.

✅ Passed checks (3 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

📝 Generate docstrings

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch codex/dark-devotion-dcp4-mtp3-globaltopk-fix-20260621

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

Actionable comments posted: 1

🧹 Nitpick comments (2)

vllm/v1/attention/backends/mla/b12x_mla_sparse.py (1)

683-706: 💤 Low value

Near-duplicate of _sync_dcp_warmup in sparse_attn_indexer.py.

This method is very similar to the module-level _sync_dcp_warmup() in sparse_attn_indexer.py, with minor differences:

• This passes device_ids=[torch.cuda.current_device()] to barrier
• This checks self.device.type == "cuda" vs current_platform.is_cuda()

Consider extracting a shared helper to avoid drift between the two implementations, or document why the differences are intentional.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@vllm/v1/attention/backends/mla/b12x_mla_sparse.py` around lines 683 - 706,
The _sync_dcp_warmup method in this file contains duplicate logic with a similar
method in sparse_attn_indexer.py, with minor differences in how device checking
is done and parameters passed to the barrier call. Extract a shared helper
function in a common utility module that can be used by both implementations to
maintain consistency and prevent future drift, passing any device-specific
parameters as arguments to the helper. Alternatively, if the differences between
the two implementations are intentional, add clear comments documenting why each
implementation differs in its device checking approach and barrier parameter
passing.

vllm/v1/attention/backends/mla/indexer.py (1)

960-966: Extract VLLM_DCP_GLOBAL_TOPK parsing to a shared utility to prevent duplication.

The VLLM_DCP_GLOBAL_TOPK parsing at lines 963-965 duplicates the logic in _dcp_global_topk_requested() from sparse_attn_indexer.py. A circular import exists (sparse_attn_indexer.py imports from this file), making a direct import infeasible. Extract this logic to a shared utility function in envs.py or a common module to avoid drift if either is updated independently.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@vllm/v1/attention/backends/mla/indexer.py` around lines 960 - 966, The
parsing logic for the VLLM_DCP_GLOBAL_TOPK environment variable is duplicated
between the keep_empty_dcp_chunk assignment in this file and the
_dcp_global_topk_requested() function in sparse_attn_indexer.py. Create a new
shared utility function in envs.py that encapsulates the VLLM_DCP_GLOBAL_TOPK
parsing logic (checking if the retrieved value is in the set of "1", "true",
"yes", "on" after lowercasing), then replace the inline
os.environ.get("VLLM_DCP_GLOBAL_TOPK", "1").lower() condition in the
keep_empty_dcp_chunk assignment with a call to this new utility function, and
update _dcp_global_topk_requested() to use the same utility to prevent drift
between implementations.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@vllm/model_executor/layers/sparse_attn_indexer.py`:
- Around line 443-446: The code directly calls dist.barrier(group=device_group)
which violates the documented pattern against using device_group for barriers
since NCCL internally creates GPU tensors that can interfere with GPU state.
Replace the dist.barrier(group=device_group) call with dcp_group.barrier()
instead, which internally uses cpu_group as the established convention dictates.
Remove the direct device_group barrier invocation and use the dcp_group barrier
method to maintain consistency with the GroupCoordinator pattern.

---

Nitpick comments:
In `@vllm/v1/attention/backends/mla/b12x_mla_sparse.py`:
- Around line 683-706: The _sync_dcp_warmup method in this file contains
duplicate logic with a similar method in sparse_attn_indexer.py, with minor
differences in how device checking is done and parameters passed to the barrier
call. Extract a shared helper function in a common utility module that can be
used by both implementations to maintain consistency and prevent future drift,
passing any device-specific parameters as arguments to the helper.
Alternatively, if the differences between the two implementations are
intentional, add clear comments documenting why each implementation differs in
its device checking approach and barrier parameter passing.

In `@vllm/v1/attention/backends/mla/indexer.py`:
- Around line 960-966: The parsing logic for the VLLM_DCP_GLOBAL_TOPK
environment variable is duplicated between the keep_empty_dcp_chunk assignment
in this file and the _dcp_global_topk_requested() function in
sparse_attn_indexer.py. Create a new shared utility function in envs.py that
encapsulates the VLLM_DCP_GLOBAL_TOPK parsing logic (checking if the retrieved
value is in the set of "1", "true", "yes", "on" after lowercasing), then replace
the inline os.environ.get("VLLM_DCP_GLOBAL_TOPK", "1").lower() condition in the
keep_empty_dcp_chunk assignment with a call to this new utility function, and
update _dcp_global_topk_requested() to use the same utility to prevent drift
between implementations.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: 6b16eb0a-c3ee-4e28-ac6c-870d4f48c478

📥 Commits

Reviewing files that changed from the base of the PR and between 4e4a0b9 and df8ad3b202c84937a23cfa9d93f7a3677da8ecde.

📒 Files selected for processing (3)

• vllm/model_executor/layers/sparse_attn_indexer.py
• vllm/v1/attention/backends/mla/b12x_mla_sparse.py
• vllm/v1/attention/backends/mla/indexer.py

[coderabbitai]

Actionable comments posted: 1

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@vllm/v1/attention/backends/mla/b12x_mla_sparse.py`:
- Around line 698-701: The barrier call at lines 698-701 is using dist.barrier()
with device_group parameter, which violates the established pattern and causes
NCCL internal GPU tensor allocation issues. Replace the
dist.barrier(group=device_group, device_ids=[torch.cuda.current_device()]) call
with dcp_group.barrier() which properly routes the barrier operation through the
CPU group instead, avoiding the problematic NCCL internal GPU tensor creation.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: 60b946da-06ad-4c19-99af-1619884fa738

📥 Commits

Reviewing files that changed from the base of the PR and between df8ad3b202c84937a23cfa9d93f7a3677da8ecde and 000807e.

📒 Files selected for processing (4)

• vllm/model_executor/layers/sparse_attn_indexer.py
• vllm/model_executor/models/deepseek_mtp.py
• vllm/v1/attention/backends/mla/b12x_mla_sparse.py
• vllm/v1/attention/backends/mla/indexer.py

🚧 Files skipped from review as they are similar to previous changes (1)

• vllm/v1/attention/backends/mla/indexer.py