본문 바로가기

SASS_Probe

수학적 연산의 SASS Lowering 과 불변성 기반 GPU 최적화 탐구

1. 연구 배경

딥러닝 모델은 일반적으로 고수준 레이어의 조합으로 표현된다. 

프레임워크 수준에서 각 레이어는 서로 다른 수학적 의미와 인터페이스를 가지며, 개발자 역시 이러한 연산자 단위로 모델을 이해한다.

그러나 GPU 가 실제로 실행하는 것은 고수준 연산자가 아니다. 

고수준 연산은 컴파일 과정에서 CUDA 코드,  PTX 와 같은 중간 표현을 거쳐 최종적으로 GPU 가 실행할 수 있는 SASS 명령어와 메모리 접근 구조로 변환된다.

전체 변환 과정을 단순화하면

  • 수학적 연산
    • 딥러닝 레이어 또는 연산자
    • CUDA 커널
    • PTX
    • SASS
    • GPU 실행

레이어 수준에서는 서로 완전히 달라 보이는 연산도 SASS 수준에서는 제한된 종류의 실행 패턴으로 표현된다.

대표적으로 다음 요소가 반복

  • global memory load, store
  • register 기반 산술 연산
  • FMA 또는 FFMA 연산
  • 비교 및 조건 선택
  • exponential, reciprocal 과 같은 특수 함수
  • warp shuffle 을 이용한 데이터 교환
  • shared memory 를 이용한 block 내 데이터 공유
  • reduction 
  • synchronization 
  • 주소 계산과 반복문 제어

핵심은 단순히 생성된 SASS 를 읽거나 명령어 수를 세는 것이 아니라

수학적 연산의 의미가 lowering 이후 어떤 구조로 남는지, 서로 다른 연산이 어떤 공통 실행 모티프로 변환되는지, 그리고 수학적 의미를 보존하면서 해당 구조를 어떻게 재배치할 수 있는지를 탐구하는 것이 목적

 

2. 연구 문제 재정의

초기 실험은 elementwsie 연산, FMA, softmax, online softmax, attention porbability materialization 등으로 확장

이 흐름은 자연스럽게 Flashattention 의 핵심 구조와 연결된다.

예를 들어 다음과 같은 변화가 관찰된다.

  • attention score 계산
    • softmax probability 생성
    • probability global memory 저장
    • probability 다시 load
    • value 와 weighted sum 계산

이를 다음과 같이 변경할 수 있다.

  • attention score 계산
    • softmax 상태 유지
    • probability 를 외부에 저장하지 않음
    • 즉시 value 와 결합

이 구조는 FlashAttention 의 핵심 아이디어와 밀접한 관련

그러나 연구 방향을 여기에만 집중하면 전체 연구가 다음과 같이 축소될 위험이 있다.

"기존 attention 구현을 Flashattention 형태로 변환하는 과정의 재현"

이런 경우 연구 질문은 오히려 특정 알고리즘에 종속될 수 있다.

따라서 연구의 중심 질문은 attention 자체가 아니라 다음과 같이 설정할 필요가 있다.

수학적으로 정의된 연산은 GPU 명령 수준에서 어떤 실행 구조로 lowering 되며, lowering 이후에도 보존되는 수학적, 구조적 불변성은 무엇인가

이후 질문

이러한 불변성을 유지하면서 SASS 수준의 데이터 흐름, 메모리 위치, 연산 순서와 중간값 저장 방식을 어떻게 변경할 수 있는가?

 

3. 연구의 핵심 관점

3.1 레이어가 아닌 lowering 된 실행 구졸르 분석 단위로 삼는다

일반적인 딥러닝 최적화는 레이어 단위로 분류된다.

  • Softmax 최적화
  • LayerNorm 최적화
  • GEMM
  • Attention
  • Activation

그러나 SASS 수준에서 이들은 더 작은 실행 패턴으로 분해된다.

Softmax 의 경우

  • input load
  • max reduction
  • intput - max
  • exponential
  • sum reudction
  • normalization
  • output store

LayerNorm 은 

  • input load
  • mean reduction
  • variance reduction
  • normalization
  • scale
  • bias
  • output store

두 연산은 수학적으로 다르지만 GPU 실행 관점에서는 공통 구조를 가진다.

  • load
  • local transform
  • reduction
  • reduction 결과 공유
  • elementwise transform
  • store

따라서 레이어 자체보다 다음과 같은 실행 모티프를 분석 단위로 사용한다

  • elementwise transform
  • affine transform
  • reduction
  • reduction 결과 broadcast
  • producer-consumer 관계
  • 중간 결과 materialization
  • register 또는 shared memory reuse
  • tile 단위 working set 유지
  • synchronization
  • memory layout 변환

이렇게 분석 단위를 바꾸면 서로 다른 레이어에서 반복되는 최적화 가능성을 하나의 원리로 설명할 수 있다.

 

3.2 SASS 를 단순한 컴파일 결과가 아닌 실행 의미의 표현으로 본다

SASS 에는 고수준 수학 연산의 이름이 직접 남아 있지  않다.

이러한 명령이 존재하지 않음

  • SOFTMAX
  • LAYERNORM
  • ATTENTION
  • RMSNORM

대신 수학적 의미는 명령어들의 관계 속에 간접적으로 나타난다.

  •  어떤 값이 반복 사용되는가
  • 어떤 값들이 하나의 값으로 집계되는가
  • 집계 결과가 몇 개의 출력에 공유된느가
  • 중간 결과가 global memory 에 저장되는가
  • 동일한 주소 또는 동일한 regsiter 값이 반복 사용되는가
  • thread 사이에 어떤 값이 전달되는가
  • 연산 사이에 어떤 의존성이 존재하는가

SASS 분석에서 중요한 것은 특정 명령어의 존재 여부만이 아니다.

예를 들어 FFMA 명령어가 존재한다는 사실보다 더 중요한 것은 다음과 같다

  • 어떤 값이 곱셈 입력으로 사용되는가
  • 누산 값이 어떤 의존 사슬을 형성하는가
  • 중간 곱셈 결과가 외부에 노출되는가
  • 동일한 accumulator 가 반복적으로 갱신되는가
  • 최종 결과가 언제 memory 로 store 되는가

즉 SASS 명령어 목록이 아니라, 수학적 연산이 하드웨어 실행 구조로 변환되 데이터 흐름 그래프로 해석해야 한다. 

 

4. 수학적 의미가 SASS 에 남는 방식

4.1 Elementwise 독립성

다음과 같은 연산을 생각할 수 있다.

  • y_i = f(x_i)

각 출력은 대응하는 입력에만 의존한다. 출력 원소 사이에는 의존성이 없다 

이 성질은 lowering 이후에도 사라지지 않는다

SASS 에서는 대체로 다음과 같은 구조로 나타난다

  • thread index 계싼
  • x_i load
  • local arithmetic
  • y_i store

각 thread 는 다음 thread 의 결과를 필요로 하지 않는다. 따라서 일반적으로 다음 명령이 필요하지 않다. 

  • warp shuffle
  • shared memory 를 통한 값 교환
  • block synchronization
  • cross-thread reduction

이는 elementwise 연산의 중요한 구조적 불변성이다. 

ReLU, clamp, sigmoid, 단순 scale, bias add 와 같은 연ㅅ낭느 모두 세부 명령은 다르지만 thread 별 독립적인 실행 구졸르 가진다.

이러한 독립성은 다음과 같은 최적화 가능성을 만든다

  • 여러 elementwise 연산의 fusion
  • GEMM 또는 convolution epilogue 와 결합
  • 중간 텐서 제거
  • vectorized load 와 store
  • 동일 thread 내 register 전달

 

4.2 Affine 연산과 FMA 의존 사슬

댜음 연산의 생각

  • y = ax + b

컴파일러는 이를 곱셈과 덧셈으로 분리할 수 있지만 GPU 에서는 일반적으로 fused multiply-add 형태로 lowering 할 수 있다.

개념적으로 다음과 같은 구조

  • load x
  • load a
  • load b
  • FFMA a, x, b
  • store y

여기서 중요한 것은 단순히 FFMA 가 생성된다는 사실이 아니다. 

수학적으로 곱셈 결과는 독립적인 출력이 아님, 최종 결과를 계산하기 위한 중간값일 뿐이다. 

따라서 중간 곱셈 결과를 별도로 register 에 오래 유지하거나 memory 에 저장할 필요가 없다. 

이 구조는 다음과 같은 다양한 연산에서 반복된다.

  • bias add
  • normalization 의 scale 과 shift
  • GEMM epilogue
  • convolution epilouge
  • residual scaling
  • polynomial approximation
  • 누산 기반 dot product

즉 affine 연산은 특정 레이어가 아니라, SASS 에서 반복되는 하나의 공통 의존성 패턴이다.

 

4.3 Reduction 의미

다음 연산을 생각할 수 있다.

  • s = SUM_i x_i

elementwise 연산과 달리 출력 s 는 여러 입력에 동시에 의존한다.

따라서 lowering 이후에는 어떤 형태로든 여러 thread 가 가진 값을 결합하는 구조가 필요하다.

예를 들어 warp 단위 reductino 은 개념적으로 다음과 같이 나타날 수 있다.

  • local accumulation
  • shuffle
  • add
  • shuffle
  • add
  • shuffle
  • add

block 단위 reductino 에서는 다음 구조가 추가될 수 있다.

  • warp-local reduction
  • shared memory store
  • blocck synchronization
  • shared memory load
  • final reduction

실제 SASS 에서는 warp shuffle, shared memory load/store, barrier, add 명령의 조합으로 나타난다.

여러 입력 원소가 하나의 집계갑으로 수렴한다.

합계, 최댓값, 평균, 분산, norm 은 서로 다른 연산이지만 모두 이와 같은 집계 구조를 가진다. 

따라서 SASS 수준에서는 reduction 연산을 다음 요소로 분석할 수 있다.

  • reduction 대상 원소 수
  • thread 당 local accumulation 범위
  • warp-level reduction 여부
  • block-level reduction 여부
  • shuffl 단계 수
  • shared memory 사용량
  • synchronization 횟수
  • reduction 결과의 저장 위치
  • reduction 결과가 소비되는 방식

 

4.4 Broadcast 와 공통 의존성

다음 연산을 생각할 수 있다.

  • y_i = x_i / sum_j x_j

모든 출력은 하나의 공통 reduction 결과에 의존한다

따라서 lowering 된 실행 구조에는 다음 단계가 존재해야 한다.

  • sum 계산
  • sum 결과 공유
  • 각 thread 가 자신의 x[i] 와 sum 을 결합

구현에 따라 공유 방법은 달라질 수 있다.

  • 같은 warp 안에서 shuffle
  • shared memory 에 reduction 결과 저장
  • register broadcast
  • 별도 kernel 에서 global memory 에 reduction 결과 저장
  • scaler load 를 통한 재사용

모든 출력이 동일한 집계값을 소비한다는 의존성은 유지된다.

이것이 Softmax, LayerNorm, RMSNorm 등에서 공통적으로 발견된다.

 

5. Lowering 이후에도 보존되는 불변성

5.1 데이터 의존성 불변성

가장 기본적인 불변성은 출력이 어떤 입력에 의존하는가이다.

컴파일러는 명령 순서를 바꾸고, 연산을 결합하고, register 를 재배치할 수 있다. 

그럴나 프로그램의 의미를 유지하려면 최종 출력에 필요한 의존성은 보존되어야 한다.

예를 들어 다음 연산에서

  • y_i = x_i - max_j x_j

각 출력은 자신의 입력 뿐 아니라 전체 입력에서 계산된 최댓값에도 의존한다.

따라서 구현이 naive reduction 이든, warp shuffle 이든, shared memory 기반이든, 다음 구조는 반드시 존재한다.

  • 전체 범위 max 계산
  • max 결과 공유
  • 개별 원소에서 max 제거

명령 배치는 달라질 수 있지만 데이터 의존성은 유지된다.

이 불변성은 SASS 를 통해 원래 연산의 의미를 추론한느 주요 단서가 된다.

 

5.2 Reduction domain 불변성

Reduction 에서는 어떤 연산을 수행하는가뿐 아니라, 어떠 ㄴ원소 집합에 대해 수행하는가가 중요핟.

예를 들어 다음 연산들은 모두 합계를 포함한다.

  • s_1 = sum_i i
  • s_2 = sum_c (n,c)
  • s_3 = sum q_k k_k

수학적으로 각각 다른 축을 reduction 하지만, SASS 에서는 텐서 축의 이름이 직접 나타나지 않는다. 대신 다음 요소를 통해 reduction domain을 추론할 수 있다.

  • loop 반복 횟수
  • thread index 증가 폭
  • memory address stride
  • block dimension
  • warp shuflle 범위
  • shared memory index
  • thread 당 처리 원소 수
  • vectorized load 폭

