Kimi K3 на графических процессорах AMD Instinct с использованием TokenSpeed

TL;DR

• Агентный вывод с поддержкой различных вендоров, в основе которого лежит AMD: TokenSpeed объединяет планировщик на C++, уровень выполнения на Python и модульные API ядер, обеспечивая унифицированное управление кэшем и выполнение на различных ускорителях.
• Лидирующая производительность Kimi K3 на 8x AMD Instinct™ MI355X GPU: 216 выходных токенов/с на пользователя при параллелизме 1, что в 1,4 раза выше, чем у ATOM, при сопоставимой производительности на более высоком уровне параллелизма и более высокой точности.
• Использование фундамента Triton и оптимизации Gluon с поддержкой агентов: специализированные ядра, слияние и коммуникация Iris обеспечивают до 5,1 раза более быстрый префилл MLA, в 2,1 раза более быстрое спекулятивное внимание и в 4,4 раза более быстрый MoE all-reduce по сравнению с базовыми показателями, описанными ниже.

TokenSpeed: Агентный вывод с AMD в основе

TokenSpeed — это современный движок для вывода LLM с открытым исходным кодом, разработанный для генерации с низкой задержкой в агентных рабочих нагрузках. Агенты, занимающиеся написанием кода, постоянно возвращаются к длинным диалогам, обрабатывают результаты новых инструментов и генерируют следующее действие, пока пользователь ожидает. Для их эффективного обслуживания требуются низкая задержка декодирования, эффективное повторное использование префиксов и предсказуемое планирование.

TokenSpeed четко распределяет задачи для каждой части этой проблемы. Уровень управления на C++ управляет жизненным циклом запросов и владением кэшем как конечный автомат, где система типов обеспечивает безопасное повторное использование ресурсов. Уровень выполнения на Python поддерживает быструю разработку моделей. Многоуровневая подсистема ядер связывает переносимый код модели со специализированными реализациями через общие API, центральный реестр и механизм плагинов для ускорителей. Кэширование префиксов, спекулятивное декодирование, выполнение графов и дезагрегированное обслуживание строятся на этих общих механизмах.

Поддержка AMD является фундаментальной для этой архитектуры. TokenSpeed обеспечил работу Kimi K3 на CDNA4 at Day 0 с кэшированием префиксов, спекуляцией, дезагрегированным обслуживанием и декодированием с захватом графа. TML Inkling также появилась на Day 0 с поддержкой AMD, контрольной точкой AMD Quark MXFP4 и специализированными ядрами Gluon. Та же система ядер теперь охватывает MHA, MLA, DSA, KDA, MoE и GEMM на AMD. Эти результаты показывают, как общая инфраструктура моделей и планирования может поддерживать новые архитектуры, в то время как специализированные ядра эффективно используют аппаратное обеспечение.

Kimi K3 — это передовая агентная модель, и в этой статье показано, как TokenSpeed обслуживает ее длинные диалоги на MI355X, особенно в сценариях с высокой интерактивностью. Приведенные ниже измерения для многоходовых диалогов показывают 216 выходных токенов/с на пользователя при параллелизме 1 на восьми графических процессорах MI355X.

Наш процесс оптимизации начинается с переносимых реализаций Triton в качестве функциональной базы и численного эталона. Затем профилирование направляет разработку специализированных ядер Gluon с помощью агентов, которые реализуют те же интерфейсы операторов и соответствуют тем же численным контрактам. Этот подход распространяется на весь стек: объединение проекций, удержание промежуточных данных на чипе и слияние вычислений в ядра Iris — расширение Triton, которое позволяет программировать ядра для нескольких GPU.

Производительность Kimi K3 в длинных агентных диалогах

Агентные рабочие нагрузки по написанию кода постоянно возвращаются к длинному диалогу, пока пользователь ждет следующего вызова инструмента. Чтобы измерить этот паттерн, мы воспроизвели многоходовые сессии кодинга SWE-smith с помощью агентного бенчмарка TokenSpeed с открытым исходным кодом. Каждый диалог начинается с промпта объемом ~50 тыс. токенов, добавляет ~800 токенов за ход в течение 10–15 ходов и генерирует 500 токенов на ответ. Существующий запуск использует восемь графических процессоров AMD Instinct MI355X, TP8, EAGLE3, FP8 KV и частоту попаданий в кэш около 90%.

