AWS, SageMaker AI의 MoE 추론 커널 최적화 3가지 공개

AWS가 Amazon SageMaker AI에서 전문가 혼합(MoE) 모델을 더 빠르게 돌리기 위해 적용한 커널 최적화 세 가지를 공개했다. 대표적인 MoE 모델을 예로 들어, 행렬 곱 방식을 바꾸는 방법, 연산을 맡길 백엔드를 바꾸는 방법, 불필요한 재컴파일을 없애는 방법을 설명했다. 측정은 NVIDIA H200 GPU에서 했다.

행렬 순서를 바꾼 최적화는 vLLM에 반영돼 0.24.0 이후 버전에서 쓸 수 있고, 백엔드 변경은 AWS Large Model Inference(LMI) v28 컨테이너에서 제공된다. 재컴파일 제거는 vLLM의 PR #42849에 해당한다.

MoE는 일부 전문가 네트워크만 활성화해 토큰당 계산량을 낮추면서 대규모 언어 모델의 용량을 키우는 구조다. MoE 연산은 전체 계산 시간의 큰 비중을 차지해, 여기서 생기는 비효율이 요청 지연 시간과 비용을 바로 끌어올린다. AWS는 그만큼 커널 최적화가 성능을 높일 기회가 가장 크다고 밝혔다.

핵심 기법인 Fused MoE는 토큰 라우팅, 데이터 재배열, 전문가 네트워크 계산을 GPU에서 하나의 배치 연산으로 묶는다. 이 과정을 따로따로 순차 반복하던 방식을 대체하며, 지연 시간에는 이점이 있지만 커널 구현과 디버깅이 복잡해진다. 이 구조에서 GPU 활용률을 높이려면 하드웨어에 맞춘 최적화가 필요하다.

Swap-AB는 NVIDIA Hopper 아키텍처의 WGMMA 명령 제약에서 출발한다. WGMMA는 128개 스레드로 이뤄진 워프 그룹이 함께 실행하며, 첫 번째 피연산자의 행 차원(M)이 64로 고정돼 있다. 하드웨어 타일 활용을 최대로 끌어올리려면 M이 64 이상이어야 한다.

디코드 단계에서는 사용자 시퀀스마다 토큰을 하나씩 처리하므로 M이 64보다 작고 1까지 내려가는 경우가 많다. 예컨대 행렬 A가 1×4096의 얇은 벡터가 되어 64행짜리 자리를 채우지 못하고 Tensor Core 용량 대부분이 비게 된다. 반면 전문가 가중치인 행렬 B는 4096×8192처럼 크다. Fused MoE에서는 선택된 전문가 수와 전문가 중간 크기를 곱한 값이 N이 되므로, topk=8이면 N이 수만에서 수십만 단위로 커진다.

AWS는 행렬 전치의 성질을 이용해 수학적 결과는 그대로 두고 곱셈 순서를 B^T × A^T로 바꿨다. 큰 전문가 차원(N=8192)이 첫 번째 피연산자 자리로 올라가 64행 조건을 쉽게 채우고, Hopper의 복사 엔진이 큰 가중치 블록을 공유 메모리에서 레지스터로 비동기 선읽기하기도 쉬워진다. 그 결과 H200의 대규모 MoE 모델에서 종단 간 출력 처리량이 2~4% 올랐다.

vLLM의 백엔드 선택은 MoE Oracle이 양자화 방식에 따라 Triton, FlashInfer+CUTLASS 등으로 연산을 보내는 방식이다. 프로파일링 결과 Hopper에서 FP8 텐서 양자화 모델은 기본값인 FLASHINFER_CUTLASS 경로보다 Triton 백엔드가 꾸준히 빨랐다. 커널 코드는 건드리지 않고 라우팅 정책만 바꿨는데, H200의 FP8 텐서 양자화 Fused MoE 모델에서 종단 간 처리량이 10% 개선됐다.

AWS는 Triton이 낮은 동시성의 MoE 작업에서 앞서는 이유를 두 가지로 들었다. 하나는 FP8 스케일링, 행렬 곱, 전문가 결과 합산을 한 번의 하드웨어 패스로 묶는다는 점이다. 데이터를 HBM에서 한 번만 읽어 GPU 내부 SRAM 레지스터에서 처리한 뒤 내보내므로, 중간 스케일링 단계마다 생기는 VRAM을 읽고 쓰는 왕복이 없다. 다른 하나는 CUTLASS가 크고 예측 가능한 행렬에 맞춰져 있어 M<64나 M=1 같은 얇은 배치에서 워프 활용이 떨어지거나 패딩 비용이 생기는 반면, Triton은 흩어진 토큰 배열을 활성 전문가 가중치에 동적으로 매핑한다는 점이다.

Triton 재컴파일 제거는 JIT 컴파일 부담을 겨냥했다. vLLM의 여러 커널 밑바탕인 Triton은 정수 인자를 기본적으로 특수화하기 때문에, 처음 보는 값이 들어올 때마다 JIT 컴파일이 다시 일어난다. 시퀀스 길이와 배치 크기가 계속 달라지는 서비스 환경에서는 이 부담이 쌓여 첫 토큰 시간을 늘린다.

회전 위치 임베딩과 FP8 양자화를 함께 처리하는 _fused_inv_rope_fp8_quant_per_head 커널이 num_tokens 값이 바뀔 때마다 다시 컴파일되고 있었다. 가변 길이 인자에 do_not_specialize 주석을 달자 재컴파일이 사라졌고, 8-way 텐서 병렬의 H200에서 첫 토큰 시간이 9% 개선됐다.

멀티 GPU 클러스터에서는 8개 GPU가 마이크로초 단위로 동기화해 커널을 실행해야 한다. GPU 하나가 토큰 배치 수 변화로 재컴파일하느라 멈추면 나머지 일곱 개도 기다린다. AWS는 클러스터 전체에서 재컴파일을 없애 GPU 8개가 멈추지 않고 동기화하게 했다.