Kimi K3.1 Pro на графических процессорах AMD Instinct

Кратко (TL;DR)

• Kimi K3.1 Pro — агентная модель с открытыми весами и нативной мультимодальностью от Moonshot AI (2,8 трлн общих и 104 млрд активных параметров, контекст 1 млн токенов) — обслуживается на 8 графических процессорах AMD Instinct™ MI355X с использованием как ATOM, так и vLLM с первого дня, непосредственно из выпущенной контрольной точки MXFP4 и с использованием черновика спекулятивного декодирования DSpark.
• K3.1 Pro сохраняет гибридную архитектуру Kimi K3: 69 слоев Kimi Delta Attention (KDA) и 24 слоя латентного внимания поверх 896-экспертной сети MXFP4 MoE. В каждый из 24 слоев внимания добавлен разреженный индексатор (DSA), который выбирает 2048 токенов истории для каждого запроса. ATOM обслуживает проиндексированные слои по общему пути DSA с кешем индексов FP8, расположенным рядом с кешем латентных KV в формате FP8.
• ATOM: нативные ядра KDA и разреженного MLA, маршрутизируемые эксперты на базе сгруппированной архитектуры MoE MXFP4 от FlyDSL, поканальное внимание FP8, применяемое при запуске, и черновик DSpark, который разделяет кеш KV целевой модели и выполняется в виде захваченного графа CUDA.
• vLLM: эффективное по памяти разреженное внимание с параллелизмом контекста декодирования (DCP) для обслуживания длинного контекста, наряду со специализированным путем декодирования для интерактивных рабочих нагрузок с низкой задержкой
• Точность по оценке верификатора Kimi Vendor Verifier от Moonshot AI при максимальном уровне рассуждений: OCRBench 0,871 (эталон Moonshot — 0,872); схема JSON для вызова инструментов — 0,991.

Введение

Kimi K3.1 Pro работает на графических процессорах AMD Instinct серии MI350 с первого дня с помощью ATOM, механизма инференса AMD, на программном обеспечении ROCm™ с ядрами AITER и FlyDSL. Модель дообучена на базе Kimi K3 с использованием крупномасштабного совместного обучения с подкреплением. Moonshot AI описывает ее как свою самую мощную модель класса 3T с окном контекста в 1 млн токенов, настраиваемым уровнем рассуждений (низкий, высокий, максимальный), нативным вызовом инструментов и вводом изображений.

В этой статье рассматриваются изменения в K3.1 Pro по сравнению с Kimi K3, особенности работы ATOM, схема обслуживания и точность в наборе тестов верификации поставщика Moonshot AI.

Обзор Kimi K3.1 Pro

Для инференса важны пять свойств:

• Гибридное внимание, теперь разреженное: из 93 слоев декодера 69 используют Kimi Delta Attention — управляемое линейное внимание с рекуррентным состоянием фиксированного размера для каждого запроса. Остальные 24 слоя представляют собой слои управляемого латентного внимания (MLA): каждый четвертый слой плюс последний. В K3.1 Pro каждый из них также имеет индексатор (32 головы размерностью 128), который оценивает историю запроса и выбирает верхние 2048 токенов для внимания. Затем внимание считывает не более 2048 токенов на запрос, независимо от роста контекста; всю историю оценивает только легковесный индексатор.
• Два небольших кеша на слой внимания: каждый слой MLA кеширует один латентный вектор шириной 576 на токен (512 сжатых KV плюс 64 размерности RoPE) и один ключ индекса на токен.
• Широкая латентная сеть MoE: имеется 896 маршрутизируемых экспертов (16 активных на токен) плюс два общих эксперта. Маршрутизируемые эксперты работают в латентном пространстве шириной 3584 (что составляет половину скрытого размера 7168) и используют активацию SiTU-GLU. Только первый слой является плотным.
• MXFP4 с учетом квантования: маршрутизируемые эксперты хранятся в формате MXFP4 (одна шкала E8M0 на 32 значения). Внимание, индексатор, общие эксперты, плотный MLP, проекции латентного пространства вверх/вниз, маршрутизатор, языковая голова (LM head) и визуальный энкодер имеют формат BF16.
• Спекулятивное декодирование DSpark: отдельный пятислойный черновик MLA считывает скрытые состояния из пяти последних слоев целевой модели и предлагает блок токенов черновика на каждом шаге.

