Add Reasoning-Aware Compression (RAC) pruning recipe for reasoning models

[PKUWZP]

Co-authored with Ryan Lucas (MIT) and Kayhan Behdin (LinkedIn).

Motivation

Compressing reasoning models with standard pruning does more damage than compressing a conventional LLM, and it can make the model slower.

One-shot pruning picks weights by minimizing a layer-wise reconstruction error against a calibration activation matrix X:

min_{W'} || W X - W' X ||_F^2    s.t.  ||W'||_0 <= S

X is conventionally built from prompt tokens (C4, or task prompts). That is a fair proxy when the prompt dominates the token count. Reasoning models invert the ratio: they emit thousands of chain-of-thought tokens per query, so nearly every forward pass the pruned model will ever run is over a token it generated itself. Calibrating on prompts alone optimizes the pruned weights for a distribution the model barely visits.

The resulting failure mode is worse than a graceful accuracy drop — the pruned model rambles, emitting more thinking tokens for a less accurate answer, so pruning increases end-to-end latency. From the paper (DeepSeek-R1-Distill-Qwen-7B, MATH-500, SparseGPT @ 50% sparsity, 1M calibration tokens):

Calibration set	acc@1	Eval wall clock
Dense (no pruning)	0.936	23.3 min
C4	0.744	135.0 min
Task prompts only	0.812	115.6 min
RAC (prompts + on-policy CoT)	0.900	35.3 min

Reasoning-Aware Compression (RAC) fixes this by calibrating on the dense model's own on-policy rollout, reconstructing prompt and decode activations jointly:

X_RAC = [ X_prompt , X_decode ]

The solver is untouched, so RAC is a drop-in calibration-set swap for any existing SparseGPT/Wanda workflow.

From Reasoning Models Can be Accurately Pruned Via Chain-of-Thought Reconstruction (Lucas, Behdin, Wang, Tang, Song, Mazumder; ICLR 2026). Reference implementation: RyanLucas3/Reasoning-Aware-Compression.

Why this belongs in SGLang, and what deliberately does not

Collecting the rollout is Phase I of the paper's Algorithm 1 and is the expensive half — the paper's budget is 1M on-policy CoT tokens per calibration set. That is batched autoregressive generation, which is what SGLang is for; the reference implementation does it with a Hugging Face generate loop.

The pruning solver itself is training-time code and is not proposed for python/sglang/srt/. This PR lands as an offline recipe under examples/usage/, matching the existing examples/usage/modelopt_quantize_and_export.py precedent, and delegates the solver to llm-compressor:

Phase	Script	Role
I	rac_collect_traces.py	sgl.Engine samples on-policy CoT → traces.jsonl
II	rac_prune.py	llm-compressor runs SparseGPT/Wanda against those activations
III	rac_serve_and_eval.py	SGLang serves the sparse checkpoint and scores MATH-500

llmcompressor is imported lazily and is not added to SGLang's dependencies — Phases I and III need only SGLang.

Also out of scope, and better as separate PRs: sparse-serving kernels (2:4 / cuSPARSELt runtime paths), and vendoring a SparseGPT solver into the engine.

Modifications

• New examples/usage/reasoning_aware_compression/ — three scripts plus a README with full reproduction commands for the paper's DeepSeek-R1-Distill-Qwen-1.5B @ 50% row.
• New docs/docs/advanced_features/reasoning_aware_compression.mdx, registered in docs/docs.json after the quantization page.

Design points worth reviewer attention:

• Token-in-token-out. Phase I runs the engine with skip_tokenizer_init=True and emits token ids; Phase II consumes them directly. The sequence the pruner reconstructs is exactly the one the model produced, with no detokenize/retokenize drift.
• Lazy, chunked corpus reads. The paper's math corpus is 220k rows while a 1M-token budget touches only a few hundred prompts, so prompts are chat-templated and rolled out per chunk rather than up front.
• Batch size 1 during calibration. Batching variable-length sequences would require padding, and pad-token activations would enter the layer-wise Hessian as if they were real — precisely the calibration contamination RAC exists to avoid.
• --calibration-mode prompt_only reproduces the paper's ablation baseline from the same prompts, so the comparison motivating the method is runnable from the shipped code.
• Phase III reports mean CoT length and wall clock next to accuracy, because accuracy alone hides the rambling failure mode above.
• Magnitude pruning is intentionally not exposed. llm-compressor's magnitude modifier is a gradual training-time modifier, not a one-shot solver; the README says so rather than shipping an option that cannot work.

Accuracy Test

The smoke path a reviewer with a GPU can run in a few minutes:

cd examples/usage/reasoning_aware_compression
pip install "llmcompressor>=0.12.0"

python rac_collect_traces.py --model-path Qwen/Qwen3-0.6B \
    --dataset open-r1/OpenR1-Math-220k --prompt-column problem \
    --target-tokens 20000 --max-new-tokens 1024 --output-dir /tmp/rac_traces
python rac_prune.py --model-path Qwen/Qwen3-0.6B \
    --calibration /tmp/rac_traces/traces.jsonl --sparsity 0.5 --output-dir /tmp/rac_pruned
python rac_serve_and_eval.py --model-path /tmp/rac_pruned --num-problems 50 --max-new-tokens 2048

Phase I should report a decode share well above 50%; Phase II a realized sparsity within a hair of the target.

Checklist

• Format your code according to the Code Formatting with Pre-Commit.
• Add documentation as needed (docs/docs/advanced_features/reasoning_aware_compression.mdx).

---

CI States

Latest PR Test (Base): ✅ Run #31775916472
Latest PR Test (Extra): ❌ Run #31775916423

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[Qiaolin-Yu]

could you fix the conflicts?

[PKUWZP]

could you fix the conflicts?

@Qiaolin-Yu I think someone with the write access should add both labels to this PR:

• run-ci (basic CI prerequisite)
• run-ci-extra (explicit opt-in for AMD extra tests)

[PKUWZP]

could you fix the conflicts?

@Qiaolin-Yu I think someone with the write access should add both labels to this PR:

• run-ci (basic CI prerequisite)
• run-ci-extra (explicit opt-in for AMD extra tests)

I think once the labels are updated, the conflicts should be resolved.