blank

LLM 엔지니어가 알아야 할 GPU 아키텍처: Ampere → Hopper → Blackwell

Sat, 11 Apr 2026 09:00:00 +0000

왜 GPU 아키텍처를 알아야 하는가

LLM 엔지니어에게 GPU는 “빠른 연산기” 이상의 의미를 가진다. 모델 크기, 배치 사이즈, 시퀀스 길이, 학습/추론 전략 등 거의 모든 엔지니어링 결정이 GPU의 제약에 의해 결정되기 때문이다.

“70B 모델을 학습하려면 GPU 몇 장이 필요한가?” → 메모리 용량
“KV cache가 얼마나 들어가는가?” → 메모리 대역폭
“FP8로 학습하면 속도가 얼마나 빨라지는가?” → 정밀도별 연산 성능
“FlashAttention을 쓰면 왜 빨라지는가?” → 메모리 계층과 IO 병목
“H100에서 A100 대비 실제로 얼마나 빠른가?” → 실제 활용률과 병목 분석

이 글에서는 현재 LLM 엔지니어가 주로 사용하는 세 가지 GPU — A100 (Ampere), H100 (Hopper), B200 (Blackwell) — 를 LLM 학습과 추론의 관점에서 비교한다.

1. 전체 스펙 비교

	A100 SXM	H100 SXM	B200 SXM
아키텍처	Ampere (2020)	Hopper (2022)	Blackwell (2025)
공정	TSMC 7nm	TSMC 4N	TSMC 4NP
트랜지스터	54.2B	80B	208B (듀얼 다이)
SM 수	108	132	~148 × 2 다이
Tensor Core	3세대	4세대	5세대
메모리	80GB HBM2e	80GB HBM3	192GB HBM3e
메모리 대역폭	2.0 TB/s	3.35 TB/s	8.0 TB/s
NVLink	600 GB/s	900 GB/s	1.8 TB/s
TDP	400W	700W	1,000W

2. LLM 학습 관점: “얼마나 빨리 학습할 수 있는가”

2.1 Tensor Core 성능: 정밀도가 핵심

LLM 학습에서 가장 많은 시간을 차지하는 연산은 행렬곱(MatMul)이다. Linear layer, Attention의 $QK^\top$와 $PV$ 모두 matmul이다. Tensor Core의 정밀도별 처리량이 곧 학습 속도를 결정한다.

정밀도	A100	H100	B200	용도
FP32	19.5 TF	67 TF	80 TF	디버깅, 정밀 연산
TF32	156 TF	495 TF	~1,000 TF	PyTorch 기본 학습
BF16	312 TF	989 TF	2,250 TF	LLM 학습 표준
FP8	—	1,979 TF	4,500 TF	대규모 학습 가속
FP4	—	—	9,000 TF	추론 최적화

실무 포인트

BF16이 LLM 학습의 사실상 표준이다. FP16 대비 dynamic range가 넓어서 loss scaling 없이도 안정적이다. A100의 312 TFLOPS → H100의 989 TFLOPS → B200의 2,250 TFLOPS로, 세대마다 약 2-3배씩 빨라진다.

FP8 학습은 H100부터 지원된다. Transformer Engine이 layer별로 FP8/BF16을 자동 전환하여, BF16 대비 약 1.5-2배 추가 speedup을 제공한다. 다만 학습 안정성이 모델과 데이터에 따라 달라서, 아직 모든 상황에서 사용하지는 않는다.

FP4는 주로 추론용이다. 학습에 사용하기에는 정밀도가 부족하지만, 양자화된 모델의 추론 속도를 극대화할 수 있다.

2.2 실제 활용률: 이론 vs 현실

이론 성능은 좋지만, 실제로 달성하는 Model FLOPs Utilization (MFU)은 다르다.

	A100	H100	B200
이론 BF16	312 TF	989 TF	2,250 TF
일반적인 MFU	40-55%	35-50%	30-45% (추정)
실제 학습 처리량	~130-170 TF	~350-500 TF	~700-1000 TF

MFU가 세대가 올라갈수록 약간 낮아지는 경향이 있다. 이는 Tensor Core가 빨라지는 속도를 메모리 대역폭과 interconnect가 따라가지 못하기 때문이다. 이것이 바로 FlashAttention 같은 IO-aware 알고리즘이 점점 더 중요해지는 이유다.

2.3 메모리 대역폭: “Tensor Core를 먹여살릴 수 있는가”

행렬곱을 수행하려면 Tensor Core에 데이터를 공급해야 한다. 아무리 Tensor Core가 빨라도 데이터가 도착하지 않으면 놀게 된다.

Arithmetic Intensity (연산 강도)를 계산해보면:

\[\text{AI} = \frac{\text{FLOPs}}{\text{Bytes}} = \frac{2mnk}{2(mk + kn + mn)} \approx \frac{mnk}{mk + kn + mn}\]

LLM의 대표적 연산인 $[B \times S, H] \times [H, 4H]$ (FFN의 첫 번째 linear):

$B \times S = 2048$ (batch × seq), $H = 4096$, $4H = 16384$
AI ≈ $\frac{2048 \times 4096 \times 16384}{2 \times (2048 \times 16384 + 16384 \times 4096 + 2048 \times 4096)} \approx 1365$

이 정도의 AI면 모든 GPU에서 compute-bound이므로, 큰 MatMul은 Tensor Core 성능에 비례한다.

하지만 Attention은 다르다. $QK^\top$는 $[B \times H, S, d] \times [B \times H, d, S]$로, $d = 128$이면 AI가 약 64로 낮다. 시퀀스가 짧으면 memory-bound가 될 수 있다. 이때 메모리 대역폭이 중요하다.

	A100	H100	B200
BF16 Tensor Core	312 TF	989 TF	2,250 TF
메모리 대역폭	2.0 TB/s	3.35 TB/s	8.0 TB/s
Ops:Byte	156	295	281

Ops:Byte가 높을수록 memory-bound 연산에서 Tensor Core가 놀 확률이 높다. H100과 B200은 A100 대비 이 비율이 2배 가까이 높아서, FlashAttention 같은 IO 최적화가 더욱 중요하다.

3. LLM 추론 관점: “얼마나 빨리 토큰을 생성하는가”

LLM 추론은 학습과 전혀 다른 특성을 가진다.

3.1 Prefill vs Decode

단계	특성	병목
Prefill	프롬프트 전체를 한 번에 처리, 큰 MatMul	Compute-bound
Decode	토큰 하나씩 생성, 배치 1의 작은 MatMul	Memory bandwidth-bound

Decode가 추론의 대부분 시간을 차지한다. Decode에서는 weight 전체를 읽어서 하나의 토큰만 생성하므로, AI가 매우 낮다 (≈ 1). 이때 성능은 순수하게 메모리 대역폭에 의해 결정된다.

3.2 Decode 처리량 계산

\[\text{Tokens/s} = \frac{\text{Memory Bandwidth}}{2 \times \text{Model Parameters}} \text{ (BF16 기준)}\]

모델	A100 (2 TB/s)	H100 (3.35 TB/s)	B200 (8 TB/s)
Llama-7B (14 GB)	143 tok/s	239 tok/s	571 tok/s
Llama-70B (140 GB)	14.3 tok/s	23.9 tok/s	57.1 tok/s
Llama-405B (810 GB)	OOM	OOM	OOM (멀티 GPU)

이 계산은 단일 GPU, 배치 1, BF16 기준이다. 실제로는:

양자화 (INT8, FP8, INT4)를 적용하면 모델 크기가 줄어 처리량 증가
배치 크기를 키우면 compute-bound 영역으로 이동하여 Tensor Core 활용 가능
KV cache가 추가 메모리를 소비

3.3 KV Cache: “시퀀스가 길면 메모리가 부족하다”

Autoregressive 생성에서 이전 토큰의 Key, Value를 캐싱한다. KV cache 크기:

\[\text{KV Cache} = 2 \times L \times 2 \times n_h \times d \times S \times B \text{ bytes (BF16)}\]

Llama-70B ($L=80, n_h=8 \text{ (GQA)}, d=128$) 기준:

시퀀스 길이	Batch 1	Batch 32
4K	0.16 GB	5.2 GB
32K	1.3 GB	41.6 GB
128K	5.2 GB	166.4 GB

A100 80GB에서는 128K 시퀀스 + 배치 32가 불가능하다 (모델 140GB + KV 166GB > 80GB). H100도 마찬가지. B200의 192GB에서야 가능해지며, 이것이 메모리 용량 증가의 실질적 의미다.

3.4 정밀도와 양자화

추론에서는 학습보다 낮은 정밀도를 사용할 수 있다.

정밀도	모델 크기 (70B)	지원 GPU	품질 영향
BF16	140 GB	A100+	기준
FP8	70 GB	H100+	거의 없음
INT8 (W8A8)	70 GB	A100+	미미
INT4 (GPTQ/AWQ)	35 GB	A100+	약간
FP4	35 GB	B200만	모델 의존적

FP8이 현재 가장 실용적인 선택이다. 모델 크기가 절반이 되어 메모리 대역폭 2배 + Tensor Core 2배로, BF16 대비 이론적으로 4배 빠르다. H100의 Transformer Engine이 FP8을 자동 관리해준다.

B200에서 추가된 FP4는 이론적으로 BF16 대비 8배 빠르지만, 품질 저하 없이 사용하려면 정교한 양자화 기법(GPTQ, AWQ, SqueezeLLM 등)이 필요하다.

4. 멀티 GPU 스케일링: “GPU를 더 쓰면 비례해서 빨라지는가”

4.1 NVLink: GPU 간 통신

LLM은 단일 GPU에 올라가지 않으므로, Tensor Parallelism (TP), Pipeline Parallelism (PP), Data Parallelism (DP)을 조합한다. 이때 GPU 간 통신 대역폭이 스케일링을 결정한다.

	A100	H100	B200
NVLink	600 GB/s	900 GB/s	1,800 GB/s
GPU 간 대역폭 / 메모리 대역폭	30%	27%	23%

NVLink 대역폭이 절대적으로는 증가하지만, 메모리 대역폭 대비 비율은 오히려 감소하고 있다. 이는 TP를 많이 쓸수록 통신 overhead가 상대적으로 커진다는 의미다.

실무 가이드라인

TP 8 이하를 유지하는 것이 효율적. TP 16 이상은 all-reduce 비용이 급증
70B 모델: TP=8이면 GPU당 약 9-10GB로 적절
405B 모델: TP=8 + PP=4 또는 TP=8 + PP=8이 일반적
DeepSpeed ZeRO-3 + FSDP: DP 방향으로 메모리 분산, 통신은 gradient 동기화만

4.2 DGX 시스템 비교

	DGX A100	DGX H100	DGX B200
GPU 수	8 × A100	8 × H100	8 × B200
총 메모리	640 GB	640 GB	1,536 GB
총 BF16 성능	2.5 PF	7.9 PF	18 PF
총 NVLink BW	4.8 TB/s	7.2 TB/s	14.4 TB/s
가격	~$200K	~$300K	~$500K

DGX B200 한 대에서 1.5TB 메모리를 사용할 수 있어, 405B 모델(BF16 810GB)도 단일 노드에서 추론이 가능하다.

5. 아키텍처별 핵심 신기능과 LLM 영향

Ampere (A100)

기능	LLM 영향
BF16/TF32	혼합 정밀도 학습의 표준 확립
MIG	하나의 A100을 최대 7개로 분할 → 소형 모델 추론 효율화
3세대 NVLink	8-GPU TP 가능

Hopper (H100)

기능	LLM 영향
FP8 + Transformer Engine	학습/추론 모두 FP8 자동 적용 → 2배 speedup
TMA (Tensor Memory Accelerator)	FlashAttention-3의 핵심 — HBM↔SMEM 전송 자동화
WGMMA (비동기 MMA)	연산과 데이터 전송을 겹쳐 실행 → GPU 활용률 향상
Warp Specialization	Producer/Consumer warp 분리로 파이프라인 최적화

Blackwell (B200)

기능	LLM 영향
FP4 네이티브	INT4 양자화 추론이 하드웨어 수준에서 가속
192GB HBM3e	70B 모델을 양자화 없이 단일 GPU에 탑재 가능
8 TB/s 대역폭	Decode 처리량 2.4배 향상 → 추론 latency 대폭 감소
Tensor Memory (256KB/SM)	MMA accumulator 전용 메모리 → register 압력 해소
2-CTA MMA	큰 타일의 MMA를 2개 CTA가 협력 실행 → SMEM 트래픽 절반
듀얼 다이	2개 다이를 10 TB/s로 연결 → 사실상 하나의 거대 GPU

6. 비대칭 스케일링: 왜 알고리즘이 중요한가

세대별로 하드웨어 발전 속도가 균일하지 않다는 것이 핵심이다.

하드웨어	A100 → H100	H100 → B200	2세대 합산
Tensor Core (BF16)	3.2×	2.3×	7.2×
메모리 대역폭	1.7×	2.4×	4.0×
MUFU (exp 등)	~1×	~1×	~1×
SMEM 대역폭	~1×	~1×	~1×

Tensor Core는 2세대에 걸쳐 7배 빨라졌지만, softmax의 exp를 계산하는 MUFU는 그대로이다. 이것이 의미하는 바:

단순히 GPU를 바꾸는 것만으로는 성능이 비례 증가하지 않는다. Tensor Core가 7배 빨라져도, exp가 병목이면 전체 속도는 제한된다.
FlashAttention 같은 IO/compute-aware 알고리즘이 점점 더 중요해진다. 하드웨어의 비대칭을 소프트웨어로 보상해야 한다.
세대별로 다른 최적화 전략이 필요하다. FA1/FA2는 HBM IO를 줄이고, FA3는 GEMM과 softmax를 겹치고, FA4는 exp를 소프트웨어로 에뮬레이션한다.

7. 실무 가이드라인

GPU 선택

상황	추천
7B 모델 fine-tuning	A100 80GB 1장이면 충분
70B 모델 학습	H100 8장 (DGX H100) 이상
70B 모델 추론 (저비용)	A100 + INT4 양자화
70B 모델 추론 (고성능)	H100 + FP8
405B 모델 학습	H100/B200 수백 장
128K+ long context 추론	B200 (192GB 메모리 필수)

