[ROCm][CI]Fix test_concat_and_cache_mla_rope_fused on ROCm (#46409)
Signed-off-by: Divakar Verma <divakar.verma@amd.com>
Divakar Verma committed
Jun 27, 2026 at 00:38 UTC
68ee8300a047db78fb52bac477daaaac7be11216
1 file changed
+39
-1
tests/kernels/core/test_rotary_embedding_mla_cache_fused.py
+39
-1
@@ -17,6 +17,36 @@ from vllm.platforms import current_platform
17
from vllm.utils.torch_utils import set_random_seed
18
19
20
+@pytest.fixture
21
+def default_vllm_config(monkeypatch):
22
+ """Enable the AITER triton rope on ROCm for fp16-consistent numerics.
23
+
24
+ The fused CUDA kernel runs native fp16 while forward_native upcasts to
25
+ fp32, so on ROCm we route through the AITER triton rope (+rotary_embedding)
26
+ to match. Its env gates are cached at import, hence refresh_env_variables().
27
+ """
28
+ from vllm._aiter_ops import rocm_aiter_ops
29
+ from vllm.config import CompilationConfig, VllmConfig, set_current_vllm_config
30
+
31
+ is_rocm = current_platform.is_rocm()
32
+ if is_rocm:
33
+ config = VllmConfig(
34
+ compilation_config=CompilationConfig(custom_ops=["+rotary_embedding"])
35
+ )
36
+ else:
37
+ config = VllmConfig()
38
+ try:
39
+ with monkeypatch.context() as m, set_current_vllm_config(config):
40
+ if is_rocm:
41
+ m.setenv("VLLM_ROCM_USE_AITER", "1")
42
+ m.setenv("VLLM_ROCM_USE_AITER_TRITON_ROPE", "1")
43
+ rocm_aiter_ops.refresh_env_variables()
44
+ yield config
45
+ finally:
46
+ if is_rocm:
47
+ rocm_aiter_ops.refresh_env_variables()
48
+
49
+
50
@pytest.mark.parametrize("dtype", [torch.half, torch.bfloat16, torch.float])
51
@pytest.mark.parametrize("is_neox_style", [False, True])
52
@pytest.mark.parametrize("seq_len", [11, 42])
@@ -151,6 +181,10 @@ def test_concat_and_cache_mla_rope_fused(
181
kv_cache_scale,
182
)
183
184
+ # ROCm neox-style Triton FMA diverges slightly from the fused kernel, so
185
+ # relax the affected tolerance: rtol for fp8 (one e4m3 ULP ~12.5%) and atol
186
+ # otherwise (bounded ~6e-4). Other paths use the CUDA defaults.
187
+ rocm_neox = current_platform.is_rocm() and is_neox_style
188
if kv_cache_dtype == "fp8":
189
result_temp = torch.empty_like(kv_cache, dtype=torch.float16)
190
ops.convert_fp8(
@@ -163,7 +197,11 @@ def test_concat_and_cache_mla_rope_fused(
197
ops.convert_fp8(
198
expected_temp, ref_kv_cache, kv_cache_scale.item(), kv_dtype=kv_cache_dtype
199
)
166
- torch.testing.assert_close(result_temp, expected_temp, atol=0.001, rtol=0.1)
200
+ torch.testing.assert_close(
201
+ result_temp, expected_temp, atol=0.001, rtol=0.15 if rocm_neox else 0.1
202
+ )
203
+ elif rocm_neox:
204
+ torch.testing.assert_close(kv_cache, ref_kv_cache, atol=1e-3, rtol=1e-3)
205
else:
206
torch.testing.assert_close(kv_cache, ref_kv_cache)
207