[ROCm][CI] aiter kernel ops - enable rope test (#52854)
Signed-off-by: Divakar Verma <divakar.verma@amd.com>
Divakar Verma committed
Aug 21, 2026 at 16:17 UTC
29b7c2f7d41f326a8d7227c710d2e1ffe79668d6
1 file changed
+7
-23
tests/kernels/core/test_rocm_aiter_ops.py
+7
-23
@@ -256,11 +256,6 @@ def test_rocm_aiter_rmsnorm_with_add_vs_torch():
256
)
257
258
259
-@pytest.mark.xfail(
260
- strict=True,
261
- raises=AssertionError,
262
- reason="AITER Triton RoPE precision issue: https://github.com/ROCm/aiter/issues/4765",
263
-)
259
def test_rocm_aiter_triton_rotary_embedding_vs_torch():
260
"""rocm_aiter_triton_rotary_embedding matches manual NeoX-style RoPE reference.
261
@@ -298,12 +293,13 @@ def test_rocm_aiter_triton_rotary_embedding_vs_torch():
293
sin_pos = sin_half[positions] # [num_tokens, half_dim]
294
295
def apply_rope_ref(t: torch.Tensor) -> torch.Tensor:
301
- t_r = t.float().view(num_tokens, num_heads, head_size)
302
- c = cos_pos.float().unsqueeze(1) # [num_tokens, 1, half_dim]
303
- s = sin_pos.float().unsqueeze(1)
296
+ # AITER triton keeps native dtype (no fp32 upcast), so reference must match.
297
+ t_r = t.view(num_tokens, num_heads, head_size)
298
+ c = cos_pos.unsqueeze(1) # [num_tokens, 1, half_dim]
299
+ s = sin_pos.unsqueeze(1)
300
x1, x2 = t_r[..., :half_dim], t_r[..., half_dim:]
301
rotated = torch.cat([x1 * c - x2 * s, x2 * c + x1 * s], dim=-1)
306
- return rotated.to(t.dtype).view(num_tokens, num_heads * head_size)
302
+ return rotated.view(num_tokens, num_heads * head_size)
303
304
ref_q = apply_rope_ref(query)
305
ref_k = apply_rope_ref(key)
@@ -320,20 +316,8 @@ def test_rocm_aiter_triton_rotary_embedding_vs_torch():
316
True, # is_neox style -> rotate_style=0
317
)
318
323
- _assert_close_budget(
324
- q_aiter.float(),
325
- ref_q.float(),
326
- label="triton_rope query",
327
- atol=1e-3,
328
- rtol=1.6e-2,
329
- )
330
- _assert_close_budget(
331
- k_aiter.float(),
332
- ref_k.float(),
333
- label="triton_rope key",
334
- atol=1e-3,
335
- rtol=1.6e-2,
336
- )
319
+ torch.testing.assert_close(q_aiter, ref_q)
320
+ torch.testing.assert_close(k_aiter, ref_k)
321
322
323
def test_rocm_aiter_act_mul_fp8_group_quant_roundtrip():