A high-throughput and memory-efficient inference and serving engine for LLMs
[ROCm][CI] fix fp8 support for test_deepep_moe (#45302)
Signed-off-by: Divakar Verma <divakar.verma@amd.com>
Divakar Verma committed
Jun 12, 2026 at 01:16 UTC
1ce3cdc5c14f656f81f91e264d557e0b6e6fea54
1 file changed
+12
-9
tests/kernels/moe/test_deepep_moe.py
+12
-9
@@ -27,6 +27,7 @@ from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEKernel
27
from vllm.model_executor.layers.quantization.utils.fp8_utils import (
28
per_token_group_quant_fp8,
29
)
30
+from vllm.platforms import current_platform
31
from vllm.utils.import_utils import has_deep_ep
32
from vllm.utils.torch_utils import set_random_seed
33
from vllm.v1.worker.workspace import init_workspace_manager
@@ -64,7 +65,7 @@ def make_weights(
65
return w1, w2, None, None
66
67
# per-out-channel weight quantization
67
- assert dtype == torch.float8_e4m3fn
68
+ assert dtype == current_platform.fp8_dtype()
69
w1 = torch.empty((e, 2 * n, k), device="cuda", dtype=torch.float16)
70
w2 = torch.empty((e, k, n), device="cuda", dtype=torch.float16)
71
@@ -105,9 +106,11 @@ class TestTensors:
106
@staticmethod
107
def make(config: TestConfig, low_latency_mode: bool) -> "TestTensors":
108
# TODO (varun) - check that float16 works ?
108
- assert config.dtype in [torch.bfloat16, torch.float8_e4m3fn]
109
+ assert config.dtype in [torch.bfloat16, current_platform.fp8_dtype()]
110
token_dtype = (
110
- torch.bfloat16 if config.dtype == torch.float8_e4m3fn else config.dtype
111
+ torch.bfloat16
112
+ if config.dtype == current_platform.fp8_dtype()
113
+ else config.dtype
114
)
115
rank_tokens = (
116
torch.randn((config.m, config.k), device="cuda", dtype=token_dtype) / 10
@@ -216,10 +219,10 @@ def deep_ep_moe_impl(
219
return expert_map.to(device=device, dtype=torch.int32)
220
221
hidden_size = test_tensors.rank_tokens.size(1)
219
- is_quantized = w1.dtype == torch.float8_e4m3fn
222
+ is_quantized = w1.dtype == current_platform.fp8_dtype()
223
q_dtype = None
224
if is_quantized:
222
- q_dtype = torch.float8_e4m3fn
225
+ q_dtype = current_platform.fp8_dtype()
226
227
out_hidden_states = torch.empty_like(test_tensors.rank_tokens)
228
total_num_tokens = test_tensors.rank_tokens.size(0)
@@ -318,7 +321,7 @@ def torch_moe_impl(
321
.to(a.dtype)
322
)
323
321
- is_quantized = w1.dtype == torch.float8_e4m3fn
324
+ is_quantized = w1.dtype == current_platform.fp8_dtype()
325
a_dtype = a.dtype
326
if is_quantized:
327
w1 = w1.to(dtype=torch.float32) * w1_scale
@@ -367,7 +370,7 @@ def _deep_ep_moe(
370
"FP8 dispatch interface is available only in low-latency mode"
371
)
372
370
- is_quantized = w1.dtype == torch.float8_e4m3fn
373
+ is_quantized = w1.dtype == current_platform.fp8_dtype()
374
device_idx = torch.accelerator.current_device_index()
375
w1 = w1.to(device=device_idx)
376
w2 = w2.to(device=device_idx)
@@ -441,7 +444,7 @@ MNKs = [
444
(222, 1024, 2048),
445
]
446
444
-DTYPES = [torch.bfloat16, torch.float8_e4m3fn]
447
+DTYPES = [torch.bfloat16, current_platform.fp8_dtype()]
448
449
450
@pytest.mark.parametrize("dtype", DTYPES)
@@ -496,7 +499,7 @@ MNKs = [
499
(64, 1024, 2560),
500
(222, 1024, 2560),
501
]
499
-DTYPES = [torch.float8_e4m3fn, torch.bfloat16]
502
+DTYPES = [current_platform.fp8_dtype(), torch.bfloat16]
503
USE_FP8_DISPATCH = [True, False]
504
505