[cpuchip]

The #86 pattern, in a second kernel. The int8 prefill attention (patches/prefill-attn-int8.patch, prefill_attn_hd256.py) loads the block id from the int32 block table and multiplies it by the cache stride in int32, in all three of its kernels (_k_stats_kernel, _k_quant_kernel, _prefill_attn_kernel). On a pool larger than 2^31 elements divided by the block stride (bf16, block 880: block id 2,383 and up) a request whose blocks sit high in the pool reads the wrong memory. The #86 review said the other Triton kernels in the series were unaudited; this is the sweep's one hit. The int4, KVarN and Mamba kernels already cast.

1. patches/prefill-attn-int8.patch: blk = blk.to(tl.int64) after each of the three block-table loads. Nothing else in the file changes; the hunk on flash_attn.py is untouched.
2. bench/test_prefill_attn_bigpool.py: the same synthetic test shape as bench/test_spec_decode_bigpool.py, with a sharper oracle: one request's 20,000 tokens are placed at the top of a 2,600-block pool (blocks 2,577 to 2,599) and again at the bottom (0 to 22) with identical bytes, and the last 2,048-token prefill chunk is run at both placements on the exact-QK path and the int8-QK path; the two outputs must be bit-identical, and both are checked against an fp32 reference. Card-agnostic; needs about 10 GB free.

Measured on two boxes. WSL2, inside the 0.29-series image with the kernel file swapped: before the cast, the top-of-pool run dies with CUDA error: an illegal memory access was encountered; after it, top equals bottom bit for bit on both paths, max error against the reference 0.0002 (exact QK) and 0.0013 (int8 QK). Native 3090, against this branch's tree with a guard that the interpreter had loaded the 239-line zero-cast module: before the cast the top-of-pool run does not crash, it returns nan on both paths while the bottom-of-pool run still matches the reference to 0.0002; after the cast, top equals bottom, max error 0.0002 / 0.0010, exit 0. So the before-arm's tell is not "did it crash": on one card the wrong read is an illegal access and on another it is a silent nan, which is the worse failure, and the test's bit-equality against the bottom placement is what catches both.

Cost of the cast, measured after the first version of this body said it was not. Paired rows under the production launch line (CTX=fast SPEC=dflash2 DFLASH_TOKENS=7 PREFIX_CACHE=1 PREFILL_ATTN=int8 INT8_ACT=int8), fresh prefills of about 24k and 47k tokens, three each, medians, same WSL2 3090, images that differ by this cast and the #91 one: 4579 / 4011 tok/s with the casts against 4615 / 4051 without, so 0.8% / 1.0%, inside the run-to-run spread of that row. On a native 3090 with the cast in, int8 prefill against FA2 on the same boot config with the prefix cache off reads 2244 against 2190 tok/s at n=3 (direction only; separate boots on that box move a row by more than the gap), so the cast does not turn the kernel's gain into a loss.

Reach: as with #86, a 24 GB card never allocates a block id above about 2,383 on this profile, so this cannot fire on the reference 3090 or on this box; it is for the 64 GB cards, where the int8 prefill path (PREFILL_ATTN=int8, on by default in the production launch line) would read past its pool at the same block ids that killed the verify kernel.

Series: in patches/series order against a pristine v0.28.0 checkout, 32 applied, 0 rejected (33 entries, dflash2-backport skipped by design); the regenerated prefill-attn-int8.patch lands with one offset hunk and no fuzz, its new-file hunk 242 lines.

