Прорыв в оптимизации SSM-моделей
Компания Jane Street, известная своими сложными количественными исследованиями, совместно с NVIDIA представила высокопроизводительную реализацию архитектуры Mamba-2. Ключевым инструментом стал NVIDIA CUDA Tile (cuTile) — Python-based DSL, который абстрагирует управление потоками и позволяет автоматически использовать Tensor Cores без написания низкоуровневого CUDA C++. Это решение позволяет обойти традиционные компромиссы между скоростью обучения (матричные умножения) и эффективностью вывода (постоянная память).
Три оптимизированных ядра (Kernels)
Для достижения максимальной производительности команда выделила три критических этапа прямого прохода (forward pass) и реализовала их как кастомные ядра:
- chunk_cumsum: Вычисляет кумулятивные суммы по токенам и головам в логарифмическом пространстве для термина затухания (decay term). Результаты сохраняются для повторного использования.
- fwd_pass_state: Использует
ct.mmaдля выполнения операций матричного умножения на Tensor Cores. Это ядро продвигает скрытое состояние (hidden state) по чанкам последовательности, применяя формулу прыжка (jump formula) с экспоненциальным затуханием. - compute_out: Фьюжит (объединяет) три операции в одно ядро: учёт перенесённого состояния, взаимодействие с причинно-следственной маской (causal mask) и skip-связь. Промежуточные матрицы хранятся в регистрах, а данные приводятся к TF32 для эффективной работы Tensor Cores.
Технические детали оптимизации
Архитектура Mamba-2 описывается рекуррентным обновлением: h_t = exp(A * dt_t) * h_{t-1} + dt_t * B_t * x_t. Прямой расчет через цикл слишком медленен для обучения, а параллельный скан (parallel scan) требует огромных объемов памяти. Решение Jane Street использует «формулу прыжка» для вычисления состояния чанками через матричное умножение, что позволяет эффективно использовать Tensor Cores.
Ключевые технические приемы, примененные в коде:
- Latency hints: Для точного моделирования трафика DRAM.
- Operand swapping: Перестановка операндов для маршрутизации данных в Tensor Memory.
- Transposed layouts: Транспонированные макеты памяти для совместимости с TMA (Tensor Memory Accelerator).
- Autotuner: Автоматический поиск оптимальных размеров блоков (block size) и параметров занятости (occupancy) для каждой формы входных данных.
Значение для индустрии
Эта работа демонстрирует, что сложные SSM-модели (State Space Models), такие как Mamba-2, могут быть реализованы на уровне производительности, сопоставимом с Transformer-архитектурами, но с меньшими требованиями к памяти при инференсе. Использование Python-DSL для написания высокопроизводительных GPU-ядер снижает порог входа для ML-инженеров, позволяя им фокусироваться на алгоритмической эффективности, а не на низкоуровневой оптимизации памяти.
Источник: NVIDIA dev blog ↗