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 4, 21:15 UTC
[AMD][CI][BugFix] Override normalize_e4m3fn_to_e4m3fnuz for fnuz machines in test_moe_layer_no_parallel (#40550)
Signed-off-by: Randall Smith <Randall.Smith@amd.com>
rasmith committed
Apr 21, 2026 at 21:21 UTC
2463f00fb690a7b182050285c0179da03aad66fe
1 file changed
+27
tests/kernels/moe/test_moe_layer.py
+27
@@ -17,6 +17,7 @@ from typing import get_args
17
import pytest
18
import torch
19
20
+import vllm.model_executor.layers.quantization.utils.w8a8_utils
21
from tests.kernels.moe.modular_kernel_tools.parallel_utils import (
22
ProcessGroupInfo,
23
_set_vllm_config,
@@ -144,6 +145,24 @@ EPLB_SUPPORTED_QUANTS: list[str | None] = [None, "fp8"]
145
EPLB_SUPPORTED_BACKENDS: list[str] = ["allgather_reducescatter"]
146
147
148
+def mock_normalize_e4m3fn_to_e4m3fnuz(
149
+ weight: torch.Tensor,
150
+ weight_scale: torch.Tensor,
151
+ input_scale: torch.Tensor | None = None,
152
+):
153
+ return weight, weight_scale, input_scale
154
+
155
+
156
+# Needed since weights will already be in e4m3fnuz format on platforms that
157
+# use the fnuz fp8 format and the normalize_e4m3fn_to_e4m3fnuz() function
158
+# is not being tested here.
159
+# NOTE: The weights are quantized by moe_quantize_weights_2d in
160
+# _quantize_fp8_halves.
161
+# NOTE: Not able to use monkeypatch because of the spawned parallel workers.
162
+def override_normalize_e4m3fn_to_e4m3fnuz():
163
+ vllm.model_executor.layers.quantization.utils.w8a8_utils.normalize_e4m3fn_to_e4m3fnuz = mock_normalize_e4m3fn_to_e4m3fnuz # noqa: E501
164
+
165
+
166
def maybe_roundup_layer_hidden_size(
167
hidden_size: int,
168
act_dtype: torch.dtype,
@@ -1471,6 +1490,11 @@ def test_moe_layer_no_parallel(
1490
if os.environ.get("VLLM_LOGGING_LEVEL") is None:
1491
monkeypatch.setenv("VLLM_LOGGING_LEVEL", "ERROR")
1492
1493
+ # Needed since weights will already be in e4m3fnuz format and the
1494
+ # normalize_e4m3fn_to_e4m3fnuz() function is not being tested here.
1495
+ if current_platform.is_fp8_fnuz():
1496
+ override_normalize_e4m3fn_to_e4m3fnuz()
1497
+
1498
test_config = MoETestConfig(
1499
m,
1500
n,
@@ -1546,6 +1570,9 @@ def _parallel_worker(
1570
1571
dp_rank = vllm_config.parallel_config.data_parallel_rank
1572
1573
+ if current_platform.is_fp8_fnuz():
1574
+ override_normalize_e4m3fn_to_e4m3fnuz()
1575
+
1576
for test_config in test_configs:
1577
cc = vllm_config.compilation_config
1578
if "from_forward_context" in cc.static_forward_context: