Pytorch 추론 가속 기법 정리 - 100-hours-a-week/5-yeosa-wiki GitHub Wiki

1. BF16 정밀도 사용

FluxKontextPipeline(
    ...
    torch_dtype=torch.bfloat16  # 또는 torch.float16
)

a. 개념

  • 딥러닝 모델의 수치는 기본적으로 32비트 부동소수점(float32) 형태로 저장
  • 동일한 모델이라도 16비트 표현 방식bfloat16 또는 float16으로 변경하면 메모리 사용량이 절반으로 감소
    • A100 GPU와 같은 최신 하드웨어에는 16비트 연산에 최적화된 Tensor Core가 내장되어 있어 속도 향상이 더욱 뚜렷함

b. 용어 풀이

용어 의미
float32 32비트 부동소수점. 가장 일반적이고 정확도가 높음
float16 16비트 부동소수점. 지수 5비트, 가수 10비트 → 표현 범위 좁음
bfloat16 16비트 부동소수점. 지수 8비트 → FP32와 동일한 표현 범위, 가수 7비트 → 정밀도 낮음
torch_dtype PyTorch에서 모델과 연산의 데이터 타입을 지정하는 설정

[ float16이 NaN 발생 위험이 있는 이유 ]

  • float16은 숫자를 16비트로 표현하는 부동소수점 방식

    구성 요소 의미
    부호 비트 (1비트) 양수/음수를 결정
    지수 비트 (5비트) 숫자의 크기(스케일)를 결정
    가수 비트 (10비트) 소수점 이하 정밀도를 결정
    • 지수 비트가 5비트 → 표현 가능한 지수 범위가 11 비트(float32) → 5 비트(float16)으로 크게 줄어듦
  • 표현 범위를 초과하는 값이 계산 중에 발생하면 NaN(Not a Number) 또는 Inf(Infinity)가 생성될 수 있음

    • float16의 표현 가능한 값의 범위:
      • 최댓값: 약 ±6.55 × 10⁴
      • 최솟값: 약 ±6.1 × 10⁻⁵
    • 계산 중 결과가 이 범위를 벗어나면:
      • 너무 크면 → Inf
      • 너무 작으면 → Underflow (0으로 처리)
      • 잘못된 연산 (예: 0으로 나누기) → NaN
  • FP16은 지수 비트가 좁아서 자주 Inf 또는 NaN으로 튀게 됨

    • 특히 Transformer, Attention처럼 수치 폭발이 자주 발생하는 모델에서는 위험

c. 원리

[ 연산 단위가 작아짐 ]

  • 32비트 → 16비트로 줄이면 한 번에 처리할 수 있는 데이터 양이 2배 증가

[ Tensor Core 활용 ]

  • A100 GPU는 FP16, BF16 연산에 최적화된 Tensor Core를 사용 → 일반적인 FP32 연산 대비 2~4배 속도 향상

[ 메모리 절약 ]

  • 데이터 크기가 작아져 GPU 메모리 사용량이 절반으로 감소 → 더 많은 데이터를 병렬 처리 가능

[ 수치 안정성 (BF16) ]

  • FP16은 표현 가능한 값의 범위가 좁아 NaN, Inf 발생 위험이 있으나, BF16은 FP32와 동일한 지수 표현으로 오버플로우, 언더플로우에 강함

d. 요약

비교 항목 float32 (FP32) float16 (FP16) bfloat16 (BF16)
비트 수 32비트 16비트 16비트
표현 범위 매우 넓음 좁음 넓음 (FP32 동일)
정밀도 매우 높음 중간 낮음
속도 (A100) 느림 빠름 빠름 (FP16 ≒ BF16)
안정성 최고 낮음 (NaN 위험) 높음 (FP32 수준)
  • A100 환경에서는 대부분 BF16이 가장 안전하고 빠른 선택

2.모델을 컴파일해서 더 빠르게 실행

# 1. 모델 컴파일 및 실행
compiled_model = torch.compile(model)
output = compiled_model(inputs)

# 2. 캐시 저장 (torch 2.7.0 이상)
artifacts = torch.compiler.save_cache_artifacts()

# 3. load 캐시해서 컴파일된 모델 초기화
torch.compiler.load_cache_artifacts(artifact_bytes)

a. 개념

  • PyTorch는 기본적으로 동적 실행 방식(Eager Execution)으로 동작
    • 매번 Python 코드가 직접 호출되고 연산 하나하나가 실행됨 → 반복 시 속도 저하
  • torch.compile()을 사용하면 모델을 미리 분석하고 하나의 최적화된 실행 계획으로 변환하여 실행
    • 결과적으로 반복적인 추론 또는 훈련에서 실행 속도를 크게 향상

