Pull requests / #2220

#2220 sm_75: QSA decode attention for INT8 KV off the shared-memory pipe (bitwise)

open · @rafatxf · 0 Kommentare · Auf GitHub

Setup & installServer & APINVIDIA / CUDAModels & quants

Beschreibung

The decode attention over the QSA selection (`attn_chunk_kernel<1>`, INT8 pools) saturates the SM's shared-memory pipe on Turing. Its time grew linearly with the queries while the SMs were half empty: per cell, every warp reads two 16-byte slices of each of the 12 query heads from shared memory (a two-way bank conflict at `lane * 8`) and runs 12 `warp_sum`s, 60 shuffles; the value loop then reads the 12 probabilities of a cell as 12 scalar loads. Loading the K/V codes earlier (software prefetch) changed nothing, so it is not memory latency.

`attn_chunk_q8_v3_kernel` does the same operations in the same order:

- the lane's slice of the 12 heads (8 dimensions each) is read from shared memory once, into registers;
- the 12 lane sums are reduced with `reduce12_gfx906` (16 shuffles; every sum is `warp_sum`'s pairs, as its comment says);
- the probabilities are stored cell-major, so the value loop reads a cell's 12 as three 16-byte broadcasts.

The chunk partials and the merged output are bit for bit the default kernel's. It is used by both entry points (`qsa_decode_attn_batch`, `qsa_decode_attn_step`), so the MTP drafter's attention gets it too. Default on sm_75 only; `STRATA_ATTN_V3=0` keeps the default kernel, `=1` takes this one on any CUDA card. Lane-cell scoring (`STRATA_ATTN_LANECELL=1`) keeps its path.

## Setup

RTX 2080 Ti 22 GB (sm_75), i7-6700, CUDA 13.0, SM 1650 MHz, 240 W. Flash-Next, canonical Q2_0 pack, `--kv int8`, `--max-context 524288`, `--serve`.

## Kernel (chunk kernel + merge, synthetic INT8 pools, permuted page table, 200 calls)

| selected cells | queries | default | v3 | |
|---:|---:|---:|---:|---:|
| 2,051 | 1 | 51.6 µs | 34.5 µs | 1.50× |
| 2,051 | 2 | 78.5 µs | 41.0 µs | 1.91× |
| 2,051 | 3 | 110.0 µs | 73.0 µs | 1.51× |
| 2,051 | 4 | 145.6 µs | 84.9 µs | 1.71× |
| 600 | 1 | 45.9 µs | 28.6 µs | 1.60× |
| 600 | 2 | 74.6 µs | 32.8 µs | 2.27× |
| 600 | 3 | 105.5 µs | 33.5 µs | 3.15× |
| 600 | 4 | 136.9 µs | 33.3 µs | 4.11× |

Partials and outputs compared with `memcmp` against the default kernel in every case: identical. With the swizzle and `reduce12` alone (the gfx906 variants, also bitwise on CUDA) the 2,051-cell / 3-query case is 78 µs; the register copy of the query is what takes it to 73 µs, and further at fewer cells. At 2,051 cells and 3-4 queries the registers allow two blocks per SM (198-264 blocks: two waves); forcing three spills.

`qsa_prompt_attn_parity` reports the same errors with `STRATA_ATTN_V3=0` and `1` (its reference is this decode kernel: int8 32K / 2,048 queries 71.5 → 36.0 ms per chunk).

## Engine

Together with the Q2_0 matvec change of #2221 (both bitwise, measured as one step), 512 new tokens, two runs each: Spanish chat 86.0 → 87.8 t/s, code 103.1 → 106.1, English at 32K 75.3 → 78.4.

Part of the Turing set tracked in #2197.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Mehr auf der Site

Links zu Install, Modellen, Releases.