[Bugfix][Model] MiMo: keep fused fp8 qkv_proj pairing state across weight-loading calls (#58142)
Signed-off-by: vllmellm <vllm.ellm@embeddedllm.com>
yhcheong (vllmellm) committed
Sep 29, 2026 at 17:37 UTC
77e52645e9baba15d6b9cd7e09a12e70628b8237
2 files changed
+12
-3
vllm/model_executor/models/mimo_v2.py
+6
-2
@@ -742,8 +742,12 @@ class MiMoV2Model(nn.Module, EagleModelMixin):
742
expert_params_mapping = self.get_expert_mapping()
743
# Pro-format fused qkv_proj arrives as two tensors (weight and
744
# weight_scale_inv). Store them per-layer so that they can be
745
- # sharded together.
746
- pending_fp8_qkv_proj: dict[str, dict[str, torch.Tensor]] = {}
745
+ # sharded together. The state must outlive this call: AutoWeightsLoader
746
+ # delegates per contiguous group of names, so a pair can straddle two
747
+ # calls and would otherwise be dropped silently.
748
+ pending_fp8_qkv_proj = getattr(self, "_pending_fp8_qkv_proj", None)
749
+ if pending_fp8_qkv_proj is None:
750
+ self._pending_fp8_qkv_proj = pending_fp8_qkv_proj = {}
751
for name, loaded_weight in weights:
752
if "rotary_emb.inv_freq" in name:
753
continue
vllm/model_executor/models/mimo_v2_mtp.py
+6
-1
@@ -266,7 +266,12 @@ class MiMoV2MTP(nn.Module):
266
267
params_dict = dict(self.named_parameters())
268
loaded_params: set[str] = set()
269
- pending_qkv_proj: dict[str, dict[str, torch.Tensor]] = {}
269
+ # The pairing state must outlive this call: AutoWeightsLoader delegates
270
+ # per contiguous group of names, so a pair can straddle two calls and
271
+ # would otherwise be dropped silently.
272
+ pending_qkv_proj = getattr(self, "_pending_qkv_proj", None)
273
+ if pending_qkv_proj is None:
274
+ self._pending_qkv_proj = pending_qkv_proj = {}
275
276
for name, loaded_weight in weights:
277
if "rotary_emb.inv_freq" in name: