IBM Spyre стал нативным устройством PyTorch: как это устроено

IBM Spyre стал нативным устройством PyTorch: как это устроено

В блоге PyTorch вышла инженерная история о том, как ускоритель IBM Spyre стал нативным устройством фреймворка. Звучит узко, но за ней стоит ответ на вопрос, который в 2026 году возникает всё чаще: как PyTorch принимает железо, устроенное не так, как GPU. Результат перехода измеримый: на инференсе Granite 3.3 8B новая интеграция ускорила prefill в 1,7 раза, а decode в 2,4 раза.

Spyre никогда не был похож на видеокарту, и именно поэтому его подключение к PyTorch оказалось нетривиальной задачей. Вместо последовательности инструкций над потоками ядер чип исполняет потоки данных между вычислителями, хранит тензоры в собственной иерархии памяти и считает в пониженной точности. Чтобы стать полноправным устройством PyTorch, ему пришлось ответить на четыре вопроса: как выглядеть в диспетчере, кто владеет памятью, как готовится запуск программы и как перекрывать независимую работу. Разбор IBM это редкая возможность заглянуть в устройство таких интеграций изнутри.

Что такое Spyre

Spyre это инференс-ускоритель IBM из линейки AIU (Artificial Intelligence Unit) для корпоративных систем: он работает рядом с данными и приложениями на IBM Z, LinuxONE и Power. На карте 32 ядра, соединённых высокоскоростным кольцом; у каждого ядра 2 МБ локальной памяти (scratchpad) и массивы вычислительных элементов. До 128 ГБ LPDDR5 хранят тензоры и программы, а данные движутся фиксированными блоками по 128 байт, которые в статье называют sticks. Вычисления в пониженной точности спроектированы под инференс: высокая пропускная способность при низком энергопотреблении. Такое железо интересно корпорациям, которым нужно гонять модели рядом со своими данными, не отправляя их в публичное облако.

Ключевая деталь для понимания всей истории: двумя уровнями памяти владеют разные стороны. LPDDR5 распределяет и ведёт рантайм, а маршруты переноса данных между LPDDR5 и scratchpad определяет компилятор ещё до запуска программы. Непрерывность и выравнивание данных влияют на скорость подкачки. PyTorch обязан уважать эту модель, а не абстрагироваться от неё.

Почему прежний путь упирался в граф

Первая интеграция Spyre тоже начиналась с FX-графа PyTorch, и граф проблемой не был: рядом с компилятором ему и место. Проблема начиналась после передачи управления. Запуски, копирования и тензорные аргументы транслировались во второй граф, уже собственный для рантайма ускорителя. Чтобы исполнить однажды скомпилированную программу, рантайм читал и десериализовал этот граф, достраивал вершины под аргументы, парсил результат и только затем выдавал операции на устройство. Всё это повторялось на каждый запуск: граф читался с диска и разбирался заново. Копирование между хостом и устройством тоже строило вокруг каждой передачи маленький граф.

Прослойка была полезна, когда рантайму нужен целый граф на входе, но она добавляла обработку графа каждому запуску, который к этому моменту уже полностью определён. Хуже того, когда скомпилированный граф копировался в собственные абстракции рантайма, между ним и PyTorch рвалась коммуникация: фреймворк не видел, где лежит тензор, и считал его CPU-тензором, даже если рантайм давно перенёс данные на Spyre.

Что значит быть нативным устройством

В PyTorch для внешних ускорителей есть механизм PrivateUse1. Регистрация устройства начинается с пары вызовов: rename_privateuse1_backend("spyre") даёт бэкенду имя, _register_device_module("spyre", ...) регистрирует его модуль. После этого tensor.to("spyre") работает так же, как для любого другого устройства, а диспетчер направляет операции в кернелы Spyre.

