Skip to content

kv: add Q4_0 KV cache mode with FWHT-256 Hadamard rotation - #21

Merged
Niko1221 merged 1 commit into
Niko1221:mainfrom
code-martin:feature/kv-cache-q4_0
Sep 27, 2026
Merged

Niko1221 merged 1 commit into
Niko1221:mainfrom
code-martin:feature/kv-cache-q4_0

Conversation

@code-martin

@code-martin code-martin commented Sep 27, 2026 •

Copy link
Copy Markdown
Contributor

What

Adds an optional 4-bit (Q4_0) quantized KV cache mode with Fast Walsh-Hadamard Transform (FWHT-256) rotation to Strata:

  • Activated via --kv q4_0 flag in the engine or server configuration.
  • Halves KV cache memory footprint compared to Q8_0 (down to ~4.5 bytes per token per layer per head), unlocking double context capacity on memory-constrained systems (e.g. 256k context).
  • Implements fast in-register FWHT-256 rotation kernel before quantization to eliminate activation outliers across the head dimension ($D=256$).
  • Maintains attention score invariance $\langle Hq, Hk \rangle = \langle q, k \rangle$ by applying the same orthogonal Hadamard matrix to query tokens.
  • Fully compatible with speculative verification (verify.cpp), multi-token prediction (mtp.cpp), and QSA decode streaming (qsa_decode_attn.cu).

Why

At 256k context, unquantized or 8-bit KV caches consume substantial memory. Standard 4-bit quantization on raw key-value vectors often suffers quality loss due to channel outliers. Applying an orthogonal Walsh-Hadamard transform distributes outlier energy uniformly across all 256 dimensions before block quantization, preserving perplexity with zero score drift.

How

  1. FWHT-256 Kernel (src/kernels/cuda/kv_q4.cu, include/strata/kernels/kv_q4.hpp):
    • In-register butterfly network across 8 warps per block with bitwise XOR shuffles.
    • Normalization factor $1/\sqrt{256} = 1/16$.
  2. Quantized Key/Value Store & Decode Attention (src/kernels/cuda/qsa_decode_attn.cu):
    • Q4_0 block layout: 1 FP16 scale + 16 bytes for 32 int4 nibbles.
    • Vectorized dequantization in warp registers during flash-attention score accumulation.
  3. Parity & Correctness Test (src/kernels/kv_q4_parity.cpp):
    • Verifies Hadamard self-inverse $H \cdot H \cdot x == x$.
    • Verifies dot-product conservation $\langle Hq, Hk \rangle == \langle q, k \rangle$.
    • Verifies quant/dequant parity and max error bounds.
  4. Tooling & Build:
    • Added kv_q4_parity target in CMakeLists.txt.
    • Fixed path resolution in tools/_paths.py, tools/pack_layer.py, tools/strata_pack.py when repo is placed near filesystem roots.

Verification

  • kv_q4_parity.exe passes 100% on CUDA (FWHT-256 bitwise match, attention score invariance, block bounds).
  • Tested end-to-end on Qwen3.8-Flash-Next 125B (swift-1.5-iq2_xs) with 256k context window via OpenAI /v1/chat/completions API.

