Pull requests / #2146

#2146 sm_75: the v2 prompt attention in one stage, three blocks per SM (bitwise, attention 1.66x; stacked on #2092)

open · @rafatxf · 0 コメント · GitHub で見る

Setup & installServer & APINVIDIA / CUDAModels & quants

本文

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 `Smem2` is 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 group `nt` (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 `__syncwarp` each group's partials overwrite its own cells. The softmax reads them across warps after the same `__syncthreads` as before, and the next chunk's `store_stage` comes 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_DUMP` gets `STRATA_QSA_DUMP_LAYER=i` and `STRATA_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)

| | ms | output hash |
|---|---:|---|
| #2092 | 113.8 | `2541296d11c7be6f` |
| this | 68.5 | `2541296d11c7be6f` |

## Engine

| prompt | before | after | first logits |
|---|---:|---:|---|
| 98K (6 chunks) | 53.3 s | 49.9 s | bitwise identical |
| 4K | 3.54 s | 3.42 s | bitwise identical |

(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](https://claude.com/claude-code)

関連リンク

インストール・モデル・リリースへの站内リンク。