[Bugfix][Attention] Preserve post-load tensors across weight reloads (#48251)
Signed-off-by: aoshen02 <aoshen@inferact.ai>
aoshen02 committed
Jul 17, 2026 at 14:15 UTC
17fdd421009092a569e9ed22346b932e55824fb5
4 files changed
+80
-4
tests/v1/attention/test_attention_backends.py
+27
@@ -646,6 +646,33 @@ def test_flashinfer_xqa_bmm1_scale_matches_decode_q_dtype():
646
assert impl.get_xqa_bmm1_scale(MockLayer, torch.float8_e4m3fn) == 3.0
647
648
649
+@pytest.mark.skipif(
650
+ AttentionBackendEnum.FLASHINFER not in BACKENDS_TO_TEST,
651
+ reason="FlashInfer is not available.",
652
+)
653
+@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
654
+def test_flashinfer_attention_sinks_refreshed_after_reload(dtype):
655
+ from vllm.v1.attention.backends import flashinfer as flashinfer_backend
656
+
657
+ source_sinks = torch.tensor([1.0, 2.0], dtype=dtype)
658
+ impl = object.__new__(flashinfer_backend.FlashInferImpl)
659
+ impl._sinks_source = source_sinks
660
+ impl.sinks = source_sinks
661
+
662
+ impl.process_weights_after_loading(dtype)
663
+
664
+ assert impl.sinks is not None
665
+ sinks_ptr = impl.sinks.data_ptr()
666
+ assert impl.sinks.dtype == torch.float32
667
+ torch.testing.assert_close(impl.sinks, source_sinks.float())
668
+
669
+ source_sinks.copy_(torch.tensor([3.0, 4.0], dtype=dtype))
670
+ impl.process_weights_after_loading(dtype)
671
+
672
+ assert impl.sinks.data_ptr() == sinks_ptr
673
+ torch.testing.assert_close(impl.sinks, source_sinks.float())
674
+
675
+
676
@pytest.mark.skipif(
677
AttentionBackendEnum.FLASHINFER not in BACKENDS_TO_TEST,
678
reason="FlashInfer is not available.",
tests/v1/attention/test_mla_backends.py
+39
@@ -105,6 +105,45 @@ if AttentionBackendEnum.TOKENSPEED_MLA in BACKENDS_TO_TEST:
105
BACKENDS_TO_TEST.remove(AttentionBackendEnum.TOKENSPEED_MLA)
106
107
108
+def test_mla_post_load_preserves_runtime_weight_addresses(monkeypatch):
109
+ layer = MLAAttention.__new__(MLAAttention)
110
+ torch.nn.Module.__init__(layer)
111
+ layer.kv_lora_rank = 2
112
+ layer.num_heads = 2
113
+ layer.qk_nope_head_dim = 3
114
+ layer.v_head_dim = 4
115
+ layer.kv_b_proj = torch.nn.Module()
116
+ layer.kv_b_proj.weight = torch.nn.Parameter(
117
+ torch.arange(28.0, dtype=torch.float16).reshape(14, 2)
118
+ )
119
+ layer.kv_b_proj.quant_method = None
120
+ layer.is_aiter_triton_fp4_bmm_enabled = False
121
+ layer.is_aiter_triton_fp8_bmm_enabled = False
122
+ layer.quant_config = None
123
+ layer.layer_name = "test"
124
+
125
+ monkeypatch.setattr(
126
+ mla_attention_module, "set_default_quant_scales", lambda *_, **__: None
127
+ )
128
+
129
+ with torch.no_grad():
130
+ layer.process_weights_after_loading(torch.float32)
131
+ assert isinstance(layer.W_UV, torch.nn.Parameter)
132
+ assert isinstance(layer.W_UK_T, torch.nn.Parameter)
133
+ w_uv_ptr = layer.W_UV.data_ptr()
134
+ w_uk_t_ptr = layer.W_UK_T.data_ptr()
135
+ old_w_uv = layer.W_UV.clone()
136
+ old_w_uk_t = layer.W_UK_T.clone()
137
+
138
+ layer.kv_b_proj.weight.add_(100)
139
+ layer.process_weights_after_loading(torch.float32)
140
+
141
+ assert layer.W_UV.data_ptr() == w_uv_ptr
142
+ assert layer.W_UK_T.data_ptr() == w_uk_t_ptr
143
+ torch.testing.assert_close(layer.W_UV, old_w_uv + 100)
144
+ torch.testing.assert_close(layer.W_UK_T, old_w_uk_t + 100)
145
+
146
+
147
# Filtered per-test via validate_configuration (capability/deps/dims).
148
PREFILL_BACKENDS_TO_TEST = [
149
MLAPrefillBackendEnum.FLASH_ATTN,
vllm/model_executor/layers/attention/mla_attention.py
+3
-2
@@ -240,6 +240,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
240
kFp8StaticTensorSym,
241
kNvfp4Dynamic,
242
)
243
+from vllm.model_executor.utils import replace_parameter
244
from vllm.platforms import current_platform
245
from vllm.utils.flashinfer import has_flashinfer
246
from vllm.utils.math_utils import cdiv, round_down
@@ -989,9 +990,9 @@ class MLAAttention(nn.Module, AttentionLayerBase):
990
)
991
else:
992
# Convert from (L, N, V) to (N, L, V)
992
- self.W_UV = W_UV.transpose(0, 1)
993
+ replace_parameter(self, "W_UV", W_UV.transpose(0, 1), prefer_copy=True)
994
# Convert from (L, N, P) to (N, P, L)
994
- self.W_UK_T = W_UK.permute(1, 2, 0)
995
+ replace_parameter(self, "W_UK_T", W_UK.permute(1, 2, 0), prefer_copy=True)
996
997
# If we should not load quant weights, we initialize the scales to 1.0
998
# as the default value. See [Note: Register q/k/v/prob scales in state dict]
vllm/v1/attention/backends/flashinfer.py
+11
-2
@@ -1573,6 +1573,8 @@ class FlashInferImpl(AttentionImpl):
1573
)
1574
1575
self.sinks: torch.Tensor | None = None
1576
+ # Keep the source so RL weight updates can refresh the runtime tensor.
1577
+ self._sinks_source = sinks
1578
if sinks is not None:
1579
if sinks.shape[0] != num_heads:
1580
raise ValueError(
@@ -1633,8 +1635,15 @@ class FlashInferImpl(AttentionImpl):
1635
1636
# FlashInfer requires attention sinks to be float32
1637
def process_weights_after_loading(self, act_dtype: torch.dtype):
1636
- if self.sinks is not None and self.sinks.dtype != torch.float32:
1637
- self.sinks = self.sinks.to(torch.float32)
1638
+ source_sinks = self._sinks_source
1639
+ if source_sinks is None:
1640
+ return
1641
+ if source_sinks.dtype == torch.float32:
1642
+ self.sinks = source_sinks
1643
+ elif self.sinks is None or self.sinks.dtype != torch.float32:
1644
+ self.sinks = source_sinks.to(torch.float32)
1645
+ else:
1646
+ self.sinks.copy_(source_sinks)
1647
1648
def get_xqa_bmm1_scale(self, layer: torch.nn.Module, q_data_type: torch.dtype):
1649
bmm1_scale = self.scale