Qwen3-4B long-CoT, pause-code compressed with a static raw tail (tail-16)
Compresses the KV of each completed 512-token block of a long chain-of-thought into 128 learned code rows, and additionally keeps the last 16 raw content KVs of that block permanently beside them — 144 rows per block, a 3.6:1 reduction against the full cache.
Trained from Qwen/Qwen3-4B-Base on long chain-of-thought math, 2.0e9 response tokens.
This is not a drop-in Qwen3
Unlike a plain SFT, loading these weights with from_pretrained and generating does not give you
the compressed behaviour. You get a Qwen3 with extra code-token embeddings and a full KV cache.
The compression lives in the serving path — block-wise code emission, the nested-prefix quota mask
and the static tail — not in the weights alone. Reproducing the numbers below needs the CompressLM
vLLM overlay, served with FLASHINFER and the FA2 prefill pin (on sm90 FlashInfer's auto backend
resolves to FA3 on the first prefill, and FA3 has no custom-mask support, so the first code step
fails).
compresslm_config.json in this repo records the serving configuration:
block_size 512, num_code_tokens 128, retain_tail_kv 16, retain_tail_policy static,
code_positions bitrev, packed_kv false.
Evaluation
pass@1 over all passes (equivalently, mean per-generation accuracy). AIME at 8 samples, temperature 0.6 / top-p 0.95 / top-k 20, seed 1234; MATH-500 greedy at 1 sample.
| Benchmark | Passes | pass@1 (all passes) | pass@k |
|---|---|---|---|
| AIME 2025 | 8 | 29.17 | 53.33 |
| AIME 2026 | 8 | 35.42 | 66.67 |
| MATH-500 (greedy) | 1 | 86.20 | – |
Held-out response CE 0.4213.
For reference, under the same protocol the uncompressed full-context SFT of the same base
(jackcai1206/qwen3-4b-longcot-sft-vanilla) scores AIME 2025 32.92 and MATH-500 87.4. On 30-problem
AIME sets the standard error is about ±6, so the AIME difference is well inside noise; the MATH-500
gap of 1.2 points is the more meaningful comparison.
Precision
Exported bf16 from fp32 master weights, with only the first 128 code rows retained — rows at or above the trained quota are never attended under nested-prefix quotas, and the bit-reversal grid of the first q rows is independent of K. This is the exact artifact the numbers above were measured on.