Pull requests / #2072

#2072 sm_75: the decode QSA selection at long context, a halving scorer and a split top-k (bitwise, decode +6-12 % at 410K)

open · @rafatxf · 0 comentários · No GitHub

BenchmarksSetup & installServer & APINVIDIA / CUDAModels & quants

Descrição

At long context a verify window's QSA selection is a fixed cost per layer that the rest of the window does not hide: at ~420K cells on an RTX 2080 Ti the 12 layers' scores and top-k took ~6 ms of a ~52 ms window (nsys, `--serve`, 256 tokens after a 422K conversation). Both kernels here produce the same bits as today's, and both have a switch.

1. **`block_scores_halve_kernel`** (decode scorer, `block_scores_multi_kernel`'s replacement on sm_75). Its time grew with the window's queries: 117 / 265 / 510 µs for 1 / 3 / 6 queries at 420K. The keys are read once, so the cost was the 5 xor shuffles per (query, head) value - 20 per block and query - and, once those were gone, the per-block shared-memory reads of the queries (ncu: `mio_throttle`). Here the 4 × queries values (padded to a power of 2) are reduced together by recursive halving: at offset o each lane keeps half its values and adds its partner's copy of them. That is the pair (lane l, lane l ^ o) the butterfly adds for that value at that offset, with the same offsets in the same order, so every value's additions are the butterfly's operand for operand (float addition commutes). The lane's query values live in registers; 5-8 queries run as two calls of up to 4. `STRATA_SCORES_HALVE=0` keeps the old kernel.
2. **`qsa_block_topk_split`** (decode top-k without clusters). Sm_90+ has `block_topk_cluster_kernel`; elsewhere a window's 1-8 queries ran `block_topk_wide_kernel` on 1-8 SMs. Here a query's blocks are cut into chunks of 2,048, and six short launches do four radix passes (each chunk adds its histogram to the query's global one), a count of each chunk's cells above / at the threshold, and the emit. Each pass's CTAs take the previous digit from the complete histogram with `block_topk_kernel`'s rule (a pure function of integer counts), so the threshold, `above`, the tie budget and the emitted cells are the same: identical ids by construction. `STRATA_TOPK_SPLIT=0` keeps the wide kernel; `=1` turns it on for any card (tests).

Both default on sm_75 only (measured there); nothing changes for other cards.

## Setup

RTX 2080 Ti 22 GB (sm_75, PCIe 3.0 x16), i7-6700, 56 GB DDR4-2133, CUDA 13.0, driver 595, SM clock locked at 1650 MHz, 240 W. `v0.1.43` built from source with `-DSTRATA_BUILD_TESTS=ON`. Model Q2_0, `--resident-experts`, `--spec 4 --spec-min-p 0.7`, `--max-context 524288` (YaRN ×2), `--kv int8 --kv-resident 32768`.

## Scorer: `QSA_BENCH_DECODE=1 qsa_select_bench CTX Q 50 524288` (new mode: a decode call, hash of the scores)

| context | queries | before | after |
|---|---:|---:|---:|
| 32K | 1 / 3 / 4 / 8 | 12.1 / 24.8 / 31.7 / 59.2 µs | 7.7 / 14.5 / 15.0 / 36.0 |
| 131K | 1 / 3 / 4 / 8 | 41.1 / 86.8 / 113.6 / 217.0 | 34.2 / 66.2 / 67.1 / 93.3 |
| 420K | 1 / 3 / 4 / 6 / 8 | 116.7 / 264.9 / 346.6 / 509.7 / 773.9 | 100.5 / 117.9 / 124.5 / 238.0 / 255.1 |

The same hash before and after in every case (1-8 queries, 32K-524K). At 3-4 queries and 420K it now reads the keys at ~390 GB/s.

## Top-k: `qsa_topk_parity CTX Q 50 524288 0` (a decode call), dispatched time

| context | 1 query | 3 | 6 |
|---|---:|---:|---:|
| 9K | 0.034 → 0.032 ms | 0.034 → 0.035 | 0.033 → 0.042 |
| 131K | 0.095 → 0.033 | 0.094 → 0.039 | 0.095 → 0.048 |
| 420K | 0.203 → 0.033 | 0.209 → 0.040 | 0.210 → 0.064 |
| 524K | 0.243 → 0.034 | 0.248 → 0.041 | 0.251 → 0.066 |

`qsa_topk_parity --selftest`: PASS (continuous, 64-level and equal scores; 2K-524K under 9K-524K capacities; decode and prompt calls).

## Engine: decode after a long conversation

`--serve`: one 390K request, then two appends of 16K tokens, 256 tokens generated after each (resumed from the conversation, so the decode sits at ~406K and ~422K):

| depth | before | after |
|---|---:|---:|
| ~406K | 52.7 tok/s | 56.0 |
| ~422K | 53.0 | 59.3 |

A one-shot 4K run writes the same tokens before and after (same MD5) at the same speed. The server here runs it since 2026-10-12.

## Tests

`qsa_parity`: 0 failures. `qsa_topk_parity --selftest`: PASS. `qsa_select_bench` with `QSA_BENCH_DECODE=1`: identical hashes with `STRATA_SCORES_HALVE` 0 and 1.

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

No site

Links install, modelos, releases.