본문 바로가기

AI Compiler Generator

Computation Structures

개요

개별 연산자를 표면적인 이름이 아니라, 그 아래에 있는 계산 구조의 관점에서 기술하기 위한 속성이다.

중요한 것은 "이 연산자가 softmax 인가, layernorm 인가, reduction 인가"가 아니다.

 

그보다 더 중요한 것은, 해당 연산이 내부적으로 어떤 계산 패턴을 가지는지, 어떤 상태를 유지하는지, 어떤 방식으로 부분 계산을 결합할 수 있는지, 그리고 어떤 종류의 실행 구조로 실현될 수 있는지를 파악하는 것이다.

 

즉 computation_structures 는 연산자를 더 일반화된 계산 구조 단위로 끌어올리는 역할을 한다.

예를 들어

  • reduction_sum 은 reduction, summary_state 구조를 가진다.
  • softmax 는 reduction, rescaled_streaming, summary_state, weighted_aggregation 구조를 가진다.
  • layernorm 는 summary_state, reduction 구조를 가진다.

연산자는 여러 구조의 결합으로 이해될 수 있다.

 

왜 필요한가

전통적으로 최적화는 연산자별로 분리된 사례 중심으로 축적되기 쉽다.

  • reduction 최적화
  • softmax 최적화
  • layernorm 최적화
  • attention 최적화

하지만 실제로는 서로 다른 연산자들 사이에도 공통 계산 구조가 반복된다.

이 공통 구조를 먼저 정의하면 다음이 가능해진다.

  • 서로 다른 연산자 사이의 공통점 발견
  • primitive realization 의 재사용
  • structural composition 규칙 정리
  • runtime path selection 의 일반화
  • kernel ssynthesis 입력의 구조화

단순 분류를 위한 속성이 아니라, 이후 모든 단계의 가장 상위 구조 명세로 작동한다.

 

이 속성이 표현하는 것

  • 이 연산은 여러 입력을 하나의 상태로 축약하는가
  • 전체 상태를 유지해야 하는가, 아니면 작은 요약 상태로 충분한가
  • streaming 방식으로 처리할 수 있는가
  • partial state 를 merge 할 수 있는가
  • nomalization 이나 rescaling 이 구조의 일부인가
  • 가중치를 동반한 aggregation 인가
  • intermediate 를 완전 물질화하지 않고도 진행할 수 있는가

무슨 수식을 계산하느냐보다 어떤 형태의 계산을 수행하느냐를 기술한다.

 

포함될 수 있는 요소의 기준

어떤 항목이 computation_structures 에 들어갈 수 있으려면, 적어도 아래 성격 중 하나를 가져야 한다.

1. 계산의 형태를 설명해야 한다.

  • reduction
  • weighted_aggregation

이건 연산이 어떤 방식으로 입력을 처리하는지를 설명한다.

 

2. 상태 유지 방식과 관련 있어야 한다

  • summary_state
  • rescaled_streaming

계산이 어떤 중간 상태를 유지하며 진행되는지를 설명한다.

 

3. 구현 세부기법보다 상위 개념이어야 한다

  • reduction 은 가능
  • warp_reduction 은 아님
  • shared_memory_staging 도 아님

즉 computation_structures 는 primitive realization 보다 더 추상적인 층이어야 한다.

 

4. 특정 하드웨어 기법이 아니라 구조적 성질이어야 한다.

  • summary_state 는 구조적 성질
  • vectorized_load 는 구조가 아님
  • tensor_core_path 도 구조가 아님

 

현재 사용 중인 computation structure 들

1. reduction

여러 입력 또는 partial result 를 하나의 요약 결과로 축약하는 계산 구조

핵심 특징

  • partial accumulation 이 가능하다
  • partial state 를 merge 할 수 있다.
  • tree 형태의 병렬 실행으로 확장되기 쉽다.
  • 순서가 수치 결과에 영향을 줄 수 있다.

현재 예시

  • reduction_sum
  • softmax
  • layernorm

many-to-one 축약이라는 점에서 summary-state, tree merge, streaming accumulation 같은 다른 구조와도 자주 결합된다.

 

2. summary_state

전체 intermediate 를 그대로 유지하지 않고, 더 작은 요약 상태만 유지하면서 계산을 이어갈 수 있는 구조

핵심 특징

  • compact state 만 유지한다
  • 전체 입력을 항상 다시 참조할 필요가 없다.
  • partial summary 를 merge 할 수 있는 경우가 많다.
  • online / streaming 알고리즘과 잘 연결된다.

흔한 예시

  • partial sum
  • running max
  • running exp-sum
  • mean / variance 통계 상태
  • mergeable statistics

현재 예시

  • reduction_sum
  • softmax
  • layernorm

왜 중요한가

메모리 사용량을 줄이고, streaming realization 을 가능하게 하며, full materialization 없이 계산을 진행할 수 있게 한다.

 

3. rescaled_streaming

입력을 streaming 방식으로 처리핳면서, 수치 안정성 또는 값 범위 관리를 위해 중간 상태를 재스케일링하는 구조.

핵심 특징

  • 모든 intermediate 를 한 번에 저장하지 않는다
  • running state 를 유지한다
  • scale 조정이 계산의 일부로 들어간다
  • numerical stability 가 구조적 핵심이다.  

예시

  • online softmax
  • streaming log-sum-exp
  • flash-attention 계열 일부 구조
  • normalized weighted accumulation

 

4. weighted_aggregation

입력 값들을 어떤 가중치 체계와 함께 누적하거나 조합하는 계산 구조

핵심 특징

  • 값과 weight 가 함께 움직인가
  • normalization 과 결합되기 쉽다
  • weighted sum, expectation, attention accumulation 같은 형태로 확장된다.
  • reduction 보다 richer 한 구조를 가진다.

softmax 자체는 직접 value accumulation 을 하지는 않지만, normalized weight 생성 구조로서 weighted aggregation 계열과 강하게 연결된다.

 

향후 후보

1. elementwise_transform

입력 각 원소를 독립적으로 변환하는 구조

  • ReLU
  • RELU
  • sigmoid

너무 넓은 단위, 구조적으로 유의미한지 고민이 필요

 

2. map_reduce

입력을 먼저 변환한 뒤, 그 결과를 reduction 하는 구조

  • softmax 내부의 exp + sum
  • variance 계산
  • norm 계산
  • loss accumulation

많은 연산이 이 구조로 설명될 수 있다.

변환 후 축약이라는 패턴이 많은 곳에 등장한다.

 

3. mergeable_statistics

부분 통계량을 따로 계산한 뒤 merge 해서 전체 통계량을 복원할 수 있는 구조

  • mean
  • variance
  • Welford 계열 online statistics
  • layernorm / batchnorm 일부 구성

summary_state 보다 더 구체적이고, normalization 계열에서 매우 중요하다.

 

4. streaming_accumulation

입력을 순차 또는 chunk 단위로 처리함녀서 accumulation state 를 유지하느 구조

  • running sum
  • online mean
  • streaming weighted sum
  • blockwise accumulation

진행 방식에 초점

 

5. decomposable

전체 계산을 여러 부분으로 분해하고, 그 결과를 다시 조합해도 의미가 유지되는 구조

  • blockwise reduction
  • tiled normalization
  • chunked attention
  • segmented accumulation

 

6. pairewise_coupled_state

하나의 상태가 아니라 서로 연관된 여러  상태를 함께 유지해야 하는 구조

  • softmax 의 running max + running exp-sum
  • normalized weighted accumulation 의 numerator + denominator
  • variance 계산의 mean + second moment

single summary state 보다 더 풍부한 상태 구조를 설명할 수 있다.

 

7. normalization_structure

축약된 통계나 누적된 상태를 이용해서 출력을 재정규화하는 구조

  • softmax
  • layernorm
  • rmsnorm
  • batchnorm

 

무엇이 computation structure 인가

computation_structures 는 개별 연산자를 구성하는 계산 형태의 상위 구조 명세이며, primitive realization 이 runtime selection 보다 앞선 수준에서 이 연산이 본질적으로 어떤 계산 패턴을 가지는가를 설명하는 속성이다.