제3부 딥러닝 라이브러리

3.7 어텐션 II — 융합 커널·실전 라이브러리

3.7.1 이 장에서 배우는 것 / 선수 지식

3.6에서 FlashAttention의 원리(online softmax·타일 융합)를 익혔다. 이 장은 그 원리가 실제로 어떤 라이브러리·백엔드로 구현돼 있고, 내 코드에서 무엇이 선택되는지·어떻게 강제·진단하는지를 다룬다. 핵심은 "직접 커널을 작성하지 말고, 올바른 융합 커널이 선택되게 하고 폴백을 검출하는 것"이다. 추론 디코딩(쿼리 길이 1)이라는 다른 형상과 GQA·sliding window 같은 변형도 본다.

배우는 것:

• PyTorch SDPA(scaled_dot_product_attention) 의 백엔드 자동선택(flash·mem-efficient·cuDNN·math)과, 무엇이 골라지고 왜 폴백되는지 진단·강제하는 법.
• 실전 라이브러리 지형 — flash-attn(FA2/FA3)·cuDNN SDPA·xFormers 의 역할 차이와 선택 기준.
• 학습(prefill)과 추론 디코딩(쿼리 1) 의 형상 차이, FlashDecoding 류 KV 병렬화가 왜 따로 필요한지(5.7~5.8 연결).
• GQA·MQA·sliding window·ALiBi 변형이 커널 지원·KV 메모리에 주는 영향.

선수 지식: 3.6(FlashAttention 원리 — 이 장은 그 적용), 3.3(cuDNN 그래프 융합 — cuDNN SDPA가 그 위), 1.13(dtype), 2.13~2.14(프로파일링 — 백엔드 실측), 3.1(GEMM). 추론 KV 캐시·서빙 본격은 5.7~5.8. 코드는 PyTorch 2.x·CUDA 12.x 가정이며 API·지원 범위는 버전마다 바뀜 — "시점 의존, 확인 필요".

3.7.2 지형 — 누가 무엇을 하나

어텐션 융합 커널은 여러 출처가 같은 원리를 구현한다. 직접 고르기 전에 지형을 안다:

핵심 멘탈모델: 대부분 SDPA 한 줄을 부르고, SDPA가 형상·dtype·마스크·하드웨어를 보고 백엔드를 고른다. 직접 flash-attn을 import하는 것은 SDPA가 못 고르는 기능(varlen 패킹·특정 최신 기능)이 필요할 때다.

마지막 math 폴백이 3.6의 표준 경로 — 이걸로 떨어지면 융합 이득이 사라진다. "어텐션 튜닝"의 실무 절반은 math 폴백을 막는 것이다.

3.7.3 워크스루 A — SDPA 백엔드 선택·강제·진단

PyTorch에서 어떤 백엔드가 실행되는지 보고, 특정 백엔드를 강제·금지한다(API 이름은 버전 의존):

import torch, torch.nn.functional as F
from torch.nn.attention import SDPBackend, sdpa_kernel   # PyTorch 2.x (경로는 버전 의존)

q = torch.randn(8, 32, 4096, 128, device="cuda", dtype=torch.bfloat16)  # [B,H,N,d]
k = torch.randn_like(q); v = torch.randn_like(q)

# 1) 기본: 자동선택
o = F.scaled_dot_product_attention(q, k, v, is_causal=True)

# 2) 특정 백엔드만 허용(컨텍스트). flash가 안 되면 에러 → 폴백 여부를 '강제로' 드러냄
with sdpa_kernel([SDPBackend.FLASH_ATTENTION]):
    o = F.scaled_dot_product_attention(q, k, v, is_causal=True)
    # 여기서 예외가 나면: 이 형상·dtype·정렬이 flash 커널 지원 밖이라는 뜻

# 3) 후보 우선순위 지정
with sdpa_kernel([SDPBackend.CUDNN_ATTENTION, SDPBackend.FLASH_ATTENTION,
                  SDPBackend.EFFICIENT_ATTENTION]):   # MATH를 빼면 폴백 금지
    o = F.scaled_dot_product_attention(q, k, v, is_causal=True)

진단 절차(math 폴백 잡기):