При одной активной сессии этот запуск показывает 216 выходных токенов/с на пользователя, или 4,62 мс на выходной токен (TPOT), и среднее время до первого токена (TTFT) 0,62 секунды. В этой ситуации высокой интерактивности, к которой активно стремится индустрия, был достигнут результат в 1,4 раза выше по сравнению с ATOM. Производительность сопоставима или немного выше, чем у ATOM, при масштабировании до параллелизма 2, 4 и 8. При параллелизме 16 (C16) каждый агент получает 39 выходных токенов/с.

Рисунок 1: Агентное обслуживание Kimi K3 на восьми графических процессорах MI355X. Скорость декодирования на пользователя рассчитывается как 1000, деленное на время на выходной токен в миллисекундах; оно исключает время до первого токена и время в очереди. Общая пропускная способность токенов учитывает входные и выходные токены, включая кэшированные промпты. Каждая точка — это отдельный запуск при одном параметре параллелизма; условия обслуживания описаны в разделе конфигурации.

ATOM лидирует при C16. При C16 ATOM переключается на другой подход к планированию, называемый объединением префиллов. Это задерживает запросы на префилл, позволяя проходить большему количеству запросов на декодирование, пока не накопится достаточно запросов на префилл для пакетной обработки. Это жертвует TTFT (среднее значение 1,44 -> 2,90 с) ради TPOT (среднее значение 25,17 -> 19,28 мс) по сравнению с планированием по умолчанию, что приводит к общему выигрышу в пропускной способности, однако p95 TTFT запросов пользователей в наших измерениях увеличивается в 3,54 раза (2,95 -> 10,46 с). ATOM также проводит более агрессивную квантование, как показано в следующей таблице:

Компонент модели

TokenSpeed

ATOM

Проекции Attention Q/K/V и выходные проекции

BF16/BF16

FP8/FP8

Проекции общих экспертов

BF16/BF16

FP8/FP8

GEMM маршрутизируемых экспертов

FP8/MXFP4 (малое декодирование); MXFP4/MXFP4 (префилл)

MXFP4/MXFP4 (префилл)

Рекуррентное состояние KDA

FP32

FP16

Кэш MLA KV

FP8

FP8

Архитектура Kimi K3 и путь выполнения на восьми GPU

Большие языковые модели генерируют текст по одному токену за раз. Каждый новый токен проходит через механизм внимания (attention), который считывает диалог, и сеть прямой связи (feed-forward), которая преобразует текущее представление. Kimi K3 имеет 93 слоя декодера: 69 используют Kimi Delta Attention (KDA), а 24 используют gated Multi-head Latent Attention (MLA). Первый слой прямой связи является плотным; остальные 92 используют Latent MoE. Остаточные связи внимания (AttnRes) соединяют представления между блоками слоев (Рисунок 2).

Рисунок 2: Архитектура Kimi K3 и путь выполнения на восьми GPU. Номера разделов показывают, где рассматривается каждая часть слоя.

Мы используем восьмикратный тензорный параллелизм для внимания и MoE, без параллелизма экспертов (attention TP8; MoE TP8/EP1). Головы внимания и промежуточная размерность каждого эксперта разделены между восемью GPU. Каждый GPU вычисляет частичный результат, а коллективные операции объединяют части перед началом зависимой работы. Все бенчмарки в этой статье были запущены на MI355X (CDNA4, gfx950).

Оптимизация компонентов и производительность

Каждый компонент ниже начинается с описания его роли в модели, затем рассматривается оптимизация ядер, слияние и распределенная работа там, где это применимо. Цифры производительности указывают на базовую линию и оборудование.

1. KDA: Удержание рекуррентного состояния на чипе

Стандартное внимание сохраняет информацию для каждого предыдущего токена и считывает эту историю для каждого нового токена, поэтому объем работы растет вместе с длиной диалога. KDA вместо этого использует рекуррентное состояние фиксированного размера. Каждый токен обновляет состояние и считывает из него. Таким образом, рекуррентность KDA имеет постоянный объем работы на шаг декодирования по мере роста контекста; периодические слои MLA модели по-прежнему считывают всю историю (Рисунок 3).

Рисунок 3: Стандартное внимание считывает растущую историю; KDA обновляет рекуррентное состояние фиксированного размера.

