OSCAR INT2 KV cache: per-head rotations, heterogeneous SWA, MLA latent, and expanded-MHA models

.gitignore

@@ -10,12 +10,26 @@ rotation/*/MATH/
 rotation/*/MMLU/
 rotation/*/Hadamard/
 rotation/*/qkv_dumps*/
-rotation/*/rotations*/
 rotation/*/_eval_*/
 rotation/*/investigation/
-
-# MiniMax-M2.7 is internal-only for now (per-investigation work-in-progress).
-rotation/MiniMax-M2.7/
+rotation/*/bench_matrix/
+rotation/*/results/
+
+# Triton / CUDA JIT caches
+triton_cache/
+deepgemm_cache/
+*.cubin
+*.ptx
+*.llir
+*.ttir
+*.ttgir
+
+# Eval run dirs (smoke tests, wandb run artifacts)
+rotation/*/_*_smoke/
+rotation/*/runs/
+
+# Anonymous mirror dir
+anonymous/
 
 # Python caches
 __pycache__/

README.md

@@ -61,13 +61,25 @@ OSCAR is built directly into the open-source SGLang framework (main branch), lla
 <details>
 <summary><b>Qwen3.5-4B, Qwen3.5-35B-A3B, MiniMax 2.7 Preview</b> </summary>
 
-Qwen3.5
-| Model | Mode | GPQA (198) | Δ vs BF16 |
-|-------|------|------------|-----------|
-| Qwen3.5-4B | baseline | **75.25%** | — |
-| Qwen3.5-4B | OSCAR | **74.75%** | −0.50 pp |
-| Qwen3.5-35B-A3B | baseline | **80.30%** | — |
-| Qwen3.5-35B-A3B | OSCAR | **82.32%** | +2.02 pp |
+Qwen3.5 — BF16 vs OSCAR INT2 KV (2-bit, sink 64 / recent 256), mean ± std over 3 seeds (35B-A3B AIME: 8 seeds, N=30 is high-variance). OSCAR quantizer per model best: 4B uniform, 35B-A3B Lloyd-Max.
+
+**Qwen3.5-4B**
+| Benchmark | BF16 | OSCAR | Δ vs BF16 |
+|---|:---:|:---:|:---:|
+| GPQA-Diamond | 76.9 ± 1.3 | **75.8 ± 1.6** | −1.2 |
+| HumanEval | 81.7 ± 1.8 | **83.9 ± 1.0** | +2.2 |
+| AIME 2025 | 47.8 ± 3.1 | **46.7 ± 0.0** | −1.1 |
+| MATH500 | 89.5 ± 0.6 | **88.0 ± 0.6** | −1.5 |
+
+**Qwen3.5-35B-A3B**
+| Benchmark | BF16 | OSCAR | Δ vs BF16 |
+|---|:---:|:---:|:---:|
+| GPQA-Diamond | 83.3 ± 1.8 | **84.0 ± 1.3** | +0.7 |
+| HumanEval | 83.9 ± 0.6 | **86.6 ± 1.8** | +2.6 |
+| AIME 2025 † | 66.7 ± 5.3 | **62.1 ± 4.7** | −4.6 |
+| MATH500 | 92.8 ± 0.2 | **91.7 ± 0.4** | −1.1 |
+
+<sub>† AIME N=30 is high-variance; measured over 8 seeds. The −4.6 gap is not statistically significant (Welch t=1.72). At 3 seeds it read −6.7, inflated by a favorable BF16 draw.</sub>
 
 MiniMax2.7
 | Benchmark | BF16 | OSCAR (LM_RATIO=1.16) | Δ |
@@ -301,6 +313,16 @@ Calibration-pipeline folders included on this branch (`rotation/<model>/`):
 | `rotation/qwen3-8B/` | `Qwen/Qwen3-8B` | 1 | 1 | |
 | `rotation/qwen3-32B/` | `Qwen/Qwen3-32B` | 2-4 | 4 | |
 | `rotation/GLM-4.7/` | `zai-org/GLM-4.7-FP8` | 8 | 8 | FP8 weights, 92 layers |
+| `rotation/gemma-4-12B-it/` | `google/gemma-4-12B-it` | 1 | 1 | `gemma4_unified` hybrid-SWA, dual head_dim (sliding 8×256 / full 1×512), all INT2; needs transformers ≥5.5; INT2 ≈ BF16 on GPQA (62.63%); optional vision via `--enable-multimodal` |
+
+### Per-model GPQA: BF16 vs OSCAR INT2
+
+Single-seed GPQA-Diamond (198) — full-precision baseline vs OSCAR INT2 KV cache, with the calibration tag used.
+
+| Model | Calibration | GPQA (BF16) | GPQA (OSCAR INT2) |
+|---|---|---|---|
+| `Qwen/Qwen3-4B-Thinking-2507` | `seq20000_prompt83_group128` | 67.27 | 67.17 |
+| `google/gemma-4-12B-it` | `seq30000_prompt134_group128` | 62.63 | 62.63 |
 
 > MiniMax-M2.7 / Qwen3.5 calibration scripts live on `zhongzhu/hybrid-model`; pre-fit rotations for several models are available on the [RotationZoo](https://huggingface.co/Zhongzhu/OSCAR-RotationZoo).
 

rotation/.gitignore

@@ -0,0 +1,15 @@
+io_log.jsonl
+server.log
+runner.log
+pod.log
+triton_cache/
+triton_cache_head/
+__pycache__/
+bench_matrix/**/lmeval/
+eval.log
+GLM-5.1-FP8/
+qwen3-32B/rotations/
+bench_bf16_vs_int2/
+investigation/*/upproj_rotations/
+investigation/*/hp128_subspace/
+investigation/*/qkv_dumps*/

