[Engineering] Core ML Decode Optimization: Removing the 1.9 GiB Per-Token Copy Caused by GQA repeat_kv

Hello, I'm Seongjae Park from OptAI.

This work began before the official release of CoreAI. Our team was directly converting EXAONE 4.0 1.2B into Core ML to run on the iPhone Neural Engine.

Conversion succeeded, but decode throughput fell far below the target. Tracing the graph revealed that repeat_kv, which matches Grouped Query Attention (GQA) KV heads to the query-head count, generated 960 MiB of intermediate output per token across 30 layers.

Including the write and the following operation read, logical graph data movement was approximately 1.88 GiB per token. These copies added nothing to the computation and only increased memory traffic.

The solution was to rearrange the 32 query heads into [8, 4] rather than replicating 8 KV heads into 32, representing GQA head sharing directly in tensor shapes.

Applying only this change under the same quantization configuration produced the following results.

Throughput increased by 2.66 times. Neither operation semantics nor weight precision changed: the same computation was re-expressed in a form better suited to the backend.

The final top8 mixed-precision build reached 44 tok/s on iPhone 16 Pro and 56 tok/s on iPhone 17 Pro. These device results serve a different purpose from the causal A/B experiment above; section 4 distinguishes their conditions.

1. KV Cache Updates Were the Initial Suspect

In GQA, multiple query heads share a KV head. EXAONE 4.0 1.2B has 32 query heads and 8 KV heads, so four query heads share each KV head.

Compared with Multi-Head Attention (MHA), where each query head has separate K and V, this reduces KV cache size and read volume to 1/4.

The K or V cache for one layer has shape [8, 4096, 64]. In fp16, each tensor is 4 MiB; K and V across all 30 layers total 240 MiB.

Adding a token updates only one position along the sequence axis. The full cache is 240 MiB, but only 60 KiB is newly written per token.

The first concern was operation count rather than data size. Writing 60 KiB was divided into 60 slice_update operations, each separately accessing state tensors.

A read-only KV prototype received the cache as input and updated it on the host, rather than as state. Decode MIL operations fell from 7,110 to 6,153, a 13.5% reduction. Following the KV path further revealed a much larger issue in cache reads than updates.

The 60 tile outputs introduced by repeat_kv totaled 960 MiB per token.

The final value is logical data movement calculated from MIL tensor shapes. Actual DRAM traffic depends on caching and fusion, so we also compare it with measured device throughput below.

2. Why Does GQA Replicate KV?

Attention score computation begins with batched matrix multiplication.

Q  : [1, 32, q_len, 64]
Kᵀ : [1,  8,    64, ctx]

The last two axes satisfy matrix-multiplication requirements.

(q_len × 64) @ (64 × ctx)

The leading batch axes are the issue. In (1, 32) and (1, 8), 32 and 8 differ and neither is 1, so they cannot broadcast.

A typical reference implementation repeats each of the 8 KV heads four times to produce 32 heads.

K: [1, 8, ctx, 64]
          ↓ repeat_kv
K: [1, 32, ctx, 64]

Q and K now both have 32 heads and can use ordinary batched matrix multiplication.

expand Is a View, but reshape Introduces a Copy

Hugging Face repeat_kv roughly follows these steps.

1. Add a repetition axis to the KV tensor.
2. Use expand to make it appear repeated four times.
3. Merge the KV-head and repetition axes with reshape.

Up to expand, the tensor is a stride-0 view without copying. Merging that stride 0 repetition axis with a neighboring axis into a contiguous head dimension cannot be expressed using strides alone. Eager execution therefore allocates and populates new storage for this reshape.

K        storage  4.00 MB   Shares storage with the original
expand   storage  4.00 MB   Shares storage with the original
reshape  storage 16.00 MB   Allocates new storage

This cost is not inevitable on GPUs. A fused attention kernel with native GQA support can bypass repeat_kv and compute the KV head for each query head inside the kernel.

MIL Graphs Have No Stride-Based Views

PyTorch tensors combine storage, shape, stride, and offset. Changing strides can express a view while retaining the same storage.

MIL, the Core ML intermediate representation, is a value-based dataflow graph. Graph authors cannot directly specify PyTorch-style strides and offsets to express reading the same buffer differently.

Converting expand → reshape semantics into MIL left repetition as tile. Each K and V expanded from 8 heads to 32, with subsequent multiplications consuming the expanded values.

