장문 컨텍스트 추론을 위한 어텐션 공동 설계

이 요약은 AI가 원문을 분석해 생성했습니다. 정확한 내용은 원문 기준으로 확인하세요.

TL;DR

장문·에이전트형 워크로드에서 어텐션이 추론 비용을 지배하게 되면서 모델 설계가 GPU 성능을 결정하는 핵심 요소가 되었다. 프리필은 전체 입력에 대해 큰 매트멀이 생겨 계산 바운드가 되고, 디코드는 토큰별로 KV 캐시를 HBM에서 읽어오는 메모리 바운드가 되므로 두 단계는 서로 다른 최적화 레버를 요구한다. FlashAttention은 Q/K/V 타일 스트리밍과 BMM1→온라인 소프트맥스→BMM2 결합으로 메모리 이동을 줄이고, 분석은 G(그룹 크기)를 키워 디코드 산술 집약도를 약 2×G로 올리는 것이 디코드 효율을 크게 개선한다고 보여준다. 따라서 실무 권장은 디코드용으로 G를 크게 하고 Hsz를 128 또는 256으로 맞추며 KV 상태를 압축·희소화하고 병렬화 전략을 KH 기준으로 설계하라는 것이다.

빠른 이해

새로운 점

어텐션 설계를 GPU 실행 특성에 맞춰 G, Hsz, KV 관리, 그리고 KH 기반 병렬 전략으로 정렬하는 실무형 코디자인 체크리스트를 제시한 점.

핵심 메커니즘

프리필은 GEMM‑M=ISL×G로 큰 매트멀이 계산을 지배하고, 디코드는 GEMM‑M=G로 KV를 HBM에서 스트리밍하는 메모리 비용이 지배적이며 FlashAttention의 BMM1→온라인 소프트맥스→BMM2 결합과 G 증가가 디코드의 산술 집약도를 약 2×G로 끌어올려 토큰당 KV 로드를 줄인다.

핵심 수치

어텐션 비용 비중: 4K → 128K: 18% → 85%- Figure 1의 DeepSeek‑R1 프리필 시간 분해에 기초
디코드 스케일링(그룹 크기): 런타임 약 2× 개선 per G 두 배- Figure 4과 산술 집약도 수식(Equations 5–7) 근거
프리필 스케일: 프리필 런타임 O(ISL²)- 모든 토큰이 서로를 참조하는 연산량 때문에 이차 스케일이 발생

섹션별 상세

어텐션 비용의 증가 문제

대화형·장문 컨텍스트 워크로드가 늘어나면서 어텐션이 전체 추론 비용에서 차지하는 비중이 급증했고, 본문은 이 변화를 수치적 근거와 커널 분석으로 연결해 어텐션 설계가 성능을 좌우한다고 밝힌다. Figure 1에서 4K에서 128K로 갈수록 어텐션 비중이 18%에서 85%로 커지는 관찰을 근거로, 단순 구현 최적화만으로는 부족하고 모델 구조 자체를 GPU 실행 특성에 맞춰 설계해야 한다는 전제가 제시된다. 이 섹션은 문제 맥락과 이후 분석의 동기와 범위를 규정한다.

근거

Figure 1에서 DeepSeek‑R1의 프리필에서 컨텍스트가 4K에서 128K로 늘어나면 어텐션이 차지하는 비중이 18%에서 85%로 증가한다. — Figure 1(프리필 시간 분해) 참조.

프리필과 디코드의 차이

프리필은 입력 전체(ISL)를 병렬로 처리해 GEMM-M이 ISL×G로 매우 커져 계산(매트멀) 바운드가 되고, 디코드는 토큰을 하나씩 생성해 GEMM-M이 G로 작아져 HBM에서 KV를 읽는 메모리 바운드가 된다는 분리된 문제 구조를 제시한다. 산술 집약도(총 FLOPs ÷ 총 바이트)가 어느 쪽 병목을 만드는지 결정하며, 프리필은 ridge 위에 있어 계산 병목이고 디코드는 ridge 아래에 있어 메모리 병목임이 수식과 도표로 확인된다. 프리필과 디코드가 서로 다른 입력 길이와 KV 접근 패턴을 가지므로 각각 별도로 최적화해야 하는 이유가 이 섹션에서 명확해진다.

근거

산술 집약도는 총 FLOPs를 총 바이트로 나눈 값이며, 프리필은 ridge 위의 계산 바운드이고 디코드는 ridge 아래의 메모리 바운드이다. — Equation 1 및 Figure 2(roofline 모델) 참조.

