OSCAR INT2 KV cache: per-head rotations, heterogeneous SWA, MLA latent, and expanded-MHA models - #20
Open
zzz0906 wants to merge 106 commits into
Open
OSCAR INT2 KV cache: per-head rotations, heterogeneous SWA, MLA latent, and expanded-MHA models#20zzz0906 wants to merge 106 commits into
zzz0906 wants to merge 106 commits into
Conversation
google/gemma-4-12B-it is model_type gemma4_unified, a hybrid-SWA model with heterogeneous per-layer KV geometry: 40 sliding layers (8 KV heads x head_dim 256, window 1024) + 8 full-attention layers (1 KV head x head_dim 512, k_eq_v, partial proportional RoPE). This adds full INT2 OSCAR KV support: - configs/gemma4_unified.py + models/gemma4_unified.py: vendored config that remaps the unified geometry to sglang convention (full=base, sliding=swa) + a Gemma4UnifiedForConditionalGeneration shim over gemma4_causal's Gemma4ForCausalLM (text-only; reuses k_eq_v + the language_model prefix). - Two-geometry-group UnifiedInt2HPKVPool for the sliding/full split, sharing one head-dim-agnostic allocator; per-group buffers + per-group fused decode-flush; per-layer OSCAR rotation loaded at per-layer head_dim. - Window sliding layers in the mixed-KV DECODE path (not only prefill); this is the fix that recovers GPQA-Diamond INT2 from 50.0% to 62.63% (= BF16 62.63%). - Hybrid-SWA enablement: mm_disabled_models, is_hybrid_swa_compress, SWAKVPool is_decode forwarding; per-layer head_dim in compute_kv_rotation.py; Triton grouped-INT2-decode constexpr-reassign fix. - rotation/gemma-4-12B-it/ pipeline scripts (dump / rotation / eval) + README. All changes are additive and arch-gated; uniform-geometry OSCAR models (qwen3/minimax/glm) take the unchanged default path.
- run_simple_eval.py: wire base HumanEval (pass@1, 1 sample/task); raise OpenAI client read timeout to 7200s for long thinking-model runs; always honor --num-threads (HumanEval hardcodes 3) so the single-stream INT2 server is not over-parallelized - eval_oscar_gpqa.sh: drop redundant SGLANG_LLOYD_MAX=0 default
…ybrid-model # Conflicts: # rotation/_eval_runner/run_simple_eval.py
…MMU Pro Adds gated multimodal support for gemma4_unified on top of the text INT2 OSCAR path, which is byte-identical when --enable-multimodal is off: - models/gemma4_unified.py: text shim -> gated MM model. Builds the encoder-free vision embedder (vision_embedder + embed_vision, checkpoint-native names) and merges image features into the text stream via general_mm_embed_routine ONLY when enable_multimodal; otherwise delegates to Gemma4ForCausalLM. Vision tokens flow through the text decoder, so the INT2 two-group KV pool is unaffected. - multimodal/processors/gemma4_unified*.py: vendored Gemma4Unified image processor (aspect-resize -> 16px patchify -> 3x3 merge -> 48x48x3 patches + 2D pos-ids) + processor + sglang processor registration (transformers 5.9 lacks the unified processor; we vendor it to keep the text-path field remap working). - configs/gemma4_unified.py: structured Gemma4UnifiedVisionConfig + processor reg. - configs/model_config.py: add arch to multimodal_model_archs (gated by --enable-multimodal; text-only by default via mm_disabled_models). INT2-specific: image tokens use causal attention (INT2 is incompatible with the bidirectional-image-mask Triton path); the HP-prefix pool must be enlarged for image-heavy prompts. MMMU Pro standard (n=150) = 43.33% through the INT2 OSCAR KV. All additive/gated; uniform-geometry OSCAR models and the text INT2 path unaffected.
…t news - All configured models: add rotation/gemma-4-12B-it/ (google/gemma-4-12B-it, hybrid-SWA dual head_dim sliding 8x256 / full 1x512, all INT2; transformers >=5.5; INT2 ~= BF16 on GPQA 62.63%; optional vision via --enable-multimodal). - Latest News: SGLang Gemma 4 12B INT2 support (first hybrid-SWA / dual-head-dim model).
- Brings in latest main (FP4 KV fix #6, latest-news dedup + the [2026-06-04] Gemma-4-12B SGLang line, llama.cpp link fix). - Resolves the README Latest-News conflict by keeping main's deduped entries (drops the duplicate gemma SGLang news line I had added). - Adds a per-model GPQA reproducibility table (single-seed GPQA-Diamond, BF16 vs OSCAR INT2) covering Qwen3-4B-Thinking-2507 and google/gemma-4-12B-it.
- MLA/NSA INT2 latent fake-quant pool (mla_int2_kv_pool.py) + pool routing and SGLANG_OSCAR_MLA_KV_* env vars; GLM-5.2 3D c_kv dump-rank fix - Rotation generators: Rcov·P·Hblock composition, kv_b_proj-Hessian HP subspace, deterministic functional-error harness - Eval harness: aime25 + math500 tasks (deterministic boxed+sympy grading), --seed, 21600s read timeout, one-server bench matrix loop + aggregators - Qwen3.5-4B/35B-A3B 4-benchmark x 3-seed results (metrics committed; per-model-best quantizer: 4B uniform, 35B Lloyd-Max) + README tables - Qwen3-32B GSM8K strict/flexible parity check vs vLLM PR (no crater)
- pack_int2: per-group(128) Lloyd-Max bucketize bit-exact with NSAInt2HPKVPool fake-quant (max_abs 5e-9), 2-bit pack; c_kv 1024B -> 160B/token (6.4x) - mla_decode_int2: fused MQA-absorb decode with inline dequant (flash-decode online softmax + split reduce); exact vs torch ref (cos=1.000000, rel_l2 2e-5) on real dumped c_kv + published rotation - identical-structure bf16 baseline kernel for fair comparison; v1 speed 0.42x of bf16 twin (element-wise unpack, no tl.dot yet) — optimization path known from the per-head INT2 kernel line - validated on H100 (results in proto_results.txt); k8s manifest included
- minimax_m3.py: M3 (VL-MoE) model with shared-expert applied before FusedMoE.experts() (in-place mutation ordering) and swigluoai chunked branch for pure-TP - gpu_flush_int2.py: Lloyd-Max branch for the decode-flush kernel (was uniform-only; bit-exact verified on M2.7/M3) - model_config.py: default bfloat16 when text_config lacks torch_dtype - fused_moe / moe_runner: M3 MoE plumbing - rotation/MiniMax-M3/rotations: per-head Hadamard K/V rotations (calibrated all-head variant measured within noise of Hadamard; GPQA 84.85 INT2 vs 88.38 BF16, see .RUD/m3oscar/PLAN.md) No overlap with hybrid-model changes: all six files were untouched on this branch since the common base (verified per-file vs 29e5c81).
PLAN, per-investigation metrics.json + eval.log + k8s manifests from .RUD/m3oscar (GPQA-198: bf16 88.38, int2-hadamard 84.85, calib 84.85, calib-allhead 85.35, lloydmax 84.85). Raw 21GB qkv dumps not archived (regenerable via 03_calib_dump manifest).
…giene
- GLM-5.2 official GPQA metrics: bf16 3 seeds (81.31 x3, each 161/198),
pure-int2 seeds 0-1 (70.2 / 69.2); truncation-driven gap (answered-only
accuracy 94.3 vs 94.4)
- Deterministic functional-error logs on real c_kv: Rcov.P.Hblock -0.72dB
vs Hadamard; +HP64 -1.53dB; +HP128 -2.74dB (GLM-5.1 curve reproduced)
- GLM-5.1-era latent analysis protos (proto_latent_{hp,scaling}.py) used
for the GLM-5.2 evidence line
- client read-timeout 21600->28800s for 64K-budget runs
- gitignore: eval.log, legacy GLM-5.1 results, HF rotation copies,
investigation heavy artifacts; drop refuted launch_parallel_seeds.sh
…96.0 vs bf16 94.2 -> gap=truncation) + bf16 math500 seed0
… 75.6±6.9 (n.s.), math500 93.8±0.3; bf16@64K gpqa seed0=87.88
…reporter-config radix-ON (16.0/18.5/11.5/12.5) - radix+mixedKV corruption confirmed; harness + leg runner
…s exonerated); +LloydMax long_ctx 18->24
… radix-cache was reporter's bug; long_ctx residual 23 vs 30
…g256 blocked by LM-kernel/head_dim; per-chan K = upgrade path
… 36.0 > 8B bf16 30.0
…(49.5 vs 48.0) - hybrid attention = smallest quant surface
…1.7/67.8/78.7 (humaneval/aime/math500); AIME -10 = truncation-tax family pattern
…ed V unrotated while decode counter-rotated output -> generation collapse on Qwen3-30B-A3B
…cause 2): fused rope-epilogue wrote raw bf16 into int2 layout on Qwen3MoE, bypassing rotate+quantize -> generation collapse
…rness - fit_perhead.py: per-(layer,kv_head) cov-eig o Hadamard rotations, format_version 2 - fq_common.py / ppl_fakequant.py: chunked teacher-forced PPL with production KV tiering (sink64 + recent256 BF16, INT2 group128, clip .96/.92); aborts if the hook quantizes nothing. Validated on Qwen3-8B: Zoo -0.6%, Hadamard +10.2% - Qwen3-30B-A3B: shared rotation +16.6% PPL vs per-head +0.35% -> per-head is required on models with heterogeneous KV heads (issue #16 root cause) - 30B calibration + per-head rotations, BFCL leg manifests/scores - dump fork: never take the fused set_kv path (no sgl_kernel support there)
Capture hardcoded (hidden_states, residual) as the only tensors crossing a pipeline stage boundary. Kimi carries block_residual instead -- a stacked prefix-sum per attention-residual block -- so every non-first rank raised KeyError: 'block_residual' inside capture. The rank died mid-capture and its peers only saw Gloo 'Connection closed by peer', which reads like a fabric fault rather than a missing buffer. Models now optionally declare per-token proxy shapes; capture allocates exactly those keys. Kimi derives block_residual's width from its own layer range (blocks open at global idx % attn_res_block_size == 0), so the shape is known without running a forward. Both the capture slice and the replay copy already index only dim 0, so the 3-D buffer needs no other change. (cherry picked from commit 20e0547)
sgl-kernel builds FlashAttention for sm8x/sm90 only, so _is_fa3_supported() is False on Blackwell and int2 prefill died with 'flash_attn at sgl-kernel is only supported on sm90 and above' -- the message says sm90-and-above but the predicate accepts capability major 8 or 9, so sm100 falls out. The SDPA varlen path added for gemma4's 512-dim heads already handles arbitrary head_dim, causal + sliding window, asymmetric K/V dims and GQA, so route to it on unsupported hardware too. SDPA cannot apply a softcap, so a capping layer now raises instead of silently dropping it. Correctness fix, not a fast path: SDPA loops per request and materializes the mask, so Blackwell int2 prefill is slower than FA3 on Hopper. (cherry picked from commit 54384ba)
The fused flush kernel derives the int2 byte layout from reshape(BLOCK_TOK, 4, HEAD_DIM // 4), so it needs a power-of-two HEAD_DIM and raised on Kimi's K=192. This is not Blackwell-specific: the guard is pure arithmetic, so Kimi int2 decode has never run on any GPU. It went unnoticed because the published Kimi PPL/KL numbers are teacher-forced -- prefill writes KV through a different writer that already handles 192, and flush only fires once generation pushes tokens out of the HP-recent ring. Padding the row to 256 does not work: the packed layout would become (b, b+64, b+128, b+192) where the real 192-wide one is (b, b+48, b+96, b+144), and decode reads back exactly what prefill wrote. So gather the flushing rows and reuse the writer prefill already uses, which gives the same layout by construction. Costs a Python loop over layers and a materialized gather; the fused kernel stays the path for power-of-two dims. Clip matches prefill's non-power-of-two path (torch sort) rather than the fused kernel's binary search, which is what decode/prefill consistency needs here. Lloyd-Max and a caller that supplies no per-layer tensors both raise instead of silently writing something else. (cherry picked from commit 7b2ed7d)
bench_overhead/diag2/diag_uniform answered one-off questions while the MLA kernel was being written and nothing references them. bench_fused (referenced) and test_mla_kernel (the kernel-vs-simulation check) stay.
test_mla_kernel checks the packed-INT2 MLA kernel against the fake-quant simulation, so it belongs with the code it guards rather than in an investigation folder; it needs a GPU, unlike the CPU checks in tests/. bench_fused was v1-vs-v2 scaffolding from writing that kernel. Leaves no rotation/investigation tree: bench results, eval logs and the K8s manifests that produced them stay off this branch.
They evaluate an already-running server, and no other model directory ships eval tooling. validate_rotations stays: it checks the artifact the workflow just produced, which is the guard against a silently-failed calibration that loads fine and serves at Hadamard quality.
Each of the three shipped models without a rotation/<model>/ folder, so nobody could reproduce a rotation for them from this repo. Each README carries the finding that actually matters for that model: - MiniMax-M3: the dump must cover every TP rank. A rank-0 dump captures 4 of 64 query heads and 1 of 4 KV heads, fits a rotation no better than Hadamard, and reports nothing. merge_qkv_allhead.py rebuilds the full head axis and refuses a partial dump. Also records that the shared per-layer rotation is saturated here (85.35 vs 84.85 Hadamard) and that Lloyd-Max is worse. - GLM-5.2-FP8: no rotation closes the 32K GPQA gap, because it is truncation -- both arms agree on 97.1 % of answered questions and INT2 just thinks 1.69x longer. Includes the measured error table showing the only variant that reaches parity costs ~6.5 bits. - MiniMax-M2.7: mixed-KV windows are mandatory (off degenerates into repetition loops), the budget must be 95K, and its BF16 column is not a paired measurement. .gitignore had rotation/MiniMax-M2.7/ excluded as internal-only; dropped so the recipe ships. bench_matrix/ and results/ are now excluded per-model instead, so artifacts stay out while sources go in.
M2.7 has 62 layers and M3 has 60, both head_dim 128 -- read from the shipped configs. Making the caller supply a number they have to go look up is a step where a wrong value silently produces a rotation checkpoint of the wrong depth.
compute_kv_rotation.py takes --dump-path and --num-layers; the wrappers passed --dump and --layers and died on argparse the first time either was run. It also rejects --dump-path with --method hadamard, which is data-free, so the dump arguments are now added only for the calibrated method. Caught by running the wrapper instead of a hand-written command in the M2.7 smoke; the smoke now invokes the shipped script so this path stays covered.
Every model directory now has eval_gpqa.sh next to its rotation recipe, so the number in the README can be reproduced rather than taken on trust. The wrappers set only per-model defaults and exec the shared rotation/eval_oscar_gpqa.sh. Adds wrappers for Qwen3-30B-A3B, MiniMax-M2.7, MiniMax-M3 and GLM-5.2-FP8, which had recipes but no eval, and restores the shared runner (run_simple_eval.py plus the HumanEval/AIME/MATH500 graders that exist because the upstream ones mis-parse thinking-model output). CUDA graph is on by default everywhere, which is what the published GPQA numbers were measured under. Gemma-4 hardcoded --disable-cuda-graph; that is now GRAPH_ARGS, defaulting to --cuda-graph-max-bs 16, with the old flags documented for reproducing the graph-off runs.
third_party/simple_evals is a submodule, so a plain git clone leaves it empty and every grader imports it. The failure surfaced only after the model had loaded -- ImportError: cannot import name 'common' from 'simple_evals' -- which is a long wait to learn that a checkout step was missed. Initialise it up front and exit with the fix in the message if it is still missing. Found by running the eval against a fresh clone of the pushed branch rather than an existing worktree, which is what a first-time user gets.
eval_oscar_gpqa.sh hardcoded SGLANG_COQUANT_ABSORB_V_ROTATION=1. Absorption folds R_v into o_proj and assumes a single rotation per layer, so it is invalid for per-head (format_version 2) checkpoints. Qwen3-30B-A3B, whose whole result depends on per-head rotations, scored 34.3 on GPQA with absorption forced on -- worse than the 43.9 a shared rotation gets, against 58.6 with it off. It is now ABSORB_V, defaulting to 0. The same hardcoding was fixed in run_bench_matrix.sh earlier; this script was missed, so the number it produced looked like a model result rather than a harness bug. Also fixes the env var name: the pool reads SGLANG_OSCAR_ABSORB_V_ROTATION, so the SGLANG_COQUANT_ prefix was being ignored entirely.
The wrapper inherited the shared-rotation filenames from eval_oscar_gpqa.sh, so running the example as shipped would have evaluated the configuration this model demonstrably collapses under (GPQA 43.9 shared vs 58.6 per-head) while looking like the headline result. Defaults to k_perhead.pt / v_perhead.pt with absorption off, and says in the file that these come from fit_perhead_rotation.py rather than compute_rotation.sh.
eval_oscar_gpqa.sh only ever set the per-head K/V rotation paths, which the MLA pool does not read -- it reads SGLANG_OSCAR_MLA_KV_ROTATION_PATH. So the GLM-5.2-FP8 wrapper shipped in the previous commit would have evaluated with no rotation applied at all while appearing to run the documented recipe. Adds MLA_ROT_PATH (empty keeps the MHA path untouched) and points the GLM-5.2 wrapper at its per-layer layer_*.pt directory, with Lloyd-Max on to match the recipe in that model's README.
GLM-5.2-FP8 aborts with 'DeepSeek DSA only supports bf16/bfloat16 or fp8_e4m3 kv_cache_dtype': MLA quantizes the shared latent through the rotation, and the KV cache itself stays bf16. The flag is now emitted only for the MHA path, keyed on MLA_ROT_PATH, so the GLM-5.2 example runs as shipped instead of failing at server start.
Removing --kv-cache-dtype int2 for MLA exposed the next layer: sglang rejects --kv-cache-quant-group-size without it. The MLA pool reads its group size from SGLANG_OSCAR_MLA_KV_GROUP_SIZE, which the script already exports, so the CLI flag is redundant there.
The MLA branch set KV_DTYPE_ARGS="" and GROUP_SIZE_ARGS=""; an empty string still expands to one empty argv entry, so sglang saw an unexpected empty positional and exited before it wrote anything to server.log -- which is why the failure produced almost no output. Empty arrays expand to zero arguments. Introduced by the two preceding MLA fixes, not a pre-existing bug.
The shipped eval had no nnodes/node-rank/dist-init-addr, so the two models that do not fit on one node -- MiniMax-M3 and GLM-5.2-FP8, both documented in this repo as 2-node TP=16 -- could not be evaluated with it at all. Trying to squeeze GLM-5.2 onto a single 8xH100 node instead ends in CUDA OOM during MoE weight creation. NNODES defaults to 1 so single-node runs are unchanged.
The MLA/MHA flag branch landed above GROUP_SIZE's default, so under set -u every MHA run died with 'GROUP_SIZE: unbound variable' before launching a server. Introduced when the multi-node block was inserted.
The script unconditionally ran 'conda activate oscar', so a model needing a different environment could not use it. Gemma-4 requires transformers >= 5.5 for Gemma4TextConfig; its wrapper activates a dedicated venv, which this then reverted to the 5.3 env. sglang swallowed the resulting import failure as 'Ignore import error when loading ... gemma4_unified' and later reported 'Gemma4UnifiedForConditionalGeneration is not a registered model' -- an error that points at the model registry rather than at the environment.
The wrapper already passed --max-total-tokens but defaulted to 262144, which OOMs an 80 GB card during model load: INT2 KV is compact enough that the pool sizes itself very large, and the second geometry group's arena lands on top of it. The recipe notes carry 65536; make that the default so the example runs as shipped.
The wrapper already passed --prefill-attention-backend triton, but the int2 prefill path keys off the global --attention-backend, which still defaulted to fa3 -- and FA3 int2 prefill asserts on sliding-window layers, of which Gemma-4 has 40. Same shape as the Blackwell case: sglang treats attention_backend, prefill_attention_backend and decode_attention_backend as three separate fields and different code paths read different ones.
The multi-node support I added ran the whole flow on every node, so rank 1 waited on an HTTP health endpoint it never exposes, gave up, and tore the process group down -- which rank 0 reported as DistNetworkError: Failed to recv, got 0 bytes, pointing at the network rather than at the harness. Non-zero ranks now wait on their server process and exit with it.
eval_oscar_gpqa.sh hardcoded --prefill-attention-backend fa3, so no caller could run a sliding-window model through it: FA3's int2 prefill asserts on local attention, which is 40 of Gemma-4's 48 layers. Appending to EXTRA_SERVER_ARGS does not help because the hardcoded flag comes first. Both the global and prefill backends are now ATTN_BACKEND / PREFILL_BACKEND (default fa3, unchanged for everyone else), and Gemma-4's wrapper sets them to triton. The global one matters because the int2 prefill path reads it, not the prefill-specific field.
eval_oscar_gpqa.sh never set SGLANG_LLOYD_MAX, so the quantizer was whatever the environment or the code default supplied, while run_bench_matrix.sh pins it to 0. Same model, same per-head rotation, same sampling: bench_matrix scores 56.57 on Qwen3-30B-A3B GPQA and this script scored 37.37 -- a 19-point gap from one unset variable, in the direction that makes per-head rotations look broken. LLOYD_MAX now defaults to 0 and is overridable for the models whose recipe wants it on (Qwen3.5-35B-A3B, GLM-5.2).
Its recipe uses Lloyd-Max -- that is what took AIME from -14.4 to -4.6 -- but the wrapper never set it, and the shared driver now correctly defaults it off. Same class as the 30B gap: the recipe lived in the notes, not in the example.
The fake-int2 quantizer computes in float32 and returns bf16, but on an FP8 checkpoint k_pe stays Float8_e4m3fn, and torch.cat will not promote across those two -- GLM-5.2-FP8 died at the first decode with 'Promotion for Float8 Types is not supported'. Cast the quantized half back to c_kv's own dtype before the concat. Only reachable on an FP8-weight MLA model, which is why earlier GLM-5.2 attempts never got here: they failed in argument validation or multi-node setup first.
My previous fix made the halves concatenable by casting the quantized c_kv down to the checkpoint's fp8 dtype. That turns one lossy step into two: GLM-5.2-FP8 then answered GPQA with one or two tokens per question and scored 5.6 %, well under the 25 % a four-way guess gets. Promote k_pe to the quantized half's dtype instead -- the buffer setter re-casts once anyway. Turning a crash into silent garbage is worse than the crash; the score is what caught it.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Brings OSCAR 2-bit KV-cache quantization from one dense-MHA model to a set of production model families, with the rotation-fitting recipes needed to reproduce each one.
Per-(layer, KV head) rotations (
format_version: 2). A single per-layer rotation shared across KV heads does not hold on some MoE models. On Qwen3-30B-A3B the 4 KV heads want near-orthogonal rotations (mean |diag| of R₀ᵀR₁ ≈ 0.07), and a shared rotation does not cost a few points, it collapses the model. WikiText-2 perplexity against BF16: shared +16.6 %, per-head +0.35 %. V1 files keep loading unchanged — dispatch is onrotation.dim().Heterogeneous KV geometry. Gemma-4-12B-it has 40 sliding-window layers (8×256) and 8 full-attention layers (1×512); the pool now carries two geometry groups and the decode path is windowed per layer.
MLA latent quantization for GLM-5.2-style shared-latent
c_kv, with a packed INT2 kernel that reproduces the fake-quant simulation to relL2 3.6e-07 on the production path.Expanded-MHA storage for Kimi-K3, whose weights are MLA but whose cache is stored as regular MHA with asymmetric K=192 / V=128 — so it uses the real packed INT2 pool rather than the latent path.