1.1 Макет состояния и перемещение данных префилла

KDA выполняет минимум арифметических операций на байт, поэтому скорость зависит от перемещения данных. Рекуррентный процесс вычисляет одно выходное значение путем свертки строки состояния по каналам ключей. Поэтому мы храним состояние в формате V-major как [значение, ключ], делая каждую строку свертки непрерывной. Это дает преимущество как при сканировании префилла (prefill), так и при декодировании. Состояние декодирования использует порядок, в котором его считывает ядро, что сокращает время обновления состояния для одного пользователя с 10,5 до 5,2 мкс.

Для префилла промпт разбивается на фрагменты по 64 токена. Большая часть матричных вычислений внутри фрагмента выполняется параллельно, в то время как сканирование переносит рекуррентное состояние от одного фрагмента к другому. Наше сканирование Gluon сохраняет это состояние в регистрах и формирует выходные данные за один проход, вместо записи контрольной точки после каждого фрагмента.

Отдельно от компоновки состояния V-major мы реорганизовали временный буфер стробируемых ключей (gated-key buffer), используемый при префилле. Мы изменили его порядок с token-major на chunk-major, разместив размерность свертки сканирования непрерывно в памяти. Индекс фрагмента на стороне устройства устраняет необходимость вычисления адреса для каждого фрагмента, а конвейер запланированной загрузки/MFMA совмещает перемещение данных в памяти с вычислениями. Это сокращает время выполнения полного вызова префилла 8K с 659 до 595 мкс, что составляет снижение на 9,7%.

1.2 Объединенное декодирование и обновление спекулятивного состояния

Шаг декодирования изначально запускал отдельные ядра свертки, рекурсии и нормализации, каждое из которых записывало данные для чтения следующим ядром. Объединенное ядро Gluon сохраняет эти промежуточные данные в кэше чипа. Оно также напрямую потребляет упакованные и выровненные по токенам выходные данные проекции, избегая лишнего копирования QKV и гейтов вокруг ядра декодирования. Проекция затухания f_b также вычисляется внутри ядра рекурсии, что исключает еще один запуск, сохраняя при этом необходимое округление BF16 перед обработкой гейтов. Спекулятивная проверка позволяет не сохранять состояние для каждой отклоненной позиции; шаг повтора (replay) фиксирует принятый префикс.

1.3 Прирост производительности KDA при префилле, декодировании и проверке

• В 1,6 раза более быстрый фрагментированный префилл Gluon по сравнению с ядром KDA с открытым исходным кодом Triton (8K токенов, 1055 -> 659 мкс)
• В 1,11 раза более быстрое время вызова префилла KDA при 8K после дополнительной реорганизации буфера стробируемых ключей (8K токенов, 659 -> 595 мкс)
• В 1,63 раза более быстрый шаг декодирования после объединения четырех ядер в одно (16 пользователей; 25,4 → 15,6 мкс на слой)
• В 1,16–1,59 раза более быстрое обычное декодирование KDA, включая проекцию f_b при B1–64.
• В 1,66 раза меньше времени на проверку спекулятивных предположений по всем 69 слоям KDA (16 пользователей, окно 4 токена, 2,75 → 1,66 мс)

2. MLA: повторное использование сжатой истории

Оставшиеся 24 слоя внимания K3 используют MLA. Вместо хранения отдельных ключей и значений для 96 голов внимания, MLA хранит общее латентное представление шириной 512 и компонент ключа шириной 64 на токен, что значительно меньше данных, чем расширенный кэш для каждой головы. Декодирование может преобразовывать запрос в сжатое пространство и проецировать результат после этого. Префилл может расширять кэшированные латентные представления для выбранного ядра внимания. K3 использует NoPE (без ротационного позиционного эмбеддинга), а хранение кэша в формате FP8 сокращает объем байтов вдвое по сравнению с BF16 (Рисунок 4).

Рисунок 4: MLA хранит сжатое представление истории внимания. Столбцы приведены в масштабе. Меньший размер кэша означает, что в памяти GPU помещается больше пользователей и более длинные контексты, а также требуется считывать меньше данных за шаг.

2.1 Конвейерный префилл и совместное чтение KV

