LLM을 MegaKernel로 컴파일하여 Low-Latency 추론 실현하기

6 hours ago 1

LLM 추론을 단일 메가커널로 자동 변환하는 컴파일러를 개발했음
MegaKernel(Persistent 커널) 방식은 LLM 추론에서 계산과 통신을 완전히 하나의 GPU 커널에 통합하여 매우 낮은 레이턴시를 가능하게 함
기존 ML 프레임워크나 커널 라이브러리의 분산 구조로 인해 전체 파이프라인의 단일 커널화가 매우 어렵다는 문제점 존재
Mirage Persistent Kernel(MPK)은 컴파일러와 런타임 시스템을 통해 자동으로 멀티-GPU LLM 추론을 고성능 megakernel로 변환함
MPK는 연산 그래프를 세분화된 태스크 그래프로 변환해, 소프트웨어 파이프라이닝과 연산-통신 겹침을 극대화함
MPK 적용 시 기존 시스템 대비 토큰 생성 지연이 줄고, GPU 수가 많아질수록 성능 향상폭이 더욱 커짐

개요 및 MegaKernel 방식의 장점

대규모 언어 모델(LLM) 추론에서 지연 시간을 줄이는 효과적인 방법 중 하나는 모든 계산과 통신 과정을 단일 megakernel(일관적 커널) 에 융합하는 방식임
이 방식은 모델 전체의 레이어별 연산, GPU 간 통신까지 모든 처리를 하나의 GPU 커널이 끊김 없이 수행함
주요 이점은 다음과 같음
- 반복적인 커널 호출을 생략해 커널 런칭 오버헤드 제거
- 레이어 전반에서 소프트웨어 파이프라이닝 실현 가능
- 연산과 통신을 동시 진행하여 레이턴시 숨김

기존 PyTorch, Triton, TVM 같은 ML 프레임워크에서는 end-to-end megakernel 자동 생성 기능을 본질적으로 지원하지 않음
실제 LLM 시스템은 NCCL/NVSHMEM(통신), FlashInfer/FlashAttention(어텐션), CUDA/Triton(커스텀 연산) 등 다양한 커널 라이브러리 조합으로 이루어져, 단일 커널로의 통합이 어려움
이러한 배경에서, CMU, UW, Berkeley, NVIDIA, Tsinghua 연구진은 Mirage Persistent Kernel(MPK) 을 개발함
- MPK는 컴파일러와 런타임을 결합하여 LLM 추론 전체 파이프라인을 자동으로 고성능 megakernel로 변환함

MPK는 커널 런칭 오버헤드를 완전히 제거하고, 계층 간 연산/데이터 로딩/통신 겹침을 최대화하여 극저지연 LLM 추론 환경 구현함
실제 테스트(39-토큰 프롬프트, 512 토큰 생성, speculative decoding 미사용)에서,
- NVIDIA A100 40GB GPU 단일 환경 기준 vLLM/SGLang 등 기존 최적화 시스템의 토큰당 디코딩 지연(14.5ms) 대비, MPK는 12.5ms까지 단축함
- 이 수치는 이론적 하한선(10ms)에 근접함(1.6TB/s 메모리 대역폭, 16GB 가중치 로딩 기준)
멀티 GPU 환경에서 연산과 통신을 완전히 통합함으로써, GPU 수가 늘어날수록 MPK의 성능 우위가 더욱 두드러짐

일반적으로 LLM 연산은 각 연산(예: 행렬 곱, 어텐션) 또는 통신 연산(예: all-reduce)이 노드이며, 데이터 의존성이 엣지인 컴퓨테이션 그래프로 표현됨
기존 설계에서는 연산자당 별도 커널 실행 방식이 보편적이나, 이는 실제 의존 데이터 단위가 아니라 커널 단위의 의존성만 반영해, 파이프라이닝 기회가 제한적임
예시: 행렬 곱 뒤에 allreduce가 있는 경우, 전체 행렬 곱이 끝나야만 allreduce 실행이 시작됨. 실제로는 데이터를 쪼개서 부분 실행/의존 관계를 활용하는 것이 가능함
MPK 컴파일러는 연산 그래프를 세분화하여, 실제 데이터 단위에 적합한 fine-grained task graph로 자동 변환함
- 각 태스크(사각형)는 개별 GPU SM에 할당되는 연산/통신 단위임
- 각 이벤트(원)는 태스크 간 동기화 지점임
- 태스크 및 이벤트 간의 간선으로 효율적인 데이터/제어 의존성 표현함
이 태스크 그래프 덕분에 MPK는 연산과 통신이 부분적으로 혹은 병렬적으로 더 중첩될 수 있음
Mirage kernel superoptimizer로 각 태스크에 맞는 고성능 CUDA 구현도 자동 생성함

각 워커는 SM 단위로 동작하며, 전용 태스크 큐를 관리함
루프 방식으로
1. 다음 태스크를 큐에서 받아옴
2. 수행(예: matmul, attention, 데이터 전송)
3. 완료 시 이벤트에 알림
4. 반복 처리
이를 통해 각 워커의 자원 활용 최적화 및 비동기적 계층 연산 가능함

태스크가 완료되면 특정 이벤트 카운터를 증가시킴. 카운터가 임계치에 도달하면 이벤트 활성화, 스케줄러 큐에 삽입됨
스케줄러는 해당 이벤트에 의존 관계가 있는 후속 태스크를 실행함
덕분에 파인그레인드 소프트웨어 파이프라이닝과 계산-통신 중첩이 자연스럽게 이루어짐
- 예: 한 레이어의 matmul과 다른 레이어의 attention이 동시 실행
- 부분적으로 완료된 matmul 결과가 나오는 즉시 allreduce 통신 시작 가능
모든 스케줄링·태스크 전환이 단일 커널 컨텍스트 내에서 벌어지므로, 태스크간 오버헤드는 1–2 마이크로초(μs) 수준으로 매우 낮음

MPK 목표: 개발자가 적은 파이썬 코드(수십 줄 내외)만 작성해도 손쉽게 LLM을 megakernel로 컴파일하고, 최대 성능을 발휘할 수 있도록 지원함
주요 발전 방향
- 최신 GPU 아키텍처 지원: 예를 들면 NVIDIA Blackwell 대상, warp 단위 특화 방식 등
- 동적 workload 처리: mixture-of-experts(MoE) 등 동적 제어 흐름이 필요한 모델을 위한 컴파일 전략 연구
- 고급 태스크 스케줄링: 우선순위 기반, 처리량 최적화 등 현대적 정책 연구 및 적용 가능성 추구
MPK는 GPU 기반 LLM 추론 작업의 컴파일·실행 방식에 근본적 전환점을 제시하며, 커뮤니티와의 협력 확대를 바람