[ROCm][CI] Fix moe layer tests for fp8 dtype compatibility (#56106)
Signed-off-by: Divakar Verma <divakar.verma@amd.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com>
Divakar Verma committed
Sep 9, 2026 at 13:48 UTC
3fb676bfad0f1c7099af6296983e739be3fe29cc
2 files changed
+11
-8
tests/kernels/moe/test_moe_layer.py
+4
-1
@@ -84,7 +84,7 @@ def on_gfx950() -> bool:
84
return False
85
86
87
-fp8_dtype = torch.float8_e4m3fn # current_platform.fp8_dtype
87
+fp8_dtype = current_platform.fp8_dtype()
88
89
SHAPE_COMBOS = [
90
(1, 128, 256),
@@ -191,6 +191,9 @@ def mock_normalize_e4m3fn_to_e4m3fnuz(
191
# NOTE: Not able to use monkeypatch because of the spawned parallel workers.
192
def override_normalize_e4m3fn_to_e4m3fnuz():
193
vllm.model_executor.layers.quantization.utils.w8a8_utils.normalize_e4m3fn_to_e4m3fnuz = mock_normalize_e4m3fn_to_e4m3fnuz # noqa: E501
194
+ vllm.model_executor.layers.quantization.fp8.normalize_e4m3fn_to_e4m3fnuz = (
195
+ mock_normalize_e4m3fn_to_e4m3fnuz # noqa: E501
196
+ )
197
198
199
def sp_wrapper(
tests/kernels/moe/utils.py
+7
-7
@@ -33,6 +33,7 @@ from vllm.model_executor.layers.fused_moe.prepare_finalize.batched import (
33
)
34
from vllm.model_executor.layers.fused_moe.router.fused_topk_router import fused_topk
35
from vllm.model_executor.layers.fused_moe.utils import moe_kernel_quantize_input
36
+from vllm.model_executor.layers.quantization.utils.fp8_utils import is_fp8
37
from vllm.model_executor.layers.quantization.utils.nvfp4_emulation_utils import (
38
ref_nvfp4_quant,
39
)
@@ -271,11 +272,10 @@ def moe_quantize_weights_2d(
272
per_token_quant: bool,
273
block_shape: list[int] | None,
274
) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None]:
274
- assert (
275
- quant_dtype == torch.float8_e4m3fn
276
- or quant_dtype == torch.int8
277
- or quant_dtype == "nvfp4"
278
- ), "only fp8/int8/nvfp4 supported"
275
+ is_fp8_dtype = is_fp8(quant_dtype)
276
+ assert is_fp8_dtype or quant_dtype == torch.int8 or quant_dtype == "nvfp4", (
277
+ "only fp8/int8/nvfp4 supported"
278
+ )
279
280
w_gs = None
281
@@ -283,7 +283,7 @@ def moe_quantize_weights_2d(
283
assert not per_token_quant
284
if quant_dtype == torch.int8:
285
w, w_s = per_block_cast_to_int8(w, block_shape)
286
- elif quant_dtype == torch.float8_e4m3fn:
286
+ elif is_fp8_dtype:
287
w, w_s = per_block_cast_to_fp8(w, block_shape)
288
elif quant_dtype == "nvfp4":
289
raise RuntimeError("blocked quantization not supported for nvfp4")
@@ -294,7 +294,7 @@ def moe_quantize_weights_2d(
294
w, w_s = ops.scaled_int8_quant(
295
w, w_s, use_per_token_if_dynamic=per_token_quant
296
)
297
- elif quant_dtype == torch.float8_e4m3fn:
297
+ elif is_fp8_dtype:
298
w, w_s = ops.scaled_fp8_quant(
299
w, w_s, use_per_token_if_dynamic=per_token_quant
300
)