b. 용어 풀이

용어 의미
컴파일(Compile) 코드를 미리 분석·변환더 빠르게 실행 가능한 형태로 만드는 과정
그래프(Graph) 연산들의 순서도(Flow)를 의미. 각 연산(Operation)이 노드(node)가 되고, 데이터 흐름이 엣지(edge)가 됨
TorchDynamo PyTorch 내부에서 Python 코드를 그래프로 변환해주는 엔진
AOT Autograd 미리 자동미분(Backward) 경로까지 계산해놓는 최적화 방식
Inference Graph 모델의 순전파(forward) 연산만 포함된 실행 그래프

[ 핵심 ]

  • 그래프는 GPU 또는 최적화 엔진이 한 번에 실행 가능한 연산 묶음으로 구성
  • 컴파일은 이 그래프를 만들어서 재사용하는 방식으로 속도를 높임

c. 원리

가. 컴파일 전 (Eager Mode: 일반 PyTorch)

  • Pytorch의 기본 방식
  • 연산 하나하나를 Python 코드(Python 인터프리터)가 직접 호출하고, 각 연산을 GPU로 보냄
  • GPU에서 연산이 끝날 때마다 CPU(Python)가 다음 연산을 다시 스케줄링 → GPU와 CPU 사이에 왕복 비용 발생
  • 이 과정에서 CPU ↔ GPU 간 컨텍스트 전환Python 오버헤드가 지속적으로 발생

나. 컴파일 후 (Graph Mode: torch.compile)

  • 모델의 전체 연산 흐름(그래프)을 한 번에 묶어서 GPU에 전달
  • CPU는 한 번만 스케줄링하고, 나머지는 GPU에서 연속적으로 처리
  • 불필요한 Python 호출 제거 → CPU 개입 최소화
  • 연산 순서, 중복, 메모리까지 자동 최적화 가능

다. 플로우 비교

단계 Eager Mode (기존) Graph Mode (torch.compile)
코드 실행 매번 Python 함수 호출 Python → 그래프 1회 변환
연산 처리 연산 1개 → GPU 1회 호출 다수 연산 → GPU 1회 호출 (묶음)
속도 느림 빠름 (최대 2~3배)

라. 핵심

항목 Eager Mode (기존) Graph Mode (torch.compile)
연산 실행 방식 Python이 매번 호출 그래프 한번 생성, GPU 단독 실행
CPU ↔ GPU 왕복 매 연산마다 발생 최초 1회 이후 GPU 단독 진행
오버헤드 거의 없음
속도 느림 빠름

3. Scaled Dot Product Attention (SDPA)

from torch.nn.functional import scaled_dot_product_attention

a. 개념

  • Transformer 모델에서 사용되는 어텐션(attention) 연산을 더 빠르고 메모리 효율적으로 계산하는 방법
  • PyTorch 2.0 이상에서는 기본적으로 scaled_dot_product_attention() 함수로 제공

b. 용어 풀이

용어 의미
Attention 입력의 여러 부분 중 중요한 정보에 집중하도록 가중치를 계산하는 방식
Attention Score (QKᵀ) 쿼리(Q)와 키(K)의 내적 결과. 입력 토큰 간 관계를 나타냄
Softmax Attention Score를 확률처럼 변환하여 가중치를 부여
FlashAttention GPU 메모리와 속도를 극적으로 개선한 어텐션 가속 기법
SDPA Scaled Dot Product Attention. FlashAttention의 아이디어를 PyTorch에 내장한 방식

c. 원리

Flash Attention 기법 설명
Recompute Trick softmax(QKᵀ) 계산을 저장하지 않고, 뒤에서 다시 계산 (메모리 ↓)
Block-wise Streaming attention 행렬을 한 줄씩 스트리밍 처리 (전체를 안 만들고도 가능)
Custom CUDA Kernel GPU에 최적화된 전용 연산 경로 사용 (레지스터 기반)
Low precision + fused ops FP16/BF16 + softmax와 matmul을 하나의 커널로 묶음

가. 기존 어텐션 (느린 방식: Eager Mode)

  1. 쿼리(Q)와 키(K)를 곱해서 Attention Score (QKᵀ) 계산 → GPU 메모리 저장
  2. 저장된 Score에 Softmax 적용 → GPU 메모리 저장
  3. Softmax 결과와 값(V)을 곱함 → GPU 메모리 저장
  4. 마지막으로 결과 출력

