[GG] MLA: preallocate absorbed projection weights before dequant scratch

[voipmonitor]

Summary

Allocate the persistent MLA absorbed projection pair (W_UV and W_UK_T) before materializing temporary dequantization storage. This keeps long-lived tensors out of the allocator holes created by per-layer dequant scratch.

Compared with #148, this clean replacement also:

• limits preallocation to the generic 16-bit absorbed path, so Aiter FP4/FP8 BMM paths do not allocate unused tensors;
• resolves the device from weight, qweight, weight_packed, wrapped base layers, parameters, or buffers instead of assuming kv_b_proj.weight;
• reuses compatible absorbed tensors on reload and keeps prefer_copy=True, preserving CUDA-graph-visible data_ptr values;
• adds focused CPU coverage for quantized layers without .weight, allocation-before-dequant ordering, first-load storage adoption, and reload pointer stability.

Why

The original allocation order materializes each persistent absorbed pair after the dequant scratch has churned the CUDA allocator. The validation attached to #148 measured reserved-but-unallocated memory falling from 0.577 to 0.225 GiB per GPU at GLM-5.2 TP4 geometry when the persistent pair is allocated first.

This PR intentionally supersedes #148 instead of building on it: #148 is stacked on unrelated #145/#146 changes. No upstream open PR matched searches for W_UK_T W_UV preallocate or MLA dequant scratch fragmentation. #154 is related but not a duplicate: it releases B12X-owned source storage after absorption, while this PR controls allocation order. The two branches apply without a source conflict.

CodeRabbit feedback from #148

Accepted and strengthened:

• safe device lookup for quantized modules without .weight;
• prefer_copy=True and reload-safe pointer preservation.

Not ported:

• the KLD table edit in kv-scales/README.md, because that file belongs to the unrelated glm52: calibrated NVFP4 MLA KV outer scales + serve wiring #145 stack and is outside this clean PR.

Validation

ruff check vllm/model_executor/layers/attention/mla_attention.py tests/v1/attention/test_mla_backends.py
ruff format --check vllm/model_executor/layers/attention/mla_attention.py tests/v1/attention/test_mla_backends.py
git diff --check

python -m pytest tests/v1/attention/test_mla_backends.py -q \
  -k 'post_load_preserves_runtime_weight_addresses or post_load_preallocates_quantized_absorbed_weights or absorbed_weight_preallocation_reuses_compatible_storage'
3 passed

Compatibility validation with #154's focused tests: 6 passed.

The predecessor's GLM-5.2 TP4/DCP4 model validation for the same allocation-order change reported:

• reserved-but-unallocated: 0.577 -> 0.225 GiB/GPU;
• 740,255-token KV pool;
• prefill ladder through 157,932 tokens passed;
• decode 106.7 tok/s and MTP acceptance 92.8%, unchanged;
• teacher-forced prefill KLD 0.1366 +/- 0.0042 over 15 runs.

No numerical operation or runtime weight layout changes in this replacement.

AI assistance

OpenAI Codex assisted with implementation and validation. The human submitter must review every changed line and be able to explain and defend the allocator, reload, and device-selection contracts before merge.

[coderabbitai]

Warning

Review limit reached

@voipmonitor, you've reached your PR review limit, so we couldn't start this review.

Next review available in: 34 minutes

Enable usage-based reviews in Billing to review now. Otherwise, wait until the next included review is available.
You're only billed for reviews past your plan's rate limits ($0.25/file).

How can I continue?

After more reviews become available, a review can be triggered using the @coderabbitai review command as a PR comment. Alternatively, push new commits to this PR.

To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based reviews.

How do review limits work?

CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan review availability.

For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, additional reviews become available more gradually as earlier reviews age out of the rolling window.

Please refer docs for additional details.

Review details

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: 215ebe56-1efc-4ab9-b102-6d07fdba0a39

📥 Commits

Reviewing files that changed from the base of the PR and between b07bef7 and b7710b5.

📒 Files selected for processing (2)

• tests/v1/attention/test_mla_backends.py
• vllm/model_executor/layers/attention/mla_attention.py

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch perf/gg-mla-preallocate-absorbed-20260721

---

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]

Rebased the single preallocation commit onto current dev/gilded-gnosis (b07bef7) and resolved the overlap with the merged direct absorbed-BMM path by applying preallocation only inside the existing materialized fallback. Focused CPU tests pass (2 passed, 2319 deselected); Ruff and format checks pass.