Разрабатывая Svod
Где мы находимся и кто мы такие
Представь, что у тебя есть идея для AI приложения или сервиса, которые ты хочешь построить и запустить. Ты находишься перед многими выборами: какой язык взять, какой фреймворк выбрать, как оркестрировать его компоненты, запускать его на клиентских устройствах, своих серверах или облаке. Выбор широк в каждом направлении, но реальность сужает его тремя вещами: сроками на разработку, доступностью специалистов и стоимостью эксплуатации.
Ты можешь даже представить себе многомерное пространство, в котором рынок нащупывает конфигурацию, которая подойдёт большинству. Какой статус-кво мы имеем сейчас? Python дал нам пологую кривую обучения для получения дешёвых разработчиков, PyTorch смог консолидировать экосистему готовых решений, NVIDIA раздавила конкурентов софтом и у нас появились облака достойные промышленной эксплуатации. Этот статус поддерживается по спирали: набрав критическую массу решение становится всё притягательнее для выбора за счёт лучшей цены или более богатой экосистемы. Крутая ли получилось конфигурация? При всём уважении к отдельным техническим достижениям и людям за ними стоящими я должен признать, что могло быть и лучше: эксплуатировать Python для высоконагруженных или требовательных к надёжности приложений очень сложно, PyTorch имеет многолетние наслоения дизайна и несёт большой груз обратной совместимости, стоимости ускорителей NVIDIA точно не пошла на пользу монополия.
Откуда я это знаю? Последние пять лет я занимаюсь созданием и эксплуатацией поисково-аналитических систем: я написал ядро поискового движка, который обрабатывает приблизительно 4% от мирового трафика, я написал систему распознавания голоса для медиа в этом трафике, я руководил созданием системы семантического поиска с вычленением NER и кластерного анализа для этого трафика, командой мы написали самый лучший для того момента морфологический анализатор для выкачиваемых текстов. Эти системы эксплуатировались в настоящем аппаратном зоопарке: какие-то различия были сделаны в силу исторических причин (напр. старые поколения CPU не поддерживали векторные инструкции определённой ширины), какие-то были сделаны в рамках разумного дизайна (некоторые архитектуры моделей не могли эффективно использовать большие серверные GPU и потому мы использовали пользовательские).
Мы не могли позволить себе аппаратной избыточности или простоя в связи с недоступностью сервисов, потому нам приходилось использовать Rust вместо Python или C++. Мы старались постоянно выжать максимум из железа, потому мы модернизировали фреймворки для повышения утилизации GPU: я добавлял поддержку Flash Attention в наш форк GGML раньше upstream, я писал несколько кастомных мега-ядер для наших моделей и т.д..
На наших задачах GPU от AMD обходятся на 30% дешевле GPU от NVIDIA и для некоторых задач ускорители от Apple обходятся в четыре раза дешевле ускорителей от AMD.
Мы выбирали Rust как компромиссное решение: для опытного инженера скорость разработки промышленного приложения будет не существенно отличаться от скорости разработки на Python или C++ (статические проверки компилятора позволяют нам чаще использовать заимствование и многопоточность, которую программисты С++ избегают неосознанно). Да, есть языки с зависимыми типами, где модель владения ещё больше защищает нас от некоторых ошибок, но их экосистема сильно меньше. Rust же обладает существенными преимуществами: статическая типизация, отсекающая ещё в compile-time множество мелких ошибок, параллелизация из коробки, уже упомянутые проверки заимствования компилятором и всё расширяющееся сообщество. Благодаря этим качествам, Rust впервые за долгое время стал вторым языком разработки Linux, система Git также двигается в сторону Rust, Android также использует Rust для реализации mission critical компонентов.
Девять месяцев назад я начал этот проект, имея опыт модернизации сторонних фреймворков, вкупе с дизайном, который Джордж Хотц придумал при реализации Tinygrad, чтобы сделать лучшее решение в экосистеме Rust. Я надеюсь, что имеющийся результат сможет сдвинуть статус-кво для принятия решения по выбору технологического стэка для следующих проектов. Сегодня я расскажу, каких успехов уже удалось достичь и что нас ждет впереди.
Про экосистему Rust
Несмотря на то, что экосистема Rust существенно развилась за последние пять лет, в ней всё ещё остаётся зияющая дыра, которая для нас абсолютно критична: наличие хорошего DL фреймворка. Один из моих сотрудников, Денис Залетаев, несколько лет назад делал доклад про DL экосистему в Rust. С тех пор мало что изменилось, но я позволю чуть больше рассказать про ограничения, которые остановили нас от использования конкретных решений As-Is и инвестирования своего времени в их развитие.
Многие, как и мы в прежних своих проектах, используют биндинги к LibTorch или ONNX Runtime, но этот подход усложняет деплой, удорожает поддержку и эксплуатацию софта, создаёт водораздел для средств интроспекции и профилирования. Другие используют нативные реализации на Rust, такие как Candle или Burn, на которых хотелось бы остановиться для понимания общей картины.
Candle
Candle старается быть минималистичной копией PyTorch, стараясь скопировать его интерфейс и некоторые архитектурные решения. Сильной стороной Candle можно назвать количество реализованных моделей, хотя качество их реализации (широта конфигурации модели, реже численная корректность) иногда ниже порога, который мы рассматриваем для промышленного применения. Слабой стороной является низкий потенциал для добавления аппаратных ускорителей и параллельного исполнения, а также низкая производительность.
Candle поддерживает ONNX импорт, но производительность и поддержка операторов (в т.ч. в разрезе дополнительных атрибутов и символьных переменных) ни разу не позволили нам запустить модели, которые мы используем в работе.
Burn
Burn представляет из себя мета-фреймворк, который поддерживает несколько тензорных фреймворков для исполнения модели. Многие возможности описания моделей строятся на явной типизации, что, с одной стороны, позволяет получить больше безопасности на этапе описания модели, но, с другой стороны, сильно усложняет определение конкретной модели на этапе работы программы и требует больших усилий для чтения кода. Также подход, при котором Burn выступает дополнительным слоем абстракции, существенно усложняет отладку и модернизацию конкретных операций.
Burn также разрабатывает свой DSL для описания ядер CubeCL, который поддерживает несколько аппаратных ускорителей и возможность описывать ядра из Rust. На мой взгляд они движутся в правильном направлении в дизайне этого DSL и его возможностей, хотя ряд приятных решений принципиально не позволит им догнать топовые ядра по производительности.
Burn поддерживает ONNX импорт, но он работает за счёт транспиляции ONNX графа в Rust-код на этапе компиляции, что существенно ограничивает горячую замену модели (это всё ещё возможно через реализацию JIT компилятора и логики загрузки через C API, но эргономика этого решения для меня была неприемлема), также при попытке добавить модели, которыми мы реально пользовались, мы получали проблемы с поддержкой операторов (или атрибутов) или численной корректностью вычислений на некоторых платформах.
Устройство
Популярные фреймворки устроены как как слоёный торт: несколько слоёв внутренних представлений, каждый из которых решает реальную задачу. Какой-то слой описывает набор тензорных операций, какой-то превращает этот набор операций в граф, какой-то производит оптимизации, какой-то перекладывает вычислительный граф на предварительно написанные ядра, какой-то позволяет запускать вычисление.
PyTorch имеет 7+ представлений (TorchDynamo, TorchIR / FX Graph, AOTAutograd, ATen IR, PrimIR, Dispatcher и, опционально, TorchInductor), какие-то из них оперируют Python, какие-то C++, какие-то MLIR, какие-то предварительно скомпилированными ядрами. Такая конфигурация, несмотря на то, что она решает реальные проблемы, является дорогой для развития проекта (200+ инженеров, во многом за счёт Meta), очень неудобной для расширения вендорами ускорителей (именно по этой причине хорошо PyTorch работает только с ускорителями NVIDIA) и сложной для модернизации пользователями.
Повторить эту конфигурацию, переписав, например, PyTorch, никогда не было хорошим выходом: очень тяжело быть PyTorch лучше самого PyTorch (Candle пытается идти этой дорогой, но ресурсы не равны), эффективное расширение аппаратной поддержки будет очень тяжёлой задачей (стоимость выполнения каждой операции может различаться от платформы к платформе, эффективная работа потребует архитектуру оптимизатора, которую мы, как отрасль, пока не придумали, и я уверен, что сложность реализации будет преступно высокой).
Именно это рассуждение привело меня к тому, чтобы искать другие, более простые архитектурные конфигурации. Я сосредоточился на поиске архитектуры, которая давала бы приемлемую скорость работы на большом количестве железа, но при этом позволяла бы очень легко выжать производительность для какого-то конкретного ускорителя.
В результате поисков мной было найдено подобное решение на Python – Tinygrad за авторством Джорджа Хотца. Небольшой проект, реализующий JIT-ориентированный DL фреймворк с простым оптимизатором, поддержкой ONNX импорта и множества бэкендов. Хотц использует одно внутреннее представление для всех этапов работы фреймворка, что позволяет ему сохранить прозрачность отладки и упростить понимание.
Один UOp, чтобы править всеми
Как представить вычислительный граф, который включает в себя весь жизненный цикл тензора? В Tinygrad Джордж использует enum для кодирования типа операции и список произвольной длины для параметров. Лаконично, но требует дополнительно валидации при конструировании. Rust решает эту проблему за счёт того, что члены его enum могут быть структурами с полями, которые гарантируют нужное количество аргументов.
Для преобразований этого графа можно использовать как обычные функции Rust, так и экосистему Svod Rewrite Engine: при помощи derive-макроса для графа генерируется дополнительная мета-информация, которая может быть использована другим derive-макросом для описания правил преобразования графа, верифицируемых на этапе компиляции.
Добавление аппаратного ускорителя также не представляет большой сложности, требуется выполнить три шага:
- Добавить кодогенератор для целевой платформы: можно использовать текстовую генерацию, можно программную, если бэкенд позволяет.
- Реализовать трейт для менеджмента буферов: аллокация, деаллокация, копирование на хост и на другие устройства такого же типа.
- Реализовать трейт для запуска ядер для платформы: загрузка ядер на устройство, передача буферов и профилирование.
Svod предоставляет средства, которые упрощают добавление аппаратных ускорителей, код для которых похож на Си или ускорителей, которые поддерживаются LLVM, но это не является принципиальным ограничением.
Слой совместимости
Совместимость – притягательное слово. Оно означает, что мы сможем получить какую-то часть глобальной экосистемы DL бесплатно. Я сосредоточился на трёх аспектах совместимости:
- API. С поправкой на семантику, я постарался сделать его максимально близко к PyTorch, я даже сохранил именованные
аргументы при помощи
#[bon], чтобы упростить освоение знакомыми с PyTorch программистами и LLM-агентами. - ONNX. Любая модель, которая экспортируется в
.onnx, должна импортироваться в набор входных тензоров, аргументов и выходных тензоров, вычисляемых на целевом ускорителе. - Бинарная совместимость. Модели с HuggingFace должно быть можно просто скачать и запустить, без дополнительных преобразований в экзотичные
бинарные форматы, также на поддерживаемых системах должна быть возможность прямой (
mmap) загрузки тензоров с диска.
ONNX
Поддержка ONNX является более сложной задачей, чем может показаться на первый взгляд: стандарт содержит 204 операции, многие из которых имеют несколько
версий реализации, опциональные параметры и символьные аргументы. Для того, чтобы поддержать все эти операции, фреймворк должен либо реализовать каждую
операцию отдельно, либо иметь универсальное представление вычислимой операции, на которое декомпозируется любая сложная операция. Я пошёл второй дорогой,
добавив поддержку 162 операторов, сознательно исключив поддержку операций связанных с обучением, сырым текстом (напр. regex), квантованием весов и
обработкой сигналов.
Оценка реального покрытия операторов и соответствие спецификации не является лёгкой задачей, потому ONNX предоставляет 1361 набор тестовых данных для каждого оператора, а также девять легковесных моделей с референсными значениями, которые можно использовать для проверки. Тестовая система Svod использует все эти данные для оценки покрытия и я планирую в ближайшее время добавить Svod в ONNX Backend Scoreboard.
Аппаратная поддержка
Я разделяю аппаратные платформы, на которых мы хотим запускать Svod, на три группы: серверные ускорители, пользовательские ускорители и встраиваемые системы.
Таким образом мы получаем три группы бэкендов под одно и то же API:
- серверные ускорители (AMD MI300/MI350/MI450, NVIDIA H100/H200/B200)
- пользовательские (AMD Ryzen AI Halo, Apple M3/M4/M5, NVIDIA RTX 30/40/50)
- встраиваемые системы (Qualcomm Snapdragon X, RockChip RK3588)
Поддержав эти три типа платформ мы можем описывать и разворачивать модели одинаковым образом на всём этапе жизненного цикла ПО, а также откроем большие перспективы для эксплуатации моделей на пользовательском железе.
На текущий момент я опубликовал только поддержку ускорителей AMD, т.к. они являются для меня целевой платформой и у меня нет свободного времени на причёсывания и публикацию остальных бэкендов (проект разрабатывается в моё свободное время без какого-либо финансирования).
AMD
Компиляция
Для получения исполняемого кода для GPU AMD использует язык HIP (диалект С++), который при помощи компилятора hipcc превращает в LLVM IR, который LLVM
компилирует в целевые представления для выбранной GPU. Я не хотел зависеть от инфраструктуры AMD в этом вопросе, потому решил генерировать
LLVM IR самостоятельно. Такой подход позволил мне упростить процесс отладки и кодогенерации, а также уменьшил размер Docker-образа, который
поставляется пользователю.
Запуск
Для запуска исполняемого кода AMD использует HIP/ROCr инфраструктуру, которую я также не хотел использовать, чтобы не иметь лишней зависимости. Потому я реализовал подход из ROCr прямо в Svod, попутно реализовав kernel launch fusion (идеи из AMD MIGraphX), сильно облегчив запуск больших графов вычислений.
Со временем я хочу доделать userspace-драйвер, который позволит производить запуск исполняемого кода полностью минуя инфраструктуру AMD, что сделает наше ПО полностью самодостаточным и существенно упростит эксплуатацию.
Профилирование
Для профилирования исполняемого кода AMD использует rocprof, который позволяет получить счётчики отдельных типов операций, позволяющие понять сколько
времени исполнения ядра мы ожидали данных, как хорошо попадали в кэш, насколько утилизирован матричный движок и т.д.. Движок запуска ядер получился
настолько удачным, что я расширил его функциональность для извлечения этих счётчиков напрямую, что сделало использование rocprof нецелесообразным.
Tiled Kernels
Долгое время существовал тренд на то, что писать собственные ядра – это плохая практика: если у вас есть хороший тензорный фреймворк, поддерживающий GPU, и ваша задача хорошо ложится на тензорную математику, то для хорошего результата нужно только описать вашу задачу в фреймворк.
Реальность, конечно же, оказалась сложнее:
- Каждый тензорный фреймворк обладает своими архитектурными особенностями, которые могут помешать ему показать достойную производительность на вашей задаче.
- Даже в ситуации, когда фреймворк демонстрирует достойную производительность на вашей задаче, остаётся ощутимый зазор по производительности, который вы не сможете преодолеть.
Это привело к тому, что у каждого серьёзного фреймворка есть возможность добавить в граф вычисления собственные ядра. Подход к написанию этих ядер также эволюционировал во времени: сначала все писали их на Cuda, после появился Triton, сейчас на фронтире используется CuTile. Несмотря на эволюцию подходов к написанию ядер одна особенность остаётся неизменной: пользовательские ядра пишутся не в тех же представлениях, что живёт остальной фреймворк, что сильно усложняет или делает невозможной интроспекцию и отладку производительности всего приложения в целом.
Svod пошёл другой дорогой: совместив подходы HazyResearch HipKittens и NVIDIA CuTile я создал диалект, который позволяет писать высокопроизводительные ядра в терминах остального фреймворка, дотягиваясь до конкретных аппаратных возможностей при помощи интринсиков. Основным прицелом при создании этого диалекта было то, чтобы он получился простым и защищённым в достаточной мере, чтобы даже LLM агент был в состоянии написать мега-ядро самостоятельно или с минимальным участием человека.
Получившийся DSL позволяет писать ядра, которые ничуть не медленнее реализации из hipBLASLt или composed_kernels, но существенно меньшим объёмом кода. Показательно минимальное ядро матричного умножения — оно умещается в несколько строк и при этом задействует матричный ускоритель GPU (MFMA на CDNA, WMMA на RDNA):
fn micro_matmul(ker: &Kernel) -> Arc<UOp> {
let w = ker.warp();
// 64×64 tiles: A and B in bf16, accumulator C in f32.
let a = ker.rt((64, 64), DType::BFloat16, Row, RT_16X16);
let b = ker.rt((64, 64), DType::BFloat16, Col, RT_16X16);
let c = ker.rt((64, 64), DType::Float32, Col, RT_16X16);
// One mma_ab call → compiles to a single matrix-core instruction.
let out = w.mma_ab(w.zero(c), &a, &b);
ker.finish(1)
}
Я реализовал некоторые примитивы для ML и анализа данных, со следующими результатами:
Модели и пайплайны
Сам тензорный фреймворк, какой хороший бы он ни был, не интересен сообществу сам по себе, если у него нет экосистемы описанных
моделей и обвязки для них, которые позволили бы быстро, как из кубиков, собрать готовые системы. Именно по этой причине
transformers является настолько популярным решением и именно по этой причине transformers предоставляют pipelines.
С этими рассуждениями я решил разместить модели и инфраструктуру для пред- и пост- обработки прямо в проект:
- Я выбрал несколько направлений, которые, как мне кажется, будут востребованы у ML инженеров: манипуляции с текстом (генерация эмбеддингов, реранкинг, классификация токенов), манипуляции с изображением (детекция объектов, сегментация, генерация эмбеддингов), манипуляции с аудио (VAD, транскрибация). Для каждого направления я портировал модели, которые являются SOTA или имеют широкое использование. Эти модели могут быть использованы сами по себе для построения RAG-систем, LLM harness, систем речевой аналитики, и т.д..
- Я добавил несколько трейтов и структур, которые позволяют сопрягать несколько моделей в один пайплайн обработки, который прячет в себя всю сложность взаимодействия моделей. Также такой подход позволяет пользователю генерализировать логику вызова моделей, что сильно упрощает сложность миграции на другую модель в дальнейшем.
Например, пайплайн речевой аналитики собирается из модели GigaAM и сегментатора FireRedVAD в несколько строк:
let model = GigaAm::from_hub_with_revision("vpermilp/GigaAM-v3", "ctc")?;
let bounds = EncoderBounds {
sample_rate: model.config.sample_rate as u32,
hop_length: model.config.hop_length,
subsampling_factor: model.config.subsampling_factor,
max_mel_frames: model.config.max_mel_frames,
recommended_target_secs: model.recommended_chunk_secs(),
};
let splitter = FireRedVadSplitter::from_hub(&bounds)?;
let mut asr = Asr::assemble(splitter, |mc| GigaAmTranscriber::new(model, options, mc))?;
let result = asr.transcribe_default(&waveform)?;
Похожесть Svod Tensor API на PyTorch API позволило переписывать модели почти автоматически при помощи LLM: GLM 5.2 почти всегда справляется с миграцией модели и последующим тюнингом производительности при правильно выстроенной петле тестирования (в т.ч. сравнении с референсными значениями) и хорошо написанном промпте. В дальнейшем я планирую оформить несколько LLM Skills для того, чтобы упростить портирование каждому пользователю.
Также схожесть API позволила получить объем кода для описания моделей, сравнимый с PyTorch.
Чтобы сравнение было корректным, я считал только код, специфичный для самой модели (архитектура + inference-обвязка).
Однако в сравнение попали и не отделяемые со стороны Svod-а собственные реализации бэкбонов
(XLM-RoBERTa для BGE-M3 — 616 LOC, Qwen3-decoder — 630 LOC), тогда как Python-сторона импортирует их из
transformers/sentence-transformers, что не учитывается в этих числах. Помимо бэкбонов,
Python-эталон импортирует и inference-обвязку (транспилляцию в JIT, лоадеры), что также оказалось вне подсчетов.
При таком определении «модельного кода» Svod укладывается в 1.2–1.9× от эталона — это сопоставимый объём для реализации с нуля на Rust.
В сухом остатке…
Проект приближается к релизу, у меня получилось приблизительно 83 KLOC кода самого фреймворка и где-то 177 KLOC тестового покрытия и инфраструктуры. Сейчас поддержано несколько CPU (x86, ARM, RISC-V, IBM), несколько операционных систем (Linux, OS X) и несколько ускорителей от AMD (RDNA 3.5, RDNA 4, CDNA 3). Расход строк получается больше, чем в Tinygrad, за счёт большей явности обработки типов, наличия документации и отсутствия насилия над форматером.
Svod существенно проигрывает по производительности PyTorch и TensorFlow на произвольной модели, но за счёт развитых инструментов отладки и профилирования позволяет быстро найти узкое место и реализовать решение под имеющееся железо при помощи Svod TK DSL: при реализации поддержки GigaAM я потратил меньше 600 строк кода (пока не опубликованного), которые позволили получить мне производительность сравнимые с RTX4090 в связке с PyTorch у референсной реализации.
Это не конечная форма фреймворка, но хорошая опорная точка, чтобы продемонстрировать его сообществу, получить обратную связь от него и изменить те вещи, которые исправить потом будет сложно.
Перспективы развития
AOT компиляция
Сейчас каждый запуск моделей с нуля оптимизирует и компилирует вычислительный граф под каждый ускоритель, что может занимать до минуты времени на моих целевых моделях, но потенциально может и больше. Мне нужно научить Svod сериализовать граф и скомпилированные ядра в бинарные представления и мочь десериализовать их для мгновенного запуска.
Также, помимо ускорения холодного старта, AOT компиляция позволит запускать Svod в местах без доступа к компилятору, например в WASM окружении на пользовательском устройстве.
Библиотека примитивов для анализа данных
За последние два года появилось большое количество статей и референсных реализаций FA-like GPU алгоритмов для анализа данных, таких как kmeans, knn, pca, svd, dbscan, hdbscan, umap, t-sne и другие. Эта инфраструктура является разрозненной (реализована на разных фреймворках и DSL) и ориентируется на экосистему NVIDIA. Я хочу написать реализацию на Svod TK, которая поддерживала бы все заявленные мной аппаратные платформы и стала бы для них SOTA реализацией.
Формальная верификация
Сейчас мы можем генерировать С-код для целевой платформы, но ничего не мешает нам генерировать верифицируемый за счёт аннотаций код, который проверяет выход за границы массива, преобразования с потерями и т.д..
Лучшие средства интроспекции
Сейчас, копируя reference-based подхода из Python, которое должно было удешевить сравнения поддеревьев, которые в большом количестве происходят в ходе работы оптимизатора, мы потеряли возможность связывать строки Rust-кода и ONNX-узлы со строками сгенерированного ядра. Я хочу переписать сравнение на hash-based.