JetSpec ускоряет инференс LLM до 9.64× с помощью параллельного tree‑drafting
Исследователи представили метод JetSpec, который применяет speculative decoding вместе с параллельным tree‑drafting, чтобы ускорить генерацию в больших языковых моделях. Подход одновременно оптимизирует стоимость и качество этапа drafting, за счет чего заметно снижается задержка генерации. Как сообщает r/LocalLLaMA (Reddit), система показывает до 9.64× ускорения end‑to‑end на бенчмарке MATH-500 и 4.58× в задачах open-ended chat, при этом результаты остаются lossless.
В работе также используются дополнительные оптимизации CUDA graph и kernel. С их помощью производительность удается довести примерно до 1000 TPS на одной GPU B200.
Авторы объясняют, что традиционные методы speculative decoding обычно упираются в компромисс. AR-style draft heads сохраняют причинность и качество, однако стоимость drafting растет по мере увеличения глубины дерева. Block-diffusion style heads, напротив, формируют черновик дешево за один проход, но ветви часто оцениваются независимо, поэтому более глубокие пути могут расходиться и становиться несогласованными.
JetSpec обходит эту проблему иначе. Система строит дерево, которое сохраняет причинность, и делает это за один проход. Такой механизм drafting поддерживает согласованность последовательности и одновременно уменьшает вычислительные затраты, что в итоге и дает заметное ускорение инференса.
Ключевые факты
JetSpec демонстрирует до 9.64× ускорение end-to-end инференса на бенчмарке MATH-500 при сохранении lossless‑качества.
В задачах open-ended chat система показывает до 4.58× ускорение генерации без потери качества.
С использованием CUDA graph и оптимизаций kernel JetSpec достигает около 1000 TPS на одном GPU B200.
Метод использует causality-preserving parallel tree drafting, формируя причинно-согласованное дерево кандидатов за один проход.