FlashAttention과 GEMM 모양

FlashAttention 커널은 전체 어텐션 행렬을 재구성하지 않고 Q/K/V 타일을 SRAM으로 스트리밍해 BMM1(점수 계산)→온라인 소프트맥스(러닝 맥스·합산)→BMM2(가중치 적용)를 단일 패스로 결합하며 이 BMM들의 (Batch, M, N, K) 모양이 산술 집약도와 성능을 결정한다. 테이블 3의 GEMM 차원으로부터 프리필과 디코드에서의 M, N, K가 어떻게 달라지는지 수치적으로 도출하고, 이 차이가 프리필을 ISL² 연산으로, 디코드를 KVSL 기반의 메모리 전송으로 만든다고 연결한다. 커널 내부에서 BMM은 Tensor Core, 소프트맥스는 특수연산유닛에서 실행되며 이 실행 경로가 병목 분석의 근거가 된다.

근거

FlashAttention은 Q/K/V 타일을 스트리밍해 BMM1→온라인 소프트맥스→BMM2를 한 번에 수행함으로써 전체 어텐션 행렬을 재구성하지 않고 메모리 접근을 줄인다. — FlashAttention 커널 도식과 BMM 단계 설명(본문 및 Figure 3) 참조.

그룹 크기(G)의 영향

그룹 크기 G는 하나의 KV head에 묶인 쿼리 헤드 수로, 프리필 단계에서는 ISL이 지배적이어서 G를 크게 해도 런타임 변화가 거의 없지만 디코드 단계에서는 GEMM‑M이 G가 되어 G를 늘리면 디코드 산술 집약도가 약 2×G로 증가해 토큰당 로드되는 KV 양이 줄고 GPU 이용률이 개선된다. 본문 수식은 프리필에서 산술 집약도가 2×ISL/(1+1/G)로 G→∞일 때 2×ISL에 수렴함을 보이며, 디코드에서는 산술 집약도가 약 2×G로 유지되어 G 증가가 디코드에 큰 이득을 준다는 수치적 증거를 제시한다. Figure 4의 실측 결과는 G를 두 배로 늘릴 때 디코드 런타임이 대략 2배 개선되는 경향을 관찰해 이 메커니즘을 뒷받침한다.

근거

디코드에서 G(그룹 크기)를 두 배로 하면 런타임이 대략 2배 개선되는 경향이 관측되며, 산술 집약도는 약 2×G에 비례한다. — Equations 5–7 및 Figure 4의 실측 결과 참조.

헤드 차원(Hsz)의 트레이드오프

헤드 차원 Hsz는 산술 집약도 비율 자체에는 영향을 주지 않지만 매트멀이 차지하는 작업량과 KV 상태 크기, 그리고 GPU 타일·메모리 정렬에 민감해 런타임에 영향을 미친다. GPU가 128바이트 전송 및 64/128 타일 크기에 최적화되어 있으므로 Hsz=128 또는 256이 타일 채움과 TMEM 한도 사이의 효율적 균형을 제공하며 Hsz=64는 사실상 128 타일 비용을 지불하게 된다는 근거가 제시된다. Figure 5의 측정은 프리필에서 Hsz 증가가 소프트맥스 고정비용 때문에 선형 증가보다 완만하고, 디코드에서는 KV 바이트 증가로 런타임이 대략 Hsz에 비례해 증가함을 보여준다.

시퀀스 길이(ISL/KVSL)의 스케일

프리필은 모든 토큰이 서로를 참조하므로 연산이 ISL²로 증가해 시퀀스 길이에 대해 이차적으로 확장하는 반면, 디코드는 각 토큰 생성 시 전체 KV 캐시를 읽기 때문에 KVSL에 대해 선형으로 런타임이 증가한다. 이 때문에 장문 입력이 많을수록 프리필 비용이 더 빠르게 커지며, 긴 KV를 가진 멀티턴·에이전트 워크로드에서는 디코드 비용이 지배적일 수 있다. Figure 6의 실측은 프리필이 ISL 증가에 대해 대체로 O(n²) 스케일을 따르고 디코드는 O(n) 스케일을 따른다는 구체적 증거를 제공한다.

병렬화 전략과 KV 헤드 수(KH)