정밀도 선택

정밀도	학습	추론	주의사항
BF16	✅ 기본	✅ 안정적	—
FP8	⚠️ H100+	✅ 추천	Transformer Engine 필요
INT8	❌	✅	GPTQ/AWQ 양자화 필요
INT4	❌	✅ 비용 효율	품질 저하 모니터링 필요
FP4	❌	⚠️ B200만	아직 초기 단계

FlashAttention 버전 선택

GPU	추천 FA 버전	이유
A100	FA2	A100에 최적화된 알고리즘
H100	FA3 (또는 FA2)	WGMMA/TMA 활용, FP8 지원
B200	FA4	Blackwell 전용 파이프라인, SW exp

마치며

GPU 아키텍처를 이해하면 “왜 이 설정에서 성능이 안 나오는지”, “왜 이 최적화가 효과적인지”를 근본적으로 이해할 수 있다. LLM 엔지니어에게 가장 중요한 insight는:

LLM 학습은 BF16 Tensor Core 성능에 의해 결정되지만, 실제 활용률은 메모리 대역폭과 IO 패턴에 의해 제한된다.
LLM 추론(decode)은 거의 순수하게 메모리 대역폭에 의해 결정된다. Tensor Core 성능은 decode에서 거의 의미가 없다.
하드웨어가 비대칭적으로 발전하면서, 소프트웨어 최적화(FlashAttention, 양자화, KV cache 관리)의 중요성이 점점 커지고 있다.

GPU별 최적화 사례가 궁금하다면: FlashAttention (A100), FlashAttention-2 (A100), FlashAttention-3 (H100), FlashAttention-4 (B200)를 참고하자.

참고 문헌

FlashAttention-4: Algorithm and Kernel Pipelining Co-Design for Asymmetric Hardware Scaling

Sat, 11 Apr 2026 03:00:00 +0000

FlashAttention-4: Algorithm and Kernel Pipelining Co-Design for Asymmetric Hardware Scaling

Introduction

FlashAttention-3는 Hopper GPU(H100)의 비동기 실행과 FP8을 활용하여 GPU 이론 성능의 75%를 달성했다. 하지만 FA3는 H100에 특화되어 있으며, AI 산업은 이미 Blackwell 기반 시스템(B200, GB200)으로 빠르게 전환하고 있다.

문제는 Blackwell이 Hopper와 근본적으로 다른 성능 특성을 가진다는 점이다. Tensor Core 처리량은 2배로 증가했지만, shared memory 대역폭과 지수함수(exponential) 유닛은 거의 그대로이다. 이런 비대칭적 스케일링(asymmetric scaling) 때문에, FA3의 알고리즘을 Blackwell에 그대로 이식하면 성능이 크게 제한된다.

저자는 이 비대칭 하드웨어 스케일링을 정면으로 다루는 FlashAttention-4를 제안한다. 핵심 기법은 다음과 같다.

파이프라인 재설계: Blackwell의 완전 비동기 MMA와 더 큰 타일 크기를 활용하는 새로운 소프트웨어 파이프라인
지수함수 병목 완화: 다항식 근사를 통한 소프트웨어 에뮬레이션 + 불필요한 softmax rescaling을 건너뛰는 조건부 rescaling
Shared memory 트래픽 감소: Tensor Memory 활용과 backward pass에서의 2-CTA MMA 모드

FlashAttention-4는 B200 GPU에서 BF16 기준 cuDNN 9.13 대비 1.3배, Triton 대비 2.7배 빠르며, 최대 1613 TFLOPS/s (71% utilization)를 달성한다.

또한 FA4는 CuTe-DSL(Python 기반)로 전체 구현하여, 기존 C++ 템플릿 대비 컴파일 시간을 20-30배 단축했다.

Background: Blackwell GPU의 비대칭 스케일링

Hopper vs Blackwell 하드웨어 비교

하드웨어	Hopper (H100)	Blackwell (B200)	스케일링
Tensor Core (BF16)	1 PFLOPS	2.25 PFLOPS	2.25×
MMA 타일 크기	64 × N	128 × N (또는 256)	2×
MUFU (exp 등)	16 ops/clock/SM	16 ops/clock/SM	1× (동일!)
SMEM 대역폭	128 bytes/clock	128 bytes/clock	1× (동일!)
Tensor Memory	없음	256 KB/SM	신규
MMA 비동기성	register writeback 필요	TMEM에 직접 쓰기	더 높은 비동기성

핵심 관찰: Tensor Core가 2배 이상 빨라졌지만, exp를 계산하는 MUFU와 SMEM 대역폭은 그대로이다. FA3에서는 ping-pong 스케줄링으로 softmax를 GEMM 뒤에 숨길 수 있었는데, Blackwell에서는 GEMM이 2배 빨라져서 softmax를 숨길 시간이 부족하다.

Roofline 분석: Forward Pass

타일 크기 $M \times N$과 head dimension $d$에 대해, forward pass의 각 리소스별 소요 사이클:

\[T_{\text{MMA}} = \frac{4MNd}{8192} \text{ cycles}\] \[T_{\text{smem}} = \frac{3MNd}{8192} \text{ cycles (대략적)}\] \[T_{\text{exp}} = \frac{MN}{16} \text{ cycles}\]

Resource	$128^3$	$256 \times 128^2$
MMA compute	1024	2048
Shared memory	768	1536
Exponential unit	1024	2048

MMA와 exponential이 동시에 병목이다. 즉, exp를 다른 하드웨어에서 실행하지 않으면 Tensor Core가 아무리 빨라도 성능이 제한된다.

이 분석에서 FA4의 설계 원칙이 도출된다.

큰 타일 크기를 사용하여 MMA와 softmax의 overlap을 극대화
exp의 처리량을 높이기 위해 FMA 유닛으로 소프트웨어 에뮬레이션
불필요한 non-matmul 연산을 줄이기 — 조건부 rescaling

Forward Pass: 파이프라인 재설계

FA3 vs FA4 파이프라인 비교

FA3는 2개 warpgroup의 ping-pong 스케줄링을 사용했다. FA4는 이를 확장하여 역할별 전문화된 warp 그룹을 사용한다.

Blackwell에서의 핵심 변화

1. Accumulator가 Tensor Memory에 저장

Hopper에서는 MMA의 accumulator가 register에 저장되어, softmax warpgroup이 register에서 값을 읽어야 했다. Blackwell에서는 MMA가 Tensor Memory(TMEM)에 직접 accumulator를 쓴다. 이 덕분에:

Softmax warpgroup이 TMEM에서 직접 값을 읽을 수 있다
Rescaling을 별도의 correction warpgroup으로 분리 가능 — critical path에서 제거

2. 타일 크기 128 × 128

Hopper의 64 × 128 대비 2배 큰 타일. 한 번의 MMA에 더 많은 연산을 수행하므로, MMA와 softmax의 overlap 기회가 늘어난다.

3. Warp 역할 분배

각 thread당 하나의 행(row)을 담당하여 128개 원소를 register에 로드한다. 구체적으로:

역할	수량	기능
MMA warpgroup	1	Tensor Core 연산 ($QK^\top$, $PV$)
Softmax warpgroup	2	max, exp, rowsum 계산
Correction warpgroup	1	Rescaling ($e^{m_{\text{old}} - m_{\text{new}}}$로 보정)
TMA (producer)	-	HBM → SMEM 데이터 로드

Softmax warpgroup과 correction warpgroup을 분리한 것이 FA3와의 핵심 차이다. Correction은 critical path 밖에서 실행된다.

지수함수 소프트웨어 에뮬레이션

문제: MUFU 병목

Blackwell에서 MUFU(Multi-Function Unit)는 clock당 SM당 16개 연산만 처리한다. Tensor Core가 clock당 8192 FLOPs를 처리하는 것과 비교하면 512배 느리다. Head dimension 128 기준으로 forward에서 matmul FLOPs는 exp 연산 대비 512배 많지만, MUFU가 512배 느리므로 exp가 matmul과 동일한 시간을 소비한다.

해결: 다항식 근사

FMA(Fused Multiply-Add) 유닛은 MUFU와 독립적으로 병렬 실행될 수 있다. 저자는 지수함수를 FMA 기반 다항식으로 근사한다.

핵심 분해:

\[2^x = 2^{\lfloor x \rfloor} \cdot 2^{x_{\text{frac}}}\]

여기서 $x_{\text{frac}} = x - \lfloor x \rfloor \in [0, 1)$이다.

정수 부분 $2^{\lfloor x \rfloor}$: IEEE 754 부동소수점의 exponent 필드를 직접 조작 (integer ALU 명령어)
소수 부분 $2^{x_{\text{frac}}}$: 다항식 근사

\[2^{x_{\text{frac}}} \approx \sum_{i=0}^{n} p_i \cdot x_{\text{frac}}^i\]

Degree 3 다항식의 경우 3번의 FMA 명령어로 계산 가능하며, BF16 정밀도에서 하드웨어 MUFU와 거의 구분 불가능한 오차를 보인다.

Method	FP32 Max Rel Err	BF16 Max Rel Err
Hardware MUFU.EX2	$1.41 \times 10^{-7}$	$3.89 \times 10^{-3}$
Degree 3 polynomial	$8.77 \times 10^{-5}$	$3.90 \times 10^{-3}$
Degree 5 polynomial	$1.44 \times 10^{-7}$	$3.89 \times 10^{-3}$

FP32 수준에서는 degree 3가 MUFU보다 약 600배 부정확하지만, BF16으로 반올림하면 양자화 오차가 지배적이어서 차이가 사라진다. Degree 3 이상에서 BF16 오차는 모두 $\sim 3.9 \times 10^{-3}$으로 동일하다.

Partial Emulation

모든 exp를 다항식으로 대체하면 register 압력이 증가하고 대역폭이 늘어난다. 따라서 softmax row의 일부(10-25%)에만 선택적으로 적용한다. 나머지는 하드웨어 MUFU를 사용한다. 적용 비율은 MMA와 exp의 처리량 비율에 따라 경험적으로 튜닝한다.

조건부 Softmax Rescaling

기존 방식의 비효율

FlashAttention의 online softmax에서, 새로운 블록을 처리할 때마다 max가 바뀌면 이전 결과를 보정해야 한다:

\[m_j = \max(m_{j-1}, \text{rowmax}(S_j))\] \[O_j = e^{m_{j-1} - m_j} O_{j-1} + e^{S_j - m_j} V_j\]

이 rescaling ($e^{m_{j-1} - m_j} O_{j-1}$) 은 매 블록마다 발생한다. 하지만 실제로 max가 크게 변하지 않으면 $e^{m_{j-1} - m_j} \approx 1$이 되어 불필요한 연산이다.

FA4의 개선: threshold $\tau$

저자는 threshold $\tau$를 도입하여, max의 변화가 충분히 클 때만 rescaling한다:

\[O_j = \begin{cases} e^{m_{j-1} - m_j} O_{j-1} + e^{S_j - m_j} V_j & \text{if } m_j - m_{j-1} > \tau \\ O_{j-1} + e^{S_j - m_{j-1}} V_j & \text{otherwise} \end{cases}\]

$\tau$가 $m_{j-1}$와 $m_j$의 차이보다 크면, 이전 max($m_{j-1}$)를 그대로 사용한다. 마지막에 true max $m_{\text{final}}$과 true normalizer $\ell_{\text{final}}$로 한 번만 보정한다.

\[\text{Output} = \frac{1}{\ell_{\text{final}}} O_{\text{final}}\]

실용적으로 $\tau = \log_2(256) = 8.0$으로 설정한다. 이렇게 하면 rescaling 횟수가 약 10배 감소한다.

Backward Pass

Roofline 분석

Backward pass는 5개의 MMA를 수행한다: $S^\top = KQ^\top$, $dP^\top = VdO^\top$, $dV = P^\top dO$, $dK = dS^\top Q$, $dQ = dS \cdot K$.

$M = N = d = 128$ 기준:

Resource	1-CTA ($M=128$)	2-CTA ($M=256$)
MMA compute	2560	2560
Total shared memory	3328	2688
Exponential unit	1024	1024

Shared memory 트래픽이 MMA compute보다 30% 더 크다. Forward와 달리 backward에서는 shared memory가 주요 병목이다.

2-CTA MMA 모드

Blackwell은 2-CTA tensor core MMA 모드를 지원한다. 같은 thread block cluster 내의 2개 CTA가 협력하여 하나의 큰 MMA를 실행한다. $M = 256$, $N = K = 128$ 크기의 타일을 사용하면:

각 CTA는 operand B의 절반만 shared memory에 로드
나머지 절반은 peer CTA의 shared memory에서 읽음
→ Shared memory 트래픽 약 절반으로 감소

또한 $dQ$의 accumulation에서 atomic add가 필요한데, 2-CTA 모드에서는 각 CTA가 $dQ$ 타일의 절반만 쓰므로 global atomic reduction 횟수도 절반이 된다.

Backward Computation Graph

FA4의 backward pass는 Prologue → Main Loop → Tail의 3단계로 구성된다. 5개 MMA + 2개 elementwise 연산이 파이프라인으로 실행된다. FA3 대비 핵심 개선:

TMEM에 accumulator를 저장하여 MMA와 softmax gradient 계산을 overlap
$dQ$와 $dK$의 MMA를 이전 iteration의 softmax 계산과 병렬 실행
2-CTA 모드로 shared memory 병목 완화

Deterministic Backward Pass

GPU의 atomic reduction은 비결정적(nondeterministic)이다. 강화학습 등 재현 가능한 학습이 필요한 경우를 위해, deterministic mode도 제공한다. Semaphore lock으로 $dQ$ reduction 순서를 고정하며, CTA swizzling으로 stall을 최소화한다. Nondeterministic 대비 약 75%의 속도를 달성한다.

Scheduling: LPT와 Causal Masking

Longest-Processing-Time First (LPT)

