Add weights padding for fp8 per-block online quantization (#44763)
Signed-off-by: Yan Ma <yan.ma@intel.com>
Yan Ma committed
Jun 23, 2026 at 23:08 UTC
547d2c40d719e35d675405195053725a40e0351d
1 file changed
+55
vllm/model_executor/layers/quantization/online/fp8.py
+55
@@ -52,6 +52,7 @@ from vllm.model_executor.parameter import ModelWeightParameter
52
from vllm.model_executor.utils import replace_parameter
53
from vllm.platforms import current_platform
54
from vllm.utils.deep_gemm import per_block_cast_to_fp8
55
+from vllm.utils.math_utils import round_up
56
57
# ---------------------------------------------------------------------------
58
# Online FP8 Linear Methods
@@ -555,10 +556,64 @@ class Fp8PerBlockOnlineMoEMethod(_Fp8OnlineMoEBase):
556
layer=layer,
557
)
558
559
+ def maybe_roundup_sizes(
560
+ self,
561
+ hidden_size: int,
562
+ intermediate_size_per_partition: int,
563
+ act_dtype: torch.dtype,
564
+ moe_parallel_config,
565
+ ) -> tuple[int, int]:
566
+ hidden_size, intermediate_size_per_partition = super().maybe_roundup_sizes(
567
+ hidden_size=hidden_size,
568
+ intermediate_size_per_partition=intermediate_size_per_partition,
569
+ act_dtype=act_dtype,
570
+ moe_parallel_config=moe_parallel_config,
571
+ )
572
+ assert self.weight_block_size is not None
573
+ block_size = self.weight_block_size[0]
574
+ return (
575
+ round_up(hidden_size, block_size),
576
+ round_up(intermediate_size_per_partition, block_size),
577
+ )
578
+
579
+ def _zero_padding(self, layer: Module) -> None:
580
+ hidden_size = layer.moe_config.hidden_dim_unpadded
581
+ intermediate_size = layer.moe_config.intermediate_size_per_partition_unpadded
582
+
583
+ w13_half_size = layer.w13_weight.shape[1] // 2
584
+ if w13_half_size > intermediate_size:
585
+ layer.w13_weight[:, intermediate_size:w13_half_size, :] = 0
586
+ layer.w13_weight[
587
+ :, w13_half_size + intermediate_size : 2 * w13_half_size, :
588
+ ] = 0
589
+ if layer.w13_weight.shape[2] > hidden_size:
590
+ layer.w13_weight[:, :, hidden_size:] = 0
591
+
592
+ if layer.w2_weight.shape[1] > hidden_size:
593
+ layer.w2_weight[:, hidden_size:, :] = 0
594
+ if layer.w2_weight.shape[2] > intermediate_size:
595
+ layer.w2_weight[:, :, intermediate_size:] = 0
596
+
597
+ if getattr(layer, "w13_bias", None) is not None:
598
+ w13_bias_half_size = layer.w13_bias.shape[1] // 2
599
+ if w13_bias_half_size > intermediate_size:
600
+ layer.w13_bias[:, intermediate_size:w13_bias_half_size] = 0
601
+ layer.w13_bias[
602
+ :, w13_bias_half_size + intermediate_size : 2 * w13_bias_half_size
603
+ ] = 0
604
+
605
+ if (
606
+ getattr(layer, "w2_bias", None) is not None
607
+ and layer.w2_bias.shape[1] > hidden_size
608
+ ):
609
+ layer.w2_bias[:, hidden_size:] = 0
610
+
611
def process_weights_after_loading(self, layer: Module) -> None:
612
if getattr(layer, "_already_called_process_weights_after_loading", False):
613
return
614
615
+ self._zero_padding(layer)
616
+
617
fp8_dtype = current_platform.fp8_dtype()
618
w13 = torch.empty_like(layer.w13_weight, dtype=fp8_dtype)
619
w2 = torch.empty_like(layer.w2_weight, dtype=fp8_dtype)