기술 요약: Ada-MK

문제 정의

온라인 광고 시스템과 같은 지연 시간에 민감한 상업적 환경에서 대규모 언어 모델 (LLM) 을 배포할 때 중요한 병목 현상에 직면합니다: 엔드투엔드 추론 지연 시간은 밀리초 범위 (1~5 ms) 내에서 엄격하게 제한되어야 합니다. 여러 연산자를 단일 영구 커널로 융합하여 커널 시작 오버헤드와 HBM 왕복을 제거하는 MegaKernel 접근 방식은 이론적 해결책을 제시하지만, 리소스가 제한된 NVIDIA Ada (L20) GPU 에서의 실제 적용은 이식성과 효율성 사이의 근본적인 긴장 관계로 인해 방해받고 있습니다.

스탠포드 대학의 구현과 같은 기존 MegaKernel 구현은 Hopper/Blackwell 아키텍처에 밀접하게 결합되어 있으며, TMA(텐서 메모리 가속기) 와 더 큰 공유 메모리 (227 KB) 와 같은 특정 하드웨어 기능에 의존합니다. Ada 아키텍처는 네이티브 TMA 지원을 제공하지 않으며 공유 메모리 (128 KB) 를 절반만 제공하므로, 파이프라인 단계를 4 에서 2 로 축소해야 하고 상당한 듀티 사이클 손실이 발생합니다. 반면, Mirage MPK 와 같은 자동 튜닝 프레임워크는 메모리 상태에 기반한 if-else 분기와 같은 런타임 동적 의사결정을 도입하여 변동성을 처리합니다. 초저지연 환경에서 이러한 런타임 분기는 용납할 수 없는 명령어 발행 페널티와 파이프라인 버블을 초래하여 시스템이 하드웨어 최대 성능에 도달하는 것을 방해합니다.

방법론

저자들은 런타임에서 컴파일 시간으로 동적 의사결정을 전환함으로써 Ada 아키텍처에 맞춰 MegaKernel 최적화를 적응시키는 Ada-MK 프레임워크를 제안합니다. 이 시스템은 두 가지 긴밀하게 결합된 단계를 통해 작동합니다:

1. 오프라인 MegaKernel 합성 (1 단계):

   • MLIR 기반 세분화 DAG 구성: 연산자를 블랙박스로 취급하는 대신, 시스템은 Torch-MLIR 을 사용하여 Transformer Decoder 와 LM Head 를 PTX 명령어 수준에서 세분화된 방향성 비순환 그래프 (DAG) 로 분해합니다. 이는 정밀한 데이터 종속성 (RAW/WAR) 과 리소스制약을 포착합니다.
   • 적응형 공유 메모리 모델링: 하드웨어 사양, 모델 아키텍처, 동적 워크로드의 3 차원 제약 모델을 사용하여 사용 가능한 공유 메모리 페이지를 계산합니다. 시스템은 K 차원 세분화 분할을 사용하여 피크 공유 메모리 수요를 50% 감소시키기 위해 K 차원 타일을 반으로 줄여, 제한된 하드웨어에서 더 많은 파이프라인 단계를 가능하게 합니다.
   • 오프라인 탐색 및 경로 고형화: 가능한 탐색 공간 내에서 시스템은 최적의 실행 흔적을 식별하기 위해 포괄적인 오프라인 탐색을 수행합니다. 전통적인 오토튜너가 연산자 수준에서 탐색하는 것과 달리, Ada-MK 는 더 세분화된 병렬성을 포착하기 위해 DAG 수준에서 탐색합니다. 최적 경로는 커널 코드에 "고형화"(직렬화) 되어 런타임 분기를 완전히 제거합니다.

2. 온라인 워프 특화 실행 (2 단계):

   • 이종 하이브리드 엔진: Ada-MK 는 TensorRT-LLM 추론 엔진에 플러그인으로 내장됩니다. 이는 이종 하이브리드 전략을 활용합니다: Prefill 단계 (계산 바운드, 대배치) 는 높은 처리량을 위해 TensorRT-LLM 의 네이티브 융합 연산자를 유지하는 반면, Decode 단계 (IO 바운드, 소배치) 는 최적화된 MegaKernel 엔진으로 전환합니다.
   • 워프 특화: MegaKernel 내부에서 SM 리소스는 공유 메모리를 통해 메모리 접근과 계산을 오버랩하는 5 개의 워프 역할 (로더, 소비자, 스토어, 컨트롤러, 런처) 로 공간적으로 분할됩니다.

주요 기여

