[Engineering] Same GQA, Different Execution Costs: From Pointer Arithmetic to Tensor Shapes

Hello, I'm Seongjae Park from OptAI.

Part 1 examined repeat_kv left in the Core ML graph, replicating large amounts of KV every token. Rearranging 32 query heads into [8, 4] instead of expanding 8 KV heads to 32 removed 60 tile operations and raised throughput from 20.39 tok/s to 54.30 tok/s with identical quantization. A top8 mixed-precision build keeping 8 sensitive layers at 8-bit and others at 4-bit reached 44 tok/s on iPhone 16 Pro and 56 tok/s on iPhone 17 Pro.

This article goes one level deeper.

All GQA implementations ultimately express the same relationship.

kv_head = query_head // group_size

The means differ by backend: GPU kernels can use pointer arithmetic, graphs can fan a value out to multiple operations, and tensor shapes can encode the mapping directly.

Even with identical mathematics, representation choices can substantially change data movement and graph size.

1. One Mapping Shared by Every Implementation

Exaone4 1.2B has 32 query heads and 8 KV heads, so four query heads share each KV head.

group_size = 32 // 8 = 4

For query-head index h, the corresponding KV head follows.

kv = h // 4

Implementation details do not change this relationship.

The question is not merely which backend supports GQA. What matters is how long sharing semantics survive and which representations the executor can use.

2. Hugging Face repeat_kv: Repeat the Tensor Directly

A conventional reference implementation repeats KV heads to match query-head count.

def repeat_kv(hidden_states, n_rep):
    batch, n_kv_heads, seq_len, head_dim = hidden_states.shape
    hidden_states = hidden_states[:, :, None, :, :].expand(
        batch, n_kv_heads, n_rep, seq_len, head_dim
    )
    return hidden_states.reshape(
        batch, n_kv_heads * n_rep, seq_len, head_dim
    )

The initial expand creates no data. Setting repetition-axis stride to 0 creates a view reading the same storage repeatedly.

The following reshape is the problem. A stride-0 repetition axis cannot merge with the KV-head axis into one contiguous head axis, so eager mode allocates new storage.

K        storage  4.00 MB   shared ptr: True
expand   storage  4.00 MB   shared ptr: True
reshape  storage 16.00 MB   shared ptr: False

Simplicity is the advantage: KV expanded to 32 heads uses the same shapes as ordinary Multi-Head Attention.

The disadvantage is lost semantics. From a tensor expanded to [1, 32, ctx, 64], it is hard to infer the original 8 KV heads, each shared by four query heads. Recovering that relationship later to remove copies also becomes harder.

3. PyTorch SDPA: Carry GQA Information to the Backend

PyTorch scaled_dot_product_attention provides enable_gqa=True. This flag does not itself guarantee copy-free execution.

The documented reference semantics repeat K and V through repeat_interleave.

if enable_gqa:
    key = key.repeat_interleave(
        query.size(-3) // key.size(-3), -3
    )
    value = value.repeat_interleave(
        query.size(-3) // value.size(-3), -3
    )

At execution time, the dispatcher selects a backend based on dtype, shape, mask, device, and capabilities. A fused backend with native GQA support may avoid physically repeating K and V. MATH fallback may incur repetition matching the reference semantics.

The following examples look similar but pass different information to the backend.

# Call SDPA after the GQA relationship has been lost
k = repeat_kv(k, n_rep)
v = repeat_kv(v, n_rep)
out = scaled_dot_product_attention(q, k, v)

# Pass GQA information while retaining the original Q and KV head counts
out = scaled_dot_product_attention(
    q, k, v, enable_gqa=True
)

The first makes original KV sharing hard to infer from inputs alone. The second preserves distinct query/KV head counts, with enable_gqa explaining their relationship.

Compiler optimization often encounters this distinction: lowering meaningful structure too early removes information needed to select an efficient implementation later.

4. Triton and vLLM: Select KV Heads by Address

A custom GPU kernel knows which query head each program handles, so it calculates the KV-head address directly without replication.

The PyTorch Inductor FlexAttention Triton template uses this form.

off_hq = tl.program_id(1) % HQ
off_hkv = off_hq // GQA_SHARED_HEADS

k_offset = (
    off_zkv * stride_kz
    + off_hkv * stride_kh
)

off_hq is the query-head index and off_hkv is the KV-head index to read.

Four programs handling query heads 4, 5, 6, 7 all obtain the following value.

off_hkv = off_hq // 4 = 1

They calculate the same k_offset and read the same K address. Four programs reference one original rather than four replicas.

vLLM Decode: Find the KV Head from the Query Head

The vLLM Triton decode kernel computes in the same direction.

cur_head = tl.program_id(1)
cur_kv_head = cur_head // kv_group_num

