SGLang-Neuron을 통한 AWS Trainium 기반 프런티어 AI 모델 배포

오늘 Amazon은 RadixArk가 Annapurna Labs와 협력하여 출시한 AWS Deep Learning Container를 통해 SGLang-Neuron을 공개적으로 사용할 수 있다고 발표했습니다. SGLang-Neuron을 사용하면 현재 사용 중인 것과 동일한 SGLang 인터페이스를 사용하여 AWS Trainium에서 모델을 서빙하고 강화 학습(RL) 롤아웃을 생성할 수 있습니다. Trainium은 대규모 AI 훈련 및 추론을 위한 강력한 경제성을 제공하는 AWS AI 칩 포트폴리오입니다. 이제 Trainium에서 SGLang을 실행하는 것은 이미 알고 있는 SGLang 워크플로를 따릅니다. 즉, SGLang의 서버 인터페이스를 통해 모델을 시작하고 동일한 생성 API를 사용하여 엔드포인트를 호출합니다. 해당 인터페이스 내부에서 SGLang-Neuron은 네이티브 PyTorch에서 실행됩니다. Out-of-tree 플러그인 설계를 통해 Trainium 지원과 SGLang 코어는 독립적으로 개발되고 더 빠르게 반복 작업이 가능합니다.

SGLang이 Trainium에서 작동하는 방식

SGLang-Neuron은 SGLang을 AWS Trainium과 연결하는 out-of-tree 플러그인입니다. 플러그인을 설치하면 SGLang의 플러그인 엔트리 포인트를 통해 neuron 플랫폼이 등록되며, 서버 시작 시 SGLANG_PLATFORM=neuron을 설정하여 이를 선택할 수 있습니다. SGLang 코어는 계속해서 API 서버, 요청 수명 주기, RadixAttention 및 프리픽스 캐시 부기(bookkeeping)를 제공합니다. 플러그인은 Trainium 전용 모델, 어텐션, 전문가 혼합(MoE) 및 그래프 러너를 구현합니다.

그림 1: AWS Trainium 기반의 SGLang-Neuron

TorchNeuron은 Trainium을 네이티브 PyTorch 디바이스로 등록합니다. SGLang-Neuron은 모델 실행 경로에 표준 PyTorch 연산을 사용하며 성능을 위해 Neuron Kernel Interface(NKI)로 작성된 커널을 통합합니다. 이러한 연산과 커널을 torch.compile을 사용하여 실행 가능한 프로그램으로 컴파일하며, 그래프 러너가 서빙 중에 이를 준비하고 실행합니다.

SGLang-Neuron은 PD 분리(disaggregation)도 지원합니다. 이 구성에서는 프리필(prefill) 서버가 프롬프트 KV 캐시를 계산하고 Elastic Fabric Adapter (EFA)를 통한 NIXL을 사용하여 디코드(decode) 서버로 전송합니다. 그런 다음 디코드 서버가 나머지 토큰을 생성합니다. 각 서버는 여전히 SGLang의 스케줄러와 요청 수명 주기를 사용하며, 플러그인은 Trainium 실행 및 KV 캐시 전송 경로를 관리합니다.