1. 적응형 공유 메모리 관리: 이 논문은 피크 공유 메모리 사용을 50% 감소시키는 다차원 리소스 모델과 K 차원 분할 기법을 소개합니다. 이는 Ada 의 제한된 128 KB 공유 메모리에서 효율적인 파이프라인을 재구성하여 기존 연구에서 관찰된 파이프라인 단계 축소를 극복합니다.
2. 세분화 DAG 기반 자동 탐색: MLIR 로우링을 활용하여 PTX 수준 종속성 DAG 를 구성함으로써, Ada-MK 는 최적 실행 경로를 고형화하기 위해 오프라인 프로파일링을 수행합니다. 이 접근법은 Mirage MPK 와 같은 이전 자동 튜닝 방법의 고유한 런타임 동적 의사결정 오버헤드 (분기 페널티) 를 제거하여 결정론적 지연 시간을 보장합니다.
3. 이종 하이브리드 추론 엔진: 저자들은 TensorRT-LLM 내부의 플러그인으로 MegaKernel 을 상업적으로 배포한 첫 사례를 제시합니다. 이 설계는 Prefill 단계에서는 네이티브 연산자의 높은 병렬성을 활용하면서, Decode 단계에서는 저지연 커널 시작 없는 MegaKernel 을 활용하여 전체 엔진 교체에 따른 엔지니어링 비용을 회피합니다.
4. 양자화 인식 최적화: 프레임워크는 K 차원 분할 및 TensorCore 인식 가중치 재배열과 양자화 (GPTQ-W4A16) 를 통합합니다. 이는 메모리 압력을 추가로 압축하고 더 깊은 계산 - 메모리 오버랩을 가능하게 합니다.

실험 결과

실험은 고정된 짧은 시퀀스와 실제 세계 데이터셋 (CSL, Human-eval) 에서 GPTQ-W4A16 양자화된 Qwen3-1.7B 및 Qwen2.5-1.5B 모델을 사용하여 NVIDIA L20 GPU 에서 수행되었습니다.

• 처리량 향상: Ada-MK 는 고정된 짧은 시퀀스 시나리오 (입력=64, 출력=12) 에서 일반 TensorRT-LLM 대비 최대 23.6%, vLLM 대비 **50.2%**의 단일 배치 처리량 향상을 달성했습니다.
• 일관성: 모든 테스트된 배치 크기 (1~16) 및 모델 버전에서 긍정적인 향상이 관찰되었습니다. 소배치 시나리오 (BS=1) 에서 Qwen3-1.7B 의 경우 vLLM 대비 50.2%, SGLang 대비 64.5% 의 개선이 있었습니다.
• 실제 세계 성능: CSL 및 Human-eval 데이터셋에서 Ada-MK 는 배치 크기 1 에서 8 까지의 가장 높은 처리량을 유지했습니다. 더 큰 배치 크기 (BS=16) 에서 vLLM 과 같은 시스템 수준 스케줄러 대비 우위는 축소되었습니다 (예: CSL 에서 3.5% 뒤처짐). 이는 저지연/소배치 설정에서는 커널 수준 융합이 지배적이지만, 고동시성/긴 컨텍스트 시나리오에서는 시스템 수준 스케줄링이 더 중요해짐을 나타냅니다.
• 지연 시간 감소: 이 접근법은 Qwen 모델 시리즈 전반에 걸쳐 vLLM, SGLang 및 일반 TensorRT-LLM 대비 엔드투엔드 지연 시간을 10%~50% 감소시켰습니다.

의의 및 주장

이 논문은 상업적 온라인 광고 시스템에서 MegaKernel 의 첫 번째 산업적 배포를 제시한다고 주장합니다. 주요 의의는 리소스가 제한된 NVIDIA Ada 와 같은 아키텍처에서 MegaKernel 채택을Previously 방해했던 이식성 - 효율성 트레이드오프를 해결하는 데 있습니다.

오프라인 DAG 기반 탐색을 통해 런타임 동적 의사결정을 컴파일 시간으로 끌어올림으로써, Ada-MK 는 자동 튜닝 솔루션의 분기 페널티 없이 결정론적이고 저지연 추론을 달성합니다. 저자들은 이 "오프라인 탐색, 온라인 재사용" 전략이 엄격한 지연 시간 제한 (1~5 ms) 이 필수적인 프로덕션 환경에서 고급 커널 융합 기술을 배포하기 위한 실현 가능한 경로를 제공한다고 강조합니다. 이 연구는 MegaKernel 이 고사양 Hopper/Blackwell 하드웨어로만 제한되지 않으며, 엄격하고 적응적인 최적화를 통해 이전 세대에 효과적으로 적응될 수 있음을 입증합니다.