[GG] MLA: release absorbed B12X MXFP8 projection weights

[voipmonitor]

Summary

B12X_MLA_SPARSE serves both decode and extend through its absorbed W_UK_T and
W_UV projections. Once those tensors are materialized, an MXFP8
kv_b_proj's checkpoint tensors and temporary B12X packed owner are no longer
used at runtime.

This change:

• adds an explicit, default-off MLA backend capability for releasing the source;
• enables it only for B12X_MLA_SPARSE;
• releases storage only when a B12X MXFP8 packed owner actually exists;
• refuses the release if an MHA prefill backend is active; and
• reconstructs the absorbed weight from a regenerated packed owner during
  sleep/wake or weight reload.

Other quantization formats and other MLA backends remain unchanged.

Scope and prior work

This supersedes only the kv_b_proj half of #146. It is an independent commit
directly on the current dev/gilded-gnosis head. Compared with #146, it does not
clear arbitrary non-BF16 parameters based solely on prefill_backend is None,
does not use .data, and covers the reload lifecycle. Searches of open upstream
and local PRs found no other implementation; #146 is the intentional predecessor
being split.
PR #148 contains #146 in its stacked history, but its unique tip preallocates the
absorbed pair to reduce allocator fragmentation. That is an orthogonal follow-up,
not a second implementation of this storage-lifetime change, and will need a
clean rebase after the predecessor replacement.

The original GLM-5.2 TP4 validation in #146 reported about 290 MiB/GPU reclaimed,
unchanged 157,932-token prefill and decode throughput, unchanged 94.1% MTP
acceptance, and teacher-forced prefill KLD 0.1360 +/- 0.0019 versus a
0.1356 +/- 0.0054 certification band. The target-path storage lifetime is the
same here, while the guards are narrower.

Validation

pytest tests/v1/attention/test_mla_backends.py -q -k \
  'mla_post_load_preserves_runtime_weight_addresses or \
   b12x_mxfp8_mla_post_load_releases_absorbed_sources or \
   release_kv_b_proj_is_inert_without_b12x_mxfp8_owner or \
   reloads_absorbed_weight_from_regenerated_packed_owner'
4 passed

pytest tests/model_executor/model_loader/test_reload.py -q -k \
  'layerwise or capture_layer'
7 passed

ruff check <changed Python files>
ruff format --check <changed Python files>
git diff --check
all passed

AI assistance was used. This PR must not be merged until the human submitter has
reviewed every changed line and can defend the behavior end to end.

[voipmonitor]

@coderabbitai review

[coderabbitai]

✅ Action performed

Review finished.

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.

[coderabbitai]

📝 Walkthrough

Walkthrough

MLA post-load processing now supports regenerating KV-b-proj weights from packed B12X owners and conditionally releasing absorbed sources. Backend capability flags, prefill safeguards, and CPU-only lifecycle tests were added.

Changes

MLA KV-b-proj lifecycle

Layer / File(s)	Summary
KV-b-proj materialization and release helpers vllm/model_executor/layers/attention/mla_attention.py	Adds helpers to materialize KV-b-proj weights from existing or packed sources and release absorbed B12X parameters after fusion.
Backend capability and post-load integration vllm/v1/attention/backend.py, vllm/v1/attention/backends/mla/b12x_mla_sparse.py, vllm/model_executor/layers/attention/mla_attention.py	Adds a default-disabled release capability, enables it for B12X sparse MLA, integrates materialization into post-load processing, and rejects release when MHA prefill is enabled.
Lifecycle validation tests/v1/attention/test_mla_backends.py	Tests runtime address preservation, source release, missing-owner no-op behavior, and packed-owner weight regeneration.

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

Sequence Diagram(s)

sequenceDiagram
  participant MLAProcess as process_weights_after_loading
  participant KVBProj as kv_b_proj
  participant QuantMethod as quant_method
  participant ReleaseHelper as _release_b12x_mxfp8_kv_b_proj

  MLAProcess->>KVBProj: materialize KV-b-proj weight
  KVBProj->>QuantMethod: regenerate from packed owner
  QuantMethod-->>MLAProcess: return materialized weight
  MLAProcess->>ReleaseHelper: release sources when capability allows
  ReleaseHelper-->>MLAProcess: clear absorbed B12X storage

Loading

Possibly related PRs

• local-inference-lab/vllm#136: Related MLA post-load handling for runtime weight address stability and weight-processing contracts.
• local-inference-lab/vllm#146: Related conditional release of absorbed KV-b-proj artifacts.
• local-inference-lab/vllm#148: Related packed KV-b-proj weight lifecycle changes.

Suggested reviewers: lukealonso, koush, yewentao256

🚥 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 accurately summarizes the main change: releasing absorbed B12X MXFP8 projection weights in MLA.

✨ Finishing Touches

📝 Generate docstrings

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch perf/gg-mla-kvb-release-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.}

[coderabbitai]

🧹 Nitpick comments (1)

vllm/model_executor/layers/attention/mla_attention.py (1)

298-304: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Add Args: and Returns: sections to docstrings.

As per coding guidelines, Python code must use Google-style docstrings with explicit Args:/Returns:/Raises: sections. Both of these new helper functions have one-line docstrings that omit these sections.

• vllm/model_executor/layers/attention/mla_attention.py#L298-L304: Add Args: (documenting layer, out_dtype, fallback_device) and Returns: sections.
• vllm/model_executor/layers/attention/mla_attention.py#L324-L325: Add Args: (documenting layer) and Returns: sections.

🤖 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/model_executor/layers/attention/mla_attention.py` around lines 298 -
304, Update the docstrings for _materialize_kv_b_proj_weight and the helper at
the second site in vllm/model_executor/layers/attention/mla_attention.py:298-304
and vllm/model_executor/layers/attention/mla_attention.py:324-325 to Google
style, adding Args: entries for each requested parameter and a Returns: entry
describing the returned tensor/value. Keep the existing behavior unchanged.

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/model_executor/layers/attention/mla_attention.py`:
- Around line 298-304: Update the docstrings for _materialize_kv_b_proj_weight
and the helper at the second site in
vllm/model_executor/layers/attention/mla_attention.py:298-304 and
vllm/model_executor/layers/attention/mla_attention.py:324-325 to Google style,
adding Args: entries for each requested parameter and a Returns: entry
describing the returned tensor/value. Keep the existing behavior unchanged.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: 38ddb06a-2776-4225-a89d-3bfe9185e1e8

📥 Commits

Reviewing files that changed from the base of the PR and between 73e4a8c and b15e560.

📒 Files selected for processing (4)

• tests/v1/attention/test_mla_backends.py
• vllm/model_executor/layers/attention/mla_attention.py
• vllm/v1/attention/backend.py
• vllm/v1/attention/backends/mla/b12x_mla_sparse.py

[voipmonitor]

Addressed the CodeRabbit Google-style docstring nit in 7724c65. The focused MLA lifecycle tests and Ruff/format checks still pass.