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).
関連リンク
インストール・モデル・リリースへの站内リンク。