[CUDA] Normalize the PER_CHANNEL K scale folded into Q for paged XQA

[Tianlei Wu (tianleiwu)]

Description

Stacked on #32515 — review that one first. The base branch is tlwu/20260909/int4_kv_per_channel, so the diff here is the single commit on top; it retargets to main once #32515 merges.

Paged XQA folds a PER_CHANNEL K scale into the query and stores the product in T, while the portable kernel keeps the same product in an FP32 shared-memory tile. A large scale saturates FP16 there, and a zero cache code then turns that infinity into NaN, so the two backends can disagree on otherwise valid input:

q = 100, k_scale = 1000  ->  XQA fold: inf (NaN once multiplied by a zero code)
                             portable: 100000.0

This addresses the correctness thread on #32515. INT4 is the most exposed because its scale spans max|K| / 7 instead of max|K| / 127, roughly 18x larger for identical data, and INT4 XQA is only eligible with PER_CHANNEL scales — but the fold itself is shared with INT8 and FP8, so this is not INT4-specific.

Approach

kCacheScale is already applied as a scalar into qkScale, once per CTA and outside the K/V loop. So the fix needs no change in the inner loop:

• Divide the fold by max|k_scale| (PagedMaxAbsScaleKernel, a single-block reduction launched on device so the step stays capturable).
• Pass that maximum to XQA as its scalar K scale; the kernel multiplies it back into qkScale.

The correction is exact, and the folded query is bounded by max|q| at any scale magnitude.

INT4 dequantizes into FP16 shared memory, so cacheElemSize == 2 and isKVCacheQuantized is false, which is why the scalar scale was previously ignored on that path. That predicate is now separate from "the cache elements are narrower than T". Both INT4 translation units define XQA_PAGED_INT4, so decode and speculative decode are covered.

Cost

No inner-loop work is added and shared memory is unchanged, so XQA eligibility is unaffected. The only new work is one small max reduction (1024 floats at H256) per node per step, amortized under CUDA-graph replay.

Validation status

Not built — there is no CUDA toolkit on the machine this was written on. This needs a GPU build and a perf comparison against #32515 before it leaves draft.

Host-side numerical checks against an FP32 reference:

scale regime	current fold	this change
q=100, k_scale=1000, zero code	NaN	18750.0, exact
`max	s	` 1e2 / 1e3
`max	s	` 1e4 / 1e5

Accuracy improves at ordinary magnitudes too, since normalizing makes better use of the FP16 mantissa.

test_int4_xqa_large_per_channel_k_scale_matches_portable was moved into the regime that previously produced NaN: it drives q * k_scale to 1e6 with every K code zero, which is the inf * 0 path. The codes are zeroed deliberately — with nonzero codes a scale that large makes softmax one-hot, and the argmax is then fp16-sensitive and flaky.

Reviewer notes

• This changes INT8 and FP8 PER_CHANNEL XQA as well, since they share the fold. Those paths previously passed a null scalar scale; they now receive the normalizer. Please exercise them alongside INT4.
• If the preference is to keep the blast radius on the new feature only, the normalizer can be passed for uint8 caches alone, leaving INT8/FP8 exactly as today.

[Copilot]

🟡 Changes recommended

The folding expression can still overflow before normalization, and affected INT8/FP8 extremes remain untested.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Pull request overview

Normalizes PER_CHANNEL K scales before paged XQA query folding to prevent FP16 saturation and NaNs.

Changes:

• Adds device-side maximum-scale reduction and normalization.
• Reapplies the normalizer inside XQA.
• Updates INT4 regression coverage and documentation.

File summaries

File	Description
test_paged_attention_int4.py	Tests large-scale INT4 parity.
xqa/mha_impl.cuh	Enables scalar scales for INT4 XQA.
paged_attention.cc	Allocates normalizer scratch memory.
paged_attention_impl.cu	Implements reduction, folding, and scale forwarding.
attention_data.h	Adds normalizer scratch pointer.
paged_attention.md	Documents normalized folding.

Review details

• Files reviewed: 6/6 changed files
• Comments generated: 3
• Review effort level: Balanced

---

💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.