# math를 제외한 백엔드만 허용해 봤을 때 예외가 나면 = 평소 math로 폴백 중이었다는 신호
try:
    with sdpa_kernel([SDPBackend.FLASH_ATTENTION, SDPBackend.EFFICIENT_ATTENTION,
                      SDPBackend.CUDNN_ATTENTION]):
        F.scaled_dot_product_attention(q, k, v, is_causal=True)
    print("융합 커널 사용 가능")
except RuntimeError as e:
    print("폴백 원인:", e)   # dtype·head_dim·정렬·마스크 형태 등을 알려줌

실측은 Nsight Systems(2.13)로 커널 이름을 본다 — flash·fmha·cudnn이 보이면 융합, 큰 bmm+softmax+bmm 3커널이 보이면 math 폴백(N×N 구현 중).

폴백을 부르는 흔한 조건(시점 의존, 확인 필요):

• dtype이 fp32(대개 fp16·bf16만 융합), 또는 head_dim이 지원 밖(64·128·256 등 외).
• 임의 attn_mask(불리언·float 텐서)를 넘김 — is_causal=True 플래그 대신 큰 마스크 텐서를 주면 mem-efficient로 가거나 폴백.
• 정렬·연속성(non-contiguous)·비정상적인 stride.

3.7.4 학습(prefill) vs 추론(decode) — 형상이 다르다

FlashAttention(3.6)은 Q·K·V가 모두 길이 N인 prefill(학습·프롬프트 인코딩)에 최적이다. 그런데 자기회귀 추론 디코딩은 형상이 정반대다:

디코딩에서 한 스텝은 토큰 하나를 생성하며 누적된 KV 캐시(T 토큰) 전체와 어텐션한다. Q가 1행이라 3.6의 "Q 블록을 threadblock에 병렬 분배"가 무력화 — SM이 유휴 상태다.

FlashDecoding(및 변형). 해결책은 KV 차원(T)을 분할해 병렬화하는 것. KV를 여러 split으로 나눠 각 split의 부분 online softmax(부분 m·ℓ·O)를 병렬로 구하고, 마지막에 3.6의 보정 점화식으로 합친다(split들의 max·합을 재보정해 결합).

FlashDecoding: T가 길수록 병렬도↑ (긴 컨텍스트 디코딩에서 SM 채움).

이것이 5.7~5.8의 LLM 추론(인플라이트 배칭·페이지드 KV)과 맞물린다 — 디코딩 어텐션은 KV 캐시 관리와 떼어 생각할 수 없다. 요지: 같은 어텐션이라도 prefill과 decode는 다른 커널·다른 병렬화이며, 추론 스택은 둘을 구분해 처리한다.

3.7.5 변형 — GQA·MQA·sliding window·ALiBi

실전 모델은 순수 MHA가 아니다. 각 변형이 커널 지원·KV 메모리에 영향을 준다:

GQA·MQA(grouped·multi-query attention). KV 헤드를 Q 헤드보다 적게 둬 KV 캐시를 줄인다(추론 메모리·대역폭의 핵심). MQA는 KV 헤드 1개, GQA는 그 중간(예 Q 32헤드·KV 8헤드).

디코딩은 KV 캐시 읽기가 메모리 바운드라, KV 헤드 축소가 직접 속도·메모리 이득이다.

융합 커널은 GQA·MQA를 1급 지원한다(KV 헤드를 Q 헤드 그룹이 공유). 직접 KV를 복제(broadcast)해 MHA처럼 실행하면 메모리 이득이 사라지니, 커널의 GQA 경로를 써야 한다.

sliding window(국소 어텐션). 각 토큰이 최근 W개만 본다 → $N\times N$이 $N\times W$로, 긴 컨텍스트의 연산·KV를 선형화. 커널이 window를 지원해야 블록 skip으로 싸진다.

ALiBi·위치 바이어스. softmax 전에 거리 기반 바이어스를 더한다. 커널이 바이어스를 융합 지원하지 않으면 큰 마스크 텐서로 넘겨 폴백·메모리 증가.

softcap·logit 제한 등도 커널 플래그로 지원되어야 융합이 유지된다.

함정 선행: 모델이 GQA·sliding window·ALiBi를 쓰는데 커널이 그 변형을 지원하지 않으면 큰 마스크·복제로 우회하다 math 폴백·메모리 폭증. 모델 구성과 커널 지원 범위를 맞춰야 한다.

3.7.6 워크스루 B — flash-attn varlen (패킹된 가변 길이)

