[KK] Release dense MLA priming buffers after preparation

[voipmonitor]

Behavior

Release the temporary query, output, lengths, page-table and scratch tensors used to prime B12X dense MLA after synchronized preparation. Published plans retain compiled launchers; serving binds its own buffers.

Keeping the priming operands in PreparedCall.owners also kept their allocations alive for the plan's lifetime. This consumes memory needed by KV storage and speculative CUDA graphs. The preparation session synchronizes the priming callback before discarding it, so these operands do not need publication ownership.

No arithmetic, precision, cache layout or public configuration changes. This applies to the B12X dense-MLA backend, not only Kimi-K3.

Validation

• Implemented: two weak-reference regression cases cover BF16 and FP8 KV. They fail before the fix and pass after it; the test retains the published state while requiring priming operands to be collectable.
• Qualified: 18 GPU attention cases cover BF16/FP8 KV, DCP1/DCP16, padded block tables, high physical page IDs, FP32 partials and CUDA Graph replay on RTX PRO 6000 Blackwell, CUDA 13.4, CUTLASS DSL 4.6.2. The installed runtime carries the identical fix at 275b2f96239f0e4330ec4a69d5b9e20eea13422b.
• Qualified serving: Kimi X4T, TP16/DCP16, DSpark K7, CC8, 5.6 GB KV/GPU, three coding generations and eight simultaneous 32k-prompt requests complete without OOM, preemption or restart. The serving composition and evidence are tracked in [KK] Kimi runtime, FP4-CSF and GLM online Spark precision: merge inventory #906.
• Unsupported claim: this patch alone does not establish CC16 feasibility; CC16 still OOMs during lazy NCCL initialization at this KV budget. It does not qualify the full model under CUTLASS 4.7.1.

Evidence: kimi-x4t-serving-20260930/evidence/mla-preparation-ownership-{before,after}.log, mla-graph-reuse-tests.json, and the CC8 receipts linked from #906. AI assistance: OpenAI Codex; no human-review claim is implied.

Summary by CodeRabbit

• Bug Fixes
  • Improved memory handling during dense MLA preparation: temporary resources are released after preparation completes, while prepared calls and their published state remain available.
  • Verified behavior with bfloat16 and float8 cache types.

[coderabbitai]

Navigate logical layers of code changes, visualize relationships, and explore their blast radius.

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Advanced

Run ID: 5b021830-8365-4943-8745-99207d7f9810

📥 Commits

Reviewing files that changed from the base of the PR and between ab86b70 and 5f8a575.

📒 Files selected for processing (3)

• .lil/changes/vllm-b12x-mla-transient-preparation.json
• tests/v1/attention/test_b12x.py
• vllm/v1/attention/backends/mla/b12x_mla.py

Included review availability: This review used your included allowance. Your plan provides up to 1 included review per hour; 0 remain after this review.

---

📝 Walkthrough

Walkthrough

Dense MLA preparation now synchronizes the priming call before discarding temporary operands. PreparedCall no longer retains those operands or the binding. Tests cover their lifetime and the published state.

Changes

Dense MLA preparation

Layer / File(s)	Summary
Synchronized preparation and transient ownership vllm/v1/attention/backends/mla/b12x_mla.py, tests/v1/attention/test_b12x.py, .lil/changes/*	The priming call is synchronized before temporary operands are discarded. PreparedCall no longer retains scratch buffers or the binding. Tests check operand lifetime with bfloat16 and float8 caches and confirm the published state remains available. The release-change record describes the fix and its compatibility statement.

Priority: ➖ Normal

Estimated code review effort: 2 (Simple) | ~10 minutes

Change: Bug fix

Merge Risk: ⚪ Minimal · up to 5f8a5

The change removes redundant retention of temporary preparation buffers while serving continues to use its own buffers. No actionable merge-blocking defect was identified; normal dependency and GPU test checks remain necessary.

Security Architecture Review

Security architecture risk: 🔵 Low · up to 5f8a5

The change shortens temporary GPU-buffer lifetimes without demonstrating new access or privileges. Remaining risk concerns whether outstanding GPU work always finishes before those buffers are released.

Retained concerns
No architecture-level concerns identified.

Security review details

Security Blast Radius

• inferred — The demonstrated change affects buffer lifetime in workers using dense MLA preparation. An incorrect completion-before-release contract could affect GPU memory safety or worker availability, but no such failure or broader tenant, service, or privilege exposure is established.

Trust Boundaries and Controls

• inferred — The inspected change does not add an attacker-controlled source, identity transition, authorization bypass, or new launch sink. Existing serving inputs still pass through the same capacity and sequence-shape checks and are bound separately from priming operands.

Resilience and Maintainability Implications

• observed — Preparation assumes serial borrowing of scratch. Workspace documentation places responsibility for preventing overlapping borrowers on callers; lane reservation is not itself mutual exclusion. This pre-existing constraint remains relevant, while external session ordering and cleanup could not be verified.

Hardening Proposals

• proposed — Validate the pinned preparation runtime's completion-before-discard and publication contracts with real-session lifecycle coverage, including partial failure, interruption, and repeated or concurrent preparation. This would close the ownership proof gap rather than address a demonstrated vulnerability.

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 6 functions across 2 files. (1 skipped: 1 …	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly and concisely describes the main change: releasing dense MLA priming buffers after preparation.
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.

Full details: Docstring Coverage

Explanation

Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 6 functions across 2 files. (1 skipped: 1 unsupported.)

• Fix all pre-merge checks with AI

✨ Finishing Touches 💡 1

📝 Generate docstrings 💡

• Commit to this branch
• Create a new PR

🧪 Generate unit tests (beta)

• Commit to this branch
• Create a new PR

• Autopilot · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

Autopilot is currently an internal CodeRabbit preview.

---

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

[voipmonitor]

Closing: this change is already in dev/karmic-kraken, through Luke's karmic integration. The remaining diff is only formatting or test annotations.