[chore] clean-up weight prepack for INT8 MoE (#50116)
Signed-off-by: Fadi Arafeh <fadi.arafeh@arm.com>
Fadi Arafeh committed
Jul 31, 2026 at 11:15 UTC
0e9b50074e1cd3c498830b9124364037899a5c98
4 files changed
+7
-26
vllm/model_executor/layers/fused_moe/experts/cpu_moe.py
+4
-19
@@ -554,23 +554,6 @@ class CPUExpertsInt4(mk.FusedMoEExpertsMonolithic):
554
# ===========================================================================
555
556
557
-def prepare_int8_moe_layer_for_cpu(
558
- w13: torch.Tensor,
559
- w2: torch.Tensor,
560
-) -> tuple[torch.Tensor, torch.Tensor]:
561
- """Prepack INT8 MoE weights for the current CPU architecture."""
562
- # SMMLA packing for AArch64
563
- if current_platform.get_cpu_architecture() == CpuArchEnum.ARM:
564
- return (
565
- cpu_prepack_moe_weight_int8(w13, "neon"),
566
- cpu_prepack_moe_weight_int8(w2, "neon"),
567
- )
568
- # VNNI packing for x86
569
- packed_w13 = torch.ops._C.convert_weight_packed(w13)
570
- packed_w2 = torch.ops._C.convert_weight_packed(w2)
571
- return packed_w13, packed_w2
572
-
573
-
557
class CPUExpertsInt8(mk.FusedMoEExpertsMonolithic):
558
"""CPU INT8 W8A8 per-channel weight / dynamic per-token activation
559
monolithic MoE experts."""
@@ -650,7 +633,8 @@ class CPUExpertsInt8(mk.FusedMoEExpertsMonolithic):
633
"""VNNI-prepack INT8 MoE weights for CPU kernel."""
634
from vllm.model_executor.utils import replace_parameter
635
653
- w13, w2 = prepare_int8_moe_layer_for_cpu(layer.w13_weight, layer.w2_weight)
636
+ w13 = torch.ops._C.convert_weight_packed(layer.w13_weight)
637
+ w2 = torch.ops._C.convert_weight_packed(layer.w2_weight)
638
replace_parameter(layer, "w13_weight", w13)
639
replace_parameter(layer, "w2_weight", w2)
640
@@ -814,7 +798,8 @@ class ArmCPUExpertsInt8(mk.FusedMoEExpertsMonolithic):
798
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
799
from vllm.model_executor.utils import replace_parameter
800
817
- w13, w2 = prepare_int8_moe_layer_for_cpu(layer.w13_weight, layer.w2_weight)
801
+ w13 = cpu_prepack_moe_weight_int8(layer.w13_weight, "neon")
802
+ w2 = cpu_prepack_moe_weight_int8(layer.w2_weight, "neon")
803
replace_parameter(layer, "w13_weight", w13)
804
replace_parameter(layer, "w2_weight", w2)
805
vllm/model_executor/layers/fused_moe/oracle/int8.py
+1
-7
@@ -266,13 +266,7 @@ def convert_to_int8_moe_kernel_format(
266
quant_config=_humming_int8_weight_schema(w13, layer.w13_weight_scale),
267
)
268
return layer.w13_weight, layer.w2_weight
269
- elif int8_backend == Int8MoeBackend.CPU:
270
- from vllm.model_executor.layers.fused_moe.experts.cpu_moe import (
271
- prepare_int8_moe_layer_for_cpu,
272
- )
273
-
274
- w13, w2 = prepare_int8_moe_layer_for_cpu(w13, w2)
275
- elif int8_backend != Int8MoeBackend.TRITON:
269
+ elif int8_backend not in (Int8MoeBackend.TRITON, Int8MoeBackend.CPU):
270
raise ValueError(f"Unsupported Int8 MoE backend: {int8_backend.value}")
271
272
return w13, w2
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w8a8_int8.py
+1
@@ -162,6 +162,7 @@ class CompressedTensorsW8A8Int8MoEMethod(CompressedTensorsMoEMethod):
162
routing_tables=layer._expert_routing_tables(),
163
layer=layer,
164
)
165
+ self.moe_kernel.fused_experts.process_weights_after_loading(layer)
166
167
def maybe_make_prepare_finalize(
168
self,
vllm/model_executor/layers/quantization/online/int8.py
+1
@@ -114,6 +114,7 @@ class Int8OnlineMoEMethod(OnlineMoEMethodBase):
114
routing_tables=layer._expert_routing_tables(),
115
layer=layer,
116
)
117
+ self.moe_kernel.fused_experts.process_weights_after_loading(layer)
118
119
def get_fused_moe_quant_config(
120
self, layer: torch.nn.Module