Repository navigation
Conversation
…cp.async (STRATA_PA_V2T=1) On Turing the int8-KV prompt attention ran v1 because v2 stages its K/V slices with cp.async. Here a warp loads its next slice into registers and stores it into its other stage after the current chunk (the same bytes in the same places); the MMAs are the m16n8k8 pairs v1 already uses. Accuracy as v1 (qsa_prompt_attn_parity vs FP64: 3.7e-6 / 4.9e-6 / 4.5e-6 / 3.8e-6 of 3.5-3.7). RTX 2080 Ti, ms per chunk v1 -> v2: 32K x 1024 queries 13.0 -> 7.6, 64K x 512 6.7 -> 4.8, 1.5K 6.2 -> 5.0, 2.1K x 256 3.2 -> 2.3. Opt-in until the engine A/B. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
On Turing the v2 int8 prompt attention (STRATA_PA_V2T=1) stages the next chunk in registers, so its second shared-memory stage bought nothing - and at 45.6 KB a block only one fit an SM (64 KB): four warps on four schedulers, every latency exposed (ncu: 12.5 % warps active, tensor pipe 8 %). Smem2T: one K/V stage, and the q.k partials of cell group nt (8 cells) written over the 512 bytes of the warp's own K slice that held those cells (read by then, a __syncwarp between) instead of their own 8.4 KB. 19.7 KB a block: three per SM (168 registers fit three). The same arithmetic in the same order - bitwise the same output. RTX 2080 Ti, a recorded 130K selection (QSA layer 5, last chunk: 15871 queries x 2051 cells, tests/cuda/pa_bench.cpp): 113.8 -> 68.5 ms, same output hash. Engine (Q2_0 pack, 16K chunks, fused experts): 98K prefill 53.3 -> 49.9 s, 4K 3.54 -> 3.42 s; first logits bitwise identical at 98K and 4K, the same generated text. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…A_QSA_DUMP_LAYER / _STEPS STRATA_QSA_DUMP_LAYER=i keeps one QSA layer, STRATA_QSA_DUMP_STEPS=1 appends the steps to each record. pa_bench: the prompt sparse attention (qsa_prompt_attn_batch) on a recorded selection - a STRATA_QSA_DUMP record of one layer and chunk with its steps appended - random q and int8 K/V; median time and an output hash for A/B of kernel variants (PA_BENCH_OUT=file writes the output). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Niko1221
pushed a commit
that referenced
this pull request
Oct 17, 2026
…itwise) On Turing the v2 int8 prompt attention (STRATA_PA_V2T=1) stages the next chunk in registers, so its second shared-memory stage bought nothing - and at 45.6 KB a block only one fit an SM (64 KB): four warps on four schedulers, every latency exposed (ncu: 12.5 % warps active, tensor pipe 8 %). Smem2T: one K/V stage, and the q.k partials of cell group nt (8 cells) written over the 512 bytes of the warp's own K slice that held those cells (read by then, a __syncwarp between) instead of their own 8.4 KB. 19.7 KB a block: three per SM (168 registers fit three). The same arithmetic in the same order - bitwise the same output. RTX 2080 Ti, a recorded 130K selection (QSA layer 5, last chunk: 15871 queries x 2051 cells, tests/cuda/pa_bench.cpp): 113.8 -> 68.5 ms, same output hash. Engine (Q2_0 pack, 16K chunks, fused experts): 98K prefill 53.3 -> 49.9 s, 4K 3.54 -> 3.42 s; first logits bitwise identical at 98K and 4K, the same generated text.
Niko1221
pushed a commit
that referenced
this pull request
Oct 17, 2026
… STRATA_QSA_DUMP_LAYER / _STEPS STRATA_QSA_DUMP_LAYER=i keeps one QSA layer, STRATA_QSA_DUMP_STEPS=1 appends the steps to each record. pa_bench: the prompt sparse attention (qsa_prompt_attn_batch) on a recorded selection - a STRATA_QSA_DUMP record of one layer and chunk with its steps appended - random q and int8 K/V; median time and an output hash for A/B of kernel variants (PA_BENCH_OUT=file writes the output).
Owner
|
Thanks rafatxf, good work. We ported this with your authorship kept. We have no Turing card, so we could not confirm the speed. What we did check: the default path is byte-identical to before (identity gate on RTX 3060, P100 and A4000), and with the new path forced on, including a build that runs the real sm_75 device code on an sm_86 card, the gate is identical too. Because we cannot measure it on sm_75 it ships as opt-in (STRATA_SM75_PERF=1, plus your own variable STRATA_PA_V2T=1 (numerics of #2092)), documented in docs/DETAILS.md. If a Turing owner confirms the gain we will consider turning it on by default. |
Contributor
Author
|
Ported in 0.1.45, thanks! Turing numbers for |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Stacked on #2092 (the Turing v2 prompt attention,
STRATA_PA_V2T=1): its first commit is #2092's.With register staging, #2092's v2 keeps the next chunk's K/V pieces in registers until the current chunk's p·v is done, so its second shared-memory stage buys nothing on Turing - but
Smem2is 45.6 KB and a Turing SM has 64 KB: one 4-warp block per SM, one warp per scheduler, every latency exposed. ncu on a 130K chunk: 12.5 % warps active, the tensor pipe at 8 %.Below sm_80 the kernel now takes
Smem2T: one K/V stage, and no separate q·k partials (part, 8.4 KB). The partials of cell groupnt(8 cells, 16 rows) are exactly 512 bytes, and the cells they come from occupy bytes[512 nt, 512 nt + 512)of the warp's own K slice (the swizzle stays inside a cell's 64 bytes), read by then - so after a__syncwarpeach group's partials overwrite its own cells. The softmax reads them across warps after the same__syncthreadsas before, and the next chunk'sstore_stagecomes after the second one, as before. 19.7 KB a block: three per SM (168 registers fit three). Same arithmetic, same order: the same bits.The second commit adds
tests/cuda/pa_bench.cpp, which replays a recorded selection (STRATA_QSA_DUMPgetsSTRATA_QSA_DUMP_LAYER=iandSTRATA_QSA_DUMP_STEPS=1), random q and int8 K/V, and prints the median time and an output hash.Setup
RTX 2080 Ti 22 GB (sm_75), i7-6700, CUDA 13.0, SM 1650 MHz, 240 W;
v0.1.43+ #2092. Q2_0,--prefill auto:16384,STRATA_PA_V2T=1.pa_bench(QSA layer 5, last chunk of a 130,559-token prompt: 15,871 queries × 2,051 cells)2541296d11c7be6f2541296d11c7be6fEngine
(with the fused experts of #2145 on; the generated text is the same.) Deep, served (
--serve, 512K context): 389,560 tokens from scratch 255.0 → 223.6 s and +16K at 406K 12.82 → 11.42 s together with the fused experts; needle 2/2 at 352,660 tokens.🤖 Generated with Claude Code