SGLangがAWS Trainium向けプラグインBetaを発表：専用AIチップ上で推論とRL Rolloutを実現

過去1年間、SGLangチームはサポートするハードウェアエコシステムの拡張を継続しながら、新規オープンモデルのリリース初日サポートや、MilesにおけるRL rollout機能の強化を進めてきた。このたびSGLangは、AWSとの連携によりAWS Trainiumへのサポートを正式に発表した。

フロンティアAIが異種ハードウェアアーキテクチャへと移行する中、AWS Trainiumは注目に値する機会を提供している。TrainiumはAWSがAIワークロード向けに専用設計したチップであり、クラウドサービスとして大規模に提供されている。高スループットシナリオにおけるより優れたコストパフォーマンスを目標としており、オープンモデルのサービング・RL rollout生成はまさにその典型的なユースケースだ。これらは継続的な稼働を必要とすることが多く、トークンあたりのコストが、チームが展開できる算力規模を直接左右する。

SGLangチームは、TrainiumでこれらのワークロードをRunする際のファーストクラスな選択肢となることを目指すと表明した。AWSとの緊密な連携およびTorchNeuronのリリースを背景に、チームは新たにリリースされたTorchNeuronを基盤とするSGLang Trainium plugin Betaを正式発表した。このBeta版はTrainium上での推論サービスとRL rolloutの両方をサポートする。

最初からTorchNeuronをベースに構築

TorchNeuronはAWSがTrainium向けに提供するネイティブPyTorchバックエンドだ。SGLangチームはTorchNeuronのプライベートプレビュー段階からNeuronバックエンドの構築に着手しており、この技術的選択が適応作業全体の方向性を決定づけた。

TorchNeuronはネイティブPyTorchの手法を採用しているため、SGLangのモデルコンポーネントとコア抽象レイヤーの大部分をTrainiumへほぼそのまま移植できる。Trainium上でモデルをサービングする際も、GPU上と同一のインターフェースを使用できるため、少ない変更でGPUとTrainium間のワークロード移行が可能だ。

さらなる性能最適化においては、NKI（Neuron Kernel Interface）により、開発者はPyTorchのワークフロー内でカスタムTrainiumカーネルを記述できる。複数のハードウェアプラットフォームをカバーするオープンソースプロジェクトとして、使い慣れたプログラミングモデルを提供しながら、パフォーマンスクリティカルなパスでは低レベルのハードウェア制御能力も維持している。

Beta版が提供する機能

今回のBeta版はAWS Trainium2およびTrainium3上で2種類のコアワークロードをサポートする。LLM推論サービングと、後学習（post-training）向けのRL rollout生成だ。初期サポート範囲は以下の通りだ：

• モデル：Qwen3（dense）、Qwen3-MoE-30B-A3B、DeepSeek-V2-Lite、GLM-4、GLM-5.2/5.3。チームはより大規模なフロンティアモデルへのサポートも進行中であると述べている。
• 機能：Data Parallel、Tensor Parallel、Expert Parallel、Radix Cache、BF16 / FP8、Async scheduler。
• RL rollout：SGLangはTrainium上で異種RL pipelineのrollout engineとして機能し、後学習プロセスを支援できる。

SGLangチームはBeta期間中もモデルと機能のサポート範囲を継続的に拡大する予定であり、今年後半にはより広範な提供に向けた取り組みを進める見通しだ。

早期プレビュー：異種RLワークフロー

推論サービングに加え、TrainiumはRL後学習における新たなデプロイ構成の可能性をもたらす。各RLステップは通常2つの異なるフェーズで構成される。1つ目はrollout generationであり、これは本質的にデコーディング主体の推論負荷だ。2つ目はpolicy updateであり、これはトレーニング負荷に当たる。多くのチームはこの2つのフェーズを同一のGPUクラスター上で実行しているが、両者のハードウェア要件は異なる。

SGLang Trainium pluginはこの2つのフェーズを異なるハードウェア上で実行することを可能にする。初回のエンドツーエンド実行では、SGLangはTrainium3上でQwen3-MoE-30B-A3BのrolloutをGenerateし、同時にSGLangチームのRL後学習フレームワークMilesがMegatronを用いて単一のH200ノード上でGRPO policyのトレーニングを実施した。MilesはTrainium上のSGLangを外部rollout engineとして接続し、安定したトレーニングRunを完了した。

このアプローチにより、チームはRL算力のプランニングにおいてより高い柔軟性を得られる。各フェーズはそれぞれのコストパフォーマンスに応じてハードウェアを選択でき、Trainium上のrollout capacityをGPUトレーニングリソースとは独立してスケールすることも可能だ。

ただし、SGLangチームはこれがまだ初期のマイルストーンであることを強調している。Trainium上の性能最適化は継続中であり、チームは11月のSGLang Summitで続報を共有する予定だ。

Betaへの参加方法

現在の初期Beta版は少数のチームに限定して公開されており、SGLangチームが各参加者と緊密に連携できる体制を維持している。このBetaは特に以下の2種類のチームに適している：

• Trainium2 / Trainium3上でLLM推論を実行中または計画中のチーム；
• 異種RLに関心を持ち、Trainium上でrolloutをGenerateしてGPU上でトレーニングを行うことを望むチーム。

アクセス権の申請を希望するチームは、公式フォームからSGLangチームに連絡できる。