텐서 병렬화(TP)는 어텐션 헤드를 GPU들에 샤딩해 GEMM 모양을 유지하지만 TP가 KV 헤드 수 KH를 초과하면 동일한 KV head가 여러 랭크에 복제되어 KV 상태가 중복되고 메모리·대역폭 오버헤드가 발생한다. 따라서 TP는 KH 이하로 유지하는 것이 효율적이며, KH가 적은 모델(Nemotron 3의 예: KH=2 또는 MQA의 KH=1)은 ADP(요청 샤딩)나 KVP(긴 KV를 장치들에 분산)같은 대안적 스케일링을 써야 한다. TensorRT-LLM은 Wide EP(ADP+EP)와 Helix Parallelism(KVP+EP)으로 이 문제를 해결하는 구현 사례를 제시한다.

근거

텐서 병렬화는 TP > KH이면 KV 캐시의 중복 복사를 초래하므로, 효율적 확장을 위해서는 TP ≤ KH를 유지해야 한다. — 섹션의 병렬화 설명 및 Figure 7(샤딩 전략) 참조.

실무용 권장 체크리스트

네 가지 실무 권장은 디코드 성능을 우선한 G를 크게 설정하고, Hsz는 128 또는 256으로 맞추며, 유효 KV 상태를 압축·희소화하거나 하이브리드 계층으로 줄이고, 병렬화 전략은 KH에 맞춰 TP ≤ KH를 지키는 것이다. 본문은 각 권장이 왜 성능·상호작용성을 동시에 개선하는지 GEMM·아키텍처 근거와 실측 결과로 연결해 보여주며, Nemotron 3 등의 하이브리드 아키텍처가 실제로 이런 설계를 따른 사례임을 언급한다. 이 섹션은 모델 설계자가 GPU 하드웨어 특성과 커널 동작을 고려해 구조적 선택을 할 때 참고할 구체적 규칙을 모아둔다.

용어 해설

산술 집약도(Arithmetic Intensity): — 산술 집약도는 GPU 작업이 얼마나 계산 중심인지 결정하는 지표로 FLOPs(연산량)를 메모리에서 읽는 바이트 수로 나눈 값이다. 이 값이 낮으면 메모리 대역폭에, 높으면 연산 유닛에 병목이 생기며 본문은 이를 통해 프리필(계산 바운드)과 디코드(메모리 바운드)를 구분한다. 수식은 총 FLOPs를 총 바이트로 나눈 형태로 표현되며, ridge point가 경계 역할을 한다.
FlashAttention: — FlashAttention은 전체 어텐션 행렬을 메모리에 올리지 않고 Q/K/V 타일을 스트리밍해 BMM1→온라인 소프트맥스→BMM2를 한 번의 패스로 결합하는 커널 설계이다. BMM은 Tensor Core에서 실행되고 소프트맥스는 특수 연산 유닛에서 처리되어 메모리 이동을 줄이면서 계산을 병렬화한다. 이 구조가 GEMM 모양과 arithmetic intensity 분석의 근거가 된다.
그룹 크기(G)(Group Size (G)): — 그룹 크기 G는 하나의 KV head를 공유하는 query head 수로 MHA는 G=1, GQA는 G=4,8,16 등이며 MQA는 G=QH이다. 디코드 단계에서 GEMM-M이 G가 되어 G가 커질수록 토큰당 로드되는 KV 데이터가 줄고 연산 이용률이 개선되어 디코드의 산술 집약도가 약 2×G로 증가한다. 반면 프리필은 ISL에 지배되어 G에 둔감하다.
KV 캐시(KV cache): — KV 캐시는 디코드 중에 각 토큰의 키와 값을 저장한 상태로, 디코드 단계에서 HBM에서 읽어오는 대상이 된다. KVSL로 표현되는 KV 캐시 길이에 따라 디코드의 메모리 트래픽이 선형으로 증가하며, 압축·희소화·슬라이딩 윈도 기법으로 유효 KV 상태를 줄이면 디코드 비용을 낮출 수 있다.
텐서 병렬화(TP)(Tensor Parallelism (TP)): — 텐서 병렬화는 어텐션 헤드를 GPU들에 샤드해 각 장치가 QH/TP 쿼리 헤드와 KH/TP KV 헤드를 소유하게 하는 방식이다. 중요한 한계는 TP가 KV 헤드 수(KH)를 초과하면 동일한 KV head가 복제되어 메모리·대역폭 오버헤드가 발생한다는 점이며 따라서 TP ≤ KH를 권장한다.

기술

FlashAttention
FP8
Tensor Cores
HBM
TensorRT-LLM
Speculative decoding

언급된 리소스

문서원문 링크