Элемент

Конфигурация

Параметры

2,8 трлн всего, 104 млрд активных

Слои декодера

93: 69 KDA, 24 управляемых DSA (MLA + индексатор); слой 1 плотный, остальные MoE

KDA

96 голов, размерность головы 128, короткая свертка шириной 4

MLA

96 голов; ранг q LoRA 1536; ранг KV LoRA 512 + 64 размерности RoPE; выходной гейт

Индексатор DSA

32 головы × 128 размерностей; топ-2048 токенов; каждый слой MLA

Эксперты

896 маршрутизируемых (top-16, сигмоидальная маршрутизация) + 2 общих; промежуточный размер эксперта 3072 в латентном пространстве шириной 3584

Словарь / контекст

163 840 / 1 048 576 токенов

Зрение

Энкодер MoonViT-V2 (401 млн параметров)

Черновик DSpark

5 слоев MLA (96 голов), чересстрочный YaRN RoPE, низкоранговая голова Маркова (ранг 512)

Обслуживание с помощью ATOM

ATOM обслуживает Kimi K3.1 Pro с использованием тензорного параллелизма на восьми графических процессорах одного узла MI355X для ввода текста и изображений, используя выпущенную контрольную точку MXFP4 без конвертации. При загруженном черновике DSpark каждый графический процессор достигает пикового значения около 199 ГБ из своих 288 ГБ во время запуска, а один запрос на 1 млн токенов помещается на одном узле.

Что мы оптимизировали

Kimi Delta Attention: префилл запускает фрагментированное ядро Kimi Delta Attention из AITER. Декодирование продвигает рекуррентное состояние каждого запроса с помощью слитного обновления по дельта-правилу с гейтингом, а состояния хранятся в пуле для каждого запроса.

Разреженный MLA на общем пути DSA: реализация DSA в ATOM, общая с моделями классов DeepSeek-V3.2 и GLM-5, обслуживает индексатор K3.1. Латентный KV и ключи индекса находятся в постраничных кешах FP8 с блоками по 128 токенов. Проекции ключей индекса и весов голов упакованы в слитную нисходящую проекцию q/kv слоя, поэтому один вызов GEMM создает запросы, латентные векторы KV и ключи индекса. Выборка топ-2048 затем управляет разреженным вниманием MLA.

Латентный MoE MXFP4 на FlyDSL: маршрутизируемые эксперты выполняются на нативном сгруппированном ядре MoE SiTUv2 от FlyDSL с весами MXFP4 и вторым этапом FP8. Два общих эксперта выполняются на втором потоке одновременно с маршрутизируемыми экспертами.

Внимание PTPC FP8: при запуске ATOM выполняет реквантование внимания BF16 и плотных линейных слоев в формат FP8 с одной шкалой на выходной канал, причем активации масштабируются на токен во время выполнения (--online_quant_config с ptpc_fp8): 559 линейных слоев на ранг, выполняется примерно за 2 секунды. Маршрутизируемые эксперты, свертки KDA, визуальный энкодер и упакованная нисходящая проекция индексатора сохраняют свои сохраненные форматы.

Спекулятивное декодирование DSpark: латентное пространство черновика шириной 576 имеет ту же форму, что и слои MLA целевой модели, поэтому ATOM связывает пять слоев черновика в пул KV целевой модели в качестве дополнительных строк, и они наследуют тип данных кеша FP8. Блок черновика, который предлагает все токены черновика за один параллельный проход, захватывается в граф CUDA.

Параллелизм контекста декодирования: для множества одновременных запросов с длинным контекстом ATOM разделяет кеш латентного KV на сегменты вдоль последовательности по всем восьми графическим процессорам (--decode-context-parallel-size 8). Каждый графический процессор производит вычисление внимания по своему сегменту, а частичные результаты объединяются с помощью операции all-to-all. Проекция запроса MLA реплицируется по всей группе, поэтому пошаговая операция all-gather для запросов не требуется. Метаданные разреженного внимания перестраиваются после каждого слоя индексатора.

Монодекодирование: для обслуживания, ограниченного задержкой, ATOM заменяет каждый слой KDA на шаге декодирования или верификации двумя слитными ядрами: K1 выполняет путь внимания через рекуррентность KDA, а K2 выполняет выходную проекцию плюс латентный MoE. Слои MLA сохраняют постраничное разреженное внимание, а также выполняют K2. Монодекодирование обслуживает шаги размером до 64 строк токенов, например, один запрос, верифицирующий семь токенов черновика.

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

