[CI][Bugfix] Relax packed_qk_rope_ correctness test to one ULP (#59008)
Signed-off-by: Stefan Koncarevic <stefan.koncarevic@amd.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com>
stefankoncarevic committed
Sep 28, 2026 at 19:29 UTC
6c300ddb9e9d9582324032ee68525aae60c9800c
2 files changed
+22
-7
tests/kernels/core/test_apply_rotary_emb.py
+18
-5
@@ -233,9 +233,15 @@ def test_apply_rotary_emb_xpu_matches_native(
233
def test_packed_qk_rope_correctness(
234
num_tokens: int, dtype: torch.dtype, default_vllm_config
235
):
236
- """packed_qk_rope_ must be bitwise identical to the per-tensor
237
- ApplyRotaryEmb path with enable_fp32_compute=True, and must leave the V
238
- slice untouched."""
236
+ """packed_qk_rope_ must match the per-tensor ApplyRotaryEmb path with
237
+ enable_fp32_compute=True to within one ULP, and must leave the V slice
238
+ untouched.
239
+
240
+ Both sides compute in fp32 on identical inputs. The only gap is which
241
+ product the backend keeps exact inside the o1 multiply-add, and that
242
+ follows the register layout, so it is decided by codegen rather than by
243
+ the source and holds neither across backends nor across Triton versions.
244
+ """
245
from vllm.model_executor.layers.rotary_embedding.common import ApplyRotaryEmb
246
from vllm.model_executor.layers.rotary_embedding.packed_qk_rope import (
247
packed_qk_rope_,
@@ -265,6 +271,13 @@ def test_packed_qk_rope_correctness(
271
xk_ref = op(xk, cos, sin)
272
273
packed_qk_rope_(xqkv, freqs_cis)
268
- torch.testing.assert_close(xqkv[:, 0], xq_ref, atol=0, rtol=0)
269
- torch.testing.assert_close(xqkv[:, 1], xk_ref, atol=0, rtol=0)
274
+
275
+ # One ULP at the largest magnitude present, not relative: o1 can cancel,
276
+ # so the error scales with the inputs to the sum, not with the result.
277
+ eps = torch.finfo(dtype).eps
278
+ q_atol = eps * xq_ref.abs().max().item()
279
+ k_atol = eps * xk_ref.abs().max().item()
280
+ torch.testing.assert_close(xqkv[:, 0], xq_ref, atol=q_atol, rtol=0)
281
+ torch.testing.assert_close(xqkv[:, 1], xk_ref, atol=k_atol, rtol=0)
282
+ # V is never read or written by the kernel, so it stays exact.
283
torch.testing.assert_close(xqkv[:, 2], xv, atol=0, rtol=0)
vllm/model_executor/layers/rotary_embedding/packed_qk_rope.py
+4
-2
@@ -46,8 +46,10 @@ def packed_qk_rope_(xqkv: torch.Tensor, freqs_cis: torch.Tensor) -> None:
46
freqs_cis: contiguous (seqlen, headdim // 2) complex64 rotary freqs,
47
read through its interleaved re/im fp32 view.
48
49
- Bitwise identical to ``ApplyRotaryEmb(enable_fp32_compute=True)`` applied
50
- to Q and K separately, but in one kernel with ~5x less memory traffic.
49
+ Matches ``ApplyRotaryEmb(enable_fp32_compute=True)`` applied to Q and K
50
+ separately to within one ULP, but in one kernel with ~5x less memory
51
+ traffic. Both compute in fp32; the gap is which product the compiler
52
+ keeps exact inside the fused multiply-add, which differs per backend.
53
Requires triton: callers must check HAS_TRITON and use the unfused path
54
otherwise. Contract violations raise.
55