Add OSCAR 2-bit KV cache (PER_GROUP) support to CPU GroupQueryAttention - #31726
Open
Hector Li (HectorSVC) wants to merge 4 commits into
Open
Add OSCAR 2-bit KV cache (PER_GROUP) support to CPU GroupQueryAttention#31726Hector Li (HectorSVC) wants to merge 4 commits into
Hector Li (HectorSVC) wants to merge 4 commits into
Conversation
|
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) |
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
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.
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.)