고수준의 텐서 축은 SASS 에서 주소 계산과 thread 협업 범위로 lowering 된다.

이 관계를 분석하면 고수준 shape 정보가 실제 하드웨어 execution geometry 로 어떻게 변환되는지 이해 가능

 

5.3 외부 관측 가능성 불변성

중간값이 최종 출력 이외의 위치에서 사용되는지를 구분하는 것은 중요

다음 구조를 생각

  • A 연산
  • 중간값 T 생성
  • T 를 global memory 에 저장
  • B 연산에서 T 를 다시 load
  • 최종 출력 Y 생성

만약 중간값 T 가 B 연산 이외의 위치에서 사용되지 않느다면, T 를 global memory 에 저장하는 것은 수학적으로 필수적이지 않다. 

다음과 같이 변경할 수 있다.

  • A 연산
  • T 를 register 또는 shared memory 에 유지
  • 즉시 B 연산 수행
  • 최종 출력 Y 저장

이때 외부에서 관측 가능한 최종 출력이 동일하다면 프로그램의 의미는 유지된다.

이 원리는 다음 최적화에서 공통적으로 나타난다.

  • bias 와 activation fusion
  • GEMM epilogue fusino
  • normalization 과 affine transform fusion
  • Softmax 와 weighted sum fusion
  • Attention probability materialization 제거
  • transpose 와 consumer 연산 결합
  • residual add 와 activation 결합

따라서 materialization 제거 가능성은 특정 attention 알고리즘의 특징이 아니라, 중간값의 외부 관측 가능성에 관한 일반적인 원리다 .

 

5.4 재사용 불변성

중간값이 여러 연산에 사용되는 경우 lowering 이후에도 그 재사용 관계가 존재한다.

구현 중 다음 하나를 선택 가능

  • 매번 global memory 에서 다시 load
  • shared memory 에 저장
  • 특정 register 에 유지
  • warp shuffle 오 broadcast
  • 별도 kernel 의 출력으로 materialize

선택은 달라질 수 있지만, 재사용 관계는 유지

SASS 에서는 다음의 형태로 재사용 가능성을 관찰 가능

  • 동일 register source 의 반복 사용
  • 동일 주소 반복 load
  • loop 밖에서 load 한 값을 loop 내부에서 재사용
  • shared memory 의 동일 위치에 대한 반복 접근
  • constant 또는 uniform load
  • broadcast 형태의 warp communication

 

5.5 연산 순서의 자유도

모든 데이터 의존성이 연산 순서를 완전히 고정하는 것이 아니다.

독립적인 연산은 순서를 바꿀 수 있으며, 일부 수학적 성질은 연산 구조를 크게 재배치할 수 있게 한다.

Online Softmax 는 기존 실행 순서를 그대로 유지하는 대신, 동일한 수학적 결과를 유지할 수 있느 상태 갱신 규칙을 정의한다.

이는 단순한 이론적 성질을 넘어 실제 메모리 이동과 실행 순서를 변경하는 근거가 된다는 것을 보여준다. 

 

6. 불변성의 종류와 수치적 동등성

SASS 수준의 최적화를 다룰 때는 동일한 결과의 의미를 구분해야 한다.

6.1 수학적 동등성

실수 연산을 기준으로 두 식이 동일한 결과를 생성하는 경우,

  • 실수 수학에서는 교환, 결합 법칙이 성립한다.

그러나 실제 GPU 는 유한 정밀도 부동 소수점 연산을 사용한다.

 

6.2 부동소수점 수치 동등성

연산 순서에 따라 반올림 위치가 달라진다. 

bitwise 하게 같지 않다. ( 교환, 결합 법칙에 따라 바꿨을 때 )

Reduction tree 를 변경하면 수치 오차가 달라질 수 있다.

또한 다음 최적화도 결과의 하위 비트에 영향을 줄 수 있다.

  • FMA 결합
  • reciprocal 근사
  • exponential 근사
  • mixed precision
  • accumulation precision 변경
  • instruction reordering
  • tensor core 사용

따라서 실험에서는 결과의 동등성을 다음과 같이 구분해야 한다.

  • Bitwise equivalent
  • Numerically equivalent within tolerance
  • Mathematically equivalent
  • Semantically equivalent

어떤 종류이ㅡ 동등성을 목표로 하는지 (아마 4단계, Semantically )

 

6.3 구조적 불변성과 수치적 불변성의 분리

구조적 불변성은 비교적 강하게 보존된다.

  • 출력 shape
  • 출력이 의존하는 입력 집합
  • reduction 대상 원소 집합
  • 외부에서 관측 가능한 값
  • elementwise 독립성
  • producter-consumer 관계

반면 수치적 불변성은 구현에 따라 근사적으로만 유지될 수 있다.

  • 덧셈 순서
  • reduction tree
  • exponential 근사
  • reciprocal precision
  • FMA 적용 여부
  • 데이터 타입

SASS 수준 최적화에서는 구조적 의미의 보존과 수치적 오차의 변화를 별도 분석해야 한다.

 

7. 연구의 기대 결과

본 연구가 목표로 하는 결과는 특정 레이어의 최적 구현 하나가 아니다.

최종적으로 다음과 같은 일반 규칙을 도출하는 것이 중요하다.

규칙 1

중간 결과가 외부에서 관측되지 않고 producer와 consumer가 동일한 실행 범위에 존재한다면, global memory materialization을 제거할 가능성이 있다.

규칙 2

여러 출력이 하나의 reduction 결과를 공유한다면, reduction 결과의 저장 위치와 broadcast 방식이 주요 최적화 대상이 된다.

규칙 3

동일한 값이 여러 산술 연산에 반복 사용된다면, global reload보다 register 또는 shared memory 유지가 유리할 수 있다.

규칙 4

수학적 불변성이 연산 순서 변경을 허용한다면, pass 수와 중간 상태를 줄이는 새로운 lowering 구조를 만들 수 있다.

규칙 5

Fusion으로 memory traffic이 감소하더라도 register pressure, occupancy 저하, synchronization 증가가 전체 성능을 악화시킬 수 있다.

규칙 6

서로 다른 레이어라도 동일한 SASS 실행 모티프를 공유한다면, 최적화 역시 레이어 단위가 아니라 실행 모티프 단위로 일반화할 수 있다.


8. 연구 방향의 최종 정의

본 연구는 딥러닝 레이어를 고수준 연산자 이름으로만 분석하지 않는다.

각 연산이 CUDA와 PTX를 거쳐 SASS 명령, 데이터 의존성, memory flow, thread communication 구조로 lowering되는 과정을 관찰한다.

그리고 lowering 이후에도 보존되는 다음 요소를 찾는다.

  • 입력과 출력의 의존 관계
  • reduction domain
  • elementwise 독립성
  • 공통값의 재사용 관계
  • 외부에서 관측 가능한 중간값
  • 연산 순서를 변경할 수 있는 수학적 자유도
  • 수치적 오차를 허용하는 범위

이를 통해 어떤 실행 구조가 수학적 의미에 의해 반드시 필요한지, 어떤 구조가 단순히 현재 구현 또는 컴파일 결과에 의해 선택된 것인지 구분한다.

전자는 보존해야 할 제약이며, 후자는 최적화 가능성이 된다.

연구 방향을 한 문장으로 정리하면 다음과 같다.

본 연구는 수학적 연산이 GPU의 SASS 명령과 데이터 흐름으로 lowering되는 과정을 분석하고, lowering 이후에도 보존되는 수학적·구조적 불변성을 식별하여, 의미를 유지하면서 실행 구조를 재배치할 수 있는 최적화 가능성을 탐구한다.

보다 구체적으로는 다음과 같이 표현할 수 있다.

레이어 이름이나 프레임워크 연산자를 최적화 단위로 삼는 대신, SASS 수준에서 나타나는 데이터 의존성, reduction, broadcast, 재사용, 중간값 materialization, synchronization과 같은 실행 모티프를 최적화 단위로 재정의한다.

결국 이 연구에서 탐구하려는 핵심은 다음 질문이다.

어떤 수학적 성질이 특정 SASS 구조를 필수적으로 만들며, 어떤 수학적 자유도가 그 SASS 구조의 변형과 최적화를 허용하는가?

현재 진행된 Softmax와 Attention 실험은 이 질문에 답하기 위한 첫 번째 사례다. 향후 LayerNorm, RMSNorm, GEMM epilogue, residual, layout transformation으로 실험을 확장하면, Attention에 종속되지 않는 보다 일반적인 GPU lowering 및 최적화 연구로 발전시킬 수 있다.