elizaos--eliza/packages/training/scripts/quantization/README.md

# Post-Training Quantization

This directory holds the post-training quantization passes used to
shrink the fine-tuned Eliza checkpoints before they leave the
training rig. Each method is a self-contained CLI; they are independent
and can be combined or compared on the same fine-tuned checkpoint.

> **Gemma 4 cutover note.** The eliza-1 base is now Gemma 4 (dense:
> alternating SWA/global, shared-KV, MQA, dual head dims 512/256, stock q8_0
> KV). Gemma geometry is the active release target. The shipping Gemma weight
> quant is stock llama.cpp `Q4_K_M` from `gguf-q4_k_m_apply.py`; TurboQuant
> and QJL are runtime KV-cache experiments, and PolarQuant is a separate
> weight-quant experiment. Those optional paths must be revalidated per tier
> before their sidecars can be cited as release provenance.

## PolarQuant

Reference: Caio Vicentino, *PolarQuant: Optimal Gaussian Weight
Quantization via Hadamard Rotation for LLM Compression*, arXiv:2603.29078
(March 2026). The arXiv PDF was withdrawn for errata fixes; the
implementation we run against lives in
[`caiovicentino/eoq-quantization`](https://github.com/caiovicentino/eoq-quantization)
@ commit `15a12160245d7d3015290c6c5b6dbb7f22094d5e`. The two source
files we depend on are vendored under `polarquant/` (see
`polarquant/LICENSE.md` for the upstream license situation — currently
no LICENSE file in the upstream repo; vendored under a documented
research-citation arrangement that must be revisited if the upstream
project publishes a restrictive license).

### Algorithm in one paragraph

For each linear's weight tensor, group the values into power-of-two
blocks (default 128). Per block: extract the L2 norm and store it as
fp16; normalize the block to the unit hypersphere; multiply by a
Walsh–Hadamard matrix so that each coordinate is an i.i.d. draw from
roughly N(0, 1/√d); quantize each coordinate against the **Lloyd–Max
optimal centroids for N(0, 1)** (precomputed via the iterative
conditional-expectation algorithm in `polarquant/polar_quant.py`),
storing the centroid index as int8; optionally append a 1-bit QJL
residual sign per block as a cheap correction. At inference, undo
each step: lookup centroid → divide by √d → inverse Hadamard → rescale
by per-block norm. The Hadamard rotation alone accounts for ~98% of
the perplexity recovery vs absmax (paper, Ablation §4.2). PolarQuant
is **data-free** — no calibration set is required — because the
post-rotation distribution is analytically Gaussian.

### Tradeoffs

- **Pros.** Data-free; near-lossless at Q5 (paper claims very small PPL
  deltas on decoder-only checkpoints vs FP16). int8 codes + fp16 per-block norms gives the
  storage payload that downstream INT4 inference kernels (torchao,
  llama.cpp, MLX) consume directly. Architecture-agnostic at the
  ``nn.Linear`` level.
- **Cons.** Reconstruction-back-into-fp16 (the path we use today) does
  **not** save inference VRAM by itself — the model still loads as a
  fp16 ``nn.Linear`` matrix; the win shows up when the codes-only
  payload is loaded by an INT4 inference kernel. The sidecar
  ``polarquant_artifacts.safetensors`` we emit is the artifact for
  that downstream path. The vendored compute kernel is a pure-Python
  PyTorch loop over linears (no Triton, no fused kernel), so the
  *quantization step itself* is slow on big models — it costs O(N)
  Hadamard products at quantization time, then is free at inference.

### Supported architectures

The vendored kernel runs on any model that exposes its weights as
``nn.Linear`` modules. We have explicitly verified the active path on:

- Gemma (`google/gemma-4-E2B`)
- Llama, Mistral, Phi-3 style decoder stacks by structural inspection

#### Gemma compatibility notes

PolarQuant operates on `nn.Linear` weights, so it quantizes the Q/K/V/O
and MLP projections that Gemma exposes through the HF model graph. Keep
these constraints in mind before adding a new tier:

1. Non-linear recurrent/state buffers, if present on a future hybrid
   tier, are **not** `nn.Linear` and must stay outside PolarQuant.
2. Vision-language variants must expose the text decoder before calling
   `quantize_checkpoint`; use the text config/model tower, not the vision
   encoder.
3. Future MoE router weights are tiny, fall under the `--min-numel`
   cutoff, and must be deliberately skipped when that line is reintroduced.

### CLI

Quantize a fine-tuned 2B checkpoint:

```bash
uv run python scripts/quantization/polarquant_apply.py \
    --model checkpoints/gemma4-e2b-eliza/final \
    --calibration data/final/val.jsonl \
    --calibration-samples 128 \
    --output checkpoints/gemma4-e2b-eliza/final-polarquant
```

The `--calibration*` flags are accepted for parity with the rest of
the quantization pipeline but are **unused** — PolarQuant is data-free
per the paper. Passing a non-existent calibration path still errors so
that misconfigured callers fail loudly.

Useful knobs:

- `--bits {2,3,4,5,6}` (default 4). Q5 is the paper's near-lossless
  point; Q4 is the practical default for INT4 inference kernels.
- `--block-size N` (default 128, must be a power of 2).
- `--no-qjl` to drop the 1-bit residual correction (slightly worse
  PPL, slightly smaller sidecar).
- `--no-artifacts` to skip writing the sidecar codes; useful when you
  only want the reconstructed fp16 model.
- `--include-lm-head` / `--include-embedding` to override the default
  skip behavior. Quantizing the LM head on weight-tied models is
  redundant; quantizing embedding tables hurts logits because they're
  looked up rather than multiplied.

### Validation

`scripts/quantization/test_polarquant.py` runs the round-trip on
`google/gemma-4-E2B`, using 5 native JSON-shaped samples
from `data/final/val.jsonl`. It asserts (a) the codes-only payload is
at least 30% smaller than the fp16 baseline checkpoint and (b) the
quantized model produces non-degenerate text on every sample.

```bash
uv run python scripts/quantization/test_polarquant.py
```

## TurboQuant

Reference: Amir Zandieh, Majid Daliri, Majid Hadian, Vahab Mirrokni
(Google Research / Google DeepMind / NYU), *TurboQuant: Online Vector
Quantization with Near-Optimal Distortion Rate*, **arXiv:2504.19874**,
ICLR 2026. Implementation:
[`vivekvar-dl/turboquant`](https://github.com/vivekvar-dl/turboquant)
shipped to PyPI as `turbokv` v0.1.0 (import name: `turboquant`). Pinned
in `pyproject.toml` under `[project.optional-dependencies] train`.

### What this is — and what it is NOT

TurboQuant is an **online, data-oblivious KV-cache quantizer**. It runs
at inference time and replaces the standard fp16 ``DynamicCache`` with
a ``TurboQuantCache`` that stores keys and values in 2- or 4-bit
packed form per coordinate.

It is **not** a weight quantizer. The ``model.safetensors`` file does
not change. ``turboquant_apply.py`` therefore writes the merged base
model unchanged and adds a ``turboquant.json`` sidecar that records the
quantizer config (`nbits`, `base_seed`, calibrated `skip_layers`,
`residual_length`) so a downstream loader can reconstruct the cache
deterministically.

### Algorithm in one paragraph

For each cached vector ``x`` of length ``head_dim`` (per attention
head, per token): (1) extract and store ``||x||`` in bf16; (2) rotate
the unit vector by a fixed-per-layer random orthogonal matrix Π drawn
from SO(d) via QR-decomposition with sign-fix — after rotation each
coordinate is approximately Beta-distributed (≈ N(0, 1/d) for large
d); (3) bucketize each coordinate against the precomputed Lloyd-Max
optimal centroids for the Beta distribution at b bits; (4) bit-pack
the indices into uint8 (2 indices/byte at 4-bit, 4 indices/byte at
2-bit). At attention time, unpack → centroid lookup → inverse rotation
``x ≈ ||x|| · ΠᵀC[i]``. The recent ``residual_length`` tokens stay in
fp16 to keep the freshly-generated context lossless.

### Tradeoffs

- **Pros.** Data-free / online — calibration is a single forward pass
  used only to detect outlier-norm layers (typically only layer 0) that
  should stay fp16. Drops naturally into ``model.generate`` via
  ``past_key_values=cache``. Works across Gemma/Llama/Phi-style decoder
  architectures without per-model code paths. Information-theoretic
  near-optimal: paper proves the rate is within ~2.7× the per-channel
  Shannon-Bennett lower bound.
- **Cons.** The reference implementation is pure PyTorch — the
  per-step quantize/dequantize is a Python-level operation per layer
  per step, which costs throughput. On a gemma-4-E2B model on a 5080 we
  observed **~5× slowdown** vs the bf16 ``DynamicCache``
  (66.8 → 12.2 tok/s). The TurboQuant paper claims faster runtime than
  the bf16 baseline because it ships **Triton kernels**; those are not
  in the `turbokv` 0.1.0 PyPI release we depend on. Until upstream
  ships Triton, this method is a *memory* win, not a *speed* win.
- The savings are concentrated in the long-context regime. At 4096-
  token prefill on gemma-4-E2B we measured **3.52× per-token KV
  reduction** (114,688 → 32,608 bytes/token) which produced a real
  274 MB peak-VRAM drop on a tiny model — the absolute savings scale
  with `num_hidden_layers × num_kv_heads × head_dim × context_length`.

### Supported architectures

`TurboQuantCache` materializes a `TurboQuantLayer` per full-attention
layer reported by the model config. Verified locally against:

- Gemma (`google/gemma-4-E2B`) using the active validation harness.

Should work, by structural inspection, on:

- Llama and Phi style full-attention decoders with GQA, the same shape
  `TurboQuantLayer` already handles.

#### Gemma hybrid-cache notes

Gemma tiers can declare per-layer `layer_types`. TurboQuant is only
meaningful for layers with a standard (B, H, T, D) KV cache; recurrent or
state-space layers have no KV tensor to quantize. Concretely on
gemma-4-E2B, 6 of 24 layers are full attention, so the analytic ceiling
on KV reduction is capped by those layers. The
`kv_bytes_per_token_analytic` helper in `test_turboquant.py` honors
`layer_types` so the reported reduction factor is correct for hybrid
models.

For vision-language Gemma variants, `TurboQuantCache(model.config, ...)`
must receive the **text decoder config** —
`model.config.get_text_config(decoder=True)` when the config provides it.
The `cache.py` in `turbokv` 0.1.0 already calls `get_text_config` when
available.

For future dense/MoE variants, TurboQuant is orthogonal to expert routing
when the KV cache shape is unchanged. Revalidate that separately before
adding those tiers back to the active release line.

### CLI

Apply to a fine-tuned 2B checkpoint (auto-merges if `--model` points
to a LoRA adapter):

```bash
uv run python scripts/quantization/turboquant_apply.py \
    --model checkpoints/gemma4-e2b-eliza/final \
    --calibration data/final/val.jsonl \
    --calibration-samples 128 \
    --output checkpoints/gemma4-e2b-eliza/final-turboquant
```

Useful knobs:

- `--nbits {2,4}` (default 4). 4-bit hits ~3.5× KV reduction with
  ~zero quality loss; 2-bit hits ~6.4× with measurable but small
  degradation per the paper.
- `--residual-length N` (default 128). Most-recent N tokens stay in
  fp16 to keep freshly-generated context lossless.
- `--base-seed N` (default 42). Layer i uses `seed = base_seed + i`.
  Pin this to the value recorded in `turboquant.json` at inference
  time; otherwise the rotation matrices will not match and dequant
  gives garbage.
- `--norm-threshold f` (default 5.0). Calibration skips layers whose
  per-token key norms exceed `f × median`.

### Inference-time use

```python
import json
from turboquant import TurboQuantCache
from transformers import AutoModelForCausalLM, AutoTokenizer

side = json.load(open("checkpoints/gemma4-e2b-eliza/final-turboquant/turboquant.json"))
model = AutoModelForCausalLM.from_pretrained(
    "checkpoints/gemma4-e2b-eliza/final-turboquant",
    torch_dtype="bfloat16", device_map="cuda",
)
tok = AutoTokenizer.from_pretrained("checkpoints/gemma4-e2b-eliza/final-turboquant")

cache = TurboQuantCache(
    model.config,
    nbits=side["nbits"],
    residual_length=side["residual_length"],
    base_seed=side["base_seed"],
    skip_layers=set(side["skip_layers"]),
)
out = model.generate(**tok("...", return_tensors="pt").to("cuda"),
                     past_key_values=cache, max_new_tokens=256)
```

### Validation

`scripts/quantization/test_turboquant.py` runs the round-trip on
`google/gemma-4-E2B`, with 5 native JSON-shaped prompts from
`data/final/val.jsonl` and a 4096-token long-context probe. It asserts
(a) the per-token KV-cache size shrinks by at least 30% and (b) every
quantized output is non-empty and not degenerate.

```bash
uv run python scripts/quantization/test_turboquant.py
```

The full numeric report is written to
`scripts/quantization/turboquant_report.json`. Last measured run on
gemma-4-E2B / 5080 (4-bit, skip={0}, 4096-token long context):

| metric | baseline (bf16 DynamicCache) | TurboQuant 4-bit | delta |
|---|---|---|---|
| KV bytes/token | 114,688 | 32,608 | **3.52× / -71.6%** |
| Peak VRAM (4096-tok prefill) | 1.783 GB | 1.509 GB | -274 MB |
| Tok/s (5 short prompts × 128 new) | 66.8 | 12.2 | -82% (no Triton) |

## Fused TurboQuant (Triton)

Reference: same paper as TurboQuant above (arXiv:2504.19874). Implementation:
[`fused-turboquant`](https://pypi.org/project/fused-turboquant/) v0.1.0 — a
re-implementation of the TurboQuant scheme in **Triton kernels** (encode,
decode, fused Q@K^T scoring directly from packed indices). The math is
identical to ``turbokv`` 0.1.0 above; the win is throughput.

Pinned in `pyproject.toml` under `[project.optional-dependencies] train`
alongside `turbokv`. The two are kept side-by-side because (a) `turbokv` is
the architecture-portable fallback when ``patch_model`` rejects a model
(non-power-of-2 head_dim, fused QKV, sliding window), and (b) `fused-turboquant`
needs a working Triton + system-dev-headers stack that ``turbokv`` does not.

### CLI

Apply to a fine-tuned 27B checkpoint (auto-merges if `--model` points to a
LoRA adapter):

```bash
uv run python scripts/quantization/fused_turboquant_apply.py \
    --model checkpoints/gemma4-e4b-eliza/final \
    --output checkpoints/gemma4-e4b-eliza/final-fused-turboquant \
    --bits 4
```

The script runs ``check_model_compatibility(model)`` first and refuses to
proceed if the architecture isn't supported (logged with a structured
report). After patching it discards the cache, unpatches the model, and
saves the **unmodified** base weights — fused-TurboQuant is a runtime
concern, the safetensors files are byte-identical to the input.

`--calibration` / `--calibration-samples` are accepted for parity with the
other quantizer CLIs but are unused: the Lloyd-Max codebooks and RHT seeds
are data-oblivious.

### Inference-time use

```python
from quantization.fused_turboquant_vendored.hf import patch_model
from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained(
    "checkpoints/gemma4-e4b-eliza/final-fused-turboquant",
    torch_dtype="bfloat16", device_map="cuda",
)
tok = AutoTokenizer.from_pretrained("checkpoints/gemma4-e4b-eliza/final-fused-turboquant")

cache = patch_model(model, bits=4, compress_v=True)  # patches model.forward in-place
out = model.generate(**tok("...", return_tensors="pt").to("cuda"),
                     past_key_values=cache, use_cache=True, max_new_tokens=256)
```

### Compatibility constraints (read this before deploying)

The Triton-kernel path is more constrained than the pure-PyTorch
``turbokv`` 0.1.0 path. The script's first step is
``check_model_compatibility(model)`` — these are the failure modes:

- **`head_dim` must be a power of 2** ∈ {64, 128, 256}. The Randomized
  Hadamard Transform in `fused_turboquant.kernels.triton_rht` is built
  around butterfly operations and has no implementation for arbitrary
  dims. Verified on Gemma text decoders with supported head_dim values.
- **Separate Q/K/V projections required.** Fused-QKV models (`qkv_proj`,
  `c_attn`) are rejected — `make_fused_attention_forward` raises with a
  clear error rather than producing garbage.
- **No sliding-window attention.** Models with `sliding_window` set on
  the config or attention module are rejected (e.g., Mistral-Sliding,
  Gemma-2). The kernel is causal-full-attention only.
- **No attention logit softcapping** (Gemma-2-style).
- **RoPE expected.** ALiBi / learned positional embeddings produce
  incorrect results; `check_model_compatibility` warns when RoPE isn't
  detected in config.
- **Hybrid decoder models**: only the full-attention layers are patched;
  recurrent/state-space layers keep their native state. The compatibility
  checker reports `compatible=True` when the Triton path can run, but the
  savings scale only with the full-attention layer count. **Note**: the bonus
  gemma-4-E2B run in our local test failed at the *baseline* generate
  step (HF `DynamicCache` is not the right cache for a hybrid model —
  it raises `has_previous_state can only be called on LinearAttention
  layers`); fused-TurboQuant is orthogonal to that issue. Hybrid models
  need a `HybridCache` wrapped around `CompressedKVCache`, which is
  upstream work. We document the gap and skip the bonus run.

### Blackwell / RTX 5080 Laptop notes

The 5080 Laptop is **sm_120** (Blackwell), CUDA 13.0, torch 2.11. Triton
3.6.0 supports this architecture. The build does, however, require Python
development headers to compile its `cuda_utils.so` shim on first use:

```bash
sudo apt install python3.12-dev
```

If the headers are missing the kernel raises a confusing error nested
inside the smoke test — see the "Last measured run" section below for the
exact failure mode. Without the headers the JIT cannot build and
`patch_model(..., verify=True)` fails on the very first forward pass.

### Validation

`scripts/quantization/test_fused_turboquant.py` runs three paths back-to-
back on `google/gemma-4-E2B` with 5 prompts × 128 new tokens at a 4096-token
prompt. It writes the full report to
`scripts/quantization/fused_turboquant_report.json` and asserts:

- fused peak VRAM ≤ pure-PyTorch turbokv peak VRAM, and
- fused tok/s ≥ 1.5x pure-PyTorch turbokv tok/s.

```bash
uv run python scripts/quantization/test_fused_turboquant.py
```

#### Last measured run (gemma-4-E2B, 5080 Laptop, 4-bit, 4096-token prompt + 128 new tokens)

| path | peak VRAM | tokens/sec | notes |
|---|---|---|---|
| baseline (bf16 DynamicCache) | 1.782 GB | 33.34 tok/s | HF default |
| pure-PyTorch turbokv 0.1.0 | 1.496 GB | 9.59 tok/s | -286 MB vs baseline; **3.48x slower** (Python per-step quant/dequant) |
| fused-turboquant 0.1.0 | **BLOCKED** | **BLOCKED** | Triton JIT could not compile its `cuda_utils.so` helper |

The fused path failed at `patch_model(..., verify=True)` during the first
single-token forward pass. The exact error from
`/usr/bin/gcc` was:

```
fatal error: Python.h: No such file or directory
    7 | #include <Python.h>
      |          ^~~~~~~~~~
```

**Fix:** install the matching Python development headers system-wide so
Triton's CUDA utility helper can compile:

```bash
sudo apt install python3.12-dev
```

After installing, re-run the test — the same script will exercise the
Triton path without modification. The pure-PyTorch numbers above are real
and confirm that `turbokv` 0.1.0 still delivers the documented memory win
(-19.2% peak VRAM vs baseline at 4096-token prompt) but at a 3.48x
throughput cost; recovering that throughput is exactly what the Triton
kernels in `fused-turboquant` exist to do.

## QJL

Reference: Amir Zandieh, Majid Daliri, Insu Han, *QJL: 1-Bit Quantized JL
Transform for KV Cache Quantization with Zero Overhead*,
**arXiv:2406.03482**, AAAI 2025
([acm dl](https://dl.acm.org/doi/10.1609/aaai.v39i24.34773)).
Implementation: vendored from
[`amirzandieh/QJL`](https://github.com/amirzandieh/QJL) @
`648b3641f96b6e95e091217220b94e4739fd4d82` under
`scripts/quantization/qjl/` (Apache 2.0 — see
`scripts/quantization/qjl/LICENSE` and `qjl/NOTICE.md`). There is **no
`pyqjl` package on PyPI** — the QJL kernel is shipped as a CUDA C++
extension that the user has to compile against their local nvcc.
`pyproject.toml` adds `pybind11>=2.12.0` to the `train` extras so the
template-binding step has its dep; `nvcc` and the matching
`pythonX.Y-dev` headers come from the system package manager.

### What this is — and what it is NOT

QJL is the **K (keys) side** companion to TurboQuant's V (values) side.
Both compressors were authored by Amir Zandieh; they are explicitly
designed to compose. The combination delivers ~10x KV-cache compression
at long context: 1-bit per JL-projected key coordinate + 4-bit per value
coordinate, plus a small per-token bf16 norm on each side and an
amortized outlier sketch shared across `group_size` consecutive tokens.

QJL is **not** a weight quantizer. ``model.safetensors`` is unchanged
after applying. ``qjl_apply.py`` writes the merged base model unchanged
plus a ``qjl_config.json`` sidecar that records the projection geometry
(`projection_dim_per_head`, `projection_seed`, `outlier_count_*`,
`initial_layers_count`, `group_size`, `buffer_size`, paired
`value_bits`) so the inference loader can deterministically reconstruct
the JL projection matrix Π and the value-side codebook.

### Algorithm in one paragraph

For each cached key vector ``k`` of length ``head_dim`` (per attention
head, per token): (1) extract ``||k||`` and store it in bf16; (2)
multiply by a fixed-per-layer Johnson–Lindenstrauss matrix Π ∈
R^{head_dim × s} drawn from N(0, 1) and orthogonalized via a chunked
QR decomposition (``QJLSketch.init_rot_dir``), giving a sketch ``s = Πᵀ
k`` of length ``s = projection_dim_per_head``; (3) take the **sign** of
each sketch coordinate, packing 8 signs into one uint8. To recover an
unbiased estimator of inner products at attention time, the query is
projected through the *same* Π and the score is reconstructed in a
custom CUDA kernel (``qjl_kernel/csrc/qjl_score_kernel.cu``). The paper
proves the resulting cosine-similarity estimator has minimal relative
distortion at 1 bit. To handle outlier coordinates (a few head_dim
indices with disproportionately large norms — common on layer 0 in
Llama/Gemma-style models), the kernel additionally stores a top-k outlier
sketch per group of ``group_size`` consecutive tokens, with its own
larger JL projection of dimension ``dim_outlier`` (256 for general
layers, 128 for the first ``initial_layers_count`` layers). The recent
``buffer_size`` tokens stay in fp16 to preserve freshly-generated
context losslessly.

### Tradeoffs

- **Pros.** Provably minimal relative distortion at 1-bit (paper
  Theorem 3). Storage is purely the packed sign sketch plus a per-token
  bf16 norm — no quantization scale/zero-point per coordinate, which
  is what gives the "zero overhead" tagline. Hand-written CUDA kernels
  for both quantize and score paths recover throughput at long context
  (the paper reports 2.0–2.5x speedup on prefill and decode for
  Llama-2/3 at 32k context). Composes cleanly with TurboQuant on the V
  side for end-to-end ~4x KV reduction whole-model.
- **Cons.** The compression unit is the JL sketch dimension, not the
  number of bits per coord — at the canonical
  ``projection_dim_per_head=256`` the K-side ratio
  ``head_dim*2 / (projection_dim/8 + 2)`` works out to **7.53x for
  head_dim=128** (Llama-3-style dense attention), not the
  marketing-headline 16x (which would assume zero norm overhead).
  Pushing to ``projection_dim_per_head=128`` recovers ~14.2x at the
  cost of attention-score quality. The kernel hard-codes
  ``EMB_DIM 128`` in ``csrc/qjl_quant_kernel.cu`` (line 7), so other
  head_dim values need a kernel rebuild with a different define.
  The CUDA kernel was written for Ampere/Hopper; Blackwell (sm_120)
  requires the PTX-fallback build flag (see "Build" below).

### Supported architectures

The vendored ``LlamaAttention_QJL`` and ``LlamaDecoderLayer_QJL``
modules under ``scripts/quantization/qjl/`` (mirrored from upstream
``models/llama3_qjl.py``) target the Llama-3 attention shape directly:
``q_proj`` / ``k_proj`` / ``v_proj`` / ``o_proj`` linear projections
plus ``LlamaRotaryEmbedding`` and grouped-query attention (the GQA
score kernel ``cuda_qjl_gqa_score`` handles
``num_attention_heads != num_key_value_heads``). Verified upstream on:

- Llama-2 7B and Llama-3 8B (the upstream ``run_longbench.py``
  evaluation set)

Gemma tiers require per-tier validation before release use. The current
kernel is authored around a 128-dim Llama-style attention path, while the
active Gemma targets can expose different text-decoder head dimensions.

#### Gemma caveat (read this)

QJL only applies to ``full_attention`` layers — there is nothing to
compress in recurrent/state-space layers. The ``qjl_apply.py``
calibration step honors `layer_types` and silently skips non-full-attention
layers. The on-disk config records ``n_full_attention_layers`` so the
inference loader knows which layers to wrap.

Vision-language Gemma variants need the text decoder extracted before
patching the attention modules.

### Build

The QJL kernel is **not** pip-installable. Build it once, in place,
inside the vendored directory:

```bash
cd scripts/quantization/qjl
# Default Ampere/Hopper build:
python setup.py build_ext --inplace
# Blackwell (RTX 50-series, sm_120) — the kernel was not authored
# against this arch; force PTX fallback so it compiles for compute_120
# via the JIT path:
TORCH_CUDA_ARCH_LIST="12.0+PTX" python setup.py build_ext --inplace
```

Prerequisites (verified missing on the local 5080 dev box at the time
this was vendored):

```bash
sudo apt install nvidia-cuda-toolkit python3.12-dev
```

The first installs `nvcc` (PyTorch wheels ship `ptxas` and `nvrtc` but
**not** the full nvcc driver — the cu13 wheel directory only contains
`include/`, `lib/`, no `bin/nvcc`). The second installs `Python.h` so
the C++ extension's pybind11 bindings can compile.

### CLI

Apply to a fine-tuned 2B checkpoint (auto-merges if `--model` points to
a LoRA adapter):

```bash
uv run python scripts/quantization/qjl_apply.py \
    --model checkpoints/gemma4-e2b-eliza/final \
    --calibration data/final/val.jsonl \
    --calibration-samples 128 \
    --output checkpoints/gemma4-e2b-eliza/final-qjl
```

Apply to a fine-tuned 27B checkpoint (same shape; calibration is
single-pass forward and fits in 16 GB only with offload):

```bash
uv run python scripts/quantization/qjl_apply.py \
    --model checkpoints/gemma4-e4b-eliza/final \
    --calibration data/final/val.jsonl \
    --calibration-samples 128 \
    --projection-dim-per-head 256 \
    --projection-dim-per-head-initial 512 \
    --initial-layers-count 15 \
    --outlier-count-general 8 \
    --value-bits 4 \
    --output checkpoints/gemma4-e4b-eliza/final-qjl
```

Useful knobs:

- `--key-bits 1` (only). QJL is a 1-bit-per-projected-coord scheme by
  construction; the flag exists for forward-compat.
- `--projection-dim-per-head N` (default 256). The JL sketch dimension
  per attention head. Smaller = more compression, lower attention
  score fidelity. Must be byte-aligned (`% 8 == 0`).
- `--projection-dim-per-head-initial N` (default 512). The first
  `--initial-layers-count` layers carry more attention mass; the paper
  recommends a larger sketch budget there. Leaving this at 2× the
  general dim matches upstream defaults exactly.
- `--initial-layers-count N` (default 15, paper's choice for Llama-7B
  with 32 layers). For different layer counts, scale ~half.
- `--outlier-count-general K` / `--outlier-count-initial-layers K`
  (default 8 each). Top-K head_dim coords promoted to the outlier
  sketch per group. The calibration step measures per-layer outlier
  norm ratio and records it in the sidecar so the inference loader
  can validate this budget.
- `--value-bits {2,4}` (default 4). Companion TurboQuant V-side bits
  recorded in the sidecar.
- `--group-size N` (default 32). Group of consecutive tokens that
  share an outlier-coord index table. Larger = smaller per-token
  outlier overhead, slightly worse outlier tracking on bursty
  activations.
- `--projection-seed N` (default 42). PRNG seed for the JL matrix.
  Pin this to the value recorded in `qjl_config.json` at inference
  time; otherwise Π won't match and the dequantized cosine scores are
  garbage.

### Validation

`scripts/quantization/test_qjl.py` runs on `google/gemma-4-E2B` (closest
text-only stand-in for `google/gemma-4-E2B` — see caveat above):

1. Attempts to build the vendored CUDA extension. If `nvcc` or
   `Python.h` is missing it records the exact remediation command
   (`sudo apt install nvidia-cuda-toolkit python3.12-dev`) and skips
   the runtime-kernel path.
2. Runs a baseline bf16 generation and records peak VRAM, tok/s, and
   sample outputs.
3. Captures the K activations from the first 4 attention layers via a
   forward hook, runs a **pure-PyTorch reference QJL quantize** on
   them (no CUDA extension needed — JL @ matmul + sign + bit-packing
   in plain PyTorch), and reports the realized K-side compression
   ratio per layer.
4. Sweeps `projection_dim ∈ {128, 256, 512}` so the report shows the
   full size-vs-quality curve at the same K activations.
5. Computes the analytic whole-model KV-bytes-per-token reduction with
   the paired V-side TurboQuant numbers via
   `qjl_apply.kv_bytes_per_token_analytic`.
6. Asserts (a) K-side ratio ≥ 7× at the canonical 256-dim setting,
   matching the closed-form `head_dim*2 / (projection_dim/8 + 2) =
   7.53×` for head_dim=128, and (b) baseline outputs are
   non-degenerate.

```bash
uv run python scripts/quantization/test_qjl.py
```

The full numeric report is written to
`scripts/quantization/qjl_report.json`. Last measured run on
gemma-4-E2B / 5080 (bf16 baseline cache, projection_dim=256, seed=42):

| metric | value | notes |
|---|---|---|
| nvcc present | **No** | system blocker; `sudo apt install nvidia-cuda-toolkit` |
| Python.h present | **No** | system blocker; `sudo apt install python3.12-dev` |
| QJL CUDA kernel built | **No** | both blockers above must be fixed first |
| K-side ratio (proj_dim=256, real activations) | **7.53×** | head_dim=128, per-token bf16 norm |
| K-side ratio (proj_dim=128, real activations) | 14.22× | smaller sketch — quality tradeoff |
| K-side ratio (proj_dim=512, real activations) | 3.88× | larger sketch — higher fidelity |
| Analytic KV bytes/token (bf16 baseline) | 114,688 B | K + V over 28 full-attention layers |
| Analytic KV bytes/token (QJL-K + TurboQuant-V) | 27,608 B | proj_dim=256 + V 4-bit, group_size=32 |
| **Whole-model KV reduction** | **4.15×** (-75.9%) | K + V combined; matches paper's headline range |
| Baseline tok/s (5 prompts × 128 new) | 28.93 | bf16 DynamicCache baseline |
| Baseline peak VRAM | 3.43 GB | single-prompt generation |
| Baseline output sample | `"<think> Okay, let's see. The user provided the terminal output and wants me to check if the task is complete..."` | non-degenerate |

### Blockers (current state on the 5080 dev box)

- **`nvcc` is not on the system.** PyTorch's bundled cu13 wheel ships
  `nvrtc` and `ptxas` but no full `nvcc` driver, and the
  `nvidia-cuda-nvcc-cu12` PyPI wheel only ships `ptxas` (verified). Fix:
  `sudo apt install nvidia-cuda-toolkit`.
- **`Python.h` is not present** (no `python3.12-dev` package
  installed). Fix: `sudo apt install python3.12-dev`.
- **Blackwell (sm_120) is not in the upstream test matrix.** The
  kernel sources are written against Ampere/Hopper. After the two
  `apt install`s above, the recommended build command is
  `TORCH_CUDA_ARCH_LIST="12.0+PTX" python setup.py build_ext
  --inplace`, which forces the JIT-PTX fallback path; the kernel will
  compile to compute_120 via PTX at first load. If the kernel still
  fails at runtime on sm_120, the immediate workaround is to run the
  pure-PyTorch reference path in `test_qjl.py:qjl_pure_pytorch_quantize`
  for measurement and ship the validated checkpoint to a
  Hopper/Ampere host for actual inference.
- **Hard-coded `EMB_DIM 128`** in
  `qjl/csrc/qjl_quant_kernel.cu:7`. The code only works for
  `head_dim == 128` out of the box. gemma-4-E2B / gemma-4-E2B / Llama-3
  all match. If we later need to apply QJL to a model with
  `head_dim != 128`, the `#define EMB_DIM` must be edited and the
  kernel rebuilt; there is no runtime arg for it.

## Abliteration

Reference: Arditi et al., *Refusal in LLMs is mediated by a single
direction* ([arXiv:2406.11717](https://arxiv.org/abs/2406.11717)).
Practical writeup: [Maxime Labonne, "Uncensor any LLM with
abliteration"](https://huggingface.co/blog/mlabonne/abliteration).

Computes the rank-1 refusal direction
`r = normalize(mean(harmful) - mean(harmless))` from residual-stream
activations at a configurable mid-stack layer, then projects `r` out of
every block's `self_attn.o_proj` and `mlp.down_proj` weights so the
model can never write the refusal direction back into the residual
stream. Destructive transform — save to a NEW directory.

```
uv run python scripts/quantization/abliteration_apply.py \
    --checkpoint google/gemma-4-E2B \
    --output checkpoints/gemma4-e2b-abliterated \
    --harmful-jsonl data/harmful.jsonl \
    --harmless-jsonl data/harmless.jsonl
```

Without `--harmful-jsonl` / `--harmless-jsonl` the script falls back to
a small built-in pair list intended for smoke testing only — pass real
prompt corpora (e.g.
[`mlabonne/harmful_behaviors`](https://huggingface.co/datasets/mlabonne/harmful_behaviors))
in production.