[GG] Fix MRV2 high-utilization startup memory profiling

[voipmonitor]

Problem

MRV2 can size KV cache before persistent DCP/JIT kernel resources exist. At
high GPU utilization this can over-allocate KV and fail later during kernel
warmup or CUDA graph capture.

A single model profile is insufficient after warmup moves earlier: it observes
either the transient activation peak or newly persistent resources, not their
simultaneous high-water mark.

Fix

• materialize persistent kernel resources once during automatic memory
  profiling, before KV allocation;
• replay the model profile after warmup so activation workspace is measured on
  top of the persistent baseline;
• keep kernel warmup idempotent and retryable after a failed warmup;
• allow pre-KV runners that have not created optional attn_groups yet;
• avoid creating the synthetic sparse-DCP profile workspace inside CUDA graph
  capture while retaining it during ordinary memory profiling.

This PR does not change the inference hot path, backend selection, or requested
GPU utilization.

Related fixes

This PR addresses persistent-resource ordering only. The full high-utilization
startup fix also needs:

• vLLM [GG] Fix MRV2 profiled CUDA graph pool lifecycle #180 for the profiled CUDA graph pool lifecycle and full graph
  high-water accounting;
• SparkInfer quant(fp8): bridge MXFP8 dense and NVFP4 experts #76 for an exact-capacity, profile-visible persistent PCIe DMA
  default-output workspace.

Those responsibilities remain separate so kernel warmup, graph-pool lifecycle,
and communication-buffer ownership can be reviewed independently.

Validation

• focused unit suites for worker ordering, sparse workspace capture behavior,
  and FlashInfer pre-KV warmup: passed;
• retry after a failed kernel warmup: covered by unit test;
• isolated full stack on 4x RTX PRO 6000 Blackwell:
  TP4/DCP4/MTP3 NF3, MNS 16, graph 64, max model length 350,208, batched tokens
  2,048, GMU 0.98 reached READY with 563,712 KV tokens and completed a
  300,017-token prefill plus decode; health remained OK;
• isolated TP8/DCP2/MTP3 A4 online-MXFP8 ring-DMA gate completed a
  300,066-token prefill plus 512 decode in 65.924 seconds and remained healthy.

The tested graph-pool profile measured 1.00 GiB total, with 0.72 GiB already
retained in the reusable pool. vLLM #180 intentionally reserves that full
high-water mark; the older gross - retained result omitted live
PyTorch-reserved pages and could overstate KV capacity.

Capacity interpretation

This change makes persistent memory visible rather than reclaiming it. A lower
reported KV capacity can therefore be the correct result. The helper's logged
GMU adjustment can recover capacity only when physical headroom exists; it
cannot compensate for genuinely persistent model or DCP allocations.

The earlier 480k/GMU 0.98 failure was not rerun with the complete #180/#76
stack. Only the 350,208-token profile above is validated here, so this PR does
not claim 480k is safe.

AI assistance disclosure

The implementation review, tests, hardware validation, and this write-up were
prepared with OpenAI Codex assistance and reviewed by the human submitter.

[coderabbitai]

📝 Walkthrough

Walkthrough

Changes

The worker now performs guarded one-time kernel warmup and replays model profiling around persistent warmup allocations. MLA sparse profiling skips synthetic workspace allocation during CUDA graph capture. FlashInfer warmup safely handles runners without optional attention metadata.

Kernel warmup and profiling

Layer / File(s)	Summary
Worker kernel warmup and profiling flow vllm/v1/worker/gpu_worker.py, tests/v1/worker/test_gpu_worker.py	Worker warmup is made retryable and idempotent, profiling is replayed around compressor and kernel warmup, and call ordering is tested.
CUDA graph-aware MLA workspace allocation vllm/model_executor/layers/attention/mla_attention.py, tests/v1/attention/test_b12x_mla_dcp_workspace.py	Sparse profile workspace allocation now checks CUDA graph capture state, with tests covering captured and uncaptured execution.
Optional FlashInfer attention metadata handling vllm/model_executor/warmup/kernel_warmup.py, tests/model_executor/test_flashinfer_autotune_cache.py	FlashInfer warmup gates use optional attn_groups access, and tests cover pre-KV runners and updated worker setup.
Estimated code review effort: 3 (Moderate)	~20 minutes

Sequence Diagram(s)