Causal masking이나 variable sequence length 상황에서 SM 간 load imbalance가 발생한다. FA4는 LPT 스케줄링을 도입한다.

SM들을 L2 cache를 공유하는 section으로 나누고
각 section 내에서 worktile을 실행 시간이 긴 순서대로 배치
Causal masking에서는 대각선 근처의 긴 worktile을 먼저 처리

이 스케줄링은 Hopper에서도 적용 가능하며, BF16 hdim 128 기준 MHA에서 4-8% FLOPS 향상, MQA 8에서 7-14% 향상을 보인다.

Language: CuTe-DSL

FA4는 CUDA C++이 아닌 CuTe-DSL(Python 기반)로 전체 구현했다. CuTe-DSL은 CUTLASS의 일부로, Python 코드를 PTX → SASS로 컴파일한다.

	FA3 (C++)	FA4 (CuTe-DSL)
Forward 컴파일	55s	2.5s
Backward 컴파일	45s	1.4s
Speedup	—	22-32×

C++ 템플릿 메타프로그래밍의 복잡한 컴파일 과정 없이, Python의 JIT 컴파일을 활용하여 빠른 iteration이 가능하다. FA2, FA3는 수백 개의 커널을 미리 컴파일해야 했지만, FA4는 필요할 때 JIT로 컴파일한다.

이 접근의 장점:

개발 생산성: C++ 템플릿 전문 지식 없이도 GPU 커널 개발 가능
모듈성: Block-sparse, FlexAttention, variable sequence length 등을 독립적인 primitive로 구현하여 자유롭게 조합
PTX escape hatch: CuTe-DSL API에 아직 없는 기능은 직접 PTX를 삽입 가능

Empirical Evaluation

B200 GPU에서 BF16 입력으로 벤치마크를 수행했다. Hidden dimension 2048, head dimension 64 또는 128, 시퀀스 길이 1K-32K.

Forward Pass

Head dim 128 기준:

Non-causal: FA4가 cuDNN 9.13.0 대비 1.1-1.3× 빠르고, Triton 대비 2.1-2.7× 빠르다
Causal: LPT 스케줄링 덕분에 특히 긴 시퀀스에서 이점이 크다
최대 1613 TFLOPS/s (이론 최대 2250 TFLOPS의 약 71%)

DeepSeek V3에서 사용하는 head dim (192, 128) 설정에서도 cuDNN 대비 일관되게 우수한 성능을 보인다.

Backward Pass

Backward에서도 cuDNN 대비 일관된 speedup을 달성한다. 2-CTA 모드가 shared memory 병목을 완화하여 특히 긴 시퀀스에서 효과적이다.

Deterministic Backward Ablation

Deterministic backward의 스케줄링 전략 비교:

SPT (Shortest-Processing-Time first): Causal에서 최적
Reverse mblock LPT: 차선
Naive: 스케줄링 없이는 성능이 크게 떨어짐

CTA swizzling과 LPT 스케줄링 조합이 deterministic 모드에서도 nondeterministic 대비 75%의 성능을 유지하게 한다.

FlashAttention 시리즈 비교

	FA1	FA2	FA3	FA4
GPU	A100	A100	H100	B200
핵심 아이디어	Tiling + Recomputation	non-matmul 감소, split-Q	Ping-pong, FP8	SW exp, conditional rescale, 2-CTA
병목	HBM IO	non-matmul FLOPs	GEMM vs softmax	exp + SMEM
MMA 명령	mma.sync	mma.sync	WGMMA	tcgen05.mma
Accumulator	Register	Register	Register	Tensor Memory
BF16 성능	—	~230 TFLOPS	~740 TFLOPS	~1613 TFLOPS
GPU 활용률	—	50-73%	75%	71%
구현 언어	CUDA C++	CUDA C++	CUDA C++	CuTe-DSL (Python)
컴파일 시간	—	~55s	~55s	~2.5s

Discussion and Conclusion

FlashAttention-4는 비대칭 하드웨어 스케일링이라는 현대 가속기의 근본적 추세를 정면으로 다룬 논문이다. Tensor Core가 다른 유닛보다 훨씬 빠르게 발전하면서, 병목이 matmul에서 shared memory 트래픽과 지수함수 처리량으로 이동했다. FA4는 이를 세 가지 방향에서 해결한다.

연산 분산: exp를 MUFU에서만 하지 않고 FMA 유닛에도 분산하여 처리량 증가
불필요한 연산 제거: 조건부 rescaling으로 보정 횟수 10배 감소
메모리 효율: TMEM과 2-CTA 모드로 shared memory 병목 완화

또한 CuTe-DSL로의 전환은 단순한 구현 선택이 아니라, attention 커널 개발의 접근성을 크게 높인 결정이다. FlexAttention, block-sparse attention 등 다양한 attention variant를 FA4 프레임워크 위에 빠르게 구현할 수 있다.

FA1이 “IO를 줄이자”, FA2가 “non-matmul을 줄이자”, FA3가 “GEMM과 softmax를 겹치자”였다면, FA4는 “하드웨어의 비대칭을 소프트웨어로 보상하자”라는 메시지를 던진다. 가속기가 더 빠르게, 더 비대칭적으로 발전할수록, 이런 하드웨어-소프트웨어 co-design의 중요성은 더 커질 것이다.

FlashAttention의 원리가 궁금하다면 FlashAttention 논문 리뷰를, 개선점이 궁금하다면 FlashAttention-2 논문 리뷰를, Hopper GPU 최적화가 궁금하다면 FlashAttention-3 논문 리뷰를, Triton으로 직접 구현하고 싶다면 Triton 05: Flash Attention을 참고하자.

참고 문헌

FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision

Thu, 09 Apr 2026 15:00:00 +0000

FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision

Introduction

Transformer의 핵심인 attention은 시퀀스 길이에 대해 $O(N^2)$의 연산량을 가지고, LLM과 long-context 애플리케이션에서 가장 큰 병목이다. FlashAttention은 tiling과 kernel fusion으로 HBM IO를 줄여서 이를 해결했고, FlashAttention-2는 non-matmul FLOPs 감소와 warp partitioning 개선으로 한 단계 더 발전했다.

하지만 FlashAttention-2는 H100 GPU에서 이론 성능의 35%밖에 활용하지 못한다. 최적화된 GEMM 커널이 80~90%를 달성하는 것과 비교하면 매우 낮은 수치다. 그 이유는 FlashAttention-2가 동기적(synchronous) 모델을 사용하고, Hopper 아키텍처의 새로운 하드웨어 기능을 활용하지 않기 때문이다.

근본적으로, FlashAttention-2의 알고리즘은 단순화된 동기 모델을 따르며 비동기성과 저정밀도를 명시적으로 활용하지 않는다. 비동기성은 하드웨어 전문화의 결과다. 행렬곱을 수행하는 Tensor Core, 메모리 로드를 담당하는 TMA 등 서로 독립적인 하드웨어 유닛이 존재하기 때문에, 이들을 동시에 활용하려면 소프트웨어도 이에 맞춰 설계해야 한다.

저자는 이 문제를 해결하기 위해 세 가지 기법을 제안한다.

Producer-Consumer 비동기 처리: Warp specialization으로 데이터 전송과 연산을 분리하고, ping-pong 스케줄링으로 GEMM과 softmax를 겹쳐서 실행한다.
Intra-warpgroup 파이프라이닝: 하나의 warpgroup 안에서도 WGMMA의 비동기 특성을 이용하여 GEMM과 softmax를 겹친다.
FP8 저정밀도 + Incoherent Processing: 블록 양자화와 Hadamard 변환으로 FP8의 정확도 손실을 최소화한다.

결과적으로 FlashAttention-3는 H100에서 FP16 기준 740 TFLOPS/s (75% 활용률), FP8 기준 1.2 PFLOPS/s를 달성하며, FlashAttention-2 대비 1.5~2.0배 빠르다.

Background

Multi-Head Attention

Query, Key, Value $Q, K, V \in \mathbb{R}^{N \times d}$에 대해 attention output $O \in \mathbb{R}^{N \times d}$는 다음과 같이 계산된다.

\[S = \alpha QK^\top \in \mathbb{R}^{N \times N}, \quad P = \text{softmax}(S) \in \mathbb{R}^{N \times N}, \quad O = PV \in \mathbb{R}^{N \times d}\]

여기서 $\alpha = 1/\sqrt{d}$이고, softmax는 row-wise로 적용된다. 실제로는 수치 안정성을 위해 $S$에서 $\text{rowmax}(S)$를 빼고 exponential을 취한다.

Multi-head attention(MHA)에서는 각 head가 자기만의 $Q, K, V$를 가지며, 여러 head와 batch에 대해 독립적으로 병렬 처리된다.

Backward Pass

손실 함수를 $\phi$, 그 gradient를 $d(\cdot) = \partial \phi / \partial (\cdot)$라고 하자. Output gradient $dO \in \mathbb{R}^{N \times d}$가 주어지면 chain rule에 따라 $dQ, dK, dV$를 다음과 같이 계산한다.

\[dV = P^\top dO \in \mathbb{R}^{N \times d}\] \[dP = dOV^\top \in \mathbb{R}^{N \times N}\] \[dS = \text{dsoftmax}(dP) \in \mathbb{R}^{N \times N}\] \[dQ = \alpha \cdot dS \cdot K \in \mathbb{R}^{N \times d}\] \[dK = \alpha \cdot dS^\top Q \in \mathbb{R}^{N \times d}\]

여기서 $ds = (\text{diag}(p) - pp^\top)dp$, 즉 $p = \text{softmax}(s)$에 대해 row-wise로 적용된다. 이를 $\text{dsoftmax}(dP)$라고 쓴다. Forward pass에서 2개의 matmul($QK^\top, PV$)이 필요했다면, backward pass에서는 5개의 matmul이 필요하다. 이 때문에 backward의 FLOPs는 forward의 2.5배이다.

GPU 하드웨어 특성과 실행 모델

FlashAttention-3를 이해하려면 Hopper GPU의 메모리 계층과 실행 모델을 알아야 한다.

메모리 계층

H100 SXM5 GPU의 메모리 계층은 다음과 같다.

하드웨어 레벨	병렬 단위	메모리	용량 & 대역폭
Chip	Grid	GMEM (HBM)	80 GiB @ 3.35 TB/s
GPC	Threadblock Cluster	L2	50 MiB @ 12 TB/s
SM	Threadblock (CTA)	SMEM	228 KiB per SM, 31TB/s per GPU
Thread	Thread	RMEM (Register)	256 KiB per SM

아래로 갈수록 빠르고 작다. FlashAttention의 핵심 아이디어는 GMEM(HBM) 접근을 최소화하고 SMEM과 RMEM에서 최대한 연산을 수행하는 것이었다.

쓰레드 계층

GPU의 실행 단위도 계층적이다.

Thread: 가장 작은 실행 단위
Warp: 32개 thread 묶음, SIMT(Single Instruction Multiple Thread) 방식으로 동시 실행
Warpgroup: 4개 warp (128 threads), Hopper에서 새로 도입된 단위. WGMMA의 실행 단위이다.
Threadblock (CTA): 같은 SM에서 실행되는 warpgroup들의 묶음. SMEM을 공유한다.
Threadblock Cluster: 같은 GPC에서 실행되는 threadblock 묶음. L2를 공유한다.
Grid: 전체 커널의 모든 threadblock

비동기 실행과 Warp Specialization

Hopper에는 두 가지 핵심 하드웨어 기능이 있다.

TMA (Tensor Memory Accelerator): HBM과 SMEM 사이의 데이터 전송을 전담하는 하드웨어 유닛이다. 기존에는 warp 내의 모든 thread가 인덱스 계산, 범위 검사, 메모리 복사를 직접 수행해야 했다. TMA는 이 모든 것을 하드웨어가 자동으로 처리한다. 덕분에 warp은 데이터 전송에서 완전히 해방되어 연산에만 집중할 수 있다.

WGMMA (Warpgroup Matrix Multiply-Accumulate): 기존 Ampere의 mma.sync를 대체하는 명령어다. mma.sync는 이름 그대로 동기적이어서, 명령을 발행하면 결과가 나올 때까지 해당 warp이 대기해야 했다. 반면 WGMMA는 비동기적이다. WGMMA 명령을 발행한 후 결과를 기다리지 않고 다른 연산(softmax의 exp, max 등)을 수행할 수 있다. 또한 WGMMA는 SMEM에서 직접 operand를 읽을 수 있어서, 레지스터로 복사하는 단계를 생략할 수 있다.

이 두 기능 덕분에 Hopper에서는 warp을 역할별로 나눌 수 있다.

Producer warp: TMA로 HBM → SMEM 데이터 로드만 담당. 레지스터를 거의 사용하지 않으므로 setmaxnreg로 레지스터를 반환한다.
Consumer warp: WGMMA로 행렬 연산만 담당. Producer가 반환한 레지스터를 가져와 더 큰 타일로 연산할 수 있다.

이 구조를 warp specialization이라고 하며, 데이터 전송과 연산이 동시에 일어나게 된다. Producer가 다음 블록의 $K, V$를 로드하는 동안 consumer는 현재 블록으로 행렬곱을 수행한다.

저정밀도 연산

Hopper의 WGMMA는 FP8(e4m3)을 지원한다. FP8 Tensor Core의 처리량은 FP16 대비 2배이다 (989 TFLOPS → 약 1978 TFLOPS). 하지만 FP8은 mantissa 3bit, exponent 4bit으로 정밀도가 낮아서 단순히 적용하면 정확도가 크게 떨어진다.

또한 FP8 WGMMA에는 레이아웃 제약이 있다. FP16 WGMMA는 mn-major와 k-major 입력 모두 지원하지만, FP8 WGMMA는 k-major만 지원한다. 이 제약은 attention처럼 연속된 두 GEMM($S = QK^\top$, $O = PV$)을 수행할 때 FP32 accumulator와 FP8 operand의 레이아웃이 충돌하는 문제를 야기한다.

Standard Attention과 FlashAttention 복습

Standard attention은 중간 행렬 $S$와 $P$를 HBM에 저장한다. 이로 인해 $O(N^2)$의 메모리가 필요하고, HBM 읽기/쓰기가 병목이 된다.

