MPK: A Compiler and Runtime for Mega-Kernelizing Tensor Programs

기술 요약: MPK - 텐서 프로그램의 메가 커널화를 위한 컴파일러 및 런타임

1. 문제 정의

대규모 언어 모델(LLM) 및 기타 텐서 프로그램을 위한 현재의 고성능 GPU 추론 시스템은 주로 **연산자당 커널 실행 모델(kernel-per-operator execution model)**에 의존합니다. 이 패러다임에서는 각 연산자(예: 행렬 곱셈, 어텐션, All-Reduce)가 별도의 GPU 커널로서 실행됩니다. 이러한 접근 방식은 프로그래밍을 단순화하지만, 성능을 저해하는 세 가지 결정적인 한계를 초래합니다:

1. 암시적 커널 장벽(Implicit Kernel Barriers): 전통적인 GPU 실행은 동일한 스트림 상의 연속적인 커널 런칭 사이에 장벽을 강제합니다. 이는 **연산자 간 소프트웨어 파이프라이닝(cross-operator software pipelining)**을 방해하여, 종속된 연산자들이 엄격하게 순차적으로 실행되도록 강제하며, 하드웨어 리소스(Tensor Cores, TMA)가 유휴 상태로 머무는 "파이프라인 버블(pipeline bubbles)"을 생성합니다.
2. 조립 단위의 오버랩(Coarse-Grained Overlap): 종속성이 전체 연산자 단위로 강제됩니다. 결과적으로, 계산과 통신 사이의 미세한(fine-grained) 오버랩이 불가능합니다. 예를 들어, AllReduce 연산은 특정 행렬 곱셈 출력의 특정 조각(fragment)만이 특정 조각에만 의존함에도 불구하고, 앞선 MatMul의 전체가 완료될 때까지 기다려야 합니다.
3. 런칭 오버헤드 및 정적 유연성 부족(Launch Overhead and Static Flexibility): 추론 반복당 수백 또는 수천 개의 커널을 실행하는 것은 상당한 런칭 오버헤드를 발생시킵니다. CUDA Graph가 이를 완화하지만, 이는 대체로 정적입니다. 제어 흐름, 텐서 형태(shape), 또는 데이터 종속성의 동적인 변화는 종종 그래프의 재인스턴스화를 요구하며, 이는 LLM 서빙과 같은 동적 워크로드의 유연성을 제한합니다.

FlashDMoE, Spector 등의 기존 솔루션(수동 메가 커널)은 성능 이점을 제공하지만, 광범위한 엔지니어링 노력과 깊은 GPU 전문 지식을 요구하며, 다양한 모델이나 하드웨어에 걸쳐 일반화되지 못합니다.

2. 방법론

저자들은 **Mirage Persistent Kernel (MPK)**을 제시합니다. 이는 멀티 GPU 모델 추론을 단일한 고성능 **메가 커널(mega-kernel)**로 자동 변환하는 최초의 컴파일러 및 런타임 시스템입니다. MPK는 다음 두 가지 주요 구성 요소를 통해 이를 달성합니다:

A. SM 레벨 그래프 표현 (tGraph)

MPK는 전체 GPU가 아닌 개별 스트리밍 멀티프로세서(SM) 단위의 계산 및 통신을 포착하는 tGraph라는 새로운 그래프 표현 방식을 도입합니다.

• 노드(Nodes): 단일 SM에서 실행되는 태스크(계산 또는 통신)를 나타냅니다.
• 엣지(Edges): 이벤트를 통해 태스크 간의 미세한 종속성을 인코딩합니다.
• 장점: 이 입도(granularity)는 커널 레벨 모델에서는 접근할 수 없는 병렬성을 노출하여, 연산자 간 소프트웨어 파이프라이닝과 계산 및 통신의 미세한 오버랩을 가능하게 합니다.

컴파일러 최적화:
MPK 컴파일러는 다음과 같은 여러 단계를 거쳐 표준 텐서 프로그램을 최적화된 tGraph로 변환합니다:

1. 연산자 분해(Operator Decomposition): 가용 SM 수를 맞추기 위해 출력 텐서 파티셔닝(예: 행렬 곱셈의 행/열 타일링)을 기반으로 연산자를 서로 소(disjoint)인 태스크 집합으로 분할합니다.
2. 종속성 분석(Dependency Analysis): 데이터 중첩이 존재하는 곳에만 동기화 이벤트를 삽 la(insert)하기 위해 프로듀서-컨슈머 쌍을 열거합니다.
3. 이벤트 퓨전(Event Fusion): 중복된 동기화 지점을 제거하기 위해 후속 집합 퓨전(successor-set fusion, 동일한 컨슈머 세트를 가진 이벤트 병합) 및 *선행 집합 퓨전(predecessor-set fusion, 동일한 프로듀서 세트를 가진 이벤트 병합)*을 사용하여 그래프 복잡도를 줄입니다.
4. 정규화(Normalization): 모든 태스크가 최대 하나의 종속 이벤트와 하나의 트리거 이벤트를 갖도록 그래프를 변환합니다. 이는 팬인(fan-in)/팬아웃(fan-out)을 처리하기 위해 더미 태스크와 이벤트를 삽입함으로써 달성되며, 균일하고 낮은 오버헤드의 표현을 보장합니다.
5. 선형화(Linearization): 동일한 이벤트에 의해 트리거되는 태스크들에 연속적인 인덱스를 할당하기 위해 BFS 기반 알고리즘을 사용합니다. 이를 통해 런타임이 명시적인 리스트 대신 단순한 인덱스 범위(첫 번째/마지막 태스크 ID)로 이벤트 팬아웃을 인코딩할 수 있게 하여 디바이스 메모리 사용량을 크게 줄입니다.
6. 태스크 구현 생성(Task Implementation Generation): 컴파일러 초최적화(superoptimization, Mirage)를 사용하여 레지스터 재사용 및 공유 메모리 레이아웃과 같은 intra-SM 최적화를 포함한 각 개별 태스크를 위한 고성능 CUDA 구현을 자동으로 생성합니다.

B. 인-커널 병렬 런타임 (In-Kernel Parallel Runtime)

MPK는 분산형 이벤트 드리븐(event-driven) 아키텍처를 사용하여 단일 지속적 메가 커널 내에서 tGraph를 실행합니다:

• 워커/스케줄러 분할(Worker/Scheduler Partitioning): GPU의 SM을 워커(Workers)(태스크 실행)와 스케줄러(Schedulers)(이벤트 큐 관리)로 정적으로 분할합니다.
• 이벤트 드리븐 실행: 스케줄러는 활성화된 이벤트(모든 전제 조건이 충족된 이벤트)를 폴링하고 그에 해당하는 태스크를 워커에게 디스패치합니다. 이는 호스트 측 동기화나 반복적인 커널 런칭을 제거합니다.
• 하이브리드 태스크 런칭: 지연 시간과 부하 분산 사이의 균형을 맞추기 위해 MPK는 두 가지 모드를 채택합니다:
  • Just-In-Time (JIT): 종속성이 충족된 후에만 태스크를 런칭합니다. 이는 가변적인 시퀀스 길이(예: 어텐션의 가변 시퀀스 길이)와 같은 동적 워크로드에 적응하여 부하 불균형을 방지합니다.
  • Ahead-Of-Time (AOT): 종속성이 충족되기 전에 태스크를 워커 큐에 미리 인큐(enqueue)합니다. 이는 예측 가능한 연산자에 대해 디스패치 지연 시간을 줄여줍니다.
• 페이지드 공유 메모리(Paged Shared Memory): 교차 태스크 소프트웨어 파이프라이닝을 가능하게 하기 위해 MPK는 공유 메모리에 대한 페이지 추상화를 도입합니다. 태스크는 고정 크기의 페이지를 획득하고 해제하며, 이를 통해 현재 태스크가 계산하는 동안 다음 태스크가 리소스 충돌 없이 데이터를 공유 메모리로 프리페칭(prefetching)할 수 있도록 합니다.

3. 주요 기여