rotation/GLM-5.2-FP8/README.md

@@ -0,0 +1,112 @@
+# GLM-5.2-FP8 — INT2 KV recipe (MLA shared latent)
+
+MLA stores one compressed latent `c_kv` (512-d) shared by all heads, plus a
+positional `k_pe`. Only `c_kv` is quantized; `k_pe` stays BF16. Per-head
+rotations do not apply here — there is nothing per-head to rotate.
+
+Rotation is `Rcov · P · Hblock` (covariance eigenvectors, bit-reversal
+permutation, per-group Hadamard) with **Lloyd-Max** on, group 128.
+
+The BF16 sink and recent windows apply here like everywhere else: sink 64,
+recent 256, on by default. That was not true until `NSAInt2HPKVPool` /
+`MLAInt2HPKVPool` learned to read `SGLANG_MIXED_KV_PREFIX_TOKENS` and
+`SGLANG_MIXED_KV_RECENT_TOKENS` — before that they quantized every latent
+token, including the attention sink and the newest ones, and this file told
+you setting those vars changed nothing. **Every number below predates the
+windows**, so read them as the no-window arm, not as the recipe.
+
+| Benchmark | BF16 | INT2 (no window) | Δ |
+|---|:---:|:---:|:---:|
+| GPQA @32K | 81.3 ± 0.0 | 68.7 ± 1.8 | −12.6 |
+| HumanEval (≤16K) | 91.5 ± 1.2 | 88.2 ± 1.5 | −3.3 |
+| AIME 2025 | 76.7 ± 3.3 | 75.6 ± 6.9 | −1.1 (n.s.) |
+| MATH500 | 94.8 ± 0.8 | 93.8 ± 0.3 | −1.0 |
+
+**The −12.6 at 32K is truncation, not wrong answers.** On questions both arms
+answer, they pick the same option 97.1 % of the time; INT2 simply thinks 1.69×
+longer, so the truncation rate goes 14 % → 28 %. BF16 itself gains 6.6 points
+moving from a 32K to a 64K budget, which confirms the metric is budget-bound.
+Give this model a generous `max_tokens` before reading a regression into it.
+The INT2 @64K cell is not measured.
+
+## Do not expect a better rotation to close it
+
+Measured on a real `c_kv` dump (functional error, lower is better):
+
+| variant | error | vs no rotation |
+|---|:---:|:---:|
+| no rotation | 0.0845 | — |
+| Hadamard | 0.0839 | ≈ 0 |
+| **`Rcov·P·Hblock`** | **0.0523** | −0.72 dB |
+| Hadamard + HP128 | 0.0447 | −2.74 dB |
+
+Only the last row reaches BF16 parity, and it costs ~6.5 bits per element, which
+defeats the point of 2-bit. The latent is a *trained* compressed representation —
+it has no outlier structure for a rotation to flatten. `Rcov·P·Hblock` is the
+best of the cheap options; take it and move on.
+
+## Steps
+
+```bash
+# 1. dump c_kv on GPQA (TP=16, 2 nodes)
+bash save_ckv_glm52.sh
+
+# 2. fit the per-layer latent rotation
+bash compute_rotation.sh     # -> rotations/layer_*.pt
+```
+
+Fitted rotations are published as
+`Zhongzhu/OSCAR-RotationZoo/GLM-5.2-FP8/c_kv_rcov_phblock_g128/` — download
+rather than re-fit unless you are changing the recipe.
+
+## Serving
+
+```bash
+SGLANG_ALLOW_OVERWRITE_LONGER_CONTEXT_LEN=1 \
+SGLANG_LLOYD_MAX=1 \
+SGLANG_OSCAR_MLA_KV_ROTATION_PATH=<rotations-dir> \
+SGLANG_OSCAR_MLA_KV_GROUP_SIZE=128 \
+python -m sglang.launch_server --model-path zai-org/GLM-5.2-FP8 \
+  --tensor-parallel-size 16 --nnodes 2 --node-rank <0|1> \
+  --dist-init-addr <head>:20000 \
+  --kv-cache-dtype bfloat16 --disable-radix-cache \
+  --mem-fraction-static 0.85
+```
+
+Three flags in that command are load-bearing and were each wrong in an earlier
+version of this file:
+
+* `SGLANG_OSCAR_MLA_KV_ROTATION_PATH` is the name the pool reads. It is also
+  what *creates* the pool — misspell it (this file said
+  `SGLANG_OSCAR_MLA_ROTATION_DIR`) and the server starts, serves, and scores
+  with a plain BF16 latent cache and no rotation at all.
+* `--kv-cache-dtype bfloat16`, not `int2`. `int2` aborts at argument validation
+  on a DSA model; the latent is fake-quantized into a float cache, so that
+  cache stays float. Leaving it at `auto` is worse than either: sglang picks
+  fp8_e4m3 on SM100+ and bfloat16 on Hopper and below, which silently changes
+  the method with the GPU generation. Pinning it is what makes a B200 run
+  comparable to the H100 numbers in the table above.
+* No `--page-size`. A DSA model forces 64; the 8 this file used to pass was
+  ignored.
+
+The BF16 windows need no flag — 64/256 is the default the pool applies when
+`SGLANG_MIXED_KV_PREFIX_TOKENS` / `SGLANG_MIXED_KV_RECENT_TOKENS` are unset,
+because the two vars' own defaults are the per-head pool's older 32/128 and
+inheriting those would put this path below the floor. Raise the recent window
+for a model that needs more (Gemma-4 and Qwen3-8B use 512). Setting both to 0
+turns the windows off; the server logs a warning when you do, since that is an
+A/B arm and not a serving configuration. Confirm the line
+
+```
+[Int2HPKVPool] BF16 latent windows: sink=64 recent=256 ...
+```
+
+in the server log — its absence means the pool fell back to quantizing every
+latent token. `SGLANG_MIXED_KV_AUDIT=1` additionally checks each live
+request's tiers against the windows every 25 decode steps (it syncs the decode
+path — never leave it on for a timed run).
+
+Set `SGLANG_OSCAR_MLA_KV_REAL_KERNEL=1` to use the packed-INT2 latent kernel
+instead of the fake-quant path: bit-identical packing, decode cosine 1.000000,
+1024 → 160 B/token (**6.4×**). It is correct but not yet fast — measure before
+enabling it for throughput.