Read original K 4 MiB
    → Write repeated K 16 MiB
    → Next matrix multiplication reads repeated K 16 MiB

One tile each for K and V across 30 layers produces 60 in total.

60 × 16 MiB = 960 MiB
Writes and reads by the next operation = 1,920 MiB ≈ 1.88 GiB

repeat_kv does not increase attention multiplication counts. The 32 query heads perform the same number of dot products before and after replication. What increases is movement to create and read KV replicas.

Why Did Ordinary Optimization Passes Not Remove It?

The tile output is consumed by the next multiplication, so it is not dead code. Different input/output shapes also prevent simple no-op elimination.

Automatic removal requires recognizing these relationships together.

• This tile matches KV-head count to query-head count.
• The consuming multiplication calculates attention scores.
• Query head h uses KV head h // 4.
• The full pattern can become GQA without KV replication.

This is attention-aware pattern optimization, not local optimization of one operation. Public coremltools 9.0 passes we inspected lacked this transformation, and the converted graph retained 60 tile operations. We cannot conclude how private compiler stages behave, but replication remained in the final MIL graph obtained.

3. Fold Q Instead of Expanding KV

There are two main ways to match 32 query heads to 8 KV heads while keeping standard batched multiplication.

During decode, q_len=1, so Q is [1, 32, 1, 64] and only 4 KiB in fp16. Crucially, splitting head axis 32 into (8, 4) repeats no elements.

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

Here, 8 is the KV-head count, and 4 is the number of query heads sharing one KV head.

Why Reshape Replaces Integer Division

Let h be the query-head index, kv the KV-head index, and g the position in the group.

Arranging 32 heads in row-major order into [8, 4] gives the following relationship.

h = 4 × kv + g      (0 ≤ g < 4)

kv = h // 4
g  = h % 4

Query heads 0~3 therefore use KV head 0, and heads 4~7 use KV head 1.

A kernel can calculate this by integer division. In MIL, reshaping into [8, 4] expresses the same mapping in advance.

Why the Computation Is Equivalent

First, group order is identical. repeat_kv arranges KV heads and repetition indices in (kv, g) order and merges them. Splitting Q into [8, 4] recovers that same (kv, g) relationship.

Second, each matrix-multiplication row is independent.

(4 × 64) @ (64 × ctx)

Each output row depends only on its Q row dot product with K. Combining 4 single-row multiplications sharing K into one 4-row multiplication changes no row computation.

Third, subsequent operations are also row-independent. Scaling and mask addition are element-wise; softmax operates on the last, context axis. prob @ V processes rows independently too.

This rearrangement would not hold if softmax crossed the head axis or an intervening operation mixed heads.

Implementation

The core decode-path code follows.

n_kv = self.num_key_value_heads
n_rep = self.num_key_value_groups

# [B, Hq, 1, D] -> [B, Hkv, G, D]
q = query_states.reshape(bsz, n_kv, n_rep, self.head_dim)

# K and V retain the stored KV head count.
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))
scores = self._score_scale_mul(scores)

if sliced_mask is not None:
    # [1, 1, 1, ctx] broadcasts over the KV-head and group axes.
    scores = scores + sliced_mask

prob = self.ANE_softmax(scores, dim=-1)
out = torch.matmul(prob, v)

# [B, Hkv, G, D] -> [B, Hq, 1, D]
out = out.reshape(bsz, self.num_heads, q_len, self.head_dim)

K and V are read at their original width of 8 heads. Q and output change shape without adding elements or changing their order.

Rather than eliminating tile with a pass, we constructed graph multiplications requiring no tile.

Is Reshape Really Free?

Logically, reshape changes neither element count nor order, so needs no new buffer. ANE has internal layouts, however, so not every reshape is necessarily free. Reshaping across a channel axis may cause backend rearrangement.

Three observations support effectiveness.

• Converted tile operations fell from 60 to 0.
• Total MIL operation count decreased.
• Physical-device decode throughput increased by 2.66 times.

4. Results in the Graph and on Devices

We converted the same model twice, changing only the implementation flag, and compared MIL graphs.

The 60 tile operations come from K and V across 30 layers. The sum of tile output sizes also matched 60 × 16 MiB = 960 MiB.

A static sum of MIL tensor sizes is not runtime peak memory or actual DRAM traffic. What is established is that 60 tile operations and their large intermediate tensors disappeared.

Identical-Condition A/B

