[Bugfix][ROCm] Use the platform FP8 range in the concat MLA q test (#58153)
Co-authored-by: Andreas Karatzas <akaratza@amd.com>
stefankoncarevic committed
Sep 23, 2026 at 01:05 UTC
9646f5306475d95de6baf03a81da4203f30928a8
1 file changed
+9
-3
tests/kernels/test_concat_mla_q.py
+9
-3
@@ -156,7 +156,11 @@ def test_concat_mla_q_fp8_nope_and_rope(rope_dim, transposed, enabled):
156
"""FP8 query packing preserves values with empty RoPE and strided BMM output."""
157
from vllm.config import CompilationConfig, VllmConfig, set_current_vllm_config
158
from vllm.model_executor.layers.attention.mla_attention import _DecodeConcatQuantFP8
159
- from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape
159
+ from vllm.model_executor.layers.quantization.utils.quant_utils import (
160
+ GroupShape,
161
+ get_fp8_min_max,
162
+ )
163
+ from vllm.platforms import current_platform
164
165
torch.manual_seed(42)
166
shape = (8, 17, 512) if transposed else (17, 8, 512)
@@ -166,8 +170,9 @@ def test_concat_mla_q_fp8_nope_and_rope(rope_dim, transposed, enabled):
170
qp = torch.randn(17, 8, rope_dim, dtype=torch.bfloat16, device="cuda") * 128
171
scale = torch.tensor(0.3, device="cuda", dtype=torch.float32)
172
joined = torch.cat((ql, qp), dim=-1)
169
- expected = (joined.float() * scale.reciprocal()).clamp(-448, 448)
170
- expected = expected.to(torch.float8_e4m3fn)
173
+ fp8_min, fp8_max = get_fp8_min_max()
174
+ expected = (joined.float() * scale.reciprocal()).clamp(fp8_min, fp8_max)
175
+ expected = expected.to(current_platform.fp8_dtype())
176
custom_ops = ["none"] + (["+mla_decode_concat_quant_fp8"] if enabled else [])
177
config = VllmConfig(
178
compilation_config=CompilationConfig(mode=0, custom_ops=custom_ops)
@@ -175,5 +180,6 @@ def test_concat_mla_q_fp8_nope_and_rope(rope_dim, transposed, enabled):
180
with set_current_vllm_config(config):
181
quant = _DecodeConcatQuantFP8(static=True, group_shape=GroupShape.PER_TENSOR)
182
actual = quant(ql, qp, scale)
183
+ assert actual.dtype == current_platform.fp8_dtype()
184
torch.testing.assert_close(actual.float(), expected.float(), atol=0, rtol=0)
185
assert actual.is_contiguous()