A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 05:30 UTC
[FIXBUG] Add return_success parameter to moe_wna16_weight_loader function (#22797)
Signed-off-by: JartX <sagformas@epdcenter.es> Co-authored-by: Michael Goin <mgoin64@gmail.com>
JartX committed
Aug 28, 2025 at 11:03 UTC
3462c1c522d214755f1dfce3d645ab5afe7f00ae
1 file changed
+19
-13
vllm/model_executor/layers/quantization/moe_wna16.py
+19
-13
@@ -124,7 +124,7 @@ class MoeWNA16Config(QuantizationConfig):
124
awq_min_capability = AWQConfig.get_min_capability()
125
126
gptq_compatible = quant_method == "gptq" and \
127
- not desc_act and num_bits in [4, 8]
127
+ not desc_act and num_bits in [4, 8]
128
awq_compatible = quant_method == "awq" and num_bits == 4 and \
129
device_capability >= awq_min_capability
130
@@ -175,11 +175,8 @@ class MoeWNA16Method(FusedMoEMethodBase):
175
quant_config: The MOE WNA16 (W8A16/W4A16) quantization config.
176
"""
177
178
- def __init__(
179
- self,
180
- quant_config: MoeWNA16Config,
181
- moe: FusedMoEConfig,
182
- ):
178
+ def __init__(self, quant_config: MoeWNA16Config,
179
+ moe: "FusedMoEConfig") -> None:
180
super().__init__(moe)
181
self.quant_config = quant_config
182
@@ -187,6 +184,7 @@ class MoeWNA16Method(FusedMoEMethodBase):
184
hidden_size: int, intermediate_size_per_partition: int,
185
params_dtype: torch.dtype, **extra_weight_attrs):
186
187
+ self.moe = layer
188
layer.quant_config = self.quant_config
189
bit8_pack_factor = self.quant_config.bit8_pack_factor
190
group_size = self.quant_config.group_size
@@ -308,7 +306,6 @@ class MoeWNA16Method(FusedMoEMethodBase):
306
logical_replica_count: Optional[torch.Tensor] = None,
307
) -> torch.Tensor:
308
assert self.fused_experts is None
311
-
309
if enable_eplb:
310
raise NotImplementedError(
311
"EPLB not supported for `MoeWNA16Method` yet.")
@@ -404,12 +401,14 @@ class MoeWNA16Method(FusedMoEMethodBase):
401
402
def moe_wna16_weight_loader(param: torch.nn.Parameter,
403
loaded_weight: torch.Tensor,
407
- weight_name: str, shard_id: str,
408
- expert_id: int):
404
+ weight_name: str,
405
+ shard_id: str,
406
+ expert_id: int,
407
+ return_success: bool = False):
408
if "g_idx" in weight_name:
410
- return
409
+ return False if return_success else None
410
if not layer.quant_config.has_zp and "qzeros" in weight_name:
412
- return
411
+ return False if return_success else None
412
413
device = get_tp_group().device
414
tp_rank = get_tensor_model_parallel_rank()
@@ -455,11 +454,18 @@ class MoeWNA16Method(FusedMoEMethodBase):
454
param.data[expert_id, :shard_size // 2] = tensor
455
else:
456
param.data[expert_id, shard_size // 2:] = tensor
457
+ return True if return_success else None
458
elif "w2_qzeros" in weight_name:
459
param.data[expert_id] = loaded_weight.view(
460
loaded_weight.size(0), layer.tp_size, -1)[:, tp_rank]
461
+ return True if return_success else None
462
else:
462
- weight_loader(param, loaded_weight, weight_name, shard_id,
463
- expert_id)
463
+ # Delegate to the original loader, passing return_success
464
+ return weight_loader(param,
465
+ loaded_weight,
466
+ weight_name,
467
+ shard_id,
468
+ expert_id,
469
+ return_success=return_success)
470
471
return moe_wna16_weight_loader