SASS 분석을 통한 연산 구조 해석과 최적화

1. 문제의식

딥러닝 모델은 보통 layer 단위로 이해된다. 이 관점은 수학적 의미를 파악하기에느느 적합하지만, 실제 GPU 에서 해당 연산이 ㅇ어떻게 실행되는지를 설명하기에는 충분하지 않다. 

GPU 는 layer 라는 추상 단위를 직접 실행하지 않는다. 실제 실행 단계에서는 각 layer 가 다음과 같은 더 작은 primitive operation 으로 분해된다.

• load
• store
• multiply
• add
• fma
• compare
• max
• sum
• exp
• reciprocal
• reduction
• barrier
• shared memory staging
• register accumulation

SASS 는 이러한 분해가 최종적으로 어떻게 기계어 수준에서 실현되었는지를 보여준다. 

SASS 는 단순한 어셈블리 코드가 아니라, 수학적 연산이 하드웨어 실행 구조로 변역된 최종 흔적으로 볼 수 있다. 

SASS 를 직접 수정하여 최적화를 수행하는 것이 아니라. layer 뒤에 숨어 있는 데이터 이동, 중간값의 생존 방식, memory materialization, fusion 가능성, online update 가능성 등을 읽어내고, 이를 바탕으로 상위 ir 또는 kernel schedule 수준에서 최적화 방향을 설계하는 것이다.

 

2. SASS 를 바라보는 관점

SASS 를 분석 도구로서의 사용, 상위 layer 수준에서는 보이지 않는 실제 실행 구조가 드러난다.

• 어떤 값이 register 에 유지되는가
• 어떤 값이 shared memory 로 내려간느가
• 어떤 값이 global memory 에 저장되는가
• 중간 텐서가 실제로 materialization 되는가
• producer 와 consumer 사이에 불필요한 store/load 가 존재하는가
• reduction 은 warp 단위인가 block 단위인가
• softmax 는 여러 pass 로 수행되난ㄱ, online update 로 수행되는가
• MMA 연산과 scalar 연산은 어떻게 interleaving 되는가

이러한 질문은 layer 의 수학적 정의가 아닌, SASS 를 보면 실제 실행 결과를 기준으로 판단할 수 있다.

 

3. Layer 중심 사고와 실행 구조 중심 사고

일반적인 딥러닝 프레임워크 관점에서는 Attention 을 다음처럼 본다.

Attention = MatMul + Softmax + MatMul

더 구체적으로 보면

S = QK^T

P = softmax(S)

O = PV

이 관점에서는 S 와 P 가 명확한 중간 텐서로 존재한다. 

먼저 attention score matrix S 를 만들고, 그다음 probability matrix P 를 만들고, 마지막으로 output O 를 만든다.

하지만 GPU 실행 구조 관점에서는 Attention 이 다음과 같이 보인다.

• global load Q/K/V
• dot product accumulation
• row-wise max reduction
• exp approximation
• sum reduction
• normalization
• value accumulation
• global store O

이 관점에서는 각 경계가 절대적이지 않다. 

중요한 것은 중간 텐서의 이름이 아니라, 값이 어디에 저장되고 언제 소비되며 얼마나 오래 살아남는가이다.

따라서 최적화의 핵심은 layer 경계를 그대로 유지하는 것이 아니라, 데이터 생존 방식과 이동 방식을 바꾸는 데 있다.

 

4. 일반 Attention 과 FlashAttetnion 의 관점 차이

일반 Attention 은 중간 결과를 완성된 데이터 객체로 물질화하며 계산

반면 FlashAttention 은 S 와 P 를 완성된 중간 텐서로 만들지 않는다. 

대신 block 단위로 score 로 계산하고, 그 score 를 즉시 소비하여 상태 변수를 갱신한다.

구조적으로는 다음과 같다

• Q block 을 고정한다
• K/V block 을 순회한다
• 현재 score block 을 계산한다
• 현재 block 의 max 를 계산한다
• 기존 max 와 비교하여 m 을 갱신한다
• normalization sum l 을 갱신한다
• output accumulator O 를 갱신한다
• score block 은 저장하지 않고 버린다.
• 다음 K/V block 으로 이됭한다.
• 최종 O 만 global memory 에 저장한다.

수학식으로 바꾸는 것이 아니라, 중간 행렬의 존재 방식을 바꾸는 것

텐서를 명시적으로 만들지 않은 채 그 효과만 상태 변수에 흡수하는 streaming state update 방식

 

5. SASS 분석을 통한 최적화 방향 도출

SASS 분석의 목적은 특정 명령어를 다른 명령ㅇ어로 치환하는 것이 아님,

더 중요한 목적은 실제 실행 구조에서 어떤 비효율이 발생하는지 관찰하고, 그로부터 상위 최적화 방향을 도출하는 것이다.

다음 패턴이 등장할 때

• STG intermediate
• LDG same intermediate

이는 중간 결과가 global memory 에 물질화되고 있음을 의미한다. 

이 경우 producer 와 consumer 사이에 fusion 가능성이 있다고 해석 가능

