Что делает StreamTensor

StreamTensor — это компилятор, который переводит графы LLM из PyTorch (GPT-2, Llama, Qwen, Gemma) в потоковые dataflow-ускорители для FPGA AMD Alveo U55C. Вместо того чтобы выполнять промежуточные тензоры как пакетные ядра с частыми обращениями к DRAM, StreamTensor прокачивает тайлы через on-chip FIFO и конвертеры, снижая количество обращений к внешней памяти и экономя время и энергию.

Ключевые идеи и поток данных

Модели проходят через Torch-MLIR, понижаются до MLIR Linalg и затем переводятся в dataflow IR. Узлы этого IR становятся аппаратными ядрами с явными stream-интерфейсами и хостовым/рантайм-обвязкой. Процесс автоматизирует генерацию ядер, вставку DMA и сопряжение рантайма без ручного написания RTL.

Важная абстракция компилятора — итеративный тензор (itensor), который кодирует порядок итераций, разбиение на тайлы и расположение данных. Явное представление порядка стрима позволяет безопасно объединять ядра и синтезировать минимальные буферы и конвертеры форматов только там, где это действительно нужно.

Иерархический поиск конфигураций

StreamTensor исследует три взаимосвязанных пространства проектирования:

• Трансформации на уровне Linalg: разбиение на тайлы, разворачивание, векторизация, перестановки
• Выборы по fusion при ограничениях памяти и ресурсов
• Распределение ресурсов и выбор ширин стримов

Иерархический DSE направлен на оптимизацию устойчивой пропускной способности при ограничениях по пропускной способности и on-chip памяти.

Размеры FIFO и предотвращение взаимоблокировок

Вместо эвристик StreamTensor формулирует задачу размера FIFO как задачу линейного программирования. LP вычисляет размеры буферов, которые предотвращают простои и дедлоки при минимальном использовании BRAM/URAM. Формальный подход обеспечивает корректность потоковой передачи между ядрами.

Результаты

Для задач декодинга LLM команда приводит сокращение средней латентности до 0.64× по сравнению с базовой GPU-конфигурацией на GPT-2 и до 0.76× по сравнению с предыдущими FPGA-ускорителями LLM. Энергоэффективность достигает до 1.99× по сравнению с A100 на отдельных моделях. Результаты зависят от модели и настроек и относятся к задачам декодинга.

Платформа и ограничения

Работа нацелена на платформа Alveo U55C (16 GB HBM2, 460 GB/s, PCIe Gen3×16 или dual Gen4×8, 2×QSFP28). Дизайн ориентирован на использование пропускной способности HBM и on-chip памяти, чтобы удерживать промежуточные тайлы в потоке и минимизировать DMA в DRAM.

Значение работы

StreamTensor показывает, что компиляторно управляемый потоковый подход способен заменить многие обращения к DRAM при декодинге LLM, снижая задержки и повышая энергоэффективность на специализированном FPGA-оборудовании. Ключевые элементы — itensor, автоматическая синтез конвертеров, иерархический DSE и LP-базированное определение размеров FIFO — делают межъядерные стримы безопасными и эффективными.

Где читать дальше

Авторы публикуют подробную статью и выкладывают репозиторий на GitHub с туториалами, кодом и ноутбуками. Результаты ограничены задачами декодинга и конкретной аппаратной платформой, но идеи itensor и формального определения размеров FIFO применимы шире.