One thing this PR does not touch, found while measuring it: the patch reads VLLM_PREFILL_ATTN with a raw os.environ.get at two sites in flash_attn.py and never registers it, so every boot with PREFILL_ATTN set prints "Unknown vLLM environment variable detected: VLLM_PREFILL_ATTN". That read predates this PR and a rename diff only shows the changed lines, which is why the gate did not flag it here. The 0.29 port (#106) registers it in envs.py and reads it through envs, together with the other launcher-exported knobs that print the same line; say if you want the 0.28 line to carry the same change in this PR or in its own.

What is still unproven: no serving run at high block ids on any card (the same gap #86 left; the synthetic test is the proof); the three kernels' other index arithmetic (pos * stride_k8n, slot * stride_ks) stays int32 by construction because those operands are bounded by the block size and chunk, not the pool, and that is argued from the source, not measured; the cost figure is one box's paired rows and the native row is a direction, not a magnitude.

[cpuchip]

Body updated with two things learned while vetting the 0.29 port, no change to the diff:

1. The cast's cost is measured now instead of "expected": paired rows under the production launch line, three fresh prefills each at about 24k and 47k tokens, medians 4579 / 4011 tok/s with the casts against 4615 / 4051 without (0.8% / 1.0%, inside that row's spread) on a WSL2 3090; on a native 3090 with the cast in, int8 prefill still reads above FA2 at n=3 with the prefix cache off (2244 against 2190, direction only).
2. This patch reads VLLM_PREFILL_ATTN with a raw os.environ.get and never registers it, so every boot with PREFILL_ATTN set prints "Unknown vLLM environment variable detected". That predates this PR and is outside its diff; vLLM 0.29.0 port is ready on our fork (patch series, KVarN, #100 and #101 carried); PR when you want it #106's port registers it with the other launcher-exported knobs. Your call whether the 0.28 line gets it here or separately.

[mhenrichsen]

Merging. The cast is right, the reproducer discriminates, and it costs nothing on our card.

Unlike #91 (where the first push was test-only), the kernel change is really in the diff here: the new-file hunk header goes 239 -> 242 and three blk = blk.to(tl.int64) lines land, one after each block-table load in _k_stats_kernel, _k_quant_kernel and _prefill_attn_kernel. That is the full coverage: every K/V pointer expression downstream reuses that same blk, and the only other pointer bases index by pos/tok, which are bounded by the chunk rather than the pool. Nothing else in the kernel multiplies a block id.

Reproduced both arms on a native 3090, production stopped, card free before each launch, power limit untouched:

# arm A: installed kernel as it is on main (239 lines, 0 casts), triton + __pycache__ wiped
$ venv/bin/python bench/test_prefill_attn_bigpool.py
int8_qk=False: top==bottom NO; max|ours-ref| bottom 0.0002 top nan  FAIL
int8_qk=True : top==bottom NO; max|ours-ref| bottom 0.0010 top nan  FAIL

# arm B: the same file with your three casts (242 lines), caches wiped again
int8_qk=False: top==bottom yes; max|ours-ref| bottom 0.0002 top 0.0002 OK
int8_qk=True : top==bottom yes; max|ours-ref| bottom 0.0010 top 0.0010 OK

Identical bytes at the top of a 2,600-block pool came back nan (we get nan rather than your illegal access; the test accepts either) and the cast is the thing that flips it. Your test is also the only unit coverage this repo has for prefill_attn_hd256.py, and its call matches the real signature, so neither arm can silently skip.

Cost, since PREFILL_ATTN=int8 is live for us: paired, one process, both module versions loaded, arms alternated per rep, n=15, production prefill geometry (chunk 2048, 24/4 heads, D=256, block 880), two independent runs.

int8_qk=True   L=24000: +0.94% / +0.32%   (run 1 / run 2, median)
int8_qk=True   L=47000: -1.99% / -1.04%
int8_qk=False  L=24000: -3.84% / -4.36%
int8_qk=False  L=47000: -3.40% / -3.08%

The production path is inside +-2% with the sign flipping between shapes in both runs, i.e. the same codegen noise we concluded on #91. The exact-QK path coming out 3-4% faster with a cast added is the tell: read it as Triton codegen luck, not as a speedup this PR delivers. Your WSL2 rows (-0.8% / -1.0%) disagree with ours only in sign, not in magnitude; both say "nothing".

Series integrity against a pristine v0.28.0 clone: bash patches/check_vllm_series.sh <pristine> gives 32 patches applied, 0 offsets, on this branch and on main alike. No new patch file, so patches/series is unchanged. Production restored afterwards: kernel back to md5 36da7ed4, health 200, 250 W.

Three things stated honestly rather than left implied:

1. Reach. This is unreachable at 24 GB: overflow needs block id > 2^31/(8804256) ~= 2,383, and our pool never gets near that. So this proves the kernel was buggy and your cast fixes it, not that our deployment was ever wrong. It is a fix for 64 GB users on PREFILL_ATTN=int8, and inert-but-free for everyone else.

2. Correcting your PR body on one point, and it is our error not yours: single-user/start_qwen.sh leaves PREFILL_ATTN empty by default, so int8 prefill is not on in a stock clone. It is on for us only via our own int8-prefill.conf service drop-in.

3. Your envs.py question: separately, not here. The raw os.environ.get predates this PR; let vLLM 0.29.0 port is ready on our fork (patch series, KVarN, #100 and #101 carried); PR when you want it #106's 0.29 port carry the registration and we will backport a one-liner to the 0.28 line ourselves if the "Unknown vLLM environment variable" line starts annoying us.

Not verified: I hand-installed the three casts into the installed .py to isolate the kernel change, so the full patch -p1 install path is covered only by the series check above; and pass 2 of check_vllm_series.sh does not evaluate cleanly in my staging of the pristine tree, but it fails identically on main, so it carries no signal about this PR. No server boot, no quality battery -- the change is kernel-local and I measured its only entry point in both arms.

#108 forward-references this fix, so it can go in next.