TokenRouter ускорил потокенную маршрутизацию между моделями до 64 раз

Команда из лаборатории thu-nics (Цинхуа) выпустила TokenRouter, систему, которая позволяет большой и малой языковым моделям делить между собой одну генерацию: каждый следующий токен выбирает, кто его напишет. Заголовочная цифра «до 64 раз» означает отношение пропускной способности TokenRouter (токенов в секунду на всех параллельных запросах) к официальным реализациям существующих алгоритмов роутинга. Выигрыш получился не из новой математики выбора модели, а из инженерии: авторы убрали накладные расходы, которые в сырых реализациях съедают почти всё время.

Роутинг запросов — зрелая тема, роутинг токенов — нет

Автор исходного поста в @gonzo_ML замечает, что направлять запросы к разным моделям умеют давно. Его компания Intento делает это почти десять лет. Поэтому его удивила оценка OpenRouter и сделка со Stripe: по данным Menlo Ventures, платёжный гигант покупает сервис-агрегатор моделей.

Такой роутер смотрит на запрос целиком. Простой вопрос уходит дешёвой модели, сложный дорогой. Примеры: коммерческие сервисы и открытый RouteLLM.

TokenRouter работает на другом уровне. Внутри одного ответа сложность неравномерна. Большинство токенов тривиальны: окончания слов, связки, шаблонный код. Лишь на отдельных развилках рассуждения нужна сильная модель. Критерием сложности служит энтропия: чем менее уверена малая модель в следующем токене, тем выше энтропия. Когда она высока, токен отдают большой модели. В остальных случаях пишет малая.

В теории это снижает стоимость инференса: тяжёлая модель работает только там, где без неё не обойтись. На практике схема упирается в обвязку.

Где пропадают токены в секунду

Наивная реализация на каждом шаге вызывает то одну модель, то другую, пересчитывает контекст и гоняет данные между процессами. Все эти операции накладны. Автор поста называет это самым интересным в работе: сопутствующая машинерия без оптимизации обходится очень дорого.

TokenRouter построен поверх SGLang, открытого движка для запуска LLM, и состоит из нескольких частей:

• API из трёх функций. Алгоритм роутинга описывается через небольшой интерфейс, поэтому разные роутеры подключаются без переписывания рантайма.
• Асинхронные подсерверы. У каждой модели свой подсервер, и они работают параллельно, не дожидаясь друг друга.
• Сохранение KV-кэша между вызовами. KV-кэш хранит уже вычисленные промежуточные состояния контекста. Если его не сохранять, при каждом переключении модели контекст пришлось бы обсчитывать заново.
• Планировщик отложенного батчинга на цепях Маркова. Видеокарта эффективнее обрабатывает пачку запросов, чем по одному. Планировщик моделирует поток обращений к каждому подсерверу и решает, стоит ли немного подождать, пока соберётся пачка.

Для R2R, одного из существующих токен-уровневых роутеров, отдельно измерили вклад CUDA-графов. Это технология NVIDIA: цепочку GPU-операций записывают один раз и запускают целиком, не тратя время на отдельные вызовы. Из общего ускорения в 2,76x на CUDA-графы пришлось 1,71x.

Цифры на олимпиадной математике и агентном коде

Систему проверили на пяти алгоритмах роутинга. Задачи двух типов: рассуждения на AIME2024 (задачи американской математической олимпиады) и агентные сценарии на SWE-Smith (набор задач по правке кода для программных агентов).

При 4 параллельных запросах:

• TokenRouter: 76,6–290,3 токен/с в зависимости от алгоритма.
• Потокенный бейзлайн на SGLang: 4,5–63,1 токен/с.
• Официальные реализации: 0,7–78,3 токен/с.
• Прирост: от 2,01x до 64,15x.

Отдельно на роутере R2R:

• 8 запросов: 134,9 → 372,5 токен/с (2,76x).
• 16 запросов: TokenRouter обгоняет R2R в 3,61x при той же параллельности.
• Против R2R с одним запросом выигрыш 18,58x, при этом скорость на одного пользователя вырастает в 1,13x.

Последняя пара цифр показывает практический смысл. Систему можно нагрузить множеством пользователей и не потерять в скорости каждого из них.

Что это даёт на практике и где подвох

Авторы показывают, что потокенный роутинг можно спрятать за единым OpenAI-совместимым эндпоинтом. Для клиента это обычный API, а внутри генерацию делят две модели. По соотношению скорости и точности такая схема может обходить роутеры уровня запросов вроде RouteLLM. Для компаний, которые платят за инференс, это новый рычаг: цена токена в доле «лёгких» токенов падает, а качество на сложных развилках сохраняется.

Есть два ограничения, которые инженеру стоит учесть до внедрения:

• Каждый подсервер вынужден хранить KV-кэш всей последовательности. Память видеокарт расходуется почти вдвое: кэш существует и для малой, и для большой модели.
• Марковский солвер планировщика проверен только до 4 параллельных запросов. Как он поведёт себя на сотнях сессий в проде, пока неизвестно. Цифры для 8 и 16 запросов получены на R2R.

Сам код выложен в открытом доступе (thu-nics/TokenRouter). Команды, которые обслуживают собственные модели, могут проверить подход на своих нагрузках. Главный урок работы прикладной: идея роутинга токенов давно известна, но без аккуратной системной инженерии её выигрыш теряется в обвязке.