RedFuser: An Automatic Operator Fusion Framework for Cascaded Reductions on AI Accelerators

이 논문은 AI 가속기에서 캐스케이드 축소 연산 (cascaded reductions) 을 자동으로 단일 루프로 융합하여 최적화된 커널을 생성하는 'RedFuser' 프레임워크를 제안하며, 기존 AI 컴파일러 대비 최대 5 배의 성능 향상을 달성함을 보여줍니다.

핵심

1. 문제: 왜 지금의 AI 는 느릴까요? (혼란스러운 도서관)

AI 가 복잡한 계산을 할 때, 보통 **'축소 (Reduction)'**라는 작업을 많이 합니다. 예를 들어, "이 책장 전체에서 가장 높은 점수를 가진 책 하나만 찾아라"거나 "모든 책의 페이지 수를 다 더해서 합계를 구해라" 같은 작업이죠.

지금까지의 AI 프로그램 (컴파일러) 은 이 작업을 할 때 다음과 같이 비효율적으로 움직였습니다.

• 과거의 방식 (여러 번 방문):
  1. 먼저 "가장 높은 점수"를 가진 책을 찾으러 책장 전체를 훑어봅니다. (1 번 방문)
  2. 그 다음, 그 결과를 바탕으로 "페이지 합계"를 계산하기 위해 다시 책장 전체를 훑어봅니다. (2 번 방문)
  3. 만약 이 과정이 여러 단계로 이어진다면 (예: 최고점 찾기 → 합계 구하기 → 평균 내기), 도서관 사서 (AI) 는 같은 책장 사이를 수십 번이나 왕복해야 합니다.

이렇게 **불필요하게 도서관을 오가는 시간 (메모리 접근)**이 너무 길어지면서 AI 가 느려지는 것입니다.

2. 해결책: RedFuser(레드퓨저) 의 등장 (한 번에 끝내는 마법)

이 논문에서 제안한 RedFuser는 이 문제를 해결하는 **'자동화 된 지능형 사서'**입니다.

• RedFuser 의 방식 (한 번에 정리):
  RedFuser 는 "아, 이 두 가지 작업은 서로 연결되어 있네?"라고 알아챕니다. 그리고 한 번의 방문으로 모든 일을 끝내버립니다.
  • 책을 한 권씩 꺼내면서 "이게 최고점인가? 아니면 합계에 더할 숫자인가?"를 동시에 계산합니다.
  • 책장 사이를 오갈 필요 없이, 손에 든 책만 가지고 모든 계산을 끝냅니다.

이걸 수학적인 용어로 말하면, 여러 단계의 계산을 하나로 합쳐서 (Fusion) 불필요한 데이터 읽기를 없애고, 계산 결과를 **점점 업데이트 (Incremental Computation)**해 나가는 방식입니다.

3. 핵심 아이디어: "점점 업데이트"하는 비법

RedFuser 의 가장 멋진 점은 **'점점 업데이트'**하는 방식입니다.

• 비유:
  • 예전 방식: "먼저 모든 학생의 점수를 다 받아서 최고점을 찾고, 그 다음 다시 모든 점수를 더해서 합계를 내야 해." (데이터를 모두 모아야 계산 시작)
  • RedFuser 방식: "학생이 하나 들어올 때마다, '지금까지 나온 최고점'과 '지금까지의 합계'를 바로바로 수정해 가자." (데이터가 들어오면 즉시 계산)

이 덕분에 컴퓨터의 **작은 메모리 (캐시)**만으로도 아주 긴 데이터도 처리할 수 있게 되었습니다. 마치 큰 창고에 다 채울 필요 없이, 손에 든 작은 상자에 필요한 것만 계속 채워가며 작업을 끝내는 것과 같습니다.

4. 실제 효과: 얼마나 빨라졌나요?

이 기술을 적용한 결과, 기존에 가장 잘 만들어진 AI 프로그램들보다 2 배에서 5 배까지 더 빨라졌습니다.

• 손으로 만든 최적화 코드: 예전에는 이 정도 속도를 내려면 수석 엔지니어가 밤새도록 손으로 코드를 짜야 했습니다.
• RedFuser 의 성과: 이제 이 기술이 자동으로 그 수준에 도달하거나, 그 이상으로 빠르게 작동합니다.

5. 요약: 왜 중요한가요?

지금 AI 모델은 점점 더 거대해지고 있습니다. 하지만 하드웨어의 속도는 그걸 따라가지 못해 병목 현상이 생깁니다.

RedFuser는 이 병목 현상을 해결하는 자동화된 열쇠입니다.

• 복잡한 계산을 간단한 한 번의 작업으로 바꿉니다.
• 불필요한 이동을 없애 전기와 시간을 아껴줍니다.
• 누구나 쉽게 사용할 수 있도록 자동화했습니다.

결론적으로, RedFuser 는 AI 가 더 똑똑해지고 빠르게 반응할 수 있도록 도와주는 **'효율적인 교통 시스템'**을 구축한 것과 같습니다. 이제 AI 는 더 이상 길에서 막히지 않고, 목적지인 '정답'으로 직행할 수 있게 되었습니다.

기술 요약

1. 문제 정의 (Problem Statement)

• 배경: AI 모델 배포에서 연산 융합 (Operator Fusion) 은 실행 효율성을 극대화하는 핵심 최적화 기법입니다. 그러나 현대 AI 컴파일러는 연쇄적 축소 (Cascaded Reductions) 연산을 처리하는 데 한계가 있습니다.
• 연쇄적 축소 (Cascaded Reductions) 의 정의: 여러 개의 루프가 서로 간의 데이터 의존성 (Data Dependency) 을 가지며 연결된 연산 패턴입니다. 예를 들어, 어텐션 메커니즘 내의 Safe Softmax (최대값 추출 $\to$ 지수 합산 $\to$ 정규화) 나 MoE 라우팅, FP8 양자화 후 GEMM 등이 이에 해당합니다.
• 기존 컴파일러의 한계:
  • 기존 컴파일러 (TVM, PyTorch Dynamo 등) 는 이러한 복잡한 의존성 구조를 분석하지 못해, 각 축소 연산을 별도의 커널로 실행하게 합니다.
  • 이로 인해 중복된 메모리 접근 (Redundant Memory Access) 이 발생하고, 연산 간의 병렬 실행이 불가능하여 지연 시간 (Latency) 이 증가합니다.
  • 수동으로 최적화된 커널 (FlashAttention 등) 은 높은 성능을 보이지만, 특정 패턴에 종속되어 일반화하기 어렵고 유지보수가 어렵습니다.

2. 방법론 (Methodology)

저자들은 연쇄적 축소를 단일 루프로 융합하고, 이를 효율적으로 실행할 수 있는 수학적 프레임워크를 제안했습니다.

가. 수학적 정의 및 융합 조건

• 연쇄적 축소 정의: $i$ 번째 축소 연산이 이전 $i-1$ 개의 축소 결과에 의존하는 형태로 정의됩니다.
• 융합 가능성 조건: 융합을 위해 다음 세 가지 조건을 만족해야 합니다.
  1. 분해 가능성 (Decomposability): 연산 함수 $F_i$ 를 입력 의존 부분 $G_i$ 와 이전 결과 의존 부분 $H_i$ 로 분리할 수 있어야 합니다 ($F_i = G_i \otimes H_i$).
  2. 대수적 구조: 연산자 $\otimes$ 가 교환 법칙, 결합 법칙, 항등원을 갖는 가환 모노이드 (Commutative Monoid) 를 이루어야 합니다.
  3. 분배 법칙 (Distributivity): 축소 연산자 $\oplus$ 가 $\otimes$ 에 대해 분배 법칙을 만족해야 합니다.
  • 예시: GEMM, Safe Softmax, Max/Sum 연산 등은 이 조건을 만족합니다.

나. 융합 표현 유도 (Fusion Expressions)

• 기존 방식은 각 축소 단계마다 이전 단계의 전체 결과를 메모리에서 다시 불러와야 했지만, RedFuser 는 동일한 레벨의 여러 축소 트리를 하나의 축소 트리로 병합합니다.
• 이를 통해 입력 데이터는 한 번만 로드되고, 중간 결과에 대한 메모리 접근이 제거됩니다.

다. 점진적 계산 형태 (Incremental Computation Form)

• 문제: 단순 융합만으로는 이전 단계의 모든 결과를 온칩 메모리 (Shared Memory) 에 캐시해야 하므로, 긴 시퀀스 처리 시 메모리 용량 제약이 발생합니다.
• 해결책: 점진적 계산 (Incremental Computation) 방식을 도입합니다.
  • 이전 단계의 결과가 완전히 계산될 때까지 기다리지 않고, 새로운 입력이 들어올 때마다 현재 단계의 결과를 즉시 업데이트합니다.
  • 이로 인해 온칩 메모리 사용량이 입력 길이에 비례하지 않고 상수 ($O(1)$) 로 유지되어 긴 시퀀스 처리가 가능해집니다.

3. 주요 기여 (Key Contributions)

1. RedFuser 프레임워크 개발:
   • 위 이론적 방법론을 구현한 자동화 프레임워크입니다.
   • 2 단계 구조:
     1. 심볼릭 추론 엔진: AST(추상 구문 트리) 를 분석하여 연산이 융합 가능한지 판단하고, $G_i, H_i$ 함수를 추출하며 융합된 수식과 점진적 계산식을 유도합니다.
     2. 하드웨어 인식 코드 생성: 융합된 수식을 타일 기반 (Tile-based) 표현으로 낮추고, GPU 아키텍처에 최적화된 고성능 커널을 생성합니다.
2. 범용성 및 자동화:
   • FlashAttention, FlashDecoding 등 수동 최적화 기법들이 RedFuser 의 특수한 경우 (Special Cases) 로 자연스럽게 도출됨을 증명했습니다.
   • 수동 최적화의 복잡성을 제거하고 일반화된 자동 컴파일러 최적화를 가능하게 합니다.
3. 성능 평가 및 검증:
   • 다양한 워크로드 (MHA, MLA, MoE 라우팅, FP8 Quant+GEMM) 에서 SOTA AI 컴파일러 대비 2 배에서 5 배의 속도 향상을 달성했습니다.
   • 수동으로 작성된 최적화된 커널 (FlashAttention 등) 과 동등하거나 더 나은 성능을 기록했습니다.

4. 실험 결과 (Results)

• 실험 환경: NVIDIA A10 및 H800 GPU, PyTorch, TVM, Dynamo 등 최신 컴파일러와 비교.
• 주요 성과:
  • MHA (Multi-Head Attention): FlashAttention2 와 동급의 성능을 내면서, PyTorch Dynamo 대비 2.8 배, TVM 대비 2.6 배 빠른 속도를 기록했습니다.
  • MLA (Multi-Latent Attention): FlashMLA 대비 102% 성능을 달성했으며, 기존 컴파일러 대비 2.4 배 ~ 8.7 배 향상되었습니다.
  • MoE 라우팅 및 FP8 Quant: 수동 최적화 라이브러리가 없는 영역에서도 PyTorch Dynamo 대비 1.7 배 ~ 3.4 배, TVM 대비 6.6 배 ~ 12.1 배의 큰 속도 향상을 보였습니다.
  • 비 ML 워크로드: 분산 (Variance) 및 관성 모멘트 계산에서도 기존 컴파일러 대비 4 배 ~ 11 배 이상의 성능 향상을 확인했습니다.

5. 의의 및 결론 (Significance)

• 이론과 실전의 가교: 연쇄적 축소라는 복잡한 계산 패턴에 대한 엄밀한 수학적 정의를 제공하여, 이를 자동화된 컴파일러 최적화로 연결하는 새로운 패러다임을 제시했습니다.
• 메모리 병목 해결: 점진적 계산 방식을 통해 온칩 메모리 용량 제한을 우회함으로써, 긴 시퀀스 처리 (LLM 추론 등) 에서 발생하는 메모리 병목 현상을 효과적으로 해결했습니다.
• 미래 지향성: 현재는 GPU 에 집중되어 있지만, 이 프레임워크는 다양한 AI 가속기로 확장 가능하며, 수동 최적화의 한계를 넘어 일반 목적의 고성능 자동 최적화를 실현하는 중요한 발걸음입니다.

요약하자면, RedFuser는 AI 가속기에서 발생하는 복잡한 연쇄적 축소 연산을 자동으로 분석하고, 수학적 유도를 통해 메모리 접근을 최소화하는 단일 커널로 변환함으로써, 수동 최적화 수준의 성능을 자동화 프레임워크로 달성한 획기적인 연구입니다.