Ускорение инференса guard-моделей: сравнение TensorRT, Triton, vLLM и Ray Serve
В системах с LLM guard-модели проверяют вход и выход на опасный контент, но их инференс часто становится узким местом. Автор протестировал пять инструментов ускорения и выяснил, что универсального решения нет, а главные проблемы — динамическая форма входа и нестандартный attention.
Инженеры, работающие над LLM-приложениями, часто забывают, что между пользователем и моделью может стоять дополнительный узел, который сам является моделью. В новой публикации разработчица рассказывает, как ускоряла guard-модель — энкодер, проверяющий текст на наличие персональных данных и опасного контента. Эта модель, хоть и весит чуть более 0,5 ГБ, вызывается дважды за каждый ход диалога: сначала проверяется запрос пользователя, затем ответ модели. Именно поэтому её скорость критична для UX, и автор решила сравнить пять инструментов: TensorRT, NVIDIA Triton, vLLM, Ray Serve, а также отдельно рассмотрела переход на Flash DeBERTa. Все эксперименты воспроизводимы, код и конфиги выложены в открытый репозиторий.
Главная сложность заключается в архитектуре современной guard-модели. Она работает zero-shot: типы сущностей (имя, адрес, телефон и т.д.) подаются в модель текстом вместе с анализируемым документом, и модель за один проход находит все сущности и выносит вердикт. Это делает модель гибкой — не нужно переобучать её под каждого клиента, но создаёт серьёзные проблемы для инференса. Во-первых, форма входа меняется от запроса к запросу, что плохо для TensorRT, который любит фиксированные размеры. Во-вторых, на выходе модели не тензор, а результат span-декодинга — отдельной питоновской логики, которая выполняется на CPU и не оптимизируется GPU-инструментами. В-третьих, в графе есть специфические операции вроде gather по парам индексов и билинейного скоринга, которые сложно экспортировать в ONNX и почти невозможно квантизовать в INT8. Наконец, бэкбон может использовать модифицированный attention, как в DeBERTa, где позиция и содержание токена обрабатываются раздельно, что тоже затрудняет оптимизацию.
Автор подчёркивает, что стоимость батча в энкодерах определяется самым длинным элементом, в отличие от декодеров, где работает continuous batching. В энкодере нельзя выйти из батча посреди прохода, поэтому один длинный документ заставляет ждать все остальные. Это привело к тому, что увеличение батча до 64 ухудшило хвостовые задержки, не улучшив пропускную способность. Кроме того, вся современная инфраструктура сервинга заточена под авторегрессионные LLM: она управляет KV-кэшем, разделяет префилл и декодинг, а для энкодеров нужен другой подход. Автор отмечает, что многие serving-движки просто не рассчитаны на такие модели, и это стало одной из причин, почему некоторые инструменты показали плохие результаты.
Результаты экспериментов оказались неоднозначными. TensorRT, несмотря на свою репутацию, столкнулся с проблемами из-за динамической формы входа: пришлось задавать профили min/opt/max и использовать паддинг, что привело к дополнительному расходу памяти. CUDA Graphs, которые могли бы ускорить выполнение, работают только для фиксированных форм, поэтому не помогли. Triton показал себя лучше, но тоже не смог полностью решить проблему span-декодинга. vLLM, который популярен для LLM, оказался неэффективен для энкодеров, так как он оптимизирован для генерации токенов, а не для классификации. Ray Serve, будучи гибким фреймворком, позволил легко масштабировать сервис, но не дал значительного ускорения самого инференса. Отдельно автор протестировала переход на Flash DeBERTa — это изменило архитектуру attention, что позволило использовать fused-ядра, но потребовало дообучения модели и не всегда оправдано.
Для российского рынка эта статья особенно актуальна, так как многие компании внедряют LLM-решения и сталкиваются с необходимостью фильтровать контент. Guard-модели становятся стандартом де-факто для безопасности, но их производительность часто недооценивают. Отрасль пока не имеет готовых решений для ускорения таких моделей, и разработчикам приходится самим экспериментировать с инструментами. Автор отмечает, что самым большим узким местом оказалась питоновская логика span-декодинга, которую невозможно перенести на GPU, и предлагает переносить её в C++ или CUDA-ядра, но это требует значительных усилий. Также она советует тщательно подбирать батч-размер, так как он сильно влияет на хвостовые задержки.
Сравнивая альтернативы, автор приходит к выводу, что идеального инструмента нет. TensorRT хорош для фиксированных форм, но не для динамических. Triton — надёжный сервер, но он не решает проблему декодинга. vLLM вообще не подходит для энкодеров. Ray Serve удобен для оркестрации, но не ускоряет сам инференс. Flash DeBERTa — перспективное направление, но требует дополнительных экспериментов. В итоге автор рекомендует комбинировать инструменты: использовать TensorRT для GPU-части и оптимизировать CPU-логику отдельно. Она также отмечает, что в её случае удалось достичь ускорения в несколько раз, но это потребовало много ручной работы и глубокого понимания внутренностей каждого инструмента.
В будущем автор планирует продолжить исследования и, возможно, попробовать другие подходы, такие как ONNX Runtime с CUDA-исполнителем или TensorFlow Serving. Она также подчёркивает, что открытые вопросы остаются: как автоматизировать оптимизацию span-декодинга, можно ли использовать квантизацию для нестандартных операций, и как адаптировать continuous batching для энкодеров. Пока же её главный совет — не полагаться на один инструмент, а тестировать каждый на своей модели и нагрузке. Эта работа, вероятно, станет отправной точкой для многих инженеров, которые ищут способы ускорить свои guard-модели, и покажет, что даже небольшие модели могут быть сложными для оптимизации.