Same GQA, Different Execution Cost Across Backends

시리즈 안내 — Core ML에서 GQA 어텐션 경로를 최적화한 과정을 두 편으로 나눠 정리합니다. (1) 이전 글은 repeat_kv가 만든 KV 복제를 제거한 과정을, (2) 이번 글은 같은 GQA 매핑을 PyTorch, Triton, vLLM, CoreAI가 표현하는 방식을 다룹니다.

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

1편에서는 Core ML 디코드 그래프에 남은 repeat_kv가 토큰마다 대규모 KV 복제를 만들고 있던 문제를 다뤘습니다. KV 헤드 8개를 32개로 늘리는 대신 쿼리 헤드 32개를 [8, 4]로 재배열하자 tile 연산 60개가 사라졌고, 동일한 양자화 구성에서 처리량이 20.39 tok/s에서 54.30 tok/s로 증가했습니다. 양자화 민감도가 높은 레이어 8개를 8-bit로, 나머지를 4-bit로 유지한 top8 mixed-precision 빌드는 실제 기기에서 iPhone 16 Pro 44 tok/s, iPhone 17 Pro 56 tok/s를 기록했습니다.

이번 글에서는 한 단계 더 들어가 보겠습니다.

GQA를 지원하는 모든 구현은 결국 같은 관계를 표현합니다.

kv_head = query_head // group_size

하지만 이를 표현하는 수단은 백엔드마다 다릅니다. GPU 커널은 포인터 산술을 사용할 수 있고, 그래프는 하나의 값을 여러 연산으로 나눠 보낼 수 있으며, 텐서 형상 자체에 매핑을 담을 수도 있습니다.

같은 수학을 구현해도 어떤 표현을 선택하느냐에 따라 실제 데이터 이동량과 그래프 크기는 크게 달라집니다.

1. 모든 구현이 공유하는 하나의 매핑

Exaone4 1.2B 모델에는 쿼리 헤드가 32개, KV 헤드가 8개 있습니다. 따라서 KV 헤드 하나를 쿼리 헤드 네 개가 공유합니다.

group_size = 32 // 8 = 4

쿼리 헤드 번호를 h라고 하면 사용할 KV 헤드는 다음과 같습니다.

kv = h // 4

쿼리 헤드	KV 헤드
0, 1, 2, 3	0
4, 5, 6, 7	1
8, 9, 10, 11	2
12, 13, 14, 15	3
16, 17, 18, 19	4
20, 21, 22, 23	5
24, 25, 26, 27	6
28, 29, 30, 31	7

구현 방식은 달라도 이 관계는 바뀌지 않습니다.

구현	매핑을 표현하는 방법	KV 복제
Hugging Face `repeat_kv` eager 경로	텐서 반복	발생
PyTorch fused SDPA	백엔드에 GQA 정보 전달	지원 커널에서는 없음
Triton·vLLM 커널	포인터 산술	없음
CoreAI iOS SDPA	그래프 fan-out	없음
이번 Core ML 구현	쿼리 텐서 형상	없음

핵심은 “어떤 백엔드가 GQA를 지원하는가”만이 아닙니다. GQA의 공유 관계가 어느 단계까지 보존되고, 최종 실행기가 그것을 어떤 수단으로 표현할 수 있는지가 중요합니다.

2. Hugging Face `repeat_kv`: 텐서를 직접 반복한다

일반적인 GQA 참조 구현은 KV 헤드를 쿼리 헤드 수에 맞게 반복합니다.

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
    )

처음의 expand는 새 데이터를 만들지 않습니다. 반복 축의 stride를 0으로 설정해 같은 storage를 여러 번 읽는 view를 만듭니다.

문제는 그다음 reshape입니다. stride 0인 반복 축을 KV 헤드 축과 합쳐 하나의 연속적인 헤드 축으로 표현할 수 없기 때문에 eager 모드에서는 새 storage가 생깁니다.

K        storage  4.00 MB   ptr 공유: True
expand   storage  4.00 MB   ptr 공유: True
reshape  storage 16.00 MB   ptr 공유: False

이 방식의 장점은 단순함입니다. KV를 32개 헤드로 늘린 뒤에는 일반적인 Multi-Head Attention과 같은 형상으로 계산할 수 있습니다.

