[GG] perf(dcp): preallocate budgeted CKV layer prefetch

[voipmonitor]

Summary

Preallocate a bounded ring of full-CKV DCP gather workspaces before KV-cache memory planning, then use it to overlap the next eligible MLA layer gather with current-layer compute.

The feature is controlled by:

• VLLM_B12X_MLA_CKV_PREFETCH_DEPTH (default 1; 0 keeps the synchronous path)
• VLLM_B12X_MLA_CKV_PREFETCH_WORKSPACE_MIB (default 1024 MiB per lane)

The requested depth is reduced automatically when the ring would exceed the configured budget. Pool exhaustion is a hard error rather than an implicit allocation.

Why this version

The earlier stacked implementation in #160 allocated lookahead buffers lazily, after vLLM had already measured non-KV memory. That made the reported KV budget optimistic and could fail later when the ring was first materialized.

This PR is a direct, standalone change against dev/gilded-gnosis. It creates the pool during backend initialization, before KV planning, and keeps target/draft plus DBO lanes isolated. It contains no replicated-indexer or sparse selected-record CKV changes.

Measured effect

GLM-5.2 NVFP4, RTX PRO 6000 Blackwell, TP8/DCP4/MTP0, A16, full-CKV gather:

Context	Depth 0	Depth 1	Delta
64k	5,656 tok/s	5,778 tok/s	+2.16%
400k	4,666.5 tok/s	5,033.1 tok/s	+7.86%

At the measured 400k profile, the bounded ring increased non-KV allocation from 410.2 to 738.4 MiB/GPU and reduced KV capacity by 25,344 tokens (1.14%). Depths 2 and 3 did not improve throughput, so depth 1 remains the default while larger depths stay explicit experiments.

After review lifecycle hardening, the final clean-branch TP8/DCP4 64k smoke reached 5,355 tok/s versus 5,367 tok/s before the review fix (-0.22%, noise), with 2,185,216 KV tokens and the two-slot depth-1 pool active.

Validation

• 63 passed:
  • tests/v1/attention/test_b12x_ckv_prefetch_policy.py
  • tests/v1/attention/test_b12x_mla_dcp_workspace.py
• ruff check and ruff format --check pass for all modified files.
• E2E TP8/DCP4/MTP0 startup, correctness response, and 64k standalone prefill pass with InstantTensor.\n- Regression coverage includes KV-cache replacement, workspace resize/retirement, pool-slot reuse, transient WorkspaceManager views, and irregular synchronous fallback ordering.

This supersedes only the prefetch portion of #160. It does not supersede #159 (replicated indexer) or #161 (selected-record decode CKV).

Summary by CodeRabbit

• Performance

  • Improved B12X MLA prefill performance by adding configurable CKV prefetching across layers.
  • Added asynchronous lookahead gathering to reduce wait time during attention processing.
  • Added safe handling for CUDA stream capture and workspace resizing scenarios.

• Configuration

  • Added environment variables to configure prefetch depth and workspace memory.
  • Prefetch settings automatically respect available workspace limits.

• Compatibility

  • Extended CKV append handling for FP8 and NVFP4 MLA configurations.

[coderabbitai]

📝 Walkthrough

Walkthrough

Adds configurable B12X MLA CKV prefetch depth and workspace budgets, persistent ring-backed registry state, asynchronous lookahead gathers in forward_mqa, expanded NVFP4 handling, and comprehensive policy and lifecycle tests.

Changes

B12X CKV Prefetch

Layer / File(s)	Summary
Configuration and workspace provisioning vllm/envs.py, vllm/v1/attention/backends/mla/b12x_mla_sparse.py	Adds environment-backed depth and workspace budgets, validates support, caps depth, and provisions persistent ring workspace storage.
Registry and metadata integration vllm/v1/attention/backends/mla/b12x_mla_sparse.py	Adds workspace identity tracking, per-lane prefetch state, ring views, gather streams, and metadata registry propagation.
Forward-path lookahead gathers vllm/v1/attention/backends/mla/b12x_mla_sparse.py	Updates CKV gathering to consume registry state, wait for pending work, schedule lookahead gathers, and support nvfp4_ds_mla.
Policy and lifecycle validation tests/v1/attention/test_b12x_ckv_prefetch_policy.py, tests/v1/attention/test_b12x_mla_dcp_workspace.py	Tests sizing, target selection, workspace isolation, stream recovery, registry retirement, pool exhaustion, capture fallback, and formatting-only adjustments.

