Qwen-Compress

(有提供繁體中文版readme)

Qwen-Compress: Learned Weighted Pooling for KV Cache Compression

A drop-in KV cache compression method that selectively compresses hidden states via learned weighted pooling between every 4th transformer layer, reducing prefill KV size by up to ~1.75× without retraining the base model.

Overview

Standard autoregressive transformers pay a KV cache cost proportional to total prefill length. Qwen-Compress mitigates this by inserting small HiddenStateCompressor modules at selected layers. Each compressor performs learned weighted pooling: it computes an importance score per token via a learned linear projection, then applies softmax over a sliding window of ratio=4 consecutive hidden states and takes the weighted sum, producing 1 compressed token from 4. This shrinks the KV cache while preserving the final window_size=128 tokens unchanged (tail retention).

The key insight: compress hidden states before they enter the attention QKV projection at the compressor layer. Downstream layers attend to compressed KV positions via correct rotary position embedding (RoPE) indices, avoiding position mismatch.

Inspiration

This work is inspired by DeepSeek's compressed attention mechanism (used in DeepSeek-V4), where inter-layer compression reduces KV cache footprint. Qwen-Compress adapts this concept to the Qwen architecture with a per-sequence compression schedule and decode-phase buffer management.

Architecture

Compressor Module

A single HiddenStateCompressor is a lightweight learned weight layer:

Input:  (B, 4, D)  — 4 consecutive hidden states
          │
    ┌─────┴─────┐
    │  wgate    │  Linear(D, 1) → per-token importance score
    │  + ape    │  Learned position bias (4, 1)
    └─────┬─────┘
          │ softmax over 4 tokens → normalized weights
          │
    Weighted sum → (B, 1, D)  — 1 compressed token

Parameters: ~D (one Linear layer + 4 learned scalars) — negligible vs. base model. No latent bottleneck or QKV attention involved.

Placement

Compressors are inserted at layers where layer_idx % 4 == 3 (every 4th layer). At non-compressor layers, attention runs unchanged.

Compression Schedule (Window + Ratio)

For each sequence at a compressor layer:

true_len   = sequence length (after left-padding)
window     = min(window_size=128, true_len)   ← kept verbatim
compressible = true_len - window
to_compress  = compressible - (compressible % ratio=4)
compressed   = to_compress / 4
residual     = compressible % ratio            ← kept verbatim (0–3 tokens)
tail         = window                          ← last 128 tokens verbatim

KV positions are tracked per-token for correct RoPE indexing. Zero-padded batch entries are masked via kv_valid throughout generation.

Decode-Phase Buffer Compression

During autoregressive decode, new tokens accumulate in an uncompressed buffer (buf_k, buf_v, buf_h). Once the buffer reaches window_size + ratio = 132 tokens, the oldest ratio=4 tokens are compressed into 1 and merged into the persistent KV cache, keeping the buffer at window_size tokens. This bounds the decode-phase KV overhead at window tokens per compressor layer.

Performance

Tested on Qwen3.6-27B (NF4, 64 layers). Compression applied at every 4th layer (i % 4 == 3), totaling 16 compressors.

GSM8K (5-shot)

IFEval (0-shot)

Key Takeaways

• IFEval is virtually unaffected — instruction following, formatting, and constraint satisfaction are preserved.
• GSM8K shows degradation — multi-step mathematical reasoning suffers under compression, especially strict format matching (#### N).
• Compression ratio: A 301-token prefill is compressed to ~172 KV entries (1.75×). Effective ratio varies with prompt length (shorter prompts see less benefit; very long prompts approach 4×).

Usage

The reference implementation is a FastAPI server (server.py) that patches a Qwen model with compressors and serves an OpenAI-compatible API.

# Install dependencies
pip install torch transformers accelerate bitsandbytes

# Start the server (loads model, patches compressors, serves on :8001)
python server.py

# Send requests via curl or any OpenAI client
curl http://localhost:8001/v1/chat/completions \
  -H "Content-Type: application/json" \
  -d '{"model": "qwen-compress", "messages": [{"role": "user", "content": "1+1="}]}'

The patching happens in build_franken_qwen() (server.py:344): every 4th layer's attention forward is replaced with get_custom_qwen_forward(compressor, window_size=128, ratio=4), and the compressor weights are loaded from a checkpoint.

Limitations

• Reasoning degradation: Complex multi-step reasoning (GSM8K) shows 16–33% accuracy drop. The compressor likely disrupts intermediate representations critical for arithmetic reasoning, even with thinking mode disabled.
• No training recovery: The compressor checkpoint is trained in isolation, not fine-tuned jointly with the base model. Post-training or LoRA-based recovery may close the gap.
• Context-independent compression: The weighted pooling treats each 4-token window independently. A cross-window or bidirectional design could better preserve global reasoning structure.