microsoft / microsoft/onnxruntime

Feature request: quantized (int8/fp8) KV cache in GroupQueryAttention with past/present shared buffer + CUDA graph capture

Open
#29,783 2 comments 0 reactions 0 assignees View on GitHub
ep:CUDA quantization stale
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.