Для длинных промптов MLA — это большая вычислительная задача. Наше ядро Gluon в формате FP8 запускает две группы потоков на каждом вычислительном блоке со смещением в один шаг: пока одна группа выполняет умножение, другая параллельно загружает следующий тайл KV.

2.2 Объединенная запись в кэш и обработка вывода

Перед запуском ядра MLA среда выполнения подготавливает строки кэша и запрос. В пути префилла одно ядро упаковывает и преобразует операнды K/V и выполняет подготовку кэша, заменяя восемь небольших запусков одним. Тот же объединенный путь подготовки запроса используется на всех этапах MLA для нормализации и проекции. При чистом декодировании он также подготавливает поглощенный запрос. Эпилог декодирования затем объединяет свертку результата внимания, проекцию из латентного представления в значение, выходное стробирование и финальную запись, избегая записи и повторного чтения промежуточного латентного представления.

2.3 Прирост производительности MLA при префилле и проверке

Рисунок 5: Обработка промпта MLA (FP8, 12 голов на GPU). Время на вызов в мкс, чем меньше, тем лучше · измерено на MI355X

В 3,8–5,1 раза быстрее, чем переносимое ядро MLA Triton на полных фрагментах 4K и 8K, и в 3,2 раза на запросе 8K с кэшированным префиксом 32K. 8-волновое исполнение Gluon в 1,13–1,18 раза быстрее предыдущего ядра Gluon в этих исторических измерениях (Рисунок 5).

Рисунок 6: Проверка спекулятивных предположений: чтение истории один раз против одного раза на позицию. Ускорение, чем выше, тем лучше, по количеству пользователей (C) · измерено на MI355X

В 1,86 раза быстрее при контексте 50K и в 2,14 раза при контексте 100K для 16 пользователей, при соблюдении численного сравнения, используемого в тестах ядер. Очень короткие истории в этом эксперименте работают медленнее (0,99–1,20x при 1K/8K для 1–16 пользователей). Это исторические показатели времени работы ядер, а не новый результат «от конца до конца» на момент отсечки (Рисунок 6).

3. AttnRes: объединение смешивания представлений

В большинстве моделей выход каждого слоя добавляется к текущему потоку остаточных связей (residual stream). Kimi K3 группирует слои в блоки по 12 и сохраняет представления завершенных блоков. Перед каждым шагом внимания и прямой связи (feed-forward) он оценивает эти представления вместе с текущим накопленным потоком и подает их взвешенную смесь на следующую операцию. Глубокие слои могут выбирать информацию из более ранних блоков (Рисунок 7).

Рисунок 7: Остаточные связи внимания (Attention Residuals) смешивают сохраненные представления блоков с текущим потоком остаточных связей.

Подвох: при наивной реализации это означает повторное чтение до девяти больших тензоров дважды за слой. И это происходит сразу после тензорно-параллельной операции all-reduce, поэтому каждая микросекунда задержки напрямую добавляется к времени шага.

3.1 Однопроходное смешивание с ограниченным использованием регистров

Наше ядро AttnRes оценивает, смешивает и нормализует за один проход по снимкам, накапливая результат в 32-битных числах с плавающей запятой. Мы также исправили тонкую проблему. Раньше ядро рассматривало один размер входных данных как константу, «зашитую» в скомпилированную программу, поэтому каждая новая длина промпта вызывала повторную компиляцию прямо во время обслуживания, по 100+ мс на каждый случай. Превращение этого параметра в значение времени выполнения сократило время одного теста с длинным промптом с 46 с до 16 с.

Отсечка также меняет цикл снимков. Чтение по одному снимку за раз не дает компилятору удерживать все кандидаты активными одновременно, что снижает нагрузку на регистры. Это делает смешивание 8 снимков для 8K токенов в 1,10 раза быстрее (283 → 257 мкс). Количество токенов и зависящие от батча шаги остаются входными данными времени выполнения, поэтому новые формы обслуживания не требуют перекомпиляции.

3.2 Совместное использование работы AttnRes между GPU

Для подходящих батчей TP8 мы распределяем смешивание после внимания по строкам токенов. Операция Pull reduce-scatter дает каждому GPU одну восьмую часть уменьшенного выхода внимания. Он обновляет остаточную связь, смешивает все сохраненные снимки для назначенных ему токенов и применяет выходную нормализацию RMSNorm. Затем операция Push all-gather собирает нормализованный вход MoE на каждом GPU, в то время как шардированная обновленная остаточная связь остается локальной до тех пор, пока ее не потребит хвост MoE (см. Раздел 6).