sequenceDiagram
  participant Worker
  participant ModelRunner
  participant Compressor
  participant KernelWarmup
  Worker->>ModelRunner: profile_run()
  Worker->>Compressor: deepseek_v4_compressor_triton_warmup()
  Worker->>KernelWarmup: _warmup_kernels_once()
  Worker->>ModelRunner: profile_run()

Loading

Possibly related PRs

• local-inference-lab/vllm#110: Both modify MLA attention DCP-related workspace control flow.
• local-inference-lab/vllm#131: Both modify sparse MLA profiling workspace allocation and regression coverage.

Suggested reviewers: lukealonso

🚥 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 clearly summarizes the main change: fixing MRV2 startup memory profiling under high utilization.

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

[malaiwah]

Data point from the same high-utilization regime, in case it is useful for sizing this fix.

On 4x RTX PRO 6000 Blackwell (96 GB, no NVLink, all-NODE), serving
madeby561/GLM-5.2-MXFP8-NVFP4-NF3-Hybrid with identical knobs on both images
(TP=4 DCP=4 dcp_comm_backend=a2a, MTP-3, gpu_memory_utilization=0.965,
max_model_len=524288, max_num_seqs=8, max_num_batched_tokens=3072, GRAPH=32,
kv_cache_dtype=nvfp4_ds_mla, OffloadingConnector 128 GiB, load_format=safetensors),
available KV drops by ~2.37 GiB/rank going v19 -> v20, and the 512K config stops booting.

v19 (gilded-gnosis-v19-vllm7ea567a-b12x4cfa530-...-20260718) — boots:

[gpu_worker.py:569] Available KV cache memory: 4.64 GiB
[kv_cache_utils.py:2237] GPU KV cache size: 537,600 tokens
[gpu_worker.py:870] Actual usage is 84.62 GiB for weight, 2.31 GiB for peak activation,
  0.07 GiB for non-torch memory, and 0.55 GiB for CUDAGraph memory.

v20 (gilded-gnosis-v20-vllm5517197-sibe0edca-...-20260725) — fails:

[model_runner.py:366] Model loading took 85.72 GiB and 176.09 seconds
[model_runner.py:968] Estimated MRV2 CUDA graph memory: 0.21 GiB additional
  (0.62 GiB captured, 0.41 GiB retained and counted as non-torch)
[gpu_worker.py:595] Available KV cache memory: 2.27 GiB
ValueError: To serve at least one request with the model's max seq len (524288),
  4.52 GiB KV cache is needed, which is larger than the available KV cache memory (2.27 GiB).

Two contributions, and the profiler is the smaller one:

1. weights +1.10 GiB/rank (84.62 -> 85.72). This is the dominant term. If persistent
   DCP resources (replicated sparse-indexer cache, CKV prefetch workspace) are now
   materialized inside model load, that would explain it — which is why this PR's
   "materialize kernel resources once during automatic memory profiling, before KV
   allocation" ordering seems directly relevant.
2. MRV2 accounting ~0.34 GiB — 0.41 GiB retained as non-torch where v19 charged
   0.07 non-torch + 0.55 CUDAGraph.

The in-log suggestion (--gpu-memory-utilization=0.9673) recovers only ~0.4 GiB of the
~2.37 GiB deficit; closing it by utilization alone would need ~0.99, which OOMs on this path.

Question: after this PR, is the expectation that the profile reports the persistent
DCP baseline more accurately (so this config should be re-tuned down from 512K), or that
some of the +1.10 GiB is transient and gets reclaimed before KV sizing? Happy to re-run
the same A/B against a build with this merged and post the delta.

Note this is specific to large-weight checkpoints already near the ceiling — the
EXL3-TR3 checkpoint on the same v20 image profiles at 77.77 GiB/rank and gets
7.42 GiB KV / 860,928 tokens, no issue.

[voipmonitor]

Answer to the capacity question after the completed investigation: this PR makes persistent kernel/DCP allocations visible before KV sizing; it does not reclaim them. The companion graph-pool fix is now #180, and it also reserves the full graph high-water mark because retained private-pool pages are PyTorch-reserved and were previously omitted by gross - retained. Therefore a lower but boot-safe KV result can be expected. The reported +1.10 GiB weight delta is larger than the graph-accounting change and should be audited separately by allocation owner; it must not be assumed transient. The complete #172/#180/SparkInfer#76 stack passed TP4/DCP4/MTP3 at maxlen 350208 and GMU 0.98 with 563,712 KV tokens, but the 480k configuration from this comment has not been revalidated, so no 480k claim is being made.