들어가며: 어텐션이 병목이 된 시대
AI 에이전트와 롱컨텍스트 워크로드가 일상이 되면서, 모델이 처리해야 하는 컨텍스트 길이는 계속 늘어나고 있습니다. 문제는 이 과정에서 어텐션 연산이 추론 시간에서 차지하는 비중이 급격히 커진다는 점이에요. 예전에는 어텐션이 '그냥 있는 레이어'였다면, 이제는 어텐션을 어떻게 설계하느냐가 모델의 추론 성능을 좌우합니다.
이 글은 NVIDIA 개발자 블로그의 분석을 바탕으로, 어텐션 성능을 결정하는 네 가지 핵심 변수—Group Size(G), Head Dimension(Hsz), Sequence Length, Tensor Parallelism—를 정리하고, 바로 적용 가능한 네 가지 가이드라인으로 압축합니다. 모델을 직접 학습시키는 사람이 아니더라도, LLM 서빙 성능을 튜닝하는 백엔드 개발자나 MLOps 엔지니어에게 실무적인 인사이트를 줄 거예요.
💡 핵심 전제: 모든 분석은 GEMM Shape 산술과 FP8 기반 Prefill/Decode 커널 측정값에 근거합니다. 감이 아니라 숫자로 이야기합니다.
Prefill과 Decode는 완전히 다른 문제다
어텐션 최적화의 첫 단추는 Prefill과 Decode를 분리해서 보는 것입니다. 많은 분들이 하나로 뭉뚱그려 생각하는데, 두 단계는 병목 지점이 정반대예요.
| 구분 | Prefill | Decode |
|---|---|---|
| Query 길이 | 전체 입력 (ISL 토큰) | 1 토큰 |
| KV 컨텍스트 | 프롬프트 (ISL 토큰) | 전체 KV 캐시 (KVSL 토큰) |
| Attention GEMM-M | ISL × G (큼) | G (작음) |
| 주요 병목 | 연산 (matmul + softmax) | HBM 대역폭 (메모리) |
Prefill은 프롬프트 전체를 병렬 처리하기 때문에 Compute-bound(연산 집약적)이고, Decode는 토큰을 하나씩 생성하기 때문에 Memory-bound(메모리 집약적)입니다. 여기서 Speculative Decoding을 쓰면 Decode의 GEMM-M이 커지면서 Compute-bound 쪽으로 이동할 수 있어요.
FlashAttention이 하는 일
FlashAttention은 전체 어텐션 행렬을 메모리에 물질화하지 않고, Q/K/V 타일을 HBM에서 SRAM으로 스트리밍하면서 세 단계를 한 번에 융합합니다.
- BMM1: Query와 Key를 스코어링
- Online Softmax: running max/sum으로 정규화
- BMM2: Value에 가중치 적용
BMM은 Tensor Core에서, softmax 지수 연산은 SFU(Special-Function Unit)에서 돌아갑니다. 이 구조가 뒤에 나오는 산술 강도 분석의 기반이에요.
Group Size(G): Decode의 핵심 레버
Group Size는 하나의 KV head를 공유하는 Query head의 수입니다. MHA는 G=1, GQA는 G=4/8/16, MQA는 G=QH입니다.
# 산술 강도(Arithmetic Intensity) 공식 (FP8 KV 캐시 가정)
# Prefill: AI = 2 * G * ISL / (G + 1) → G가 커지면 2*ISL에 수렴
# Decode: AI = 2 * G * KVSL / (G + KVSL) ≈ 2 * G (KVSL >> G일 때)
def arithmetic_intensity_prefill(G, ISL):
return 2 * G * ISL / (G + 1)
def arithmetic_intensity_decode(G, KVSL):
return 2 * G * KVSL / (G + KVSL)
# 예시: G를 8에서 16으로 올려도 Prefill은 6% 미만 개선
print(arithmetic_intensity_prefill(8, 32768)) # ~15,990
print(arithmetic_intensity_prefill(16, 32768)) # ~15,998
# 반면 Decode는 G에 거의 선형 비례
print(arithmetic_intensity_decode(8, 32768)) # ~15.99
print(arithmetic_intensity_decode(16, 32768)) # ~31.99
핵심 포인트: Prefill은 G에 둔감하고, Decode는 G에 민감합니다. G를 두 배로 올리면 Decode 산술 강도도 두 배가 되면서 메모리 트래픽이 줄고 GPU 활용률이 올라가요. 실제로 G=1(MHA)에서 G=8(GQA)로 바꾸면 Decode 런타임이 약 2배 빨라집니다.
Head Dimension(Hsz): 128 또는 256이 정답
Hsz는 산술 강도 자체를 바꾸지 않습니다. FLOPs와 Bytes가 동시에 두 배가 되니까요. 그런데도 런타임이 늘어나는 이유는 커널이 세 가지 작업을 서로 다른 비율로 수행하기 때문입니다.
- Matmul: Hsz에 비례하지만, 128 정렬 단위로 처리되므로 Hsz=64는 128 타일 비용을 그대로 지불합니다.
- Memory(KV state): HBM이 128바이트 단위로 움직이므로 Hsz가 128의 배수일 때 효율적입니다.
- Softmax: Hsz와 무관! Attention-score 행렬(Query × Key)에만 의존합니다.
# Softmax 연산량은 Hsz와 무관
def softmax_ops_prefill(PB, QH, ISL):
return PB * QH * ISL ** 2
def softmax_ops_decode(DB, QH, KVSL):
return DB * QH * KVSL
# Hsz를 64→128로 올려도 softmax 비용은 그대로
print(softmax_ops_prefill(1, 32, 8192)) # 동일
Prefill에서는 Hsz가 커질수록 matmul 비중이 늘면서 softmax가 상대적으로 숨어버립니다. Decode에서는 Hsz가 커지면 KV 바이트가 늘어나 런타임이 증가하죠. 그래서 128과 256이 스위트 스팟입니다. Hsz=64는 128 타일 비용을 지불하고, Hsz≥512는 TMEM 용량 한계에 부딪힙니다.
🔗 컨텍스트 엔지니어링 관점에서 모델 설계를 보고 싶다면 백그라운드 코딩 에이전트를 위한 콘텐츠 엔지니어링 가이드도 함께 읽어보세요.
Sequence Length: Prefill은 제곱, Decode는 선형
시퀀스 길이는 Prefill과 Decode에 전혀 다른 방식으로 들어갑니다.
- Prefill: ISL² 작업 (모든 토큰이 모든 토큰을 봄). KV 트래픽은 ISL에 비례. → 산술 강도가 ISL에 선형으로 증가 → Compute-bound 유지
- Decode: 각 스텝마다 KVSL 전체를 읽음. 산술 강도는 2×G 근처 고정 → Memory-bound 유지
즉, ISL을 두 배로 늘리면 Prefill은 대략 4배 느려지고, KVSL을 두 배로 늘리면 Decode는 2배 느려집니다. 짧은 시퀀스에서는 setup/post-processing 오버헤드 때문에 이 비율보다 완만하지만, 길이가 커지면 이론값에 수렴해요.
Tensor Parallelism: KH가 전략을 결정한다
TP는 어텐션 head를 GPU에 분산시킵니다. 각 GPU는 QH/TP개의 Query head와 KH/TP개의 KV head를 갖죠. 문제는 TP > KH가 되는 순간입니다. 하나의 그룹이 여러 랭크에 걸치면서 KV head 복제가 발생하고, 메모리와 대역폭 오버헤드만 늘어납니다.
# TP 제약 조건
def validate_tp_strategy(num_kv_heads, tp_size):
if tp_size > num_kv_heads:
print(f"⚠️ TP={tp_size} > KH={num_kv_heads}: KV head 복제 발생")
print(" → ADP(Attention Data Parallelism) 또는 KVP(KV Parallelism) 고려")
return False
print(f"✅ TP={tp_size} ≤ KH={num_kv_heads}: 정상 샤딩")
return True
validate_tp_strategy(2, 4) # ⚠️ 복제 발생
validate_tp_strategy(8, 4) # ✅ OK
Nemotron 3처럼 KV head가 2개뿐인 모델은 TP를 금방 소진합니다. 이때는 Wide EP(ADP + EP)나 Helix Parallelism(KVP + EP)으로 전환해야 합니다. TensorRT-LLM이 이 두 전략을 모두 지원합니다.
⚠️ 주의사항과 한계
- Speculative Decoding의 함정: 초안 토큰 수(D)이 충분히 커지면 Decode가 Compute-bound로 넘어가면서 기존 최적화 가정이 깨질 수 있습니다.
- Prefix Caching 환경: 에이전트/멀티턴 앱에서 흔한데, 짧은 ISL + 긴 prefix cache 조합에서는 Prefill이 Decode처럼 동작합니다. 이 경우 G와 Hsz 튜닝 전략을 다시 짜야 해요.
- 측정 환경 의존성: 이 글의 수치는 FP8 KV 캐시와 특정 GPU 아키텍처 기준입니다. BF16이나 다른 아키텍처에서는 최적점이 이동할 수 있습니다.
- 모델 재학습 비용: G나 Hsz를 바꾸려면 모델을 다시 학습시켜야 합니다. 서빙 단계에서 바꿀 수 있는 건 TP 전략과 KV 캐시 압축 정도예요.
🔗 데이터 레이어 일관성과 컨텍스트 아키텍처 설계에 관심 있다면 AI 에이전트의 침묵의 독 데이터 레이어 일관성, 이제는 컨텍스트 아키텍트가 될 때를 참고하세요.
결론: 네 가지 가이드라인 요약
이 글의 모든 분석은 다음 네 가지 실전 체크리스트로 압축됩니다.
| # | 가이드라인 | 핵심 이유 |
|---|---|---|
| 1 | G는 Decode 기준으로 높게 | Prefill은 G에 둔감, Decode 산술 강도 ≈ 2×G |
| 2 | Hsz = 128 또는 256 | GPU 타일/128B 전송 정렬 + TMEM 예산 |
| 3 | 유효 KV 상태 줄이기 | Prefill은 ISL², Decode는 KVSL 선형 증가 |
| 4 | KH가 병렬화 전략 결정 | TP ≤ KH 유지, 부족하면 Wide EP/Helix |
한국 개발 생태계에서의 적용 맥락
국내 LLM 서빙 환경은 크게 두 갈래로 나뉩니다. 대기업/클라우드 사업자는 H100/H200 클러스터를 보유하고 TensorRT-LLM이나 vLLM 기반 최적화를 직접 수행할 수 있지만, 스타트업과 중소 SI는 A100이나 L40S 같은 제한된 GPU로 멀티테넌트 서빙을 해야 하는 경우가 많아요. 후자의 경우 G를 높인 GQA 모델(예: Llama 3 계열)을 선택하는 것만으로도 동일 하드웨어에서 Decode 처리량을 크게 끌어올릴 수 있습니다. 또한 국내 금융/의료 도메인처럼 롱컨텍스트 + 온프레미스 조합이 필요한 곳에서는 KV 캐시 압축과 sliding-window attention 적용이 선택이 아닌 필수가 되고 있어요.
다음 단계 학습 방향
- Sparse Attention: 이 글은 dense attention 기준입니다. Sparse attention이 등장하면 병목 구조가 다시 바뀝니다.
- vLLM/TensorRT-LLM 실습: 실제 서빙 프레임워크에서 GQA/MQA 설정이 어떻게 반영되는지 코드 레벨로 확인해 보세요.
- Roofline 모델 직접 그려보기: 자신이 서빙하는 모델의 산술 강도를 계산해서 ridge point 대비 어디에 있는지 파악하면, 어떤 최적화가 유효한지 감이 옵니다.
어텐션은 이제 '구현 디테일'이 아니라 아키텍처 설계 결정입니다. 모델을 만드는 사람이든 서빙하는 사람이든, 이 네 가지 축을 알고 있으면 같은 GPU에서 더 많은 토큰을 더 빠르게 뽑아낼 수 있어요. 🚀
근거자료: NVIDIA Developer Blog - Co-Designing AI Model Attention for Fast, Interactive Long-Context Inference
함께 보면 좋은 글