Pull requests / #2075

#2075 QSA decode attention: PR #540's two-cells-per-step chunk kernel on every CUDA card and a bounded merge (STRATA_ATTN_DAILY=1, opt-in, bit-exact)

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

BenchmarksAMD / HIPNVIDIA / CUDAModels & quantsWindows

本文

## Title
Issue: none (an opt-in speed change, bit-exact).

## Summary
`STRATA_ATTN_DAILY=1` (opt-in, off by default) runs the QSA layers' decode attention - the verify window's and the MTP
drafter's, single query and batch - through PR #540's chunk kernel (`attn_chunk_kernel_pre75`: two cells per step
with a 12-shuffle reduction, cell-major probabilities, V loaded four at a time) on **every** CUDA card, and merges the
chunks with `attn_merge_v2_kernel`, a merge that skips the chunks the chunk kernel marked empty. Upstream compiles the
pre75 kernel only with `STRATA_EXPERIMENTAL_SM60` and picks it below sm_75; here it is compiled in every CUDA build
and selected by the variable. The output is **bit for bit** the default kernels' (see below). HIP is untouched.

## Why
The default `attn_chunk_kernel` reduces each cell with 5 `warp_sum` shuffles per head (60 per cell pair); the pre75
body does two cells per step with `reduce12` (16 shuffles) and reads the V rows four at a time. nsys on an RTX 3060
(stage 0 of a layer split, 131K context): `attn_chunk_kernel` **99.8 us** a call against **46.4 us** with the pre75
body; the same ratio on an RTX 2080 Ti. The decode attention is ~8% of the per-row cost of a window at 131K on these
cards, so the end-to-end gain is small but free.

## What changed
- `src/kernels/cuda/qsa_decode_attn.cu`: `attn_chunk_kernel_pre75<M, NC, ...>` compiled in every CUDA build;
  `daily_attn()` reads `STRATA_ATTN_DAILY`; the kernel choice `STRATA_ATTN_CHUNK_SEL(M, PRE)` takes it on any card
  when set (below sm_75 in an experimental build as before). `attn_merge_v2_kernel` with the chunk-empty marks;
  `STRATA_ATTN_MERGE_V2=0|1` forces the merge alone, `STRATA_ATTN_PRE75=0` turns the chunk kernel off (A/B).
  `spt` shares `sq`'s shared memory, so the kernel uses 15,872 B like the default one (68-78 registers against 37-41,
  256 threads a block).
- `include/strata/kernels/qsa_decode_attn.hpp`: the selection helpers.
- `src/kernels/qsa_merge_daily_parity.cpp` (tests): memcmp of the chunk kernel and the merge against upstream's on
  four KV formats, widths at the chunk borders and non-resident pages; `--bench` times the two chunk kernels back to
  back. Registered as ctest `qsa_merge_daily_parity_default` and `_daily`.
- `CMakeLists.txt`: the harness.

## Bit for bit
`reduce12` adds at every level the same two pairs as `warp_sum`'s butterfly (xor 16, 8, 4, 2, 1; float addition
commutes), the V cells fold in the same cell order with the same `fmaf`, the masked cells stay excluded; the merge v2
does the same `fmaf` in the same order and skips exactly the chunks the chunk kernel marks empty. `qsa_merge_daily_parity` on an RTX 2080 Ti (v0.1.43 build, tests on): **1,168 cases, 0 differ** (fp16 290, int8 291, q4_0 294, k8v4 293; 21,190,308 nonzero output values), with `STRATA_ATTN_DAILY` unset and `=1`. In the engine the greedy text and the drafter's acceptance are unchanged (deterministic
131K bench, forced tokens and windows).

## Measurements
RTX 2080 Ti 22 GB (sm_75) + RTX 3060 12 GB (sm_86), Windows 11, CUDA 12.6, layer split, Qwen3.8-Flash-Next IQ3_XXS,
131K context, `--kv int8`.

`qsa_merge_daily_parity --bench` on the RTX 2080 Ti (CUDA events, 200 calls per arm, B-A-B: Daily kernel, upstream
kernel, Daily kernel again; int8 KV, 32,768-cell capacity, window of n_q queries):

| n_q | width (cells) | chunk kernel: upstream | PR #540 body | ratio | merge: upstream | merge v2 | ratio |
|---|---|---|---|---|---|---|---|
| 1 | 32,768 | 448.9 us | **264.3 us** | 0.59 | 182.6 us | **45.6 us** | 0.25 |
| 2 | 32,768 | 849.8 | **483.9** | 0.57 | 186.6 | **60.9** | 0.33 |
| 3 | 32,768 | 1,259.8 | **727.3** | 0.58 | 202.3 | **85.1** | 0.42 |
| 4 | 32,768 | 1,692.3 | **969.0** | 0.57 | 208.4 | **106.5** | 0.51 |
| 4 | 2,051 | 132.7 | **86.6** | 0.65 | 41.3 | **9.9** | 0.24 |

fp16, q4_0 and k8v4 KV show the same ratios (0.57-0.63 for the chunk kernel, 0.24-0.51 for the merge). In the engine
(nsys, RTX 3060, 131K context) `attn_chunk_kernel` went from 99.8 us to 46.4 us a call.

End to end it is part of a stack that matches its reference within noise; the attention alone is ~8% of the per-row
cost, so no separate end-to-end claim is made here.

Built and started on v0.1.43 with `STRATA_ATTN_DAILY=1` (RTX 2080 Ti + RTX 3060 layer split, Windows): the engine serves as before, same answer.

## Risks / limits
Opt-in: nothing changes with the variable unset. The kernel body is PR #540's (sskver) and the merge is Hardin22's
`attn_merge_v2_kernel`, both carried in this fork since 0.1.40 and credited in the source. #1888 (split attention on
the int8 mma path, sm_80+) is a different kernel and not bit-exact; the two do not overlap (this one serves Turing).

Part of a series from a public patch registry (`stack/` in the fork imanu86/Strata-2080Ti).

関連リンク

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