[Engineering] Core ML 디코딩 최적화: GQA의 repeat_kv가 만든 토큰당 1.9 GiB 복사를 없앤 방법

안녕하세요, 옵트에이아이 박성재입니다.

이 작업은 CoreAI가 정식으로 출시되기 전에 시작됐습니다. 당시 저희 팀은 EXAONE 4.0 1.2B를 Core ML 모델로 직접 변환해 iPhone의 Neural Engine에서 구동하는 작업을 진행하고 있었습니다.

모델 변환에는 성공했지만 디코드 처리량이 목표에 크게 미치지 못했습니다. 그래프를 따라가며 원인을 분석한 결과, Grouped Query Attention(GQA)의 KV 헤드를 쿼리 헤드 수에 맞추기 위해 사용하는 repeat_kv가 30개 레이어에서 토큰마다 합계 960 MiB의 중간 출력을 만들고 있다는 사실을 발견했습니다.

이 출력을 쓰고 다음 연산에서 다시 읽는 양까지 합치면 그래프상 데이터 이동은 토큰당 약 1.88 GiB입니다. 모델의 계산 결과에는 아무것도 더하지 않고 메모리 트래픽만 늘리는 복사였습니다.

해결 방법은 KV 헤드 8개를 32개로 복제하는 대신, 쿼리 헤드 32개를 [8, 4]로 재배열하는 것이었습니다. GQA의 헤드 공유 관계를 텐서 형상에 직접 담은 것입니다.

같은 양자화 구성에서 이 변경만 적용한 결과는 다음과 같습니다.

처리량은 2.66배로 증가했습니다. 연산의 의미나 가중치 정밀도를 바꾼 것이 아니라, 같은 연산을 백엔드가 처리하기 좋은 형태로 다시 표현한 결과입니다.

최종 top8 mixed-precision 빌드를 실제 기기에서 측정했을 때는 iPhone 16 Pro에서 44 tok/s, iPhone 17 Pro에서 56 tok/s를 기록했습니다. 이 기기별 결과와 위의 원인 규명용 A/B는 측정 목적이 다르며, 자세한 조건은 4절에서 구분해 설명합니다.

1. 처음 의심한 곳은 KV cache 갱신이었다

GQA에서는 여러 쿼리 헤드가 하나의 KV 헤드를 공유합니다. EXAONE 4.0 1.2B는 쿼리 헤드 32개와 KV 헤드 8개를 사용하므로, KV 헤드 하나를 쿼리 헤드 네 개가 공유합니다.

모든 쿼리 헤드가 별도의 K와 V를 갖는 Multi-Head Attention(MHA)과 비교하면 KV cache의 크기와 이를 읽는 양을 4분의 1로 줄일 수 있습니다.

이 모델에서 레이어 하나의 K 또는 V cache 형상은 [8, 4096, 64]입니다. fp16을 사용하므로 텐서 하나의 크기는 4 MiB이고, 30개 레이어의 K와 V를 모두 합치면 240 MiB입니다.

새 토큰 하나를 추가할 때 실제로 갱신하는 것은 시퀀스 축의 한 칸뿐입니다. 전체 cache는 240 MiB지만, 토큰마다 새로 쓰는 데이터는 60 KiB입니다.

여기서 먼저 눈에 들어온 것은 데이터 크기가 아니라 연산 수였습니다. 60 KiB를 쓰는 작업이 slice_update 60개로 나뉘어 있었고, 각 연산이 상태 텐서에 별도로 접근했습니다.

캐시를 상태가 아니라 입력으로 받고 호스트에서 갱신하는 read-only KV 프로토타입에서는 디코드 함수의 MIL 연산 수가 7,110개에서 6,153개로 13.5% 줄었습니다. 하지만 KV 경로를 더 따라가 보니, 갱신보다 훨씬 큰 문제가 캐시를 읽는 쪽에 있었습니다.

repeat_kv가 만든 60개 tile 출력의 합계가 토큰당 960 MiB에 달했습니다.

마지막 값은 MIL 그래프의 텐서 형상으로 계산한 논리적 데이터 이동량입니다. 실제 DRAM 트래픽은 캐시와 연산 융합의 영향을 받을 수 있으므로, 뒤에서 실제 기기 처리량과도 대조하겠습니다.

