fix(ds4): reuse MHC autotuning output buffers

[voipmonitor]

Behavior and reason

Status: implemented; native correctness and startup preparation qualified.

DeepSeek V4 MHC autotuning owns one output-buffer set per candidate and reuses
it during priming, timing and CUDA graph capture. The existing native kernels
already accept caller-provided outputs. Serving dispatch, arithmetic,
precision, scratch ownership and the candidate search space are unchanged.

Without these buffers, functional output allocation creates additional copies
inside each candidate's graph pool. The tuner budgets primed residency before
those copies exist. On TP2 RTX PRO 6000 Workstation, DSpark K5 and a 4096-token
budget, post_pre.m4096 exhausted VRAM with 18 measured candidates plus a
carried candidate: each primed candidate held 336,199,680 bytes, then each graph
capture reserved another 379,584,512 bytes. A 128 MiB output allocation failed
with 74.88 MiB driver-free. Reusing the outputs allows all 160 weight-stage
declarations to prepare without changing the tuner or reducing the search.

Validation

Native SM120, CUDA 13.4.1, PyTorch 2.14, B12X f832e568c078:

python -m pytest -q tests/model_executor/kernels/test_b12x_linear.py \
  -k 'mhc_preparation_reuses_outputs or mhc_declares_the_operations'

• 13 tests pass: four operations at 3/4096 tokens and existing declaration checks.
• The unmodified vLLM source fails the output-pointer reuse reproducer.
• Reused outputs exactly match the same native kernel with functional allocation.
• Changed-input graph replay overwrites poisoned outputs, keeps storage pointers
  stable and adds no live allocation; comparisons use zero tolerance.
• All pre-commit hooks pass for the changed files.

The DS4 TP2 startup check also needs a separate collective-owner lifetime fix
and the TileLang runtime dependency. Weight preparation succeeds, but serving
is not qualified: the 1M-token KV admission check has only 6.14 GiB available
against 6.95 GiB required. No end-to-end speedup or memory-capacity claim is made.

No existing open PR was found for this preparation-output ownership correction.
The change is separate from serving-scratch reuse and graph-aware MoE tuning.
AI-assisted implementation and validation; submitted for maintainer review.

Summary by CodeRabbit

• Performance

  • Improved CUDA execution for B12x mHC workloads by reusing prepared output buffers across repeated operations.
  • Reduced unnecessary memory allocation and copying during optimized execution paths.

• Reliability

  • Improved consistency between standard execution and CUDA graph replay, including stable output handling across supported token counts.

[coderabbitai]

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: b5b1cb1d-caaf-4983-bd86-847ebc9be6c1

📥 Commits

Reviewing files that changed from the base of the PR and between 10c53de and 2b2fa50.

📒 Files selected for processing (2)

• tests/model_executor/kernels/test_b12x_linear.py
• vllm/models/deepseek_v4/nvidia/b12x.py

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

---

📝 Walkthrough

Walkthrough

B12x mHC preparation now preallocates output tensors and passes them to the relevant implementations. New CUDA tests cover all mHC operations, two token counts, repeated preparation, and CUDA graph replay.

Changes

mHC output reuse

Layer / File(s)	Summary
Preallocated mHC outputs vllm/models/deepseek_v4/nvidia/b12x.py	B12xMHCResidual._prepare_call allocates output tensors and forwards them to the pre, post_pre, post_pre_bf16, and post implementations.
CUDA reuse regression coverage tests/model_executor/kernels/test_b12x_linear.py	Parameterized CUDA coverage checks output addresses, native results, memory allocation, and exact CUDA graph replay results for four operations and two token counts.

Priority: ⬇️ Low

Estimated code review effort: 3 (Moderate) | ~25 minutes

Change: Bug fix

Suggested reviewers: lukealonso

Merge Risk: ⚪ Minimal · up to 2b2fa

The output-reuse change preserves the preparation operation contracts across all covered branches, with no concrete merge-blocking risk identified.

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

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

✅ Passed checks (4 passed)

Check name	Status	Explanation
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 clearly and concisely describes the main change: reusing MHC autotuning output buffers in DeepSeek V4.

• Fix all pre-merge checks with AI

✨ Finishing Touches 💡 1

📝 Generate docstrings 💡

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch fix/ds4-mhc-preparation-outputs

---

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]

@coderabbitai review

[coderabbitai]

⚠️ Action not completed

Already reviewed the last commit. Use @coderabbitai full review to rerun a review of the entire changeset.

Note: CodeRabbit is an incremental review system and does not re-review already reviewed commits. This command is applicable only when automatic reviews are paused.

[voipmonitor]

Combined serving qualification with #783 and #784 passes on stock RTX PRO 6000 Workstation GPUs, TP2, DSpark K5 and a 4096-token batch budget. MHC preparation completes without OOM; full/piecewise capture, four arithmetic/prefix checks and all six warmed C1/C8 cells pass. Medians: C1 194.82 tok/s, C8 679.58 tok/s, uncached 32K prefill 13,910 tok/s. This is qualification of the source composition, not an isolated kernel speedup claim. DS4 Vision remains under test.