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, 04:00 UTC
[CI][BugFix][AMD][FP8] Fix test_rms_norm so it runs correctly on ROCm (#32372)
Signed-off-by: Randall Smith <ransmith@amd.com> Co-authored-by: Randall Smith <ransmith@amd.com>
rasmith committed
Jan 15, 2026 at 05:05 UTC
8853a50af2a3d2a12700740aa6def37fed3142da
1 file changed
+7
-5
tests/kernels/core/test_fused_quant_layernorm.py
+7
-5
@@ -14,9 +14,10 @@ from vllm.model_executor.layers.quantization.utils.fp8_utils import (
14
from vllm.model_executor.layers.quantization.utils.int8_utils import (
15
per_token_group_quant_int8,
16
)
17
+from vllm.platforms import current_platform
18
19
DTYPES = [torch.bfloat16, torch.float]
19
-QUANT_DTYPES = [torch.int8, torch.float8_e4m3fn]
20
+QUANT_DTYPES = [torch.int8, current_platform.fp8_dtype()]
21
VEC_HIDDEN_SIZES = [1024, 1025, 1027, 1029]
22
# Avoid combinatorial explosion with full Cartesian product
23
NUM_TOKENS_HIDDEN_SIZES = [
@@ -61,14 +62,14 @@ def ref_dynamic_per_token_or_block_quant(
62
group_size: list[int] | None,
63
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:
64
if scale_ub is not None:
64
- assert quant_dtype == torch.float8_e4m3fn
65
+ assert quant_dtype == current_platform.fp8_dtype()
66
67
# Norm
68
torch_out, residual = ref_rms_norm(rms_norm_layer, x, residual)
69
70
# Quant
71
if group_size is not None:
71
- if quant_dtype == torch.float8_e4m3fn:
72
+ if quant_dtype == current_platform.fp8_dtype():
73
torch_out, scales = per_token_group_quant_fp8(
74
torch_out, group_size=group_size[1], use_ue8m0=False
75
)
@@ -78,7 +79,7 @@ def ref_dynamic_per_token_or_block_quant(
79
torch_out, group_size=group_size[1]
80
)
81
else:
81
- if quant_dtype == torch.float8_e4m3fn:
82
+ if quant_dtype == current_platform.fp8_dtype():
83
torch_out, scales = ops.scaled_fp8_quant(
84
torch_out, scale_ub=scale_ub, use_per_token_if_dynamic=True
85
)
@@ -162,6 +163,7 @@ def test_rms_norm(
163
if torch.cuda.is_available():
164
torch.cuda.manual_seed(seed)
165
torch.set_default_device(device)
166
+ torch.cuda.set_device(device)
167
168
if group_size is not None and hidden_size % group_size[1] != 0:
169
# skip
@@ -171,7 +173,7 @@ def test_rms_norm(
173
# blockwise baseline doesn't support scale_ub
174
return
175
174
- if has_scale_ub and quant_dtype != torch.float8_e4m3fn:
176
+ if has_scale_ub and quant_dtype != current_platform.fp8_dtype():
177
# skip
178
return
179