base_offs_k = (
    cur_kv_head * stride_buf_kh
    + offs_d[:, None]
)

Each program handles one query head and selects the KV head through integer division.

vLLM Prefill: Expand Query Heads from a KV Head

A prefill kernel can use the same relationship in reverse.

kv_head_idx = tl.program_id(1)
query_head_idx = (
    kv_head_idx * num_queries_per_kv
    + tl.arange(0, num_queries_per_kv_padded)
)

Each program now handles one KV head and computes the four query heads sharing it.

kv = h // 4        # Find the KV head from the query head
h  = kv × 4 + g    # Expand query heads from the KV head

These are forward and inverse forms of one mapping. The preferred direction depends on query length, parallelization granularity, and K/V reuse.

5. CoreAI iOS SDPA: Share One Value Across Operations

Graphs without direct pointer arithmetic can avoid replication by fanning one KV head out to multiple multiplication inputs.

The CoreAI iOS path conceptually uses this structure.

kv_group_size = len(queries) // len(keys)

for head_idx in range(n_heads):
    kv_idx = head_idx // kv_group_size

    q = queries[head_idx]
    k = keys[kv_idx]
    score = q @ k

Multiplications for query heads 0~3 all take keys[0]. One K value connects to four multiplications without making four copies.

This naturally expresses sharing in a value-based graph. Expanding heads individually, however, grows the graph. With 32 query heads, each layer requires 32 score multiplications and 32 output multiplications, plus transpose, permute, and concat to separate and recombine heads.

Fan-out avoids copies while increasing graph size and operator count.

6. Tensor Shapes Without Custom Kernels

We wanted standard batched multiplication in Core ML rather than separate heads. Tensor shape was the principal means available.

The decode Q shape follows.

[B, Hq, 1, D] = [1, 32, 1, 64]

Split head axis 32 into 8 KV heads and group size 4.

[B, Hq, 1, D]
       ↓ reshape
[B, Hkv, G, D] = [1, 8, 4, 64]

In row-major order, this reshape gives the following relationship.

h = kv × G + g
kv = h // G
g = h % G

Shape [Hkv, G] encodes the same relation ahead of time without evaluating h // G at runtime.

q = query_states.reshape(
    batch_size,
    num_kv_heads,
    group_size,
    head_dim,
)

scores = torch.matmul(q, k.transpose(-1, -2))
prob = softmax(scores, dim=-1)
out = torch.matmul(prob, v)

One batched multiplication handles 8 KV heads, placing the four query heads sharing each KV head as rows inside its batch.

PyTorch Flex Decoding Rearranges Q in the Same Direction

The short-query-specific PyTorch Inductor Flex Decoding path also rearranges Q before its kernel call.

# [B, Hq, Mq, D] -> [B, Hkv, G, Mq, D]
gqa_query_shape = (
    B,
    Hkv,
    gqa_shared_heads,
    seq_len_q,
    qk_head_dim,
)

query = lowerings[aten.as_strided](
    query,
    gqa_query_shape,
    gqa_query_stride,
)

PyTorch creates an as_strided view for a dedicated Triton kernel. We pass rearranged Q directly into standard batched multiplication.

The rearrangement direction matches, but execution mechanisms differ.

During decode, query length M is 1, leaving one row on M. Arranging the G shared query heads as rows naturally expresses K/V reuse.

7. Merge Groups with the Sequence Axis in Prefill

Decode with q_len=1 can represent Q as [B, Hkv, G, D]. With larger q_len, applying the same idea to prefill produces rank-5 tensors.

For q_len=64, the shapes follow.

Q  : [1, 8, 4, 64, 64]
Kᵀ : [1, 8, 1, 64, 4096]

This multiplication requires 4 to 1 broadcasting on the group axis. It is logically valid, but a backend may materialize K expanded along that axis, reintroducing a similar copy elsewhere.

Instead, merge group axis G into query sequence axis M rather than retaining it as a separate batch axis.

[B, Hq, M, D]
= [B, Hkv, G, M, D]
       ↓ Merge G and M into a single row axis
[B, Hkv, G×M, D]

With 32 query heads and prefill length 64, this becomes the following.

[1, 32, 64, 64]
       ↓ reshape
[1, 8, 256, 64]

Let r denote row index, g the query head within a group, and s sequence position.

r = g × q_len + s

The 64 tokens of four query heads sharing a KV head become 256 consecutive rows.

Restore the Original Shape to Add the Mask

The multiplication result has the following shape.

[B, Hkv, G×M, ctx]

A causal mask is normally [B, 1, M, ctx] or broadcastable to it. Adding it directly misaligns the folded G and M semantics.

