From 7f88221f35b6b77a249986d972f853abf4982797 Mon Sep 17 00:00:00 2001 From: Dudi Lester Date: Wed, 14 Oct 2026 14:35:05 +0300 Subject: [PATCH] [Bugfix][Test] Rotate KV into the RHT domain in the INT4 unified attention test test_triton_unified_attention_per_token_head_scale filled the INT4 cache with unrotated K/V. unified_attention_int4 assumes the cache holds RHT-rotated data, as written by reshape_and_cache_int4: it rotates Q and inverse-rotates the output. The kernel output was therefore uncorrelated with the reference. The test passed only while every element stayed within atol=rtol=0.5, and failed on 1-2 elements for some seeds and platforms. This is why Hopper needed atol=1.1 and Intel CI deselected the int4 seq_lens1 case. Rotate K/V with single_rht before the asymmetric quantization, and inverse-rotate the dequantized reference back to the original domain, like the other INT4 tests in this file. Both runs now see the same KV, so INT4 uses the same 5e-2 tolerance as INT8/FP8. Drop the Hopper special case and the Intel CI deselect. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Dudi Lester --- .buildkite/intel_jobs/quantization_intel.yaml | 1 - tests/quantization/test_per_token_kv_cache.py | 38 ++++++++++--------- 2 files changed, 21 insertions(+), 18 deletions(-) diff --git a/.buildkite/intel_jobs/quantization_intel.yaml b/.buildkite/intel_jobs/quantization_intel.yaml index a51c6ebfc715..7f9fd884d35f 100644 --- a/.buildkite/intel_jobs/quantization_intel.yaml +++ b/.buildkite/intel_jobs/quantization_intel.yaml @@ -26,7 +26,6 @@ steps: bash .buildkite/scripts/hardware_ci/run-intel-test.sh 'VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s tests/quantization/test_per_token_kv_cache.py - --deselect="tests/quantization/test_per_token_kv_cache.py::test_triton_unified_attention_per_token_head_scale[int4-16-128-num_heads0-seq_lens1]" && VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s tests/quantization/test_auto_awq.py && VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s tests/quantization/test_fp8_per_channel.py && VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s diff --git a/tests/quantization/test_per_token_kv_cache.py b/tests/quantization/test_per_token_kv_cache.py index 2647e8b5a8e2..46bb9a29b22a 100644 --- a/tests/quantization/test_per_token_kv_cache.py +++ b/tests/quantization/test_per_token_kv_cache.py @@ -578,9 +578,10 @@ def test_triton_unified_attention_per_token_head_scale( value_cache_bf16 = torch.randn_like(key_cache_bf16) if is_int4: - # Asymmetric quantization reference (matches the Triton kernel). - kf = key_cache_bf16.float() - vf = value_cache_bf16.float() + # The INT4 cache lives in the RHT domain: rotate K/V before the + # asymmetric quantization, like reshape_and_cache_int4 does. + kf = single_rht(key_cache_bf16.float()).to(torch.bfloat16).float() + vf = single_rht(value_cache_bf16.float()).to(torch.bfloat16).float() k_min = kf.amin(dim=-1) k_max = kf.amax(dim=-1) v_min = vf.amin(dim=-1) @@ -597,11 +598,22 @@ def test_triton_unified_attention_per_token_head_scale( (vf / v_scale_cache[..., None] + v_zp[..., None]).round().clamp(0, 15) ) - # Dequantized reference: x_hat = (q - zp) * scale - key_cache_deq = (key_cache_q_full - k_zp[..., None]) * k_scale_cache[..., None] - value_cache_deq = (value_cache_q_full - v_zp[..., None]) * v_scale_cache[ - ..., None - ] + # Dequantized reference: x_hat = (q - zp) * scale, then IRHT/d back + # to the original domain the reference attention runs in. + key_cache_deq = ( + single_rht( + (key_cache_q_full - k_zp[..., None]) * k_scale_cache[..., None], + inverse=True, + ) + / head_size + ) + value_cache_deq = ( + single_rht( + (value_cache_q_full - v_zp[..., None]) * v_scale_cache[..., None], + inverse=True, + ) + / head_size + ) # Pack two uint4 values into one byte def _pack_int4(data_float): @@ -725,12 +737,4 @@ def _pack_int4(data_float): softmax_segm_expsum=softmax_segm_expsum, ) - # Coarser quantization → wider tolerance. - if is_int4: - # Hopper's attention reduction order can move a few BF16 elements by - # just over 1.0 after INT4 quantization. - atol = 1.1 if current_platform.is_device_capability_family(90) else 0.5 - rtol = 0.5 - else: - atol, rtol = 5e-2, 5e-2 - torch.testing.assert_close(output_q, output_ref, atol=atol, rtol=rtol) + torch.testing.assert_close(output_q, output_ref, atol=5e-2, rtol=5e-2)