배치 내 시퀀스 길이가 제각각이면, 패딩해서 [B,H,N_max,d]로 만들면 패딩 토큰 계산이 낭비다. flash-attn의 varlen API는 패딩 없이 펼쳐 이어 붙인(packed) 시퀀스를 cu_seqlens(누적 길이)로 처리한다(직접 import, SDPA가 아직 못 하는 경우):

from flash_attn import flash_attn_varlen_func   # 별도 설치: pip install flash-attn (빌드 김)

# 배치 시퀀스 길이 [3, 5, 2] → 패딩 없이 이어붙임(총 10토큰)
# q,k,v: [총토큰=10, H, d]  (배치·패딩 차원 없음)
cu_seqlens = torch.tensor([0, 3, 8, 10], dtype=torch.int32, device="cuda")  # 누적 경계
max_seqlen = 5

o = flash_attn_varlen_func(
    q, k, v,
    cu_seqlens_q=cu_seqlens, cu_seqlens_k=cu_seqlens,
    max_seqlen_q=max_seqlen, max_seqlen_k=max_seqlen,
    causal=True,
)   # 각 시퀀스가 자기 경계 안에서만 어텐션, 패딩 연산 0

이득(개략): 패딩 방식은 N_max에 맞춰 다 계산해 짧은 시퀀스가 많으면 절반 이상이 낭비. varlen은 실제 토큰만 계산 — 가변 길이 배치(특히 학습 데이터·인플라이트 배칭)에서 처리량을 크게 높인다(5.8). SDPA에서도 NestedTensor 등으로 일부 지원되나, 최신·세밀한 제어는 flash-attn varlen이 앞선다(시점 의존).

3.7.7 중간 점검 (체크포인트)

• [ ] SDPA가 백엔드를 자동선택함을 알고, sdpa_kernel 컨텍스트로 강제·진단해 math 폴백을 검출한다.
• [ ] prefill(Q 길이 N)과 decode(Q 길이 1)의 형상 차이, decode에 FlashDecoding 류 KV 병렬화가 왜 필요한지 설명한다.
• [ ] GQA·MQA가 KV 캐시를 줄여 디코딩 메모리·속도에 주는 효과를 말하고, 커널의 GQA 경로를 써야 함을 안다.
• [ ] 어떤 조건(dtype·head_dim·마스크 형태·정렬)이 융합을 깨고 폴백을 부르는지 나열한다.

증상→의심: 어텐션이 의외로 느리고 메모리 큼 → math 폴백(sdpa_kernel로 강제 진단·Nsight로 커널 확인). 디코딩만 GPU 활용률 낮음 → Q=1 형상, FlashDecoding/추론 스택 필요. GQA 모델인데 KV 메모리가 MHA만큼 큼 → KV 복제 경로(GQA 커널 미사용). 커스텀 마스크 넣자 느려짐 → 큰 마스크 텐서로 폴백(가능하면 is_causal·window 플래그로).

3.7.8 트레이드오프 — 무엇을 부를 것인가

견해: 학습·연구는 SDPA를 기본으로 두고, math 폴백만 감시한다. 가변 길이 배치·최신 커널 기능이 필요하면 flash-attn varlen. LLM 추론 배포는 어텐션 커널만 따로 고르지 말고 TensorRT-LLM·서빙 스택(5.7~5.14)이 prefill·decode·KV·배칭을 통합 관리하게 한다 — 디코딩 어텐션은 KV 캐시 정책과 분리 불가.

비NVIDIA 정직: vLLM·SGLang(5.15)은 PagedAttention(KV를 페이지로 관리)과 자체 융합 커널로 같은 문제를 푼다. FlashAttention 알고리즘은 공개라 이식돼 있고, 백엔드 경쟁은 활발하다. "어느 커널이 최신·최고인가"는 분기마다 바뀌므로 벤치마크 수치를 외우지 말고 그때그때 실측한다.

3.7.9 흔한 함정