FlashAttention의 핵심은 $S$와 $P$를 HBM에 쓰지 않는 것이다. Tiling을 통해 $Q, K, V$를 블록 단위로 SMEM에 올리고, on-chip에서 local softmax를 계산한 뒤 결과를 점진적으로 누적한다. 이때 online softmax 알고리즘을 사용하여 전체 row를 보지 않고도 softmax를 정확히 계산한다. Backward에서는 $(m, l)$ 통계량만 저장해두고 $S, P$를 recomputation한다.

FlashAttention-3도 이 기본 구조는 동일하다. 차이점은 어떻게 실행하느냐에 있다.

FlashAttention-3: Algorithm

기본 구조: Warp Specialization 적용

FlashAttention-3의 forward pass는 batch, head, query sequence length에 대해 embarrassingly parallel하다. 각 CTA는 query의 한 타일 $Q_i \in \mathbb{R}^{B_r \times d}$를 담당하여 output 타일 $O_i$를 계산한다. $Q$를 $T_r = \lceil N/B_r \rceil$개의 블록으로, $K, V$를 $T_c = \lceil N/B_c \rceil$개의 블록으로 나눈다.

Producer warp과 consumer warp이 $s$-stage circular SMEM buffer를 통해 협력한다. 이 버퍼는 $K, V$ 블록을 $s$개까지 미리 로드해둘 수 있는 원형 큐이다.

Algorithm 1: Forward Pass (Warp Specialization, without intra-consumer overlapping)

Producer warpgroup:

setmaxnreg로 레지스터를 반환한다 (producer는 레지스터가 거의 필요 없다).
$Q_i$를 HBM에서 SMEM으로 로드하고, consumer에게 알린다.
$j = 0, \ldots, T_c - 1$에 대해:
- 버퍼의 $(j \bmod s)$번째 스테이지가 소비될 때까지 대기한다.
- $K_j, V_j$를 HBM에서 SMEM의 해당 스테이지에 로드한다.
- 로드 완료를 consumer에게 알린다.

TMA 덕분에 로드 명령은 비동기적이며, 처음 $s$번의 이터레이션까지는 대기 없이 연속으로 발행할 수 있다.

Consumer warpgroup:

setmaxnreg로 추가 레지스터를 확보한다.
On-chip에서 초기화: $O_i = 0 \in \mathbb{R}^{B_r \times d}$, $l_i = 0 \in \mathbb{R}^{B_r}$, $m_i = -\infty \in \mathbb{R}^{B_r}$
$Q_i$가 SMEM에 로드되기를 기다린다.
$j = 0, \ldots, T_c - 1$에 대해:
- $K_j$가 SMEM에 로드되기를 기다린다.
- SS-GEMM: $S_i^{(j)} = Q_i K_j^\top \in \mathbb{R}^{B_r \times B_c}$. Commit and wait.
- $m_i^{\text{old}} = m_i$를 저장한다.
- \[m_i = \max(m_i^{\text{old}}, \text{rowmax}(S_i^{(j)}))\]
- \[\tilde{P}_i^{(j)} = \exp(S_i^{(j)} - m_i) \in \mathbb{R}^{B_r \times B_c}\]
- \[l_i = \exp(m_i^{\text{old}} - m_i) \cdot l_i + \text{rowsum}(\tilde{P}_i^{(j)})\]
- $V_j$가 SMEM에 로드되기를 기다린다.
- 이전 결과 보정: $O_i = \text{diag}(\exp(m_i^{\text{old}} - m_i))^{-1} \cdot O_i$
- RS-GEMM: $O_i = O_i + \tilde{P}_i^{(j)} V_j$. Commit and wait.
- 버퍼의 $(j \bmod s)$번째 스테이지를 producer에게 반환한다.
최종 정규화: $O_i = \text{diag}(l_i)^{-1} O_i$
$L_i = m_i + \log(l_i)$를 계산한다 (backward용 logsumexp 저장).
$O_i$와 $L_i$를 HBM에 기록한다.

여기서 SS-GEMM은 두 operand 모두 SMEM에서 오는 GEMM이고, RS-GEMM은 한쪽이 Register(accumulator)에서 오는 GEMM이다.

이 구조만으로도 FlashAttention-2의 350 TFLOPS에서 540~570 TFLOPS로 성능이 향상된다. Producer가 다음 $K_{j+1}, V_{j+1}$을 로드하는 동안 consumer가 현재 $K_j, V_j$로 연산하기 때문이다.

Ping-Pong 스케줄링: GEMM과 Softmax 겹치기

여기까지만 해도 빨라졌지만, 아직 해결되지 않은 문제가 있다. Consumer warpgroup 안에서 GEMM과 softmax가 순차적으로 실행된다는 점이다.

왜 이게 문제일까? H100의 처리량을 보자.

연산	처리량 (FP16)	비고
행렬곱 (WGMMA)	989 TFLOPS	Tensor Core
지수함수 (exp 등)	3.9 TFLOPS	Multi-function Unit

행렬곱이 지수함수보다 약 256배 빠르다. 그런데 attention forward pass에서 matmul FLOPs와 exponential FLOPs의 비율을 계산해보면, head dimension $d = 128$ 기준:

matmul FLOPs: $4 \times N^2 \times d$ ($QK^\top$와 $PV$, 각각 $2N^2d$)
exponential FLOPs: $\sim N^2$ (softmax의 exp 연산)
비율: $4d = 512$

matmul이 512배 더 많지만, exponential이 256배 느리니까, softmax가 전체 사이클에서 차지하는 비중은:

\[\frac{1}{1 + 512/256} = \frac{1}{3} \approx 33\%\]

즉, 아무리 GEMM을 빠르게 해도 softmax 때문에 전체 시간의 약 1/3은 Tensor Core가 놀게 된다. FP8에서는 더 심각하다. GEMM 처리량이 2배로 늘어나지만 exp 속도는 그대로이므로, softmax 비중이 더 커진다.

해결책은 softmax를 GEMM 뒤에 숨기는 것이다. WGMMA가 Tensor Core에서 실행되고, exp는 Multi-function Unit(SFU)에서 실행되므로, 둘은 서로 다른 하드웨어 유닛이다. 동시에 실행할 수 있다면 softmax의 비용이 사라진다.

저자는 2개의 consumer warpgroup을 사용하는 ping-pong 스케줄링을 제안한다. 동기화 배리어(bar.sync)를 사용하여 다음과 같이 강제한다.

Warpgroup 1의 GEMM1($PV$)과 다음 iteration의 GEMM0($QK^\top$)이 먼저 스케줄링된다.
이 GEMM들이 Warpgroup 2의 GEMM들 이전에 스케줄링되도록 배리어로 순서를 강제한다.
결과적으로 Warpgroup 2가 GEMM을 하는 동안, Warpgroup 1은 softmax를 실행한다.
그 다음에는 역할이 뒤바뀐다.

그림에서 같은 색은 같은 iteration을 의미한다. 핵심은 한 warpgroup이 softmax를 하는 시간이 다른 warpgroup의 GEMM 시간에 완전히 가려진다는 것이다.

이 기법으로 570 TFLOPS → 620~640 TFLOPS로 성능이 향상된다 (FP16, hdim=128, seqlen=8192 기준).

Intra-Warpgroup 2-Stage 파이프라이닝

Ping-pong이 warpgroup 사이의 겹침이었다면, 2-stage 파이프라이닝은 하나의 warpgroup 안에서의 겹침이다.

Algorithm 1의 consumer 내부 루프를 보면, 한 iteration 안에서 다음과 같은 의존성이 있다.

GEMM0 ($S = QK^\top$) → 결과 $S$가 나와야
Softmax ($\tilde{P} = \exp(S - m)$ 등) → 결과 $\tilde{P}$가 나와야
GEMM1 ($O += \tilde{P}V$) → 다음 iteration 시작

이 세 단계가 순차적으로 실행되기 때문에, WGMMA가 비동기적이더라도 wait가 필요하다.

하지만 iteration을 넘나들면 겹칠 수 있다. 핵심 관찰은 이것이다.

Iteration $j$의 GEMM0 결과($S_{\text{next}}$)는 iteration $j$의 softmax에만 필요하다. 한편 iteration $j-1$의 GEMM1($\tilde{P}_{\text{cur}} V_{j-1}$)은 iteration $j-1$의 softmax 결과에만 의존한다.

따라서 다음과 같이 겹칠 수 있다.

Algorithm 2: Consumer Warpgroup Forward Pass (2-Stage)

$O_i = 0, l_i = 0, m_i = -\infty$ 초기화
$Q_i$와 $K_0$가 로드되기를 기다린다.
$S_{\text{cur}} = Q_i K_0^\top$ (WGMMA). Commit and wait.
$S_{\text{cur}}$에서 $m_i, \tilde{P}_{\text{cur}}, l_i$ 계산하고 $O_i$ rescale.
$j = 1, \ldots, T_c - 1$에 대해:
- $K_j$가 로드되기를 기다린다.
- $S_{\text{next}} = Q_i K_j^\top$ (WGMMA). Commit but do not wait. ← 핵심!
- $V_{j-1}$이 로드되기를 기다린다.
- $O_i = O_i + \tilde{P}_{\text{cur}} V_{j-1}$ (WGMMA). Commit but do not wait.
- $S_{\text{next}}$의 WGMMA를 기다린다.
- $S_{\text{next}}$에서 $m_i, \tilde{P}_{\text{next}}, l_i$ 계산하고 $O_i$ rescale.
- $\tilde{P}_{\text{cur}} V_{j-1}$의 WGMMA를 기다린다.
- $O_i$ rescale 반영.
- \[S_{\text{next}} \to S_{\text{cur}}, \tilde{P}_{\text{next}} \to \tilde{P}_{\text{cur}}\]
마지막 $V_{T_c-1}$ 처리
최종 정규화: $O_i = \text{diag}(l_i)^{-1} O_i$

그림에서 같은 색은 같은 iteration을 나타낸다.

WGMMA0 ($QK^\top$)이 iteration 1을 계산하는 동안
Softmax는 iteration 0의 $S$를 처리하고
WGMMA1 ($\tilde{P}V$)은 iteration 0의 softmax 결과를 사용한다.

5번 단계에서 $S_{\text{next}}$와 $\tilde{P}_{\text{cur}} V_{j-1}$을 동시에 발행하고 나중에 각각 기다리는 것이 핵심이다. 두 WGMMA가 동시에 실행되는 동안 softmax 연산이 끼어든다.

트레이드오프: 레지스터 압력

이 기법의 대가는 $S$를 두 개 동시에 보관해야 한다는 것이다 ($S_{\text{cur}}$와 $S_{\text{next}}$). 추가 레지스터 사용량은 $B_r \times B_c \times \text{sizeof(float)}$이다. 블록 크기를 키우면 memory IO는 줄지만 레지스터 압력이 커져서, GPU마다 최적 블록 크기를 조정해야 한다.

컴파일러 주의사항

이 의사코드는 이상적인 실행 순서를 나타내지만, NVCC 컴파일러가 최적화를 위해 WGMMA 명령의 순서를 재배치할 수 있다. 이 경우 의도한 파이프라이닝이 깨질 수 있으므로, SASS 코드를 확인하여 컴파일러가 올바른 순서를 생성하는지 검증해야 한다.

이 기법으로 620 TFLOPS → 640~661 TFLOPS까지 향상된다.

Backward Pass

Backward pass도 forward와 유사한 warp specialization 구조를 사용한다. 다만 한 가지 추가 역할이 필요하다. Forward에서는 각 CTA가 $Q_i$를 담당하여 $O_i$를 독립적으로 계산했지만, backward에서는 $dQ$의 누적이 필요하다. 여러 CTA가 같은 $dQ_i$에 값을 더해야 하므로, 메모리 경합(contention)이 발생한다.

이를 해결하기 위해 dQ-writer warp이라는 세 번째 역할을 추가한다.

Producer warp: $K_j, V_j, Q_i, dO_i$ 등을 HBM에서 SMEM으로 로드
Consumer warp: WGMMA로 $dV_j, dK_j, dQ_i^{(\text{local})}$ 계산 후 SMEM에 기록
dQ-writer warp: $dQ_i^{(\text{local})}$를 SMEM에서 읽어 HBM의 $dQ_i$에 semaphore를 이용하여 원자적으로 누적

이 구조 덕분에 dQ 누적의 메모리 경합이 consumer의 연산을 블로킹하지 않는다.

Low-precision FlashAttention: FP8

FP8 WGMMA의 레이아웃 문제

FP8로 FlashAttention-3를 구현할 때 가장 큰 기술적 난관은 레이아웃 충돌이다.

GEMM에서 행렬 $A \times B$를 계산할 때, $A$나 $B$가 mn-major (outer dimension이 연속)인지 k-major (inner dimension이 연속)인지에 따라 WGMMA 명령이 달라진다. FP16에서는 둘 다 지원하지만, FP8에서는 k-major만 지원한다.

Attention에서는 두 개의 연속된 GEMM이 있다.

GEMM0: $S = QK^\top$ → $S$는 FP32 accumulator에 저장됨
GEMM1: $O = \tilde{P}V$ → $\tilde{P}$는 $S$에서 softmax를 취한 결과

문제는 FP32 accumulator의 레지스터 레이아웃(아래 그림)이 FP8 operand A의 레이아웃(아래 그림)과 다르다는 것이다.

FP16에서는 이 레이아웃 차이를 mn-major 모드로 우회할 수 있었지만, FP8에서는 k-major만 지원하므로 불가능하다.

저자는 두 가지 방법으로 이를 해결한다.

V의 in-kernel transpose: LDSM(Load Shared Memory to Register)과 STSM(Store Register to Shared Memory) 명령어를 이용하여 $V$ 타일을 SMEM에서 읽어 transpose한 후 다시 SMEM에 쓴다. 이 과정은 producer warp에서 실행되며, 다음 $V$ 타일을 TMA로 로드하는 시간에 숨길 수 있다.
$\tilde{P}$의 byte permute: Accumulator의 레이아웃을 FP8 operand A 형식에 맞추기 위해 byte permute 명령어를 사용한다. 구체적으로 {d0 d1 d4 d5 d2 d3 d6 d7} 순서로 재배열한다.

