▲

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

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

개요 및 MegaKernel 방식의 장점

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

기존 한계와 MPK의 등장

  • 기존 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의 핵심 가치

  • 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의 성능 우위가 더욱 두드러짐

MPK 동작 구조 상세

Part 1. 컴파일러 – LLM 연산 그래프 → 태스크 그래프 변환

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

Part 2. 런타임 – 메가커널 내부에서 태스크 그래프 실행

  • MPK 런타임은 태스크 그래프를 GPU 하나의 커널(메가커널) 내부에서만 완전히 실행하는 방식임
  • GPU의 모든 SM(Streaming Multiprocessors)을 정적으로 워커와 스케줄러 역할로 분할함

워커

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

스케줄러

  • 분산 스케줄러가 각 SM 내 단일 warp 단위로 동작, 최대 4개 스케줄러 동시 실행 가능
  • 각 스케줄러는 활성화된 이벤트 큐를 관리하고, 조건이 충족된 태스크들을 워커에 할당함
  • 이로써 중앙 집중형 동기화 오버헤드 없이 대규모 태스크 분산 처리 가능함

이벤트 기반 실행 방식

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

미래 방향

  • MPK 목표: 개발자가 적은 파이썬 코드(수십 줄 내외)만 작성해도 손쉽게 LLM을 megakernel로 컴파일하고, 최대 성능을 발휘할 수 있도록 지원함

  • 주요 발전 방향

    • 최신 GPU 아키텍처 지원: 예를 들면 NVIDIA Blackwell 대상, warp 단위 특화 방식 등
    • 동적 workload 처리: mixture-of-experts(MoE) 등 동적 제어 흐름이 필요한 모델을 위한 컴파일 전략 연구
    • 고급 태스크 스케줄링: 우선순위 기반, 처리량 최적화 등 현대적 정책 연구 및 적용 가능성 추구

  • MPK는 GPU 기반 LLM 추론 작업의 컴파일·실행 방식에 근본적 전환점을 제시하며, 커뮤니티와의 협력 확대를 바람

추가 자료