Repeating the mask G times can recreate data repetition in the mask path. Restore scores to the original query-head shape before adding the mask instead.

q = query_states.reshape(
    bsz,
    n_kv,
    n_rep * q_len,
    self.head_dim,
)

k = K_layer_cache.reshape(
    bsz, n_kv, -1, self.head_dim
)
v = V_layer_cache.reshape(
    bsz, n_kv, -1, self.head_dim
)

scores = torch.matmul(q, k.transpose(-1, -2))

# [B, Hkv, G×M, ctx] -> [B, Hq, M, ctx]
scores = scores.reshape(
    bsz,
    self.num_heads,
    q_len,
    -1,
)

scores = self._score_scale_mul(scores)
scores = scores + causal_mask[..., :q_len, :ctx]
prob = self.ANE_softmax(scores, dim=-1)

# Merge G and M again immediately before the V matrix multiplication.
prob = prob.reshape(
    bsz,
    n_kv,
    n_rep * q_len,
    -1,
)

out = torch.matmul(prob, v)
out = out.reshape(
    bsz,
    self.num_heads,
    q_len,
    self.head_dim,
)

The group can merge only into row axis M.

• Merging into head_dim sums dot products from different heads together.
• Merging into context changes the softmax normalization domain.
• Merging into M preserves independent multiplication and softmax for every row.

The wrong axis may still pass shape checks. Validate head mapping and output equality as well as shapes.

Why Decode Was Optimized First

The original decode and prefill graphs both executed 60 tile operations per call, but handled different numbers of tokens.

If prefill reads K/V at maximum context length, tile size per call remains unchanged. A call handles 64 tokens, however, making per-token cost 1/64.

This motivated optimizing decode first. Prefill Q rearrangement reduced time from 1,393 ms to 1,219 ms, or 12.5%. Unlike decode, this was not an identical-build A/B changing only one flag, so treat it as a reference result.

8. Choosing a Representation for the Backend

This case yields five principles.

1. Preserve Meaningful Structure as Long as Possible

Expanding KV to 32 heads through repeat_kv loses the fact that 8 KV heads are shared by four query heads each. Retain original head counts where possible and carry GQA semantics to the executor.

2. Rearrange the Smaller Tensor

Q is 4 KiB during decode, while one layer K or V cache is 4 MiB. Reshaping smaller Q is preferable and requires no logical copy here.

3. Kernels and Graphs Have Different Representations

Pointer arithmetic natural to GPU kernels cannot transfer directly to value-based graphs. Graphs can instead express the same relation through fan-out or tensor shapes.

Backend constraints often mean expressing the same math differently, not abandoning computation.

4. Design Decode and Prefill Separately

Both perform attention but have different query lengths. Decode q_len=1 uses groups directly as rows; prefill must consider group and sequence axes together.

Mechanically applying the same transformation can increase rank or create unnecessary broadcasts.

5. Measure Both the Graph and the Device

Removing tile from MIL does not prove a speedup. Reshape may cause backend rearrangement, and caching or fusion can differ from expectations.

Device throughput alone cannot explain a speedup either. Compare graphs, tensor sizes, outputs, and physical-device A/B measurements together to connect cause and effect.

Conclusion

Every GQA implementation solves the same problem.

The KV head for query head h is h // group_size.

The Hugging Face reference uses tensor repetition. Fused SDPA carries GQA information to the backend. Triton and vLLM use pointer arithmetic; CoreAI iOS fans one KV value out to multiple multiplications.

We express the relationship in tensor shapes. Decode splits query heads into [Hkv, G]; prefill merges G with sequence axis M into [Hkv, G×M].

No approach is universally correct. The best choice depends on custom kernels, graph-size constraints, query length, and backend reshape/broadcast behavior.

Identical mathematics does not imply identical execution cost. Preserve semantics in a form the backend can understand and execute efficiently throughout.

References

1. Ainslie, J., et al. (2023). *GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints.* EMNLP. arXiv:2305.13245
2. Dao, T., et al. (2022). *FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness.* NeurIPS. arXiv:2205.14135
3. PyTorch. FlexAttention Inductor lowering, torch/_inductor/kernel/flex_attention.py, torch/_inductor/kernel/flex_decoding.py. PyTorch 2.8.0 installed version.
4. PyTorch. scaled_dot_product_attention, torch/nn/functional.py, torch/_decomp/decompositions.py.
5. Hugging Face Transformers. SDPA attention integration, transformers/integrations/sdpa_attention.py. Transformers 4.57.3 version.
6. vLLM. vllm/v1/attention/ops/triton_decode_attention.py, vllm/v1/attention/ops/chunked_prefill_paged_decode.py.
7. Apple. CoreAI iOS-optimized SDPA, coreai_models/primitives/ios/sdpa.py.