Pull requests / #241

#241 Faster grouped and per-hit Q2_0 expert kernels, bitwise identical

closed · @gputier · 0 commentaires · Sur GitHub

BenchmarksSetup & installServer & APIAMD / HIPNVIDIA / CUDAModels & quantsWindows

Description

## Faster grouped and per-hit Q2_0 expert kernels, bitwise identical

Rebased on v0.1.29. One commit, four files: `src/kernels/cuda/s2_expert_grouped.cu`, its header, a new parity test, and the test's CMake entry.

### What it does

The grouped kernels (verify window, MTP layer) staged each entry's activations as `xs_q[k][c * 8 + j]`, an 8-way bank conflict on every read, and recomputed the chunk sum `hx` with eight dp4a per row. The per-hit kernels (token graph) read each chunk with 44 byte loads.

New kernels:

- activations are regrouped once per chunk, so the codes expand with shift and mask and `hx` is summed once per chunk;
- grouped: word-major shared layout plus `(dx, hx)` pairs, conflict-free on the staging writes and the reads, two rows per warp;
- per hit: nine aligned word loads per activation chunk, `uint2` code loads, 16-bit scale loads, one warp per (gate, up) pair or pair of down rows.

`s`, `hx` and the float expression `dw * dx * (float) (s - hx)` are unchanged, so every output is the same bits as before. Anything the wide loads cannot take (unaligned activation row, odd stride, odd blob arena) falls back to the previous kernels, and `moe_grouped_last_path()` reports which ones ran.

### Default path is identical to the byte

Measured on an RTX 5090 (sm_120, Windows, MSVC 2022, CUDA 13.4), Release build, v0.1.29 as the base.

- `s2_expert_grouped_parity` (new, in ctest): runs every entry point with the previous kernels, then the new ones, on synthetic Q2_0 blobs and Q8_0 activations. Output and scratch match byte for byte in all 11 cases (fp32 scales, fp16 d, `_dev` with a partial count, `_multi`, `moe_grouped_s2` with and without resident groups, 2-byte aligned activations, slot stride BLOB + 4, arena 4 mod 8). It also checks against a double-precision host reference (worst 1.2e-07 of sum|term|). Result: 0 failures.
- Full ctest, same machine: v0.1.29 alone 35 of 38 pass, this branch 36 of 39. The same three fail in both, and only because the Q2_0 model files are not installed there: `ple_parity` (Q2_0 PLE table), `expert_parity` and `pool_test` (`pack/full/experts.bin`).
- End to end greedy A/B against v0.1.29 alone, IQ3_S model, `--spec 4 --expert-cache 8000 --prefill 2048 --max-new 256`, three prompts (32, 2 520 and 23 020 tokens), two passes each: all 256 output tokens are identical in every run, and the same with `STRATA_OLD_GROUPED=1` on the shortest and longest prompt. This model is not Q2_0, but reading `src/core/mtp.cpp` the MTP layer's resident experts call `moe_grouped_s2`, so the drafting path should run through the new kernels (from the code, not confirmed by a counter).

### Gain

Kernel micro-benchmark (`s2_expert_grouped_parity --bench`, RTX 5090, idle GPU, 312 MB of blobs cycled to defeat L2):

| call | previous | new |
| --- | --- | --- |
| `moe_hit_grouped_s2`, 9 hits | 50.4 us | 39.6 us (x1.27) |
| `moe_hit_grouped_s2`, 1 to 5 hits | 36.5 to 39.3 us | 35.2 to 39.2 us (x1.00 to x1.06) |
| `moe_grouped_s2`, 4 tokens, no sharing | 58.5 us | 52.9 us (x1.11) |
| `moe_grouped_s2`, 8 tokens, 30% shared | 76.0 us | 64.9 us (x1.17) |
| `moe_grouped_s2`, 8 tokens, all shared | 62.9 us | 54.8 us (x1.15) |

End to end, the gain is not measurable on the IQ3_S model: decode is 142.8 tok/s for v0.1.29 and 142.4 / 142.0 for this branch on the 23K prompt (137.5 / 136.9 against 138.1 / 137.0 on the 2.5K one), within the run to run noise. For that model the new kernels are reached through the MTP layer only. A Q2_0 model is where the verify window and the token graph use them, and none was available for this run.

### Switch back

`STRATA_OLD_GROUPED=1` (or `moe_grouped_select_old(1)`) selects the previous kernels. `STRATA_GROUPED_PAIR_MIN_HITS=N` keeps the previous per-hit kernels below N hits of capacity.

### Not tested

- HIP / ROCm: not built, not run.
- Two GPUs or more.
- A Q2_0 model end to end (verify window and token graph paths): the Q2_0 pack was not installed, so `expert_parity`, `pool_test` and `ple_parity` could not run.
- Quants other than IQ3_S end to end, and long-running or sampled (non-greedy) generation.
- Other GPU architectures than sm_120.

### Unrelated, spotted while syncing

`serve/server.py` refuses an empty API key (`STRATA_API_KEY=""` or `--api-key ""`) only after `StrataEngine(...)` has loaded, so the user waits the whole load before exit code 2, and `engine`/`vision` are not closed. Moving that check right after `parse_args()` fixes both.

Sur le site

Liens install, modèles, releases.