[Bugfix][Attention] Preserve post-load tensors across weight reloads

[aoshen02]

Purpose

Preserve attention runtime tensors derived by process_weights_after_loading() across layerwise weight reloads. This fixes stale FlashInfer attention sinks and prevents standard MLA CUDA graphs from retaining obsolete W_UV/W_UK_T addresses.

Root cause

Two post-load paths replaced tensors instead of refreshing their existing storage:

• Standard MLAAttention reassigned derived W_UV and W_UK_T views. CUDA graphs captured before a reload continued to reference the old allocations.
• FlashInfer converted BF16/FP16 sink parameters with self.sinks = self.sinks.to(torch.float32). Later reloads updated the source parameter but not the detached FP32 tensor used by the backend.

Level-2 sleep makes the MLA failure severe because the stale weights allocation is remapped without restoring its contents. Level 1 can mask the corruption by restoring the old values, but it still does not make a captured graph observe updated derived weights.

Changes

• Reuse replace_parameter(..., prefer_copy=True) for standard MLA W_UV and W_UK_T, preserving compatible storage across post-load processing.
• Retain the FlashInfer sink source and refresh its FP32 runtime tensor in place.
• Add focused regressions for both value refresh and stable storage addresses.

Validation

Focused tests on the rebased PR head:

CUDA_VISIBLE_DEVICES='' /home/aoshen/vllm/.venv/bin/python -m pytest \
  tests/v1/attention/test_mla_backends.py::test_mla_post_load_preserves_runtime_weight_addresses \
  tests/v1/attention/test_attention_backends.py::test_flashinfer_attention_sinks_refreshed_after_reload -q
# 4 passed

pre-commit run --files on all four changed files passed, including Ruff and mypy.

End-to-end MLA evaluation used Moonlight-16B-A3B-Instruct-FP8 on 8x H200 with VIME level-2 sleep, layerwise weight update, and FULL/PIECEWISE CUDA graphs:

Configuration	train/rollout logprob abs diff	Output
Before fix, level 2	1.765822	corrupted/gibberish
Diagnostic level-1 control	0.062213	coherent
After fix, level 2	0.058275	coherent
Matching slime baseline	0.065026	coherent

A minimal CUDA graph A/B also reproduced token divergence before the fix, no divergence in eager mode, and no divergence with CUDA graphs after the fix.

FlashInfer sink refresh is covered by the focused dtype-parametrized regression. A Blackwell end-to-end sink-model run was not available locally.

Duplicate-work check

Searched open and historical vLLM PRs/issues for FlashInfer sink reload, MLA post-load tensor rebinding, and process_weights_after_loading reload fixes. No matching fix was found. PR #45648 releases and recaptures CUDA graphs around full sleep/wake, but does not fix reloads outside that lifecycle and adds recapture overhead. PR #35956 addresses KV-cache allocation ownership, not weight tensor address identity.

Notes

AITER FP8/FP4 branches are intentionally unchanged because their platform-specific derived tensors require separate hardware coverage.

AI assistance: OpenAI Codex was used to investigate, implement, and test this change. The submitter reviewed the resulting diff and validation evidence.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[mergify]

Hi @aoshen02, the pre-commit checks have failed. Please run:

uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.