rotation/GLM-5.2-FP8/compute_rotation.sh

@@ -0,0 +1,27 @@
+#!/usr/bin/env bash
+# Fit the per-layer MLA latent rotation (Rcov . P . Hblock) for GLM-5.2-FP8.
+#
+# Input is a c_kv dump: one tensor per layer of the 512-d compressed latent,
+# captured while serving. Point DUMP_PATH at it. Unlike the MHA models there is
+# no per-head structure here -- the latent is shared across heads.
+set -euo pipefail
+SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
+COMPUTE_SCRIPT="${SCRIPT_DIR}/../compute_mla_oscar_rotation.py"
+
+DUMP_PATH="${DUMP_PATH:?set DUMP_PATH to the c_kv dump directory}"
+OUTPUT_DIR="${OUTPUT_DIR:-${SCRIPT_DIR}/GPQA/_rcov_phblock/rotations}"
+KV_LORA_RANK="${KV_LORA_RANK:-512}"
+GROUP_SIZE="${GROUP_SIZE:-128}"
+MAX_TOKENS="${MAX_TOKENS:-8192}"
+mkdir -p "${OUTPUT_DIR}"
+
+python3 "${COMPUTE_SCRIPT}" \
+    --dump-path "${DUMP_PATH}" \
+    --output-dir "${OUTPUT_DIR}" \
+    --kv-lora-rank "${KV_LORA_RANK}" \
+    --group-size "${GROUP_SIZE}" \
+    --max-tokens "${MAX_TOKENS}" \
+    "$@"
+
+echo "rotations -> ${OUTPUT_DIR}"
+ls "${OUTPUT_DIR}" | head

rotation/GLM-5.2-FP8/eval_gpqa.sh

