Pull requests / #2698

#2698 qwen35/q36: --kv - int8 K/V caches (codes + a FP16 scale per 64 values)

open · @TorBere · 0 comentarios · En GitHub

BenchmarksSetup & installServer & APINVIDIA / CUDAModels & quantsWindows

Descripción


Implements the int8 K/V cache for the q36 (qwen35moe) engine's own attention path as `--kv` (strata's name
for the stores: `fp32` default, `int8` for both stores, `k8v32` / `k32v8` for one of the two).  The q36 caches
were FP32 (2 KiB per cell per store), so a 128K context ran out of VRAM for expert slots: 0/256 experts in
VRAM, 34 tok/s.  This is the same effect issue #2464 reports for the V100 (the q36 decode falling from 154 to
56 tok/s as the context fills; llama.cpp quantizes the cache, this engine does not).

## What it does

- `--kv int8`: K and V stored as int8 codes with one FP16 scale per 64 values (group absmax / 127), ~3.9x
  smaller cache (528 B/cell vs 2048 B at head_dim 256); `k8v8` is an alias; `k8v32` / `k32v8` quantize one
  store only; `fp32` (the default) runs bit for bit the old kernels.
- `STRATA_Q36_KV=int8` stays as the environment fallback (`--kv` wins).
- Decode (`k_attn_part`), prompt chunk (`k_kv_f16` -> cuBLAS) and the MTP block dequantize on the fly; the
  verify windows go through the same kernels, the commit/rollback of the cache is unchanged.
- `k_attn_prep` / `k_attn_part` / `k_kv_f16` are `template <bool KI8, bool VI8>`; a small `kv4()` dispatcher
  in `Engine` launches the right specialization, so `kc/vc` stay FP32 pointers and the int8 codes + scales
  live in `kc8/vc8/ks8/vs8`.  `kKVG = 64`; `head_dim` must be a multiple of it (the host says so and stays
  FP32 otherwise).

## Measurements

RTX 3060 12 GB, Qwen3.6-35B-A3B UD-Q4_K_XL (experts 18.75 GiB, streaming), `--mtp 2`, same machine, two
runs each, 120 tokens per run:

| variant | 64K | 128K | experts in VRAM (64K / 128K) |
| --- | --- | --- | --- |
| installed engine (CUDA 13), FP32 | 50.3 tok/s | 35.8 tok/s | 18 / 0 of 256 |
| this source, FP32 (default) | 50.1 tok/s | 33.9 tok/s | 18 / 0 of 256 |
| this source, `--kv int8` | **59.7 tok/s** | **54.8 tok/s** | **46 / 31 of 256** |

At 128K the freed VRAM (~4 GiB) brings 31 experts back into the VRAM, which is where the +53% comes from.
The FP32 build reproduces the installed engine (50.1 vs 50.3 / 33.9 vs 35.8), so the difference is the patch,
not the toolchain.

Needle-in-the-haystack, 52.8K-token prompt, 128K mode, max_tokens 250, temperature 0: both FP32 and int8
find the hidden number; the int8 run took 91 s vs 159 s end to end (1.75x).  A chat answer at the same
context is qualitatively equal between the two.

`--kv k8v32` / `k32v8` smoke tested at 64K: each quantizes one store, so they land between FP32 and int8 —
32 experts in VRAM (FP32: 18, int8: 46) and 54.8 / 53.8 tok/s (FP32: 50.1, int8: 59.7).

## Not in this PR

`--kv q4_0` and `k8v4` (PR #21's Hadamard rotation over head_dim 256) are not wired into the q36 kernels
yet; `--kv` refuses them with a note.  The FWHT-256/Q4_0 pieces already exist in the repo
(`src/kernels/cuda/kv_q4.cu`, `include/strata/kernels/kv_q4.hpp`), so a follow-up can add them the same way
(rotate K after RoPE and V, rotate the queries, rotate the output back).

## How it was tested

- Built from this tree with the setup toolchain (CUDA 12.9, sm_86) into a separate build folder; the
  installed engine was not touched.
- A/B runs against the installed engine and the same source without the patch (fp32 path), same prompts.
- `--kv int8` end to end through `serve/server.py` (the config's args carry `--kv int8`).
- The `q36ref` fp32 path is untouched (the FP32 template instantiation is the previous kernel code).

En el sitio

Enlaces a install, modelos, releases.