From 77e52645e9baba15d6b9cd7e09a12e70628b8237 Mon Sep 17 00:00:00 2001 From: "yhcheong (vllmellm)" Date: Tue, 29 Sep 2026 17:37:20 +0800 Subject: [PATCH] [Bugfix][Model] MiMo: keep fused fp8 qkv_proj pairing state across weight-loading calls (#58142) Signed-off-by: vllmellm --- vllm/model_executor/models/mimo_v2.py | 8 ++++++-- vllm/model_executor/models/mimo_v2_mtp.py | 7 ++++++- 2 files changed, 12 insertions(+), 3 deletions(-) diff --git a/vllm/model_executor/models/mimo_v2.py b/vllm/model_executor/models/mimo_v2.py index dfbed25655..6267a37dc0 100644 --- a/vllm/model_executor/models/mimo_v2.py +++ b/vllm/model_executor/models/mimo_v2.py @@ -742,8 +742,12 @@ class MiMoV2Model(nn.Module, EagleModelMixin): expert_params_mapping = self.get_expert_mapping() # Pro-format fused qkv_proj arrives as two tensors (weight and # weight_scale_inv). Store them per-layer so that they can be - # sharded together. - pending_fp8_qkv_proj: dict[str, dict[str, torch.Tensor]] = {} + # sharded together. The state must outlive this call: AutoWeightsLoader + # delegates per contiguous group of names, so a pair can straddle two + # calls and would otherwise be dropped silently. + pending_fp8_qkv_proj = getattr(self, "_pending_fp8_qkv_proj", None) + if pending_fp8_qkv_proj is None: + self._pending_fp8_qkv_proj = pending_fp8_qkv_proj = {} for name, loaded_weight in weights: if "rotary_emb.inv_freq" in name: continue diff --git a/vllm/model_executor/models/mimo_v2_mtp.py b/vllm/model_executor/models/mimo_v2_mtp.py index a4e46977ec..c413b65ba9 100644 --- a/vllm/model_executor/models/mimo_v2_mtp.py +++ b/vllm/model_executor/models/mimo_v2_mtp.py @@ -266,7 +266,12 @@ class MiMoV2MTP(nn.Module): params_dict = dict(self.named_parameters()) loaded_params: set[str] = set() - pending_qkv_proj: dict[str, dict[str, torch.Tensor]] = {} + # The pairing state must outlive this call: AutoWeightsLoader delegates + # per contiguous group of names, so a pair can straddle two calls and + # would otherwise be dropped silently. + pending_qkv_proj = getattr(self, "_pending_qkv_proj", None) + if pending_qkv_proj is None: + self._pending_qkv_proj = pending_qkv_proj = {} for name, loaded_weight in weights: if "rotary_emb.inv_freq" in name: