GQA In-Place KV Cache: Investigation Report

[justinchuby]

GQA In-Place KV Cache: Investigation Report

Executive Summary

The com.microsoft.GroupQueryAttention (GQA) operator supports in-place KV cache update where past_key and present_key share the same pre-allocated buffer at runtime. This is NOT an ONNX graph-level mechanism — it's a runtime optimization managed by ORT GenAI via IO binding, controlled by the past_present_share_buffer flag in genai_config.json.

Mobius gap: The graph patterns mobius emits are already correct. The only actionable gap is that past_present_share_buffer defaults to False when it should be True for GQA models. There are also minor attribute gaps (local_window_size not propagated through GQA rewrite).

---

Part 1: GQA's KV Cache Mechanism

1.1 Operator Spec Overview

Source: onnxruntime/docs/ContribOperators.md#GroupQueryAttention

Inputs (7-14):

#	Name	Type	Description
0	query	T	(B, S, hidden) or packed QKV
1	key	T (opt)	(B, kv_S, kv_hidden)
2	value	T (opt)	(B, kv_S, kv_hidden)
3	past_key	T_CACHE (opt)	BNSH format: (B, kv_heads, cache_S, head_dim)
4	past_value	T_CACHE (opt)	BNSH format: (B, kv_heads, cache_S, head_dim)
5	seqlens_k	M (req)	(B,) — equivalent to total_seq_len - 1 per item
6	total_sequence_length	M (req)	Scalar — max past_seq + new_seq across batch
7	cos_cache	T (opt)	(max_S, head_dim/2)
8	sin_cache	T (opt)	(max_S, head_dim/2)
9	position_ids	int64 (opt)	(B, S)
10	attention_bias	T (opt)	`(B
11	head_sink	T (opt)	(num_heads,)
12	k_scale	T_KV_SCALE (opt)	KV cache quantization scale
13	v_scale	T_KV_SCALE (opt)	KV cache quantization scale

Outputs (3-4):

#	Name	Type	Description
0	output	T	(B, S, hidden)
1	present_key	T_CACHE	BNSH — see buffer modes below
2	present_value	T_CACHE	BNSH — see buffer modes below
3	output_qk	T (opt)	QK values for debugging

1.2 Buffer Sharing Modes

The GQA spec describes two distinct KV cache modes:

Mode A — Separate Buffers (no sharing):

• past_key shape: (B, kv_heads, past_seq_len, head_dim)
• present_key shape: (B, kv_heads, total_seq_len, head_dim) where total_seq_len = past_seq_len + kv_seq_len
• The kernel internally concatenates past_key with new keys to produce present_key
• Each decode step allocates a new, larger buffer

Mode B — Shared Buffer (in-place update):

• past_key and present_key are the SAME tensor at runtime
• Shape: (B, kv_heads, max_sequence_length, head_dim) — pre-allocated to max length
• The kernel writes new KV entries at the correct position (determined by seqlens_k)
• Zero-copy: no allocation, no concatenation per decode step

1.3 Key Attributes

Attribute	Type	Description
num_heads	int (req)	Number of Q heads
kv_num_heads	int (req)	Number of KV heads (GQA ratio)
scale	float	1/√head_dim
do_rotary	int	Fuse RoPE into kernel (0/1)
rotary_interleaved	int	RoPE interleaving pattern (0/1)
local_window_size	int	Sliding window size (-1 = disabled)
softcap	float	Logit soft-capping (0 = disabled)
smooth_softmax	int	Smooth factor in softmax

Note: There is NO kv_share_buffer attribute on the GQA op itself. Buffer sharing is a runtime behavior.

1.4 How ORT GenAI Manages KV Buffers

The ORT GenAI runtime (C++ side) manages buffer sharing through:

1. genai_config.json → search.past_present_share_buffer: true
2. At initialization, the runtime allocates KV buffers of size (B, kv_heads, max_seq_len, head_dim)
3. Via IO binding, both the past input and present output of each GQA node are bound to the same buffer
4. The GQA kernel uses seqlens_k to know where to write new KV entries
5. After each decode step, seqlens_k is incremented — no reallocation needed

This is the standard pattern for efficient autoregressive inference.

---

Part 2: ORT GenAI Model Builder Reference

2.1 GQA Selection Logic

File: ort-genai/src/python/py/models/builders/base.py

def make_attention_init(self):
    if self.is_gqa_supported():
        self.attention_attrs["op_type"] = "GroupQueryAttention"
        self.attention_attrs["use_packed_matmul"] = (
            self.ep not in ["dml"]
            and not self.matmul_attrs["use_lora"]
            and not self.attention_attrs["q_norm"]
            and not self.attention_attrs["k_norm"]
        )
        self.attention_attrs["use_rope_in_attn"] = self.ep not in ["dml"]
        if self.attention_attrs["use_rope_in_attn"]:
            del self.input_names["position_ids"]  # GQA handles RoPE internally
    # ...
    self.past_present_share_buffer = (
        self.attention_attrs["op_type"] == "GroupQueryAttention"
    )

Key insight: GQA always enables past_present_share_buffer. This is not optional.

2.2 GQA Node Creation

def make_group_query_attention(self, name, **kwargs):
    inputs = [
        kwargs["q_path"],
        kwargs["k_path"],
        kwargs["v_path"],
        kwargs.get("past_k", ""),        # past_key_values.{i}.key
        kwargs.get("past_v", ""),        # past_key_values.{i}.value
        kwargs.get("seqlens_k", ""),     # Computed from attention_mask
        kwargs.get("total_seq_len", ""), # Computed from attention_mask
        kwargs.get("cos_cache", ""),
        kwargs.get("sin_cache", ""),
        "",  # position_ids (unused when do_rotary=1)
        "",  # attention_bias
        kwargs.get("sinks", ""),
    ]
    output = f"{name}/output_0"
    outputs = [output, kwargs.get("present_k", ""), kwargs.get("present_v", "")]
    self.make_node(
        "GroupQueryAttention",
        inputs=inputs, outputs=outputs,
        name=name, domain="com.microsoft",
        num_heads=self.num_attn_heads,
        kv_num_heads=self.num_kv_heads,
        scale=self.attention_attrs["scale"],
        local_window_size=self.window_size,    # ← Set from config
        softcap=self.attention_attrs["softcap"],
        do_rotary=self.attention_attrs["use_rope_in_attn"],
        rotary_interleaved=self.rope_attrs["interleaved"],
    )

2.3 genai_config Generation

genai_config = {
    "model": {
        "decoder": {
            "inputs": {
                "past_key_names": "past_key_values.%d.key",
                "past_value_names": "past_key_values.%d.value",
            },
            "outputs": {
                "present_key_names": "present.%d.key",
                "present_value_names": "present.%d.value",
            },
        },
    },
    "search": {
        "past_present_share_buffer": self.past_present_share_buffer,  # True for GQA
    },
}

2.4 KV Cache Shapes

self.input_shapes = {
    "past_key_values.key": ["batch_size", num_kv_heads, "past_sequence_length", head_size],
    "past_key_values.value": ["batch_size", num_kv_heads, "past_sequence_length", head_size],
}
self.output_shapes = {
    "present.key": ["batch_size", num_kv_heads, "total_sequence_length", head_size],
    "present.value": ["batch_size", num_kv_heads, "total_sequence_length", head_size],
}

Note: Shape dimension names are symbolic — actual buffer allocation at runtime uses max_sequence_length when buffer sharing is enabled.

---

Part 3: Mobius Current State

3.1 KV Cache Architecture

Mobius has two KV cache modes:

Dynamic KV Cache (default)

Graph inputs: past_key_values.{i}.key — shape [B, kv_heads, past_seq, head_dim]
Graph outputs: present.{i}.key — shape inferred during optimization

Flow:

past_key_values.0.key ──┐
                         ├──→ Attention op ──→ present.0.key
past_key_values.0.value ─┘                 └──→ present.0.value

After GQA rewrite:

past_key_values.0.key ──┐
                         ├──→ GroupQueryAttention ──→ present.0.key
past_key_values.0.value ─┘   (+ seqlens_k,       └──→ present.0.value
                               total_seq_len,
                               cos_cache, sin_cache)

Files:

• src/mobius/tasks/_causal_lm.py — creates graph inputs/outputs
• src/mobius/tasks/_base.py:138-154 — _register_kv_cache_outputs() names outputs
• src/mobius/components/_attention.py — _apply_attention() calls op.Attention

Static KV Cache (static_cache=True)

Graph inputs: key_cache.{i}, value_cache.{i}, write_indices, nonpad_kv_seqlen
Graph outputs: updated_key_cache.{i}, updated_value_cache.{i}

Uses TensorScatter for in-place writes to pre-allocated buffers. This is an alternative in-place mechanism that does NOT use GQA's built-in buffer sharing.

3.2 GQA Rewrite Rule

File: src/mobius/rewrite_rules/_group_query_attention.py

Three rewrite rules:

1. RotaryAttentionToGQA (primary) — matches RotaryEmbedding(q) + RotaryEmbedding(k) + Attention pattern
2. AttentionToGQA (fallback) — matches any Attention with graph-input past_key/past_value
3. PackQKVForGQA — fuses separate Q/K/V MatMuls into packed QKV

GQA attributes set:

{
    "num_heads": q_num_heads,
    "kv_num_heads": kv_num_heads,
    "scale": scale,
    "do_rotary": 1,                    # RotaryAttentionToGQA only
    "rotary_interleaved": interleaved,  # RotaryAttentionToGQA only
    "softcap": softcap,                # If non-zero (Gemma2)
}

seqlens_k computation (from attention_mask):

seqlens_k = Cast(Sub(ReduceSum(attention_mask, axis=1), 1), to=INT32)
total_seq_len = Cast(Gather(Shape(attention_mask), 1), to=INT32)

3.3 genai_config Generation

File: src/mobius/integrations/ort_genai/genai_config.py

def _default_search_params():
    return {
        "past_present_share_buffer": False,  # ← DEFAULT IS FALSE
        # ...
    }

def _default_decoder_inputs(*, is_vlm: bool):
    return {
        "past_key_names": "past_key_values.%d.key",
        "past_value_names": "past_key_values.%d.value",
    }

def _default_decoder_outputs():
    return {
        "present_key_names": "present.%d.key",
        "present_value_names": "present.%d.value",
    }

---

Part 4: Gap Analysis

4.1 What Mobius Does Right ✅

Aspect	Status	Details
KV cache I/O naming	✅ Correct	past_key_values.%d.key / present.%d.key matches ORT GenAI
GQA node structure	✅ Correct	3 outputs (attn_out, present_key, present_value)
seqlens_k computation	✅ Correct	Derived from attention_mask correctly
total_seq_len computation	✅ Correct	Derived from attention_mask shape
cos/sin cache sharing	✅ Correct	Shared across layers via cached values
softcap attribute	✅ Correct	Propagated from Attention → GQA
rotary_interleaved	✅ Correct	Propagated from RotaryEmbedding → GQA
Graph pattern	✅ Correct	past_key → GQA → present_key flow is correct

4.2 Gaps ⚠️

Gap	Severity	Details
past_present_share_buffer defaults to False	High	Should be True when GQA is used. Without this, ORT GenAI allocates new buffers per decode step → significant memory waste and performance loss
local_window_size not set on GQA node	Medium	The rewrite rule doesn't propagate sliding window config. Models like Mistral/Gemma2 (which use sliding window) won't get the optimization. ORT GenAI builder sets local_window_size=self.window_size
No position_ids input removal for fused RoPE	Low	When do_rotary=1 (RoPE fused into GQA), position_ids graph input is unnecessary. ORT GenAI builder removes it. Mobius keeps it — harmless but wasteful
No attention_mask removal for GQA	Low	GQA uses seqlens_k/total_seq_len instead of attention_mask. The mask input could be removed after GQA rewrite. ORT GenAI builder omits it entirely

4.3 Not Gaps (Design Differences)

Aspect	Mobius	ORT GenAI Builder	Assessment
Static cache mode	TensorScatter-based	Not supported	Mobius's static cache is for opset-24 Attention op, orthogonal to GQA buffer sharing
Graph-level buffer aliasing	Not done	Not done either	Both rely on runtime IO binding for actual buffer sharing
KV cache output shapes	Inferred	Symbolic dims	Both valid — shape inference handles it

---

Part 5: Specific Changes Needed

Change 1: Enable past_present_share_buffer for GQA (High Priority)

File: src/mobius/integrations/ort_genai/genai_config.py

The GenaiConfigGenerator should detect whether the model uses GQA and set past_present_share_buffer: True.

Option A — Detect from ONNX model:
After optimization (which applies GQA rewrite), inspect the model for GroupQueryAttention nodes. If found, set the flag.

Option B — Detect from EP capabilities:
If the EP optimization pipeline applies GQA rewrite, set the flag. This can be determined from the EP registry.

Option C — Pass through from build pipeline:
build() → optimize_model() → returns whether GQA was applied → propagate to genai_config.

Recommended: Option A — simplest and most reliable. Something like:

# In GenaiConfigGenerator.generate():
has_gqa = any(
    node.op_type == "GroupQueryAttention"
    for node in model.graph
)
search_params["past_present_share_buffer"] = has_gqa

Change 2: Propagate local_window_size to GQA (Medium Priority)

File: src/mobius/rewrite_rules/_group_query_attention.py

The rewrite rule should accept a local_window_size parameter and set it on the GQA node. This requires the optimization pipeline to pass model config (sliding window size) to the rewrite rule.

Approach:

1. Add local_window_size: int = -1 parameter to RotaryAttentionToGQA.__init__
2. Set gqa_attrs["local_window_size"] = self.local_window_size
3. In _get_optimization_passes(), read sliding_window from the architecture config and pass it to the rule

Change 3 (Optional): Remove unused inputs after GQA rewrite

After GQA rewrite, position_ids (when do_rotary=1) and attention_mask could be removed from graph inputs since GQA doesn't use them directly. However, attention_mask IS used to compute seqlens_k/total_seq_len, so it must remain. Only position_ids removal is safe when RoPE is fused.

This is low priority — the inputs are simply ignored by the runtime.

---

Part 6: EP Capability Considerations

EP	GQA Support	Buffer Sharing	local_window_size	Notes
CUDA	✅ (fp16, bf16)	✅	✅	Primary target
CPU	✅ (fp32)	✅	✅	For testing/debug
DML	✅ (fp16)	✅	❓	No packed QKV, no fused RoPE
WebGPU	✅ (fp16, fp32)	✅	❓	Limited feature set
TRT-RTX	✅ (fp16, bf16)	✅	✅	Uses SparseAttention for some models

All EPs that support GQA also support buffer sharing. The past_present_share_buffer flag should be EP-agnostic.

---

Appendix: Architecture Diagram

┌─────────────────────────────────────────────────────────────────────────┐
│                        ONNX Model Graph                                │
│                                                                        │
│   Graph Inputs:                        Graph Outputs:                  │
│   ├─ input_ids [B, S]                 ├─ logits [B, S, V]             │
│   ├─ attention_mask [B, total_S]      ├─ present.0.key [B, H, ?, D]   │
│   ├─ position_ids [B, S]             ├─ present.0.value [B, H, ?, D] │
│   ├─ past_key_values.0.key           ├─ present.1.key ...             │
│   ├─ past_key_values.0.value         └─ ...                           │
│   ├─ past_key_values.1.key                                            │
│   └─ ...                                                               │
│                                                                        │
│   ┌─────────────────────────────────────────────────────┐              │
│   │  Layer 0:                                           │              │
│   │  Embedding → LayerNorm → GQA → MLP → LayerNorm     │              │
│   │                           │                         │              │
│   │  past_key_values.0.key ───┤                         │              │
│   │  past_key_values.0.value ─┤                         │              │
│   │  seqlens_k ───────────────┤  GroupQueryAttention    │              │
│   │  total_seq_len ───────────┤  (do_rotary=1,         │              │
│   │  cos_cache ───────────────┤   kv_num_heads=N,      │              │
│   │  sin_cache ───────────────┘   local_window_size=W) │              │
│   │                           │                         │              │
│   │                    ┌──────┴──────┐                  │              │
│   │                    │present.0.key│                  │              │
│   │                    │present.0.val│                  │              │
│   │                    └─────────────┘                  │              │
│   └─────────────────────────────────────────────────────┘              │
└─────────────────────────────────────────────────────────────────────────┘

┌─────────────────────────────────────────────────────────────────────────┐
│                     ORT GenAI Runtime (C++)                             │
│                                                                        │
│   When past_present_share_buffer = true:                               │
│                                                                        │
│   ┌─────────────────────────────────────┐                              │
│   │  Pre-allocated KV Buffer            │                              │
│   │  Shape: [B, H, max_seq, D]         │                              │
│   │                                     │                              │
│   │  IO Bind: past_key_values.0.key ───►│                              │
│   │  IO Bind: present.0.key ───────────►│  ← Same buffer!             │
│   │                                     │                              │
│   │  Step 0: Write at position 0        │                              │
│   │  Step 1: Write at position 1        │                              │
│   │  Step N: Write at position N        │                              │
│   │  (seqlens_k tracks current pos)     │                              │
│   └─────────────────────────────────────┘                              │
│                                                                        │
│   When past_present_share_buffer = false:                              │
│                                                                        │
│   Step 0: Alloc [B,H,1,D], copy                                       │
│   Step 1: Alloc [B,H,2,D], copy                                       │
│   Step N: Alloc [B,H,N+1,D], copy  ← O(N²) memory, O(N) per step     │
│                                                                        │
└─────────────────────────────────────────────────────────────────────────┘

---

Summary of Recommendations

1. [High Priority] Set past_present_share_buffer: True in genai_config when GQA nodes are present
2. [Medium Priority] Propagate local_window_size from model config through GQA rewrite rule
3. [Low Priority] Consider removing position_ids from graph inputs when GQA fuses RoPE
4. [No Change Needed] Graph I/O patterns and GQA node structure are already correct