@@ -0,0 +1,22 @@
+#!/usr/bin/env bash
+# GPQA eval wrapper for zai-org/GLM-5.2-FP8, CUDA graph on (the configuration all published
+# OSCAR GPQA numbers were measured under). Override CUDA_GRAPH_MAX_BS to change it.
+set -euo pipefail
+SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
+
+export MODEL="${MODEL:-zai-org/GLM-5.2-FP8}"
+# MLA: the rotation is per-layer latent files (layer_*.pt), not a k/v pair, and
+# is passed through MLA_ROT_PATH -- ROT_DIR/K_ROT_FILENAME do not apply here.
+export MLA_ROT_PATH="${MLA_ROT_PATH:-${SCRIPT_DIR}/GPQA/_rcov_phblock/rotations}"
+export ROT_DIR="${ROT_DIR:-${MLA_ROT_PATH}}"
+export RUN_DIR="${RUN_DIR:-$(dirname "${ROT_DIR}")/_eval_gpqa_oscar}"
+export TP_SIZE="${TP_SIZE:-16}"
+export GROUP_SIZE="${GROUP_SIZE:-128}"
+export K_CLIP="${K_CLIP:-0.96}"
+export V_CLIP="${V_CLIP:-0.92}"
+export CUDA_GRAPH_MAX_BS="${CUDA_GRAPH_MAX_BS:-32}"
+export LLOYD_MAX="${LLOYD_MAX:-1}"   # GLM-5.2 recipe uses Lloyd-Max
+export NAME="${NAME:-gpqa_oscar_glm_5_2}"
+# MLA latent path: rotations are per-layer c_kv files, not k/v pairs.
+export MEM_FRACTION_STATIC="${MEM_FRACTION_STATIC:-0.85}"
+exec bash "${SCRIPT_DIR}/../eval_oscar_gpqa.sh"

rotation/MiniMax-M2.7/README.md

@@ -0,0 +1,59 @@
+# MiniMax-M2.7 — INT2 KV recipe
+
+Dense MHA. Shared per-layer rotation (`qqt_r_h_pbr` for K, `sst_r_h_pbr` for V)
+with the **uniform** quantizer — Lloyd-Max off.
+
+| Benchmark | BF16 | INT2 | Δ |
+|---|:---:|:---:|:---:|
+| GPQA | 78.28 | 80.3 | +2.0 |
+| HumanEval | 88.17 | 92.1 | +3.9 |
+| AIME 2025 (95K budget) | 76.67 | 90.0 | +13.3 |
+| MATH500 | 93.79 | 94.6 | +0.8 |
+| SWE-bench-Verified | – | 70.8 | – |
+| LiveCodeBench v6 (95K) | – | 68.4 | – |
+
+Read the deltas with care: the BF16 column came from the model authors under a
+different harness and seed count, so these are **not** paired measurements. The
+honest claim is that INT2 lands in the same band as BF16 on all four, not that
+it beats it. LCB 68.4 is a hard-window effect, not quantization — GLM-4.7 in
+BF16 scores 65.5 on the same window.
+
+## Two settings that are not optional
+
+- **Mixed-KV windows must be on.** With `SGLANG_ENABLE_MIXED_KV_WINDOWS=0` this
+  model degenerates into repetition loops.
+- **Budget 95K generation tokens.** M2.7 is a long-thinking model; a smaller
+  budget measures truncation, not accuracy. AIME and LCB above both use 95K.
+
+## Steps
+
+```bash
+# 1. dump post-RoPE Q/K/V on GPQA (TP=4)
+bash save_qkv_m27.sh
+
+# 2. fit
+bash compute_rotation.sh    # qqt_sst / r_h_pbr
+```
+
+## Serving
+
+```bash
+SGLANG_ENABLE_MIXED_KV_WINDOWS=1 SGLANG_ALLOW_OVERWRITE_LONGER_CONTEXT_LEN=1 \
+SGLANG_MIXED_KV_PREFIX_TOKENS=64 SGLANG_MIXED_KV_RECENT_TOKENS=256 \
+SGLANG_MIXED_KV_HP_DTYPE=bfloat16 SGLANG_MIXED_KV_SCALE_DTYPE=float32 \
+SGLANG_OSCAR_ABSORB_V_ROTATION=0 SGLANG_LLOYD_MAX=0 \
+SGLANG_OSCAR_K_CLIP_RATIO=0.96 SGLANG_OSCAR_V_CLIP_RATIO=0.92 \
+SGLANG_OSCAR_K_ROTATION_PATH=<out>/k_rotation_qqt_r_h_pbr.pt \
+SGLANG_OSCAR_V_ROTATION_PATH=<out>/v_rotation_sst_r_h_pbr.pt \
+python -m sglang.launch_server --model-path MiniMaxAI/MiniMax-M2.7 \
+  --tensor-parallel-size 4 \
+  --kv-cache-dtype int2 --page-size 8 --disable-radix-cache \
+  --tool-call-parser minimax-m2
+```
+
+## Evaluating it
+
+If you serve with `--cuda-graph-max-bs 1`, drive the eval client with **one**
+thread. Two threads collapse throughput from ~115 tok/s to ~8 tok/s, because
+every concurrent batch misses the single captured graph shape and falls back to
+eager. This looks like a hung eval rather than a configuration mistake.

