Pull requests / #1888
#1888 QSA decode attention: split each query's cells across blocks on the int8 mma path (STRATA_QSA_SPLIT_ATTN=1, opt-in)
open · draft · @KarlGrier · 0 Kommentare · Auf GitHub
BenchmarksSetup & installServer & APIAMD / HIPNVIDIA / CUDAModels & quantsWindowsLinux
Beschreibung
## Title Issue: none (an opt-in speed change) ## Summary `STRATA_QSA_SPLIT_ATTN=1` (opt-in, off by default) runs the QSA layers' decode attention - the verify window's and the draft layer's - through the prompt path's int8 tensor-core kernel, with each query's cells split among up to 20 blocks (about two blocks per SM in one wave) and the splits merged with the log-sum-exp rescale. It applies to int8 KV on sm_80 and newer; everywhere else, and with the switch unset, `qsa_decode_attn_batch` runs as before. The idea is from Eddoursul's fork (fd23d91); the kernel here is written against this tree's `qsa_prompt_attn.cu` and its helpers. It sums in another order than `qsa_decode_attn_batch`, so the output is not bit-identical with the switch on (a rounding-level change); with it off nothing changes. **Measured** 2026-10-10 on one PC: RTX 5080 16 GB + RTX 4060 Ti 8 GB (driver 610.57.04, CUDA 13.3), Ryzen 9 9950X3D (16 cores), 96 GB DDR5, Linux 7.0; Qwen3.8-Flash-Next GSQ-RCO IQ3_S at a 131,072-token context. Served, stock against stock + this patch only, A B B A with a fresh server per arm (the model files evicted from the page cache before each), both arms with `--expert-cache auto --prefill auto:16384 --spec 4 --spec-min-p 0.5 --mtp --max-context 131072 --kv int8 --kv-resident 32768 --pcie-frac 0 --ple-inflight 128 --vram-reserve-mib 713` with `STRATA_IO_THREADS=64 CUDA_MODULE_LOADING=LAZY STRATA_GR_V3=1 STRATA_PF_FUSED=1 STRATA_MMVQ_IL=1 STRATA_ADAPT_LAG=2 STRATA_DISJOINT_ADAPT=1` (the stock flags of my community benchmark report, #1681); the patch's arm adds `STRATA_QSA_SPLIT_ATTN=1`. Requests carry no sampling fields, so the server's configured defaults apply: temperature 1.0, top-p 0.95, top-k 20, min-p 0.05, reasoning effort high. Decode screen: long reasoning math and code (4,096 tokens each), five short questions, a 32K prompt with 768 tokens out; decode figures are the engine's own (accepted output tokens / decode time). Prompt-read screen: time to first token for prompts of 1K x3, 4K x3, 32K x2 and 120K x2 tokens (each a different text, 768 tokens out; its decode column comes from short sampled answers and varies more than the decode screen). In short (this patch against the same stock build, means of the two A B B A arms each): | setup | decode | ms per verify round | |---|---|---| | RTX 5080 alone, v0.1.42, upstream's defaults (tail skip on) | +4.2 % | -5.1 % | | RTX 5080 alone, v0.1.42, `STRATA_ROUTE_TAIL_SKIP=0` | +2.3 % | -3.6 % | | RTX 5080 + RTX 4060 Ti helper, v0.1.41 (before the rebase) | +2.4 % | -2.7 % | *The RTX 5080 alone, against v0.1.42 `61b3fb5d` with upstream's defaults: v0.1.42 turns on `STRATA_ROUTE_TAIL_SKIP=7` for single-GPU serve, so both arms skip low-ranked missed experts (what single-card users get).* | metric | v0.1.42: arm 1 / arm 4 -> mean | v0.1.42 + this patch: arm 2 / arm 3 -> mean | patch vs v0.1.42 | |---|---|---|---| | decode (tok/s) | 131.8 / 130.2 -> **131.0** | 137.6 / 135.5 -> **136.6** | +4.2 % | | ms per verify round | 15.31 / 15.06 -> **15.19** | 14.56 / 14.27 -> **14.41** | -5.1 % | | tokens per round | 2.018 / 1.961 -> **1.990** | 2.004 / 1.933 -> **1.969** | -1.1 % | | acceptance | 0.647 / 0.644 -> **0.645** | 0.655 / 0.642 -> **0.648** | +0.5 % | | long reasoning math (tok/s) | 149.0 / 147.4 -> **148.2** | 151.1 / 151.4 -> **151.2** | +2.0 % | | long reasoning code (tok/s) | 119.9 / 118.3 -> **119.1** | 128.6 / 124.7 -> **126.6** | +6.3 % | | TTFT 32K (s) | 7.72 / 7.74 -> **7.73** | 7.72 / 7.72 -> **7.72** | -0.1 % | | all requests, total time (s) | 89.4 / 91.6 -> **90.5** | 86.2 / 88.4 -> **87.3** | -3.5 % | | prompt size | v0.1.42: TTFT s / decode tok/s | v0.1.42 + this patch: TTFT s / decode tok/s | TTFT, patch vs v0.1.42 | decode, patch vs v0.1.42 | |---|---|---|---|---| | 1K (6 runs per side) | 0.82 / 124.4 | 0.81 / 127.2 | -0.7 % | +2.2 % | | 4K (6 runs per side) | 1.23 / 123.0 | 1.23 / 126.0 | -0.4 % | +2.4 % | | 32K (4 runs per side) | 7.66 / 127.0 | 7.66 / 128.7 | -0.0 % | +1.3 % | | 120K (4 runs per side) | 29.64 / 129.8 | 29.64 / 146.8 | -0.0 % | +13.1 % | *The RTX 5080 alone, against v0.1.42 with `STRATA_ROUTE_TAIL_SKIP=0` in both arms (the tail skip off: the answers of a build without it).* | metric | v0.1.42: arm 1 / arm 4 -> mean | v0.1.42 + this patch: arm 2 / arm 3 -> mean | patch vs v0.1.42 | |---|---|---|---| | decode (tok/s) | 115.3 / 114.5 -> **114.9** | 117.7 / 117.4 -> **117.5** | +2.3 % | | ms per verify round | 17.42 / 17.31 -> **17.37** | 16.72 / 16.78 -> **16.75** | -3.6 % | | tokens per round | 2.009 / 1.983 -> **1.996** | 1.968 / 1.970 -> **1.969** | -1.3 % | | acceptance | 0.659 / 0.648 -> **0.654** | 0.640 / 0.645 -> **0.642** | -1.7 % | | long reasoning math (tok/s) | 132.6 / 127.8 -> **130.2** | 136.4 / 134.2 -> **135.3** | +3.9 % | | long reasoning code (tok/s) | 107.2 / 110.4 -> **108.8** | 110.1 / 110.0 -> **110.0** | +1.2 % | | TTFT 32K (s) | 7.74 / 7.73 -> **7.74** | 7.73 / 7.74 -> **7.74** | -0.0 % | | all requests, total time (s) | 100.0 / 99.1 -> **99.6** | 100.2 / 100.3 -> **100.2** | +0.7 % | | prompt size | v0.1.42: TTFT s / decode tok/s | v0.1.42 + this patch: TTFT s / decode tok/s | TTFT, patch vs v0.1.42 | decode, patch vs v0.1.42 | |---|---|---|---|---| | 1K (6 runs per side) | 0.82 / 107.7 | 0.82 / 107.3 | -0.3 % | -0.4 % | | 4K (6 runs per side) | 1.24 / 104.3 | 1.24 / 107.8 | -0.2 % | +3.4 % | | 32K (4 runs per side) | 7.67 / 101.3 | 7.67 / 104.8 | -0.1 % | +3.4 % | | 120K (4 runs per side) | 29.65 / 113.9 | 29.64 / 114.5 | -0.0 % | +0.5 % | *The RTX 5080 with the RTX 4060 Ti as the helper expert cache (`--expert-cache-device1 3300`), measured on v0.1.41 `fb58e0db` with this change on v0.1.41, before the rebase. The tail skip does not apply with a helper GPU.* | metric | v0.1.41: arm 1 / arm 4 -> mean | v0.1.41 + this patch: arm 2 / arm 3 -> mean | patch vs v0.1.41 | |---|---|---|---| | decode (tok/s) | 125.0 / 126.5 -> **125.7** | 128.6 / 128.9 -> **128.8** | +2.4 % | | ms per verify round | 16.09 / 15.85 -> **15.97** | 15.30 / 15.77 -> **15.54** | -2.7 % | | tokens per round | 2.011 / 2.004 -> **2.008** | 1.969 / 2.033 -> **2.001** | -0.3 % | | acceptance | 0.660 / 0.665 -> **0.663** | 0.642 / 0.644 -> **0.643** | -3.0 % | | long reasoning math (tok/s) | 140.1 / 140.1 -> **140.1** | 139.6 / 144.7 -> **142.2** | +1.5 % | | long reasoning code (tok/s) | 116.1 / 114.7 -> **115.4** | 121.9 / 118.0 -> **120.0** | +4.0 % | | TTFT 32K (s) | 7.72 / 7.74 -> **7.73** | 7.72 / 7.73 -> **7.73** | -0.1 % | | all requests, total time (s) | 92.9 / 94.8 -> **93.9** | 92.8 / 88.0 -> **90.4** | -3.7 % | | prompt size | v0.1.41: TTFT s / decode tok/s | v0.1.41 + this patch: TTFT s / decode tok/s | TTFT, patch vs v0.1.41 | decode, patch vs v0.1.41 | |---|---|---|---|---| | 1K (6 runs per side) | 0.80 / 124.1 | 0.81 / 128.2 | +0.3 % | +3.3 % | | 4K (6 runs per side) | 1.23 / 124.8 | 1.22 / 121.2 | -0.6 % | -2.9 % | | 32K (4 runs per side) | 7.67 / 119.1 | 7.67 / 125.1 | +0.0 % | +5.1 % | | 120K (4 runs per side) | 29.69 / 132.9 | 29.68 / 134.0 | -0.0 % | +0.8 % | ## What changed - `src/kernels/cuda/qsa_prompt_attn.cu`: `decode_attn_i8_split_kernel` (the int8 q.k / p.v of one split of one query's cells, the same fragments and cp.async pipeline as the prompt kernel, partials in the decode scratch) and `decode_attn_merge_kernel` (the splits of each query and head merged in split order); `qsa_decode_attn_split()` picks the split count and returns false (nothing launched) unless int8 KV on sm_80+, or when its partials would not fit the caller's scratch. HIP: returns false before any launch (the kernels use the file's `mma16816` / `cp_async16`, which compile to traps there, like the existing prompt kernels). - `include/strata/kernels/qsa_prompt_attn.hpp`: its declaration. - `src/core/verify.cpp`, `src/core/mtp.cpp`: with the switch set, the verify window's and the draft layer's QSA attention try `qsa_decode_attn_split()` first and keep `qsa_decode_attn_batch()` when it returns false. Left alone: the prompt path, KV formats other than int8 (K8V4, q4), the GDN layers, the switch-off path (no new work, the same launches). ## Extra Notes **How the default path was checked.** exact7 (the CLI, seven fixed computations: `--adapt-swaps 0 --suffix-draft 0 --expert-cache 3500`) on v0.1.42 `61b3fb5d` and on this branch, both cards, the measurements' flags (above) plus `STRATA_PREFILL_CPU_SHARE=0` on both (with the CPU prompt-read share on, short-prompt outputs depend on timing and the stock build does not repeat itself; with it off, stock v0.1.42 run twice gave the same seven hashes; the CLI does not use v0.1.42's tail skip, whose default is set only at serve start-up). With the switch unset: **7/7 identical to v0.1.42**. With `STRATA_QSA_SPLIT_ATTN=1`: 3/7 identical; the other 4 differ because the attention sums in another order (a rounding-level change; the scoring check below). | exact7 item | v0.1.42 `61b3fb5d` | this branch, switch unset | switch on | |---|---|---|---| | math, greedy, 1,024 tokens | `02deb3f1` | `02deb3f1` (=) | `2dbacf64` | | math, sampled (seed 11), 1,024 tokens | `6c9b2d10` | `6c9b2d10` (=) | `0cd91a5c` | | code, greedy, 1,024 tokens | `ae3b089b` | `ae3b089b` (=) | `2d7b9115` | | code, sampled (seed 11), 1,024 tokens | `38e671a2` | `38e671a2` (=) | `c917d6d8` | | 2K-token prompt + 64 tokens | `3e219376` | `3e219376` (=) | `3e219376` (=) | | 32K-token prompt + 64 tokens | `651f56c8` | `651f56c8` (=) | `651f56c8` (=) | | 112K-token prompt + 64 tokens | `a393ca0a` | `a393ca0a` (=) | `a393ca0a` (=) | **The patch's own checks.** - Teacher-forced scoring (the same fixed token sequences scored through the verify window with the switch on and with `STRATA_QSA_SPLIT_ATTN=0`, on a local build of this branch with a scoring hook that is not part of the change): mean NLL on minus off, math -0.0005 nats, a text corpus -0.0016 nats. My bound for a rounding-level change is 0.005. - Needle in a haystack (`tools/needle_bench.py`) with the switch on, 32K and 128K prompts at depths 10, 50 and 90 %: 6/6 found. **Upstream's tests** (ctest, the whole suite, among them verify_parity, verify_batch_parity, spec_verify_parity, qsa_parity and spec_prob_test): v0.1.42 109/112 pass, this branch 109/112 pass; the same 3 fail on both (expert_parity, ple_parity, pool_test): they read model files from fixed paths that this PC does not have. Not tested: - **HIP and SYCL builds.** This PC has no ROCm or oneAPI, and installing them would be a system change, so it is a draft. On HIP `qsa_decode_attn_split()` returns false before any launch (`#if defined(__HIPCC__)`, like the file's prompt kernels), but the file was not compiled for HIP here. SYCL is its own tree and is not touched. - Other GPUs: only the RTX 5080 (sm_120) ran the kernel. The RTX 4060 Ti holds experts and runs no attention. - KV formats other than int8 (K8V4, q4): these keep `qsa_decode_attn_batch()` and were not run. - Several requests at once (`--batch-groups`), Windows, and other models. Measured and drafted with Claude Code.
Mehr auf der Site
Links zu Install, Modellen, Releases.