2. GQA는 왜 KV를 복제하는가

어텐션 점수를 계산하는 첫 단계는 배치 행렬곱입니다.

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

뒤의 두 축은 행렬곱 조건을 만족합니다.

(q_len × 64) @ (64 × ctx)

문제는 앞의 배치 축입니다. (1, 32)와 (1, 8)에서 32와 8은 서로 같지 않고 어느 쪽도 1이 아니므로 브로드캐스트할 수 없습니다.

일반적인 참조 구현은 KV 헤드 8개를 각각 네 번 반복해 32개로 늘립니다.

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

이제 Q와 K의 헤드 축이 모두 32이므로 평범한 배치 행렬곱으로 계산할 수 있습니다.

expand는 view지만 reshape에서 복사가 생긴다

Hugging Face의 repeat_kv는 대략 다음 순서로 동작합니다.

1. KV 텐서에 반복 축을 하나 추가합니다.
2. expand로 그 축을 네 번 반복한 것처럼 보이게 합니다.
3. KV 헤드 축과 반복 축을 reshape로 합칩니다.

expand까지는 stride가 0인 view라서 데이터를 복사하지 않습니다. 하지만 stride 0인 반복 축을 옆 축과 합쳐 연속적인 헤드 축으로 만들 때는 stride만으로 같은 배치를 표현할 수 없습니다. eager 모드에서는 이 reshape가 새 storage를 할당하고 데이터를 채웁니다.

K        storage  4.00 MB   원본과 storage 공유
expand   storage  4.00 MB   원본과 storage 공유
reshape  storage 16.00 MB   새 storage 할당

GPU에서 이 비용이 항상 나타나는 것은 아닙니다. 융합 어텐션(fused attention) 커널이 GQA를 직접 지원하면 repeat_kv 경로를 거치지 않고, 각 쿼리 헤드가 사용할 KV 헤드를 커널 안에서 계산할 수 있습니다.

MIL 그래프에는 stride 기반 view가 없다

PyTorch 텐서는 storage, shape, stride, offset의 조합으로 볼 수 있습니다. 같은 storage를 유지한 채 stride만 바꿔 view를 표현할 수 있습니다.

반면 Core ML의 중간 표현인 MIL은 값 중심의 데이터플로우 그래프입니다. 그래프 작성자가 PyTorch와 같은 stride와 offset을 지정해 “같은 버퍼를 다른 방식으로 읽는 view”를 직접 표현할 수 없습니다.

기존 expand → reshape의 의미를 MIL로 변환하는 과정에서 반복은 tile 연산으로 남았습니다. 최종 그래프에서 각 K와 V는 8개 헤드에서 32개 헤드로 늘어났고, 뒤의 행렬곱은 이 확장된 값을 입력으로 받았습니다.

원본 K 읽기 4 MiB
    → 반복된 K 쓰기 16 MiB
    → 다음 행렬곱이 반복된 K 읽기 16 MiB

레이어마다 K와 V에 하나씩, 30개 레이어에서 총 60개의 tile이 생깁니다.

60 × 16 MiB = 960 MiB
쓰기와 다음 연산의 읽기 = 1,920 MiB ≈ 1.88 GiB

repeat_kv는 어텐션의 곱셈 횟수를 늘리지 않습니다. 쿼리 헤드 32개가 수행하는 내적의 수는 복제 전후가 같습니다. 추가되는 것은 KV 복제본을 만들고 읽는 데이터 이동입니다.

일반적인 최적화 패스로는 왜 사라지지 않았을까

tile 출력은 바로 다음 행렬곱이 실제로 사용하므로 죽은 코드가 아닙니다. 입력과 출력의 형상이 달라서 단순한 no-op 제거 대상도 아닙니다.

이를 자동으로 없애려면 컴파일러가 다음 관계를 한꺼번에 알아야 합니다.

• 이 tile은 KV 헤드 수를 쿼리 헤드 수에 맞추기 위한 연산이다.
• 소비자인 행렬곱은 어텐션 점수를 계산한다.
• 쿼리 헤드 h는 h // 4번 KV 헤드를 사용한다.
• 따라서 KV를 복제하지 않는 GQA 형태로 전체 패턴을 바꿀 수 있다.