rotation/MiniMax-M2.7/compute_rotation.sh

@@ -0,0 +1,42 @@
+#!/usr/bin/env bash
+# Fit the shared per-layer K/V rotation for MiniMaxAI/MiniMax-M2.7.
+#
+#   METHOD=qqt_sst (default)  calibrated from a Q/K/V dump (DUMP_PATH)
+#   METHOD=hadamard           data-free, no dump needed
+#
+# Outputs into OUTPUT_DIR:
+#   qqt_sst  -> k_rotation_qqt_r_h_pbr.pt, v_rotation_sst_r_h_pbr.pt
+#   hadamard -> k_rotation_hadamard.pt,    v_rotation_hadamard.pt
+set -euo pipefail
+SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
+COMPUTE_SCRIPT="${SCRIPT_DIR}/../compute_kv_rotation.py"
+
+METHOD="${METHOD:-qqt_sst}"
+COMPOSITION="${COMPOSITION:-r_h_pbr}"
+HEAD_DIM="${HEAD_DIM:-128}"
+NUM_LAYERS="${NUM_LAYERS:-62}"   # MiniMax-M2.7
+
+DATASET="${DATASET:-GPQA}"
+CALIB_DIR="${CALIB_DIR:-${SCRIPT_DIR}/${DATASET}/latest}"
+DUMP_PATH="${DUMP_PATH:-${CALIB_DIR}/qkv_dumps_merged}"
+OUTPUT_DIR="${OUTPUT_DIR:-${CALIB_DIR}/rotations}"
+mkdir -p "${OUTPUT_DIR}"
+
+ARGS=(--method "${METHOD}"
+      --num-layers "${NUM_LAYERS}"
+      --head-dim "${HEAD_DIM}"
+      --output-dir "${OUTPUT_DIR}")
+
+# hadamard is data-free; only the calibrated method reads a dump.
+if [[ "${METHOD}" != "hadamard" ]]; then
+    if [[ ! -d "${DUMP_PATH}" ]]; then
+        echo "no dump at ${DUMP_PATH} -- run save_qkv_m27.sh first" >&2
+        exit 1
+    fi
+    ARGS+=(--dump-path "${DUMP_PATH}" --composition "${COMPOSITION}")
+fi
+
+python3 "${COMPUTE_SCRIPT}" "${ARGS[@]}" "$@"
+
+echo "rotations -> ${OUTPUT_DIR}"
+ls -la "${OUTPUT_DIR}"

rotation/MiniMax-M2.7/eval_gpqa.sh

@@ -0,0 +1,20 @@
+#!/usr/bin/env bash
+# GPQA eval wrapper for MiniMaxAI/MiniMax-M2.7, CUDA graph on (the configuration all published
+# OSCAR GPQA numbers were measured under). Override CUDA_GRAPH_MAX_BS to change it.
+set -euo pipefail
+SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
+
+export MODEL="${MODEL:-MiniMaxAI/MiniMax-M2.7}"
+export ROT_DIR="${ROT_DIR:-${SCRIPT_DIR}/rotations}"
+export RUN_DIR="${RUN_DIR:-$(dirname "${ROT_DIR}")/_eval_gpqa_oscar}"
+export TP_SIZE="${TP_SIZE:-4}"
+export GROUP_SIZE="${GROUP_SIZE:-128}"
+export K_CLIP="${K_CLIP:-0.96}"
+export V_CLIP="${V_CLIP:-0.92}"
+export CUDA_GRAPH_MAX_BS="${CUDA_GRAPH_MAX_BS:-1}"
+export NAME="${NAME:-gpqa_oscar_minimax_m27}"
+# M2.7 is a long-thinking model: budget generation, and with a single captured
+# graph shape the eval client must be single-threaded (115 tok/s vs 8 at two).
+export MAX_NEW_TOKENS="${MAX_NEW_TOKENS:-95000}"
+export NUM_WORKERS="${NUM_WORKERS:-1}"
+exec bash "${SCRIPT_DIR}/../eval_oscar_gpqa.sh"