• math 폴백을 모르고 방치. 가장 흔한 손실. dtype·head_dim·마스크·정렬이 안 맞아 조용히 N×N 경로로 떨어진다. sdpa_kernel로 math 제외하고 예외를 드러내 점검(3.7.3).
• 큰 마스크 텐서 남발. is_causal=True 플래그 대신 [N,N] 불리언/float 마스크를 만들어 넘기면 그 자체가 $O(N^2)$ 메모리고 융합을 깨 폴백시킨다. causal·window는 플래그/전용 인자로.
• GQA 모델에 KV 복제. GQA를 모르고 KV를 Q 헤드 수로 broadcast하면 KV 캐시·대역폭 이득이 사라진다. 커널의 GQA 경로 사용.
• prefill 커널로 decode 기대. 디코딩(Q=1)에 prefill FA를 그대로 쓰면 병렬도가 1로 축소되어 SM이 유휴 상태다. 추론은 FlashDecoding·서빙 스택(5.7).
• 버전·빌드 불일치. flash-attn은 CUDA·PyTorch·아키텍처에 맞춰 빌드해야 하고, 안 맞으면 import 실패·런타임 폴백. FA3는 특정 세대(Hopper) 전용 경로가 있음(4.2). 설치 매트릭스 확인.
• 수치 차이를 버그로 오해. 융합·블록 누적은 표준과 비트가 다를 수 있다(수학적 동일, 부동소수 순서). 정상. 검증은 허용오차로(2.2).
• 벤치마크에 워밍업·폴백 섞임. 첫 호출이 커널 선택·컴파일로 느리거나, 측정 형상이 폴백을 타면 잘못된 결론. 워밍업 후·실제 형상으로 Nsight 확인(2.13).

3.7.10 연습문제

연습 1 (폴백 진단). 어떤 모델의 어텐션이 긴 시퀀스에서 느리고 메모리가 크다. sdpa_kernel을 써서 "지금 math로 폴백 중인지"를 확정하는 절차와, 그다음 원인을 좁히는 단계를 쓰라.

연습 2 (decode 병렬도). 컨텍스트 T=8192로 한 토큰을 디코딩한다. prefill용 FlashAttention(Q 블록 병렬)을 그대로 쓰면 왜 GPU가 유휴 상태인지, FlashDecoding이 무엇을 분할해 해결하는지 설명하라.

연습 3 (GQA KV 메모리). 모델 Q 헤드 32, head_dim 128, fp16, 컨텍스트 T=16384, 배치 1. (a) MHA(KV 헤드 32)와 (b) GQA(KV 헤드 8), (c) MQA(KV 헤드 1)의 KV 캐시 크기를 각각 구하고(K·V 둘 다), 디코딩 속도에 주는 함의.

연습 4 (마스크 함정). 인과 어텐션을 구현하며 [N,N] float 마스크(상삼각 -inf)를 만들어 attn_mask로 SDPA에 넘겼더니 긴 시퀀스에서 메모리가 부족해지고 느려졌다. 원인과 수정.

연습 5 (배포 결정). 사내 LLM을 추론 서빙한다. "어텐션이 느리니 flash-attn을 직접 import해 커널만 바꾸자"는 제안의 한계와, 더 나은 접근.

3.7.11 요약 · 다음 장

• 대부분 PyTorch SDPA 한 줄을 부르고 SDPA가 flash·cuDNN·mem-efficient·math 백엔드를 자동선택한다. 실무의 핵심은 math 폴백(N×N 구현)을 sdpa_kernel·Nsight로 검출하고 dtype·head_dim·마스크·정렬을 맞춰 융합을 유지하는 것.
• prefill(Q=N)과 decode(Q=1)는 다른 형상·다른 병렬화 — 디코딩은 KV 차원을 분할하는 FlashDecoding 류가 필요하고, KV 캐시 관리(5.7~5.8)와 분리 불가.
• GQA·MQA는 KV 캐시를 줄여 디코딩 메모리·속도에 직접 이득 — 커널의 GQA 경로를 써야 실현된다. sliding window·ALiBi·varlen은 커널 지원 여부가 융합·메모리를 가른다.
• 배포는 어텐션 커널만 교체하지 말고 서빙 스택이 prefill·decode·KV·배칭을 통합 관리하게 한다.

다음 장: 3.8 데이터 로딩·전처리 가속 — DALI·CV-CUDA·nvJPEG. 모델이 아무리 빨라도 CPU 데이터 파이프라인이 GPU를 기아 상태로 만들면 소용없다. 디코딩·증강·전처리를 GPU로 옮겨 입력 병목을 없애는 도구로 제3부를 마무리한다.