[ROCm][CI] Fix MLA RoPE fused-kernel tests for TheRock image (#57112)
Signed-off-by: Matthew Wong <Matthew.Wong2@amd.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com>
Matt committed
Sep 15, 2026 at 22:57 UTC
b3b13c12923f0b68f36ee515ebbc3efc7679eb4f
1 file changed
+121
-16
tests/kernels/core/test_rotary_embedding_mla_cache_fused.py
+121
-16
@@ -17,6 +17,88 @@ from vllm.platforms import current_platform
17
from vllm.utils.torch_utils import set_random_seed
18
19
20
+def assert_rope_close_rocm(
21
+ actual: torch.Tensor,
22
+ ref: torch.Tensor,
23
+ rope_input: torch.Tensor,
24
+ cos_sin_cache: torch.Tensor,
25
+ positions: torch.Tensor,
26
+ is_neox: bool,
27
+ *,
28
+ max_fail_fraction: float | None = None,
29
+ atol: float | None = None,
30
+ rtol: float | None = None,
31
+) -> None:
32
+ """assert_close for RoPE fused-vs-reference on ROCm, with a per-element
33
+ product-ULP allowance for catastrophic-cancellation elements.
34
+
35
+ The fused and reference kernels may contract the RoPE multiply-add into an
36
+ FMA with either product pre-rounded to the working dtype. The two valid
37
+ orderings differ by at most 2 ULPs of the product at that element, which
38
+ cancellation (x*cos ~= -y*sin) can amplify past the default tolerance.
39
+
40
+ Tier 1: |actual - ref| <= atol + rtol * |ref| (default tolerance).
41
+ Tier 2: elements failing tier 1 must each satisfy
42
+ |actual_e - ref_e| <= 2 * eps * |product_e|, where product_e is the larger
43
+ of the two products combined at that element, and their fraction must be
44
+ <= max_fail_fraction (1e-4 for fp16; 1e-2 for bf16, whose coarser eps
45
+ puts 1-ULP divergences past the default tolerance far more often).
46
+ """
47
+ if atol is None:
48
+ atol = get_default_atol(ref)
49
+ if rtol is None:
50
+ rtol = get_default_rtol(ref)
51
+ if max_fail_fraction is None:
52
+ max_fail_fraction = 1e-2 if ref.dtype == torch.bfloat16 else 1e-4
53
+ eps = torch.finfo(ref.dtype).eps
54
+
55
+ # per-element product magnitude, in the (seq, heads, rot_dim) layout
56
+ seq_len, rot_dim = rope_input.shape[0], rope_input.shape[-1]
57
+ half = rot_dim // 2
58
+ inp = rope_input.reshape(seq_len, -1, rot_dim).to(torch.float32)
59
+ cs = cos_sin_cache[positions].to(torch.float32) # (seq_len, 2*half)
60
+ cos = cs[:, :half].unsqueeze(1) # broadcast over heads
61
+ sin = cs[:, half:].unsqueeze(1)
62
+ x, y = inp[..., :half], inp[..., half:]
63
+ if is_neox:
64
+ # out = [x*cos - y*sin, y*cos + x*sin]
65
+ prod = torch.cat(
66
+ [
67
+ torch.maximum((x * cos).abs(), (y * sin).abs()),
68
+ torch.maximum((y * cos).abs(), (x * sin).abs()),
69
+ ],
70
+ dim=-1,
71
+ )
72
+ else:
73
+ # out[2i] = x[2i]*cos[i] - x[2i+1]*sin[i], interleaved
74
+ xe, xo = inp[..., 0::2], inp[..., 1::2]
75
+ prod = torch.stack(
76
+ [
77
+ torch.maximum((xe * cos).abs(), (xo * sin).abs()),
78
+ torch.maximum((xo * cos).abs(), (xe * sin).abs()),
79
+ ],
80
+ dim=-1,
81
+ ).flatten(-2)
82
+ prod = prod.reshape(rope_input.shape)
83
+
84
+ diff = (actual.to(torch.float32) - ref.to(torch.float32)).abs()
85
+ tol = atol + rtol * ref.to(torch.float32).abs()
86
+ fail = diff > tol
87
+ n_fail = int(fail.sum())
88
+ if n_fail == 0:
89
+ return
90
+ assert n_fail <= max_fail_fraction * diff.numel(), (
91
+ f"{n_fail} elements beyond default tolerance "
92
+ f"({n_fail / diff.numel():.2%} > {max_fail_fraction:.2%}); "
93
+ f"max diff {diff.max().item()}"
94
+ )
95
+ excess = diff[fail] - 2 * eps * prod[fail]
96
+ assert (excess <= 0).all(), (
97
+ f"{int((excess > 0).sum())} cancellation elements exceed 2 product-ULPs; "
98
+ f"worst excess {excess.max().item()}"
99
+ )
100
+
101
+
102
@pytest.fixture
103
def default_vllm_config(monkeypatch):
104
"""Enable the AITER triton rope on ROCm for fp16-consistent numerics.
@@ -99,6 +181,8 @@ def test_concat_and_cache_mla_rope_fused(
181
182
k_pe = torch.flatten(key[..., :qk_rope_head_dim], start_dim=1).to(device=device)
183
kv_c = torch.flatten(key[..., qk_rope_head_dim:], start_dim=1).to(device=device)
184
+ k_pe_orig = k_pe.clone()
185
+ query_orig = query.clone()
186
187
if current_platform.is_rocm():
188
# We use forward_hip for the same numerics as the fused custom kernel on ROCm
@@ -181,12 +265,9 @@ def test_concat_and_cache_mla_rope_fused(
265
kv_cache_scale,
266
)
267
184
- # On ROCm the AITER Triton rope diverges by ~1 ULP from the fused kernel,
185
- # which the tight CUDA tolerances don't cover for the low-precision paths:
186
- # - bf16: up to ~1 bf16 ULP (0.0156-0.03125 at values ~2-3), so relax the
187
- # kv-cache atol (bounded ~0.02) and the query atol (~0.04).
188
- # - neox-style fp8: one e4m3 ULP (~12.5%), so relax rtol.
189
- # Other paths use the CUDA defaults.
268
+ # On ROCm the AITER Triton rope diverges by up to 2 product-ULPs from the
269
+ # fused kernel, which the default tolerance misses at the rare
270
+ # catastrophic-cancellation elements. fp8 keeps its calibrated tolerance.
271
rocm_neox = current_platform.is_rocm() and is_neox_style
272
rocm_bf16 = current_platform.is_rocm() and dtype == torch.bfloat16
273
if kv_cache_dtype == "fp8":
@@ -207,16 +288,40 @@ def test_concat_and_cache_mla_rope_fused(
288
atol=0.004 if rocm_bf16 else 0.001,
289
rtol=0.15 if rocm_neox or rocm_bf16 else 0.1,
290
)
210
- elif rocm_bf16:
211
- torch.testing.assert_close(kv_cache, ref_kv_cache, atol=0.02, rtol=1e-3)
212
- elif rocm_neox:
213
- torch.testing.assert_close(kv_cache, ref_kv_cache, atol=1e-3, rtol=1e-3)
291
+ elif current_platform.is_rocm() and dtype in (torch.half, torch.bfloat16):
292
+ # gather the rope region of the cache back to token order
293
+ block_idx = slot_mapping // block_size
294
+ block_offset = slot_mapping % block_size
295
+ fused_k_rope = kv_cache[block_idx, block_offset, kv_lora_rank:]
296
+ ref_k_rope_gathered = ref_kv_cache[block_idx, block_offset, kv_lora_rank:]
297
+ assert_rope_close_rocm(
298
+ fused_k_rope,
299
+ ref_k_rope_gathered,
300
+ k_pe_orig,
301
+ rope.cos_sin_cache,
302
+ positions,
303
+ is_neox_style,
304
+ )
305
+ # the nope region is a byte-exact copy
306
+ torch.testing.assert_close(
307
+ kv_cache[..., :kv_lora_rank], ref_kv_cache[..., :kv_lora_rank]
308
+ )
309
else:
310
torch.testing.assert_close(kv_cache, ref_kv_cache)
311
217
- torch.testing.assert_close(
218
- query,
219
- ref_q_pe,
220
- atol=0.04 if rocm_bf16 else get_default_atol(query),
221
- rtol=get_default_rtol(query),
222
- )
312
+ if current_platform.is_rocm() and dtype in (torch.half, torch.bfloat16):
313
+ assert_rope_close_rocm(
314
+ query,
315
+ ref_q_pe,
316
+ query_orig,
317
+ rope.cos_sin_cache,
318
+ positions,
319
+ is_neox_style,
320
+ )
321
+ else:
322
+ torch.testing.assert_close(
323
+ query,
324
+ ref_q_pe,
325
+ atol=get_default_atol(query),
326
+ rtol=get_default_rtol(query),
327
+ )