On the same physical device and quantization configuration, we changed only the decode flag. Both builds used FFN 4-bit, 8 layers promoted to 8-bit, lm_head 6-bit, embedding 4-bit, context length 4096. Results average 30 prompts.

Weights and quantization stayed identical; only GQA representation changed. Throughput rose 2.66 times, approximately 166%.

Device Results: iPhone 16 Pro and 17 Pro

The preceding A/B isolates repeat_kv. Separately, to compare practical configurations, we measured throughput across quantization configurations on iPhone 16 Pro and iPhone 17 Pro.

top8 and top12 are mixed-precision configurations retaining higher precision for selected layers. top8 promotes the 8 most sensitive layers to 8-bit and keeps the rest at 4-bit. top12 promotes 12 layers by the same criterion to 8-bit. A larger number retains more 8-bit layers. top12 preserves more precision but reads more weights, giving top8 slightly higher device throughput.

top8 mixed precision led on both devices: 44 tok/s on iPhone 16 Pro and 56 tok/s on iPhone 17 Pro.

This does not compare only repeat_kv removal. base, top8, and top12 have different quantization configurations; moving from base to top8 also changed processing chunks from 2 to 1. The full 15→44 tok/s and 21→56 tok/s differences must not be attributed to repeat_kv removal.

Only the identical-condition A/B of 20.39→54.30 tok/s is directly attributable to repeat_kv. The device table shows final-candidate end-to-end throughput across two iPhone generations. The tables have different purposes and experimental groupings; even 54.30 versus 56 must not be directly compared.

Graph Estimates and Throughput Consistency

We did not directly measure DRAM traffic with hardware counters. Instead, throughput and per-token read volume across builds were approximated with a simple bandwidth model.

tok/s ≈ 55,900 MB/s ÷ effective reads per token (MB)

Using this constant to infer effective read volumes gives the following.

Q rearrangement       : 55,900 / 54.30 = 1,030 MB/token
retain repeat_kv : 55,900 / 20.39 = 2,742 MB/token
Difference           :                  1,712 MB/token

Graph-shape tile writes and reads total 1,920 MiB, approximately 2,013 MB. The throughput-derived difference, 1,712 MB, is approximately 15% smaller.

This cross-check establishes scale and direction, not an independent DRAM measurement. If caching absorbs some rereads, both values indicate additional movement of the same order.

5. Limitations and Next Steps

Graph-Level 1.88 GiB Is Not Actual DRAM Traffic

The 1.88 GiB estimate calculates tile writes and subsequent reads from tensor shapes. Hardware caching, layouts, and fusion affect actual traffic; treat this as a graph-level estimate.

Prefill Needs Separate Design

With q_len=1, the group axis can directly represent rows. Larger q_len in prefill requires rank-5 broadcasting and mask alignment. We merged group and sequence axes to retain rank 4. Part 2 explains the principles and backend comparison.

Query Heads Must Be Ordered by KV Group

This reshape assumes contiguous query heads per KV group: heads 0~3 share KV head 0, and 4~7 share KV head 1.

Interleaved heads cannot express this mapping with reshape alone. Additional transpose or gather operations may be needed, changing optimization costs.

KV Cache Update Operations Remain

The optimized path reads KV into attention. The 60 per-token state writes updating K and V remain.

The read-only KV prototype verified graph counts and output agreement, but physical-device performance remains unvalidated. The update path is next.

Conclusion

GQA itself was not the problem; the graph representation of its head sharing was.

Fused GPU attention can derive the KV head from the query-head index. Our MIL graph lacked that kernel: repeat_kv remained as tile, creating large replicas each token.

Splitting 32 query heads into [8, 4] instead of expanding KV from 8 to 32 expresses the same sharing without copying.

All 60 decode-path tile operations disappeared. Throughput rose from 20.39 tok/s to 54.30 tok/s under unchanged quantization, with exactly matching tested outputs.

Further weight compression can reduce movement at a quality cost. This optimization instead removes redundant representations of identical data, reducing backend work while preserving quality and semantics.

Part 2 compares the same GQA mapping across PyTorch, Triton, vLLM, and CoreAI, and explains why decode and prefill need different shapes.

Glossary

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. LG AI Research. *EXAONE 4.0.* Model card and config.json.
4. PyTorch. torch/nn/functional.py, torch/_decomp/decompositions.py.
5. Hugging Face Transformers. transformers/integrations/sdpa_attention.py.
6. Apple. coremltools MIL optimization passes, coremltools 9.0 version.