이 두 변환 모두 다른 연산에 숨길 수 있어서 추가 비용이 거의 없다.

Block Quantization

전체 텐서에 하나의 스케일링 값을 사용하는 per-tensor scaling은 이상치 하나가 전체 텐서의 양자화 범위를 지배할 수 있다. 이를 완화하기 위해 블록 단위 양자화를 사용한다.

$Q, K, V$ 각각을 $B_r \times d$ 또는 $B_c \times d$ 크기의 블록으로 나누고, 블록별로 하나의 스케일링 값을 유지한다.

\[s_Q = \frac{\max(|Q_{\text{block}}|)}{448}, \quad Q_{\text{fp8}} = \text{round}\left(\frac{Q_{\text{block}}}{s_Q}\right)\]

여기서 448은 FP8 e4m3의 최대 표현값이다.

FlashAttention-3는 이미 블록 단위로 $Q, K, V$를 처리하므로, 각 블록의 $S$에 대해 $s_Q \cdot s_K$를 곱해주면 된다. 이 양자화는 rotary embedding 같은 memory-bound 연산에 fuse할 수 있어서 추가 slowdown이 없다 (rotary embedding 자체가 memory-bandwidth bounded이므로).

Incoherent Processing

블록 양자화만으로는 부족하다. LLM의 활성화값에는 이상치(outlier)가 존재하는데, 전체 원소의 0.1% 정도가 나머지보다 매우 큰 값을 가진다. 이 소수의 큰 값이 블록의 양자화 범위를 지배하면 나머지 99.9%의 값들이 좁은 범위에 몰려서 정밀도가 크게 떨어진다.

저자는 양자화 문헌에서 사용되는 incoherent processing 기법을 도입한다. 아이디어는 양자화 전에 $Q$와 $K$에 랜덤 직교행렬 $M$을 곱해서 이상치를 분산시키는 것이다.

$M$이 직교행렬이므로 $MM^\top = I$이고:

\[(QM)(KM)^\top = QMM^\top K^\top = QK^\top\]

따라서 attention 결과는 전혀 변하지 않지만, $QM$의 각 원소는 원래 $Q$의 원소들의 가중합이 된다. 이상치의 영향이 여러 원소로 분산되어, 각 원소의 크기가 균등해진다.

왜 분산될까?

직관적으로 이해하면, $Q$의 한 행 $q = [1, 1, 1, 100, 1, \ldots]$처럼 이상치가 있다고 하자. $M$을 곱하면 $qM$은 $q$의 모든 원소를 섞은 값이 된다. 100이라는 큰 값이 다른 원소들과 합쳐져서 각 원소가 대략 비슷한 크기가 된다.

실제 구현

실제로 $M$을 $N \times N$ 임의의 직교행렬로 사용하면 $O(d^2)$의 연산이 필요하다. 대신 $M = HD$로 구성한다.

$D$: $\pm 1$을 원소로 가지는 랜덤 대각행렬 ($O(d)$)
$H$: Hadamard 행렬 ($O(d \log d)$, Fast Walsh-Hadamard Transform)

$HD$의 곱은 $O(d \log d)$에 계산 가능하다. 또한 이 변환은 rotary embedding과 같은 memory-bound 연산에 fuse할 수 있어서 사실상 공짜다. Rotary embedding은 어차피 HBM에서 $Q, K$를 읽고 써야 하므로 memory-bandwidth bounded인데, Hadamard 변환을 추가해도 메모리 대역폭은 변하지 않기 때문이다.

수치 오차 검증

이상치가 있는 분포에서 수치 오차를 비교한다. $Q, K, V$의 원소를 다음과 같이 생성한다.

\[\mathcal{N}(0, 1) + \mathcal{N}(0, 100) \cdot \text{Bernoulli}(0.001)\]

즉, 대부분은 표준정규분포이지만 0.1%의 원소에 표준편차 10의 이상치가 추가된다. FP64 구현을 기준(ground truth)으로 RMSE를 측정한 결과:

방법	RMSE
Baseline FP16 (standard attention)	3.2e-4
FlashAttention-2 FP16	1.9e-4
FlashAttention-3 FP16	1.9e-4
Baseline FP8 (per-tensor scaling)	2.4e-2
FlashAttention-3 FP8 (block quant + incoherent)	9.1e-3
FlashAttention-3 FP8 (block quant만, no incoherent)	9.3e-3
FlashAttention-3 FP8 (no block quant)	2.4e-2

주목할 점:

FP16에서 FlashAttention-2/3는 standard attention보다 1.7배 더 정확하다. 중간 결과(softmax)를 FP32로 유지하기 때문이다.
FP8에서 block quantization + incoherent processing을 적용하면 baseline 대비 2.6배 오차가 줄어든다.
Block quantization만으로도 대부분의 개선이 이루어지고, incoherent processing이 추가적인 개선을 제공한다.

Empirical Validation

H100 80GB SXM5 GPU에서 벤치마크를 수행했다. Hidden dimension은 2048, 시퀀스 길이는 512~16K, head dimension은 64, 128, 256으로, total token 수가 16K가 되도록 batch size를 조절했다. Forward FLOPs는 다음과 같이 계산한다.

\[\text{FLOPs} = 4 \times \text{seqlen}^2 \times \text{head\_dim} \times \text{num\_heads}\]

Causal masking이 있으면 약 절반만 계산하므로 2로 나눈다. Backward FLOPs는 forward의 2.5배이다 (forward에 matmul 2개, backward에 5개).

FP16 Forward Pass

FlashAttention-3는 FlashAttention-2 대비 1.5~2.0배 빠르다.

Head dim 64: Non-causal, seqlen 16K 기준 — Standard 73, FA-2 332, cuDNN 412, FA-3 497 TFLOPS
Head dim 128: Non-causal, seqlen 8K 기준 — Standard 133, FA-2 370, cuDNN 610, FA-3 649 TFLOPS
Head dim 256: Non-causal, seqlen 8K 기준 — FA-2 581, cuDNN 581, FA-3 746 TFLOPS (최대)
Standard attention 대비 3~16배 빠르다
cuDNN(NVIDIA 자체 최적화 라이브러리)보다도 대부분의 설정에서 빠르거나 비슷한 성능을 보인다

Head dimension이 클수록 matmul 비중이 커져서 GPU 활용률이 높아진다. Head dim 256에서 최대 756 TFLOPS, 이론 최대 989 TFLOPS의 약 75%에 도달한다.

FP16 Backward Pass

Backward pass에서도 FlashAttention-2 대비 1.5~1.75배 빠르다. Forward보다 speedup이 약간 낮은데, backward에는 5개의 matmul과 더 복잡한 데이터 의존성이 있어서 파이프라이닝 효과가 상대적으로 줄어들기 때문이다.

FP8 Forward Pass

FP8에서는 head dim 256, non-causal, seqlen 16K 기준 최대 1,171 TFLOPS에 달하며, 1.2 PFLOPS/s에 근접한다. FP16 대비 약 1.5~2배의 추가 speedup이다.

다만 짧은 시퀀스와 causal masking 조합에서는 FP8 cuDNN이 더 빠른 경우도 있다. 이는 FP16 FlashAttention-3가 persistent kernel과 load balancing 전략을 사용하는 반면 FP8 버전은 아직 이를 적용하지 않았기 때문이다.

Ablation: 파이프라이닝 효과

Non-causal FlashAttention-3 (FP16, batch=4, seqlen=8448, nheads=16, hdim=128) 기준:

Configuration	Time	TFLOPs/s
FlashAttention-3 (full)	3.538 ms	661
No GEMM-Softmax pipelining, with warp specialization	4.021 ms	582
GEMM-Softmax pipelining, no warp specialization	4.105 ms	570

Warp specialization만으로 570 → 582 TFLOPS, GEMM-softmax pipelining을 추가하면 582 → 661 TFLOPS로, 각 기법이 단계적으로 성능을 끌어올리는 것을 확인할 수 있다.

FlashAttention 시리즈 비교

	FlashAttention	FlashAttention-2	FlashAttention-3
핵심 아이디어	Tiling + Recomputation	non-matmul FLOPs 감소, warp partitioning	비동기 실행, FP8
타겟 GPU	A100	A100	H100 (Hopper)
주요 명령어	mma.sync	mma.sync	WGMMA + TMA
정밀도	FP16	FP16	FP16 + FP8
GPU 활용률	—	35% (H100 기준)	75% (H100)
FP16 성능	—	~370 TFLOPS	~740 TFLOPS
FP8 성능	—	—	~1.2 PFLOPS
Softmax 처리	순차 실행	순차 실행	GEMM과 겹침
Backward 특이점	—	outer loop을 K, V로 변경	dQ-writer warp 추가

Discussion, Limitations, Conclusion

FlashAttention-3는 Hopper GPU의 하드웨어 특성을 적극적으로 활용하여 attention 성능을 크게 향상시켰다. 특히 WGMMA의 비동기 특성을 이용한 ping-pong 스케줄링과, FP8의 정확도 문제를 해결하는 incoherent processing이 인상적이다.

FlashAttention-2가 “어떤 GPU에서든 작동하는 범용 최적화”였다면, FlashAttention-3는 “Hopper의 능력을 극한까지 끌어내는 하드웨어 특화 최적화”라고 할 수 있다. GPU 아키텍처가 발전할수록 소프트웨어도 이에 맞춰 진화해야 한다는 것을 보여주는 좋은 사례다.

Limitations

저자가 언급한 한계는 다음과 같다.

Hopper 전용: 현재 구현은 H100에 특화되어 있다. 다만 비동기 실행과 warp specialization이라는 개념 자체는 비슷한 하드웨어 특성을 가진 다른 GPU에도 적용 가능하다.
FP8 커널 설계의 복잡성: Persistent kernel과 FP8을 통합하는 것이 아직 남은 과제이다.
FP8 학습의 불확실성: 추론에서 FP8의 효과는 검증되었지만, 학습에서 저정밀도가 안정적인지는 추가 연구가 필요하다.

FlashAttention-3와 관련된 연구 방향들도 간략히 정리한다.

Distributed attention: Ring Attention 등은 FlashAttention을 여러 GPU로 확장하여 최대 100만 토큰까지 처리할 수 있다. FlashAttention-3의 개선은 이런 분산 attention 방법에도 그대로 적용된다.
Alternative architectures: Mamba, RWKV, RetNet 등 linear attention 계열 모델이 등장하고 있지만, 대형 모델(Jamba, Zamba 등)에서도 여전히 attention layer를 포함하고 있어 FlashAttention의 최적화가 유효하다.
KV cache quantization: QuIP, KIVI 등은 KV cache를 4bit, 2bit까지 양자화하여 추론 효율을 높인다. FlashAttention-3의 incoherent processing 기법은 이런 양자화 연구에서 영감을 받았다.

FlashAttention의 원리가 궁금하다면 FlashAttention 논문 리뷰를, 개선점이 궁금하다면 FlashAttention-2 논문 리뷰를, Blackwell GPU에서의 최적화가 궁금하다면 FlashAttention-4 논문 리뷰를, Triton으로 직접 구현하고 싶다면 Triton 05: Flash Attention을 참고하자.

참고 문헌

Triton 05: Flash Attention — 종합 프로젝트

Wed, 01 Apr 2026 15:00:00 +0000

개요

지금까지 배운 모든 기법을 종합하여 Flash Attention을 구현합니다. LLM 추론/학습에서 가장 중요한 최적화 기법 중 하나입니다.

Flash Attention의 원리와 논문 내용이 궁금하다면 FlashAttention 논문 리뷰를 먼저 읽어보는 것을 추천한다.

핵심 개념

Attention 수식

\[O = \text{softmax}\!\left(\frac{Q \cdot K^T}{\sqrt{d}}\right) \cdot V\]

$Q, K, V$: Query, Key, Value 행렬 (각각 $N \times d$)
$\sqrt{d}$: head dimension의 제곱근으로 나눠서 스케일링
$\text{softmax}$: 행(row) 단위로 적용 → 확률 분포로 변환

Standard Attention의 문제

시퀀스 길이 N=4096, float16이면:

S 행렬 크기: 4096 × 4096 × 2 bytes = 32MB
N=16384이면: 512MB — 시퀀스가 길어질수록 VRAM 폭발

Flash Attention의 핵심 아이디어

S 행렬을 전체 생성하지 않는다!

타일 단위로 Q, K, V를 처리하면서 결과를 점진적으로 누적합니다. 이를 위해 Online Softmax 알고리즘이 필요합니다.

Online Softmax

데이터를 청크(블록) 단위로 받으면서 점진적으로 업데이트합니다.

청크 1 처리 후 ($S_1$ = 첫 번째 K 블록과의 attention score):

\[m^{(1)} = \max(S_1)\] \[l^{(1)} = \sum_j e^{S_{1,j} - m^{(1)}}\] \[O^{(1)} = \text{diag}(l^{(1)})^{-1} \cdot e^{S_1 - m^{(1)}} \cdot V_1\]

청크 2 처리 후 — 보정 계수 (핵심!):

\[\alpha = e^{m^{(1)} - m^{(2)}}\]

이전 결과를 새로운 max 기준으로 보정:

\[l^{(2)} = l^{(1)} \cdot \alpha + \sum_j e^{S_{2,j} - m^{(2)}}\] \[O^{(2)} = O^{(1)} \cdot \alpha + e^{S_2 - m^{(2)}} \cdot V_2\]

왜 보정 계수 $\alpha$가 필요한가?

max가 바뀌면 이전에 계산한 exp 값들이 틀어집니다:

청크 1: max=5, exp(3-5) = exp(-2) = 0.135 청크 2: max=10, exp(3-5)는 틀림! exp(3-10) = exp(-7) = 0.0009여야 함 보정: 0.135 × exp(5-10) = 0.135 × exp(-5) ≈ 0.0009 ✓ α = exp(m_old - m_new) 