1. 최초의 자동화된 메가 커널 시스템: MPK는 임의의 멀티 GPU 텐서 프로그램을 단일 메가 커널로 자동 변환하는 최초의 시스템으로, 수동 커널 퓨전의 필요성을 제거했습니다.
2. SM 레벨 추상화: tGraph 표현의 도입은 종속성 추적을 연산자 레벨에서 SM-태스크 레벨로 전환하여, 커널당 연산자 모델에서는 불가능했던 최적화(파이프라이닝, 미세한 오버랩)를 가능하게 했습니다.
3. 엔드 투 엔드 컴파일러 및 런타임: 이 시스템은 태스크 그래프를 최적화하고 CUDA 코드를 생성하는 컴파일러와, 단일 커널 내에서 이러한 태스크를 비동기적으로 실행하는 런타임을 통합하여 컨티뉴어스 배칭(continuous batching) 및 페이지드 어텐션(paged attention)과 같은 동적 워크로드를 지원합니다.
4. 하이브리드 실행 전략: JIT와 AOT 태스크 런칭의 결적인 조합을 통해, MPK는 LLM 추론의 동적 특성(가변 배치/시퀀스 길이)과 정적 연산자의 효율성을 모두 처리할 수 있습니다.

4. 평가 결과

저자들은 Qwen3 및 Llama-3 변형을 포함한 5개의 널리 사용되는 LLM에 대해 세 세대의 NVIDIA GPU(A100, H100, B200)에서 MPK를 평가하였으며, 이를 SGLang, vLLM 및 vanilla PyTorch와 같은 최첨단 커널당 연산자 시스템과 비교하였습니다.

• 성능 향상: MPK는 SGLang 및 vLLM 대비 1.0× ~ 1.7× 낮은 엔드 투 엔드 추론 지연 시간을 달성했습니다. 개선 효과는 더 작은 모델과 최신 하드웨어 세대(예: B200)에서 가장 두드러졌는데, 이는 계산 시간 대비 커널 런칭 오버헤드와 파이프라인 버블이 더 중요해지기 때문입니다.
• 처리량(Throughput): 멀티 GPU 설정(최대 8개의 H100)에서 MPK는 최적화된 베이스라인 대비 1.1× ~ 1.4× 향상된 처리량을 보여줍니다.
• 지연 시간 한계: A100에서 MPK는 Qwen3-8B의 토큰당 디코딩 지연 시간을 vLLM/SGLang의 14.5ms에서 하드웨어 추정 하한선인 약 10ms에 근접한 12.5ms로 단축했습니다.
• 절제 연구(Ablation Studies):
  • 교차 태스크 파이프라이닝: 태스크 실행 시간을 1.2~1.3× 단축했습니다.
  • 계산-통신 오버랩: 반복당 지연 시간을 1.1× 감소시켰습니다.
  • 커널 런칭 감소: 베이스라인 시스템에서 토큰당 293회의 커널 런칭과 관련된 약 1.1ms의 오버헤드를 제거했습니다.
• 확장성: MPK는 2, 4, 8개의 GPU에 걸쳐 효과적으로 확장되며, 베이스라인 대비 속도 향상을 유지합니다.

5. 의의 및 주장

본 논문은 MPK가 실행 파이프라인을 단일 메가 커널로 통합함으로써 GPU 추론 시스템의 패러다임 전환을 나타낸다고 주장합니다. 그 의의는 다음과 같습니다:

• 커널 장벽 돌파: 동기화를 SM 레벨로 이동함으로써, MPK는 현재 커널당 연산자 시스템에서는 접근할 수 없는 하드웨어 활용 잠재력을 끌어냅니다.
• 자동화 대 수동 노력: 전문가의 수동 설계가 필요했던 이전의 메가 커널 접근 방식과 달리, MPK는 프로세스를 자동화하여 최소한의 코드 변경만으로 표준 PyTorch 개발자가 고성능 메가 커널 실행을 사용할 수 있도록 합니다.
• 하드웨어 근접성: MPK는 소프트웨어 오버헤드(런칭, 스케줄링, 동기화)를 최소화하고 하드웨어 오버랩을 극대화함으로써 LLM 추론 성능을 하드웨어의 이론적 한계치에 가깝게 밀어붙입니다.

저자들은 MPK가 기존 프로그래밍 모델의 유연성을 유지하면서도 하드웨어 한계에 근접한 성능을 제공하는 고성능 컴파일러 기반 추론 시스템을 향한 새로운 길을 열었다고 결론짓습니다.