iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Bugfix][ROCm] Keep MLA derived weights in place on weight reload (#60389)

Browse Source 
Signed-off-by: Matthew Wong <Matthew.Wong2@amd.com>
Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
This commit is contained in:
Matt
2026-10-07 20:38:06 -05:00
committed by GitHub
co-authored by Claude Opus 5.5
parent 8733a21d55
commit 9dba339d07
2 changed files with 68 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/model_executor/model_loader/test_reload.py
+42 -1
View File
@@ -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
vllm/model_executor/layers/attention/mla_attention.py
+26 -12
View File
@@ -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