→ 즉, 계산 → 메모리 → 계산 → 메모리의 패턴이 반복

→ 계산이 끝날 때마다 GPU 메모리에 중간결과가 쌓임

→ 메모리 낭비 + 대기시간 발생

나. SDPA (빠른 방식: Graph Mode)

[ 핵심 1: Attention Score를 레지스터에만 보관 (On-the-fly 계산) ]

  • QKᵀ 전체를 미리 계산하지 않고, 필요한 부분만 계산 → 바로 다음 연산으로 넘김
  • 이때 사용하는 것이 GPU의 레지스터 (Register) — 매우 빠르고 작은 임시 저장소

→ 결과적으로 Attention Score 전체를 GPU 메모리에 쓰지 않음

→ 메모리 사용량 O(L)

[ 핵심 2: Softmax와 곱셈의 스트리밍 처리 ]

  • 기존 : Score = softmax(QKᵀ)Output = Score × V
  • SDPA : Softmax와 Value 곱셈을 분리하지 않고, 연속된 한 과정으로 처리
  • 이렇게 하면 softmax 값도 메모리에 저장하지 않고, 바로 다음 연산으로 연결 → 중간 메모리 사용량 대폭 감소

[ 정리 ]

단계 기존 방식 SDPA 방식 (Flash)
QKᵀ 계산 전체 메모리 저장 계산 → 바로 소프트맥스 → 바로 곱셈 (레지스터)
Softmax 별도로 계산 후 메모리 저장 소프트맥스 값도 바로 사용 (저장 안 함)
곱셈 softmax 결과와 V 곱셈 (메모리 필요) 곱셈까지 한 번에 처리

다. GPU Register의 작은 용량을 극복한 방법

[ GPU 레지스터의 특징 ]

  • GPU 레지스터(Register): GPU 안에서 가장 빠르고 가장 가까운 임시 저장 공간
  • 용량: 아주 작음 (몇 MB 수준)
  • 속도: GPU 메모리(GDDR6/HBM)보다 수십~수백 배 빠름
저장소 속도 용량
레지스터 가장 빠름 아주 작음
공유 메모리 빠름 중간 (수십 KB)
글로벌 메모리 느림 (GDDR) 큼 (GB급)

[ SDPA의 핵심 아이디어 : 큰 작업을 작은 조각(block)으로 쪼개어 처리 ]

  1. 기존 방식
    • 시퀀스 길이 L이면, L × L 전체 Attention Score (QKᵀ) 행렬을 한 번에 계산 → 메모리에 저장
    • O(L²) 메모리 요구
  2. SDPA 방식:
    • 작은 Block 단위로 쪼개어 계산 (예: 64×64, 128×128)
    • 이 작은 블록 하나는 GPU 레지스터/공유메모리에 충분히 담을 수 있음
    • 블록 단위로 Attention Score → Softmax → Value 곱까지 한 번에 처리 후 버림
    • 다음 블록 진행

→ 전체 시퀀스를 한 번에 계산하는 것이 아니라, 부분 부분 처리해서 메모리 없이 이어붙이는 방식