이는 단일 연산을 보는 지역 최적화가 아니라 어텐션 구조를 인식하는 패턴 최적화입니다. 확인한 coremltools 9.0의 공개 최적화 패스에는 이 변환이 없었고, 실제 변환 결과에도 tile 60개가 남았습니다. 비공개 컴파일러 단계의 동작까지 단정할 수는 없지만, 적어도 우리가 얻은 최종 MIL 그래프에서는 복제가 제거되지 않았습니다.

3. KV를 늘리는 대신 Q를 접는다

표준 배치 행렬곱 형태를 유지하면서 32개의 쿼리 헤드와 8개의 KV 헤드를 맞추는 방법은 크게 두 가지입니다.

디코드에서는 q_len=1이므로 Q의 형상은 [1, 32, 1, 64]입니다. 크기는 fp16 기준 4 KiB에 불과합니다. 더 중요한 점은 원소를 반복하지 않고 헤드 축 32를 (8, 4)로 나누기만 한다는 것입니다.

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

여기서 8은 KV 헤드 수이고, 4는 KV 헤드 하나를 공유하는 쿼리 헤드 수입니다.

reshape가 정수 나눗셈을 대신하는 이유

쿼리 헤드 번호를 h, KV 헤드 번호를 kv, 그룹 안의 위치를 g라고 하겠습니다.

헤드 32개를 row-major 순서로 [8, 4]에 배치하면 다음 관계가 생깁니다.

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

kv = h // 4
g  = h % 4

따라서 0~3번 쿼리 헤드는 KV 헤드 0을, 4~7번 쿼리 헤드는 KV 헤드 1을 사용합니다.

커널에서는 이 관계를 정수 나눗셈으로 계산할 수 있습니다. MIL 그래프에서는 헤드 축을 [8, 4]로 나누는 reshape가 같은 매핑을 미리 표현합니다.

계산 결과가 같은 이유

첫째, 헤드 그룹의 순서가 같습니다. repeat_kv는 KV 헤드와 반복 인덱스를 (kv, g) 순서로 배치한 뒤 하나의 헤드 축으로 합칩니다. 반대로 Q의 헤드 축을 [8, 4]로 나누면 동일한 (kv, g) 관계가 복원됩니다.

둘째, 행렬곱의 각 행은 서로 독립적입니다.

(4 × 64) @ (64 × ctx)

출력의 각 행은 해당 Q 행과 K의 내적으로만 결정됩니다. 같은 K를 사용하는 한 행짜리 행렬곱 4개를 4행짜리 행렬곱 하나로 묶어도 각 행의 계산은 바뀌지 않습니다.

셋째, 뒤따르는 연산도 모두 행 단위로 독립적입니다. scale과 mask 덧셈은 원소별 연산이고, softmax는 context 축인 마지막 축에서 수행됩니다. prob @ V 역시 각 행을 독립적으로 처리합니다.

반대로 softmax가 헤드 축을 가로지르거나 서로 다른 헤드의 값을 섞는 연산이 중간에 있다면 이 재배열은 성립하지 않습니다.

구현

디코드 경로의 핵심 코드는 다음과 같습니다.

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와 V는 저장된 KV 헤드 수를 유지한다.
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]가 KV 헤드 축과 그룹 축으로 브로드캐스트된다.
    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와 V는 8개 헤드인 원래 폭으로 읽힙니다. Q와 출력은 원소를 추가하거나 순서를 바꾸지 않고 형상만 바꿉니다.

tile을 최적화 패스로 지운 것이 아니라, 그래프에서 tile이 필요하지 않은 행렬곱을 만든 셈입니다.

reshape는 정말 무료인가

논리적인 관점에서 reshape는 원소 수와 순서를 바꾸지 않으므로 새 버퍼가 필요하지 않습니다. 하지만 ANE는 내부 레이아웃을 사용하므로 모든 reshape가 항상 무료라고 단정할 수는 없습니다. 특히 채널 축을 가로지르는 reshape는 백엔드에서 재배치를 만들 가능성이 있습니다.

이번 변경이 실제로 유효했다는 근거는 세 가지입니다.

