QSA decode attention: split each query's cells across blocks on the int8 mma path (STRATA_QSA_SPLIT_ATTN=1, opt-in) by KarlGrier · Pull Request #1888 · Niko1221/Strata

KarlGrier · 2026-10-10T20:17:10Z

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.

…nt8 mma path (STRATA_QSA_SPLIT_ATTN=1, opt-in) With STRATA_QSA_SPLIT_ATTN=1 the verify window's and the draft layer's QSA attention run through the prompt path's int8 tensor-core kernel with each query's cells split among up to 20 blocks, and the splits are merged with the log-sum-exp rescale. qsa_decode_attn_split() returns false (nothing launched) unless the KV is int8 on sm_80 or newer, or when its partials would not fit the caller's scratch, and on HIP; the caller then keeps qsa_decode_attn_batch(). The sums run in another order, so with the switch on the output is a rounding-level change. Unset, nothing changes. The idea is from Eddoursul's fork (fd23d91); the kernel is written against this tree's qsa_prompt_attn.cu.

…y's cells across blocks on the int8 mma path (STRATA_QSA_SPLIT_ATTN=1, opt-in) With STRATA_QSA_SPLIT_ATTN=1 the verify window's and the draft layer's QSA attention run through the prompt path's int8 tensor-core kernel with each query's cells split among up to 20 blocks, and the splits are merged with the log-sum-exp rescale. qsa_decode_attn_split() returns false (nothing launched) unless the KV is int8 on sm_80 or newer, or when its partials would not fit the caller's scratch, and on HIP; the caller then keeps qsa_decode_attn_batch(). The sums run in another order, so with the switch on the output is a rounding-level change. Unset, nothing changes. The idea is from Eddoursul's fork (fd23d91); the kernel is written against this tree's qsa_prompt_attn.cu. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_0157fjixQD2nRxvCvtXq66DD

…1888 #1889 #1418 #1370 #1367 #1262 #1368 #1107), measured on sm_86; STRATA_DEQUANT_WIDE default on sm_86

Niko1221 · 2026-10-17T07:13:39Z

Thanks, this is a good one. Ported as an opt-in (STRATA_QSA_SPLIT_ATTN=1). On an RTX 3060 and an RTX A4000, decode after a 30K prompt is +2.8% and +3.5%; teacher-forced KL against the default is 0.0016 where the default against itself is 0.0012, and 3 of 3 needles at 32K were found. Greedy text differs (as you say, another summation order), so it stays opt-in.

Upstream 0.1.45 took in ten of the PRs we carried: Niko1221#567 (prompt encoder), Niko1221#1164 (sibling borrow), Niko1221#1166, Niko1221#1368, Niko1221#1367, Niko1221#1120 (split overlap), Niko1221#1912, Niko1221#1888, Niko1221#1949, Niko1221#1305. Their versions are kept and our copies dropped: verify.cpp's duplicate prelaunch/handoff block, verify_kernels.hpp, the qsa prompt kernels, gr_upmix_parity, the tokenizer and its tests (tools/test_strata_tokenizer_upstream.py is gone: upstream's test file is the one now), the second quantize_scatter declaration. README credits updated. Pool parts kept on top: - conversation cache: upstream renamed the borrow guard (pin -> guard); clear() resets it. - prefill: the remote_next hand-off beside upstream's hand-off pieces/worker (a pool stage has no pieces: the whole chunk goes to the workers). - generate.cpp: --pool-affinity measure over a node's own layers; the node's dense-range load beside STRATA_HC_REQ8; the carve on primary_device over the pool range; checkpoints keep pool_id (STRATA_CKPT_POOL/THREADS' parallel save is skipped on a pool PC); the workers' retain after eviction and after a resume; the borrow/restore of parked conversations on the workers with upstream's --conversation-cache-park-gain; Niko1221#652's emit count in upstream's merged windows. - expert_source: our budget retry beside upstream's lend budget. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_0157fjixQD2nRxvCvtXq66DD

Niko1221 · 2026-10-19T16:58:12Z

Thanks for the pull request.

Your change is already in main for 0.2.0-beta, with your credit, so there is nothing left for you to do here. We're closing it because we're resetting the tracker for 0.2.0-beta: there are more than 1,300 open issues and pull requests, and we're closing everything that isn't being worked on right now.

If part of your work did not make it in and you want it considered, open a new pull request on top of main that follows CONTRIBUTING.md, or comment /reopen after fixing the title area prefix and the template.

imanu86 mentioned this pull request Oct 12, 2026

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) #2075

Closed

yi-sudo mentioned this pull request Oct 13, 2026

Batch decode: one launch for all rows in QSA attention and GDN recurrence (opt-in) #2234

Closed

Niko1221 added a commit that referenced this pull request Oct 17, 2026

docs: the CUDA opt-in switches ported for 0.1.45 (#2058 #2059 #2075 #…

abce866

…1888 #1889 #1418 #1370 #1367 #1262 #1368 #1107), measured on sm_86; STRATA_DEQUANT_WIDE default on sm_86

Niko1221 closed this Oct 19, 2026

Navigation Menu

Repository navigation

QSA decode attention: split each query's cells across blocks on the int8 mma path (STRATA_QSA_SPLIT_ATTN=1, opt-in) - #1888

QSA decode attention: split each query's cells across blocks on the int8 mma path (STRATA_QSA_SPLIT_ATTN=1, opt-in)#1888
KarlGrier wants to merge 1 commit into
Niko1221:mainfrom
KarlGrier:qsa-split-decode-attn

KarlGrier commented Oct 10, 2026

Uh oh!

Niko1221 commented Oct 17, 2026

Uh oh!

Niko1221 commented Oct 19, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

KarlGrier commented Oct 10, 2026

Title

Summary

What changed

Extra Notes

Uh oh!

Niko1221 commented Oct 17, 2026

Uh oh!

Niko1221 commented Oct 19, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants