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

[justinchuby]

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.