fix(moe_wna16): access tp_size via moe_config for RoutedExperts compatibility (#45404)
Signed-off-by: Oxygen <1391083091@qq.com> Signed-off-by: Willow Lopez <100782273+Oxygen56@users.noreply.github.com>
Willow Lopez committed
Jun 23, 2026 at 23:46 UTC
f3410b3bb16b1b0f33468a65f260148565c9948c
1 file changed
+4
-3
vllm/model_executor/layers/quantization/moe_wna16.py
+4
-3
@@ -426,7 +426,6 @@ class MoeWNA16Method(FusedMoEMethodBase):
426
427
device = get_tp_group().device
428
tp_rank = get_tensor_model_parallel_rank()
429
- tp_size = layer.moe_config.moe_parallel_config.tp_size
429
loaded_weight = loaded_weight.to(device)
430
shard_size = layer.intermediate_size_per_partition
431
@@ -465,7 +464,9 @@ class MoeWNA16Method(FusedMoEMethodBase):
464
)
465
466
if "w13_qzeros" in weight_name:
468
- tensor = loaded_weight.view(tp_size, -1, loaded_weight.size(1))[tp_rank]
467
+ tensor = loaded_weight.view(
468
+ layer.moe_config.tp_size, -1, loaded_weight.size(1)
469
+ )[tp_rank]
470
if shard_id == "w1":
471
param.data[expert_id, : shard_size // 2] = tensor
472
else:
@@ -473,7 +474,7 @@ class MoeWNA16Method(FusedMoEMethodBase):
474
return True if return_success else None
475
elif "w2_qzeros" in weight_name:
476
param.data[expert_id] = loaded_weight.view(
476
- loaded_weight.size(0), tp_size, -1
477
+ loaded_weight.size(0), layer.moe_config.tp_size, -1
478
)[:, tp_rank]
479
return True if return_success else None
480
else: