Memory Optimization Pattern Catalog

AICF Memory Domain Design Philosophy

GPU 기반 AI 시스템에서 성능 병목은 단순히 연산량에 결정되지 않고, 실제 실행 환경에서는 대부분의 지연이 데이터 이동, HBM 접근에서 발생한다.

따라서 메모리 최적화의 핵심 질문은 다음과 같다.

어떤 계산을 수행할 것인가가 아니라, 어떤 데이터를 저장하고, 어떤 데이터를 이동시키지 않을 것인가

이러한 관점을 바탕으로 메모리 최적화를 구조적 패턴 structural patterns 으로 분류한다.

특정 연산에 종속된 트릭이 아니라, 컴파일러가 수학적 성질을 인식하고 자동 적응할 수 있는 일반적 구조로 정의된다.

 

The Three Core Questions

AICF Memory Optimization 은 세 가지 질문으로 요약된다.

1. 어떤 값은 저장하지 않아도 되는가?

어떤 intermediate 는 필요할 때 다시 계산 하는 것이 더 저렴하다

이 질문은 다음 패턴으로 이어진다.

Rematerializable Intermediate

 

2. 어떤 reduction 은 streaming 으로 처리할 수 있는가?

일부 reeduction 은 전체 데이터를 materialize 하지 않고도 incremental state update 를 통해 계산할 수 있다.

이 질문은 다음 두 패턴으로 이어진다.

Online Reducible Norm
Streaming Weighted Reduction

 

3. 어떤 연산은 tile 내부에서 닫힐 수 있는가?

데이터가 shared memory / registers 에 있는 동안 가능한 많은 연산을 수행하면 HBM traffic 을 크게 줄일 수 있다.

이 질문은 다음 패턴으로 이어진다.

Tile-Compatible Compute

 

The Four Core Patterns

메모리 최적화를 네 가지 구조적 패턴으로 정의한다.

• Online Reducible Norm 
  • 통계 계산을 multi-pass -> streaming single-pass 로 변환
• Streaming Weighted Reduction 
  • 정규화된 weighted reduction 을 streaming 구조로 변환
• Rematerializable Intermediate
  • intermediate tensor 를 ㅓㅈ장 대신 재계싼
• Tile-Compatible Compute
  • 연산을 tile 내부에 닫히도록 스케줄링

이 패턴들을 서로 독립적인 것이 아니라 동시에 결합될 수 있다.

예를 들어 FlashAttention 은 다음 패턴들의 결합이다.

Tile-Compatible Compute
+
Streaming Weighted Reduction
+
Rematerialization (partial)

 

Pattern 1

Online Reducible Norm

일반적인 통계 계산은 여러 번의 데이터 순회를 요구한다.

예

mean = sum(x)/N
var  = sum((x-mean)^2)/N

이 방식은 multi-pass reduction 이다.

하지만 Welford 알고리즘을 사용하면 평균과 분산을 single-pass streaming 방식으로 계산할 수 있다.

핵심 성질

mergeable statistics state

즉 partial statistics 를 합칠 수 있다.

이  구조는 다음 장점을 가진다.

• single-pass reduction
• streaming statistics
• distributed merge 가능

대표 적용

LayerNorm
BatchNorm
dataset sttatistics

 

Pattern 2

Streaming Weighted Reduction

Attention 은 다음 형태를 가진다.

softmax(QK^T)V

일반 구현에서는 다음 intermediate 가 materialize 된다.

score matrix
probability matrix

이것은 매우 큰 HBM traffic 을 발생시킨다.

하지만 online softmax 의 rescaling invariant 성질을 이용하면 softmax-weighted reduction 을 streaming 방식으로 계산할 수 있다.

핵심 상태

running max
normalized denominator
weighted accumulator

이 구조 덕분에 다음이 가능해진다

score matrix materialization 제거
probability matrix materialization 제거

대표 적용

FlashAttention
memory-efficient attention

 

Pattern 3

Rematerializable Intermediate

많은 intermediate tensor 는 다음 형태를 가진다.

Y = f(X)

이 경우 Y 를 저장하지 않고 필요할 때 다시 계산할 수 있다.

이 전략을 rematerialization 이라고 한다.

핵심 trade-off

memory cost vs compute cost

GPU 에서는 다음이 종종 성ㄹ비

HBM load cost > recompute cost

따라서 intermediate 를 저장하는 대신

재계산하는 것이 더 빠를 수 있다.

대표 적용

activation checkpointing
temporary tetnsor elimination
fused elementwise ops

 

Pattern 4

Tile-Compatible Compute

GPU 성능의 핵심은 data reuse 이다.

shared memory / registers 에 데이터가 있는 동안 가능한 많은 연산을 수행해야 한다.

이를 위해 연산이 다음 조건을 만족해야 한다.

tile 내부 dependency closure
partial accumulation 가능

대표 예

GEMM
convolution
attention

결과

HBM traffic 감소
on-chip reuse 증가

 

Compiler Perspective

이 네 가지 패턴의 중요한 특징은 단순한 구현 트릭이 아니라 컴파일러가 인식 가능한 수학적 구조라는 점이다.

MCIR 에서는 다음과 같은 property 로 표현할 수 있다.

online_reducible_norm
weighted_stsreaming_reduction
rematerializable_intermediate
tile_compatible_compute

컴파일러는 computation graph 에서 이러한 property 를 탐지하고 다음 변환을 수행할 수 있다.

multi-pass reduction
→ streaming reduction

materialized intermediate
→ recompute

naive schedule
→ tiled schedule

즉 메모리 최적화는

low-level kernel trick

이 아니라

graph-level structural rewrite

로 정의된다.

 

Memory Optimization as Structural Transformation

AICF 에서 메모리 최적화는 다음과 같이 정의된다.

memory optimization =
data residency planning
+
intermediate elimination
+
streaming reduction

즉 핵심은

어떤 값은 streaming
어떤 값은 recompute
어떤 값은 tile residency

로 재구성하는 것이다.

이 과정은 execution structure 자체를 변경한다.

 

Summary

AICF Memory Domain 은 다음 네 가지 패턴으로 구성된다.

Online Reducible Norm
Streaming Weighted Reduction
Rematerializable Intermediate
Tile-Compatible Compute

이 패턴들은 다음 질문을 해결한다.

어떤 값은 streaming 가능한가
어떤 값은 저장할 필요가 없는가
어떤 연산ㅇ느 tile 내부에 닫힐 수 있는가

결국 메모리 최적화의 본질은 다음과같다

계산을 줄이는 것이 아니라 불필요한 데이터 이동을 제거하는 것