Задержка одиночного запроса при монодекодировании: монодекодирование снижает среднее время генерации одного токена (TPOT) на 33,6% по сравнению с тем же сервером на стандартном пути декодирования.

Метрика

Стандартное декодирование

Монодекодирование

Улучшение

Среднее TPOT

5,78 мс

3,84 мс

33,6%

Пропускная способность с параллелизмом контекста декодирования.

Входные / выходные токены

Общая пропускная способность (токенов/с)

Пропускная способность вывода (токенов/с)

Медианное TTFT (с)

Медианное TPOT (мс)

8K / 1K

7,484

831.6

0.821

16.36

128K / 100

14,005

10.68

27.9

1,146

Как запустить

Запустите контейнер:

Запустите сервер (TP8 с DSpark):

Примечания:

• Требуется параметр --block-size 128: это макет MLA, а строки черновика находятся на страницах цели.
• Без параметра --max-model-len сервер использует модель с 1 048 576 токенами.
• Без спекулятивного декодирования удалите последнюю строку.
• Уровень рассуждений задается для каждого запроса с помощью поля верхнего уровня reasoning_effort (low, high, max). Диалоги с несколькими ходами и вызовами инструментов должны возвращать полное сообщение ассистента, включая reasoning_content и tool_calls.

Быстрая проверка:

Рассуждения возвращаются в reasoning_content, а вызов get_weather({"city": "Paris"}) — в tool_calls.

Точность

Оценено по ATOM на 8 графических процессорах MI355X в режиме размышлений с максимальным уровнем усилий рассуждений, если не указано иное. MMMU-Pro, OCRBench, BEAM и проверка схемы вызовов инструментов взяты из верификатора поставщика Kimi от Moonshot AI (KVV).

Бенчмарк

Параметры

ATOM

MMMU-Pro (стандартный, 10 вариантов), 1730 вопросов

максимальный уровень усилий, температура 1.0, top_p 0.95, макс. токенов 98 304

0.842 (±0.009)

OCRBench, 1000 вопросов

максимальный уровень усилий, температура 1.0, top_p 0.95, макс. токенов 16 384

0.871 (±0.011)

BEAM (1M), 700 вопросов

максимальный уровень усилий, температура 1.0, top_p 0.95; оценка выполнена с помощью GPT-4.1

0.416

GSM8K 5-shot, 1319 вопросов

жадный поиск

0.78 гибкий / 0.977 строгий (±0.004)

Схема JSON для вызова инструментов, 428 проверок

максимальный уровень усилий, макс. токенов 2048, одна попытка на проверку

0.997 (427 / 428)

• Тесты MMMU-Pro и проверка схемы вызовов инструментов выполнялись на сервере в разделе «Как запустить».
• Тест OCRBench выполнялся на TP8/DCP8 с кэшем KV в формате BF16, индексным кэшем в формате FP8 и семью токенами черновика.
• Тесты GSM8K и BEAM выполнялись на TP8 без спекулятивного декодирования; в BEAM использовался контекст на 1 млн токенов с кэшированием префиксов.
• Оценка BEAM производилась с помощью GPT-4.1, а не модели-судьи, используемой в качестве эталона Moonshot AI, поэтому эти два балла несопоставимы по шкале.

Обслуживание с помощью vLLM

vLLM поддерживает Kimi K3.1 Pro на графических процессорах AMD Instinct MI355X с контрольной точкой MXFP4 и спекулятивным декодированием DSpark. Реализация ROCm, созданная на базе ядер AITER и FlyDSL, объединяет разреженное внимание с оптимизацией памяти и параллелизм контекста декодирования (DCP) для обслуживания длинного контекста, а также специализированный путь декодирования для интерактивных рабочих нагрузок с низкой задержкой.

Что мы оптимизировали

Индексация Native MXFP4. Индексатор сканирует контекст для выбора до 2048 токенов для разреженного внимания. MXFP4 уменьшает размер каждого ключа индекса со 132 байт в формате FP8 до 68 байт, включая шкалы, сокращая объем памяти индексного кэша на 48,5%. Для подсчета очков используются собственные матричные инструкции FP4 с накоплением и редукцией с весовыми коэффициентами голов в FP32. Во время спекулятивной верификации несколько строк запроса используют одни и те же ключи для сокращения числа повторных чтений.