[ 왜 가능할까? ]

  • Attention 연산의 블록 단위 병렬성 덕분

  • Softmax의 성질(선형 연산 아님)도 부분 계산 → 결합이 가능하도록 설계됨

  • GPU 아키텍처(A100, H100)는 이런 블록 처리에 최적화된 Tensor Core 보유

  • Softmax와 블록 단위 Attention 계산

    [ FlashAttention 방식 (Memory-efficient Softmax) ]

    • 핵심 트릭
    1. Softmax는 선형이 아니지만, Max-TrickPartial Sum으로 쪼갤 수 있다.
    2. 블록별로 계산한 값을 Incremental Reduction (누적)할 수 있다.
    • 단계
      • Block 1 → S₁, Softmax numerator₁, denominator₁ 계산 (exp, sum)
      • Block 2 → S₂, Softmax numerator₂, denominator₂ 계산 (exp, sum)
      • 각 블록의 결과를 합산해서 전체 Softmax의 정확한 결과 복원

    → 즉, 블록마다 exp()sum()을 미리 구해놓고, 마지막에 정확한 Softmax normalization을 위해 global maxglobal sum을 사용

    [ 작은 4×4 행렬로 보는 블록 계산 예시 ]

    • 4×4 Attention Score 행렬
    K1 K2 K3 K4
    Q1 s11 s12 s13 s14
    Q2 s21 s22 s23 s24
    Q3 s31 s32 s33 s34
    Q4 s41 s42 s43 s44
    • 블록 나누기 (2×2)
    Block 1 Block 2
    s11 s12 s13 s14
    s21 s22 s23 s24
    Block 3 Block 4
    s31 s32 s33 s34
    s41 s42 s43 s44
    • Block 1 (Q1의 일부 Scores: s11, s12)

      1. 계산:

        $\text{exp}(s11),\quad \text{exp}(s12)$

      2. Partial sum:

        $\text{sum}_1 = \text{exp}(s11) + \text{exp}(s12)$

      3. Softmax numerator 준비:

        $numerator_1 = exp(s11) × V1, exp(s12) × V2$

      ⇒ 모두 레지스터에서만 유지 → 메모리 미사용

    • Block 2 (Q1의 나머지 Scores: s13, s14)

      1. 계산:

        $\text{exp}(s13), \quad \text{exp}(s14)$

      2. Partial sum:

        $\text{sum}_2 = \text{exp}(s13) + \text{exp}(s14)$

      3. Softmax numerator 준비:

        $numerator_2 = exp(s13) × V3, exp(s14) × V4$

      ⇒ 역시 레지스터에서만 유지

    • 모든 블록 완료 후 → Softmax 복원

      • Q1의 전체 $sum = sum₁ + sum₂$

      • Q1의 전체 $numerator = numerator₁ + numerator₂$

      • 최종 Softmax 결과 (Q1):

        $\text{Output}_{Q1} = \frac{\text{numerator}_1 + \text{numerator}_2}{\text{sum}_1 + \text{sum}_2}$

      ⇒ Softmax의 정확한 분모/분자를 블록별로 나눠 계산 후 합산 → 수학적으로 정확히 동일

d. 플로우 비교 (어텐션 계산 방식)

단계 기존 어텐션 (기본) SDPA (PyTorch 2.0+)
Attention Score 계산 메모리 전체 저장 (QKᵀ) 필요할 때만 계산 (레지스터 활용)
Softmax 적용 별도 실행 곱셈과 함께 스트리밍 처리
메모리 사용량 O(n²) O(n) (시퀀스 길이에 따라)
속도 느림 2~4배 빠름 (특히 긴 시퀀스에서 효과)

e. 요약

  • SDPA는 PyTorch에서 자동으로 선택되어 사용되므로 별도 설치 없이도 가속 효과를 얻을 수 있음
  • 특히 A100 GPU에서는 bfloat16, SDPA 조합으로 최고의 속도와 효율 달성
  • 참고 : https://arxiv.org/abs/2205.14135

Graph Mode?

torch.compile의 Graph Mode vs SDPA의 Graph Mode

가. 공통점 (왜 둘 다 Graph라는 말을 쓰나?)

  • 둘 다 계산 흐름(연산의 순서도)을 하나의 덩어리로 만들어 최적화한다는 공통점
    • 즉, 기존의 하나하나의 개별 연산 호출연속된 계산 흐름(그래프) 으로 변환한다는 점에서 모두 Graph Mode라고 부름

나. 차이점

구분 torch.compile() (모델 레벨) SDPA (연산 레벨)
적용 범위 모델 전체 (Layer, Module, Forward 전체) Attention 연산 하나 (QKᵀ, softmax, matmul)
최적화 대상 모델의 모든 연산, 루프, 조건문 등 Scaled Dot Product Attention의 세부 연산
그래프 생성 시점 Python → TorchDynamo → Inductor 단계 이미 호출된 어텐션 함수 내부
목적 Python 오버헤드 제거, 연산 결합, 메모리 최적화 특정 연산의 속도/메모리 최적화

다. torch.compile()의 Graph Mode → "모델 실행 그래프"

  • 입력 → 여러 Layer → 여러 연산 → 출력까지의 전체 연산 흐름을 하나의 그래프로 컴파일
  • 기존의 Python 함수 호출 → GPU 연산 호출 → Python 복귀 → 이런 왕복을 제거
  • CPU ↔ GPU 오버헤드, 반복 호출 최적화
model = torch.compile(model)  # 모델 전체가 그래프 실행으로 변환

라. SDPA의 Graph Mode → "어텐션 연산 최적화 그래프"

  • scaled_dot_product_attention() 함수 내부의 계산이 GPU에서 하나의 최적화된 커널로 실행
  • 중간 결과 없이 Softmax, Matmul, Scoring을 모두 묶어서 GPU가 직접 연산
