[LZ-QWQ]

Background

Triton attention can overflow int32 slot * stride before pointer addition, producing wrong reads or illegal accesses at element offsets >= 2**31 (4 GiB into an FP16/BF16 K or V buffer). This depends on the physical cache address, not context length.

On H100 80GB, Qwen/Qwen3.5-0.8B with its default pool (5.70 GiB per K/V buffer) crashes on three short requests when Graph padding reads the pool-tail dummy slot. The fix completes the same requests with the same pool, matching the bounded-pool reference's tokens and recorded top-5 IDs/logits.

Fix

• Bound offsets using both K/V shapes and strides: int32 when the maximum offset is below 2**31, int64 otherwise. Small pools retain int32 arithmetic.
• Cast slots before multiplication in paged, grouped decode, extend and both split-extend prefix loops. Page-table storage stays int32.
• Add eight numerical regression cases covering four paths at signed overflow and full 32-bit wraparound, using an independent PyTorch reference.

Validation

Against main 9cacd0896c9b, all five GPUs below show 8 numerical failures before -> 8 passes after, plus 32 layout checks each. The int32/int64 implementation passes the related suites: 76 passed/4 optional-dependency skips on each AMD GPU; 78 passed/2 skips on each NVIDIA GPU.

Kernel latency change versus main, positive = slower. BF16, Q=128 (decode Q=1), KV=8192; pools select int64, while low-slot accesses keep main correct. Eight interleaved Graph timing rounds after warmup:

GPU	Paged	Decode	Extend	Split-extend
W7900D	+0.76%	-0.25%	+0.35%	-0.35%
R9700	-1.12%	+0.17%	-0.01%	-1.08%
RTX4090	+3.87%	+1.79%	+1.05%	+6.45%
RTX5090	+2.46%	-0.22%	+9.13%	+7.54%
H100	+6.00%	+2.92%	+5.93%	+9.19%

At Q=1, H100 paged attention is +13.61%. Small-pool deltas range from -1.25% to +1.94%; small differences may be noise.

H100 end-to-end, same Qwen checkpoint and 5,976,310-slot pool, identical output tokens. Output tokens/s; main Graph crashes, so only the fixed Graph throughput is reported:

Input / batch / output	Main eager -> fix eager	Fix Graph
128 / 4 / 128	379.17 -> 370.97	2226.06
8192 / 1 / 128	99.27 -> 97.43	514.38

Eager throughput is 2.16% and 1.85% lower on these workloads. AMD end-to-end impact remains unmeasured; these are focused checks, not full CI certification.

Minimal reproduction and environment

Copy this PR's test file into the baseline checkout for the before run; use the PR checkout for the after run. Each needs an installed FreeToken environment and approximately 9 GiB of available GPU memory.

PYTHONPATH=python python -m pytest tests/kernels/test_triton_attention.py \
  -k high_physical_slots -o addopts= -q

Expected: eight numerical failures on main, eight passes with the fix.

Tested with Torch 2.14.0 / Triton 3.8.0; CUDA 13.0 on NVIDIA; ROCm 10.1.0 Torch build / HIP 7.16.26385 on AMD.

Model A/B used greedy sampling with ignore_eos=True, cache flush per repetition, one excluded warmup, synchronized timing, and four repetitions in each of two processes per configuration.

[jason-fxz]

Thanks for the fix. The uint32 tier is not needed. Please just use int32 when the max offset is below 2**31 and int64 otherwise:

 def _kv_offset_dtype(k_cache: torch.Tensor, v_cache: torch.Tensor):
-    # Bound physical offsets, not request lengths. Nonnegative offsets can use
-    # all 32 bits before falling back to int64, even when slot IDs fit in int32.
+    # Bound physical offsets, not request lengths.
     max_offset = max(
         sum((size - 1) * stride for size, stride in zip(cache.shape, cache.stride()))
         for cache in (k_cache, v_cache)
     )
-    if max_offset < 1 << 31:
-        return tl.int32
-    return tl.uint32 if max_offset < 1 << 32 else tl.int64
+    return tl.int32 if max_offset < 1 << 31 else tl.int64

Please also update the description to match.