Это устраняет семь восьмых дублирующейся работы по смешиванию и нормализации post-attention. Те же проверки количества строк и совместимости применяются к prefill, decode и смешанным пакетам. Например, совместимый пакет проверки EAGLE3 из 64 строк дает каждому GPU по восемь строк для смешивания.

3.3 Подготовка оценок и объединение финального микса

Это самый наглядный пример того, почему важно владеть каждым ядром. Снимки состояния меняются только каждые 12 слоев, поэтому большую часть смешивания можно вычислить заранее, как побочный продукт более раннего матричного умножения. Только последний шаг, включение новой суммы, должен ждать. Затем мы включаем этот последний шаг в саму операцию all-reduce (Рисунок 8):

Рисунок 8: Объединение all-reduce, обновления остатков (residual update), смешивания AttnRes и нормализации устраняет промежуточный трафик памяти.

3.4 Ядро AttnRes и прирост производительности обслуживания

Рисунок 9: Ядро AttnRes против стандартной версии PyTorch. Время на вызов в мкс, логарифмическая шкала, чем меньше, тем лучше · измерено на MI355X

В 4,1–5,9 раза быстрее при размерах decode и в 23 раза быстрее на фрагменте промпта в 8 тыс. токенов (Рисунок 9).

• −27% времени для объединенного шага all-reduce + AttnRes при 16 токенах (21,7 → 15,8 мкс; 8 GPU)
• +7–14% токенов/с в сквозном режиме для 1–4 пользователей благодаря этому объединению и связанным изменениям

4. Latent MoE: снижение накладных расходов на экспертов и проекции

Сеть feed-forward использует 896 маршрутизируемых экспертов в 92 из 93 слоев декодера K3; первый слой является плотным. Маршрутизатор выбирает 16 экспертов на токен, и два общих эксперта работают на каждом токене. Маршрутизируемые эксперты работают в латентном пространстве шириной 3584, а не в скрытом пространстве шириной 7168, а затем проецируют свой объединенный результат обратно на полную ширину. Веса экспертов MXFP4 сокращают хранение и трафик весов; маршрутизируемые и общие ветви сохраняют свои отдельные проекции (Рисунок 10).

Рисунок 10: Latent MoE направляет каждый токен выбранным экспертам. Сетка — это эскиз (224 из 896 экспертов); 16 выделенных ячеек — это эксперты, выбранные для одного токена.

4.1 Повторное использование весов экспертов и меньшие вычислительные блоки

При 896 экспертах большинство экспертов получают лишь несколько токенов на любом шаге. Блоки матричного умножения GPU любят аккуратные блоки по 64 или 128 строк, поэтому старые ядра дополняли несколько токенов каждого эксперта до полного блока, тратя большую часть работы впустую. Мы атаковали эту проблему с обеих сторон:

• Decode: для небольших пакетов проходите по выбранным маршрутам каждого токена напрямую, вместо того чтобы создавать отсортированные группы. Активации используют FP8 с весами MXFP4. При 32–64 строках в поддерживаемой конфигурации K3 (attention TP8; MoE TP8/EP1) мы вводим moe_sorting для warp-decode MoE, который сортирует маршруты по экспертам, чтобы несколько маршрутов повторно использовали каждую загруженную панель весов.
• Prefill: используйте экспертные блоки по 16 и 32 строки там, где разреженная маршрутизация привела бы к потере большего блока, и объединяйте комбинирование вывода экспертов, когда формат активации и размер пакета делают это выгодным.

4.2 Объединение редукций маршрутизируемых и общих экспертов

Каждый GPU завершает шаг вычисления экспертов MoE частичными результатами маршрутизируемых и общих экспертов, которые должны быть просуммированы по всем восьми GPU. Наши ядра записывают частичные результаты BF16 маршрутизируемых и общих экспертов непосредственно в один повторно используемый буфер Iris. Одна коллективная операция редуцирует обе ветви по восьми GPU, сохраняя их результаты раздельно и избегая копирования при конкатенации. Крошечные пакеты используют протокол push, чьи поступающие данные также сигнализируют о готовности (Раздел 6).