output = scaled_dot_product_attention(query, key, value)
  • 여기는 특정 연산 단위에서의 Graph Mode → 주로 FlashAttention 커널 사용

마. 비유로 정리

비유 의미
torch.compile 요리 레시피 전체를 미리 계획해서 효율적으로 실행 (모델 전체)
SDPA 칼질-볶기-양념을 한 번에 묶어서 빠르게 처리 (특정 작업 최적화)

torch.compile은 모델 전체 스케줄링 관점

SDPA는 개별 연산의 미시적 가속 관점

바. 요약

항목 torch.compile (Graph Mode) SDPA (Graph Mode)
적용 대상 모델 전체 (Module, Layer) 어텐션 연산 하나
속도 개선 포인트 Python 호출 제거, 연산 결합 GPU 커널 최적화
메모리 최적화 가능 매우 강력 (Flash)
함께 사용 가능? Yes Yes

4. channels_last 메모리 포맷 사용

model = model.to(memory_format=torch.channels_last)
input = input.to(memory_format=torch.channels_last)

a. 개념

  • 이미지나 영상 데이터는 기본적으로 (Batch, Channels, Height, Width) 형태로 저장
  • 이걸 (Batch, Height, Width, Channels) 형태로 채널을 마지막에 두는 방식(channels_last) 으로 변경하면, GPU에서 데이터를 더 빠르고 효율적으로 처리할 수 있음
  • 연산 결과의 정확도에는 영향을 주지 않으며, 메모리 배치 구조만 바뀌는 것

b. 용어 풀이

용어 의미
Channels First (NCHW) (Batch, Channels, Height, Width) → PyTorch 기본 메모리 포맷
Channels Last (NHWC) (Batch, Height, Width, Channels) → GPU 최적화된 메모리 포맷
memory_format PyTorch에서 텐서의 메모리 배치 방식을 지정하는 설정
  • N: Batch 크기
  • C: 채널 (예: RGB 3채널)
  • H: 이미지 높이
  • W: 이미지 너비

c. 예시

가. 가정: 1장의 2×2 크기 RGB 이미지

  • 높이(Height) = 2 픽셀
  • 너비(Width) = 2 픽셀
  • 채널(Channels) = 3 (R, G, B)

→ 즉, 2×2 이미지에 각 픽셀마다 R/G/B 값이 있는 상태

나. NCHW 포맷 (Channels First)

Shape: (1, 3, 2, 2)
→ (Batch, Channels, Height, Width)
채널 데이터 (픽셀 값) 의미
R [1, 2], 3, 4 빨간색(Red) 채널의 밝기값 (0~255 가정)
G [9, 10], 11, 12 초록색(Green) 채널의 밝기값
B [17, 18], 19, 20 파란색(Blue) 채널의 밝기값
  • [1,2]는 이미지의 (0,0), (0,1) 두 픽셀의 빨간색 밝기
  • [3,4]는 이미지의 (1,0), (1,1) 두 픽셀의 빨간색 밝기

→ 즉, 숫자는 각 픽셀 위치에서의 색상별 밝기값

다. NHWC 포맷 (Channels Last)

Shape: (1, 2, 2, 3)
→ (Batch, Height, Width, Channels)
픽셀 위치 (H,W) [R, G, B] 값 실제 의미
(0,0) [1, 9, 17] → 빨강=1, 초록=9, 파랑=17 좌측 상단 픽셀의 RGB 밝기
(0,1) [2, 10, 18] 우측 상단 픽셀
(1,0) [3, 11, 19] 좌측 하단 픽셀
(1,1) [4, 12, 20] 우측 하단 픽셀
  • 즉, 이제는 한 픽셀당 RGB 값이 함께 묶여서 메모리에 저장
    • 여기서 1,9,17 같은 값은 0~255 범위의 이미지 밝기값 (uint8, 실제 딥러닝에서는 보통 0~1로 정규화되어 들어감)

라. 정리

포맷 메모리 구조 픽셀 접근 방식
NCHW (Channels First) 픽셀 색상 정보가 분리되어 저장 픽셀 하나의 R/G/B를 각각 다른 위치에서 읽어야 함
NHWC (Channels Last) 픽셀 색상 정보가 함께 저장 픽셀 하나의 R/G/B를 연속적으로 읽어올 수 있음

c. 원리