Дальше бэкенд обязан наполнить контракты фреймворка: интерфейс аллокатора (кто выдаёт память устройства и следит за временем жизни), хуки потоков (что именно делает torch.Stream(device="spyre")), поверхность torch.accelerator для переносимого кода. Нативное устройство это не косметика под API, а настоящий резидентный тензор, памятью которого владеет PyTorch. Иначе каждый мелкий шаг в eager-режиме превращался бы в границу «скопировать на ускоритель и обратно».

Ещё свойство нативного дизайна: компилируемый и eager-путь используют одну декомпозицию. Первый eager-вызов операции компилируется и кэшируется, последующие переиспользуют готовое, а внутри torch.compile тот же код захватывается в общий FX-граф. У этого есть цена: компиляция при первом использовании, возможные фолбэки на неподдерживаемых операциях и новые артефакты под другие размеры тензоров. Зато исправления и регрессии действуют сразу в обоих режимах.

Четыре проекции PyTorch на железо Spyre

Самая содержательная часть разбора это четыре отображения, которые связывают абстракции фреймворка с физикой чипа.

Потоки: от описания программы к типизированным операциям

Работе нужно место, где она упорядочена, асинхронна и привязана к устройству. В PyTorch это поток. Для Spyre важно не то, что железо ведёт себя как поток PyTorch, а то, что упорядоченная потоком работа понижается в типизированные операции рантайма, а зависимости этих операций рантайм сохраняет сам. Сначала компилятор превращает артефакт в launch plan: рецепт, в котором записаны хостовая работа, перемещения данных, вычисления на устройстве и синхронизации. После отправки в поток нижний уровень видит операции с явными операндами и зависимостями, а не описание программы, которое надо переинтерпретировать. Переинтерпретировать нечего: описание до него не доходит.

Память: регионы и ручки вместо плоского пула

Аллокаторы ускорителей обычно представляют память как плоский пул, адресуемый указателями. У Spyre под этим интерфейсом лежат регионы, непрерывные куски памяти устройства, каждый со своей ручкой (handle). Регион вмещает размещения для множества тензоров, а бюджет ручек ограничен, особенно когда карта делится между несколькими арендаторами. Поэтому аллокатор должен упаковывать много тензоров в небольшое число регионов. Для eager-режима это не оптимизация, а условие жизни: без резидентных тензоров каждый мелкий шаг создавал бы новую границу переноса данных.

Запуск: три подхода и патчинг на месте

У любого ускорителя с ahead-of-time компиляцией есть разлом: компилятор знает структуру вычисления, но не знает, где окажутся данные вызывающего, и до запуска эти картины надо согласовать. Spyre прошёл три подхода. Сначала адресами владела скомпилированная задача: входы копировались во внутренние буферы, результаты копировались наружу. Такое до сих пор годится для памяти, которой владеет сама задача, то есть промежуточных и временных буферов компилятора, но не для тензоров, которые размещает аллокатор PyTorch. Потом появился промежуточный вариант: таблица трансляции (xlat) с окнами адресов, и каждое окно привязывалось к одному тензору. Патчить ничего не требовалось, но упор упирался в число входов. Сегодня рантайм использует только третий подход, runtime patching: адреса тензоров попадают в скомпилированную программу прямо в момент запуска.

Конкурентность: подготовка не ждёт вычисление

Передачи и вычисления идут по разным аппаратным пайплайнам, но один поток заставляет их ждать друг друга. В одиночной последовательности хост готовит данные коррекции, передача несёт их на устройство, вычисление потребляет результат, и только потом начинается следующий шаг. В новом дизайне подготовка итерации N+1 стартует на отдельном потоке, пока устройство считает итерацию N: адреса к моменту запуска уже известны. Событие связывает потоки ровно в точке зависимости, вычисление N+1 дожидается своей передачи, а подготовка хоста не дожидается вычисления N. Хост и пайплайн передач перестают простаивать.

Что дал переход в цифрах