단점은 GQA의 중요한 의미가 사라진다는 것입니다. 한 번 [1, 32, ctx, 64]로 확장된 텐서만 보면, 원래 KV 헤드가 8개였고 각 헤드를 네 개의 쿼리 헤드가 공유했다는 사실을 알기 어렵습니다. 이후 백엔드가 이 관계를 다시 복원해 복제를 제거하기도 어려워집니다.

3. PyTorch SDPA: GQA 정보를 백엔드까지 보낸다

PyTorch의 scaled_dot_product_attention은 enable_gqa=True 옵션을 제공합니다. 중요한 점은 이 플래그 자체가 복사 없는 실행을 보장하지는 않는다는 것입니다.

문서에 제시된 참조 의미는 K와 V를 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
    )

실제 실행에서는 dispatcher가 입력의 dtype, shape, mask, 기기와 지원 기능을 보고 사용할 백엔드를 고릅니다. 선택된 fused backend가 GQA를 직접 지원하면 K와 V를 물리적으로 반복하지 않을 수 있습니다. 반면 MATH fallback에서는 참조 의미에 해당하는 반복 비용이 나타날 수 있습니다.

따라서 다음 두 코드는 겉보기에는 비슷하지만 백엔드에 전달하는 정보가 다릅니다.

# GQA 관계가 사라진 뒤 SDPA 호출
k = repeat_kv(k, n_rep)
v = repeat_kv(v, n_rep)
out = scaled_dot_product_attention(q, k, v)

# Q와 KV의 원래 헤드 수를 유지한 채 GQA 정보 전달
out = scaled_dot_product_attention(
    q, k, v, enable_gqa=True
)

첫 번째 코드에서는 백엔드가 입력만 보고 원래의 KV 헤드 공유 관계를 알아내기 어렵습니다. 두 번째 코드에서는 쿼리 헤드와 KV 헤드 수의 차이가 그대로 남아 있고, enable_gqa가 그 관계의 의미를 알려 줍니다.

이 차이는 컴파일러 최적화에서 자주 등장합니다. 의미 있는 구조를 낮은 수준의 연산으로 너무 일찍 풀어 버리면, 뒤 단계에서 더 효율적인 구현을 선택할 정보가 사라집니다.

4. Triton과 vLLM: KV 헤드를 주소로 선택한다

커스텀 GPU 커널은 각 프로그램 인스턴스가 어떤 쿼리 헤드를 처리하는지 알고 있습니다. 따라서 KV 텐서를 복제하지 않고 사용할 KV 헤드의 주소를 직접 계산할 수 있습니다.

PyTorch Inductor의 FlexAttention Triton 템플릿은 다음과 같은 형태를 사용합니다.

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가 쿼리 헤드 번호이고 off_hkv가 사용할 KV 헤드 번호입니다.

예를 들어 쿼리 헤드 4, 5, 6, 7을 맡은 네 프로그램은 모두 다음 값을 얻습니다.

off_hkv = off_hq // 4 = 1

네 프로그램이 같은 k_offset을 계산해 동일한 K 주소를 읽습니다. KV 데이터를 네 벌로 복사한 것이 아니라, 네 프로그램이 같은 원본을 참조한 것입니다.

vLLM 디코드: 쿼리 헤드에서 KV 헤드를 찾는다

vLLM의 Triton 디코드 커널도 같은 방향의 계산을 사용합니다.

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]
)

프로그램 하나가 쿼리 헤드 하나를 맡고, 정수 나눗셈으로 읽을 KV 헤드를 선택합니다.

vLLM prefill: KV 헤드에서 쿼리 헤드를 펼친다

prefill 커널에서는 반대 방향으로 같은 관계를 사용할 수 있습니다.

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)
)

이번에는 프로그램 하나가 KV 헤드 하나를 맡습니다. 그 안에서 해당 KV 헤드를 공유하는 쿼리 헤드 네 개를 계산합니다.

kv = h // 4        # 쿼리 헤드에서 KV 헤드 찾기
h  = kv × 4 + g    # KV 헤드에서 쿼리 헤드 펼치기

