microsoft / microsoft/onnxruntime
Feature request: quantized (int8/fp8) KV cache in GroupQueryAttention with past/present shared buffer + CUDA graph capture
- Dominant language
- C++
- Stars
- 21.9k
- Forks
- 4.2k
- Avg merge
- 4d 11h
- Merged PRs (30d)
- 184
Description
### Summary
Feature request: allow the **KV cache (`past_key_values` / `present`) to be stored quantized** (per-token int8 or fp8) in the `GroupQueryAttention` contrib op when using the past/present **shared buffer** (`past_present_share_buffer=1`) together with `IoBinding` and `enable_cuda_graph`, with the kernel dequantizing on read (and quantizing the newly-appended token on write).
### Motivation
In the onnxruntime-genai-style decode path, the runtime owns a single fixed-capacity KV buffer allocated once as `[batch, kv_heads, max_len, head_dim]` and bound in place as both `past_key_values.*` (input) and `present.*` (output). On memory-constrained GPUs the **fp16 KV cache dominates VRAM** — for long `max_len` it can exceed the model weights themselves. Quantizing KV to int8 (or fp8) halves/quarters that footprint.
**Per-token (per-position) quantization** is particularly attractive here because each token's K/V is scaled independently, so appending a new token never requires re-quantizing existing entries — it preserves the append-only, in-place aliasing that the shared-buffer + CUDA-graph path relies on. The quant/dequant themselves are fixed-shape elementwise ops, so they should be CUDA-graph-capturable.
### Current limitation
Today `GroupQueryAttention` consumes fp16 `past_key_values` and produces fp16 `present`. To use a quantized KV store behind this kernel, a caller would have to **materialize a full fp16 `past` before every decode step** (the kernel reads the entire past each step), which is an O(context) dequant per token plus a full fp16 copy — defeating both the memory saving and the in-place aliasing. There is no way to feed quantized KV directly to the attention kernel.
### Request / questions
1. Is quantized KV-cache support in `GroupQueryAttention` (or a variant) existing or planned?
2. Would a **fused dequant-attention** that reads int8/fp8 `past` (with per-token scales) and appends a quantized new token be in scope, so the shared-buffer + CUDA-graph decode path can keep KV quantized end-to-end?
3. If not planned, is there a recommended pattern (e.g. a quantized-KV attention op, or a documented way to bind quantized `past`/`present` tensors) to achieve this without a per-step fp16 materialization?
### Environment
- Decode path: past/present shared buffer, `IoBinding`, `enable_cuda_graph`, CUDA EP, GQA contrib op, fp16 KV.
- Use case: single-sequence captured decode on memory-limited GPUs (e.g. 8 GB), where fp16 KV sized to the model's full context length is the dominant VRAM consumer.
Contributor guide
Research direction
Start by reviewing the GroupQueryAttention contrib op and its past_key_values/present handling, then trace how IoBinding, past_present_share_buffer, and enable_cuda_graph interact in the CUDA EP. Determine the design for quantized per-token KV storage, fused dequantization, and quantized appends; done requires an agreed implementation scope and validation of shared-buffer CUDA-graph decode without per-step fp16 materialization.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp
- Domain
- machine-learning, performance
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Active
- Clarity
- Needs clarification
- Newbie friendliness
- 25/100