Pull requests / #2234

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

open · @yi-sudo · 0 comentários · No GitHub

BenchmarksMulti-GPUAMD / HIPNVIDIA / CUDAModels & quantsDocumentationWindowsLinux

Descrição

## Title
Issue: none (an opt-in speed change)

## Summary
With `--batch N`, the verify window ran QSA block scoring, top-k and decode attention as one launch per row, and the
GDN recurrence (in the window and in the 1-row commit) as one launch per slot. With 4-8 live slots, these serial
launches are most of what a window costs beyond the expert reads. Two opt-in switches run each of them as one launch
for all rows:

- `STRATA_BATCH_ATTN_ROWS=1`: block scores, top-k and decode attention for all of the window's rows (up to 8), one
  launch each. Every row still reads its own slot's indexer keys and K/V pools. These are passed by value as kernel
  params, so the graph captured per row layout stays valid. A slot that streams KV (`kv_mode` 1) keeps the per-row
  path.
- `STRATA_BATCH_GDN_GROUPS=1`: the GDN recurrence in the window, and the commit of 1-row groups, for all slots in one
  launch (`grid.y` = slot group). The conv commits stay per slot.

**Default path:** with both switches unset, the engine runs the same launches as before. Only kernel params with a
`{}` default were added, and they are unused when `n = 0`.

**Exactness with the switches on:** the arithmetic per row is the same. I sent 4 concurrent temperature-0 prompts
(`parity.py`) and the patched build produced byte-identical output to upstream `a0116b2`, with both `--batch 4` and
`--batch 8`.

**Measured** 2026-10-13, same day for both arms, on one RTX PRO 6000 Blackwell Workstation (96 GB, driver
610.57.04, CUDA 13.0, Linux). Settings: IQ3_S with all 24,576 experts in VRAM, `--kv int8`, 262144 context,
`--mtp --spec 4 --spec-min-p 0.5`, `--vision`. The load went through the HTTP server: C different story prompts sent
at once, 1024 tokens each, temperature 0.7, mean of 2 rounds. Total tok/s:

| Concurrent | `--batch 4` | + both switches | `--batch 8` (reserve 6144 MiB) | + both switches |
| ---: | ---: | ---: | ---: | ---: |
| 1 | 179.2 | 179.4 | 186.6 | 179.7 |
| 2 | 204.1 | 218.7 | | |
| 3 | 243.2 | 276.6 | | |
| 4 | 268.4 | 311.7 | 267.7 | 311.6 |
| 6 | | | 297.7 | 354.8 |
| 8 | | | 308.4 | 374.8 |

The window time, from `strata batch:` in the engine log:
- 4 rows: 14.3 ms → 12.2 ms
- 8 rows: 24.6 ms → 19.9 ms

That is about 1.7-1.9 ms per added row instead of 2.4 ms. The same patch on 0.1.43 (`7695859`) gave 268.5 → 312.9
(`--batch 4`, 4 clients) and 308.5 → 377.8 (`--batch 8`, 8 clients).

Scripts, raw JSON and engine logs: https://github.com/yi-sudo/strata-batch-boost (`lab/`, `lab/results/latest-a0116b2`).

## What changed
- `src/kernels/cuda/qsa_decode_attn.cu`, `include/strata/kernels/qsa_decode_attn.hpp`:
  - `attn_chunk_kernel` takes an optional `QsaAttnPoolsRows`; query `blockIdx.z` uses its own pools.
  - New `qsa_decode_attn_batch_rows`.
- `src/kernels/cuda/qsa_select.cu`, `include/strata/kernels/qsa_select.hpp`:
  - `block_scores_multi_kernel` takes an optional `QsaScoreRows` (per-query pooled/dead keys).
  - New `qsa_block_scores_rows`.
- `src/kernels/cuda/verify_kernels.cu`, `include/strata/kernels/verify_kernels.hpp`:
  - `gdn_step_norm_multi_kernel` takes an optional `GdnGroups` (per-group state, keep, first row, token count).
  - New `gdn_step_norm_groups`, for the window and for the commit.
- `src/core/verify.cpp`:
  - `record_window` uses the row launches in the `batch_rec_` QSA and GDN branches when the switches are set.
  - `capture_commit_batch` collects 1-row groups into one GDN commit launch.
- `docs/BATCHING.md`: a short section with the switches and the table above.

What this leaves alone: the solo path, `--batch-mtp`, layer splits (not tested), KV streaming (falls back), the conv
commits, the indexer appends, and every default.

## Extra Notes
- **Backends:** built and run with CUDA (sm_120) only. The `.cu` files are also part of the HIP build, which I could
  not build here (no ROCm). The SYCL port has its own copies of these files and is untouched.
- **Not tested:** a layer split / multi-GPU, experts on PCIe or the CPU, and AMD.
- **Related PRs:** #1524 also targets concurrent throughput, but by a different route (wider speculative batches,
  shared experts). #1888 changes the QSA decode attention kernel for the single-query case. This PR is independent of
  both and only merges launches.
- **Also tried, not included:** more MTP drafts per slot with `--batch-mtp`. Drafts are accepted 50-57% of the time
  in batch windows, but a row only pays above ~67%, so it was slower (raw data in the repo above).

No site

Links install, modelos, releases.