- Implement Q4_0 KV quantization and dequantization kernels with scale/offset
- Implement fast Walsh-Hadamard transform (FWHT-256) rotation for outlier suppression
- Add Q4_0 attention decode support in qsa_decode_attn
- Add kv_q4_parity parity test verifying CPU vs CUDA numerical accuracy
- Support --kv q4_0 in generate CLI and runtime layer allocation
- Update .gitignore for build and repository hygiene
Niko1221 added a commit that referenced this pull request Sep 27, 2026
…with KV streaming; accurate expert-cache note (#23)

PR #21 (code-martin): a 4-bit K/V format (orthonormal FWHT-256 rotation of K, V and Q, then q4_0 blocks). On merge:
- Q4 is a full third format of the KV-streaming code (host copy, block residency, prompt-path staging, MTP ring,
  STATE_HASH): with --kv-resident the PR wrote K/V through the residency map's -1 entries. kv_stream_parity now
  checks Q4 streamed vs resident, bitwise.
- The per-token path ran attention twice (once before the blocks were resident); once now.
- The q4 group maximum breaks |x| ties the same way in every lane (one scale per block).
- Dropped: the duplicate path fixes (in 0.1.7), personal .gitignore entries, the re-encoded CMakeLists.
- Measured (bench/results/2026-09-27-kv-q4): int8 matches fp16; q4_0 costs precision on long documents
  (perplexity +8-12%), needles 5/5, +3.7% decode at 128K. Opt-in: setup asks above 8K context (int8 recommended),
  START-HERE.bat --kv q4_0, or --kv q4_0 in the config.

#23: the start-up warning "the GPU hit path is NOT CORRECT" dated from a long-fixed fault. Measured teacher-forced
on 2,557 tokens (bench/results/2026-09-27-cache-parity): cache on vs off, 95-98% same top-1 and equal perplexity
(-0.005 +- 0.005 nats); the warning is now one accurate line.

Co-Authored-By: code-martin <code-martin@users.noreply.github.com>
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0154bwiad8gY7wRehCXy27pN
@Niko1221
Niko1221 merged commit 89ba2dc into Niko1221:main Sep 27, 2026
@Niko1221

Copy link
Copy Markdown
Owner

Merged, thanks! Nice work on the FWHT kernel and the parity test. It's in v0.1.8 as --kv q4_0 (opt-in; int8 stays the default).

A few things on top, so it works with everything that landed since:

  • KV streaming: the Q4 blocks go through the same host pool / VRAM block cache as int8 (write-through, device-side resolve), and the MTP ring too.
  • Deterministic quantization: q4_group breaks ties on the max-abs element the same way every time (a > amax || (a == amax && v > mval)), so the conversation cache's restore/re-read and STRATA_STATE_HASH stay bit-exact.
  • The per-token path rotates Q and the output as well (it was only on the windowed paths). kv_stream_parity now covers int8, fp16 and q4.

Measured on Q2_0: document perplexity +8% at 1K and +12% at 8K vs int8, needles 5/5 with Q4 + streaming, +3.7% rounds/s at 128K, and KV RAM 0.83 vs 1.52 GiB. That's why it's a choice rather than the default. Setup asks above 8K context.

Your branch was rebased after I pulled it; GitHub shows it merged, and the code is the same as what shipped.

eddoursul referenced this pull request in eddoursul/Lamina Oct 1, 2026
…ernels

Upstream's formats (PR #21, PR Niko1221#120): each head's K and V rotated by the orthonormal Walsh-Hadamard transform, then
4-bit codes in blocks of 32 values with an FP16 scale each (ggml's q4_0 rounding, against the stored scale).  q4_0
stores K and V that way (576 bytes a token and QSA layer, INT8 1,056), k8v4 V only, with K in INT8 (816).  This
branch's form:
- the pools keep the codes and the scales apart, as the INT8 pools do, so the attention's fetch warps copy both in
  16-byte pieces (upstream's interleaved 18-byte blocks cannot be copied that way);
- the tensor-core attention (the decode's splits and the prompt path) is templated on K's and V's format: a 4-bit chunk
  goes through the same ldmatrix loads, its nibbles become signed bytes in registers, and the sums are scaled per
  32-value block; a 4-bit scale has its block's sign, so V's p * scale is signed fixed point (its top limb signed, as
  q's is); K's block scales are read a word per group and V's fragments built a nibble half at a time, so the 4-bit
  kernels spill no more than the INT8 ones;
- the rotations stay inside: the appends rotate K and V, the kernel rotates q into shared memory (q4_0) and the output
  back (the decode's merge, before the gate, or the prompt kernel's epilogue), so callers pass and get the model's
  vectors in every format;
- a KV format per QSA state (core::KvFormat) with the helpers the verify window, the prompt path and the prompt cache's
  stash use (qsa_attn_pools, qsa_kv_append_steps, qsa_kv_arrays); the MTP draft layer keeps INT8; the canonical
  (non-native) decode path refuses the 4-bit formats.
verify_parity: the 4-bit appends bitwise a host reference of the rotation and rounding (both forms), the INT8 append
of K alone; the attention in all four formats against FP64 from the stored values (2.6e-6) and the decode against the
prompt form (1.5e-6); a round trip from float K/V (INT8 0.6%, the 4-bit formats 8.5% RMS on Gaussian data).  INT8's
window hashes are unchanged on both models; the prompt cache's regenerate and stash round trip reproduce the output in
both 4-bit formats (--adapt-every 0).

Measured (RTX 3090 + RTX 5070 Ti, 200K context): the 3090's cache holds IQ3_S 7,546 / 7,835 / 8,116 experts (INT8,
k8v4, q4_0), UD-Q4_K_XL 4,539 / 4,722 / 4,908; generation on fixed text k8v4 0.1-1.8% less time a pass (UD after 16K
1.5% more), q4_0 0.4-1.9% less; the prompt's attention 4.7 / 3.5 / 3.3 ms for 2,048 queries of 2,051 cells.  The
distributions along llama.cpp's texts 1.3-3.5x as far from llama.cpp's as INT8's (mean KL, IQ3_S thinking 2.1e-3 /
4.7e-3 / 7.4e-3).  The configs keep INT8.

Claude-Session: https://claude.ai/code/session_01GgM6qCebHZfQmgkcxNY1U5
MistVVK referenced this pull request in MistVVK/XeStrata Oct 3, 2026
…with KV streaming; accurate expert-cache note (#23)

PR #21 (code-martin): a 4-bit K/V format (orthonormal FWHT-256 rotation of K, V and Q, then q4_0 blocks). On merge:
- Q4 is a full third format of the KV-streaming code (host copy, block residency, prompt-path staging, MTP ring,
  STATE_HASH): with --kv-resident the PR wrote K/V through the residency map's -1 entries. kv_stream_parity now
  checks Q4 streamed vs resident, bitwise.
- The per-token path ran attention twice (once before the blocks were resident); once now.
- The q4 group maximum breaks |x| ties the same way in every lane (one scale per block).
- Dropped: the duplicate path fixes (in 0.1.7), personal .gitignore entries, the re-encoded CMakeLists.
- Measured (bench/results/2026-09-27-kv-q4): int8 matches fp16; q4_0 costs precision on long documents
  (perplexity +8-12%), needles 5/5, +3.7% decode at 128K. Opt-in: setup asks above 8K context (int8 recommended),
  START-HERE.bat --kv q4_0, or --kv q4_0 in the config.

#23: the start-up warning "the GPU hit path is NOT CORRECT" dated from a long-fixed fault. Measured teacher-forced
on 2,557 tokens (bench/results/2026-09-27-cache-parity): cache on vs off, 95-98% same top-1 and equal perplexity
(-0.005 +- 0.005 nats); the warning is now one accurate line.

Co-Authored-By: code-martin <code-martin@users.noreply.github.com>
Hugua700 pushed a commit to Hugua700/Strata that referenced this pull request Oct 6, 2026
MistVVK referenced this pull request in MistVVK/XeStrata Oct 10, 2026
…or pull request (#21); README: 2x CMP 170HX with Maya-M in the users' table (#22), STRATA_GLM_MTP_GGUF under Tuning, a link to CONTRIBUTING.md
FlareP1 pushed a commit to FlareP1/Strata that referenced this pull request Oct 13, 2026
…ay F, refused outside (0, 1)

The pull request's copy of Niko1221#2's commit resolved to this branch's 1c1f0fe; the decay on this branch's tier, with a
usage line.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
FlareP1 pushed a commit to FlareP1/Strata that referenced this pull request Oct 13, 2026
…ko1221#14, Niko1221#16, Niko1221#19, Niko1221#20 and Niko1221#21: their content is in the merges above

txiki739 rebased the series onto e3727cb after they were merged here (the same changes: Niko1221#19's residency uploads
through res_put, Niko1221#12 without its gr_parity commit).  Merged with -s ours: the tree is unchanged; the commit records the
pull requests as merged.

Claude-Session: https://claude.ai/code/session_01H6vhr2MCptzyCJCipfLCCM
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants