Представьте: у вас есть LLM-сервинг на NVIDIA H100, и вы хотите добавить поддержку AMD MI355X. Сколько кода придётся переписать? Команда LightSeek в новом посте на PyTorch Blog показывает, что ответ может быть «почти ноль» — если ядерная система спроектирована правильно. TokenSpeed-kernel, открытая подсистема для LLM-инференса, отделяет public API от hardware-specific реализаций: один и тот же вызов `mha_prefill` работает на AMD через Gluon-ядра и на NVIDIA через TRT-LLM, без единой строчки изменений в модели или рантайме.

Проблема: ядра становятся лабиринтом

Современный LLM-сервинг — это не одна быстрая операция, а десятки конкурирующих путей. Attention, MoE-маршрутизация, expert GEMM, квантизация, коммуникация между устройствами — всё это работает на GPU-ядрах, и каждое ядро выбирается под конкретную задачу. Проблема в том, что «лучшее ядро» — это не фиксированный ответ. Оптимальный выбор зависит от архитектуры модели, формы тензоров, формата квантизации, поколения GPU, наличия vendor-библиотек и от того, идёт ли трафик на prefill или decode.

Без чёткой системы управления ядрами движок накапливает пути: in-tree ядра, врапперы vendor-библиотек, экспериментальные ядра, архитектурно-специфичные fast paths, исторические fallback'ы. Добавить новую модель — значит затронуть несвязанные рантамные пути. Добавить новый silicon — значит провести device checks через все слои модели. Ядерная разработка превращается в考古 work, потому что поведение модели, dispatch логика, выбор backend'а и реализация ядра переплетены за неявной границей.

Когда в инференс-движке нет чёткого ядерного слоя, каждое новое поколение GPU или каждый новый вендор приносит свои библиотеки и форматы. Допустим, вышла новая AMD-карта с новыми tensor core фичами. В классической архитектуре нужно писать conditional checks по всей модели, перебирать пути в runtime, возможно форкать целые ветки serving-кода. Инженеры тратят недели на integration и regression testing вместо того, чтобы писать ядра.

Архитектура: три слоя вместо лабиринта

TokenSpeed-kernel строится вокруг трёх практических принципов. Первый: мультисиликоновая поддержка должна быть фундаментальной, не добавочной. Вторая система должна понимать возможности платформы напрямую, а не как разрозненные conditionals — одна и та же операция может иметь несколько решений для разных silicon targets, и все должны конкурировать через единый selection system. Вторая: портативность и производительность должны сосуществовать. Новой модели нужна портативная точка входа, которая работает сразу, а затем может постепенно переходить на оптимизированные ядра. Третий: быстрая итерация ядер требует guardrails. Путь от идеи до production-дployment должен быть коротким, с минимумом зависимостей и понятными точками бенчмаркинга.

Из этих принципов вытекает архитектура сверху вниз. Рантам входит через generic public API — `mha_prefill`, `mha_decode_with_kvcache`, `moe_apply`. Эти вызовы платформо-agnostic: никакого «AMD-пути» или «Triton-пути» в коде модели. Селектор фильтрует зарегистрированные ядра по platform capability, формату тензоров и traits, ранжирует оставшиеся и возвращает callable. Backend-реализации регистрируются через декоратор `@register_kernel`, объявляя семейство операторов, имя решения, требования к платформе, поддерживаемые tensor signatures, traits и приоритет. Все ядра — и портативные Triton-пути, и высокопроизводительные Gluon для AMD, и вендорные врапперы — конкурируют через одну систему селекции.

Registry — единый источник правды о доступных реализациях. Каждое зарегистрированное ядро описывается метаданными: семейство операторов и режим, принадлежность к решению, capability платформы, форматы тензоров, traits, приоритет относительно других кандидатов. При запросе ядро фильтруется по platform capability, формату и traits, оставшиеся ранжируются. Для фиксированной модели, платформы и форматов результат обычно стабилен, и TokenSpeed-kernel кэширует resolved callable. При этом можно принудительно указать конкретное решение для отладки и бенчмаркинга — например, чтобы проверить, действительно ли кастомное ядро быстрее стандартного.

Numerics, бенчмарки и плагины

Ядерная система — это не только dispatch. Она даёт авторам ядер workflow для безопасной и быстрой итерации: проверки численности, standalone-бенчмарки и профилирование. Reference implementations дают общую цель для correctness, бенчмарки дают timing и reporting path вне полного сервера, профилирование делает выбранные kernel names и key parameters видимыми в end-to-end traces.

Тот же boundary поддерживает out-of-tree плагины. Плагин регистрирует ядра через тот же декоратор, назначает свой приоритет и участвует в normal selection наравне с in-tree имплементациями. Это держит core package чистым, при этом оставляя место для hardware vendors, researchers и deployment-команд, чтобы привносить специализированные ядра без форка всей системы. Эти эргономические решения важны так же, как dispatch. Пакет остаётся pip-installable и dependency-conscious: специализированные ядра легко установить, проверить, замерить и заменить.

TokenSpeed-kernel предоставляет и CLI, и programmatic interfaces для основных задач разработки: численная верификация и standalone бенчмаркинг. Эти инструменты переиспользуют тот же registry metadata, что serving использует для kernel selection, поэтому зарегистрированное ядро может быть проверено против reference implementation, замерено на стандартных или кастомных shapes, опционально профилировано, а затем автоматически выбрано, когда его capabilities и traits совпадают с runtime-запросом.

Числа: Gluon на AMD MI355X

GPT-OSS 120B — хороший стенд для валидации дизайна: это современная LLM, которую можно запустить на одной GPU, при этом она нагружает и attention, и MoE. Attention использует регулярный MHA с attention sinks и миксом sliding-window и full-attention слоёв. AMD-деплоймент использует MXFP4 expert weights и FP8 activation flow для MoE. Именно те детали, которые обычно «протекают» в рантамный код, если граница между ядрами и runtime недостаточно чёткая.

Для AMD-пути ключевые ядра реализованы на Gluon — Triton-family DSL с явным контролем производительности. Gluon даёт прямой доступ к CDNA4-фичам: async copies, shared-memory layouts, scaled MFMA для FP8/MXFP форматов. Все фичи — явные программные примитивы, не скрытые компиляторные оптимизации. Ядерные авторы могут выбирать layouts (BlockedLayout или DistributedLinearLayout), аллоцировать shared memory с SwizzledSharedLayout или PaddedSharedLayout для избежания bank conflict, выбирать AMD matrix-core layouts через AMDMFMALayout. AMD Gluon modules экспортируют операции, которые близко маппятся к hardware: `mfma`, `mfma_scaled`, `buffer_load`, `buffer_store`, async global-or-buffer loads into shared memory.

Attention prefill на Gluon — текущий fastest evaluated MI355X backend на 14 из 15 измеренных GPT-OSS prefill shapes, 1.4–2.3× быстрее Triton baseline. Против AITER — 1.1–1.3× uplift. На графике с TFLOP/s для prefill throughput при BF16 Q/K/V, head_dim=64, 64 Q heads, 8 KV heads, full causal prefill — Gluon-версия стабильно выигрывает по всему диапазону sequence lengths от 1K до 8K при batch sizes от 1 до 16.

MoE — где разделение слоёв окупается ещё больше. GPT-OSS MoE layer — это не одна dense matrix multiply, а структура: маршрутизация токенов к экспертам, gather/dispatch, expert GEMMs, применение активации, комбинация top-k выходов с весами маршрутизации. AMD Gluon MoE path построен вокруг этой полной структуры, а не как две изолированные GEMM. Для prefill ключевой вызов — держать CDNA4 Compute Units занятыми при неравномерном распределении токенов по экспертам. Имплементация использует ragged block schedules, где работа следует за реальным распределением экспертов, потом выбирает tile shapes из логического количества токенов и per-expert slice size.

При самых малых batch size Gluon-ядра работают в 1.7–2.1× быстрее Triton и в 1.1–1.6× быстрее AITER. В medium decode band AITER слегка вырывается вперёд, но Gluon остаётся в пределах 0.9× от fastest, при этом всё ещё 1.3–1.4× быстрее Triton. На full-MoE latency (routing + оба GEMM + SwiGLU + combine) при MXFP4 weights и FP8 activations на MI355X — Gluon выигрывает по всему диапазону от M=1 до M=8192.

End-to-end: 3.6× на AMD MI355X

Результат end-to-end: на AMD MI355X Gluon-backed path показывает 1.6×–3.6× ускорение output throughput над портативным Triton path на всех измеренных точках. Это ускорение получено без отдельного AMD-специфичного serving path. AMD performance приобретена через имплементацию тех же public attention и MoE contracts со специализированными Gluon-ядрами, регистрацию их platform и shape constraints, и dispatch через selector. Модель не знает ничего про MI355X архитектурные детали, как MXFP4 scales следует расположить для CDNA4, или какое AMD-ядро fastest на конкретном prefill/decode attention case. Оно только передаёт правильные тензоры и metadata в public API.

На NVIDIA пути в текущей GPT-OSS Blackwell конфигурации attention использует TRT-LLM backend через FlashInfer-экспоз TensorRT-LLM wrappers, а MXFP4 MoE использует flashinfer_trtllm solution. Рантам по-прежнему вызывает исключительно `mha_prefill`, `mha_decode_with_kvcache` и `moe_apply`. Мультисиликоновая поддержка — это не два несвязанных стека. AMD и NVIDIA support — sibling implementations за одним и тем же kernel API, registry и selection model.

Почему это важно для инфраструктуры

Гетерогенность в LLM-инференсе — это уже не будущее, а настоящее. Blackwell, CDNA4, будущие архитектуры — инфраструктуре придётся работать с несколькими поколениями GPU одновременно. Без системы, которая отделяет ядра от рантамной логики, каждое добавление нового silicon target означает рефакторинг serving-кода и перепроверку всех моделей. TokenSpeed-kernel делает ядра first-class subsystem вместо коллекции спрятанных fast paths.

Практический выхлоп: ядра, оптимизированные для AMD, уже подхвачены vLLM как отдельный пакет tokenspeed-kernel-amd — без зависимости от полного TokenSpeed serving stack. Это означает, что vendor'ы, researchers и deployment-команды могут привносить специализированные ядра без форка всей системы. Clean separation также позволила опубликовать TokenSpeed-kernel как standalone packages, которые можно установить и использовать отдельно — целиком или по отдельности для разных ядер. Это делает kernel packages полезными для всего ecosystem, не только для TokenSpeed.

FAQ

Чем это отличается от vLLM или SGLang?

vLLM и SGLang — это serving engines с встроенными ядерными решениями. TokenSpeed-kernel — это ядерный слой, который может работать поверх любого engine. Он не заменяет SGLang или vLLM, а предоставляет стандартизированное API, через которое эти engines получают доступ к hardware-specific ядрам на AMD и NVIDIA без переписывания dispatch логики.

Почему AMD, а не только NVIDIA?

NVIDIA доминирует в LLM-инференсе, но AMD MI355X предлагает конкурентную стоимость на FLOPS и активно развивается. Если inference становится более гетерогенным (а он становится), иметь AMD как first-party target, а не third-party fallback — значит получить redundancy и потенциально лучшую стоимость на dollar-per-TFLOP для определённых конфигураций.

Gluon — это Triton?

Gluon относится к Triton-семейству DSL, то есть имеет похожий programming model с block-level контролем, но с прямым доступом к AMD-специфичным фичам CDNA4 (async copies, MFMA, bank conflict avoidance через swizzling). На NVIDIA аналогом будет TRT-LLM через FlashInfer-экспоз. Оба — DSL-подходы, но с разными hardware targets как primary citizens.

Что это значит для разработчика ядер?

Разработка ядра для нового silicon сейчас обычно означает integration в конкретный serving engine, с неявными зависимостями от dispatch логики и модельного кода. TokenSpeed-kernel делает путь короче: написал ядро, проверил на reference, замерил на бенчмарках, зарегистрировал с metadata — и оно автоматически участвует в selection. Не нужно форкать serving engine или переписывать dispatch.

Итог

TokenSpeed-kernel решает конкретную проблему: fragmentation ядерных реализаций при мультисиликоновом inference. Архитектура с registry и selection механизмом позволяет добавлять новые hardware targets не через fork и переписывание serving-кода, а через регистрацию новых ядер с metadata и traits. AMD MI355X с Gluon-ядрами демонстрирует, что это не теория: 3.6× ускорение на реальном production workload на всех измеренных точках. По мере того как рынок инференса становится гетерогенным, такой подход становится не nice-to-have, а необходимостью. Оригинал на PyTorch Blog.