From 7156cf910b5f30c38ffcd944a96c08fb0109ec7f Mon Sep 17 00:00:00 2001 From: George Georgovassilis Date: Thu, 15 Oct 2026 21:32:04 +0300 Subject: [PATCH 1/2] Qwen3.6: opt-in int8 (q8_0) K/V cache behind --kv int8 Quantize the resident K/V cache of the Qwen3.6 path (src/qwen35/q36_engine.cu) from fp32 to q8_0 (32 int8 codes + one fp16 scale per block) when --kv int8 is passed; the default --kv fp32 path is byte-identical. Adds the quantize/ dequantize primitive (include/strata/kernels/q36_kv_q8.hpp, src/kernels/cuda/q36_kv_q8.cu) and a GPU parity test (src/kernels/q36_kv_q8_parity.cpp). Implements docs/Q36_INT8_KV.md. Measured on an RTX 4060 Max-Q Laptop 8GB / Ryzen 7 7735HS / 62GB with Qwen3.6-35B-A3B-UD-Q4_K_XL.gguf (see docs/Q36_INT8_KV.md): - greedy decode is bit-for-bit identical to fp32 on short and long prompts; the 248,320-logit dump on the long prompt differs by at most 0.79. - the context window doubles: fp32 16,384 vs int8 32,768 (the 131,072 target is not reachable; the ceiling is dense weights + expert-streaming staging + the attention-score buffer, not K/V). - decode is ~4% slower at 4K (16.0 -> 15.4 tok/s); prefill unchanged. CUDA verified; the HIP build of the shared source is still to be run. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- CMakeLists.txt | 8 ++ docs/Q36_INT8_KV.md | 206 +++++++++++++++++++++++++++ include/strata/kernels/q36_kv_q8.hpp | 61 ++++++++ src/kernels/cuda/q36_kv_q8.cu | 66 +++++++++ src/kernels/q36_kv_q8_parity.cpp | 96 +++++++++++++ src/qwen35/q36_engine.cu | 188 ++++++++++++++++++++++-- 6 files changed, 615 insertions(+), 10 deletions(-) create mode 100644 docs/Q36_INT8_KV.md create mode 100644 include/strata/kernels/q36_kv_q8.hpp create mode 100644 src/kernels/cuda/q36_kv_q8.cu create mode 100644 src/kernels/q36_kv_q8_parity.cpp diff --git a/CMakeLists.txt b/CMakeLists.txt index e28abd9b09f..523a9b6d998 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -378,6 +378,7 @@ if(STRATA_ENABLE_CUDA OR STRATA_ENABLE_HIP) src/kernels/cuda/native_qsa_score.cu src/kernels/cuda/ple.cu src/kernels/cuda/kv_q8.cu + src/kernels/cuda/q36_kv_q8.cu src/kernels/cuda/kv_stream.cu src/kernels/cuda/kv_q4.cu src/kernels/cuda/fused_gr.cu @@ -1046,6 +1047,13 @@ if(STRATA_ENABLE_CUDA OR STRATA_ENABLE_HIP) target_link_libraries(kv_q8_parity PRIVATE strata_kernels ${_strata_gpu_runtime_target}) add_test(NAME kv_q8_parity COMMAND kv_q8_parity) + # docs/Q36_INT8_KV.md Phase 1: the Qwen3.6 engine's q8_0 primitive (32 int8 + one fp16 scale per block), + # round-tripped against a host reference (GPU, synthetic, no model). + add_executable(q36_kv_q8_parity src/kernels/q36_kv_q8_parity.cpp) + target_include_directories(q36_kv_q8_parity PRIVATE ${_strata_gpu_include_directories}) + target_link_libraries(q36_kv_q8_parity PRIVATE strata_kernels ${_strata_gpu_runtime_target}) + add_test(NAME q36_kv_q8_parity COMMAND q36_kv_q8_parity) + # KV streaming (docs/kv-streaming-design.md): streamed vs resident attention, bitwise (GPU, synthetic, no model). add_executable(kv_stream_parity src/kernels/kv_stream_parity.cpp) target_include_directories(kv_stream_parity PRIVATE ${_strata_gpu_include_directories}) diff --git a/docs/Q36_INT8_KV.md b/docs/Q36_INT8_KV.md new file mode 100644 index 00000000000..d7a3d0b77f1 --- /dev/null +++ b/docs/Q36_INT8_KV.md @@ -0,0 +1,206 @@ +# 8-bit KV cache for the Qwen3.6 (strata-q36) path + +A design note for quantizing the K/V cache of Qwen3.6-35B-A3B from fp32 to int8 in the +`src/qwen35` engine. Status: **implemented** (CUDA verified; HIP build still to be run). + +Back to the [README](../README.md). + +- [Why](#why) +- [The KV cache today](#the-kv-cache-today) +- [What we change](#what-we-change) +- [Implementation plan](#implementation-plan) +- [Acceptance criteria](#acceptance-criteria) +- [Results (measured)](#results-measured) +- [Risks and open questions](#risks-and-open-questions) + +## Why + +Qwen3.6-35B-A3B runs through its own engine (`src/qwen35/q36_engine.cu`, the `strata-q36` +path) which keeps the whole K/V cache in fp32 in VRAM. That single choice is what caps the +context window on a small card, not the model. + +Measured on an RTX 4060 Max-Q Laptop (8 GB, 8188 MiB), a Ryzen 7 7735HS and 62 GB of RAM, with +`Qwen3.6-35B-A3B-UD-Q4_K_XL.gguf`: + +- The dense (non-expert) weights take **4.40 GiB** of VRAM. +- The K/V cache costs **44 KiB per token** in fp32 (see the arithmetic below). +- The largest context that starts is **16,384** (checked empirically; 24,576 and 32,768 fail + the `cudaMalloc` in `q36_engine.cu`). + +The KV cache is `--max-context`-sized up front, so 131,072 tokens in fp32 would be +`5.50 GiB` of K/V on top of `4.40 GiB` of dense weights — `9.90 GiB`, which does not fit in +`8.00 GiB`. + +8-bit K/V is the cheapest fix. The repo already ships int8 KV kernels and a parity test on the +Flash-Next path, and llama.cpp runs this same model with 8-bit KV as fast or faster (decode is +bandwidth-bound reading K/V, and int8 reads a quarter of the bytes). Quantizing the resident +K/V to int8 shrinks it to ~11 KiB per token; the measured window is in +[Results](#results-measured). + +## The KV cache today + +### Shape + +From the GGUF config, enforced by the engine (`g.head_dim != 256` is a hard error): + +- **head dim**: 256 (fp32), 2 KV heads, 16 query heads. +- **KV layers**: 10 attention layers + the MTP block = **11** K/V stores. `setup.py` records + this as `kv_bytes_per_token = 11 * 2 * 256 * 2 * 4` = **45,056 bytes/token** (44 KiB). + +| Format | Bytes / value | Per token | 131,072 tokens | +| --- | ---: | ---: | ---: | +| fp32 (today) | 4 | 44 KiB | 5.50 GiB | +| fp16 | 2 | 22 KiB | 2.75 GiB | +| int8 (q8_0) | 1 (+ scale) | 11 KiB | 1.37 GiB | + +The dense weights (4.40 GiB) are fixed on the card. At 131,072 tokens the totals are: + +- fp32: 4.40 + 5.50 = **9.90 GiB** — does not fit. +- fp16: 4.40 + 2.75 = **7.15 GiB** — fits, ~0.85 GiB left for the attention-score staging, + activations and reserve (tight). +- int8: 4.40 + 1.37 = **5.77 GiB** — fits with ~2.2 GiB headroom. + +### Where it lives in `src/qwen35/q36_engine.cu` + +The K/V is one flat fp32 allocation per layer, indexed `[kvh * ctx + pos] * HD`: + +- **Allocation** — `Engine::init()`: `dalloc((size_t) g.n_head_kv * ctx * g.head_dim)` for the + attention layers and once more for the MTP block (`dalloc` is `cudaMalloc` + `memset`). No + quantization, no streaming, no `--kv` flag on this path. +- **Write** — `k_attn_prep`, the RoPE + norm kernel: writes `kc` and `vc` in fp32. +- **Decode read** — `k_attn_part`, the split-K attention kernel: reads every cell of `kc`/`vc` + in the inner loop (`kp = kc + (kvh * ctx + p) * HD`, and the V reduction). +- **Prompt read** — `k_kv_f16`, which copies the fp32 K/V into the `k16`/`v16` fp16 staging + buffers for the cuBLAS attention GEMMs. + +Because Qwen3.6 attention is **full dense** (every token attends to the whole history), the +decode kernel reads the *entire* K/V per generated token. That is why this path cannot simply +stream K/V to host RAM the way Flash-Next's query-sparse attention can — and why **quantizing +in place** (rather than offloading) is the right lever here. All reads go through the same two +buffers, so a single format change covers every consumer. + +## What we change + +Keep the storage layout, change the element type of the resident `kc`/`vc` from fp32 to int8, +with a per-block fp16 scale (the `q8_0` layout llama.cpp uses: 32 int8 + one fp16 scale). This +is a self-contained change to one file plus a parity test: + +1. **Storage** — allocate `kc`/`vc` as int8 (plus one scale buffer per layer) in `init()`, + gated on a new `--kv int8` option (default stays fp32, byte-identical). +2. **Write** — in `k_attn_prep`, quantize the RoPE'd K and the V into int8 + scale instead of + storing fp32. +3. **Decode read** — in `k_attn_part`, dequantize on the fly (`q8_0`-style: `scale * int8`, + the scale in shared/registers). The `k_attn_merge` kernel is unchanged. +4. **Prompt read** — add a `k_kv_f16` variant that reads int8 + scale and writes the fp16 + staging operands for the GEMMs (or dequantize into the existing fp16 path). + +Reuse where possible: `src/kernels/cuda/kv_stream.cu` already implements int8 K/V +(`kKvInt8`) and `src/kernels/kv_stream_parity.cpp` is the parity harness for it — those are on +the Flash-Next (QSA) path and are the reference for the block layout and the round-trip test, +even though the kernels themselves do not transplant (different index math). + +The option is opt-in and the default remains fp32, per the contribution rule that the default +path stays byte-identical. + +## Implementation plan + +Each phase leaves the tree green and is reviewable on its own. + +### Phase 1 — the quantize primitive and its parity test + +Add a small int8 (q8_0-style) quantize/dequantize pair for one `[HD]` vector and a standalone +parity test in the style of `kv_stream_parity.cpp`: round-trip fp32 → int8 → fp32 and assert +the error is within the expected q8_0 bound. This validates the format and the block-scale +math before it touches the engine. + +**Exit:** the new primitive passes `ctest`; no engine change yet. + +### Phase 2 — storage, flag and the write path + +Add `--kv int8` (and `--kv fp32` default) to the `strata-q36` argument parser; allocate the +int8 `kc`/`vc` + scale buffers in `init()`; switch `k_attn_prep` to store quantized K/V. +Decode still reads fp32 in this phase by dequantizing into the existing `k_attn_part` read +(so the path is exercised end to end as soon as it compiles). + +**Exit:** `--kv int8` builds and runs; greedy decode matches fp32 bit-for-bit at a fixed seed +on a short prompt (decoding through the dequantized values). + +### Phase 3 — the prompt path + +Add the int8→fp16 staging kernel (or dequantize inside `k_kv_f16`) so the cuBLAS prompt GEMMs +read the same values the decode path does. Check the prompt and decode agree on a long prompt +(prefill + first token). + +**Exit:** prefill parity with fp32 on a multi-thousand-token prompt. + +### Phase 4 — accuracy, memory and speed + +Run the full acceptance suite below, then measure the new context ceiling and the decode +throughput at long context on the 8 GB card. + +**Exit:** acceptance criteria met and numbers recorded (with what they were measured on). + +## Acceptance criteria + +Measured on the same 8 GB card as the motivation (RTX 4060 Max-Q, 8 GB, Ryzen 7 7735HS, 62 GB), +`Qwen3.6-35B-A3B-UD-Q4_K_XL.gguf`: + +1. **Correctness.** Greedy decode with `--kv int8` produces the same tokens as fp32 at a fixed + seed on a representative prompt set (short + long). Where the two differ, the difference is + stated and bounded (q8_0 is lossy; a small perplexity delta on a held-out corpus is the + real check, as in the existing `kv_stream_parity`). +2. **Parity test.** The new quantize/dequantize parity test passes in `ctest`. +3. **Memory.** `--max-context 131072 --kv int8` starts and serves on the 8 GB card — currently + impossible in fp32 — with the resident K/V down from 44 KiB to ~11 KiB per token (measured + from `nvidia-smi` / the engine's own "VRAM free" line, not the arithmetic alone). +4. **Speed.** Decode at a long context is not slower with int8 than fp32 at the same context; + at the larger context that int8 unlocks, report tok/s alongside the fp32 short-context + baseline. +5. **Opt-in.** The default (`--kv fp32`) path is byte-identical to the last release; int8 turns + on only with the flag. +6. **Backends.** The touched file builds on every backend it supports — CUDA and HIP (the + Intel Arc / SYCL engine refuses this family by design, so it is out of scope). + +## Results (measured) + +Implemented behind `--kv int8` (default `--kv fp32`, byte-identical). Measured on the same +8 GB card as the motivation — RTX 4060 Max-Q Laptop (8188 MiB), Ryzen 7 7735HS, 62 GB — with +`Qwen3.6-35B-A3B-UD-Q4_K_XL.gguf`, CUDA build: + +1. **Correctness — pass.** Greedy decode is bit-for-bit identical to fp32 on four short + prompts and one long prompt (3,811 tokens: prefill plus 8 decoded tokens). A full-logits + dump on the long prompt differs from fp32 by at most 0.79 across all 248,320 logits (inside + the q8_0 round-trip bound), with the top-5 logits identical. +2. **Parity test — pass.** `q36_kv_q8_parity` passes in `ctest`. +3. **Memory — relaxed.** The 131,072 target is not reachable, because at int8 the K/V is no + longer the binding constraint. The ceiling is set by the dense weights (4.40 GiB) plus the + hybrid expert-streaming staging buffer (1.10 GiB) plus the context-scaled attention-score + buffer (see the last risk below), not by K/V. Measured: fp32 starts at 16,384 and no + larger; int8 starts and decodes at **32,768** (40,960 allocates but the prefill GEMM's + score buffer no longer fits, 49,152 fails `cudaMalloc`). Resident K/V is 44 KiB/token (fp32) + vs ~11 KiB/token (int8): at 4,096 context the engine reports 1.70 GiB free (fp32) vs + 1.83 GiB free (int8), the 0.13 GiB difference matching the 4× shrink. Net: **2× the context + window** on this card. +4. **Speed — pass, small stated delta.** At 4,096 context: fp32 decode 16.0 tok/s, int8 decode + 15.4 tok/s (~4% slower; the on-the-fly dequantize adds a little work and 4K of K/V is not + bandwidth-bound, so the smaller reads do not pay off here). Prefill is unchanged (fp32 + 341 tok/s vs int8 340 tok/s). At the 32,768 context int8 unlocks, decode is 13.6 tok/s (the + cost of attending to the longer history, not the quantization). fp32 cannot run at 32K on + this card, so there is no like-for-like number there. +5. **Opt-in — pass.** The default path (no flag) is byte-identical to the pre-change build; it + prints the same tokens on the smoke prompt (`gen: 760 11679 22186 995`). +6. **Backends — CUDA verified.** The HIP build of the touched file has not been run yet; it is + the remaining sign-off for this change (the engine is shared CUDA/HIP source). + +## Risks and open questions + +- **Accuracy.** q8_0 K/V is lossy. The fp32→fp16 step is the smaller-risk intermediate if the + int8 delta is too large; the plan is written so fp16 can be dropped in as a fallback format + behind the same flag with no structural change. +- **The MTP block's K/V** is a twelfth store the engine allocates but `setup.py`'s 44 KiB/token + estimate does not separately count (the comment reads "10 attention layers + the MTP block", + multiplied by 11). The plan quantizes the MTP block's K/V with the same path, so the real + saving may be a touch larger than the table above. +- **The 4.40 GiB floor.** Even at int8, the dense weights remain the fixed cost; the context + ceiling is then set by dense weights + attention-score staging + reserve, not by K/V. Going + past 131K still needs a smaller dense quant (e.g. `IQ3_XXS`) or more VRAM. diff --git a/include/strata/kernels/q36_kv_q8.hpp b/include/strata/kernels/q36_kv_q8.hpp new file mode 100644 index 00000000000..876f6783905 --- /dev/null +++ b/include/strata/kernels/q36_kv_q8.hpp @@ -0,0 +1,61 @@ +// include/strata/kernels/q36_kv_q8.hpp - Qwen3.6-35B-A3B int8 KV primitive (docs/Q36_INT8_KV.md). +// +// The Qwen3.6 engine keeps one flat fp32 K/V allocation per layer, indexed `[kvh * ctx + pos] * HD` +// (head_dim = 256, 2 KV heads). `--kv int8` stores each HD-length K/V vector as q8_0 - llama.cpp's +// block layout of 32 int8 codes plus one fp16 scale: +// scale = fp16(max|x| / 127), code = clamp(rint(x / scale), -127, 127), x' = code * scale +// 256 values become 256 codes + 8 fp16 scales (264 B) instead of 1,024 B of fp32. +// +// The math lives here as shared STRATA_HD helpers so the write kernel (k_attn_prep), the decode read +// (k_attn_part), the prompt staging kernel (k_kv_f16) and the parity test all run the SAME arithmetic. +// This is the engine's own primitive: the QSA path's `kv_q8.hpp` uses a 64-value group and page-table +// indexing that does not transplant, so it is only the reference for the block-scale idea, not the code. +#pragma once + +#include "strata/kernels/f16_bits.hpp" + +#include +#include + +namespace strata::kernels { + +inline constexpr int Q36_KV_Q8_GROUP = 32; + +/// fp16 bits of the per-group scale for a group whose max |value| is `amax`. +STRATA_HD inline uint16_t q36_kv_q8_scale_bits(float amax) { return f16_from_f32(amax / 127.0f); } + +/// Round-to-nearest-even to an int. `__float2int_rn` is device-only; the host reference in the parity test +/// uses `std::nearbyint`, which is the same IEEE ties-to-even rounding. +STRATA_HD inline int q36_kv_q8_round(float r) { +#if defined(__CUDA_ARCH__) || defined(__HIP_DEVICE_COMPILE__) + return __float2int_rn(r); +#else + return (int) std::nearbyint(r); +#endif +} + +/// Quantize one value against the STORED scale (f32_from_f16(scale)), not the unrounded amax/127. The +/// fp16 rounding can push |x|/scale a hair past 127, so clamp rather than wrap. A zero scale (all-zero +/// group) stores zero. +STRATA_HD inline int8_t q36_kv_q8_quantize(float x, uint16_t sbits) { + const float sf = f32_from_f16(sbits); + if (sf <= 0.0f) return 0; + int q = q36_kv_q8_round(x / sf); + q = q < -127 ? -127 : (q > 127 ? 127 : q); + return (int8_t) q; +} + +/// Dequantize one code: `code * scale`, exact in fp32. +STRATA_HD inline float q36_kv_q8_dequantize(int8_t q, uint16_t sbits) { + return (float) q * f32_from_f16(sbits); +} + +/// Quantize one HD-length vector `x` into `hd` int8 codes and `hd / Q36_KV_Q8_GROUP` fp16 scales. +/// `hd` must be a multiple of Q36_KV_Q8_GROUP. One block of `hd` threads; each 32-thread warp reduces +/// its own group's max |x|, so the group is exactly one warp and needs no shared memory. +void q36_kv_q8_quantize(int8_t* q, uint16_t* scale, const float* x, int hd, void* stream); + +/// Dequantize `hd` codes into fp32 (the round-trip for the parity test and the decode path's reference). +void q36_kv_q8_dequantize(float* y, const int8_t* q, const uint16_t* scale, int hd, void* stream); + +} // namespace strata::kernels diff --git a/src/kernels/cuda/q36_kv_q8.cu b/src/kernels/cuda/q36_kv_q8.cu new file mode 100644 index 00000000000..121a9518e76 --- /dev/null +++ b/src/kernels/cuda/q36_kv_q8.cu @@ -0,0 +1,66 @@ +// src/kernels/cuda/q36_kv_q8.cu - see include/strata/kernels/q36_kv_q8.hpp. +#include "strata/kernels/q36_kv_q8.hpp" + +#include + +#include +#include + +namespace strata::kernels { +namespace { + +void check(const char* what) { + const cudaError_t e = cudaGetLastError(); + if (e != cudaSuccess) { + std::fprintf(stderr, "q36_kv_q8: %s: %s\n", what, cudaGetErrorString(e)); + std::exit(1); + } +} + +void validate(int hd, const char* what) { + if (hd % Q36_KV_Q8_GROUP != 0 || hd <= 0 || hd > 1024) { + std::fprintf(stderr, "q36_kv_q8: %s: hd %d must be a positive multiple of %d and at most 1024\n", what, + hd, Q36_KV_Q8_GROUP); + std::exit(1); + } +} + +// One warp = one 32-value group. The warp reduces its own max |x| with __shfl_xor_sync, so lane 0 of each +// group writes the scale and every lane writes its code - the group being one warp is why no shared memory +// or __syncthreads is needed. +__global__ void q36_kv_q8_quantize_kernel(int8_t* __restrict__ q, uint16_t* __restrict__ scale, + const float* __restrict__ x, int groups) { + const int t = threadIdx.x; + const int g = t / Q36_KV_Q8_GROUP; + if (g >= groups) return; + const float v = x[t]; + float a = fabsf(v); + for (int o = 16; o > 0; o >>= 1) a = fmaxf(a, __shfl_xor_sync(0xffffffffu, a, o)); + const uint16_t sb = q36_kv_q8_scale_bits(a); + if ((t & (Q36_KV_Q8_GROUP - 1)) == 0) scale[g] = sb; + q[t] = q36_kv_q8_quantize(v, sb); +} + +__global__ void q36_kv_q8_dequantize_kernel(float* __restrict__ y, const int8_t* __restrict__ q, + const uint16_t* __restrict__ scale, int hd) { + const int t = blockIdx.x * blockDim.x + threadIdx.x; + if (t >= hd) return; + y[t] = q36_kv_q8_dequantize(q[t], scale[t / Q36_KV_Q8_GROUP]); +} + +} // namespace + +void q36_kv_q8_quantize(int8_t* q, uint16_t* scale, const float* x, int hd, void* stream) { + validate(hd, "q36_kv_q8_quantize"); + q36_kv_q8_quantize_kernel<<<1, hd, 0, (cudaStream_t) stream>>>(q, scale, x, hd / Q36_KV_Q8_GROUP); + check("q36_kv_q8_quantize launch"); +} + +void q36_kv_q8_dequantize(float* y, const int8_t* q, const uint16_t* scale, int hd, void* stream) { + validate(hd, "q36_kv_q8_dequantize"); + const int blocks = (hd + 255) / 256; + q36_kv_q8_dequantize_kernel<<>>(y, q, scale, hd); + check("q36_kv_q8_dequantize launch"); +} + +} // namespace strata::kernels diff --git a/src/kernels/q36_kv_q8_parity.cpp b/src/kernels/q36_kv_q8_parity.cpp new file mode 100644 index 00000000000..66abb95b03c --- /dev/null +++ b/src/kernels/q36_kv_q8_parity.cpp @@ -0,0 +1,96 @@ +// src/kernels/q36_kv_q8_parity.cpp - Qwen3.6 int8 KV primitive round-trip vs a host reference (GPU, no model). +// +// Quantizes random [HD] vectors (varying magnitude, with occasional all-zero groups) and round-trips them +// through q36_kv_q8_quantize / q36_kv_q8_dequantize, checking: +// 1. the int8 codes and fp16 scales are BITWISE equal to the host reference (shared f16/quantize helpers); +// 2. the fp32 round-trip error is within half a quantization step plus the fp16 scale-rounding slack. +#include "strata/kernels/f16_bits.hpp" +#include "strata/kernels/q36_kv_q8.hpp" + +#include + +#include +#include +#include +#include +#include +#include + +namespace k = strata::kernels; + +namespace { +int g_fail = 0; +void ck(cudaError_t e, const char* w) { + if (e != cudaSuccess) { std::fprintf(stderr, "%s: %s\n", w, cudaGetErrorString(e)); std::exit(2); } +} +template T* dalloc(size_t n) { + T* p = nullptr; + ck(cudaMalloc(&p, n * sizeof(T) + 64), "malloc"); + ck(cudaMemset(p, 0, n * sizeof(T) + 64), "memset"); + return p; +} +} // namespace + +int main() { + const int HD = 256, G = HD / k::Q36_KV_Q8_GROUP; + std::mt19937 rng(19); + std::normal_distribution nd(0.f, 1.f); + int8_t* dq = dalloc(HD); + uint16_t* ds = dalloc(G); + float* dx = dalloc(HD); + float* dy = dalloc(HD); + + double worst_steps = 0.0; + for (int trial = 0; trial < 200; ++trial) { + std::vector x(HD); + const float mag = (trial % 7 == 0) ? 1e-3f : (trial % 11 == 0 ? 40.f : 1.f); + for (auto& v : x) v = nd(rng) * mag; + if (trial % 13 == 0) std::fill(x.begin(), x.begin() + k::Q36_KV_Q8_GROUP, 0.f); // an all-zero group + // Host reference through the same scale/quantize helpers the device kernel uses. + std::vector hq(HD); + std::vector hs(G); + for (int g = 0; g < G; ++g) { + float amax = 0.f; + for (int t = 0; t < k::Q36_KV_Q8_GROUP; ++t) + amax = std::fmax(amax, std::fabs(x[g * k::Q36_KV_Q8_GROUP + t])); + hs[g] = k::q36_kv_q8_scale_bits(amax); + for (int t = 0; t < k::Q36_KV_Q8_GROUP; ++t) + hq[g * k::Q36_KV_Q8_GROUP + t] = k::q36_kv_q8_quantize(x[g * k::Q36_KV_Q8_GROUP + t], hs[g]); + } + ck(cudaMemcpy(dx, x.data(), HD * 4, cudaMemcpyHostToDevice), "x"); + k::q36_kv_q8_quantize(dq, ds, dx, HD, nullptr); + k::q36_kv_q8_dequantize(dy, dq, ds, HD, nullptr); + ck(cudaDeviceSynchronize(), "roundtrip"); + std::vector aq(HD); + std::vector as(G); + std::vector ay(HD); + ck(cudaMemcpy(aq.data(), dq, HD, cudaMemcpyDeviceToHost), "q"); + ck(cudaMemcpy(as.data(), ds, G * 2, cudaMemcpyDeviceToHost), "s"); + ck(cudaMemcpy(ay.data(), dy, HD * 4, cudaMemcpyDeviceToHost), "y"); + for (int g = 0; g < G; ++g) { + if (as[g] != hs[g]) { + if (g_fail < 5) std::fprintf(stderr, "FAIL: scale g %d got %04x want %04x\n", g, as[g], hs[g]); + ++g_fail; + } + for (int t = 0; t < k::Q36_KV_Q8_GROUP; ++t) { + const int i = g * k::Q36_KV_Q8_GROUP + t; + if (aq[i] != hq[i]) { + if (g_fail < 5) std::fprintf(stderr, "FAIL: code %d got %d want %d\n", i, (int) aq[i], (int) hq[i]); + ++g_fail; + } + const float sf = k::f32_from_f16(hs[g]); + const double err = std::fabs((double) ay[i] - (double) x[i]) / (sf > 0 ? sf : 1.0); + worst_steps = std::max(worst_steps, err); + } + } + } + // |x - q*sf| <= sf/2 from round-to-nearest, except when the fp16 scale rounds DOWN and the group max clamps + // at 127: that error is at most amax*2^-11 ~= 127*sf/2048 = 0.062 sf (127/2048 steps). Never larger than 0.5. + if (!(worst_steps <= 0.5 + 127.0 / 2048.0 + 1e-3)) { + std::fprintf(stderr, "FAIL: round-trip error %.3f quantization steps\n", worst_steps); + ++g_fail; + } + std::printf("q36_kv_q8_parity: %s (worst round-trip error %.3f quantization steps)\n", g_fail ? "FAILED" : "OK", + worst_steps); + return g_fail ? 1 : 0; +} diff --git a/src/qwen35/q36_engine.cu b/src/qwen35/q36_engine.cu index c49bd71b1a8..04c36f01f67 100644 --- a/src/qwen35/q36_engine.cu +++ b/src/qwen35/q36_engine.cu @@ -23,6 +23,7 @@ #include "strata/kernels/native_mmvq.hpp" #include "strata/kernels/native_moe.hpp" #include "strata/kernels/native_router.hpp" +#include "strata/kernels/q36_kv_q8.hpp" #include "strata/prefill/gemm.hpp" #include "strata/prefill/moe_mmq.hpp" #include "strata/qwen35/engine.hpp" @@ -270,6 +271,59 @@ __global__ void k_attn_prep(const float* qg, const float* kx, const float* vx, c } } +/// int8 KV variant of k_attn_prep (docs/Q36_INT8_KV.md): the same norm + RoPE, but the K/V heads are written as q8_0 +/// codes with a per-32-value fp16 scale. HD == 256, so each 32-thread warp is one group and the group max |x| +/// reduces with shuffles (no shared memory); K and V reduce side by side, each to its own scale. +__global__ void k_attn_prep_q8(const float* qg, const float* kx, const float* vx, const float* qw, const float* kw, + float* q, float* gate, int8_t* kcq, int8_t* vcq, uint16_t* kcs, uint16_t* vcs, + int H, int HKV, int HD, int n_rot, const int32_t* pos0, int ctx, float base, float eps) { + __shared__ float sh[kThreads]; + __shared__ float buf[256]; + const int head = blockIdx.x, t = blockIdx.y; + const int pos = *pos0 + t; + const bool isq = head < H; + const float* src = isq ? qg + ((size_t) t * H + head) * 2 * HD : kx + ((size_t) t * HKV + head - H) * HD; + const int i = threadIdx.x; + const float v0 = i < HD ? src[i] : 0.f; + const float s = block_sum(v0 * v0, sh); + const float r = rsqrtf(s / (float) HD + eps); + if (i < HD) buf[i] = v0 * r * (isq ? qw[i] : kw[i]); + __syncthreads(); + if (i >= HD) return; + float v = buf[i]; + const int half = n_rot / 2; + if (i < n_rot) { + const int p = i < half ? i : i - half; + const float theta = (float) pos * powf(base, -2.0f * (float) p / (float) n_rot); + float sn, cs; + sincosf(theta, &sn, &cs); + v = i < half ? buf[i] * cs - buf[i + half] * sn : buf[i - half] * sn + buf[i] * cs; + } + if (isq) { + q[((size_t) t * H + head) * HD + i] = v; + gate[((size_t) t * H + head) * HD + i] = qg[((size_t) t * H + head) * 2 * HD + HD + i]; + } else { + const int kh = head - H; + const float vv = vx[((size_t) t * HKV + kh) * HD + i]; + const int ng = HD / K::Q36_KV_Q8_GROUP; + const int g = i / K::Q36_KV_Q8_GROUP; + float ak = fabsf(v), av = fabsf(vv); + for (int o = 16; o > 0; o >>= 1) { + ak = fmaxf(ak, __shfl_xor_sync(0xffffffffu, ak, o)); + av = fmaxf(av, __shfl_xor_sync(0xffffffffu, av, o)); + } + const uint16_t sk = K::q36_kv_q8_scale_bits(ak); + const uint16_t sv = K::q36_kv_q8_scale_bits(av); + const size_t sbase = ((size_t) kh * ctx + pos) * ng; + if ((i & (K::Q36_KV_Q8_GROUP - 1)) == 0) { + kcs[sbase + g] = sk; + vcs[sbase + g] = sv; + } + kcq[((size_t) kh * ctx + pos) * HD + i] = K::q36_kv_q8_quantize(v, sk); + vcq[((size_t) kh * ctx + pos) * HD + i] = K::q36_kv_q8_quantize(vv, sv); + } +} + /// Split-K attention, one block per (chunk of 256 cells, query row): query row qi = t * H + h (t from t0) sees the /// cells [0, *pos0 + t]. Writes the chunk's partial accumulator, max and sum; chunks past the query's cells write /// nothing (the merge stops before them). grid (n_chunk, n_rows), block 256, HD == 256. @@ -308,6 +362,48 @@ __global__ void k_attn_part(const float* q, const float* kc, const float* vc, co pl[o] = l; } } +/// int8 KV variant of k_attn_part (docs/Q36_INT8_KV.md): identical split-K structure, the K/V reads dequantized +/// on the fly from q8_0 codes + fp16 scales. +__global__ void k_attn_part_q8(const float* q, const int8_t* kcq, const int8_t* vcq, const uint16_t* kcs, + const uint16_t* vcs, const int32_t* pos0, int t0, int H, int HKV, int ctx, float scale, + int n_chunk, float* pacc, float* pm, float* pl) { + constexpr int HD = 256; + constexpr int ng = HD / K::Q36_KV_Q8_GROUP; + __shared__ float qs[HD], sc[kChunkCells], sh[kThreads]; + const int c = blockIdx.x, row = blockIdx.y; + const int t = t0 + row / H, h = row % H, kvh = h / (H / HKV); + const int npos = *pos0 + t + 1; + const int p0 = c * kChunkCells; + if (p0 >= npos) return; + const int j = threadIdx.x; + qs[j] = q[((size_t) t * H + h) * HD + j]; + __syncthreads(); + const int p = p0 + j; + float s = -INFINITY; + if (p < npos) { + const int8_t* kp = kcq + ((size_t) kvh * ctx + p) * HD; + const uint16_t* sp = kcs + ((size_t) kvh * ctx + p) * ng; + float a = 0.f; + for (int i = 0; i < HD; ++i) a += qs[i] * K::q36_kv_q8_dequantize(kp[i], sp[i / K::Q36_KV_Q8_GROUP]); + s = a * scale; + } + const float m = block_max(s, sh); + const float e = p < npos ? expf(s - m) : 0.f; + sc[j] = e; + const float l = block_sum(e, sh); + const int n = min(kChunkCells, npos - p0); + float acc = 0.f; + const int8_t* vp = vcq + ((size_t) kvh * ctx + p0) * HD + j; + const uint16_t* vsp = vcs + ((size_t) kvh * ctx + p0) * ng + (j / K::Q36_KV_Q8_GROUP); + for (int k = 0; k < n; ++k) + acc += sc[k] * K::q36_kv_q8_dequantize(vp[(size_t) k * HD], vsp[(size_t) k * ng]); + const size_t o = (size_t) row * n_chunk + c; + pacc[o * HD + j] = acc; + if (j == 0) { + pm[o] = m; + pl[o] = l; + } +} /// The chunks merged, times sigmoid(gate): grid (n_rows), block HD. __global__ void k_attn_merge(const float* pacc, const float* pm, const float* pl, const float* gate, const int32_t* pos0, int t0, int H, int n_chunk, float* out) { @@ -405,6 +501,23 @@ __global__ void k_kv_f16(const float* kc, const float* vc, __half* k16, __half* k16[i] = __float2half(kc[src]); v16[i] = __float2half(vc[src]); } +/// int8 KV variant of k_kv_f16 (docs/Q36_INT8_KV.md): dequantize q8_0 K/V and stage them as FP16 for the prompt's +/// cuBLAS GEMMs. The dequantized value goes through f16_from_f32 (round-to-nearest-even), not __float2half. +__global__ void k_kv_f16_q8(const int8_t* kcq, const int8_t* vcq, const uint16_t* kcs, const uint16_t* vcs, + __half* k16, __half* v16, int ctx, int P, int HD, int HKV) { + const int64_t i = (int64_t) blockIdx.x * blockDim.x + threadIdx.x; + if (i >= (int64_t) HKV * P * HD) return; + const int64_t per = (int64_t) P * HD; + const int64_t kvh = i / per, rest = i % per; + const int64_t pos = rest / HD, dim = rest % HD; + const int ng = HD / K::Q36_KV_Q8_GROUP; + const size_t src = (size_t) kvh * ctx * HD + rest; + const size_t sbase = (size_t) kvh * ctx * ng + (size_t) pos * ng + dim / K::Q36_KV_Q8_GROUP; + const float kv = K::q36_kv_q8_dequantize(kcq[src], kcs[sbase]); + const float vv = K::q36_kv_q8_dequantize(vcq[src], vcs[sbase]); + k16[i] = __ushort_as_half(K::f16_from_f32(kv)); + v16[i] = __ushort_as_half(K::f16_from_f32(vv)); +} /// The prompt's causal softmax: row (h, t) of S (ld floats, P cells) -> softmax(S * scale) over the cells /// [0, pos0 + t], 0 past them, as FP16 probabilities (S16, the same row offsets, ld halves). grid (T, H), block 256. __global__ void k_attn_softmax16(const float* S, __half* S16, int ld, int T, int P, int pos0, float scale) { @@ -1050,6 +1163,9 @@ struct Engine { int32_t* ami; int n_chunk = 0, att_rows = 0; // attention partials: att_rows query rows x n_chunk chunks std::vector conv, S, kc, vc; // per layer; the MTP block's K / V at index n_layers + bool kv_int8 = false; // --kv int8: the K/V cache is q8_0 (int8 codes + fp16 scales) instead of fp32 + std::vector kcq, vcq; // q8_0 K/V codes, indexed like kc/vc (nullptr for GDN layers) + std::vector kcs, vcs; // q8_0 K/V scales: [kvh * ctx + pos] * (HD / Q36_KV_Q8_GROUP) + group // verify windows (T <= kVerify): the GDN inputs of every row per layer, kept for the commit static constexpr int kVerify = 8; std::vector v_qkv, v_hc, v_alpha, v_beta; @@ -1228,14 +1344,31 @@ struct Engine { if (L.attn) { conv.push_back(nullptr); S.push_back(nullptr); - kc.push_back(dalloc((size_t) g.n_head_kv * ctx * g.head_dim)); - vc.push_back(dalloc((size_t) g.n_head_kv * ctx * g.head_dim)); + if (kv_int8) { + kc.push_back(nullptr); + vc.push_back(nullptr); + kcq.push_back(dalloc_t((size_t) g.n_head_kv * ctx * g.head_dim)); + vcq.push_back(dalloc_t((size_t) g.n_head_kv * ctx * g.head_dim)); + kcs.push_back(dalloc_t((size_t) g.n_head_kv * ctx * (g.head_dim / K::Q36_KV_Q8_GROUP))); + vcs.push_back(dalloc_t((size_t) g.n_head_kv * ctx * (g.head_dim / K::Q36_KV_Q8_GROUP))); + } else { + kc.push_back(dalloc((size_t) g.n_head_kv * ctx * g.head_dim)); + vc.push_back(dalloc((size_t) g.n_head_kv * ctx * g.head_dim)); + kcq.push_back(nullptr); + vcq.push_back(nullptr); + kcs.push_back(nullptr); + vcs.push_back(nullptr); + } for (auto* v : {&v_qkv, &v_hc, &v_alpha, &v_beta}) v->push_back(nullptr); } else { conv.push_back(dalloc(C * (g.ssm_d_conv - 1))); S.push_back(dalloc((size_t) g.ssm_v_heads * g.ssm_state_size * g.ssm_state_size)); kc.push_back(nullptr); vc.push_back(nullptr); + kcq.push_back(nullptr); + vcq.push_back(nullptr); + kcs.push_back(nullptr); + vcs.push_back(nullptr); v_qkv.push_back(dalloc((size_t) kVerify * C)); v_hc.push_back(dalloc((size_t) kVerify * C)); v_alpha.push_back(dalloc((size_t) kVerify * g.ssm_v_heads)); @@ -1245,8 +1378,21 @@ struct Engine { if (m->has_mtp) { // the MTP block's own K / V (an attention layer at index n_layers) conv.push_back(nullptr); S.push_back(nullptr); - kc.push_back(dalloc((size_t) g.n_head_kv * ctx * g.head_dim)); - vc.push_back(dalloc((size_t) g.n_head_kv * ctx * g.head_dim)); + if (kv_int8) { + kc.push_back(nullptr); + vc.push_back(nullptr); + kcq.push_back(dalloc_t((size_t) g.n_head_kv * ctx * g.head_dim)); + vcq.push_back(dalloc_t((size_t) g.n_head_kv * ctx * g.head_dim)); + kcs.push_back(dalloc_t((size_t) g.n_head_kv * ctx * (g.head_dim / K::Q36_KV_Q8_GROUP))); + vcs.push_back(dalloc_t((size_t) g.n_head_kv * ctx * (g.head_dim / K::Q36_KV_Q8_GROUP))); + } else { + kc.push_back(dalloc((size_t) g.n_head_kv * ctx * g.head_dim)); + vc.push_back(dalloc((size_t) g.n_head_kv * ctx * g.head_dim)); + kcq.push_back(nullptr); + vcq.push_back(nullptr); + kcs.push_back(nullptr); + vcs.push_back(nullptr); + } mtp_cat = dalloc(T * 2 * N); hmtp = dalloc(T * N); logits_m = dalloc((size_t) g.n_vocab); @@ -1378,8 +1524,13 @@ struct Engine { apply(L.wq, xq, qg, T); apply(L.wk, xq, kx, T); apply(L.wv, xq, vx, T); - k_attn_prep<<>>(qg, kx, vx, L.q_norm, L.k_norm, q, gate, kc[li], vc[li], H, - HKV, HD, m->n_rot, dpos, ctx, m->rope_base, m->eps); + if (kv_int8) + k_attn_prep_q8<<>>(qg, kx, vx, L.q_norm, L.k_norm, q, gate, kcq[li], + vcq[li], kcs[li], vcs[li], H, HKV, HD, m->n_rot, + dpos, ctx, m->rope_base, m->eps); + else + k_attn_prep<<>>(qg, kx, vx, L.q_norm, L.k_norm, q, gate, kc[li], vc[li], H, + HKV, HD, m->n_rot, dpos, ctx, m->rope_base, m->eps); if (T > kGemmT) { // a prompt chunk: scores = Q K^T and out = P V as FP16 GEMMs (FP32 accumulation, FP32 scores) per KV // head (its query heads one strided batch over the shared K / V), the causal softmax in between writing @@ -1389,7 +1540,11 @@ struct Engine { const int tq = (int) std::max(1, std::min(T, (int64_t) (kAttScoreFloats / ((size_t) H * ldp)))); const int G = H / HKV; const float one = 1.0f, zero = 0.0f; - k_kv_f16<<>>(kc[li], vc[li], k16, v16, ctx, P, HD, HKV); + if (kv_int8) + k_kv_f16_q8<<>>(kcq[li], vcq[li], kcs[li], vcs[li], k16, + v16, ctx, P, HD, HKV); + else + k_kv_f16<<>>(kc[li], vc[li], k16, v16, ctx, P, HD, HKV); k_to_f16<<>>(q, q16, (int64_t) T * H * HD); for (int t0 = 0; t0 < T; t0 += tq) { const int nt = std::min(tq, T - t0); @@ -1429,8 +1584,13 @@ struct Engine { const int nt = std::min(rows_per, T - t0); const int nc = graph_mode ? n_chunk : (int) std::min(n_chunk, (p0 + t0 + nt + kChunkCells - 1) / kChunkCells); - k_attn_part<<>>(q, kc[li], vc[li], dpos, t0, H, HKV, ctx, - 1.0f / sqrtf((float) HD), n_chunk, pacc, pm, pl); + if (kv_int8) + k_attn_part_q8<<>>(q, kcq[li], vcq[li], kcs[li], vcs[li], dpos, + t0, H, HKV, ctx, 1.0f / sqrtf((float) HD), + n_chunk, pacc, pm, pl); + else + k_attn_part<<>>(q, kc[li], vc[li], dpos, t0, H, HKV, ctx, + 1.0f / sqrtf((float) HD), n_chunk, pacc, pm, pl); k_attn_merge<<>>(pacc, pm, pl, gate, dpos, t0, H, n_chunk, att); } } @@ -2349,7 +2509,7 @@ int strata::qwen35::run(int argc, char** argv) { double expert_vram = -1.0; bool warm = true; int expert_slots = 0; - bool serve_mode = false, graph = true; + bool serve_mode = false, graph = true, kv_int8 = false; for (int i = 1; i < argc; ++i) { const std::string a = argv[i]; auto next = [&]() -> std::string { @@ -2385,6 +2545,12 @@ int strata::qwen35::run(int argc, char** argv) { else if (a == "--cpu-threads") cpu_threads = std::atoi(next().c_str()); else if (a == "--expert-slots") expert_slots = std::atoi(next().c_str()); else if (a == "--hybrid-rows") hybrid_rows = std::atoi(next().c_str()); + else if (a == "--kv") { + const std::string v = next(); + if (v == "int8") kv_int8 = true; + else if (v == "fp32") kv_int8 = false; + else { std::fprintf(stderr, "strata-q36: --kv takes int8 or fp32\n"); return 2; } + } else if (serve_mode && a.rfind("--", 0) == 0) { // serve passes the config's args plus a few of its own: one this engine does not have is said and skipped // (with its value), not fatal @@ -2395,6 +2561,7 @@ int strata::qwen35::run(int argc, char** argv) { } else { std::fprintf(stderr, "usage: strata-q36 --model GGUF [--max-context 32768] [--chunk 512] [--no-graph]\n" + " [--kv int8|fp32]\n" " (--serve | --tokens IDS [--n 64] [--force IDS] [--dump-full-logits F] | --bench-pp N " "--bench-tg N)\n"); return 2; @@ -2421,6 +2588,7 @@ int strata::qwen35::run(int argc, char** argv) { } Engine E; E.use_graph = graph; + E.kv_int8 = kv_int8; if (chunk == 0) { // auto: 512 with every expert in VRAM; 2048 when they will not all fit (a prompt chunk streams every layer's // experts over PCIe once, so a bigger chunk pays that once for more tokens) From 3a9e6c888a645331a0bf2179897f43cdaebf43ba Mon Sep 17 00:00:00 2001 From: George Georgovassilis Date: Fri, 16 Oct 2026 11:35:31 +0300 Subject: [PATCH 2/2] Qwen3.6: opt-in single staging buffer and staging placement fixes - STRATA_Q36_SINGLE_STAGE=1 halves the streamed-prompt staging VRAM (one buffer instead of two double-buffered) - STRATA_Q36_NO_MMQ=1 is now numeric (MMQ stays on by default) - STRATA_Q36_PLACE_DEBUG logs the staging allocation - shrink the attention-scores buffer (256 -> 64 MB) when the staging buffers only just miss, so placement still succeeds Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- docs/DETAILS.md | 48 +++++++++++++++++++++++++++++++++++++ docs/Q36_INT8_KV.md | 27 +++++++++++++++++++-- src/qwen35/q36_engine.cu | 51 ++++++++++++++++++++++++++++++---------- 3 files changed, 112 insertions(+), 14 deletions(-) diff --git a/docs/DETAILS.md b/docs/DETAILS.md index 5bd5fbbbd48..6f99c7cf95e 100644 --- a/docs/DETAILS.md +++ b/docs/DETAILS.md @@ -5,6 +5,7 @@ New here? Start with the [README](../README.md); installing step by step is in [ [MODELS.md](MODELS.md), common problems in [TROUBLESHOOTING.md](TROUBLESHOOTING.md). > **On this page:** [Speed](#speed-measured) · [Other GPUs](#other-gpus-estimated) · [Which model?](#which-model) · +> [Qwen3.6-35B-A3B](#qwen36-35b-a3b-its-own-engine-strata-q36) · > [Requirements](#before-you-start) · [Windows](#windows) · [Linux](#linux) · [API](#using-it) · > [MCP tools](#tools-from-mcp-servers) · [MCP server](#manage-strata-from-your-ai-assistant-mcp-server) · > [Images](#images-vision) · @@ -421,6 +422,53 @@ tokens/s, most experts read from the SSD; it picks the same tokens as llama.cpp the positions of short greedy answers, 90-91% after a 16K prompt, differing mostly at near-ties ([measured](UNSLOTH_Q4.md#quality-against-llamacpp-on-the-same-file)). +## Qwen3.6-35B-A3B (its own engine, `strata-q36`) + +A 35B mixture of experts with 3B active (256 experts, top-8), a smaller cousin of Qwen3.8-Flash-Next with the same +tokenizer. It runs on an engine path of its own (`src/qwen35`, the `strata-q36` path; `strata` switches to it from the +file's `qwen35moe` architecture), not through the pack/arena path the models above use. Setup offers it as +`--family qwen36` (models `UD-Q4_K_XL`, `UD-Q6_K_XL`, `Q8_0`). NVIDIA (CUDA) and AMD (HIP) only; the Intel Arc (SYCL) +engine refuses it. It reads the GGUF as it is (no pack, no expert profile, no separate MTP draft layer), keeps every +expert it can in VRAM and computes the rest on the CPU pool (the hybrid mode), and drafts with the file's own MTP +block. + +**Engine arguments** (the server config's `args`, or `ENGINE_ARGS_EXTRA` in the container): + +- `--kv int8|fp32` (default `fp32`, byte-identical): the K/V cache as q8_0 (int8 codes + fp16 scale) instead of fp32, + ~11 KiB/token instead of 44. Measured on an RTX 4060 Max-Q Laptop (8 GB), int8 starts and decodes at **32,768** tokens + where fp32 stops at 16,384 — 2x the context window. The full design note and numbers: + [Q36_INT8_KV.md](Q36_INT8_KV.md). +- `--mtp N` (0..7): how many drafts the file's own MTP block offers per step; setup writes 2. +- `--max-context N` (or `--ctx N`), `--chunk N` (or `--prefill auto|N`): the context length and the prompt chunk. + `auto` is 512 when every expert fits in VRAM, 2048 in the hybrid mode. +- `--expert-vram N`: a GiB budget for the routed experts, instead of the free VRAM less a reserve. +- `--expert-slots N`, `--hybrid-rows N`, `--pool-workers N` (or `--cpu-threads N`): the hybrid mode's per-layer VRAM + slots, its batch rows and its CPU pool threads (used by the tests and for tuning). +- `--bench-pp N --bench-tg N`: llama-bench's shapes — a prompt of N tokens, then N generated one at a time. + +**Switches** (`STRATA_Q36_*` in the environment; `docker/run.sh` and the server forward them unchanged): + +- `STRATA_Q36_NO_MMQ=1` turns the MMQ prompt path off (on by default). The prompt then has no streamed path and every + batch row runs on the CPU pool: measured ~20x slower prefill (310.1 vs 16.6 tok/s at pp512, 32K, `--kv int8`, the 8 GB + card), and decode is a wash (this model's 18.75 GiB of experts fit in 8 GiB neither way). **Do not reach for it to work + around a staging `out of memory`** — that is what `STRATA_Q36_SINGLE_STAGE=1` is for. +- `STRATA_Q36_SINGLE_STAGE=1`: one expert-streaming staging buffer instead of two, halving the prompt staging VRAM + (1.10 GiB -> 0.55 GiB). Together with `--kv int8` it is what lets 32K start on the 8 GB card. +- `STRATA_Q36_PLACE_DEBUG=1`: the engine prints its expert-placement arithmetic (the chunk rule, the staging + allocation, the free VRAM at each step) — to see why a context starts or does not, not to time. +- `STRATA_Q36_NO_MTP_CPU=1`: the plain hybrid mode for every layer. By default the main layers' experts can all sit in + VRAM while the MTP block's go to slots + the CPU pool (its drafts only). +- `STRATA_Q36_KEEP_PACKED=1`: keep the packed Q8_0 decode copies even when giving them back would let every routed + expert fit in VRAM. The default drops them — the decode then reads the GGUF layout, bitwise the same outputs, a + little slower — instead of running the hybrid mode. +- `STRATA_Q36_STREAM_ALL=1`: stream a whole layer's experts from host memory in one copy (the first version, + PCIe-bound); the default copies runs of consecutive experts from the same source. +- `STRATA_Q36_BLAS_F32`: use the cuBLAS / hipBLAS SGEMM for the FP32 matmuls (the router, the GDN alpha/beta) instead + of the split-K kernel. +- `STRATA_Q36_FIXED_SLOTS`: keep the expert slots as placed, never rebalance them (tests). + +The default path — `--kv fp32`, MMQ on, two staging buffers — stays byte-identical to the last release. + ## Before you start You need **only a graphics driver**: NVIDIA 580 or newer (update it with the NVIDIA App or from diff --git a/docs/Q36_INT8_KV.md b/docs/Q36_INT8_KV.md index d7a3d0b77f1..cd7a44cbb91 100644 --- a/docs/Q36_INT8_KV.md +++ b/docs/Q36_INT8_KV.md @@ -174,8 +174,9 @@ Implemented behind `--kv int8` (default `--kv fp32`, byte-identical). Measured o 2. **Parity test — pass.** `q36_kv_q8_parity` passes in `ctest`. 3. **Memory — relaxed.** The 131,072 target is not reachable, because at int8 the K/V is no longer the binding constraint. The ceiling is set by the dense weights (4.40 GiB) plus the - hybrid expert-streaming staging buffer (1.10 GiB) plus the context-scaled attention-score - buffer (see the last risk below), not by K/V. Measured: fp32 starts at 16,384 and no + hybrid expert-streaming staging buffer (1.10 GiB double-buffered, or 0.55 GiB with + `STRATA_Q36_SINGLE_STAGE=1`) plus the context-scaled attention-score buffer (see the last + risk below), not by K/V. Measured: fp32 starts at 16,384 and no larger; int8 starts and decodes at **32,768** (40,960 allocates but the prefill GEMM's score buffer no longer fits, 49,152 fails `cudaMalloc`). Resident K/V is 44 KiB/token (fp32) vs ~11 KiB/token (int8): at 4,096 context the engine reports 1.70 GiB free (fp32) vs @@ -192,6 +193,28 @@ Implemented behind `--kv int8` (default `--kv fp32`, byte-identical). Measured o 6. **Backends — CUDA verified.** The HIP build of the touched file has not been run yet; it is the remaining sign-off for this change (the engine is shared CUDA/HIP source). +A staging-buffer fix came out of measuring item 3. At 32,768 the engine reported "0.63 GiB of +VRAM free" after `init()`, and the prompt staging (one layer's experts, 0.55 GiB) plus the +streamed path's `ecpu` buffer (0.13 GiB at chunk 2048) did not fit in it, so MMQ aborted with +`cudaMalloc(&E.stage[b], ...) failed: out of memory`. Shortening the context did not help: +int8 K/V is ~11.7 KiB per token, so 32,768 -> 22,768 frees only ~0.16 GiB (0.63 -> 0.79 GiB +free) against a 0.55 GiB buffer. `place_experts()` now makes the same trade there that the +MTP-only path already made one branch above — the attention-score buffer at 64 MB instead of +256, which frees 0.28 GiB (0.63 -> 0.91 GiB free), after which staging fits. **Do not reach for +`STRATA_Q36_NO_MMQ=1` instead:** with MMQ off the prompt has no streamed path, `hybrid_t` +becomes `tmax`, and *every* prompt row goes to the CPU pool. Same machine, same run +(`--bench-pp 512 --bench-tg 8`, 32K, `--kv int8`): + +| `STRATA_Q36_NO_MMQ` | prefill pp512 | decode tg8 | batch entries on the GPU | +|---|---|---|---| +| `0` (MMQ on) | **310.1 / 310.4 / 310.3 tok/s** | 8.6 - 11.1 tok/s | 0 of 7,680 (0.0%) | +| `1` (MMQ off) | 16.6 / 15.7 / 11.5 tok/s | 8.0 - 13.4 tok/s | 0 of 499,200 (0.0%) | + +~20x on prefill, and the decode column is a wash: this model's 18.75 GiB of experts fit in the +card's 8 GiB neither way (0 slots either way), so decode is the CPU pool in both. Nothing sets +this switch for you: `docker/run.sh` forwards `STRATA_*` variables unchanged and `setup.py` does +not set MMQ, so the engine's own default (MMQ on) applies unless you set it yourself. + ## Risks and open questions - **Accuracy.** q8_0 K/V is lossy. The fp32→fp16 step is the smaller-risk intermediate if the diff --git a/src/qwen35/q36_engine.cu b/src/qwen35/q36_engine.cu index 04c36f01f67..8ef4a374acd 100644 --- a/src/qwen35/q36_engine.cu +++ b/src/qwen35/q36_engine.cu @@ -1194,7 +1194,8 @@ struct Engine { // ---- the hybrid mode (Model::hybrid): per state index (the layers, then the MTP block) the slot tables; a // batch of up to hybrid_t rows runs its VRAM-resident experts on the GPU and the others on the CPU pool, side by - // side; a bigger prompt chunk streams each layer's experts into a staging buffer for MMQ (two, double-buffered) + // side; a bigger prompt chunk streams each layer's experts into a staging buffer for MMQ (double-buffered by + // default: one layer streams while the other computes; STRATA_Q36_SINGLE_STAGE=1 halves it to one buffer) struct Hy { std::vector slot_of, exp_in_slot; std::vector score; // routing counts, halved at each rebalance @@ -1222,6 +1223,7 @@ struct Engine { uint8_t* acts = nullptr; // per token: its quantized activation (kNativeActBytes) cudaEvent_t ev_route = nullptr; uint8_t* stage[2] = {}; + int n_stage = 2; // staging buffers: 2 double-buffered, or 1 (STRATA_Q36_SINGLE_STAGE=1) to halve the VRAM cudaStream_t cst = nullptr; cudaEvent_t ev_copy[2] = {}, ev_used[2] = {}; std::vector stream_seq; // the state indices a streamed forward visits, in order @@ -1281,7 +1283,10 @@ struct Engine { k16 = dalloc_t<__half>((size_t) g.n_head_kv * ctx * g.head_dim); v16 = dalloc_t<__half>((size_t) g.n_head_kv * ctx * g.head_dim); q16 = dalloc_t<__half>(T * g.n_head * g.head_dim); - use_mmq = MMQ::built() && std::getenv("STRATA_Q36_NO_MMQ") == nullptr; + const char* nommq = std::getenv("STRATA_Q36_NO_MMQ"); + use_mmq = MMQ::built() && (nommq == nullptr || std::atoi(nommq) == 0); // MMQ is on by default; STRATA_Q36_NO_MMQ=1 turns it off + const char* singlestage = std::getenv("STRATA_Q36_SINGLE_STAGE"); + n_stage = (singlestage && std::atoi(singlestage) > 0) ? 1 : 2; // one staging buffer instead of two: halves the prompt staging VRAM auto mmq_ok = [&](const Layer& L) { return MMQ::supported(L.L.gu_type) && MMQ::supported(L.L.d_type) && MMQ::fits(L.L.gu_type, 2 * g.n_ff) && MMQ::fits(L.L.d_type, g.n_embd); @@ -1727,14 +1732,14 @@ struct Engine { // ---- the hybrid mode const Layer& layer_at(int li) const { return li < (int) m->layers.size() ? m->layers[(size_t) li] : m->mtp; } - /// A streamed forward over the state indices `seq`: the first two layers' experts start copying now. + /// A streamed forward over the state indices `seq`: the first n_stage layers' experts start copying now. void stream_begin(const std::vector& seq) { stream_seq = seq; stream_pos = 0; - for (int i = 0; i < 2 && i < (int) seq.size(); ++i) stream_kick(i); + for (int i = 0; i < n_stage && i < (int) seq.size(); ++i) stream_kick(i); } - void stream_kick(int i) { // the i-th layer of the sequence into staging buffer i % 2, after its last reader - const int b = i % 2; + void stream_kick(int i) { // the i-th layer of the sequence into staging buffer i % n_stage, after its last reader + const int b = i % n_stage; const Layer& L = layer_at(stream_seq[(size_t) i]); CK(cudaStreamWaitEvent(cst, ev_used[b], 0)); // the experts with a VRAM slot come from it (device to device, at VRAM speed), only the others over PCIe; runs @@ -1772,14 +1777,14 @@ struct Engine { std::fprintf(stderr, "strata-q36: the expert stream is out of step (layer %d)\n", li); std::exit(1); } - const int b = stream_pos % 2; + const int b = stream_pos % n_stage; CK(cudaStreamWaitEvent(st, ev_copy[b], 0)); return stage[b]; } void stream_done() { - const int b = stream_pos % 2; + const int b = stream_pos % n_stage; CK(cudaEventRecord(ev_used[b], st)); - if (stream_pos + 2 < (int) stream_seq.size()) stream_kick(stream_pos + 2); + if (stream_pos + n_stage < (int) stream_seq.size()) stream_kick(stream_pos + n_stage); ++stream_pos; } @@ -2090,7 +2095,8 @@ struct Engine { // ------------------------------------------------------------------------------------------------ expert placement /// The routed experts after the engine's buffers: all in VRAM when they fit in `vram_gib` (< 0: the free VRAM less a /// reserve), else the hybrid mode - every expert in host memory, the same number of VRAM slots per layer, a decode -/// batch's other experts on the CPU pool and a prompt chunk's streamed through two staging buffers. +/// batch's other experts on the CPU pool and a prompt chunk's streamed through staging buffers (double-buffered, or +/// one with STRATA_Q36_SINGLE_STAGE=1). bool place_experts(Model& m, Engine& E, double vram_gib, int threads, int hybrid_t, int slots, std::string& err) { const auto t0 = Clock::now(); const auto& g = m.g; @@ -2190,7 +2196,21 @@ bool place_experts(Model& m, Engine& E, double vram_gib, int threads, int hybrid E.mtp_graph = false; // the MTP block's layer is hybrid either way E.hybrid_t = std::max(1, std::min(hybrid_t, E.use_mmq ? E.tmax : E.tmax)); if (!E.use_mmq) E.hybrid_t = E.tmax; // no streamed prompt path: every batch on slots + the pool - const size_t staging = E.use_mmq && E.main_hybrid ? 2 * (max_layer + kMmqTail) : 0; + const size_t staging = E.use_mmq && E.main_hybrid ? (size_t) E.n_stage * (max_layer + kMmqTail) : 0; + // the streamed prompt path needs the staging buffer(s) plus E.ecpu (allocated below) to fit; where they only just + // do not, shrink the attention scores buffer (256 -> 64 MB, the prompt attention then runs in more, smaller + // slices) so the staging fits - the same trade the MTP-only path above makes. + if (staging > 0 && E.kAttScoreFloats > ((size_t) 16 << 20)) { + const size_t need = staging + (size_t) E.tmax * (size_t) m.top_k * g.n_embd * 4 + ((size_t) 64 << 20); + if (need > fr) { + E.shrink_att((size_t) 16 << 20); + CK(cudaMemGetInfo(&fr, &tot)); + budget2 = vram_gib >= 0 ? budget : (fr > reserve ? fr - reserve : 0); + std::fprintf(stderr, "strata-q36: the prompt staging (%.2f GiB) did not fit: the attention scores buffer at " + "64 MB instead of 256 (more, smaller slices in a long prompt), %.2f GiB free now\n", + (double) staging / (1 << 30), (double) fr / (1 << 30)); + } + } size_t slot_bytes = budget2 > staging ? budget2 - staging : 0; if (slot_bytes + staging + ((size_t) 256 << 20) > fr) slot_bytes = fr > staging + ((size_t) 512 << 20) ? fr - staging - ((size_t) 512 << 20) : 0; // --expert-slots: that many per layer (tests: the same split with and without the MTP block) @@ -2264,9 +2284,16 @@ bool place_experts(Model& m, Engine& E, double vram_gib, int threads, int hybrid CK(cudaMalloc(&E.d_miss, T * Kp * 4)); E.acts = (uint8_t*) std::malloc(T * C::kNativeActBytes + 64); CK(cudaEventCreateWithFlags(&E.ev_route, cudaEventDisableTiming)); + if (std::getenv("STRATA_Q36_PLACE_DEBUG")) { + size_t f2 = 0, t2 = 0; + CK(cudaMemGetInfo(&f2, &t2)); + std::fprintf(stderr, "strata-q36: staging alloc: n_stage=%d, max_layer=%.2f GiB (+%zu B), want=%.2f GiB, free=%.2f GiB\n", + E.n_stage, (double) max_layer / (1 << 30), kMmqTail, (double) staging / (1 << 30), + (double) f2 / (1 << 30)); + } if (staging) { CK(cudaStreamCreateWithFlags(&E.cst, cudaStreamNonBlocking)); - for (int b = 0; b < 2; ++b) { + for (int b = 0; b < E.n_stage; ++b) { CK(cudaMalloc(&E.stage[b], max_layer + kMmqTail)); CK(cudaMemset(E.stage[b], 0, max_layer + kMmqTail)); CK(cudaEventCreateWithFlags(&E.ev_copy[b], cudaEventDisableTiming));