Amazon anunció que SGLang-Neuron ya está disponible públicamente a través de un AWS Deep Learning Container. RadixArk lo publicó en colaboración con Annapurna Labs. El plugin conecta el framework de servicio SGLang con AWS Trainium, la familia de chips de IA de la compañía.
Qué hace
Con el plugin, los equipos pueden servir modelos en Trainium y generar rollouts de aprendizaje por refuerzo. Lanzan un modelo mediante la interfaz de servidor de SGLang y llaman al endpoint con las mismas APIs de generación de siempre. Por debajo de esa interfaz, el código se ejecuta sobre PyTorch nativo.
Cómo está construido
SGLang-Neuron es un plugin externo al proyecto principal (out-of-tree). Una vez instalado, registra una plataforma neuron, y basta con definir la variable de entorno SGLANG_PLATFORM=neuron para seleccionarla al arrancar el servidor. Como queda fuera del proyecto principal, el soporte de Trainium y el núcleo de SGLang pueden desarrollarse por separado, lo que según AWS permite que ambos avancen más rápido.
El reparto del trabajo es este:
- El núcleo de SGLang conserva el servidor de API, el ciclo de vida de las peticiones, RadixAttention y la gestión de la caché de prefijos.
- El plugin aporta los modelos específicos para Trainium, la atención, el soporte de mixture-of-experts y el graph runner.
- TorchNeuron hace que Trainium aparezca como un dispositivo nativo de PyTorch. La ejecución del modelo usa operaciones estándar de PyTorch más kernels escritos en Neuron Kernel Interface. Todo se compila con torch.compile.
El plugin también admite la desagregación de prefill y decode. Un servidor calcula la caché KV del prompt y se la pasa a un segundo servidor, que genera el resto de los tokens. La transferencia usa NIXL sobre Elastic Fabric Adapter.
Qué muestra el tutorial
AWS hace la demostración en una instancia Amazon EC2 con Trainium3. Se inicia el contenedor, se lanza el servidor de SGLang con el modelo gpt-oss-120b de OpenAI y se envía una petición de chat completion a un endpoint local. La publicación remite a una guía de inicio y a la documentación del repositorio para consultar modelos, funciones y configuraciones validadas.
Lo que la fuente no dice
El anuncio no incluye precios, ni comparaciones de rendimiento o coste con otro hardware, ni una lista de modelos compatibles más allá del ejemplo. Si Trainium resulta más barato o más rápido para una carga concreta no está confirmado y hay que comprobarlo con pruebas.