Skip to content

[Bugfix][Test] Rotate KV into RHT domain in the INT4 unified attn test #62881

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
jikunshang merged 1 commit into from
Nov 4, 2026
+21 −18
Merged

[Bugfix][Test] Rotate KV into RHT domain in the INT4 unified attn test #62881

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 0 additions & 1 deletion .buildkite/intel_jobs/quantization_intel.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
38 changes: 21 additions & 17 deletions tests/quantization/test_per_token_kv_cache.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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):
Expand Down Expand Up @@ -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)
Loading