From 32c163dd99830dea5706f02be5d66e208f84f8f1 Mon Sep 17 00:00:00 2001 From: Michael Stufflebeam Date: Sun, 13 Sep 2026 17:09:54 -0500 Subject: [PATCH] prefill-attn-int8: cast the block id to int64 before the stride multiply (the #86 pattern, second kernel) All three kernels in prefill_attn_hd256.py (_k_stats_kernel, _k_quant_kernel, _prefill_attn_kernel) loaded the block id from the int32 block table and multiplied it by the cache stride in int32, so on a pool larger than 2^31 / stride_kb blocks (bf16, block 880: block id 2,383 and up) a request high in the pool read the wrong memory, exactly as #86 in the split-KV verify kernel. The #86 review called the other kernels unaudited; this is the sweep's one hit (the int4, KVarN and Mamba kernels already cast). bench/test_prefill_attn_bigpool.py is the synthetic proof, card-agnostic: one request's 20,000 tokens at the top of a 2,600-block pool and again at the bottom with identical bytes, the last 2,048-token chunk run at both placements on the exact-QK and int8-QK paths; outputs must match bit for bit and both are checked against an fp32 reference. Before the cast the top-of-pool run dies with an illegal memory access; after it, top equals bottom on both paths, max error 0.0002 (exact QK) and 0.0013 (int8 QK) against the reference. Unreachable on a 24 GB card, as with #86; it is for the 64 GB cards running PREFILL_ATTN=int8. --- bench/test_prefill_attn_bigpool.py | 56 ++++++++++++++++++++++++++++++ patches/prefill-attn-int8.patch | 5 ++- 2 files changed, 60 insertions(+), 1 deletion(-) create mode 100644 bench/test_prefill_attn_bigpool.py diff --git a/bench/test_prefill_attn_bigpool.py b/bench/test_prefill_attn_bigpool.py new file mode 100644 index 00000000..b0e4db97 --- /dev/null +++ b/bench/test_prefill_attn_bigpool.py @@ -0,0 +1,56 @@ +"""The #86 pattern in the int8 prefill kernel (patches/prefill-attn-int8.patch): the block id is loaded from +the int32 block table and multiplied by the cache stride in int32, in three kernels (_k_stats_kernel, +_k_quant_kernel, _prefill_attn_kernel). On a pool larger than 2^31 elements / stride_kb (bf16, block 880: +901,120 elements per block, so block ids above 2,383) a request whose blocks sit high in the pool reads the +wrong memory. Places one request's blocks at the TOP of a 2,600-block pool and compares the kernel's output +with the same request placed at the BOTTOM of the same pool (identical K/V bytes, so the two must agree to +the bit whatever the int8 rounding does), plus a loose check against an fp32 reference on the exact-QK path. +Run inside the vLLM venv after the patches are applied (any card with ~10 GB free): + venv/bin/python bench/test_prefill_attn_bigpool.py +Exit code 0 = all OK; a CUDA illegal access kills the process (that is the 'before' result).""" +import sys, torch +from vllm.v1.attention.backends.prefill_attn_hd256 import prefill_attn +torch.manual_seed(0); dev = "cuda" +Hq, Hkv, D, BS = 24, 4, 256, 880 +scale = D ** -0.5 +NB = 2600 # blocks in the pool: NB*BS*Hkv*D = 2.34e9 elements > 2^31 +L, Q = 20000, 2048 # context tokens and the last chunk's query rows (chunked prefill: q_start_pos = L - Q) +nb_req = (L + BS - 1) // BS +kc = torch.empty(NB, BS, Hkv, D, device=dev, dtype=torch.bfloat16).normal_() +vc = torch.empty(NB, BS, Hkv, D, device=dev, dtype=torch.bfloat16).normal_() +top = torch.arange(NB - nb_req, NB, device=dev, dtype=torch.int32) # the highest block ids +bottom = torch.arange(0, nb_req, device=dev, dtype=torch.int32) # the same request, low in the pool +kc[bottom.long()] = kc[top.long()]; vc[bottom.long()] = vc[top.long()] # identical bytes at both placements +q = torch.randn(Q, Hq, D, device=dev, dtype=torch.bfloat16) + +def run(bt, int8_qk): + out = torch.empty_like(q) + prefill_attn(q, kc, vc, out, bt, L, L - Q, scale, int8_qk=int8_qk) + torch.cuda.synchronize() + return out + +def reference(): + k = kc[bottom.long()].reshape(-1, Hkv, D)[:L].float(); v = vc[bottom.long()].reshape(-1, Hkv, D)[:L].float() + G = Hq // Hkv; k = k.repeat_interleave(G, 1); v = v.repeat_interleave(G, 1) + s_ = torch.einsum("qhd,khd->hqk", q.float(), k) * scale + qpos = torch.arange(L - Q, L, device=dev)[:, None]; kpos = torch.arange(L, device=dev)[None, :] + s_ = s_.masked_fill((kpos > qpos)[None], float("-inf")); p = torch.softmax(s_, -1) + return torch.einsum("hqk,khd->qhd", p, v) + +print(" (pre-fix expectation: the top-of-pool arm dies with a CUDA illegal memory access, or returns nan, or", flush=True) +print(" disagrees with the bottom-of-pool arm; measured both ways on two cards. That IS the bug, not a broken test)", flush=True) +ok = True +ref = reference() +for int8_qk in (False, True): + lo = run(bottom, int8_qk) + hi = run(top, int8_qk) + same = torch.equal(lo, hi) + e_lo = (lo.float() - ref).abs().max().item() + e_hi = (hi.float() - ref).abs().max().item() + tol = 0.05 if not int8_qk else 0.5 # the int8-QK path is approximate by design; the bit-equality is the sharp check + row_ok = same and e_lo < tol and e_hi < tol + ok = ok and row_ok + print(f" int8_qk={int8_qk}: pool {NB} blocks, request on blocks {NB-nb_req}..{NB-1} vs 0..{nb_req-1}, L={L}, chunk {Q}: " + f"top==bottom {'yes' if same else 'NO'}; max|ours-ref| bottom {e_lo:.4f} top {e_hi:.4f} {'OK' if row_ok else 'FAIL'}", flush=True) +print("PREFILL BIGPOOL OK" if ok else "PREFILL BIGPOOL FAIL", flush=True) +sys.exit(0 if ok else 1) diff --git a/patches/prefill-attn-int8.patch b/patches/prefill-attn-int8.patch index b9f5798a..227e4202 100644 --- a/patches/prefill-attn-int8.patch +++ b/patches/prefill-attn-int8.patch @@ -44,7 +44,7 @@ Written against vLLM 0.27.1. --- /dev/null +++ b/v1/attention/backends/prefill_attn_hd256.py -@@ -0,0 +1,239 @@ +@@ -0,0 +1,242 @@ +"""Prefill attention for head_dim 256 on sm86: GQA-packed, paged-KV, causal, +with an int8-QK^T mode (SageAttention-style) that runs the score GEMM on int8 +tensor cores at 2x the fp16 rate. @@ -85,6 +85,7 @@ Written against vLLM 0.27.1. + pos = start + tl.arange(0, CHUNK) + ok = pos < kv_len + blk = tl.load(bt_ptr + pos // BLOCK_SIZE, mask=ok, other=0) ++ blk = blk.to(tl.int64) # the stride multiply below overflows int32 past block id 2,383 (#86) + slot = pos % BLOCK_SIZE + ptrs = k_ptr + blk[:, None] * stride_kb + slot[:, None] * stride_ks \ + + kvh * stride_kh + d[None, :] @@ -110,6 +111,7 @@ Written against vLLM 0.27.1. + ok = pos < kv_len + d = tl.arange(0, D) + blk = tl.load(bt_ptr + pos // BLOCK_SIZE, mask=ok, other=0) ++ blk = blk.to(tl.int64) # the stride multiply below overflows int32 past block id 2,383 (#86) + slot = pos % BLOCK_SIZE + ptrs = k_ptr + blk[:, None] * stride_kb + slot[:, None] * stride_ks \ + + kvh * stride_kh + d[None, :] @@ -177,6 +179,7 @@ Written against vLLM 0.27.1. + pos = start + tl.arange(0, BLOCK_N) + k_ok = pos < kv_len + blk = tl.load(bt_ptr + pos // BLOCK_SIZE, mask=k_ok, other=0) ++ blk = blk.to(tl.int64) # the stride multiply below overflows int32 past block id 2,383 (#86) + slot = pos % BLOCK_SIZE + + if INT8_QK: