diff --git a/tests/model_executor/model_loader/test_reload.py b/tests/model_executor/model_loader/test_reload.py index 85ddea3946..aaad16778f 100644 --- a/tests/model_executor/model_loader/test_reload.py +++ b/tests/model_executor/model_loader/test_reload.py @@ -1051,6 +1051,37 @@ def test_hpc_rope_norm_kernel_sees_refit_norm_weights(monkeypatch, hpc_rope_norm assert not hasattr(rnorm, "qnorm_weight") +def _device_tensor_ptrs(worker) -> dict[str, int]: + """Addresses of every device tensor held by a module; graphs capture them.""" + ptrs = {} + for name, module in worker.get_model().named_modules(): + tensors = {**vars(module), **module._parameters, **module._buffers} + for attr, tensor in tensors.items(): + if isinstance(tensor, torch.Tensor) and tensor.device.type != "cpu": + ptrs[f"{name}.{attr}"] = tensor.data_ptr() + return ptrs + + +def _moved_tensors(llm, ptrs: list[dict[str, int]]) -> list[list[str]]: + after = llm.collective_rpc(_device_tensor_ptrs) + return [ + sorted(n for n, p in old.items() if new.get(n) != p) + for old, new in zip(ptrs, after) + ] + + +@pytest.mark.parametrize( + "use_aiter", + [ + False, + pytest.param( + True, + marks=pytest.mark.skipif( + not current_platform.is_rocm(), reason="AITER is ROCm-only" + ), + ), + ], +) @pytest.mark.parametrize( "tp_size", [pytest.param(1), pytest.param(2, marks=[pytest.mark.slow_test])] ) @@ -1094,13 +1125,19 @@ def test_hpc_rope_norm_kernel_sees_refit_norm_weights(monkeypatch, hpc_rope_norm ), ], ) -def test_reload_weights(base_model, mul_model, add_model, tp_size, vllm_runner): +def test_reload_weights( + base_model, mul_model, add_model, tp_size, use_aiter, vllm_runner, monkeypatch +): if current_platform.device_count() < tp_size: pytest.skip(reason="Not enough CUDA devices") if "FP8" in base_model and _fp8_reload_unsupported(): pytest.skip(reason="Requires FP8 support") + monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1") + if use_aiter: + monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1") + with vllm_runner( model_name=base_model, tensor_parallel_size=tp_size, @@ -1109,12 +1146,16 @@ def test_reload_weights(base_model, mul_model, add_model, tp_size, vllm_runner): max_model_len=16, max_num_seqs=1, ) as llm: + # Reloading must update tensors in place: captured graphs keep pointers. + ptrs = llm.collective_rpc(_device_tensor_ptrs) llm.collective_rpc("reload_weights", kwargs={"weights_path": mul_model}) + assert _moved_tensors(llm, ptrs) == [[]] * len(ptrs) mul_perp = llm.generate_prompt_perplexity(["3 4 = 12"], mask=["3 4 ="])[0] add_perp = llm.generate_prompt_perplexity(["3 4 = 7"], mask=["3 4 ="])[0] assert mul_perp < add_perp llm.collective_rpc("reload_weights", kwargs={"weights_path": add_model}) + assert _moved_tensors(llm, ptrs) == [[]] * len(ptrs) mul_perp = llm.generate_prompt_perplexity(["3 4 = 12"], mask=["3 4 ="])[0] add_perp = llm.generate_prompt_perplexity(["3 4 = 7"], mask=["3 4 ="])[0] assert add_perp < mul_perp diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index f550dc3d90..4286502ea9 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -1272,23 +1272,34 @@ class MLAAttention(nn.Module, AttentionLayerBase): quark_quantize_weight_to_mxfp4, ) - self.W_K, self.W_K_scale = quark_quantize_weight_to_mxfp4(W_UK) - # Convert from (L, N, P) to (N, L, P) - self.W_K = self.W_K.transpose(0, 1) - self.W_K_scale = self.W_K_scale.transpose(0, 1) - - self.W_V, self.W_V_scale = quark_quantize_weight_to_mxfp4( - W_UV.permute(1, 2, 0) - ) + W_K, W_K_scale = quark_quantize_weight_to_mxfp4(W_UK) + W_V, W_V_scale = quark_quantize_weight_to_mxfp4(W_UV.permute(1, 2, 0)) + # Copy in place on reload so captured CUDA graphs stay valid. + # Convert W_K from (L, N, P) to (N, L, P). + for name, tensor in ( + ("W_K", W_K.transpose(0, 1)), + ("W_K_scale", W_K_scale.transpose(0, 1)), + ("W_V", W_V), + ("W_V_scale", W_V_scale), + ): + replace_parameter(self, name, tensor, prefer_copy=True) elif self.is_aiter_triton_fp8_bmm_enabled: W_K = W_UK.transpose(0, 1) # 16 512 128 W_V = W_UV.permute(1, 2, 0) # 16 128 512 - self.W_K, self.W_K_scale = dynamic_per_batched_tensor_quant( + W_K, W_K_scale = dynamic_per_batched_tensor_quant( W_K, dtype=current_platform.fp8_dtype() ) - self.W_V, self.W_V_scale = dynamic_per_batched_tensor_quant( + W_V, W_V_scale = dynamic_per_batched_tensor_quant( W_V, dtype=current_platform.fp8_dtype() ) + # Copy in place on reload so captured CUDA graphs stay valid. + for name, tensor in ( + ("W_K", W_K), + ("W_K_scale", W_K_scale), + ("W_V", W_V), + ("W_V_scale", W_V_scale), + ): + replace_parameter(self, name, tensor, prefer_copy=True) # The kernel operates on non-padded inputs. Hence, pre-compiling # triton kernel to avoid runtime compilation for unseen batch sizes @@ -1327,8 +1338,11 @@ class MLAAttention(nn.Module, AttentionLayerBase): # Convert from (L, N, P) to (N, P, L) replace_parameter(self, "W_UK_T", W_UK.permute(1, 2, 0), prefer_copy=True) if self.dcp_q_replicate: - self.W_UK_T_dcp_qrep = get_dcp_group().all_gather( - self.W_UK_T.contiguous(), dim=0 + replace_parameter( + self, + "W_UK_T_dcp_qrep", + get_dcp_group().all_gather(self.W_UK_T.contiguous(), dim=0), + prefer_copy=True, ) # If we should not load quant weights, we initialize the scales to 1.0