▲
GN⁺ 11시간전  [-]
Hacker News 의견

  • 저자님께, on-GPU 인터프리터 방식이 굉장히 유망한 미래 방향처럼 보이는 점이 흥미로움. 거의 동일한 접근을 보이는 다른 연구도 있으니 관련 글 참고 권장. CUDA의 근본적인 프로그래밍 모델(예: 커널 런치)이 미세한 작업 기반 병렬화를 위해 우회되고 있는데, 이 방식이 하드웨어의 활용도를 더 높이는 걸 직접 목격. CUDA가 우리를 여러 면에서 잡아두고 있던 건 아닌지 궁금증. 저자의 연구가 PyTorch의 실험적 백엔드로 들어올 가능성에 대한 기대감. 그리고, 첫 번째 부분의 두 문단이 거의 동일하니 사소한 오타 지적.

    • 피드백 감사 인사 및 Stanford의 MegaKernel 프로젝트도 유사 도전을 하고 있다는 언급. 하지만 MPK는 사용자들이 PyTorch 수준에서 LLM을 표현하면 컴파일러가 자동으로 최적화된 megakernel로 변환해주는 방식 추구. 목표는 megakernel 프로그래밍을 누구나 쉽게 만들 수 있게 하는 것. CUDA가 특히 레이턴시에 민감한 워크로드에 한계로 작용하는 것에 전적으로 동의. GPU가 커지고 빨라지면서 작은 배치에서도 하드웨어 자원을 충분히 활용하는 독립 커널 작성이 점점 더 어려워지는 현실. PyTorch와 협업 하에 MPK가 megakernel 생성 지원에 도움 줄 수 있다는 방향으로 적극 탐구 중. 중복 문단도 지적해줘서 감사.

  • vLLM과 SGLang에서 일정 기간 밀접하게 작업해왔는데, 이 프로젝트가 바로 후속 프로젝트의 이상적인 모습이라고 확신. 연산 의존성 그래프를 분석하고, 연산을 퓨징하거나 더 똑똑하게 태스크를 스케줄링하는 내용이 인상적. 팀에 축하 인사.

    • 긍정적인 피드백에 대한 감사 인사. MPK가 특히 저지연 LLM 서빙 분야에서 기존 LLM 시스템을 확장하는데 기여 가능성에 대해 큰 기대감. 앞으로 다양한 협업과 방향성 탐구에 의욕.

  • 글과 github README 훑어보고 정말 멋진 프로젝트라고 생각. 이런 최적화 방안이 추론뿐 아니라 학습 단계에까지 적용될 수 있는지 궁금증. 특히 backward 연산과 gradient 통신의 퓨징이 도전 과제임을 인지. 현재는 dynamic workload(예: MoE)를 지원하지 않는 것으로 알고 있는데, 최근에 MoE를 한 커널에서 처리하는 논문 FlashDMoE: Fast Distributed MoE in a Single Kernel 언급.

    • 글, README까지 읽어줘서 감사 인사. 학습단계 지원도 가능하지만, 대체로 학습 커널이 더 커서 커널 런치 오버헤드가 크게 문제되지 않기 때문에 추론(특히, 저지연)이 더 큰 수혜 대상. 공유해 준 FlashDMoE 논문도 흥미깊게 봤고, MoE 모델 지원도 다음 목표로 삼고 있음을 강조.

    • 개인적으로 gradient 기반 학습 최적화에 시간 투자하는 것에 다소 회의적 시각. 실제 많은 학습 태스크들은 이산적 값의 특성을 가져서, gradient 기반 학습으론 잘 다룰 수 없다고 생각.

  • 다음 단계로는 바로 Verilog로 컴파일해서 aliexpress에서 LLM 하드웨어 직접 구입하는 게 꿈.

    • Chisel 등 하드웨어 기술을 소개하는 글 공유. AI, GPU 등장이전에는 이러한 소프트웨어에서 하드웨어 직접 전환 아이디어가 유망한 접근이었음. CPU 발전이 정체 상태이고, 소프트웨어와 하드웨어 중간 계층을 더 최적화하려는 열망은 꾸준하지만, GPU 스타일의 병렬 컴퓨팅이 주류 가속 방식을 이어갈 가능성이 높음. 일반용 CPU는 결국 GPU를 관리하는 작은 뇌 역할로 남아 있을 전망. 다만 소프트웨어에서 바로 하드웨어로 전환하는 방식이 주류로 가기는 어렵다는 예상.

    • 5~10년 뒤 LLM의 구조가 안정화되면, 하드웨어에 바로 매핑하는 게 실질적이 될 수도 있다는 예상. 현재 기술로 수백억 파라미터도 1.5비트 근방 초저정밀 논리 게이트만 사용하여 단일 웨이퍼에 들어갈 가능성 언급. 정밀도가 더 높아질수록 게이트 수가 기하급수적으로 증가하니, 현재로썬 가중치 메모리 유지 및 계산 유닛 공유 방식이 효율적. 미래엔 초저정밀 LLM 개발이 필수 과제.

    • 학습 비용이 이미 높은데 마스크 비용도 추가하면 상황이 더 어려워짐에 대한 유머와, 사실상 AI 하드웨어 스타트업들이 이런 방향의 시도는 오래 해오고 있다는 냉철한 평가.

    • LLM-in-a-box 방식이 실제로 존재한다면 상당히 매력적이라는 평. 조만간 오프라인(air-gap) 환경에서 일할 기회가 있는데 저런 솔루션이 굉장히 유용할 것 같다는 기대.

  • 직접 Modal GPU 환경에서 코드를 돌려봤는데, 연구에서 주장하는 성능 향상 수치가 실제로 재현됨. mirage 프로젝트 결과 코드 공유. Triton + FlashInfer 조합에서 1토큰당 레이턴시 19.2ms 수준, MPK에서는 동일 조건에서 7.7ms로 대폭 개선 경험.

    • 직접 결과 재현해준 데에 감사 표시.

  • 예전에 작은 CUDA 대회를 한 경험. 이미지나 비전 분야의 병렬 알고리즘이었는데, 여기서 똑똑해 보이려고 중간 결과를 메모리로 캐싱했음. 공모전 결과 확인 후, 남들은 나보다 훨씬 빠른 코드 제출함에 놀람. 이유를 보니, 중간 결과 따위 캐싱하지 않고 계속 다시 계산하는 방식이었음. 메모리 왕복보다 연산 비용이 훨씬 작았던 것. 이 프로젝트도 아마 비슷할 거라 추정. megakernel로 컴파일하면서 layer 경계가 사라져서 중간 결과 공유는 줄고 연산량은 늘어나는데, 전체적으로 보면 메모리 왕복이 줄어들어 큰 이득. 특히 컨볼루션 네트워크에서는 sweet spot이 있을 텐데, megakernel에서 이 부분을 어떻게 처리하는지는 모르겠음.

  • 지금도 LLM에 대한 새로운 비유가 계속 등장. 혹시 LLM을 트랜지스터처럼 여길 수도 있지 않을까 생각. 지금은 마치 천공카드로 곱셈만 하던 방 크기의 컴퓨터 단계와 비슷하다고 상상. 1백만 개의 o3-pro 쿼리를 동시에 돌릴 수 있다면 어떤 일이 펼쳐질지 상상하는 재미.

  • 이 프로젝트는 CMU(카네기멜론) 출신. Stanford의 Hazy Research에서도 megakernel을 다룬 블로그 No Bubbles 언급. 이 분야에서 경쟁이 활발히 벌어지는 모습을 보는 것이 인상적. (추가) "mirage" 프로젝트의 더 큰 그림을 다룬 논문도 있으나 megakernel 접근법은 다루지 않음 논문 링크

    • 게시글 작성자 본인이 직접 답변. Stanford와의 연구가 병렬적으로 이뤄지고 있음에 동의. 주요 차이점은 자동화된 megakernel 생성 컴파일러에 주력한다는 점.

    • Hazy Research의 ThunderKittens도 굉장히 쿨한 라이브러리임을 언급. 최근 NVIDIA GPU 모델을 최대한 이용하기 위해 포멀화, 파이프라이닝, 분할정복, 효율 극대화, 그리고 전용 컴파일러/DSL 개발에 큰 노력이 집중되고 있다는 평가.

  • Qwen 8B 성능 수치는 검증된다면 상당히 인상적. 이전 megakernel 방식보다 더 실용적으로 느껴짐. 각 SM당 하나씩 유지되는 이러한 커널 방식이 과거 Larrabee를 연상시킴. 기존 CUDA가 아닌, 전통적인 프로세스-스레드-SIMD 경로를 갔으면 지금 세상이 어떻게 됐을까 궁금증.

  • 소프트웨어 기반 추론보다 순수 ASIC 방식으로 고정형 LLM을 만드는 것에 대한 아이디어. 비용 이점? 소프트웨어에서 추가로 다루거나 미세 조정 가능한 계층 제공 가능성? 실제로 ‘충분히 좋은’ 수준에 거의 도달한 만큼, 앞으로 2~4년간 전문화된 칩으로 고정해 사용하는 결정을 할지 모를 상황. 초특화 하드웨어가 제공할 이점이 도대체 어느 시점에서 빛을 발할지 의문.

    • 이어지는 추가 질문:
      1. 오토컴플리트, 키워드 라우팅, 음성 인식 등 특정 태스크에서 ASIC과 megakernel GPU 각각의 지연과 전력 소모 차이, 그리고 엣지 디바이스에서 고정형 함수 방식 채택을 정당화할 정도인지에 대한 의문.
      2. ASIC은 재학습이 어렵지만, 기본 모델만 하드웨어에 박고, LoRA처럼 작은 학습 가능한 모듈만 범용 코프로세서에 띄우는 하이브리드 구상 가능성.
      3. 트랜스포머의 고정 토폴로지가 ASIC 설계에서 공간 재사용에 적합한지, 아니면 GPT-3 급 모델 크기로 인해 과감한 프루닝이나 양자화 없이는 여전히 ASIC화가 어렵지 않을지에 대한 탐구.
답변달기