두 식은 같은 매핑의 정방향과 역방향입니다. 어느 방향을 택할지는 쿼리 길이, 병렬화 단위, K와 V의 재사용 방식에 따라 달라집니다.

5. CoreAI iOS SDPA: 하나의 값을 여러 연산이 공유한다

포인터 산술을 직접 쓸 수 없는 그래프에서도 KV를 복제하지 않는 방법이 있습니다. KV 헤드 하나를 여러 행렬곱의 입력으로 나눠 보내는 fan-out입니다.

CoreAI의 iOS 경로는 개념적으로 다음과 같은 형태를 사용합니다.

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

0~3번 쿼리 헤드의 행렬곱은 모두 keys[0]을 입력으로 받습니다. 하나의 K 값이 네 개의 행렬곱으로 연결되지만, K 자체를 네 번 복제하지는 않습니다.

값 중심 그래프에서 공유를 표현하는 자연스러운 방법입니다. 다만 헤드를 하나씩 펼치면 그래프가 커집니다. 쿼리 헤드가 32개인 경우 레이어마다 어텐션 점수 행렬곱 32개와 출력 행렬곱 32개가 필요합니다. 헤드를 분리하고 다시 모으기 위한 transpose, permute, concat도 추가됩니다.

즉 fan-out은 데이터 복사를 피하지만 그래프 크기와 연산자 수를 늘리는 선택입니다.

6. 커스텀 커널 없이 텐서 형상으로 표현하기

이번 Core ML 구현에서는 헤드를 하나씩 펼치지 않고 표준 배치 행렬곱을 유지하고 싶었습니다. 사용할 수 있는 핵심 수단은 텐서 형상이었습니다.

디코드의 Q 형상은 다음과 같습니다.

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

헤드 축 32를 KV 헤드 수 8과 그룹 크기 4로 나눕니다.

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

row-major에서 이 reshape는 다음 인덱스 관계를 만듭니다.

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

런타임에 h // G를 계산하지 않아도 [Hkv, G]라는 형상이 같은 관계를 미리 담고 있습니다.

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)

한 번의 배치 행렬곱이 KV 헤드 8개를 처리하고, 각 KV 헤드의 배치 안에는 이를 공유하는 쿼리 헤드 네 개가 행으로 들어갑니다.

PyTorch Flex Decoding도 Q를 같은 방향으로 재배열한다

PyTorch Inductor의 짧은 쿼리 전용 Flex Decoding 경로도 커널 호출 전에 Q를 다음과 같이 재배열합니다.

# [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에서는 as_strided로 view를 만든 뒤 전용 Triton 커널에 넘깁니다. 이번 구현에서는 재배열한 Q를 표준 배치 행렬곱에 직접 넣습니다.

재배열 방향은 같지만 실행 수단이 다릅니다.

구현	Q 재배열 이후
PyTorch Flex Decoding	전용 Triton 커널이 GQA 그룹을 처리
이번 Core ML 구현	표준 배치 행렬곱이 그룹을 행으로 처리

디코드에서는 쿼리 길이 M이 1이어서 M축에 한 행만 있습니다. 같은 KV 헤드를 공유하는 G개의 쿼리 헤드를 이 행 방향에 배치하면 K와 V의 재사용 관계를 자연스럽게 표현할 수 있습니다.

7. Prefill에서는 그룹을 시퀀스 축에 합친다

디코드에서는 q_len=1이므로 Q를 [B, Hkv, G, D]로 만들 수 있습니다. prefill에서는 q_len이 더 크므로 같은 아이디어를 그대로 적용하면 rank-5 텐서가 됩니다.

예를 들어 q_len=64라면 다음과 같습니다.

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

이 행렬곱은 그룹 축에서 4 대 1 브로드캐스트를 요구합니다. 논리적으로는 유효하지만, 백엔드에 따라 K가 그룹 축으로 펼쳐진 중간 버퍼를 만들 수 있습니다. 디코드에서 제거한 것과 비슷한 복사가 다른 위치에 다시 생길 수 있는 셈입니다.

이를 피하기 위해 그룹 축 G를 별도 배치 축으로 두지 않고 쿼리의 시퀀스 축 M에 합칩니다.