Меньший объем индексного кэша оставляет больше памяти GPU доступной для кэшированного контекста, что повышает эффективность обслуживания длинного контекста при высокой параллельности. В тестах AgentX конфигурация с индексатором MXFP4 обеспечила на 5,3% большую общую пропускную способность рабочей нагрузки при параллельности 48.

Предварительное заполнение (Prefill): разреженное внимание, параллельное по контексту. Во время предварительного заполнения GPU собирают ключи индексов и распределяют строки запросов между собой. Каждый GPU оценивает назначенные ему строки по всему контексту и напрямую выбирает глобальный топ-k, избегая объединения кандидатов для каждого ранга. Основное внимание может отдельно собирать KV, когда это снижает нагрузку на коммуникацию и укладывается в бюджет рабочей памяти, а затем обрабатывать все строки запроса для локальных голов.

Оптимизации предварительного заполнения DCP позволили улучшить пропускную способность и задержку в бенчмарках разработки, причем наибольший прирост наблюдался на длинных входных данных.

Входные / выходные токены

Общая пропускная способность (токенов/с)

Пропускная способность вывода (токенов/с)

Медианное TTFT (с)

Медианное TPOT (мс)

8K / 1K

5 357,8 → 6 015,4 (+12%)

595,3 → 668,4

2,462 → 1,875

24,571 → 22,063

128K / 100

6 950 → 17 678 (+154%)

5,30 → 13,48

268,7 → 105,5

331,5 → 129,6

Разреженный MLA FlyDSL считывает выбранные токены напрямую из кэша KV FP8. Каждая рабочая группа обрабатывает одну строку запроса по своим локальным головам внимания, обеспечивая параллелизм для множества строк запросов на шаге предварительного заполнения. Меньшие батчи сохраняют путь Triton, который обеспечивает дополнительный параллелизм внутри каждой строки. В отдельном контролируемом сравнении на восьми графических процессорах MI355X включение разреженного MLA FlyDSL увеличило пропускную способность вывода на 7,0% для рабочей нагрузки со входом 128K.

Реализация

Длительность бенчмарка (с)

Медианное TTFT (с)

Пропускная способность вывода (токенов/с)

Triton

432,8

96,2

14,79

FlyDSL

404,4

89,8

15,83

Улучшение

на 6,6% ниже

на 6,6% ниже

на 7,0% выше

Декодирование: слитное декодирование с низкой задержкой. Индексация DSA и разреженное внимание выполняются отдельно, после чего запускается персистентное моноядро, которое объединяет проекцию вывода внимания, AttnRes и полный скрытый MoE, включая все операции all-reduce с тензорным параллелизмом, в один запуск. Слои KDA также объединяют входные проекции, короткую свертку, обновления рекуррентного состояния и нормирование с управляемым доступом. Специализированный путь поддерживает батчи верификации размером до восьми строк токенов, включая один запрос с семью токенами черновика.

Метрики

Обычный путь

Специализированное декодирование

Улучшение

среднее TPOT

6,02 мс

4,40 мс

27,0%

Связь: эффективный межсетевой обмен данными между GPU. Во время декодирования каждый GPU обменивается оценками кандидатов, сохраняя идентификаторы токенов локальными. Общий порог и согласованная обработка совпадений определяют, какие локальные кандидаты выбираются. Для 2048 кандидатов простой обмен оценками сокращает объем полезной нагрузки кандидатов с 16 до 8 КиБ на запрос для каждого GPU.

AITER ускоряет операцию all-gather для DCP с помощью выделенного ядра GPU, которое считывает буферы узлов через HIP IPC и записывает данные напрямую в требуемый макет вывода. Это сочетает сбор данных и преобразование макета, избегая отдельного копирования для изменения порядка. В автономных микробенчмарках на восьми графических процессорах MI355X путь AITER сократил задержку all-gather на 53–61% по сравнению с RCCL для протестированных форм тензоров запросов и кандидатов.

Как запустить

Следующая конфигурация использует восемь графических процессоров AMD Instinct MI355X с TP8/DCP8, кэшем KV MLA FP8, ключами индекса MXFP4 и DSpark с пятью токенами черновика. Установите пути к контрольным точкам локальных копий Kimi K3.1 Pro и соответствующей ей модели черновика DSpark.

