Когда ты листаешь ленту Instagram или ленту X, за каждым постом стоит сложная математика: модель оценивает сотни или тысячи кандидатов и выбирает, что тебе показать. При этом профиль пользователя — его история просмотров, предпочтения, контекст — одинаков для всех кандидатов в рамках одного запроса. Казалось бы, очевидная оптимизация: не копировать одни и те же данные 70 раз. Но на практике именно это происходило в промышленных рекомендательных системах до недавнего времени.

Meta вместе с командой PyTorch опубликовала подробный разбор подхода под названием IKBO — In-Kernel Broadcast Optimization. Это не просто ещё одна оптимизация CUDA-ядра. Это пример того, как ко-дизайн на трёх уровнях — ядро, модель, инфраструктура — позволяет устранить фундаментальное узкое место, которое раньше считалось неизбежным. Результат: до 4× ускорения на H100 SXM5 и сокращение вычислительной латентности на две трети в production-моделях Meta.

Что такое broadcast-проблема в RecSys

В рекомендательных системах входные данные делятся на две категории. Первые — признаки пользователя: история взаимодействий, профиль, геолокация, время суток. Эти данные одинаковы для всех кандидатов в одном запросе. Вторые — признаки кандидатов: ID поста, категория, статистика вовлечённости. Они уникальны для каждого элемента.

На уровне взаимодействия эти два потока сходятся в так называемых interaction-слоях — линейных проекциях, кросс-фичах, attention-механизмах. Проблема в том, что эти слои требуют тензоров с совпадающими размерностями батча. Если в запросе 1024 кандидата и примерно 15 пользователей, эмбеддинги пользователя нужно реплицировать — broadcast — примерно 70 раз, чтобы сопоставить размерность с кандидатами.

Это не мелкая неэффективность. При батчах от десятков до десятков тысяч кандидатов репликация масштабируется линейно и съедает значительную часть пропускной способности памяти и вычислительных ресурсов. Причём чем богаче взаимодействие между пользователем и кандидатом — а современные архитектуры от DLRM до HSTU и Phoenix движутся именно в сторону более глубокого взаимодействия — тем выше эта плата.

Почему broadcast — это проблема раскладки данных, а не вычислений

Ключевой инсайт IKBO формулируется просто: broadcast — это проблема раскладки данных, а не вычислительная необходимость. Если посмотреть на операцию математически, то репликация одинаковых эмбеддингов не добавляет никакой новой информации. Это чисто механическая подготовка данных к формату, который ожидает ядро.

Стандартные подходы до сих пор работали вокруг проблемы, а не устраняли её. Системный broadcast материализует реплицированный тензор перед отправкой на GPU — просто, но расточительно. Разделение сети на RO и NRO подсети уменьшает избыточную работу, но ограничивает архитектуру модели и добавляет накладные расходы при малом размере батча пользователей.

IKBO действует иначе: broadcast устраняется на уровне вычислительного примитива. Ядро принимает входные данные пользователя и кандидатов в их естественных, несовпадающих размерностях батча и обрабатывает broadcast внутри себя. Реплицированные тензоры никогда не материализуются — ни в памяти, ни на шине, ни в кэше.

Как устроен ко-дизайн IKBO

Реализация затрагивает три уровня стека. На уровне ядер — кастомные GPU-ядра, которые принимают несовпадающие размерности RO и NRO и обрабатывают broadcast внутри. На уровне компилятора — спецификация динамических диапазонов форм для выбора подходящих ядер. На уровне инференса — рантайм передаёт mapping кандидатов к пользователям вместо материализации broadcast.

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

В production IKBO развёрнут на всей воронке ранжирования Meta — от ранних стадий до поздних, на GPU и на MTIA, собственном ускорителе Meta. Это не экспериментальная разработка, а инфраструктурный слой, на котором работают LLM-масштабные модели в реальном времени.

Deep Dive: Linear Compression — 4× ускорение за четыре стадии

Первое ядро, которое разбирают авторы, — Linear Compress Embedding. Это широко используемая в моделях Meta операция сжатия входных эмбеддингов через обучаемую проекцию. Базовая версия вычисляет один батчевый матмул по всем кандидатам. IKBO проходит четыре стадии прогрессивной оптимизации.

Стадия первая — декомпозиция матмула. Поскольку матрица весов W не зависит от батча, операцию можно разложить по линейности: отдельно вычислить пользовательскую и кандидатную части, а сложить результаты в конце. Это сразу даёт 28.5% ускорения — с 1.944 мс до 1.389 мс. При этом дедупликация устраняет более 98% пользовательской работы: батч с 1024 элементов сокращается до примерно 15 уникальных пользователей.

Стадия вторая — оптимизация раскладки памяти. Анализ показывает, что после декомпозиции узким местом становится L1/TEX-конвейер, загруженный на 84%, в то время как DRAM используется лишь на 19%. Причина — неоптимальное выравнивание: каждая загрузка копирует по 4 байта вместо одной 128-битной операции. Решение — дополнить K до следующего кратного 8 нулями. Это математически эквивалентно исходной операции, но позволяет использовать TMA-загрузки, полностью обходя L1/TEX. Латентность падает с 1.389 мс до 0.798 мс — ещё 42.5%.

Стадия третья — fusion broadcast в ядро кандидатного GEMM. Оставшиеся unfused broadcast и сложение оказываются связаны пропускной способностью памяти: запись результата кандидатного GEMM в HBM, чтение обратно вместе с пользовательским результатом, сложение и запись снова. IKBO устраняет это, встраивая broadcast в эпилог кандидатного GEMM: после накопления каждого тайла эпилог загружает предвычисленный пользовательский результат, складывает в регистрах и записывает финальную сумму. Промежуточный тензор никогда не материализуется. Результат: 0.798 мс → 0.580 мс, ещё 27.4%.