[B, Hq, M, D]
= [B, Hkv, G, M, D]
       ↓ G와 M을 하나의 행 축으로 결합
[B, Hkv, G×M, D]

32개 쿼리 헤드와 길이 64의 prefill이라면 다음과 같습니다.

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

행 인덱스를 r, 그룹 안의 쿼리 헤드를 g, 시퀀스 위치를 s라고 하면 다음 관계가 됩니다.

r = g × q_len + s

같은 KV 헤드를 공유하는 네 쿼리 헤드의 64개 토큰이 256개의 연속된 행으로 배치됩니다.

mask를 더할 때는 원래 형상으로 돌아간다

행렬곱 결과는 다음 형상입니다.

[B, Hkv, G×M, ctx]

하지만 causal mask는 일반적으로 [B, 1, M, ctx] 또는 이에 브로드캐스트할 수 있는 형상으로 준비됩니다. 접힌 점수 텐서에 그대로 더하면 G와 M의 의미가 맞지 않습니다.

mask를 G번 복제하면 없애려던 데이터 반복이 mask 쪽에 생길 수 있습니다. 대신 점수 텐서를 원래 쿼리 헤드 형상으로 되돌린 뒤 mask를 더합니다.

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)

# V 행렬곱 직전에 G와 M을 다시 합친다.
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,
)

여기서 그룹을 합칠 수 있는 축은 행을 나타내는 M축입니다.

head_dim에 합치면 서로 다른 쿼리 헤드의 내적이 하나로 더해집니다.
context 축에 합치면 softmax의 정규화 범위가 달라집니다.
M축에 합치면 각 행의 행렬곱과 softmax가 계속 독립적으로 유지됩니다.

잘못된 축을 골라도 shape 검사는 통과할 수 있습니다. 따라서 형상뿐 아니라 헤드 매핑과 출력 일치 테스트를 함께 확인해야 합니다.

decode를 먼저 최적화한 이유

기존 그래프에서는 decode와 prefill 모두 호출당 60개의 tile을 실행했습니다. 하지만 한 번의 호출이 처리하는 토큰 수가 다릅니다.

경로	호출당 처리 토큰	호출당 `tile` 출력	토큰당 `tile` 출력
decode	1	960 MiB	960 MiB
prefill, `q_len=64`	64	960 MiB	15 MiB

prefill도 K와 V를 최대 context length에 맞춰 읽는 그래프라면 호출당 tile 크기는 줄지 않습니다. 다만 한 번의 호출이 64개 토큰을 처리하므로 토큰당 비용은 64분의 1로 나뉩니다.

이 차이 때문에 decode를 먼저 최적화했습니다. prefill 경로에도 Q 재배열을 적용한 측정에서는 1,393 ms에서 1,219 ms로 12.5% 감소했습니다. 다만 이 수치는 디코드 실험처럼 하나의 플래그만 바꾼 완전히 동일한 빌드 A/B가 아니므로, 현재로서는 참고 결과로 보는 편이 정확합니다.

8. 백엔드에 맞는 표현을 선택하는 법

이번 사례에서 얻은 원칙은 다섯 가지입니다.

1. 의미 있는 구조를 가능한 오래 보존한다

repeat_kv로 KV를 32개 헤드로 늘리는 순간 “8개 KV 헤드를 네 쿼리 헤드가 공유한다”는 정보가 사라집니다. 가능하다면 쿼리 헤드 수와 KV 헤드 수를 원래대로 유지하고 GQA라는 의미를 실행기까지 전달하는 편이 좋습니다.

2. 작은 텐서를 재배열한다

디코드에서 Q는 4 KiB이고 레이어 하나의 K 또는 V cache는 4 MiB입니다. 둘 중 하나를 변형해야 한다면 작은 Q의 형상을 바꾸는 편이 훨씬 유리합니다. 이번 reshape는 논리적으로 데이터 복사도 필요하지 않았습니다.

3. 커널과 그래프는 서로 다른 표현 수단을 갖는다

GPU 커널에서 자연스러운 포인터 산술을 값 중심 그래프에 그대로 옮길 수는 없습니다. 반대로 그래프에서는 하나의 값을 여러 연산이 공유하는 fan-out이나 텐서 형상으로 같은 관계를 표현할 수 있습니다.