Этот рецепт использует стандартное отклонение (rejection), соответствующее оценкам точности ниже. Выгрузка KV на CPU резервирует примерно 1,64 ТиБ памяти хоста на восемь рангов в дополнение к другим требованиям среды выполнения. Ограничение контекста в 1 млн токенов и ограничение на 48 последовательностей не гарантируют емкость для 48 одновременных запросов по 1 млн токенов.

Запустите контейнер:

Запустите сервер (DCP8 с DSpark):

Точность vLLM

В приведенных ниже оценках сравниваются индексаторы FP8 и MXFP4 с DSpark при использовании стандартного отклонения.

Бенчмарк

Индексатор vLLM FP8

Индексатор vLLM MXFP4

GSM8K 5-shot, T = 0.7

1295/1319 (98.18%)

1291/1319 (97.88%)

Passkey, контекст 8K–500K, 25 случаев

25/25

25/25

LongBench-v2 short, жадный поиск

125/180 (69.44%)

124/180 (68.89%)

OCRBench, макс. усилия, температура 1.0, top_p 0.95, макс. токенов 16 384

0.871 ± 0.0106

0.870 ± 0.0106

Схема JSON для вызова инструментов, макс. усилия, макс. токенов 2048

407/408 (99.75%)

407/408 (99.75%)

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

Поддержка Kimi K3.1 Pro в день выпуска (day-0) стала результатом совместных усилий разработчиков моделей, движков выводов и аппаратных сообществ.

Мы благодарим команду Moonshot AI за создание Kimi K3.1 Pro, обмен деталями архитектуры до выпуска весов, участие в первоначальной интеграции модели и тесное сотрудничество в вопросах проверки корректности и готовности к продакшену.

Мы благодарим команду Inferact за интеграцию модели в vLLM, реализацию семантики обслуживания и мультимодальной поддержки, создание инструкций по развертыванию и повышение комплексной производительности.

Мы благодарим Цзинь Линпэна, Ван Итин, Чжу Юхуа, Янь Чживэя, Чжан Лижун и У Хэтти за поддержку AMD ATOM в первый день выпуска (day-0), Билла Хэ, Ичао Чжу, Сынрока Чжуна, Хёкджуна Ли, Хунся Яна, Ричарда Ли за поддержку AMD vLLM в первый день выпуска (day-0), а также Карлуса Хуана, Гилберта Лея, Пэн Суна и Энди Луо за межкомандное сотрудничество.

Самое главное, мы благодарим широкие сообщества разработчиков с открытым исходным кодом и AMD ROCm за ожидание, тестирование и отзывы, которые уже сопровождают Kimi K3.1 Pro на графических процессорах AMD Instinct. Мы с нетерпением ждем возможности предоставить вам веса, а также поддержку механизмов выводов ATOM и vLLM для графических процессоров AMD Instinct MI355X.

Конфигурация тестирования и отказ от ответственности

Результаты ATOM были измерены на 8 графических процессорах AMD Instinct MI355X в среде разработки ROCm ATOM с изменениями Kimi K3.1 в ноябре 2026 года. Контрольная точка: выпуск Kimi K3.1 Pro (маршрутизируемые эксперты MXFP4) с черновиком Kimi K3.1 DSpark. Точность оценена с помощью инструмента Kimi Vendor Verifier. Результаты на другом аппаратном обеспечении, версиях программного обеспечения, настройках или рабочих нагрузках будут отличаться. Эти показатели приведены для наглядности и не являются гарантией производительности.

Сноски

Приложение: переменные среды

Переменная

Установлено в

Описание

AITER_SITUV2_A4W4

Нативный SiTUv2 MXFP4 с группировкой MoE от FlyDSL для маршрутизируемых экспертов

AITER_FLYDSL_STAGE2_FP8

Второй этап FP8 в FlyDSL MoE

ATOM_USE_FLYDSL_FP8_PREFILL_ATTN

Предварительное заполнение внимания FlyDSL FP8 на слотах MLA

ATOM_GDN_SSM_DTYPE

fp16

dtype пула состояний KDA для каждого запроса

ATOM_STATE_CHECKPOINT_DEMAND

только контрольные точки состояния на конце промпта

AITER_LOG_LEVEL

WARNING

более тихий журнал запуска AITER