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 7, 21:30 UTC
[Quantization] Support FP8 MoE bias for models like GPT-OSS (#34906)
Signed-off-by: jasperjiaguo <jasperg662@gmail.com> Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Jia Guo committed
Feb 23, 2026 at 19:07 UTC
ec85340531e83962877fa683e23f1d5d17cd786a
1 file changed
+54
-1
vllm/model_executor/layers/quantization/fp8.py
+54
-1
@@ -758,6 +758,25 @@ class Fp8MoEMethod(FusedMoEMethodBase):
758
layer.register_parameter("w2_weight", w2_weight)
759
set_weight_attrs(w2_weight, extra_weight_attrs)
760
761
+ # BIASES (for models like GPT-OSS that have biased MoE)
762
+ if self.moe.has_bias:
763
+ w13_bias = torch.nn.Parameter(
764
+ torch.zeros(
765
+ num_experts,
766
+ 2 * intermediate_size_per_partition,
767
+ dtype=layer.orig_dtype,
768
+ ),
769
+ requires_grad=False,
770
+ )
771
+ layer.register_parameter("w13_bias", w13_bias)
772
+ set_weight_attrs(w13_bias, extra_weight_attrs)
773
+ w2_bias = torch.nn.Parameter(
774
+ torch.zeros(num_experts, hidden_size, dtype=layer.orig_dtype),
775
+ requires_grad=False,
776
+ )
777
+ layer.register_parameter("w2_bias", w2_bias)
778
+ set_weight_attrs(w2_bias, extra_weight_attrs)
779
+
780
# WEIGHT_SCALES
781
if not self.block_quant:
782
# For per-tensor quant, the scales are per expert and weight.
@@ -939,7 +958,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):
958
a1_scale = layer.w13_input_scale
959
a2_scale = layer.w2_input_scale
960
942
- return make_fp8_moe_quant_config(
961
+ quant_config = make_fp8_moe_quant_config(
962
fp8_backend=self.fp8_backend,
963
w1_scale=w1_scale,
964
w2_scale=w2_scale,
@@ -948,6 +967,18 @@ class Fp8MoEMethod(FusedMoEMethodBase):
967
block_shape=self.weight_block_size,
968
)
969
970
+ # Inject biases into the quant config if the model has them
971
+ # (e.g. GPT-OSS biased MoE)
972
+ if quant_config is not None and self.moe.has_bias:
973
+ w13_bias = getattr(layer, "w13_bias", None)
974
+ w2_bias = getattr(layer, "w2_bias", None)
975
+ if w13_bias is not None:
976
+ quant_config._w1.bias = w13_bias
977
+ if w2_bias is not None:
978
+ quant_config._w2.bias = w2_bias
979
+
980
+ return quant_config
981
+
982
@property
983
def supports_eplb(self) -> bool:
984
return True
@@ -1168,6 +1199,28 @@ class Fp8OnlineMoEMethod(Fp8MoEMethod):
1199
# stash the correct device for `patched_weight_loader`
1200
layer._load_device = torch.get_default_device()
1201
1202
+ # BIASES (for models like GPT-OSS that have biased MoE)
1203
+ if self.moe.has_bias:
1204
+ # Use the original weight_loader (not patched) for biases
1205
+ orig_extra_weight_attrs = dict(extra_weight_attrs)
1206
+ orig_extra_weight_attrs["weight_loader"] = weight_loader
1207
+ w13_bias = torch.nn.Parameter(
1208
+ torch.zeros(
1209
+ num_experts,
1210
+ 2 * intermediate_size_per_partition,
1211
+ dtype=layer.orig_dtype,
1212
+ ),
1213
+ requires_grad=False,
1214
+ )
1215
+ layer.register_parameter("w13_bias", w13_bias)
1216
+ set_weight_attrs(w13_bias, orig_extra_weight_attrs)
1217
+ w2_bias = torch.nn.Parameter(
1218
+ torch.zeros(num_experts, hidden_size, dtype=layer.orig_dtype),
1219
+ requires_grad=False,
1220
+ )
1221
+ layer.register_parameter("w2_bias", w2_bias)
1222
+ set_weight_attrs(w2_bias, orig_extra_weight_attrs)
1223
+
1224
# WEIGHT_SCALES
1225
# Allocate 2 scales for w1 and w3 respectively.
1226
# They will be combined to a single scale after weight loading.