가. GPU 메모리의 데이터 접근 방식

  • GPU는 데이터를 연속적인 메모리 블록(linear memory) 에서 읽을 때 가장 빠름
  • Channels First (NCHW) 포맷에서는 같은 픽셀의 R, G, B 값이 서로 멀리 떨어진 메모리 위치에 저장됨
    • 예: (0,0) 픽셀의 R → G → B 값이 떨어져 저장
  • Channels Last (NHWC) 포맷에서는 같은 픽셀의 R, G, B 값이 메모리상에서 연속적으로 저장
    • 예: (0,0) 픽셀의 [R, G, B] 값이 붙어있음

→ 이렇게 붙어 있어야 GPU가 한 번에 읽어들이는 메모리 대역폭(memory coalescing) 을 최대로 활용 가능

→ 반대로 멀리 떨어져 있으면 여러 번의 메모리 접근이 필요해 느려짐

나. 연산 속도 개선 이유

  • 현대 GPU (특히 A100, H100)는 대부분의 컨볼루션 연산이 channels_last 포맷에 최적화

    • NVIDIA Tensor Core는 내부적으로 NHWC (channels_last) 형태의 데이터 배치를 요구
    • 따라서 입력이 NCHW (channels_first) 포맷일 경우:
      • GPU는 자동으로 NHWC 형태로 데이터를 재배열
      • 또는 별도의 임시 버퍼에 데이터를 복사

    → 이 과정이 Layout Transformation 또는 Memory Reordering

    • 재배열 과정의 문제점

      항목 결과
      데이터 복사 비용 추가 메모리 사용 + 복사 연산 (bandwidth 소모)
      캐시 효율 저하 GPU의 연속적인 메모리 접근 불가
      Tensor Core 활용 불가 최적화된 kernel을 사용할 수 없음

다. 메모리 대역폭 효율

  • GPU는 여러 스레드(쓰레드 워프)가 동시에 인접한 메모리 주소에 접근할 때 가장 빠름 → 이를 Coalesced Memory Access라고 부름

  • 메모리에서 데이터를 읽어올 때, 연속적인 데이터한 번에 묶어서 읽기 가능 → 대역폭 효율이 높아짐

  • NCHW vs NHWC 비교 (픽셀 중심)

    포맷 메모리상 배치 예시 메모리 읽기
    NCHW [R R R R …][G G G G …][B B B B …] 픽셀 하나당 여기저기서 따옴
    NHWC [[R,G,B], [R,G,B], [R,G,B], …] (픽셀 단위) 픽셀 하나 = 연속 메모리
    • NCHW에서는 같은 위치의 픽셀이라도 R, G, B 값을 서로 다른 메모리 영역에서 찾아야 함
    • NHWC에서는 픽셀 하나의 모든 정보(RGB) 가 연속 → 한 번에 읽기 가능

    → 즉, 한 번의 메모리 호출으로 한 픽셀의 R, G, B 값을 모두 읽을 수 있는 구조가 GPU 성능에 최적

  • GPU는 64바이트, 128바이트 단위로 데이터를 한 번에 읽음

  • NHWC에서는 픽셀 단위 데이터가 이 블록에 딱 맞게 배열 → 불필요한 메모리 읽기가 줄어듦

  • **캐시 적중률(Cache hit rate)**도 상승 → 속도 향상

라. 데이터 재사용 효율(Cache Locality)

  • 컨볼루션 연산에서는 같은 입력 텐서의 데이터가 여러 번 재사용됨 (예: 커널 슬라이딩)

  • 데이터 재사용캐시(Local cache, L2, shared memory) 에 데이터가 얼마나 오래 남아있는가에 달려있음

  • NCHW vs NHWC 재사용 관점

    포맷 재사용 효율성
    NCHW 픽셀마다 R → G → B 가져오기 → 캐시 오염 ↑
    NHWC 픽셀당 R,G,B 함께 캐시 → 다음 슬라이딩에서도 바로 사용 가능
    • NHWC에서는 공간적으로 인접한 데이터들이 함께 캐시에 들어가서, 컨볼루션 슬라이딩 시 다음 계산에 바로 재사용 가능 → 메모리 대역폭 부담 감소

마. 요약

항목 Channels First (NCHW) Channels Last (NHWC)
메모리 읽기 방식 분산 → Coalescing 안 됨 연속 → Coalescing 가능
대역폭 효율 낮음 높음
데이터 재사용 캐시 오염 많음 캐시 적중률 높음 → 재사용 효율 ↑
GPU 연산 최적화 비효율적 Tensor Core 최적화 대응