Стадия четвёртая — warp-специализированная многостадийная fusion через TLX. Triton Low-Level Extensions открывают доступ к warp-специализации, TMA, mbarriers и именованным барьерам Hopper, сохраняя Python DSL Triton и инфраструктуру автотюнинга. Проблема предыдущей стадии — низкая occupancy: 6.25% с одним варпом на планировщик. TLX решает это через функциональное разделение вместо дополнительных варпов. Пользовательский GEMM и кандидатный GEMM с broadcast-фusion выполняются конкурентно в разных warp-группах одного CTA, с синхронизацией через release-acquire семантику. Итоговая латентность: 0.482 мс. Кумулятивное ускорение относительно базовой линии — примерно 4×.

Deep Dive: Flash Attention — от IO-bound к compute-bound

Второе ядро — IKBO-адаптированный Flash Attention. В рекомендательных системах attention занимает примерно 40% латентности инференса при длине последовательности 1024. Проблема в том, что стандартный attention с broadcast K и V оказывается сильно связан пропускной способностью памяти: арифметическая интенсивность около 60 FLOPs/Byte при пике H100 SXM5 около 495 FLOPs/Byte.

IKBO решает это через амортизацию доступов к K/V-памяти across кандидатов, разделяющих одного пользователя. Арифметическая интенсивность растёт с ~60 до ~833 FLOPs/Byte при соотношении кандидатов к пользователям 70:1, и ядро переходит в compute-bound режим. Для максимизации эффекта реализуется переупорядочивание сетки запуска threadblock: batch_size_candidate идёт перед num_heads, что обеспечивает конкурентное выполнение threadblock'ов с разными кандидатами но одним пользователем и улучшает reuse L2-кэша.

Сравнение с cutting-edge реализациями Flash Attention на Hopper показывает: Triton IKBO FA2 достигает 425 TFLOPs/s при 487 GB/s IO и 0.321 мс латентности. TLX IKBO FA3 с персистентностью — 621 TFLOPs/s при 455 GB/s и 0.230 мс. Для сравнения, CuTeDSL FA4 Hopper без IKBO даёт 518 TFLOPs/s, но с дополнительными 0.912 мс на broadcast K и V. В пересчёте на полный pipeline IKBO обеспечивает 2.4× ускорения по чистому ядру и 6.4× с учётом broadcast.

Self + Target Attention Fusion через ко-дизайн модели

Естественный вопрос: можно ли слить self-attention и target-attention в одно ядро? Они разделяют один источник K/V — пользовательскую последовательность. Разница только в query: self-attention запрашивает из пользовательской стороны, target-attention — из кандидатской. Разделяя K/V-проекции между ними, IKBO обеспечивает горизонтальную fusion внутри одного запуска.

Это даёт три преимущества. Во-первых, снижение накладных расходов на запуск ядер. Во-вторых, экономия SMEM за счёт совместного использования K/V. В-третьих, смягчение wave quantization — эффекта, при котором неполное число warp-групп на SM оставляет ресурсы незагруженными. В production разделяемые K/V-проекции дают дополнительную экономию на стоимости линейных проекций, аналогичную reuse KV-кэша в LLM-инференсе.

Что это значит за пределами Meta

IKBO демонстрирует принцип, который применим далеко за пределами рекомендательных систем. Любая задача, где один набор данных разделяется между множеством вычислительных элементов, потенциально содержит аналогичную broadcast-избыточность. Это может быть batch processing в компьютерном зрении, множественные запросы к одному контексту в RAG-системах, или параллельная оценка гипотез в автономных агентах.

Ко-дизайн как методология также заслуживает внимания. Оптимизация на одном уровне — будь то только ядро, только модель или только инфраструктура — даёт локальные улучшения. Но фундаментальные прорывы происходят, когда оптимизации на разных уровнях усиливают друг друга. Декомпозиция матмула бессмысленна без модельной поддержки выравнивания. Warp-специализация неэффективна без инфраструктурного mapping кандидатов к пользователям.

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

Часто задаваемые вопросы

Работает ли IKBO только на H100?

Нет. IKBO развёрнут и на GPU, и на MTIA — собственном ускорителе Meta. В статье фокус на H100 для иллюстрации принципов, поскольку Hopper предоставляет TMA и warp-специализацию, которые максимизируют выигрыш. Ядра реализованы на Triton и TLX, что обеспечивает переносимость на разные архитектуры.

Нужно ли переписывать модель для использования IKBO?

Не обязательно. Прямое внедрение требует изменений в коде модели, но inference-time transformation автоматически заменяет стандартные операции на IKBO-эквиваленты во время инференса без изменений в модели. Это позволяет применять IKBO к существующим production-моделям.

Применима ли эта идея к LLM-инференсу?

Прямое применение ограничено, поскольку в LLM нет явного разделения на RO и NRO батчи в том же смысле. Однако принцип устранения избыточного broadcast и ко-дизайна ядер с моделью релевантен для оптимизации grouped query attention, speculative decoding и других сценариев, где один контекст разделяется между несколькими запросами.

Итог

IKBO — это не просто оптимизация CUDA-ядра. Это демонстрация того, как фундаментальное переосмысление проблемы на трёх уровнях — вычислительном, модельном и инфраструктурном — позволяет устранить узкое место, которое раньше считалось неизбежным. До 4× ускорения на ключевых операциях, две трети сокращения вычислительной латентности, и всё это в production-масштабе Meta.

Для разработчиков ИИ главный урок в том, что broadcast — это проблема раскладки данных, а не вычислений. И когда вы начинаете видеть подобные избыточности в своих системах, правильный вопрос не «как сделать broadcast быстрее», а «как сделать так, чтобы broadcast не происходил вообще».