백엔드의 제약은 계산을 포기해야 한다는 뜻이 아니라, 같은 수학을 다른 형태로 써야 한다는 뜻일 때가 많습니다.

4. decode와 prefill을 별도로 설계한다

둘은 같은 어텐션을 수행하지만 쿼리 길이가 다릅니다. q_len=1인 decode에서는 그룹 축을 직접 행처럼 사용할 수 있고, prefill에서는 그룹 축과 시퀀스 축을 함께 고려해야 합니다.

같은 변환을 두 경로에 기계적으로 적용하면 rank 증가나 불필요한 브로드캐스트가 생길 수 있습니다.

5. 그래프와 실제 기기를 모두 측정한다

MIL에서 tile이 사라졌다는 사실만으로 성능 향상을 단정할 수는 없습니다. reshape가 백엔드 내부에서 재배치를 만들 수도 있고, 캐시나 연산 융합이 예상과 다르게 작동할 수도 있습니다.

반대로 실제 기기 처리량만 보면 왜 빨라졌는지 설명하기 어렵습니다. 그래프 비교, 텐서 크기 계산, 출력 일치 테스트, 실제 기기 A/B를 함께 봐야 원인과 결과를 연결할 수 있습니다.

결론

GQA 구현들이 풀어야 하는 문제는 같습니다.

쿼리 헤드 h가 사용할 KV 헤드는 h // group_size이다.

Hugging Face의 참조 구현은 이 관계를 텐서 반복으로 표현합니다. fused SDPA는 GQA 정보를 백엔드에 전달합니다. Triton과 vLLM은 포인터 산술을 사용하고, CoreAI iOS 경로는 하나의 KV 값을 여러 행렬곱으로 fan-out합니다.

이번 Core ML 구현은 같은 관계를 텐서 형상에 담았습니다. 디코드에서는 쿼리 헤드 축을 [Hkv, G]로 나누고, prefill에서는 G를 시퀀스 축 M과 합쳐 [Hkv, G×M]으로 만들었습니다.

어느 방법이 항상 정답인 것은 아닙니다. 커스텀 커널을 사용할 수 있는지, 그래프 크기가 중요한지, 쿼리 길이가 얼마인지, 백엔드가 reshape와 broadcast를 어떻게 처리하는지에 따라 최선의 표현이 달라집니다.

중요한 것은 수학식만 같다고 실행 비용까지 같지는 않다는 점입니다. 모델의 의미를 백엔드가 이해하고 효율적으로 실행할 수 있는 형태로 끝까지 보존해야 합니다.

참고 문헌

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

1. 모든 구현이 공유하는 하나의 매핑

2. Hugging Face repeat_kv: 텐서를 직접 반복한다

3. PyTorch SDPA: GQA 정보를 백엔드까지 보낸다

4. Triton과 vLLM: KV 헤드를 주소로 선택한다

vLLM 디코드: 쿼리 헤드에서 KV 헤드를 찾는다

vLLM prefill: KV 헤드에서 쿼리 헤드를 펼친다

5. CoreAI iOS SDPA: 하나의 값을 여러 연산이 공유한다

6. 커스텀 커널 없이 텐서 형상으로 표현하기

PyTorch Flex Decoding도 Q를 같은 방향으로 재배열한다

7. Prefill에서는 그룹을 시퀀스 축에 합친다

mask를 더할 때는 원래 형상으로 돌아간다

decode를 먼저 최적화한 이유

8. 백엔드에 맞는 표현을 선택하는 법

1. 의미 있는 구조를 가능한 오래 보존한다

2. 작은 텐서를 재배열한다

3. 커널과 그래프는 서로 다른 표현 수단을 갖는다

4. decode와 prefill을 별도로 설계한다

5. 그래프와 실제 기기를 모두 측정한다

결론

참고 문헌

관련 글

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

[Tech Review] Core AI: 온디바이스 생성형 모델을 위한 변환 스택의 재설계

[Engineering] FlashDecoding++ 이식기 (1) — 비동기 Softmax의 원리

2. Hugging Face `repeat_kv`: 텐서를 직접 반복한다