SGLang-TorchTPU가 TPU v7에서 8160억 파라미터 MoE를 가속한 방법

들어가며

대규모 모델 추론 속도는 행렬 곱셈 성능만으로 결정되지 않는다. 특히 배치가 작은 디코드 단계에서는 연산 장치가 빠르게 작업을 끝내도 HBM에서 가중치가 도착하기를 기다리거나, KV 캐시를 읽거나, 장치 간 통신이 끝나기를 기다리는 시간이 발생한다. 레이어마다 별도 커널을 실행하면 커널 경계에서 가중치 스트리밍과 계산의 중첩도 끊어진다.

PyTorch의 기술 블로그는 SGLang-TorchTPU가 이 문제를 다루는 방식을 설명했다. SGLang은 스케줄링, KV 캐시, 모델 서비스 기능을 담당하고 TorchTPU는 PyTorch 백엔드로 사용된다. 주요 경로에는 Pallas 커널을 배치했으며, 4개 TPU v7 칩으로 구성된 호스트 한 대에서 Meta의 Muse Spark Open을 실행한다. 이 호스트는 8개 디바이스로 노출되고 모델은 텐서 병렬 방식으로 8분할된다.

Muse Spark Open은 총 8160억 개의 파라미터를 가진 Mixture-of-Experts 모델이다. 토큰 하나를 처리할 때 약 420억 개의 파라미터가 활성화된다. 62개 디코더 레이어가 모두 MoE 구조이며, 각 레이어는 256개 전문가 중 8개를 선택한다. 어텐션은 최근 2048개 토큰을 보는 슬라이딩 윈도와 최대 약 100만 토큰까지 처리하는 전체 컨텍스트 어텐션을 함께 사용한다. 따라서 연산 성능뿐 아니라 가중치 이동과 KV 캐시 용량도 중요한 변수가 된다.

핵심 최적화

• 전체 단계를 하나의 Megakernel로 통합했다. 기존 경로는 레이어마다 여러 커널과 집단 통신을 실행했다. 새 Pallas 프로그램은 62개 레이어와 마지막 정규화를 디바이스별 단일 호출에 넣는다. 임베딩, 초안 모델, LM 헤드, 샘플링은 바깥에 남지만, 핵심 디코드 또는 검증 단계는 하나의 연속적인 실행 영역이 된다.
• 데이터 이동을 파이프라인으로 처리한다. 밀집 가중치와 전문가 가중치, KV 데이터, 부분 합, 장치 간 전송을 실제 계산보다 앞서 DMA로 시작한다. 전송 자체의 시간을 줄이는 것뿐 아니라 MXU와 벡터 유닛의 계산과 겹쳐 대기 시간을 핵심 경로에서 숨기는 것이 목표다.
• 디코드 경로를 세부적으로 개선했다. 집단 통신 접기, KV 청크를 미리 처리하는 max-free 어텐션 루프, 일반 경로에서 레지스터 스필을 줄이는 전문가 루프를 사용한다. DFlash2 추측 디코딩에서는 작은 모델이 여러 토큰을 먼저 제안하고 전체 모델이 한 번에 검증한다.
• KV 캐시를 계층화했다. Megakernel은 시퀀스를 레인에 배치하는 KV 레이아웃을 사용한다. TPU의 L2 HiCache는 호스트 DRAM을 추가 캐시 계층으로 활용하며, 세션을 고려한 제거 정책으로 긴 세션이나 공유 프리픽스에 유용한 데이터를 보존하도록 설계됐다.
• Prefill 경로도 최적화했다. 융합 MoE 커널과 시퀀스 병렬화를 사용해 긴 입력을 처리하는 prefill 성능을 개선했다. 토큰 단위 디코드와는 병목이 다르기 때문에 별도의 최적화가 필요하다.

성능 결과와 해석

이전 Muse Spark 체크포인트에서 DFlash2 검증 경로를 측정한 결과, Megakernel을 사용한 단일 요청의 출력 토큰당 시간은 1.95밀리초였다. 기존 경로는 3.78밀리초로, 약 1.94배의 개선이다. 동시 요청이 2개일 때의 개선 폭은 1.85배였다. 동시성이 높아지면 차이가 줄어드는데, 배치화만으로도 하드웨어 활용률이 높아져 커널 경계의 상대적인 비용이 감소하기 때문이다.

검증 단계의 시간은 해당 단계에 필요한 고유 가중치를 HBM에서 읽는 데 걸리는 시간의 약 1.5~1.7배 수준까지 낮아졌다. 이는 레이어별 실행과 집단 통신 오버헤드가 성능을 지배하던 상태에서 메모리 대역폭이 주요 한계가 되는 상태에 가까워졌다는 의미다.

팀은 에이전트 코딩 세션을 기반으로 한 100만 토큰 AgentX 트레이스 세트에서도 평가했다. 4칩 TPU v7 호스트에서 FP8 연산과 NVFP4 전문가 가중치를 사용했으며, 높은 상호작용성을 중시하는 낮은 총 처리량 조건부터 요청당 출력 속도는 낮지만 전체 처리량이 높은 조건까지 여러 운용 지점을 확인했다. 다만 비용과 GB200 비교 결과는 특정 하드웨어, 모델, 비용 산정 방식에 따른 것이므로 모든 클라우드 환경에 적용되는 보편적인 순위로 해석해서는 안 된다.

이 사례의 가장 큰 의미는 추론 최적화의 단위를 개별 연산자에서 전체 실행 단계로 확장했다는 점이다. MoE와 초장문 컨텍스트 모델에서는 가중치, KV 데이터, 통신, 스케줄러 사이의 경계가 함께 성능을 결정한다. 이들을 하나의 파이프라인 가능한 프로그램에 배치하면 실제 서비스 병목에 더 직접적으로 대응할 수 있다.

물론 대가는 구현 복잡성이다. TPU의 벡터 레이아웃과 VMEM 용량을 맞춰야 하고, 페이지 기반 KV 캐시와 동적인 요청 흐름도 지원해야 한다. 따라서 이 방법이 더 많은 모델로 확장되려면 컴파일러, 런타임, 모델 구현, 서빙 프레임워크 간의 긴밀한 연동이 필요하다. 그럼에도 이번 결과는 추론 성능이 개별 커널의 최고 속도보다 시스템 전체의 유휴 시간을 얼마나 줄이느냐에 달려 있음을 보여준다.

출처: PyTorch Blog