Аргумент производительности здесь про то, что уходит с пути запуска. Отправка типизированных операций в очередь вместо построения графа убирает из каждого запуска чтение с диска, десериализацию, сшивание аргументов и компиляцию с разбором. На модели Granite 3.3 8B при batch size 1 и длине последовательности 1024 одна только смена маршрута для передач данных дала 9,9% на prefill и 13,6% на decode. Когда тот же переезд касается диспетчеризации вычислений, суммарный выигрыш доходит до 1,7x на prefill и 2,4x на decode, и примерно три четверти его приносит путь вычислений. Ещё одна статья экономии это конкурентность из прошлого раздела: каждая часть подготовки, убранная с пути вычислений, сокращает простой устройства.

Что PyTorch не даёт и что IBM вернула

Абстракция событий у PyTorch есть, но кросс-поточный порядок приходится обеспечивать там, где операции реально стоят в планировщике рантайма, в слое, который видит их след в памяти устройства. Поэтому рантайм Spyre строит программные примитивы событий и производные зависимости сам, а вывести их наружу в публичные хуки это отдельная работа на будущее.

Возврат в экосистему тоже часть истории. IBM поддерживает openreg, открытый референсный бэкенд и симулятор, показывающий, как внешнее устройство регистрируется через PrivateUse1 без правок ядра PyTorch. Интеграция Spyre подключена к Cross-Repository CI Relay: изменение в апстриме прогоняется против Spyre, а результат попадает в CI-дашборд PyTorch. Плюс инструменты отбора и адаптации апстрим-тестов, которые команда планирует отдать сообществу. Бэкенд, который только потребляет интерфейс, оставляет интерфейс не лучше, чем нашёл.

Почему это важно за пределами Spyre

Выигрыш тут двусторонний. Пользователи Spyre получают знакомую поверхность: размещение тензоров, потоки, компиляцию. PyTorch получает доказательство, что его точки расширения выдерживают железо, устроенное не как GPU. В 2026 году это уже не одиночный случай: TPU, Trainium и другие ускорители идут в PyTorch тем же маршрутом, а мультисиликоновая тема стала одной из главных на PyTorch Conference, которая пройдёт 20 и 21 октября. Для команды, которой предстоит подключить следующее необычное железо, опыт torch-spyre сокращает путь: контракты PrivateUse1, уроки про аллокатор и запуск теперь описаны, а не открыты заново.

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

Что такое IBM Spyre простыми словами?

Это ускоритель IBM для инференса в корпоративных системах. Внутри 32 ядра, соединённых кольцом, у каждого своя локальная память (2 МБ), а до 128 ГБ LPDDR5 хранят тензоры и программы. Вычисления идут в пониженной точности, а программа задаёт поток данных между ядрами, а не последовательность инструкций, как на GPU.

Чем нативное устройство отличается от кастомного бэкенда?

Нативное устройство реализует контракты PyTorch: аллокатор, потоки, диспетчер, компиляцию. Тензоры живут в памяти устройства под управлением PyTorch, а один путь исполнения обслуживает и eager, и компилируемый режим. Кастомный бэкенд решает частную задачу, но тащит собственную машинерию и собственную коммуникацию с фреймворком, как показала история со вторым рантайм-графом.

Даёт ли это ускорение на практике?

Да, и оно измеримо. На Granite 3.3 8B при batch size 1 и длине последовательности 1024 смена дизайна запуска ускорила prefill в 1,7 раза, а decode в 2,4 раза. Часть прироста пришла только от маршрута перемещений данных: 9,9% и 13,6% соответственно.

Итог

История Spyre это учебник по подключению необычного железа к зрелой экосистеме. Сначала определить, какие абстракции фреймворка ложатся на физику чипа, а потом довести каждую до контракта, а не до совместимой имитации. IBM не заставила Spyre притворяться видеокартой, и PyTorch не пришлось заводить отдельные правила для каждого ускорителя. Если следите за мультисиликоновым сдвигом, посмотрите репозиторий torch-spyre и блог PyTorch: разборы архитектур новых бэкендов выходят там регулярно.

← Все записи