WORLDTECHAmazon macht SGLang-Neuron für AWS Trainium öffentlich verfügbar
Amazon hat angekündigt, dass SGLang-Neuron nun öffentlich über einen AWS Deep Learning Container verfügbar ist. Entwickler können damit Modelle auf Trainium-Chips bereitstellen und Rollouts für Reinforcement Learning erzeugen. RadixArk hat das Plugin in Zusammenarbeit mit Annapurna Labs veröffentlicht.
Was die Veröffentlichung bietet
SGLang-Neuron ist ein außerhalb des Kerns entwickeltes Plugin, das das Framework SGLang mit AWS Trainium verbindet. Amazon beschreibt Trainium als Reihe von KI-Chips, die für Training und Inferenz in großem Maßstab wirtschaftlich attraktiv sind. Der Zugang läuft über einen AWS Deep Learning Container, den RadixArk in Zusammenarbeit mit Annapurna Labs herausgegeben hat.
Nach der Installation registriert das Plugin über die Einstiegspunkte von SGLang eine neuron-Plattform. Mit SGLANG_PLATFORM=neuron wird diese Plattform beim Start des Servers ausgewählt.
Aufteilung zwischen SGLang-Kern und Plugin
Der SGLang-Kern übernimmt weiterhin den API-Server, den Lebenszyklus der Anfragen, RadixAttention und die Verwaltung des Präfix-Caches. Das Plugin liefert die Trainium-spezifischen Modelle, Attention, Komponenten für Mixture of Experts (MoE) und den Graph-Runner.
Darunter läuft SGLang-Neuron auf nativem PyTorch. TorchNeuron meldet Trainium als natives PyTorch-Gerät. Die Modellausführung stützt sich auf Standard-PyTorch-Operationen sowie auf Kernel, die mit dem Neuron Kernel Interface (NKI) für bessere Leistung geschrieben wurden. Diese Operationen und Kernel werden mit torch.compile in ausführbare Programme übersetzt, die der Graph-Runner während des Betriebs vorbereitet und ausführt.
Einbindung in den bekannten SGLang-Workflow
Der Betrieb von SGLang auf Trainium behält den Workflow bei, den Entwickler bereits kennen. Modelle werden über die Server-Schnittstelle von SGLang gestartet, und der Endpunkt wird mit denselben Generierungs-APIs angesprochen.
Amazon beschreibt außerdem eine getrennte Konfiguration: Ein Prefill-Server erstellt den KV-Cache des Prompts und schickt ihn über NIXL und Elastic Fabric Adapter (EFA) an einen Decode-Server, der die restlichen Tokens erzeugt.