Pull requests / #2806
#2806 sycl: decode QSA selection and attention in fewer, shorter launches (opt-in, bitwise)
open · @hotschmoe · 0 comments · View on GitHub
BenchmarksSetup & installMulti-GPUAMD / HIPNVIDIA / CUDAModels & quantsDocumentationWindows
Description
## Title
Issue: Resolves #2805
## Summary
On an Arc Pro B70 the QSA selection (block scores, top-k) and attention of a decode window are small, latency-bound
launches. This adds four opt-in switches that do the same arithmetic in fewer launches or with less serial work. Each
gives exactly the bits of the path it replaces. All are off by default, and each is a branch in front of the existing
code that tests its flag first, so the default path is unchanged.
| switch | what it does |
| --- | --- |
| `STRATA_BATCH_QSA_ATTN=1` | A batch window (`--batch`) ran the scores, top-k and attention one row at a time, because each slot has its own K/V. Now the scores run once per slot, the top-k once per window, and the attention once per window over every slot's pools. |
| `STRATA_QSA_TOPK_SCAN=1` | The wide top-k kernel (a capacity above ~135K cells) found each radix pass's digit with one thread walking 256 bins. Now it is a work-group scan. Integers only: the same ids. |
| `STRATA_QSA_SCORES_TRANSPOSED=1` | The decode block scores reduce a block's (query, head) sums together: 16-31 shuffles per block instead of 20 per query, adding the same pairs in the same order. |
| `STRATA_ATTN_MERGE_ACTIVE=1` | The attention merge of a call of 8 rows or fewer folds only the chunks its selection reaches, with loads issued 8 at a time. The chunks past it hold no cells and were skipped anyway. The MTP drafter's 32K window is 512 chunks. |
Engine, main 296b59d2 vs this branch with all four on. Setup: 2x Arc Pro B70 32 GB, kernel 7.1, Compute Runtime
26.22.38646.4, oneAPI 2026.1, AOT bmg-g31. Model: Qwen3.8-Flash-Next ISTA GSQ-RCO IQ3_XXS. Flags:
`--layer-split 30 --prefill 8192 --kv int8 --spec 4 --mtp <BF16 draft layer> --mtp-max-t 4 --spec-min-p 0
--stream-experts --expert-cache auto`. Greedy, fixed prompts.
| | without | with the switches |
| --- | ---: | ---: |
| one stream, `--max-context 65536`, tok/s, 3 runs (without = main) | 81.8 / 79.5 / 78.7 | 82.2 / 80.0 / 79.1 (+0.6%) |
| one turn on a 40,796-token context, tok/s (without = main) | 76.0 | 77.1 (+1.4%) |
| two streams in one window (`"parallel": 2 --batch-groups 1 --max-context 16384 --vram-reserve-mib 512`), tok/s together, 3 runs (without = switches unset) | 56.78 / 57.84 / 58.23 | 56.41 / 58.99 / 59.14 (median +2%) |
The outputs were identical in every row: text, and for one stream the draft counts too, which match main. With the
switches unset the branch matched main, and prompt reading was unchanged. Coherence and per-card health passed after
every run.
Kernels (`qsa_verify_batch_parity --bench`, one B70, synthetic layers: 24 q heads, INT8 K/V, a 200,000-cell
capacity; graph replay, us; 1K / 16K / 64K of context):
| | off | on |
| --- | --- | --- |
| `STRATA_BATCH_QSA_ATTN`, 2 slots x 1 row, selection + attention per layer | 102.8 / 184.0 / 252.6 | 62.3 / 107.9 / 156.2 |
| `STRATA_BATCH_QSA_ATTN`, 2 slots x 2 rows | 183.6 / 365.0 / 498.2 | 78.3 / 138.2 / 205.0 |
| `STRATA_QSA_TOPK_SCAN`, one query (16K / 64K; at 1K every cell is selected) | 36.3 / 57.9 | 19.5 / 42.1 |
| `STRATA_QSA_SCORES_TRANSPOSED`, four queries | 5.5 / 18.7 / 63.1 | 4.6 / 14.1 / 43.2 |
| `STRATA_ATTN_MERGE_ACTIVE`, MTP drafter call (one row, 32K window) | 86.9 / 223.7 / 460.9 | 41.0 / 157.2 / 335.9 |
What it leaves alone:
- the default path;
- the prompt path: prefill calls take the active-count kernels, and the merge switch only acts on calls of 8 rows or
fewer;
- `qsa_decode_attn_step` (decode without MTP);
- a batch window's K/V and indexer appends, which stay per row;
- batch windows with a streamed K/V, which keep the per-row path because their blocks are made resident row by row;
- CUDA and HIP: no shared file is touched.
## What changed
- `sycl/src/core/verify.cpp`: adds `STRATA_BATCH_QSA_ATTN`. In `record_window`'s QSA branch, a new branch before the
per-row batch loop runs the scores per slot (a slot's rows are contiguous), one `qsa_block_topk` and one
`qsa_decode_attn_multi` for the window.
- `sycl/include/strata/qsa_decode_attn_multi.hpp` (new, SYCL only) and `sycl/src/kernels/cuda/qsa_decode_attn.dp.cpp`:
- `qsa_decode_attn_multi` is `qsa_decode_attn_batch` over up to four pool sets per launch. Row z's work-group picks
its set; the chunk and merge kernels are the existing ones. Mixed KV formats or `STRATA_ATTN_LANECELL=1` fall back
to one `qsa_decode_attn_batch` per set.
- `STRATA_ATTN_MERGE_ACTIVE`: a new `attn_merge_active_kernel`, and a check in front of `qsa_decode_attn_batch`'s
merge launch.
- `sycl/src/kernels/cuda/qsa_select.dp.cpp`:
- `STRATA_QSA_TOPK_SCAN`: `block_topk_wide_kernel` gets `template <bool SCAN = false>`. The scan is an
`if constexpr` branch that only the new launch instantiates; the existing launch compiles the same code.
- `STRATA_QSA_SCORES_TRANSPOSED`: a new `block_scores_multi_t_kernel<NQ>`, launched from `qsa_block_scores`'s
multi-query branch.
- `sycl/src/kernels/qsa_verify_batch_parity.cpp` (new) and `sycl/CMakeLists.txt` add two ctest entries:
- `qsa_verify_batch_parity`, without switches, checks in process that:
- the batch-window form equals the per-row calls (2x1, 2x2, 4x2, 6x1 and 1x4 windows at 1K / 16K / 64K);
- the selections equal `qsa_block_topk_ref`;
- the attention equals `qsa_decode_attn_step`;
- 1-8-query scores equal one call per query;
- the MTP drafter call equals `qsa_decode_attn_step`.
It also writes a hash of each of its 45 outputs.
- `qsa_verify_batch_parity_switches` repeats this with the three kernel switches set and requires the same hashes.
- `docs/INTEL.md`: a short section with the numbers above.
## Extra Notes
- Prepared with AI assistance (Claude Code): the code, the measurements and this text were produced by Claude on the submitter's two Arc Pro B70s, at the submitter's request.
- How the bitwise claim was checked:
- both ctest entries pass, and each switch alone passes against the same hashes;
- the engine outputs above are identical with the switches on, off and unset, and match main.
No tolerance is used anywhere.
- Built: SYCL (AOT bmg-g31), `strata` and `qsa_verify_batch_parity`. No CUDA or HIP file is touched.
- `STRATA_QSA_TOPK_SCAN` did not act in the engine runs. At `--max-context 65536` the wide top-k is not used.
`--max-context 200000` (where it is) does not fit every expert on a plain two-card layer split: 710 were refused.
Its effect is shown by the kernel bench only.
- On a layer split, main refuses `--batch-mtp`. So a two-stream window there is one row per slot, without MTP drafts.
The default `--batch-groups auto` runs one slot per stage, where `STRATA_BATCH_QSA_ATTN` changes nothing. With
`--batch-mtp` on one GPU a window is 2 rows per slot: the 2 x 2 kernel row.
- Not included (possible follow-ups), all measured on a development tree:
- a value-pipelined chunk kernel for decode-sized calls. It needs a kernel variant main does not have; it was 7-10%
faster on the MTP call.
- per-slot K/V and indexer appends in batch windows: 16 -> 9 us per layer at 2 x 2, nothing at one row per slot.
- the same digit scan in the register top-k kernel: not measured.
🤖 Generated with [Claude Code](https://claude.com/claude-code)
Related on this site
Quick links to related pages.