[Bugfix][Model][CPU] Fix Mamba2 quantized in_proj weight and scale loading for TP >1 (#58083)
Signed-off-by: Akash kaothalkar <akash.kaothalkar@ibm.com> Co-authored-by: Akash kaothalkar <akash.kaothalkar@ibm.com> Co-authored-by: Antigravity <antigravity@google.com> Co-authored-by: Li, Jiang <jiang1.li@intel.com>
Akash kaothalkar committed
Sep 30, 2026 at 15:29 UTC
d71f6626061af38dad1c3f27830366cf7f279fe4
1 file changed
+12
-10
vllm/model_executor/layers/mamba/mamba_mixer2.py
+12
-10
@@ -232,13 +232,14 @@ def mamba_v2_sharded_weight_loader(
232
# - the ignore is for a mundane mypy error as it does not
233
# seem to handle slices well.
234
# https://github.com/python/mypy/issues/2410
235
+ target_slice = param.data[boundary : (boundary + take), ...]
236
param.data[
237
boundary : (boundary + take), ... # type: ignore[misc]
238
] = loaded_weight[
239
loaded_start_idx : (
240
loaded_start_idx + take
241
) # type: ignore[misc]
241
- ] # type: ignore[misc]
242
+ ].view_as(target_slice)
243
244
# move indexing boundaries
245
boundary += shard_size
@@ -440,17 +441,18 @@ class MambaMixer2(MambaBase, PluggableLayer):
441
tp_rank,
442
)
443
443
- # Apply the custom weight loader to in_proj.weight
444
+ # Apply the custom weight loader to in_proj.weight and its scales.
445
# Works for both non-quantized (Parameter) and quantized
446
# (ModelWeightParameter which extends BasevLLMParameter)
446
- if isinstance(self.in_proj.weight, BasevLLMParameter):
447
- # For BasevLLMParameter subclasses (quantized layers like FP8)
448
- # These have a weight_loader property that can be directly set
449
- self.in_proj.weight.weight_loader = mamba_loader
450
- else:
451
- # For standard Parameter (non-quantized layers)
452
- delattr(self.in_proj.weight, "weight_loader")
453
- set_weight_attrs(self.in_proj.weight, {"weight_loader": mamba_loader})
447
+ for attr_name in ("weight", "weight_scale", "input_scale"):
448
+ param = getattr(self.in_proj, attr_name, None)
449
+ if param is not None:
450
+ if isinstance(param, BasevLLMParameter):
451
+ param.weight_loader = mamba_loader
452
+ else:
453
+ if hasattr(param, "weight_loader"):
454
+ delattr(param, "weight_loader")
455
+ set_weight_attrs(param, {"weight_loader": mamba_loader})
456
457
# unsqueeze to fit conv1d weights shape into the linear weights shape.
458
# Can't do this in `weight_loader` since it already exists in