Почему Mamba-2 ускоряется в cuTile: три ядра и пределы сравнения
Разбираем, как chunk_cumsum, fwd_pass_state и compute_out превращают рекуррентный Mamba-2 в работу для Tensor Cores — и почему 1,14–1,82× нельзя переносить на обучение целиком.
NVIDIA и инженеры Jane Street показали не новую версию Mamba-2, а конкретную реализацию её прямого прохода на CUDA Tile, или cuTile. В ней вычисление разделено на три ядра: одно заранее считает коэффициенты затухания, второе переносит скрытое состояние между блоками последовательности, третье собирает выходы токенов. По данным NVIDIA, этот набор ядер оказался в 1,14–1,82 раза быстрее сравниваемой реализации на Triton. Но цифра относится к микроизмерению forward pass на одной NVIDIA B300 и не описывает скорость обучения или приложения целиком.
Ценность разбора не в самом диапазоне ускорения. Пример хорошо показывает, как преимущество архитектуры превращается в производительность только после работы с разбиением на блоки, раскладкой тензоров и перемещением данных. Для разработчика это полезнее рекламного вывода «Mamba быстрее attention»: видно, какую именно работу выполняет GPU и где заканчивается применимость сравнения.
Зачем Mamba-2 разбивает последовательность на блоки
Обычное causal attention сопоставляет каждый новый токен с предшествующими. При обучении и prefill это удобно выражается матричными умножениями, но объём попарных взаимодействий растёт квадратично с длиной последовательности. При авторегрессионной генерации вычисление на очередном шаге и память под KV-кэш растут вместе с контекстом. У state-space model другая форма: прошлое сжимается в скрытое состояние фиксированного размера, которое обновляется рекуррентно. Поэтому состояние не удлиняется с каждым токеном.
Такое отличие не означает, что SSM автоматически лучше attention. Последовательная рекуррентная формула удобна для генерации, но плохо загружает параллельный GPU во время обучения и prefill. Mamba-2 использует state-space duality: ту же динамику можно разложить на блоки токенов и внутри блока выразить через матричные операции, похожие на линейное attention. Между блоками переносится только состояние на границе. В результате длинная последовательность обрабатывается матричными умножениями, пригодными для Tensor Cores, без сохранения отдельного большого состояния для каждого шага.
Затухание памяти задаётся коэффициентом, который умножает прежнее состояние. Если перемножать такие коэффициенты последовательно, появляется длинная зависимая цепочка. Реализация переводит их в логарифмическое пространство: произведение превращается в накопленную сумму G. Для перехода через блок достаточно знать сумму к его концу и разности внутри блока. Это и есть связка между первым ядром, которое готовит G, и двумя следующими, которые используют её в матричных формулах.
Первое ядро: chunk_cumsum готовит затухание
chunk_cumsum получает коэффициент затухания для каждой головы и временной шаг для каждого токена. Оно перемножает их и считает накопленную сумму вдоль токенов внутри одного блока. Результат сохраняется, потому что он нужен и для обновления состояния, и для вычисления выходов. Это не scan по всей последовательности: каждая сумма начинается заново на границе блока, а связь между блоками обеспечивает скрытое состояние.
В коде результат дополнительно масштабируется на log2(e), чтобы следующие ядра могли применять exp2 вместо exp. Такая замена не меняет математический смысл коэффициента, но использует более подходящую операцию GPU. У этого ядра есть и цена: подготовленный тензор приходится записать в глобальную память, а затем снова прочитать. Поэтому его раскладка должна соответствовать тому, как данные потребят следующие этапы.
Второе ядро: fwd_pass_state переносит состояние
fwd_pass_state не вычисляет ответы для токенов. Оно начинает со скрытого состояния перед первым блоком, сохраняет его как начальное состояние блока, затем одним «прыжком» получает состояние после блока. В формуле старое состояние затухает, а матричное произведение добавляет вклад всех токенов. После последнего блока итоговое состояние можно вернуть следующему вызову при инференсе; при обучении оно обычно не нужно.
Зависимость по блокам остаётся последовательной, но она локализована внутри одного экземпляра программы cuTile. Одновременно GPU обрабатывает независимые срезы по batch, головам, части head dimension и части state dimension. Крупное произведение выполняет ct.mma на Tensor Cores. Это важная граница формулировки: алгоритм не устраняет рекуррентность целиком, а уменьшает число последовательных шагов с количества токенов до количества блоков и параллелит остальное.
Третье ядро: compute_out не выгружает промежуточную матрицу
compute_out получает начальное состояние каждого блока и формирует выходы его токенов. В первой фазе оно одновременно строит взаимодействия C·Bᵀ и проецирует перенесённое состояние. Во второй применяет попарное затухание и нижнетреугольную причинную маску, умножает результат на вход X и добавляет skip connection D·X. По форме центральная часть действительно напоминает attention без softmax, но коэффициенты затухания и перенос состояния сохраняют специфику SSM.
Прямая реализация на PyTorch могла бы по отдельности посчитать слагаемые и записать в глобальную память матрицу взаимодействий размером block_size × block_size для каждого блока и головы. cuTile объединяет этапы: матрица C·Bᵀ и аккумулятор выхода остаются в регистрах, а причинная маска создаётся через индексы и условный выбор, без отдельного тензора. Финальное умножение переводит операнды в TF32, чтобы задействовать Tensor Cores при аккумуляции в FP32-подобном режиме. Выход затем приводится к заданному типу.
В блоге NVIDIA отдельно оговорена численная деталь. Для элементов выше диагонали сначала могут вычислиться экспоненты положительных разностей, хотя эти элементы затем обнуляет причинная маска. Авторы реализации утверждают, что немедленное маскирование не даёт потенциальным переполнениям повлиять на сохраняемый результат. Это свойство конкретного порядка операций; его нельзя переносить на любую реализацию Mamba-2.
Почему раскладка памяти важна не меньше формул
Три оптимизации показывают, где cuTile получает выигрыш. Во-первых, подсказки latency для load и store помогают компилятору планировать обращения к DRAM и глубину конвейера. Во-вторых, тензоры decay_cumsum и dt переставлены в порядок [batch, head, token]. При исходном шаговом доступе компилятор получал неудобные LDG-загрузки; непрерывное расположение по токенам позволяет использовать Tensor Memory Accelerator. Это изменение раскладки данных, а не модели.
В-третьих, на Tensor Cores Blackwell операнды MMA несимметричны: операнд A может размещаться в Tensor Memory, тогда как B должен находиться в shared memory. Реализация транспонирует произведение по тождеству BᵀAᵀ=(AB)ᵀ, чтобы составной операнд с уже применёнными dt и затуханием попал в позицию A. Так снижается давление на shared memory. Наконец, exhaustive_search подбирает размер блока и occupancy отдельно для формы входа. Следовательно, опубликованный результат включает не один универсальный вариант ядра, а выбор настройки под конкретную конфигурацию.
Что измерил бенчмарк NVIDIA
NVIDIA сравнила cuTile forward pass с открытой Triton-реализацией Mamba-2 из mamba_ssm 2.3.2.post1. Измерения проводились на одной B300 SXM6 Blackwell Ultra с 268 ГиБ HBM, лимитом мощности 1100 Вт и максимальной частотой SM 1920 МГц. В программной среде указаны cuda-tile 1.5.0, Triton 3.7.1, PyTorch 2.13.0, CUDA 13.3, драйвер 595.58.03, Python 3.12.3 и Ubuntu 24.04.4.
Пример из публикации использует batch 4, 16 384 токена, восемь голов, размер состояния 128, head size 64, block size 128, входы bfloat16 и состояние float32. В исходном сравнении NVIDIA проверены несколько длин последовательности и размеров состояния. В этих условиях NVIDIA сообщает диапазон ускорения 1,14–1,82 раза; преимущество росло вместе с длиной и state size. Это результаты авторов и поставщика платформы, а не наши измерения.
Где заканчивается сравнение
Бенчмарк измеряет только прямой проход трёх ядер. Он не включает backward pass, оптимизатор, обмен между GPU, загрузку данных, остальные слои модели и задержки приложения. NVIDIA прямо пишет, что результат обучения целиком будет зависеть и от реализации обратного прохода. Диапазон также нельзя автоматически переносить с B300 на H100, потребительские GPU или другие версии CUDA и Triton: часть оптимизации опирается на особенности Tensor Memory и Tensor Cores Blackwell.
Сравнение не доказывает и превосходство Mamba-2 над Transformer по качеству модели. Исходная работа Mamba-2 связывает SSM и варианты attention через structured state space duality и отдельно оценивает архитектуру. Публикация NVIDIA отвечает на более узкий вопрос: можно ли эффективно выразить уже известный forward pass в новой tile-ориентированной системе программирования. Её ответ положителен для протестированных форм и оборудования.
Практический вывод состоит из двух частей. Архитектура Mamba-2 даёт возможность обрабатывать длинную последовательность блоками и переносить состояние фиксированного размера. Реальная скорость появляется лишь тогда, когда промежуточные матрицы не уходят в HBM, загрузки становятся непрерывными, операнды попадают в подходящую память Tensor Core, а параметры ядра подбираются под форму. Поэтому 1,82× — верхняя точка конкретного сравнения, а не новое свойство любой Mamba-2 и не обещание ускорить модель или обучение целиком.