• 변환 후 tile 연산이 60개에서 0개로 줄었습니다.
• 전체 MIL 연산 수가 감소했습니다.
• 실제 기기에서 디코드 처리량이 2.66배 증가했습니다.

4. 그래프와 실제 기기에서 확인한 결과

같은 모델에서 구현 플래그만 바꿔 두 번 변환한 뒤 MIL 그래프를 비교했습니다.

tile 60개는 30개 레이어의 K와 V에서 하나씩 나온 값입니다. 형상으로 계산한 tile 출력의 합계도 60 × 16 MiB = 960 MiB로 일치했습니다.

MIL 텐서 크기의 정적 합은 런타임 peak memory나 실제 DRAM 트래픽과 같은 지표가 아닙니다. 따라서 여기서 확실히 말할 수 있는 것은 최종 그래프에서 60개의 tile과 그에 따른 큰 중간 텐서가 사라졌다는 사실입니다.

동일 조건 A/B

동일한 실제 기기와 동일한 양자화 구성에서 디코드 구현 플래그만 바꿔 측정했습니다. 두 빌드 모두 FFN 4-bit, 8개 레이어 8-bit 승격, lm_head 6-bit, embedding 4-bit, context length 4096를 사용했습니다. 결과는 프롬프트 30개의 평균입니다.

가중치와 양자화 구성은 같고, 달라진 것은 디코드의 GQA 표현 방식뿐입니다. 처리량은 2.66배, 증가율로는 약 166% 높아졌습니다.

iPhone 16 Pro·17 Pro 실제 기기 결과

위 표는 repeat_kv 변경의 효과만 분리하기 위한 A/B입니다. 이와 별도로, 실제 사용 구성을 비교하기 위해 iPhone 16 Pro와 iPhone 17 Pro에서 양자화 구성별 디코드 처리량도 측정했습니다.

여기서 top8과 top12는 일부 레이어만 더 높은 정밀도로 유지한 혼합 정밀도(mixed precision) 양자화 구성입니다. top8은 양자화 민감도가 높은 레이어 8개를 8-bit로 승격하고 나머지 레이어는 4-bit로 유지합니다. top12는 같은 기준으로 선택한 레이어 12개를 8-bit로 승격한 구성입니다. 즉 숫자가 클수록 8-bit로 유지하는 레이어가 많습니다. top12가 정밀도를 더 많이 보존하는 대신 읽어야 할 가중치도 늘어나므로, 실제 기기 처리량은 top8이 조금 더 높았습니다.

두 기기 모두 top8 mixed-precision 구성이 가장 높은 처리량을 보였습니다. iPhone 16 Pro에서는 44 tok/s, iPhone 17 Pro에서는 56 tok/s였습니다.

다만 이 표는 repeat_kv 제거만 비교한 실험이 아닙니다. base, top8, top12는 서로 다른 양자화 구성이고, base에서 top8로 넘어가는 과정에는 처리 chunk를 2에서 1로 바꾼 변경도 포함됩니다. 따라서 15→44 tok/s와 21→56 tok/s의 차이 전체를 repeat_kv 제거 효과로 해석해서는 안 됩니다.

repeat_kv 변경에 직접 귀속할 수 있는 수치는 앞의 동일 조건 A/B인 20.39→54.30 tok/s입니다. 이 기기별 표는 최종 후보 구성이 두 세대의 실제 iPhone에서 어느 정도 처리량을 내는지 보여 주는 end-to-end 결과입니다. 두 표는 측정 목적과 실험 묶음이 다르므로 54.30과 56의 세부 차이를 직접 비교해서도 안 됩니다.

그래프 수준 추정과 처리량의 정합성

별도의 하드웨어 카운터로 DRAM 트래픽을 직접 측정한 것은 아닙니다. 대신 여러 빌드에서 관찰한 처리량과 토큰당 읽는 데이터의 관계를 단순 대역폭 모델로 근사했습니다.

tok/s ≈ 55,900 MB/s ÷ 토큰당 유효 읽기량(MB)

이 상수로 두 빌드의 유효 읽기량을 역산하면 다음과 같습니다.

Q 재배열       : 55,900 / 54.30 = 1,030 MB/token
repeat_kv 유지 : 55,900 / 20.39 = 2,742 MB/token
차이           :                  1,712 MB/token

