Pull requests / #2514

#2514 HIP: decode QSA top-k on the register kernel at 262K capacity

open · @rindi-ismail · 0 comentários · No GitHub

BenchmarksAMD / HIPNVIDIA / CUDAModels & quantsDocumentationWindows

Descrição

## Title
HIP: decode QSA top-k on the register kernel at 262K capacity

## Summary
On HIP with `--max-context 262144`, every decode window runs the QSA top-k on the one-workgroup reference kernel. The `TK_PER_MAX` register kernel exists on AMD for exactly this capacity (its comment says so), but only the prompt path reached it. This sends decode there too.

Measured with the repo's own tests (`qsa_topk_parity`):
- **Same ids:** identical selections on gfx1201 and gfx1100.
- **Kernel time:** 2.1-2.4x faster at 98K-262K on the R9700 and 1.7-1.8x on the 7900 XTX.

End to end, on an R9700 + 7900 XTX split, decode at 157K context goes from 66.8 / 64.7 to 69.8 / 68.6 tok/s (two pairs), and short context is unchanged.

## What changed
`src/kernels/cuda/qsa_select.cu`, `qsa_block_topk`, HIP only (9 lines):
- **Old behaviour:** a call without a block count (decode; its graphs are captured) kept `fit = TK_T * TK_PER` = 33,792 blocks. With a 262,144 capacity (`max_blocks` = 65,538), every decode call went to `qsa_block_topk_ref`. That kernel is one 256-thread workgroup per query, rereading the keys on each of its four radix passes, so a decode window with 2-6 queries ran its selection on 2-6 compute units.
- **New behaviour:** on HIP, such a call uses `TK_T * TK_PER_MAX` (66 per thread, 67,584 blocks), so decode runs `block_topk_reg_kernel<TK_PER_MAX>`. That kernel reads only the live `n_bid + 1` blocks and returns the same ids.
- **Escape hatch:** `STRATA_TOPK_DECODE_REG=0` restores the old dispatch. It is read once, like the other switches here.

Left alone:
- the CUDA dispatch (the change is inside `#if defined(__HIPCC__)`; the `#else` branch is the original line)
- the prompt path's `counted` rule and `kRegMinBlocks`
- both kernels and the cluster path

## How I checked
**Ids.** `src/kernels/qsa_topk_parity.cpp` was built against this change and run on each card with `HIP_VISIBLE_DEVICES`:
- `--selftest`: PASS on both.
- Decode-shaped calls (4 queries, `COUNT=0`, capacity 262,144) at 2,500 / 4,096 / 8,192 / 16,384 / 24,576 / 32,768 / 98,304 / 160,965 / 262,144: PASS, with identical ids for continuous, 64-level and all-equal scores.

**Kernel time** (the parity test's `dispatched` time, 4 queries, before -> after; 30-50 reps):

| Context | R9700 (gfx1201) | RX 7900 XTX (gfx1100) |
|---|---|---|
| 2,500 | 0.038 -> 0.042 ms | 0.036 -> 0.039 ms |
| 8,192 | 0.041 -> 0.040 ms | 0.038 -> 0.038 ms |
| 24,576 | 0.059 -> 0.042 ms | 0.050 -> 0.039 ms |
| 32,768 | 0.063 -> 0.040 ms | 0.051 -> 0.037 ms |
| 98,304 | 0.171 -> 0.078 ms | 0.125 -> 0.073 ms |
| 160,965 | 0.218 -> 0.102 ms | 0.162 -> 0.096 ms |
| 262,144 | 0.383 -> 0.158 ms | 0.275 -> 0.154 ms |

Below about 4K cells the register kernel's fixed cost is 3-4 µs more per call. Over 12 QSA layers that is under 0.2% of a ~26 ms window, and short-context decode measured the same.

**Where it showed up.** A decode profile (`STRATA_VERIFY_PROFILE=1 STRATA_DECODE_TIMING=1 STRATA_SPLIT_TIMING=1`) at 1K / 64K / 157K showed the same ~2 tokens per window at every length. The QSA `scores+topk` per window grew as follows:

| Stage | 1K | 64K | 157K | 157K with this PR |
|---|---|---|---|---|
| R9700 | 0.10 ms | 1.31 ms | 2.05 ms | 1.30 ms |
| XTX | 0.08 ms | 0.98 ms | 1.47 ms | 1.11 ms |

**End to end.**
- **Machine:** Radeon AI PRO R9700 32 GB + RX 7900 XTX 24 GB, layer split 26 (R9700 first), `--pipeline-windows 2` (the default), Core Ultra 7 270K Plus, 125 GiB RAM, Ubuntu 26.04.1, ROCm 7.2.4.
- **Model and settings:** Flash-Next GSQ-RCO IQ3_S, `--max-context 262144`, int8 KV, `--kv-resident 32768`, `--spec 4 --mtp`.
- **Method:** one restart per arm, run in the order off, on, off, on. Four fixed questions with no prefix, then a 156,996-token prefix read once and the same four questions on it from the prompt cache. 300 tokens each, greedy. ms/window comes from the engine's `decode timing` line.

| Arm | Short tok/s | Short ms/window | 157K tok/s | 157K ms/window |
|---|---|---|---|---|
| `STRATA_TOPK_DECODE_REG=0` (before) | 79.5 / 82.1 | 27.53 / 26.98 | 66.8 / 64.7 | 32.91 / 32.25 |
| this PR | 82.4 / 82.6 | 27.53 / 27.52 | 69.8 / 68.6 | 31.10 / 31.87 |

Tokens per window were 2.17-2.38 in every arm.

**Builds:**
- HIP built (`gfx1100;gfx1201`).
- CUDA was not built, but its preprocessed source is unchanged.
- SYCL has its own file.

## Extra Notes
- Related: #669 found the same cost of the reference top-k above 135K cells on CUDA.
- A further change I measured but left out of this PR: a parallel digit search inside `block_topk_reg_kernel`. It is another 10-25% on the kernel and showed no visible end-to-end change. It is on my fork's [`xtx-r9700-tuning`](https://github.com/rindi-ismail/Strata/tree/xtx-r9700-tuning) branch with the raw logs and scripts ([docs/XTX_R9700_SPLIT.md](https://github.com/rindi-ismail/Strata/blob/xtx-r9700-tuning/docs/XTX_R9700_SPLIT.md)).

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

No site

Links install, modelos, releases.