AWS는 Amazon SageMaker AI에서 혼합 전문가(MoE) 모델을 제공하기 위한 최적화 3종을 공개했으며, 처리량은 최대 10%, 첫 토큰 생성 시간(TTFT)은 최대 9% 개선됐다고 밝혔습니다. vLLM 커뮤니티와 함께 개발한 작업은 공개 소프트웨어로 배포됐고, SageMaker AI의 LMI 딥러닝 컨테이너에 적용됐습니다.

첫 번째 최적화는 행렬 곱셈의 피연산자 순서를 바꿔 NVIDIA Hopper GPU의 WGMMA 명령어 배치에 맞춥니다. WGMMA는 M 차원이 64행으로 고정돼 있지만, MoE 디코딩은 사용자 시퀀스당 토큰 하나를 처리하는 경우가 많아 M이 1에 머물고 텐서 코어 실행 영역 상당수가 유휴 상태가 됩니다. 융합 MoE에서 topk=8 설정은 선택된 전문가들의 중간 차원을 큰 N 차원 하나로 합칩니다. 곱셈을 전치하면 더 큰 차원이 첫 번째 피연산자에 놓여 64행 배치에 맞고, 비동기 사전 가져오기도 가능해집니다. AWS는 H200 GPU에서 대규모 MoE 모델의 전체 출력 처리량이 2~4% 개선됐다고 보고했습니다. 이 변경은 vLLM 0.24.0 이상에 반영됐습니다.

두 번째 변경은 Hopper GPU에서 FP8 텐서 양자화 MoE 연산을 기본 FlashInfer+CUTLASS 경로 대신 Triton으로 보냅니다. AWS의 프로파일링 결과, 낮은 동시성 작업에서는 Triton이 스케일 조정, 행렬 곱셈, 전문가 집계를 GPU 한 번의 처리 과정으로 묶어 중간 결과의 전역 메모리 읽기와 쓰기를 피하면서 더 나은 성능을 보였습니다. 동적 스케줄링도 M < 64 또는 M = 1 같은 작고 크기가 변하는 토큰 묶음을 CUTLASS의 기본 타일링보다 효율적으로 처리합니다. AWS는 커널 코드를 바꾸지 않고 경로 설정만 변경해 H200의 FP8 양자화 융합 MoE 모델에서 전체 처리량을 10% 높였다고 밝혔습니다. 이 변경은 AWS 대형 모델 추론(LMI) v28 컨테이너에서 이용할 수 있습니다.

세 번째 최적화는 토큰 수가 새로 나타날 때마다 Triton이 커널을 다시 컴파일하는 일을 막습니다. _fused_inv_rope_fp8_quant_per_head 커널은 num_tokens 값이 바뀔 때마다 재컴파일하고 있었으며, 길이가 변하는 매개변수에 do_not_specialize 주석을 추가해 반복되는 적시 컴파일을 없앴습니다. AWS는 H200의 8-way 텐서 병렬화에서 TTFT가 9% 개선됐으며, 한 GPU가 재컴파일하는 동안 다른 GPU들이 기다리는 일도 방지한다고 밝혔습니다. 이 최적화는 vLLM 0.22.0 이상에 반영됐습니다. 세 변경 사항을 적용한 SageMaker AI의 LMI 컨테이너에서는 처리량이 최대 10%, TTFT가 최대 9% 개선됩니다.

AWS는 Amazon SageMaker AI에서 혼합 전문가(MoE) 모델을 제공하기 위한 최적화 3종을 공개했으며, 처리량은 최대 10%, 첫 토큰 생성 시간(TTFT)은 최대 9% 개선됐다고 밝혔습니다. vLLM 커뮤니티와 함께 개발한 작업은 공개 소프트웨어로 배포됐고, SageMaker AI의 LMI 딥러닝 컨테이너에 적용됐습니다.

첫 번째 최적화는 행렬 곱셈의 피연산자 순서를 바꿔 NVIDIA Hopper GPU의 WGMMA 명령어 배치에 맞춥니다. WGMMA는 M 차원이 64행으로 고정돼 있지만, MoE 디코딩은 사용자 시퀀스당 토큰 하나를 처리하는 경우가 많아 M이 1에 머물고 텐서 코어 실행 영역 상당수가 유휴 상태가 됩니다. 융합 MoE에서 topk=8 설정은 선택된 전문가들의 중간 차원을 큰 N 차원 하나로 합칩니다. 곱셈을 전치하면 더 큰 차원이 첫 번째 피연산자에 놓여 64행 배치에 맞고, 비동기 사전 가져오기도 가능해집니다. AWS는 H200 GPU에서 대규모 MoE 모델의 전체 출력 처리량이 2~4% 개선됐다고 보고했습니다. 이 변경은 vLLM 0.24.0 이상에 반영됐습니다.

두 번째 변경은 Hopper GPU에서 FP8 텐서 양자화 MoE 연산을 기본 FlashInfer+CUTLASS 경로 대신 Triton으로 보냅니다. AWS의 프로파일링 결과, 낮은 동시성 작업에서는 Triton이 스케일 조정, 행렬 곱셈, 전문가 집계를 GPU 한 번의 처리 과정으로 묶어 중간 결과의 전역 메모리 읽기와 쓰기를 피하면서 더 나은 성능을 보였습니다. 동적 스케줄링도 M < 64 또는 M = 1 같은 작고 크기가 변하는 토큰 묶음을 CUTLASS의 기본 타일링보다 효율적으로 처리합니다. AWS는 커널 코드를 바꾸지 않고 경로 설정만 변경해 H200의 FP8 양자화 융합 MoE 모델에서 전체 처리량을 10% 높였다고 밝혔습니다. 이 변경은 AWS 대형 모델 추론(LMI) v28 컨테이너에서 이용할 수 있습니다.

세 번째 최적화는 토큰 수가 새로 나타날 때마다 Triton이 커널을 다시 컴파일하는 일을 막습니다. _fused_inv_rope_fp8_quant_per_head 커널은 num_tokens 값이 바뀔 때마다 재컴파일하고 있었으며, 길이가 변하는 매개변수에 do_not_specialize 주석을 추가해 반복되는 적시 컴파일을 없앴습니다. AWS는 H200의 8-way 텐서 병렬화에서 TTFT가 9% 개선됐으며, 한 GPU가 재컴파일하는 동안 다른 GPU들이 기다리는 일도 방지한다고 밝혔습니다. 이 최적화는 vLLM 0.22.0 이상에 반영됐습니다. 세 변경 사항을 적용한 SageMaker AI의 LMI 컨테이너에서는 처리량이 최대 10%, TTFT가 최대 9% 개선됩니다.