Amazon сообщила, что SGLang-Neuron публично доступен в виде AWS Deep Learning Container. Выпустила его RadixArk совместно с Annapurna Labs. Плагин связывает фреймворк SGLang для обслуживания моделей с AWS Trainium - семейством AI-чипов компании.
Что он умеет
С этим плагином команды могут запускать модели на Trainium и генерировать rollout-ы для обучения с подкреплением. Модель поднимается через серверный интерфейс SGLang, а эндпоинт вызывается теми же API генерации, что и раньше. Под этим интерфейсом код работает на нативном PyTorch.
Как это устроено
SGLang-Neuron - внешний плагин (out-of-tree). После установки он регистрирует платформу neuron, а выбрать её при старте сервера можно переменной окружения SGLANG_PLATFORM=neuron. Поскольку плагин лежит вне основного проекта, поддержку Trainium и ядро SGLang можно развивать по отдельности. По словам AWS, так обе части двигаются быстрее.
Работа распределена так:
- Ядро SGLang сохраняет API-сервер, жизненный цикл запросов, RadixAttention и учёт префиксного кэша.
- Плагин добавляет модели, внимание (attention), поддержку mixture-of-experts и graph runner, рассчитанные именно на Trainium.
- TorchNeuron показывает Trainium как обычное устройство PyTorch. Модель выполняется на стандартных операциях PyTorch и ядрах, написанных на Neuron Kernel Interface. Всё компилируется через torch.compile.
Плагин также поддерживает разделение prefill и decode. Один сервер считает KV-кэш для промпта и передаёт его второму, который генерирует остальные токены. Передача идёт через NIXL поверх Elastic Fabric Adapter.
Что показано в примере
AWS демонстрирует настройку на инстансе Amazon EC2 с Trainium3. Нужно запустить контейнер, поднять сервер SGLang с моделью OpenAI gpt-oss-120b и отправить запрос chat completion на локальный эндпоинт. В статье есть ссылки на руководство для начала работы и на документацию в репозитории: там перечислены поддерживаемые модели, функции и проверенные конфигурации.
О чём источник умалчивает
В анонсе нет цен, нет сравнения производительности или стоимости с другим оборудованием и нет списка поддерживаемых моделей, кроме примера. Будет ли Trainium дешевле или быстрее для конкретной задачи, не подтверждено, это нужно проверять тестом.