메모리 복잡도

방식	메모리	RTX 4080 (16GB)에서 최대 seq_len
Standard	O(N²)	~8K (float16)
Flash	O(N)	수십만+

커널 동작 원리

단계별 의사코드

Causal Masking

Autoregressive 모델(GPT 등)에서는 미래 토큰을 볼 수 없습니다:

코드 라인별 설명

Online Softmax 변수 초기화

m_i: 행별 최대값 추적 (처음엔 -inf → 점점 커짐)
l_i: 행별 softmax 분모 추적 (처음엔 0 → 점점 커짐)
acc: 최종 출력 누적기 (처음엔 0 → P@V 결과가 점점 누적)
이 세 변수가 Online Softmax의 핵심 — 전체 S 행렬 없이 softmax 계산

내부 루프 — Online Softmax 업데이트 (핵심!)

각 K/V 블록에 대해 다음을 수행합니다:

K 블록 로드 → S = Q @ K^T * scale 계산 (attention score 타일)
Causal mask 적용 → 미래 토큰 차단 (-inf로 마스킹)
Online Softmax 업데이트:
- m_new = max(m_old, max(S)) — 전체 최대값 갱신
- alpha = exp(m_old - m_new) — 이전 결과 보정 계수 (max가 바뀌면 이전 exp 값이 틀어지므로)
- l_i = l_i * alpha + sum(exp(S - m_new)) — 분모 업데이트
- acc = acc * alpha — 이전 출력 보정
V 블록 로드 → acc += P @ V 누적
p.to(v.dtype): FP32 → FP16 변환 (tl.dot은 같은 타입 필요)

매 반복마다 acc에 결과가 누적되므로 S 전체를 저장할 필요가 없습니다.

최종 정규화

l_i: 각 행의 softmax 분모 (Σ exp) → 마지막에 한 번만 나눔
FP32 → FP16 변환 후 저장

전체 튜토리얼과의 연결

개념	어디서 배웠나	Flash Attention에서의 역할
`tl.load`, mask	01 Vector Add	Q, K, V 블록 로드
reduction, `tl.exp`	02 Softmax	Online Softmax의 max, sum, exp
stride, 다중 포인터	03 RMSNorm	batch, head, seq, dim 차원 접근
`tl.dot`, 2D 타일링	04 MatMul	S = Q@K^T, O += P@V
K 차원 루프	04 MatMul	K/V 블록 순회 (내부 루프)
Online Softmax	신규	SRAM 제한 극복의 핵심

벤치마크 결과

정확도: PyTorch standard attention과 거의 동일한 결과
속도: 시퀀스 길이가 길수록 (1024+) 큰 속도 향상
메모리: O(N²) → O(N)으로 극적인 메모리 절약

전체 코드

Triton 04: Matrix Multiplication — 2D 타일링과 Autotune

Wed, 01 Apr 2026 15:00:00 +0000

개요

딥러닝의 핵심 연산인 행렬 곱셈(GEMM)을 Triton으로 구현합니다. 2D 타일링, tl.dot, triton.autotune 등 고급 기능을 학습합니다.

핵심 개념

행렬 곱셈이 왜 중요한가

딥러닝의 거의 모든 연산이 행렬 곱셈:

Linear layer: y = xW + b
Attention: QK^T, PV
MLP: 모든 Feed-Forward 블록

나이브 vs 타일링

나이브: 출력의 각 원소마다 Global Memory에서 행/열 전체를 읽음 → 같은 데이터를 반복 로드

타일링: 행렬을 작은 블록으로 나누어 SRAM에 올리고, 블록 단위로 계산

커널 동작 원리

2D 그리드

이전 튜토리얼은 1D 그리드(행 단위)였지만, MatMul은 2D 그리드를 사용합니다:

K 차원 루프

행렬 곱셈 C = A × B에서 A(M×K), B(K×N)일 때, K가 크면 한 번에 SRAM에 못 올립니다. 그래서 K를 BLOCK_SIZE_K씩 잘라서 반복하며, 부분 결과를 누적합니다.

L2 캐시 최적화 (Swizzling)

Swizzling = “같은 B 블록을 쓰는 프로그램들을 묶어서 실행”

`triton.autotune` 이란?

블록 크기에 따라 성능이 크게 달라집니다. Autotune은 여러 설정을 실행해보고 가장 빠른 것을 선택합니다:

코드 라인별 설명

K 차원 루프 (핵심)

이전 튜토리얼과의 차이점

	01~03	04 MatMul
그리드	1D (행 수)	1D (M타일 × N타일)
데이터	1D 벡터/행	2D 블록 (타일)
루프	없음	K 차원 루프
핵심 연산	`+`, `exp`, `sum`	`tl.dot` (텐서 코어)
파라미터 튜닝	수동 BLOCK_SIZE	`triton.autotune`

벤치마크 결과

cuBLAS(torch.matmul)는 수십 년간 최적화된 라이브러리입니다. Triton으로 cuBLAS의 80~90% 성능에 도달하는 것이 목표입니다.

전체 코드

Triton 03: RMSNorm — LLM에서 쓰이는 실전 커널

Wed, 01 Apr 2026 15:00:00 +0000

개요

LLaMA, Mistral, Gemma 등 최신 LLM에서 사용하는 RMSNorm을 Triton으로 구현합니다. Softmax와 유사한 패턴이지만, 학습 가능한 가중치(gamma)가 추가됩니다.

핵심 개념

LayerNorm vs RMSNorm

LayerNorm: y = (x - mean(x)) / sqrt(var(x) + ε) * γ + β RMSNorm: y = x / sqrt(mean(x²) + ε) * γ 

RMSNorm이 LLM에서 선호되는 이유:

mean 계산이 필요 없음 → 연산량 감소
bias(β) 없음 → 파라미터 수 감소
실험적으로 LayerNorm과 성능이 비슷

수식 분해

1. 제곱합: sum_sq = Σ(x_i²) 2. RMS: rms = sqrt(sum_sq / n + ε) 3. 정규화: x_norm = x / rms 4. 스케일링: y = x_norm * γ 

커널 동작 원리

코드 라인별 설명

PyTorch 참조 구현

커널 함수

래퍼 함수

02 Fused Softmax와의 차이점

	02 Softmax	03 RMSNorm
reduction	`max` + `sum` (2번)	`sum` (1번)
수치 안정성	max 빼기	eps 더하기
범위 밖 채움	`-inf`	`0.0`
추가 입력	없음	가중치 γ
입력 shape	2D만	3D/4D → 2D 변환

벤치마크 결과

PyTorch의 수동 RMSNorm 구현 대비 커널 퓨전으로 인한 성능 향상이 나타납니다. hidden_size가 클수록(2048, 4096 등) 차이가 명확합니다.

전체 코드

Triton 02: Fused Softmax — 커널 퓨전과 Reduction

Wed, 01 Apr 2026 15:00:00 +0000

개요

Softmax를 하나의 커널로 퓨전(fusion)하여 메모리 접근을 최소화합니다. 커널 퓨전이 왜 중요한지, reduction 연산을 어떻게 처리하는지 학습합니다.

핵심 개념

Softmax 수식

softmax(x_i) = exp(x_i - max(x)) / Σ exp(x_j - max(x))

max(x)를 빼는 이유: exp는 큰 값에서 오버플로우가 발생합니다. 최대값을 빼면 모든 지수가 0 이하가 되어 안정적으로 계산됩니다.

왜 커널 퓨전인가?

Reduction 연산

전체 데이터에서 하나의 값을 계산하는 연산:

max: 최대값
sum: 합계
mean: 평균

Triton에서는 tl.max(x, axis=0), tl.sum(x, axis=0) 으로 간단하게 수행합니다.

커널 동작 원리

입력 행렬의 각 행(row) 을 하나의 프로그램이 처리합니다.

코드 라인별 설명

커널 함수

핵심: max → exp → sum → 나누기를 전부 SRAM 안에서 처리. PyTorch는 이 4단계를 각각 별도 커널로 실행하므로 매번 Global Memory를 왕복합니다.

래퍼 함수

01 Vector Add와의 차이점

	01 Vector Add	02 Fused Softmax
처리 단위	1D 벡터의 청크	2D 행렬의 행
프로그램당 연산	덧셈 1번	max+exp+sum+나누기
퓨전 효과	없음 (연산이 1개)	4개 연산을 1커널로
새로운 기능	-	`tl.max`, `tl.sum`, `tl.exp`, stride

벤치마크 결과

커널 퓨전 덕분에 메모리 대역폭을 절약하여, 특히 열(column) 수가 클수록 PyTorch 대비 성능 향상이 눈에 띕니다.

전체 코드

Triton 01: Vector Addition — Triton 커널 기초

Wed, 01 Apr 2026 15:00:00 +0000

개요

가장 간단한 GPU 커널인 벡터 덧셈을 구현합니다. 이 튜토리얼에서 Triton의 핵심 개념을 모두 배울 수 있습니다.

핵심 개념

GPU 병렬 프로그래밍

CPU는 순차적으로 빠르게 처리하고, GPU는 수천 개의 코어로 동시에 처리합니다.

CUDA vs Triton

구분	CUDA	Triton
언어	C/C++	Python
메모리 관리	수동 (shared memory 직접 관리)	자동 (컴파일러가 처리)
스레드 관리	warp/thread 단위	block(프로그램) 단위
난이도	높음	낮음
성능	최고	CUDA의 90%+ 달성 가능

Triton 핵심 용어

커널(Kernel): GPU에서 실행되는 함수
프로그램(Program): 커널의 하나의 인스턴스 (CUDA의 thread block에 해당)
그리드(Grid): 프로그램 인스턴스의 총 개수
BLOCK_SIZE: 각 프로그램이 처리하는 데이터 크기

커널 동작 원리

길이 N인 벡터를 BLOCK_SIZE 크기의 청크로 나누고, 각 프로그램이 하나의 청크를 담당합니다.

단계별 분석

사용된 Triton 기능

기능	설명
`@triton.jit`	함수를 Triton 커널로 컴파일
`tl.program_id(axis)`	현재 프로그램의 ID (어떤 청크를 처리할지 결정)
`tl.arange(start, end)`	연속 정수 벡터 생성 (numpy의 arange와 유사)
`tl.load(ptr, mask)`	Global Memory에서 데이터 읽기
`tl.store(ptr, value, mask)`	Global Memory에 데이터 쓰기
`tl.constexpr`	컴파일 타임 상수 (BLOCK_SIZE처럼 컴파일 시 결정되는 값)

마스크(Mask)란?

벡터 길이가 BLOCK_SIZE의 배수가 아닐 때 경계 처리가 필요합니다.

그리드 설정

래퍼 함수

포인터(Pointer)란?

C/CUDA 경험이 없으면 포인터가 낯설 수 있습니다:

벤치마크 결과

Vector Add는 메모리 대역폭 바운드(memory-bound) 연산입니다. 연산량이 적고 데이터 이동이 대부분이라, Triton과 PyTorch의 성능 차이가 크지 않습니다. 하지만 이 패턴은 이후 모든 커널의 기초가 됩니다.

전체 코드

Triton 00: GPU 기초 — Triton을 시작하기 전에 알아야 할 것들

Wed, 01 Apr 2026 15:00:00 +0000

GPU란 무엇인가?

GPU(Graphics Processing Unit)는 원래 그래픽 렌더링을 위해 만들어졌지만, 지금은 딥러닝과 과학 계산의 핵심 장치입니다.

CPU vs GPU

딥러닝은 “쉬운 연산(곱하기, 더하기)을 엄청나게 많이” 하는 작업이라 GPU가 유리합니다.

RTX 4080 스펙 이해하기

스펙	값	의미
CUDA 코어	9,728개	동시에 연산할 수 있는 유닛 수
SM (Streaming Multiprocessor)	76개	CUDA 코어들을 묶은 그룹 (128코어/SM)
VRAM (Global Memory)	16GB GDDR6X	GPU 전용 메모리 (모델/데이터 저장)
메모리 대역폭	~717 GB/s	1초에 읽을 수 있는 데이터 양
L2 캐시	64MB	자주 쓰는 데이터 임시 저장
SRAM (Shared Memory)	64KB / SM	각 SM 내부의 초고속 메모리
텐서 코어	4세대	행렬 곱셈 전용 하드웨어 (매우 빠름)
Compute Capability	8.9 (Ada Lovelace)	GPU 아키텍처 세대

GPU 메모리 계층 (매우 중요!)

GPU 프로그래밍에서 가장 중요한 개념이 메모리 계층입니다. 성능 최적화의 핵심은 “느린 메모리 접근을 줄이는 것”입니다.

속도 비교 (대략적)

메모리	대역폭	HBM 대비
HBM (Global Memory)	~717 GB/s	1x
L2 Cache	~3-5 TB/s	~5x
SRAM (Shared Memory)	~19 TB/s	~26x
레지스터	훨씬 빠름	~100x+

핵심: HBM에서 데이터를 읽는 것이 병목입니다. 데이터를 한 번 SRAM으로 올리면, 그 안에서 여러 연산을 하는 게 훨씬 빠릅니다. 이것이 커널 퓨전의 핵심 원리입니다.

SM (Streaming Multiprocessor) 이해하기

SM은 GPU의 “미니 프로세서”입니다. RTX 4080에는 76개의 SM이 있습니다.

Warp란?

스레드(Thread): GPU에서 실행되는 가장 작은 실행 단위
Warp: 32개 스레드의 묶음 (항상 32개가 동시에 같은 명령 실행)
Thread Block: 여러 Warp의 묶음 (최대 1024 스레드)

Thread Block ├── Warp 0: [thread 0~31] ← 32개가 동시에 같은 명령 실행 ├── Warp 1: [thread 32~63] ├── Warp 2: [thread 64~95] └── ... 

Triton이 좋은 이유: CUDA에서는 이 모든 것을 직접 관리해야 하지만, Triton에서는 블록(프로그램) 단위로 생각하면 됩니다. Warp 관리는 컴파일러가 처리!

Compute-Bound vs Memory-Bound

GPU 연산은 두 가지 유형으로 나뉩니다:

Memory-Bound (메모리 병목)

연산량이 적고, 데이터 이동이 대부분인 경우 예: Vector Add, RMSNorm, Softmax → 원소 하나당 덧셈 1번만 하면 됨 → 대부분의 시간이 HBM에서 데이터 읽기/쓰기에 소비 최적화 전략: 커널 퓨전 (메모리 접근 횟수 줄이기) 

Compute-Bound (연산 병목)

연산량이 많고, 계산에 시간이 걸리는 경우 예: Matrix Multiplication (행렬 곱셈) → 원소 하나를 계산하려면 K번의 곱셈+덧셈 필요 → 연산량이 데이터 양보다 훨씬 많음 최적화 전략: 텐서 코어 활용, 타일링으로 데이터 재사용 

Arithmetic Intensity (연산 강도)

AI = 연산 횟수 / 메모리 접근 바이트 수 Vector Add: AI = 1 (낮음) → Memory-Bound Matrix Multiply: AI = N (높음) → Compute-Bound 

연산 강도가 높을수록 GPU의 계산 능력을 효과적으로 활용합니다.

텐서 코어 (Tensor Core)

행렬 곱셈을 위한 전용 하드웨어입니다.

일반 CUDA 코어: 텐서 코어: 한 클럭에 1번의 FMA 한 클럭에 4x4x4 행렬 곱 (a*b + c) (64번의 FMA를 한 번에!) 속도: 1x 속도: ~16x (fp16 기준) 

Triton에서 tl.dot(a, b)를 사용하면 자동으로 텐서 코어가 활용됩니다.

데이터 타입 — FP32, FP16, BF16, INT8

GPU 연산에서 데이터 타입 선택은 속도와 메모리 모두에 영향을 줍니다.

부동소수점이란?

컴퓨터에서 소수점 있는 숫자를 표현하는 방식입니다.

3.14159를 저장하려면? 부호(sign) | 지수(exponent) | 가수(mantissa) ± | 크기 범위 | 정밀도 비유: 과학적 표기법과 같습니다 3.14159 = 3.14159 × 10^0 31415.9 = 3.14159 × 10^4 ^^^^^^^^ ^^^^ 가수 지수 

왜 중요한가?

같은 모델을 다른 타입으로 저장하면: Llama 3 8B (80억 파라미터): FP32: 80억 × 4바이트 = 32GB → RTX 4080에 안 들어감! FP16: 80억 × 2바이트 = 16GB → 빠듯하게 들어감 INT8: 80억 × 1바이트 = 8GB → 여유 있게 들어감 INT4: 80억 × 0.5바이트 = 4GB → 넉넉! 

Triton에서의 데이터 타입

# 커널 내에서 타입 변환 x_fp32 = x.to(tl.float32) # 정밀한 중간 계산용 result = acc.to(tl.float16) # 저장 시 FP16으로 변환  # tl.dot은 내부적으로 FP32로 누적 (정밀도 유지) acc += tl.dot(a, b) # a, b는 FP16이지만 acc는 FP32 

대역폭 계산법 — 내 커널이 얼마나 효율적인지 측정하기

GPU 커널의 성능을 평가할 때 가장 중요한 지표가 메모리 대역폭 활용률입니다.

기본 공식

실효 대역폭 (GB/s) = (읽은 바이트 + 쓴 바이트) / 실행 시간(초) / 10^9 

예제: Vector Add

# 크기 N인 float32 벡터 2개를 읽고, 1개를 쓰는 경우 N = 10_000_000 # 1천만  읽기: x (N × 4바이트) + y (N × 4바이트) = 80MB 쓰기: output (N × 4바이트) = 40MB 총 데이터 이동: 120MB 실행 시간이 0.2ms 걸렸다면: 실효 대역폭 = 120MB / 0.0002초 = 600 GB/s RTX 4080 최대 대역폭: 717 GB/s 활용률: 600 / 717 = 83.7% ← 꽤 좋은 편! 

Occupancy — GPU를 얼마나 바쁘게 유지하는가

Occupancy는 SM이 실행할 수 있는 최대 warp 수 대비 실제 활성 warp의 비율입니다.

왜 중요한가?

이것을 Latency Hiding (지연 숨기기) 이라고 합니다. Warp가 많을수록 GPU가 쉬지 않고 일할 수 있습니다.

Triton에서는?

Triton 컴파일러가 자동으로 레지스터와 SRAM 사용을 최적화하므로, 대부분의 경우 Occupancy를 직접 걱정할 필요가 없습니다. 하지만 BLOCK_SIZE를 너무 크게 잡으면 SRAM이 부족해져 Occupancy가 떨어질 수 있습니다.

커널 실행의 전체 흐름

Python에서 Triton 커널을 호출하면 실제로 무슨 일이 일어나는지 봅시다.

Roofline Model — 내 커널의 한계를 이해하기

Roofline Model은 “이 GPU에서 이 커널이 이론적으로 얼마나 빠를 수 있는가”를 알려줍니다.

연산 강도(AI) < 68 → Memory-Bound (메모리가 병목) — Vector Add, Softmax, RMSNorm
연산 강도(AI) > 68 → Compute-Bound (연산이 병목) — MatMul, Flash Attention

Triton이 이 모든 것을 어떻게 단순화하는가

CUDA에서 직접 해야 하는 것	Triton에서는?
Shared memory 할당/관리	자동 (컴파일러가 처리)
Warp 동기화 (`__syncthreads`)	불필요
Memory coalescing 최적화	`tl.arange` + `tl.load`로 자동
텐서 코어 호출 (WMMA API)	`tl.dot`으로 자동
Thread block 크기 결정	`BLOCK_SIZE`만 지정
레지스터 압박 관리	컴파일러가 최적화

결론: Triton은 “GPU의 성능은 거의 다 뽑아내면서, CUDA보다 10배 쉽게 작성”할 수 있는 도구입니다.

용어 정리

용어	의미	비유
HBM / VRAM / Global Memory	GPU의 메인 메모리 (16GB)	큰 창고
SRAM / Shared Memory	SM 내부의 고속 메모리 (64KB)	작업대 위 선반
레지스터	연산 유닛 바로 옆의 메모리	손에 들고 있는 것
SM	코어 묶음 + 메모리를 가진 미니 프로세서	공장 하나
Warp	32개 스레드 묶음 (항상 같이 움직임)	32명이 한 줄로 행진
커널 퓨전	여러 연산을 하나의 커널로 합침	왕복 줄이기
타일링	큰 데이터를 작은 블록으로 나눠 처리	퍼즐 조각

LoRA vs Full Fine-tuning: An Illusion of Equivalence

Sat, 28 Dec 2024 15:00:00 +0000

LoRA vs Full Fine-tuning: An Illusion of Equivalence

Introduction

Pre-trained 모델을 downstream task에 fine-tuning하는 것은 computation-, data-efficient한 방법이다. 하지만 full fine-tuning은 모든 파라미터를 업데이트해야 하므로 시간과 비용적으로 부담이 크다. 이를 해결하기 위해 LoRA(Low-Rank Adaptation)와 같은 PEFT(Parameter-Efficient Fine-Tuning) 방법이 제시되었다. LoRA는 전체 파라미터의 극히 일부만 학습하면서도 full fine-tuning에 필적하는 성능을 보여주어 널리 사용되고 있다.

하지만 성능이 비슷하다고 해서 두 방법이 정말 같은 solution을 학습하는 걸까? 저자는 이 질문에 대해 weight matrix의 spectral properties(singular value decomposition)를 분석하여 답한다. 결론적으로, 같은 성능을 내더라도 LoRA와 full fine-tuning은 구조적으로 매우 다른 모델을 만든다.

핵심 발견:

LoRA는 intruder dimensions을 도입한다 — pre-trained weight의 singular vector와 거의 직교(orthogonal)하는 새로운 high-ranking singular vector가 나타난다. Full fine-tuning에서는 이런 현상이 없다.
Intruder dimensions은 forgetting을 유발한다 — intruder dimensions의 singular value를 줄이면 pre-training distribution의 모델링이 크게 개선되고, downstream 성능 저하는 미미하다.
Continual learning에서 LoRA가 더 취약하다 — 여러 task를 순차적으로 학습할 때 intruder dimensions이 누적되어 성능이 떨어진다.

Background

LoRA: Low-Rank Adaptation

Pre-trained weight matrix $W_0 \in \mathbb{R}^{m \times n}$에 대해, full fine-tuning은 모든 원소를 업데이트하여 $W = W_0 + \Delta W$를 학습한다. 이때 학습 가능한 파라미터 수는 $mn$이다.

LoRA는 $\Delta W$를 두 개의 low-rank 행렬의 곱으로 분해한다.

\[\Delta W = \frac{\alpha}{r} BA, \quad B \in \mathbb{R}^{m \times r}, \; A \in \mathbb{R}^{r \times n}\]

여기서 $r \ll \min(m, n)$이 rank이다. 학습 가능한 파라미터 수는 $(m + n)r$로, $mn$에 비해 매우 작다. 예를 들어 $m = n = 4096$이고 $r = 16$이면, 파라미터가 $16.8M$에서 $131K$로 약 128배 줄어든다.

초기화 시 $B = 0$, $A \sim \mathcal{N}(0, \sigma^2)$으로 설정하여 학습 시작 시 $\Delta W = 0$이 되도록 한다. Inference 시에는:

\[Y = W_{\text{tuned}} X = \left(W_0 + \frac{\alpha}{r} BA\right) X\]

Scaling Parameter $\alpha$

$\alpha$는 LoRA update의 크기를 조절하는 파라미터이다. 많은 실무에서 $\alpha = 2r$로 설정하는데, 이는 rank가 바뀌어도 update 크기가 일정하게 유지되도록 하기 위함이다. 이 논문에서는 $\alpha$의 선택이 intruder dimensions의 수와 forgetting에 큰 영향을 미친다는 것을 보여준다.

Singular Value Decomposition (SVD)

임의의 행렬 $M \in \mathbb{R}^{m \times n}$은 다음과 같이 분해할 수 있다.

\[M = U \Sigma V^\top\]

$U \in \mathbb{R}^{m \times m}$: 왼쪽 singular vectors (열벡터 $u_1, \ldots, u_m$)
$\Sigma \in \mathbb{R}^{m \times n}$: singular values $\sigma_1 \geq \sigma_2 \geq \ldots \geq 0$ (대각 행렬)
$V \in \mathbb{R}^{n \times n}$: 오른쪽 singular vectors

Singular value가 큰 singular vector일수록 행렬에서 더 중요한 방향을 나타낸다. 이 논문에서는 fine-tuning 전후의 singular vector 변화를 분석하여 LoRA와 full fine-tuning의 구조적 차이를 밝힌다.

Effective Rank

Effective rank는 singular value가 얼마나 집중되어 있는지를 측정한다. Frobenius norm의 90%를 차지하는 데 필요한 singular value의 개수로 정의한다. Effective rank가 낮으면 정보가 소수의 차원에 집중되어 있다는 뜻이다.

Structural Differences: Intruder Dimensions

핵심 관찰

저자는 Sharma et al. (2024)의 SVD 기반 pruning에서 영감을 얻어, fine-tuned weight matrix의 singular vector를 pre-trained weight matrix의 singular vector와 비교했다. 구체적으로, fine-tuned 모델의 각 singular vector $y_j$와 pre-trained 모델의 모든 singular vector $x_i$ 사이의 최대 cosine similarity를 측정했다.

Fig. 2는 이 분석의 핵심 결과를 보여준다.

(a) LoRA는 $W_0 + BA$로, full fine-tuning은 $W_0 + \Delta W$로 weight를 업데이트한다. 각각의 SVD를 구해서 pre-trained weight의 SVD와 비교한다.
(b) Full fine-tuning의 singular vector는 pre-trained singular vector와 높은 cosine similarity를 가진다 (대각선 구조). 반면 LoRA($r = 64$)의 singular vector는 일부가 매우 낮은 cosine similarity를 보인다.
(c) Intruder dimension은 모든 pre-trained singular vector와 낮은 cosine similarity를 가진다 (빨간 점).

Intruder Dimension의 정의

Definition 3.1. Fine-tuned weight matrix $W_{\text{tuned}}$의 singular vector $y_j$가 intruder dimension이라 함은, pre-trained weight matrix $W_0$의 모든 singular vector $x_i$에 대해 $\max_i(\cos(y_j, x_i)) < \epsilon$을 만족하는 것이다. 여기서 $\epsilon$은 similarity threshold이다.

직관적으로, intruder dimension은 pre-trained 모델이 전혀 알지 못하던 새로운 방향이다. Full fine-tuning은 기존 방향을 미세하게 조정하는 반면, LoRA는 완전히 새로운 방향을 추가한다.

Fig. 3은 이 차이를 시각적으로 보여준다. Full fine-tuning (왼쪽)의 similarity matrix는 깨끗한 대각선 구조를 보인다. 즉, fine-tuned singular vector $i$가 pre-trained singular vector $i$에 잘 대응된다. 반면 LoRA (오른쪽)에는 빈 열(empty column)이 존재하는데, 이것이 intruder dimensions이다 — 어떤 pre-trained singular vector와도 대응되지 않는 새로운 방향이다.

실험 모델과 데이터셋

저자는 두 가지 모델로 실험했다.

RoBERTa-base (125M): Encoder-only 모델. MNLI, QQP, SST-2, SiQA, WinoGrande, FEVER 등 6개 분류 태스크에서 fine-tuning
LLaMA2-7B / LLaMA-7B: Decoder-only 모델. Alpaca(instruction tuning), MetaMathQA(수학), Magicoder(코드) 등에서 fine-tuning

실험 결과

1. LoRA는 high-ranking intruder dimensions을 가지지만, full fine-tuning은 그렇지 않다

Top-$k$ singular vectors에 대해 Algorithm 1을 적용하면:

LoRA: $r \leq 16$에서 모든 $\epsilon$에 대해 지속적으로 intruder dimensions이 존재한다.
Full fine-tuning: 심지어 $\epsilon = 0.6 \sim 0.9$처럼 관대한 threshold에서도 intruder dimensions이 거의 없다.
Rank가 올라갈수록 intruder dimensions이 줄어든다. $r = 2048$에서는 full fine-tuning과 유사해진다.

