Разработка собственного рантайма для инференса больших языковых моделей (LLM) — это не просто упражнение в программировании, а способ полностью контролировать стек вычислений. Автор проекта annotated-llm-runtime Anubhab Banerjee опубликовал детальное руководство по созданию кастомного C++/CUDA движка для модели Qwen2.5-Coder-7B-Instruct на архитектуре NVIDIA Hopper (H100, sm_90). Главная цель — не побить существующие решения, а понять, как работает «bare-metal AI» изнутри.
Ключевые метрики производительности
Самым важным открытием стало влияние CUDA Graphs на производительность. Изначальный «жадный» (eager) декодинг работал крайне медленно — около 119 мс на токен. Обернув стабильный этап декодирования в захваченный CUDA Graph, автор снизил задержку до ~17 мс на токен. Это дает прирост в 7 раз без изменения самих ядер вычислений, просто за счет уменьшения накладных расходов драйвера.
Сравнение с эталонной реализацией llama.cpp (коммит b3040, квантование Q4_K_M) показывает текущее состояние кастомного рантайма:
| Метрика | Кастомный Runtime (FP16/INT4) | llama.cpp (Q4_K_M) |
|---|---|---|
| TTFT (512 токенов) | 128 мс | 43 мс |
| Decode ITL (steady-state) | 16.7 мс/токен (~60 tok/s) | 4.95 мс/токен (~200 tok/s) |
| Первый токен (FTL) | ~128 мс | Информация недостаточна |
Архитектурные решения и оптимизации
Рантайм использует специфические параметры модели Qwen2.5-Coder-7B: 28 слоев, скрытый размер 3584, 28 заголовков внимания (GQA с коэффициентом 7:1). Квантование весов выполнено в формате symmetric group-wise INT4 (группа 128), при этом операции RMSNorm и lm_head остаются в FP16 для минимизации ошибок численной стабильности.
Автор выделяет три критические ошибки, которые пришлось исправлять:
- __syncthreads() в ветвлениях: Использование синхронизации внутри warp-specialized ветвей является неопределенным поведением (UB) и приводило к коррупции данных в хвостах страниц KV.
- Prmt.b32 vs __dp4a: На архитектуре Hopper инструкция
prmt.b32для перестановки целых чисел оказалась эффективнее__dp4aдля форм GEMV в гейтах (gate/up/down), что позволило отказаться от пути активации INT8. - Выравнивание памяти: Страницы KV (16 токенов) упакованы так, чтобы идеально совпадать с 128-байтными линиями кэша L2 Hopper (4 KiB на срез головы).
Почему это важно?
Хотя llama.cpp остается золотым стандартом для продакшена, этот проект демонстрирует, как можно «открыть черный ящик». Понимание того, как упаковывать веса, управлять страницами KV и захватывать графы вычислений, необходимо для внедрения кастомных квантований, новых вариантов внимания или поддержки новых архитектур GPU, где готовые решения еще не оптимизированы.
Источник: Towards Data Science ↗