fix(ds4): reserve full B12X compressed MLA workspace envelope - #227
Conversation
Assisted-by: Claude Assisted-by: OpenAI Codex Signed-off-by: joninco <jon@bullpoint.org>
Share capacity calculations across reservation and metadata paths, cover the split-contract row envelope, and tighten the real-planner regression test. Assisted-by: Claude Fable Assisted-by: OpenAI Codex Signed-off-by: joninco <jon@bullpoint.org>
|
👋 Hi! Thank you for contributing to the vLLM project. 💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in PRs do not trigger a full CI run by default. Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging. To run CI, PR reviewers can either: Add If you have any questions, please reach out to us on Slack at https://slack.vllm.ai. Agent GuidelinesIMPORTANT: 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. 🚀 |
📝 WalkthroughWalkthroughChangesCompressed MLA workspace geometry
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant B12x
participant GeometryHelpers
participant SparkinferPlanner
participant WorkspaceManager
B12x->>GeometryHelpers: calculate compressed widths and query chunks
B12x->>SparkinferPlanner: request scratch plan
SparkinferPlanner-->>B12x: return runtime workspace requirements
B12x->>WorkspaceManager: reserve workspace envelope
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
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. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (1)
vllm/v1/attention/backends/mla/compressor_utils.py (1)
16-54: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd Args/Returns sections to the new public helpers.
get_c128a_topk_width,get_dspark_swa_index_width, andget_compressed_mla_max_q_chunkshave one-line docstrings withoutArgs:/Returns:sections.get_compressed_mla_split_caphas no docstring at all. These four functions form the new shared geometry contract consumed bysparse_mla.py,sparse_swa.py, andb12x.py, so documenting their parameters and return values helps downstream maintainers.📝 Proposed docstring additions
def get_c128a_topk_width(max_model_len: int, compress_ratio: int) -> int: - """Return C128 indexed width padded for FlashMLA B_TOPK divisibility.""" + """Return C128 indexed width padded for FlashMLA B_TOPK divisibility. + + Args: + max_model_len: Maximum model context length in tokens. + compress_ratio: Compression ratio applied to the KV cache. + + Returns: + The compressed top-k width, rounded up to the alignment boundary. + """ compressed_width = cdiv(max_model_len, compress_ratio) return cdiv(compressed_width, _C128A_TOPK_ALIGNMENT) * _C128A_TOPK_ALIGNMENT def get_dspark_swa_index_width( window_size: int, num_speculative_tokens: int, ) -> int: - """Return the padded width of non-causal DSpark SWA indices.""" + """Return the padded width of non-causal DSpark SWA indices. + + Args: + window_size: Sliding-window attention size. + num_speculative_tokens: Number of speculative tokens for DSpark. + + Returns: + The padded non-causal index width. + """ width = max(int(window_size), 0) + max(int(num_speculative_tokens), 0) return cdiv(width, _DSPARK_SWA_INDEX_ALIGNMENT) * _DSPARK_SWA_INDEX_ALIGNMENT def get_compressed_mla_split_cap(width: int) -> int: + """Return the maximum split-chunk cap for a given compressed-MLA width. + + Args: + width: Combined SWA and indexed width. + + Returns: + The split-chunk cap. + """ return max(1, cdiv(max(int(width), 1), _COMPRESSED_MLA_SPLIT_ALIGNMENT))As per coding guidelines, "Use Google-style docstrings in Python code, with
Args:/Returns:/Raises:sections instead of reStructuredText/Sphinx fields such as:param:,:return:, and:rtype:."🤖 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/compressor_utils.py` around lines 16 - 54, Update the docstrings for get_c128a_topk_width, get_dspark_swa_index_width, get_compressed_mla_split_cap, and get_compressed_mla_max_q_chunks to use Google-style Args and Returns sections describing each parameter and the integer value returned; preserve the existing behavior and calculations.Source: Coding guidelines
🤖 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.
Nitpick comments:
In `@vllm/v1/attention/backends/mla/compressor_utils.py`:
- Around line 16-54: Update the docstrings for get_c128a_topk_width,
get_dspark_swa_index_width, get_compressed_mla_split_cap, and
get_compressed_mla_max_q_chunks to use Google-style Args and Returns sections
describing each parameter and the integer value returned; preserve the existing
behavior and calculations.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: ab88d2ff-4211-4446-889a-a834c3d24409
📒 Files selected for processing (5)
tests/models/deepseek_v4/test_b12x_compressed_mla_workspace.pyvllm/models/deepseek_v4/nvidia/b12x.pyvllm/models/deepseek_v4/sparse_mla.pyvllm/v1/attention/backends/mla/compressor_utils.pyvllm/v1/attention/backends/mla/sparse_swa.py
|
Functional review is clean. I reproduced the reported planner failure exactly: the old reservation plans 321.50 MiB at 2,048 rows, while the runtime peak is 480.40 MiB at 232 rows because the split count rises to 66. The PR test set passes 16/16 against the installed SparkInfer planner, and a TP2/DCP1/K5 long-prompt E2E run completed without a workspace assertion. I also addressed the remaining CodeRabbit docstring nit in commit voipmonitor@6d8687d. GitHub did not permit maintainer push to the contributor fork, so this commit can be cherry-picked onto the PR branch; it is documentation-only. |
3 participants
Closes #224.
Summary
Prevent the DeepSeek-V4 B12X compressed-MLA path from requesting a larger v1 workspace after
lock_workspace().The profiling call already reached
_reserve_dummy_compressed_mla_scratch; the failure was not a skipped warmup. The reserve planned only the singlerows=max_num_batched_tokenspoint, while Sparkinfer's split contract can require more total q-chunk scratch at a smaller runtime row count. This change reserves the complete row/split envelope before graph capture and keeps runtime workspace growth disabled.It also makes the reservation geometry match runtime for:
Confirmed root cause
Instrumentation against the exact reporter image showed that the reserve branch ran:
The failing runtime shape used the same width but a different split regime:
Sparkinfer's scratch contains both row-sized buffers and split intermediates. The latter are governed by:
That product is not maximized at
rows=max_num_batched_tokens. At width 4224, the split count falls from 66 for small chunks to 5 at 2048 rows, so the original single-point reserve under-allocated even though runtime and reserve widths matched.For the 1,048,576-token configuration in #224, real-planner instrumentation gives:
The reported 607.70 MiB and 917.67 MiB requests are both inside the newly reserved envelope.
Changes
max_q_chunkscapacity to the Sparkinfer plan.reserve_bytes >= max(runtime_bytes)over every row count for C1/C4/C128, causal/DSpark, and DCP2 cases.No workspace is unlocked or resized after capture. Kernel inputs, attention math, and output numerics are unchanged.
Validation
Commands run from the PR branch:
The compressed-MLA test was also run inside the exact reporter image with the four patched source files bind-mounted over site-packages:
That image run exercised the real
sparkinfer.attention.compressed_mla.plan()on CPU and swept all runtime row counts without allocating GPU scratch.Model/serving evaluation
A full post-fix DeepSeek-V4 long-context serving run was not performed because the available GPUs were already occupied at approximately 91 GiB/GPU. This patch changes allocation capacity only; it does not change model outputs or numerical kernels. The exact-image real-planner envelope test above validates the failure mechanism and the post-lock allocation bound. A human-run long-context smoke test should be completed before marking this PR ready.
Duplicate-work check
Searched open PRs for
compressed MLA workspace,DeepSeek V4 B12X workspace,workspace locked b12x, and references to #224. No PR implements this fix.PR #161 references the issue number in search results but is materially different: it adds selected-record sparse-CKV decode under DCP and does not modify the DeepSeek-V4 compressed-MLA planner or reserve path.
AI assistance and accountability
AI assistance was used for instrumentation, implementation, test construction, and review: OpenAI Codex performed the primary work, and Anthropic Claude/Fable provided independent review. The final Fable review found no substantive issues after follow-up revisions.
Before moving this PR out of draft, the human submitter must review every changed line, understand and be able to defend the change end-to-end, and run the relevant long-context serving validation.
Summary by CodeRabbit
Bug Fixes
Tests