Pull requests / #2804
#2804 sycl: KREG small-batch mmvq kernels for Q4_K/Q5_K/Q6_K/IQ4_XS at 1-8 columns (opt-in STRATA_MMVQ_KREG=1)
open · @hotschmoe · 0 Kommentare · Auf GitHub
BenchmarksMulti-GPUAMD / HIPNVIDIA / CUDAModels & quantsDocumentation
Beschreibung
## Title
Issue: Resolves #2803
## Summary
On the Arc Pro B70 the SYCL port's decode GEMVs for Q4_K, Q5_K, Q6_K and IQ4_XS (`native_mmvq`, 1-8 columns) are
bound by instruction issue, not by memory: the wide kernels re-read every column's activations for every row. This
adds an opt-in kernel family, KREG: one work-group covers a row's whole K, each lane loads its slice of every
column's activations once into registers, and the group streams tiles of 2-8 rows against them.
`STRATA_MMVQ_KREG=1` turns it on. Unset, the default path is unchanged (output identical to main, measured below).
Upstream main 296b59d2 against this branch: 2x Intel Arc Pro B70 32 GB, kernel 7.1, Compute Runtime 26.22.38646.4,
oneAPI 2026.1, AOT bmg-g31; Qwen3.8-Flash-Next ISTA GSQ-RCO IQ3_XXS, `--layer-split 30 --prefill 8192
--max-context 65536 --kv int8 --spec 4 --mtp <BF16 draft layer> --mtp-max-t 4 --spec-min-p 0 --stream-experts
--expert-cache auto`, one slot, greedy, a fixed prompt with a 512-token coding reply, 3 runs each, the container
environment of docs/INTEL.md.
| | decode tok/s | tokens per verify window | ms per window |
|---|---|---|---|
| main | 81.8 / 79.5 / 78.7 | 3.22 / 3.14 / 3.12 | 39.4 / 39.5 / 39.7 |
| this branch, flag unset | 81.9 / 79.4 | the same as main | - |
| this branch, `STRATA_MMVQ_KREG=1` | 95.0 / 93.1 / 96.3 | 3.35 / 3.20 / 3.33 | 35.2 / 34.4 / 34.5 |
- Median decode 79.5 -> 95.0 tok/s (+19.5%). The part that does not depend on the drafts is the verify window,
39.5 -> 34.5 ms (-12%); the rest is the changed output drafting a little better on this prompt.
- Flag unset: output and draft counts identical to main on both runs and on a turn at a 40,796-token context.
- At a 40,796-token context: 76.0 -> 83.4 tok/s. Prompt reading unchanged: 4K 1,189 -> 1,191 tok/s, 32K
1,956-2,004 -> 1,961-2,008.
Kernel time, us per call, default -> KREG (`mmvq_kreg_parity --bench`, one B70, random weights, distinct copies per
graph so L2 never holds a matrix):
| format, shape (n_in x n_out) | 1 column | 2 | 4 | 6 |
|---|---|---|---|---|
| Q6_K 2560 x 10240 | 55.7 -> 41.0 | 63.0 -> 41.5 | 73.4 -> 46.0 | 88.9 -> 57.7 |
| Q5_K head 2560 x 248320 | 1063 -> 759 | 1374 -> 775 | 2013 -> 893 | 2691 -> 1175 |
| Q4_K 2560 x 10240 | 44.5 -> 28.8 | 55.6 -> 29.6 | 78.8 -> 31.9 | 107.9 -> 40.5 |
| IQ4_XS 2560 x 10240 | 42.9 -> 33.5 | 47.9 -> 37.8 | 55.9 -> 41.2 | 67.6 -> 49.3 |
## What changed
- `sycl/src/kernels/cuda/native_mmvq.dp.cpp`
- KREG kernels: `native_mmvq_kreg_kernel` with traits `KregQ4K`, `KregQ5K`, `KregQ6K`, `KregIQ4XS` (the wide
kernels' lane split and dot expressions; K-quant scales decoded branch-free; IQ4_XS's codebook from a
256-entry local-memory table), a transposed sub-group butterfly and a double-buffered local-memory sum across
sub-groups. Rows per tile and the 256-register mode per format and column count come from a sweep on the B70.
- `try_kreg<F>()`: its first statement reads the flag (latched once, one stderr line when on). It declines n_in
above 32,768 and Q4_K/Q5_K (IQ4_XS) matrices that are not 16- (8-) byte aligned; the existing kernels then run
as before.
- One guarded hook per entry: `native_q4_k_mmvq`, `native_q5_k_mmvq`, `native_iq4_xs_mmvq` before `try_wide`;
`native_q6_k_mmvq` after the `STRATA_Q6_PACKED` check. The output head goes through `native_mmvq`, so it is
covered.
- `native_mmvq_set_kreg(int)` for tests (1 on, 0 off, -1 the env latch, the default).
- `sycl/include/strata/kernels/native_mmvq.hpp`: its declaration (the port's copy of the header).
- `sycl/src/kernels/mmvq_kreg_parity.cpp` (new), in the `sycl/CMakeLists.txt` parity loop (`--selftest`). For each
format, 5 shapes and 1-8 columns: KREG and the default kernels against a float64 reference over the same weight
and Q8_1 bytes; KREG differs from the default kernels (the call reached it); a T-column KREG call bitwise equal to
T one-column calls; a power check (one scale step in one sub-block must fail the bound). `--bench` times both on
random weights (constant fills read too fast on Xe2: memory compression).
- `docs/INTEL.md`: a section for the flag with these numbers.
Left alone: the other formats (Q8_0, IQ4_NL, IQ3_S, Q2_0, ...), the multi/small/wave kernels, tensors packed by
`STRATA_Q6_PACKED`, `include/`, `src/` and the CUDA/HIP sources. Only the SYCL backend is touched and built.
## 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.
- Default path: byte-identical by construction (with the flag unset each hook reads a latched bool and returns; no
shared code changed), and measured identical output and draft counts against main (above).
- `mmvq_kreg_parity --selftest` on one B70: 0 failures over 160 cases (4 formats x 5 shapes x 1-8 columns). KREG and
the default kernels are both within 5e-8 of sum |w x| of the reference, every multi-column output is bitwise
equal to the one-column calls, KREG differs from the default kernels in 144 cases, no non-finite output.
- Not bitwise equal to the default kernels (another FP32 summation order), so greedy outputs change after some
tokens and MTP acceptance moves with them; the per-window time is the fair comparison. Teacher-forced NLL on 3
texts (4,762 tokens) through the decode-window path, measured in a development tree (expert parallelism over the
two cards, not upstream main): mean NLL 0.013 nats lower with the flag, inside the noise band of other
summation-order changes. Not re-run on main.
- In that development tree the same kernels gave single-stream decode 107.2 -> 131.7 tok/s and two streams
63.9 -> 73.1 tok/s each.
- IQ4_XS at 1 column on small matrices (2560 x 512 / 640) is 0.90-0.98x of the default kernel. KREG still takes it,
so that whether a call uses KREG never depends on the column count and multi-column calls stay bitwise equal to
one-column calls.
- Not tested: two streams on main (main refuses `--batch-mtp` on a layer split), `--max-context 200000` (a plain
two-card layer split of main does not fit all 24,576 experts there), the Arc A-series cards.
🤖 Generated with [Claude Code](https://claude.com/claude-code)
Mehr auf der Site
Links zu Install, Modellen, Releases.