Для подходящих более крупных пакетов reduce-scatter дает каждому GPU оба редуцированных результата для одной восьмой строк токенов. Он применяет маршрутизируемую RMSNorm и восходящую проекцию локально, затем добавляет общий результат и остаток во время push all-gather. Каждый GPU выполняет только одну восьмую часть работы по нормализации и проекции, а сбор только объединенного вывода сокращает трафик all-gather на треть.

4.3 Объединение входных проекций и выходных дополнений

Каждый слой MoE применяет три проекции к одному и тому же входу: оценки маршрутизатора, маршрутизируемые латентные входы и значения gate/up общих экспертов. Мы храним их матрицы весов последовательно и вычисляем их с помощью одного матричного умножения. Там, где это выгодно, его эпилог также применяет активацию общих экспертов, избегая промежуточной записи и чтения.

Объединение вывода следует за вычислением выбранных экспертов. Некоторые ядра для малых пакетов объединяют вклады выбранных экспертов напрямую; expert-sorted 32–64-row path записывает частичные результаты маршрутов и суммирует их в FP32. В конце MoE поддерживаемые ядра восходящей проекции на 2–32 строки также добавляют вывод общих экспертов и текущий остаток, экономя отдельное ядро сложения.

4.4 Прирост производительности экспертов, проекций и коммуникации

Рисунок 11: Вычисление экспертов до и после оптимизации, по количеству токенов. Время на вызов в мкс, чем меньше, тем лучше · измерено на MI355X при 7e2f1deb

В совокупности эти изменения делают вычисление экспертов на одном GPU в 2,2–2,3 раза быстрее при 32–64 токенах и в 1,5–2,0 раза быстрее на фрагментах промпта в 128–2048 токенов. При 4096–8192 токенах прирост составляет 3–4%. Сравнение отключает оптимизации вычисления экспертов главы на той же ревизии кода, изолируя их совокупный эффект (Рисунок 11).

Рисунок 12: Шаг ввода MoE: одно упакованное ядро против трех GEMM плюс SiTU. Время на вызов в мкс, чем меньше, тем лучше · измерено на MI355X при 7e2f1deb

В 2,0–3,0 раза быстрее от 1 до 8192 токенов, включая 2,26 раза при одном токене и 2,95 раза при 8192 токенах. Ранее сообщенный сквозной результат увеличил генерацию для одного пользователя с 64,1 до 77,6 токенов/с (+21%, сообщается) (Рисунок 12).

Работа по проекции также выигрывает от дополнений Gluon GEMM. Работа Bucketed Gluon decode GEMM сообщает о до 3,02 раза более быстрой проекции MLA kv_b decode, чем torch.mm при измеренной форме TP8. Реализация Gluon GEMM с большим M расширяет GEMM с большим M до неровных ширин prefill и маршрутизирует измеренный набор проекций K3. По 160 отложенным случаям проекции диспетчеризация сокращает общее время GEMM примерно на 4%. Объединение латентной восходящей проекции сокращает время латентной восходящей проекции плюс сложение остатков на 11–16% при 2–32 строках.

5. EAGLE3: проверка нескольких токенов за проход

Спекулятивное декодирование использует небольшую быструю черновую модель, чтобы предложить несколько токенов, а затем просит целевую модель проверить их за один проход. В трехшаговой конфигурации EAGLE3, показанной на Рисунке 13, Kimi K3 проверяет четыре позиции, принимает совпадающий префикс и предоставляет целевой токен в точке отклонения или после того, как все предложения приняты. При жадной выборке принятие проверяет жадный выбор цели. Точное равенство токенов при различных конфигурациях пакетирования или численных конфигурациях требует отдельной сквозной проверки.

Рисунок 13: EAGLE3 проверяет несколько черновых токенов за один проход целевой модели.

Спекуляция меняет форму работы каждого компонента. Например, при 16 одновременных запросах и четырех позициях проверки цель получает 64 строки. Соответствующие оптимизации описаны там, где они происходят: повторное использование KV по позициям запроса (Раздел 2.1), повторное использование весов экспертов и редукция top-k (Раздел 4.1– Раздел 4.2) и шардирование строк внимания (Раздел 6.2). Описанная здесь трехшаговая конфигурация использует цепочку (topk=1); поддержка дерева, объединенная отсечением, не меняет эту конфигурацию.