그래프 형상에서 계산한 tile 쓰기·읽기는 1,920 MiB, 즉 약 2,013 MB입니다. 처리량에서 역산한 차이는 1,712 MB로 약 15% 작았습니다.

이 계산은 독립적인 DRAM 트래픽 측정이 아니라 크기와 방향을 확인하는 교차 점검입니다. 캐시가 일부 재읽기를 흡수했을 가능성을 고려하면 두 값은 같은 규모의 추가 데이터 이동을 가리킵니다.

5. 한계와 다음 단계

그래프의 1.88 GiB가 실제 DRAM 트래픽을 뜻하지는 않는다

1.88 GiB는 tile 출력의 쓰기와 다음 연산의 읽기를 텐서 형상으로 계산한 값입니다. 실제 하드웨어에서는 캐시, 내부 레이아웃, 연산 융합의 영향을 받습니다. 따라서 이 숫자는 그래프 수준의 트래픽 추정치로 읽어야 합니다.

prefill은 별도의 설계가 필요하다

q_len=1인 디코드에서는 그룹 축을 그대로 행 축으로 사용할 수 있습니다. prefill처럼 q_len이 커지면 rank-5 브로드캐스트와 mask 정렬까지 고려해야 합니다. 이 경우에는 그룹 축을 시퀀스 축에 합쳐 rank 4를 유지했습니다. 자세한 원리와 백엔드 비교는 2편에서 다룹니다.

쿼리 헤드가 KV 그룹 순서로 배치되어 있어야 한다

이번 reshape는 쿼리 헤드가 KV 그룹별로 연속 배치돼 있다는 전제에서 성립합니다. 즉 0~3번 쿼리 헤드가 KV 헤드 0을, 4~7번이 KV 헤드 1을 공유해야 합니다.

헤드가 교차 배치된 모델이라면 단순 reshape만으로 같은 매핑을 만들 수 없습니다. transpose나 gather 같은 추가 연산이 필요할 수 있으며, 이 경우 최적화의 비용 구조도 달라집니다.

KV cache 갱신 연산은 아직 남아 있다

이번 변경은 KV cache를 읽어 어텐션에 넣는 경로를 최적화했습니다. 토큰마다 K와 V를 갱신하는 상태 쓰기 연산 60개는 그대로입니다.

read-only KV 프로토타입에서 그래프 연산 수와 출력 일치는 확인했지만, 실제 기기 성능 검증은 남아 있습니다. 다음 최적화 대상은 이 갱신 경로입니다.

결론

문제는 GQA 자체가 아니라 GQA의 헤드 공유를 그래프에 표현한 방식이었습니다.

GPU의 fused attention 커널은 쿼리 헤드 번호로부터 사용할 KV 헤드를 직접 계산할 수 있습니다. 반면 우리가 사용한 MIL 그래프에는 그런 커널이 없었고, repeat_kv가 tile 연산으로 남아 매 토큰마다 큰 KV 복제본을 만들었습니다.

KV 헤드를 8개에서 32개로 늘리는 대신 쿼리 헤드 32개를 [8, 4]로 나누자 같은 공유 관계를 복사 없이 표현할 수 있었습니다.

그 결과 디코드 경로의 tile 연산 60개가 모두 사라졌고, 같은 양자화 구성에서 처리량이 20.39 tok/s에서 54.30 tok/s로 증가했습니다. 테스트한 입력에서 출력도 정확히 일치했습니다.

가중치를 더 압축하면 데이터 이동을 줄일 수 있지만 품질 저하가 생길 수 있습니다. 이번 최적화는 같은 데이터를 여러 번 표현하던 중복을 제거했습니다. 품질과 연산의 의미를 유지하면서 백엔드가 실제로 처리해야 할 데이터만 줄였다는 점이 가장 중요합니다.

2편에서는 같은 GQA 헤드 매핑을 PyTorch, Triton, vLLM, CoreAI가 각각 어떻게 표현하는지 살펴보고, 디코드와 prefill에서 서로 다른 텐서 형상이 필요한 이유를 정리합니다.

용어 정리

참고 문헌

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.* 모델 카드 및 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 기준.