Estimated code review effort: 4 (Complex) | ~60 minutes

Sequence Diagram(s)

sequenceDiagram
  participant forward_mqa
  participant CKVPrefetchStateRegistry
  participant CKVPrefetchWorkspacePool
  participant GatherStream
  forward_mqa->>CKVPrefetchStateRegistry: resolve layer prefetch state
  CKVPrefetchStateRegistry->>CKVPrefetchWorkspacePool: acquire ring workspace
  forward_mqa->>GatherStream: schedule lookahead gather
  GatherStream-->>forward_mqa: record completion event
  forward_mqa->>CKVPrefetchStateRegistry: wait or consume pending gather

Loading

Possibly related PRs

• local-inference-lab/vllm#160: Implements the same B12X native CKV prefetch registry, ring storage, and related tests.
• local-inference-lab/vllm#110: Introduces the CKV gather path that this prefetch flow extends.
• local-inference-lab/vllm#164: Updates the same metadata and speculative execution path in b12x_mla_sparse.py.

Suggested reviewers: lukealonso, koush

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
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.
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title accurately summarizes the main change: budgeted preallocation of CKV layer prefetch for DCP.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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

[coderabbitai]

Actionable comments posted: 4

🤖 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 1605-1608: Update reset_kv_cache_binding_state() to reset the
builder-owned _CKVPrefetchStateRegistry, clearing its layer_caches and
pending_layers entries when KV caches are replaced. Ensure subsequent
for_workspace calls cannot reuse state referencing profiling or freed cache
tensors, while preserving the existing legacy attribute cleanup.
- Around line 364-369: The close method releases ckv_workspace_slot before
guaranteeing completion of side-stream writes. Update close to synchronize the
relevant completion event, or record the slot view on the owning stream, before
calling workspace_pool.release; keep begin_step on the normal path and preserve
clearing ckv_workspace_slot and ckv_workspace after safe release.
- Around line 2050-2059: The staging layout in the ring-buffer setup must avoid
sharing local_buffer across buf_idx values. Update the records/local_buffer
addressing around _dcp_gather_ckv so each ring slot stages into its own
_ckv_local_capacity region, while preserving the per-slot gathered_buffer
layout; alternatively, ensure the synchronous fallback gather is explicitly
ordered after pending side-stream work.
- Line 316: The CKV prefetch state currently tracks liveness through per-call
storage wrappers and private storage identity. Update the workspace tracking
around `_CKVPrefetchState` and its callers to retain `weakref.ref(workspace)`
directly, and derive identity from documented tensor/storage metadata instead of
`storage._cdata`; apply the change at both the workspace registration logic and
the `workspace_storage_ref` assignment so state remains live while the workspace
allocation is tracked.

🪄 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: cfe8c514-e86b-4e72-ac5d-b7a1fc0d5e3c

📥 Commits

Reviewing files that changed from the base of the PR and between 89b4a98 and cb5d9ee.

📒 Files selected for processing (4)

• tests/v1/attention/test_b12x_ckv_prefetch_policy.py
• tests/v1/attention/test_b12x_mla_dcp_workspace.py
• vllm/envs.py
• vllm/v1/attention/backends/mla/b12x_mla_sparse.py

[voipmonitor]

Review follow-up is in affff57c0d.

• Clears builder-owned prefetch registries on KV-cache replacement.
• Synchronizes pending side-stream writers before a persistent pool slot is recycled; the normal path remains device-side Event.wait().
• Orders irregular synchronous fallback gathers behind pending users of the shared local staging area without increasing persistent VRAM.
• Replaces private storage._cdata identity with public storage/tensor metadata. Backing-storage liveness is retained intentionally because WorkspaceManager returns transient borrowed tensor views.

Validation:

• 63 passed across CKV prefetch policy and DCP workspace tests.
• Ruff check and format clean.
• TP8/DCP4/MTP0, InstantTensor, depth 1, 64k prefill: 5,355 tok/s versus 5,367 tok/s before the review fix (-0.22%, noise).
• KV capacity: 2,185,216 tokens.
• End-to-end correctness response: 80235; no runtime errors.