5.1 Производительность

Мы оцениваем производительность EAGLE3 на синтетическом случайном наборе данных с 50 000 входных токенов и 500 выходными токенами на запрос, сравнивая выполнение с EAGLE3 и без него (Рисунок 14). Это отдельная оценка, отличная от многоходовой рабочей нагрузки SWE-smith, используемой при сквозном сравнении обслуживания. EAGLE3 сокращает время на один выходной токен при любой протестированной параллельности: с 12,9 до 4,2 мс при одном одновременном пользователе и с 28,8 до 11,4 мс при 16 одновременных пользователях.

Рисунок 14: Время на один выходной токен (TPOT) с EAGLE3 и без него на восьми графических процессорах MI355X с использованием синтетического случайного набора данных с 50 тыс. входных токенов и 500 выходными токенами на запрос. Чем ниже значение, тем лучше; C1–C16 обозначают количество одновременных пользователей.

6. Iris: объединение коммуникации и вычислений

Тензорный параллелизм распределяет каждый слой между восемью графическими процессорами MI355X, поэтому их частичные выходные данные должны быть суммированы до начала зависимых операций. Наша реализация Kimi K3 использует Iris для создания протоколов связи в нативном Gluon и объединения вычислений там, где это уместно. Симметричная куча предоставляет буферы, которые ядра могут напрямую считывать или записывать на одноранговых графических процессорах. Для декодирования с низкой задержкой и предварительного заполнения с высокой пропускной способностью мы реализовали оптимизированные ядра, основанные на размере сообщения и специфических для K3 вычислениях, окружающих каждую операцию редукции.

Каждый слой имеет две операции all-reduce: одну после проекции выходных данных внимания и одну после вычислений экспертов. Со стороны внимания остатки внимания смешивают текущий остаток с сохраненными состояниями из предыдущих блоков слоев, используя обученные веса, зависящие от токенов, с последующей RMSNorm. Операция all-reduce для MoE объединяет две частичные суммы. Маршрутизируемые эксперты создают латентный вектор шириной 3584, который мы нормализуем и проецируем в скрытую размерность модели шириной 7168. Общий эксперт уже выдает выходные данные полной ширины. Мы размещаем оба частичных результата последовательно в одном входном буфере, получая общую рабочую нагрузку BF16 в 21 КиБ на строку.

6.1 Малые/средние сообщения

Со стороны внимания подходящие слои объединяют одноэтапную операцию push all-reduce с обновлением остатков, объединением AttnRes и финальной RMSNorm. Объединение этих шагов устраняет отдельные запуски эпилогов и промежуточный трафик памяти, уменьшая фиксированные накладные расходы, что наиболее важно для малых сообщений. Со стороны MoE мы реализовали all-reduce без барьеров, который в 1,2–3 раза быстрее, чем RCCL для размеров сообщений, для которых он используется. Поступающие полезные нагрузки сигнализируют о готовности, позволяя каждому тайлу выполнить редукцию, как только поступят его входные данные, без отдельного барьера. Сообщения среднего размера используют оптимизированную двухэтапную операцию pull all-reduce, которая также обрабатывает более крупные сообщения, не подходящие для построчно-сегментированных вычислений, обсуждаемых далее (Рисунок 15).

6.2 Большие сообщения

Рисунок 15: Реплицированные и построчно-сегментированные коммуникационные «хвосты»: где каждый графический процессор вычисляет и обменивается результатами.

Обычно каждый ранг TP повторяет вычисления после all-reduce по всем M. Для более крупных сообщений (M >= 40 для MoE, M >= 56 для внимания) мы используем pull reduce-scatter, назначаем M/8 полных строк каждому из восьми графических процессоров и откладываем push all-gather до завершения «хвоста» внимания или MoE. Эти «хвосты» передают активации BF16; редукции и нормализация используют арифметику FP32 внутри.

6.3 Задержка коллективных операций и прирост производительности обслуживания

На рисунках 16 и 17 сравнивается Iris с базовыми показателями RCCL для коммуникационных «хвостов» MoE и внимания на восьми графических процессорах MI355X.

Рисунок 16: Задержка коммуникационного «хвоста» MoE: Iris и базовые показатели RCCL на восьми графических процессорах MI355X.

Заключение

Благодарности

Сноски

Конфигурация системы