2. 수학/코드 같은 어려운 태스크에서도 intruder dimensions이 존재한다

LLaMA2-7B를 MetaMathQA(수학)와 Magicoder(코드)로 fine-tuning한 경우에도 LoRA는 intruder dimensions을 보인다. 이는 pre-training domain과 target domain의 차이에서 비롯된다. 특히 코드처럼 pre-training과 매우 다른 도메인에서는 full fine-tuning도 일부 intruder dimensions을 보이지만, LoRA가 여전히 훨씬 많다.

3. Full fine-tuning은 LoRA보다 높은 effective rank를 가진다

$r = 768$인 full-rank LoRA조차도 실제 effective rank는 약 300에 그친다. 즉, LoRA는 파라미터화된 capacity $r$을 전부 활용하지 못한다. 반면 full fine-tuning의 $\Delta W$는 더 높은 effective rank로 업데이트를 수행한다. 이 차이는 코딩 같은 어려운 태스크에서 더 두드러진다.

4. Intruder dimensions은 singular value의 크기와 무관하게 존재한다

High singular value(중요한 방향)뿐 아니라 low singular value 영역에서도 LoRA는 full fine-tuning보다 항상 더 많은 intruder dimensions을 가진다.

5. $\alpha = 2r$로 설정하면 intruder dimensions이 줄고 effective rank가 늘어난다

많은 논문에서 $\alpha = 2r$로 설정한다. 저자는 $\alpha = 2r$과 $\alpha = 8$(고정)을 비교한 결과:

$\alpha = 8$: 모든 rank에서 intruder dimensions이 많고, effective rank가 낮다
$\alpha = 2r$: intruder dimensions이 적고, effective rank가 높다 — generalization도 더 좋다

이는 $\alpha = 2r$ 설정의 중요성을 뒷받침한다.

6. Fine-tuning 데이터가 많을수록 intruder dimensions이 늘어난다

$r = 8$에서 여러 데이터셋을 학습시키면 intruder dimensions이 누적된다. $r = 1$에서는 표현력 한계 때문에 일정하지만, 충분한 rank가 있으면 데이터가 많을수록 새로운 intruder dimensions이 추가된다.

Intruder Dimensions의 진화

Intruder dimensions은 학습 초기부터 점진적으로 나타난다. 학습이 진행됨에 따라:

Intruder dimension의 rank가 점점 높아진다 (더 중요한 위치를 차지한다)
Singular value가 점점 커진다 (영향력이 증가한다)
Pre-trained singular vectors와의 cosine similarity는 계속 낮게 유지된다

Model Differences: Forgetting과 Out-of-Distribution Generalization

LoRA는 full fine-tuning보다 덜 잊지만, 그 forgetting은 intruder dimensions에 집중된다

Fig. 6에서 forgetting(pre-training distribution의 pseudo loss 증가)을 측정한 결과:

LLaMA2-7B: LoRA는 rank가 작을수록 더 많이 잊는다. Full fine-tuning이 가장 적게 잊는다.
RoBERTa-base: 마찬가지로 LoRA의 rank가 낮을수록 pseudo loss가 크다.

놀라운 점은, learning rate가 클수록 intruder dimensions이 더 많이 생기고 (왼쪽), 그에 비례하여 forgetting도 증가한다는 것이다 (오른쪽). Intruder dimensions과 forgetting 사이의 Spearman correlation은 $\rho = 0.971$ ($p \ll 0.001$)로 매우 강한 상관관계를 보인다.

반면 intruder dimensions과 test accuracy 사이에는 통계적으로 유의미한 상관관계가 없다. 이는 intruder dimensions이 성능에는 불필요하지만 forgetting을 유발한다는 것을 의미한다.

Pre-training pseudo loss의 U-shaped curve

같은 test accuracy로 fine-tuning했을 때, pre-training pseudo loss가 rank에 대해 U자 곡선을 그린다. 이는:

Rank가 너무 낮으면: intruder dimensions의 영향으로 forgetting이 크다
Rank가 너무 높으면: overparameterization으로 target task에 overfitting

따라서 downstream task마다 최적의 rank가 존재한다.

Intruder Dimensions이 Forgetting을 야기하는가: Causal Intervention

상관관계를 넘어 인과관계를 확인하기 위해, 저자는 intruder dimensions에 직접 개입(intervention)하는 실험을 수행한다.

실험 방법

각 weight matrix에서 가장 high-ranking인 intruder dimension의 singular value를 조절한다. 구체적으로:

\[W = W_0 + \Delta W + (\lambda - 1) u_i \sigma_i v_i^\top\]

여기서 $i$는 top intruder dimension의 인덱스이다.

$\lambda = 0$: intruder dimension 완전 제거
$\lambda = 1$: 변화 없음 (원본)
$\lambda > 1$: intruder dimension 증폭

결과

Fig. 8에서 $\lambda$를 0부터 1까지 변화시키면:

Forgetting (빨간선): intruder dimensions을 줄이면 forgetting이 크게 감소한다
Test accuracy (초록선): 거의 변하지 않거나 아주 약간 감소한다

구체적 수치:

LLaMA2-7B (MetaMath, $r = 256$): $\lambda = 0.3$에서 test accuracy -0.1%, forgetting -33.3%
RoBERTa (QQP, $r = 8$): $\lambda = 0.7$에서 test accuracy 동일, forgetting -33.2%

일부 경우에는 $\lambda > 1$ (intruder dimension 증폭)에서 오히려 test accuracy가 향상되면서 forgetting이 악화되는 현상도 관찰된다.

중요한 점은, 이 효과가 intruder dimensions에만 해당한다는 것이다. Pre-trained singular vector에 가까운 normal dimensions의 singular value를 같은 방식으로 조절하면 forgetting에 큰 영향이 없고, 오히려 test accuracy가 떨어진다.

Continual Learning에서의 영향

실험 설정

Intruder dimensions이 누적되면 해로울 것이라는 가설을 검증하기 위해, RoBERTa를 6개 태스크에 순차적으로 fine-tuning한다: MNLI → QQP → SST-2 → SiQA → WinoGrande → FEVER.

각 태스크 학습 후 LoRA weights를 모델에 merge하고, adapter를 재초기화한 후 다음 태스크를 학습한다. 각 시점에서 모든 태스크의 성능을 측정한다.

결과

Fig. 9(a)에서:

Full fine-tuning: 6개 태스크를 순차 학습해도 이전 태스크 성능이 비교적 잘 유지된다
LoRA ($r = 1, 8, 64$): 모든 rank에서 full fine-tuning보다 빠르게 이전 태스크를 잊는다
Low rank($r = 1$)이 가장 빠르게 성능이 저하된다

Fig. 9(b)에서 LoRA($r = 8$)의 similarity matrix를 보면, 각 태스크가 학습될 때마다 새로운 intruder dimensions이 추가되는 것을 확인할 수 있다. MNLI(1번째) → QQP(2번째) → … → FEVER(6번째)로 갈수록 intruder dimensions이 누적된다.

반면 Fig. 9(c)에서 full fine-tuning의 similarity matrix는 6개 태스크를 거쳐도 대각선 구조가 유지된다. Pre-trained structure를 보존하면서 적응하기 때문이다.

왜 Intruder Dimensions이 발생하는가

1. Low-rank constraint

LoRA의 update $\Delta W = BA$는 rank $r$로 제한된다. Pre-trained weight의 주요 방향과 잘 align되지 않는 경우, LoRA는 제한된 rank 안에서 task를 풀기 위해 기존에 없던 새로운 방향을 만들어낸다. Rank가 충분히 높으면 ($r \geq 2048$) 이 문제가 완화되어 full fine-tuning과 유사해진다.

2. Learning rate와 gradient projection

LoRA는 일반적으로 full fine-tuning보다 더 큰 learning rate를 사용한다. 또한 gradient가 low-rank space로 projection되므로, 업데이트 방향이 full-rank gradient와 달라진다. 이 두 요소가 합쳐져 pre-trained structure에서 크게 벗어나는 업데이트가 발생한다.

3. Product parameterization ($BA$)

두 행렬의 곱은 spectral differences를 증폭시킨다. $B$만 학습시키고 $A$를 고정하면 intruder dimensions이 줄어드는 것이 이를 뒷받침한다.

Practical Implications

이 논문의 발견은 LoRA를 실무에서 사용할 때 다음과 같은 시사점을 준다.

Rank 선택: 가능하면 높은 rank를 사용하라. Low rank는 intruder dimensions을 유발하고 forgetting을 악화시킨다.
$\alpha$ 설정: $\alpha = 2r$이 $\alpha = \text{constant}$보다 intruder dimensions이 적고 generalization이 좋다.
Continual learning 주의: LoRA로 여러 태스크를 순차 학습하면 intruder dimensions이 누적되어 성능이 저하된다. 가능하면 adapter를 combine하지 말고 별도로 유지하라.
모델 선택 기준: 같은 test accuracy를 가진 두 모델 중, intruder dimensions이 적은 모델이 out-of-distribution에서 더 robust하다.
Post-hoc mitigation: 이미 학습된 LoRA 모델에서 intruder dimensions의 singular value를 줄이면 ($\lambda \approx 0.3 \sim 0.7$), 성능 저하 없이 forgetting을 크게 줄일 수 있다.

Conclusion

저자는 LoRA와 full fine-tuning이 같은 성능을 내더라도 weight matrix의 spectral properties에서 구조적으로 매우 다르다는 것을 보여주었다. LoRA는 pre-trained singular vectors와 orthogonal한 intruder dimensions을 도입하며, 이것이 forgetting의 직접적 원인이 된다. Full fine-tuning은 pre-trained spectral structure를 유지하면서 효과적으로 적응한다.

이 발견은 “LoRA가 full fine-tuning과 동등하다”는 통념에 의문을 제기하며, 성능 지표만으로는 두 방법의 차이를 포착할 수 없다는 점을 강조한다. 특히 continual learning처럼 여러 task를 순차 학습하는 현실적 시나리오에서 intruder dimensions의 누적은 심각한 문제가 될 수 있다.

역할	수량	기능
MMA warpgroup	1	Tensor Core 연산 (\(QK^\top\), \(PV\))
Softmax warpgroup	2	max, exp, rowsum 계산
Correction warpgroup	1	Rescaling (\(e^{m_{\text{old}} - m_{\text{new}}}\)로 보정)
TMA (producer)	-	HBM → SMEM 데이터 로드

Method	FP32 Max Rel Err	BF16 Max Rel Err
Hardware MUFU.EX2	\(1.41 \times 10^{-7}\)	\(3.89 \times 10^{-3}\)
Degree 3 polynomial	\(8.77 \times 10^{-5}\)	\(3.90 \times 10^{-3}\)
Degree 5 polynomial	\(1.44 \times 10^{-7}\)	\(3.89 \times 10^{-3}\)

blank

LLM 엔지니어가 알아야 할 GPU 아키텍처: Ampere → Hopper → Blackwell

왜 GPU 아키텍처를 알아야 하는가

1. 전체 스펙 비교

2. LLM 학습 관점: “얼마나 빨리 학습할 수 있는가”

2.1 Tensor Core 성능: 정밀도가 핵심

실무 포인트

2.2 실제 활용률: 이론 vs 현실

2.3 메모리 대역폭: “Tensor Core를 먹여살릴 수 있는가”

3. LLM 추론 관점: “얼마나 빨리 토큰을 생성하는가”

3.1 Prefill vs Decode

3.2 Decode 처리량 계산

3.3 KV Cache: “시퀀스가 길면 메모리가 부족하다”

3.4 정밀도와 양자화

4. 멀티 GPU 스케일링: “GPU를 더 쓰면 비례해서 빨라지는가”

4.1 NVLink: GPU 간 통신

실무 가이드라인

4.2 DGX 시스템 비교

5. 아키텍처별 핵심 신기능과 LLM 영향

Ampere (A100)

Hopper (H100)

Blackwell (B200)

6. 비대칭 스케일링: 왜 알고리즘이 중요한가

7. 실무 가이드라인

GPU 선택

정밀도 선택

FlashAttention 버전 선택

마치며

참고 문헌

FlashAttention-4: Algorithm and Kernel Pipelining Co-Design for Asymmetric Hardware Scaling

Introduction

Background: Blackwell GPU의 비대칭 스케일링

Hopper vs Blackwell 하드웨어 비교

Roofline 분석: Forward Pass

Forward Pass: 파이프라인 재설계

FA3 vs FA4 파이프라인 비교

Blackwell에서의 핵심 변화

지수함수 소프트웨어 에뮬레이션

문제: MUFU 병목

해결: 다항식 근사

Partial Emulation

조건부 Softmax Rescaling

기존 방식의 비효율

FA4의 개선: threshold \(\tau\)

Backward Pass

Roofline 분석

2-CTA MMA 모드

Backward Computation Graph

Deterministic Backward Pass

Scheduling: LPT와 Causal Masking

Longest-Processing-Time First (LPT)

Language: CuTe-DSL

Empirical Evaluation

Forward Pass

Backward Pass

Deterministic Backward Ablation

FlashAttention 시리즈 비교

Discussion and Conclusion

참고 문헌

FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision

Introduction

Background

Multi-Head Attention

Backward Pass

GPU 하드웨어 특성과 실행 모델

메모리 계층

쓰레드 계층

비동기 실행과 Warp Specialization

저정밀도 연산

Standard Attention과 FlashAttention 복습

FlashAttention-3: Algorithm

기본 구조: Warp Specialization 적용

Algorithm 1: Forward Pass (Warp Specialization, without intra-consumer overlapping)

Ping-Pong 스케줄링: GEMM과 Softmax 겹치기

Intra-Warpgroup 2-Stage 파이프라이닝

Algorithm 2: Consumer Warpgroup Forward Pass (2-Stage)

트레이드오프: 레지스터 압력

컴파일러 주의사항

Backward Pass

Low-precision FlashAttention: FP8

`triton.autotune` 이란?