Tokasaurus — это движок инференса для языковых моделей в режиме высоконагруженных задач. Он максимизирует пропускную способность при работе с LLM, предлагает поддержку API OpenAI, эффективно управляет памятью и оптимизирует вычисления в сценариях, где важно одновременно обрабатывать множество запросов без задержек.
Архитектура Tokasaurus разделена на 3 компонента: веб-сервер, менеджер и модельные воркеры.
🟢Веб-сервер отвечает за взаимодействие с клиентами, принимая запросы и отправляя ответы.
🟢Менеджер, запущенный в отдельном процессе, управляет планированием задач, KV-кешем и группировкой последовательностей с общими префиксами.
🟢Модельные воркеры выполняют прямые запросы к подключенным LLM. Компоненты обмениваются данными асинхронно через очереди, и это позволяет держать GPU загруженным без простоев.
Проект учитывает растущую потребность в масштабировании и предлагает 3 типа параллелизма: дата-параллелизм (
dp_size), пайплайн (pp_size) и тензорный (tp_size) с поддержкой AsyncTP.Async Tensor Parallelism в PyTorch — это техника ускорения распределенных вычислений для LLM, где операции связи (all-gather/reduce-scatter) разбиваются на асинхронные части и перекрываются с матричными умножениями (matmul) с помощью чередующихся CUDA-потоков: пока один поток вычисляет фрагмент matmul, другой параллельно передаtт данные для следующего фрагмента через P2P-копирование (NVLink + copy engines), минимизируя простои GPU.
При использовании нескольких GPU, например,
dp_size=2 и pp_size=4, система задействует 8 GPU, создавая 2 дублирующиеся группы по 4 GPU каждая. При этом параметры управления памятью (kv_cache_size_num_tokens, max_seqs_per_forward) применяются к каждой дата-параллельной группе отдельно. Это позволяет тонко управлять ресурсами, исходя из контекста конкретных нагрузок.Tokasaurus поддерживает модели семейств Llama3 и Qwen2, использует технологию Hydragen для ускорения внимания над общими префиксами последовательностей.
⚠️ Проект пока молодой, поэтому некоторые функции могут быть нестабильными. Разработчики активно работают над улучшениями.
📌 Лицензирование: Apache 2.0 License.
🖥GitHub
@ai_machinelearning_big_data
#AI #ML #LLM # #Tokasaurus #Github