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 10, 20:15 UTC
[Bugfix][ROCm] Fix Static Quant Issue (#31502)
Signed-off-by: Robert Shaw <robshaw@redhat.com> Co-authored-by: Robert Shaw <robshaw@redhat.com> Co-authored-by: TJian <tunjian.tan@embeddedllm.com>
Robert Shaw committed
Dec 29, 2025 at 16:27 UTC
56f516254c5ff024e3986a4d625f698f3ec2a288
2 files changed
+14
-14
vllm/model_executor/layers/fused_moe/rocm_aiter_fused_moe.py
+4
-1
@@ -325,8 +325,11 @@ class AiterExperts(mk.FusedMoEPermuteExpertsUnpermute):
325
expert_tokens_meta: mk.ExpertTokensMetadata | None,
326
apply_router_weight_on_input: bool,
327
):
328
+ # TODO(rob): rocm_aiter_fused_experts uses self.quant_config's
329
+ # a_scales for static quantization. Update this to fit better
330
+ # with the interface once all quant integrations are complete.
331
assert a1q_scale is None
329
- assert a2_scale is None
332
+ assert a2_scale == self.quant_config.a2_scale
333
assert expert_tokens_meta is None
334
335
result = rocm_aiter_fused_experts(
vllm/model_executor/layers/quantization/fp8.py
+10
-13
@@ -1046,35 +1046,32 @@ class Fp8MoEMethod(FusedMoEMethodBase):
1046
w2_weight = layer.w2_weight
1047
w13_weight_scale = getattr(layer, f"w13_{self.weight_scale_name}")
1048
w2_weight_scale = getattr(layer, f"w2_{self.weight_scale_name}")
1049
+ w13_input_scale = layer.w13_input_scale
1050
+ w2_input_scale = layer.w2_input_scale
1051
1052
# MI300x and MI325x use FNUZ format for FP8. Convert if needed.
1053
if current_platform.is_fp8_fnuz():
1052
- w13_weight, w13_weight_scale, layer.w13_input_scale = (
1054
+ w13_weight, w13_weight_scale, w13_input_scale = (
1055
normalize_e4m3fn_to_e4m3fnuz(
1054
- w13_weight, w13_weight_scale, layer.w13_input_scale
1056
+ w13_weight, w13_weight_scale, w13_input_scale
1057
)
1058
)
1057
- w2_weight, w2_weight_scale, layer.w2_input_scale = (
1058
- normalize_e4m3fn_to_e4m3fnuz(
1059
- w2_weight, w2_weight_scale, layer.w2_input_scale
1060
- )
1059
+ w2_weight, w2_weight_scale, w2_input_scale = normalize_e4m3fn_to_e4m3fnuz(
1060
+ w2_weight, w2_weight_scale, w2_input_scale
1061
)
1062
1063
# Per tensor kernels require single activation scale. Use the max.
1064
if self.quant_config.activation_scheme == "static":
1065
assert not self.block_quant
1066
- assert layer.w13_input_scale is not None
1067
- assert layer.w2_input_scale is not None
1068
- if not all_close_1d(layer.w13_input_scale) or not all_close_1d(
1069
- layer.w2_input_scale
1070
- ):
1066
+ assert w13_input_scale is not None and w2_input_scale is not None
1067
+ if not all_close_1d(w13_input_scale) or not all_close_1d(w2_input_scale):
1068
logger.warning_once(
1069
"Found input_scales that are not equal for "
1070
"fp8 MoE layer. Using the maximum across experts "
1071
"for each layer."
1072
)
1076
- replace_parameter(layer, "w13_input_scale", layer.w13_input_scale.max())
1077
- replace_parameter(layer, "w2_input_scale", layer.w2_input_scale.max())
1073
+ replace_parameter(layer, "w13_input_scale", w13_input_scale.max())
1074
+ replace_parameter(layer, "w2_input_scale", w2_input_scale.max())
1075
1076
# Per tensor kernels require single weight scale for w13 per expert, but
1077
# on disk there is a scale for w1 and w3. Use the max to requantize.