TokenSpeed-kernel: открытая подсистема для переносимых и производительных LLM‑ядер на разных типах чипов

Разработчики представили TokenSpeed-kernel, автономную open-source подсистему для упрощения бэкенда при инференсе LLM. Она предлагает многоуровневый API и систему реестров. Эти компоненты разделяют высокоуровневый runtime и низкоуровневый код, который зависит от конкретного оборудования. В итоге один и тот же интерфейс можно использовать на разных платформах, при этом разработчики ядер сохраняют возможность глубоко оптимизировать код под конкретные устройства.

Как сообщает PyTorch Blog, в архитектуре TokenSpeed runtime занимается выполнением модели, хранит метаданные планирования, управляет page table и состоянием маршрутизации. TokenSpeed-kernel отвечает за другой слой: API операторов, регистрацию и выбор бэкендов, численные аспекты, а также инструменты бенчмаркинга и профилирования. Платформенно-специфичная оптимизация вынесена в отдельные ядра. В качестве примера приводится GPT-OSS. Runtime вызывает одни и те же публичные API TokenSpeed-kernel независимо от платформы, а производительность на AMD и NVIDIA достигается за счет подключаемых ядер. Для AMD GPT-OSS 120B такая схема показывает top-of-the-line производительность при использовании Gluon kernels.

Подсистема распространяется в виде отдельных пакетов. Их можно устанавливать и применять независимо от остальной части TokenSpeed, полностью или только для отдельных ядер. Авторы рассчитывают, что эти пакеты окажутся полезны более широкой экосистеме как переносимый и производительный набор ядер для multi-silicon инференса. В подобных системах именно ядра определяют скорость и эффективность стека. Здесь выполняются attention, MoE routing, expert GEMMs, communication, quantization и sampling. TokenSpeed-kernel собирает связанную с этим сложность на одном уровне архитектуры.