d. 적용 대상

적용 가능 적용 불가
✅ Conv2D 기반 Vision 모델 (CNN, ViT, Diffusion) ❌ NLP, LLM, Transformer (비비전)
✅ 이미지 입력 모델 ❌ 순차 데이터, 텍스트 데이터
  • Conv2D와 Pixel-wise 연산은 channels_last 포맷 적용 시 성능 향상
  • 비전 모델에서 속도가 최대 30~50% 향상하는 사례도 확인됨

e. 요약

비교 항목 Channels First (NCHW) Channels Last (NHWC)
메모리 배치 순서 Batch → Channel → H → W Batch → H → W → Channel
GPU 접근 속도 느림 빠름
적용 대상 CPU, 디버깅 중심 GPU, 고속 연산 중심
속도 효과 (A100) 보통 최대 수십 % ↑
  • A100 환경에서는 torch.channels_last 설정 시 특히 Conv 연산이 최대 30~50%까지 빨라질 수 있음
  • 정확도 영향 없음 → 속도만 개선

5. cudnn.benchmark = True

torch.backends.cudnn.benchmark = True

a. 개념

  • GPU에는 같은 연산(예: Conv2D) 을 처리하는 여러 가지 알고리즘이 존재
    • 적용 연산(링크)
      • Conv2D: 다양한 구현 알고리즘(GEMM, Winograd, FFT 등)에 대해 벤치마크 수행
      • Pooling 등 cuDNN 지원 연산 일부도 빠른 알고리즘 선택 대상
    • 적용되지 않는 연산
      • Linear, Attention, LayerNorm, Softmax 등은 cuDNN이 아닌 일반 CUDA/CPU 커널로 수행되므로 영향 없음
      • 결과적으로 Transformer 기반 연산(ViT, Self-Attention)에는 효과가 없음
  • cudnn.benchmark = True를 설정하면 PyTorch가 현재 입력 크기와 데이터에 맞는 가장 빠른 알고리즘을 자동으로 찾아 사용
  • 입력 데이터가 반복적인 형태일 때 특히 효과적

b. 용어 풀이

용어 의미
cuDNN NVIDIA의 딥러닝 연산 라이브러리. Conv2D, ReLU, BatchNorm 등 GPU에서 빠르게 처리
benchmark 여러 알고리즘의 속도를 미리 측정하고 가장 빠른 알고리즘을 선택하는 과정
알고리즘 Conv 연산을 수행하는 다양한 방식 (GEMM, FFT, Winograd 등)

c. 원리

  • Conv2D는 커널 크기, 입력 크기, 배치 크기, 패딩 여부에 따라 최적의 알고리즘이 달라짐
  • benchmark = True를 켜면 PyTorch는 첫 실행 시 다양한 알고리즘을 실제로 실행하여 속도를 측정 → 이후 동일한 입력 크기에는 최적화된 알고리즘만 사용
  • 결과적으로:
    • 반복적인 입력 → ✅ 속도 향상
    • 매번 입력 크기가 다르면 → ❌ 오히려 비효율 (매번 새로 선택)

d. 적용 대상과 주의사항

상황 설정 추천
입력 크기가 항상 동일 (예: 이미지 고정 크기) benchmark = True
입력 크기가 매번 달라짐 (예: 다양한 크기의 입력) benchmark = False

[ 적용 대상 ]

  • Diffusion, ViT, CNN 등 고정된 배치/입력 크기에서 성능 향상
  • NLP, 유동적인 입력 길이에서는 오히려 느려질 수 있음

[ 주의 사항 ]

  • 최초 실행 시에는 여러 알고리즘의 성능을 실제로 측정하기 때문에 오히려 첫 추론은 느릴 수 있음
    • 두 번째 실행부터 캐시된 최적 알고리즘이 사용되어 속도가 빨라짐 → 반드시 1회 실행 이후부터 효과
  • 입력 크기가 자주 변하면 매번 새로 측정 → 효율 저하 (반복성이 없는 입력에는 부적합)
  • 모델 추론이 반복적이고 입력 크기가 고정인 경우에만 확실한 효과

⇒ 따라서 서비스 배포 전 또는 서버 기동 시 미리 한 번 dummy inference(warm-up)을 돌려두면 안정적

e. 그럼 FLUX에는?

가. 전체 구조

  • FLUX.1 Kontext는 Stable Diffusion 기반의 Flow Matching 모델로, UNet 구조를 중심으로 구성됨
  • 구조 요소:
    • VAE Encoder / Decoder: Conv2D 레이어로 이미지 ⇄ Latent 변환
    • UNet 본체:
      • Downsampling & Upsampling 섹션: 대부분 Conv2D
      • 중간(Mid) 및 CrossAttention 블록 포함 (Attention은 일부)

다. Conv2D 레이어 비중

  • UNet 구조의 핵심은 Residual/Skip합성 Layer들로 구성된 Conv2D 연속 처리

  • 예:

    Conv2D → GroupNorm → Swish → Conv2D → GroupNorm → Swish
    ↓
    Cross-Attention (잠깐 포함)
    ↓
    Conv2D → Upsample...
    
  • 즉, 모델 전체의 80–90% 이상이 Conv2D 실행에 해당하며, Attention은 일부 블록에만 사용

f. 요약 (보강 버전)

항목 설정 없음 (False) 설정함 (True)
알고리즘 선택 방식 매번 일반적인 기본 알고리즘 사용 가장 빠른 알고리즘 자동 선택 (캐시)
속도 보통 반복 시 빨라짐
추천 상황 입력 크기 자주 바뀔 때 입력 크기 고정일 때
FLUX 적용 가능성 ❌ 성능 개선 없음 ✅ 고정된 Latent 크기 + Conv2D 중심 → 효과적
  • FLUX.1-Kontext-dev는 고정된 Latent Space와 Conv2D 기반 UNet 구조를 사용 → cudnn.benchmark = True 설정이 반복 실행 시 성능 최적화에 유리

6. CUDA Graph

import torch.cuda
g = torch.cuda.CUDAGraph()

a. 개념

  • GPU에서 자주 반복하는 연산들을 **한 번 실행 계획(그래프)으로 기록(record)**하고, 이후에는 그 계획을 그대로 재실행(replay) 하여 속도를 높이는 방법
  • 컴파일/그래프 기반 런타임 엔진과는 달리, 호출 단계의 오버헤드를 줄이는 것이 핵심

b. 용어 풀이

용어 의미
CUDA Graph GPU 커널 실행 흐름을 기록하고 재사용하는 기능
record-replay 한 번 실행 계획을 저장한 뒤 동일한 연산을 반복 실행

c. 원리

  • 일반적으로 GPU는:
    • 매번 연산 실행 → 매번 Launch 호출 → 매번 GPU와 CPU 동기화
  • CUDA Graph는:
    • 연산 계획을 한 번만 기록(record)
    • 이후에는 CPU 개입 없이 바로 GPU에서 실행(replay)
  • 결과:
    • Launch 오버헤드 감소 → 지연 시간(latency) 감소
    • 특히 반복적이고 빠른 inference에서 효과

d. 적용 대상

상황 효과
모델의 입력/출력 크기 고정 ✅ 매우 효과적
Batch Inference 반복 ✅ Latency 감소
입력 크기 유동적 ❌ 적용 불가
  • 예시: Diffusion Inference, LLM Serve, ViT 배치 추론

e. 주의사항

  • 입력 크기가 고정되어야 함 (그래프는 입력 크기까지 포함해서 기록됨)
  • torch.compile병행 사용 가능 → 두 방법은 서로 다른 최적화 계층 (Graph는 실행, compile은 계산 최적화)

f. 요약

항목 CUDA Graph 사용 안 함 CUDA Graph 사용
실행 계획 매번 새로 호출 한번 기록 후 재사용
지연 시간 (latency) 높음 낮음
적용 조건 입력 크기 유동적 입력 크기 고정
  • 특히 FLUXDiffusion 모델처럼 고정된 크기와 반복적인 inference에는 CUDA Graph가 매우 효과적

g. torch.compile / SDPA / CUDA Graph — 연관성과 차이점

항목 torch.compile SDPA (FlashAttention) CUDA Graph
적용 범위 모델 전체 (Python → 그래프 변환) 어텐션 연산 단위 (커널 최적화) 실행 단계 (GPU 실행 흐름)
최적화 방식 연산 결합, Python 오버헤드 제거 메모리/속도 최적화 (커널 개선) 실행 계획 고정, 호출 오버헤드 제거
상태 모델 컴파일 전 모델 컴파일 또는 실행 중 모델 컴파일 후, 실행 중
연관성 ✅ torch.compile과 병행 사용 가능 ✅ 내부적으로 SDPA가 포함될 수 있음 ✅ torch.compile과 병행 사용 가능