최적화 방향

• kernel fusion 
• epilogue fusion
• intermediate materialization 제거
• register / shared memory reuse
• online update 구조 도입

또 다른 예로, SASS 에서 arithmetic instruction 대비 memory instruction 이 지나치게 많다면 memory bound 가능성이 크다.

최적화 방향은 다음과 같다.

• tiling
• shared memory staging
• vectorized load
• coalesced access
• data reuse 증가
• global memory traffic 감소

softmax 주변에서 다음과 같은 패턴이 분리되어 보인다면

• max reduction
• exp
• sum rduction
• normalization
• store

그리고 그 결과가 다시 다음 kernel 에서 load 된다면,  softmax 가 독립적인 중간 텐서 생성 단계로 실행되고 있을 가능성이 높다. 

이 경우 최적화 방향은 다음과 같다.

• online softmax
• attention fusion
• normalization 과 value aggregation 결합
• FlashAttetnion-like schedule

 

6. Attention 에서 FlashAttention 으로의 최적화 해석

이는 SASS 명령어 몇 개를 바꾸는 문제가 아니다. 이는 전체 실행 스케줄을 바꾸는 문제다

일반 Attention 의 실행 구조

• QK 전체 계싼
• S 저장
• S load 
• softmax 계산
• P 저장
• P load 
• PV 계산
• O 저장

FlashAttention 의 실행 구조

• Q block 고정
• K/V block 순회 
• score block 계산
• row-wise max 갱신
• normalizatoin sum 갱신
• output accumulator 갱신
• 중간 score / probability 미저장
• 최종 O 저장

따라서 최적화의 본질은 다음과 같다.

• 중간 tensor materialization 제거
• block-wise tiling
• online softmax
• QK^T, softmax, PV 의 fusion 
• register / shared memory 중심의 accumulator 구조
• global memory traffic 감소

 

7. 연구 방향 정리

본 연구에서 SASS는 직접적인 최적화 언어가 아니다. SASS는 수학적 연산이 GPU에서 어떤 primitive execution pattern으로 분해되는지를 보여주는 분석 계층이다.

Layer 수준에서는 MatMul, Softmax, Attention처럼 추상화되어 보이는 연산들이 SASS 수준에서는 다음과 같은 실행 요소로 나타난다.

load/store
FMA
comparison
reduction
exponential approximation
reciprocal
accumulator update
memory synchronization
shared memory movement
global memory materialization

이 분해 결과를 통해 다음을 판단할 수 있다.

중간 텐서가 실제로 저장되는가?
값이 register에 유지되는가?
shared memory를 통해 재사용되는가?
global memory traffic이 병목인가?
fusion 가능한 producer-consumer 관계가 있는가?
online update로 바꿀 수 있는가?
tile size나 accumulator 수를 조정해야 하는가?

따라서 SASS 분석은 다음 목적을 가진다.

수학적 layer의 실제 실행 구조 복원
데이터 생존 방식 분석
memory hierarchy 이동 분석
중간 materialization 여부 판단
operator fusion 가능성 판단
streaming update 가능성 판단
상위 IR/kernel schedule 최적화 방향 설계
컴파일 결과 검증

---

8. 결론

SASS를 직접 수정하여 Attention을 FlashAttention으로 바꾸는 것은 현실적으로 매우 복잡하고 비용이 큰 작업이다. 그러나 SASS를 분석하면 수학적 layer가 실제 GPU에서 어떻게 실행되는지, 어떤 중간 데이터가 물질화되는지, 어떤 값이 register나 shared memory에서 재사용되는지, 어떤 연산들이 fusion될 수 있는지 확인할 수 있다.

따라서 SASS는 최적화의 직접적인 구현 언어라기보다, 최적화 설계를 위한 역방향 분석 도구로 사용해야 한다.

핵심은 다음과 같다.

SASS로 최적화를 직접 수행한다

가 아니라,

SASS를 통해 실행 구조를 읽어낸다
그 구조에서 병목과 중간 materialization을 찾는다
상위 IR 또는 kernel schedule 수준에서 최적화 방향을 설계한다
다시 SASS로 그 결과를 검증한다

Attention과 FlashAttention의 차이도 이 관점에서 이해할 수 있다. 일반 Attention은 QKᵀ와 softmax(QKᵀ)를 중간 텐서로 물질화하는 방식이고, FlashAttention은 이 중간 텐서들을 만들지 않은 채 block 단위 score를 즉시 소비하여 max, sum, output accumulator를 온라인으로 갱신하는 방식이다.

즉 FlashAttention은 Attention의 수학식을 바꾼 것이 아니라, Attention을 실행하는 인식 단위를 바꾼 것이다.

일반 Attention이 중간 행렬을 실체화하여 계산을 진행한다면, FlashAttention은 중간 행렬을 사건처럼 취급하고 그 효과만 상태 변수에 흡수한다.

SASS 분석의 가치는 바로 이 차이를 기계어 수준에서 관찰하고, 그 관찰을 상위 최적화 설계로 되돌려 보내는 데 있다.