Skip to content

Add OSCAR 2-bit KV cache (PER_GROUP) support to CPU GroupQueryAttention - #31726

Open
Hector Li (HectorSVC) wants to merge 4 commits into
microsoft:mainfrom
HectorSVC:kv_cache_2bits
Open

Add OSCAR 2-bit KV cache (PER_GROUP) support to CPU GroupQueryAttention#31726
Hector Li (HectorSVC) wants to merge 4 commits into
microsoft:mainfrom
HectorSVC:kv_cache_2bits

Conversation

@HectorSVC

@HectorSVC Hector Li (HectorSVC) commented Aug 7, 2026

Copy link
Copy Markdown
Contributor

Description

Summary
This PR adds an opt-in 2-bit KV cache path to the GroupQueryAttention (GQA) contrib op on the CPU EP, implementing the OSCAR scheme (per-group asymmetric INT2 codec + a sink/recent high-precision window + optional spectral rotation). It extends the existing quantized-KV-cache machinery (which already supports INT8/INT4) rather than adding a new op, so existing GQA graphs are unaffected. All new behavior is gated behind new attributes / session-config entries and defaults to off.

Reference

OSCAR: Offline Spectral Covariance-Aware Rotation for 2-bit KV Cache Quantization. 28 May, 2026
https://arxiv.org/abs/2605.17757

Motivation and Context

For long-context / large-batch serving, the KV cache dominates memory. INT8→INT4 halves it; going to 2 bits needs outlier handling. OSCAR keeps a small sink+recent window in full precision, quantizes the long tail of history to 2 bits with per-group scale/zero, and (optionally) applies an orthogonal rotation so the history quantizes with much lower error while leaving QK scores invariant.

What's included
New quantization path (k_quant_type/v_quant_type = "PER_GROUP", kv_cache_bit_width = 2)

Per-token, per-group asymmetric INT2 codec (4 codes/byte) with per-group scale + zero-point stored inline in the packed cache row (head_size/4 + num_groups2meta_bytes). Scales are computed dynamically at append time, so no k_scale/v_scale inputs are required.
kv_quant_group_size attribute selects the group size.
k_quant_rho / v_quant_rho: magnitude-percentile outlier clip before computing per-group min/max (1.0 = disabled).
kv_quant_metadata_fp16: store inline scale/zero as fp16 instead of fp32, shrinking the row (e.g. 48B→40B for D=128,G=64, i.e. 3.0→2.5 effective bits/elem) with no measured accuracy loss.
Mixed-precision sink/recent window (OSCAR "Option C")

New optional I/O: inputs past_hp_key/past_hp_value (16, 17), outputs present_hp_key/present_hp_value (4, 5) — the first sink and last recent tokens are kept unquantized (type T), the middle history is 2-bit.
New session-config entries gqa.kv_quant.sink / gqa.kv_quant.recent select the window sizes.
Optional spectral rotation

New optional inputs oscar_rotation_k/oscar_rotation_v (18, 19), per-kv-head orthogonal matrices (kv_num_heads, head_size, head_size) applied to post-RoPE K/Q (and V) before 2-bit quantization; the V output is un-rotated by the transpose.
fp16 compute support

MLFloat16 is registered for GQA; the OSCAR path bridges fp16 Q/K/V (and the hp window) to the float codec and converts outputs back, so an fp16 model can drive the 2-bit cache. (attention_bias and output_qk are not supported on the fp16 2-bit path and return a clear error.)

@azure-pipelines

Copy link
Copy Markdown
Azure Pipelines:
There may be pipelines that require an authorized user to comment /azp run to run.

def _make_rotation(kv_num_heads, head_size, seed):
"""Per-kv-head random orthogonal matrices R [kv_num_heads, head_size, head_size] (float32)."""
rng = np.random.default_rng(seed)
R = np.zeros((kv_num_heads, head_size, head_size), dtype=np.float32)
k_bnsh = key_input.reshape(batch_size, seq_len, kv_num_heads, head_size).transpose(0, 2, 1, 3)
v_bnsh = value_input.reshape(batch_size, seq_len, kv_num_heads, head_size).transpose(0, 2, 1, 3)

R_K = _make_rotation(kv_num_heads, head_size, seed=1)
v_bnsh = value_input.reshape(batch_size, seq_len, kv_num_heads, head_size).transpose(0, 2, 1, 3)

R_K = _make_rotation(kv_num_heads, head_size, seed=1)
R_V = _make_rotation(kv_num_heads, head_size, seed=2)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants