BASS: Пакетная оптимизация внимания при спекулятивном сэмплировании

BASS: Batched Attention-optimized Speculative Sampling

April 24, 2024

Авторы: Haifeng Qian, Sujan Kumar Gonugondla, Sungsoo Ha, Mingyue Shang, Sanjay Krishna Gouda, Ramesh Nallapati, Sudipta Sengupta, Xiaofei Ma, Anoop Deoras

cs.AI

Аннотация

Спекулятивная декодировка стала мощным методом улучшения задержки и пропускной способности при работе с большими языковыми моделями. Однако большинство существующих реализаций сосредоточены на генерации одной последовательности. Прикладные генеративные ИИ-приложения часто требуют нескольких ответов, и вопрос о том, как выполнять спекулятивную декодировку в пакетном режиме, сохраняя при этом преимущества в задержке, представляет собой нетривиальную задачу. В данной статье описывается система пакетной спекулятивной декодировки, которая устанавливает новый уровень в области задержки генерации нескольких последовательностей и демонстрирует более высокое использование GPU, а также качество генерации в рамках отведенного времени. Например, для модели размером 7,8 млрд на одном GPU A100 с размером пакета 8 каждая последовательность генерируется со средней скоростью 5,8 мс на токен, общая пропускная способность составляет 1,1 тыс. токенов в секунду. Эти результаты представляют собой современный уровень задержки и ускорение в 2,15 раза по сравнению с оптимизированным обычным декодированием. В рамках временного бюджета, в котором обычное декодирование не завершается, наша система способна генерировать последовательности с HumanEval Pass@First на уровне 43% и Pass@All на уровне 61%, что значительно превышает то, что возможно с помощью спекулятивной декодировки одной последовательности. Наш пиковый уровень использования GPU во время декодирования достигает 15,8%, более чем в 3 раза выше, чем у обычного декодирования, и около 10 раз выше, чем у спекулятивной декодировки одной последовательности.

English

Speculative decoding has emerged as a powerful method to improve latency and throughput in hosting large language models. However, most existing implementations focus on generating a single sequence. Real-world generative AI applications often require multiple responses and how to perform speculative decoding in a batched setting while preserving its latency benefits poses non-trivial challenges. This paper describes a system of batched speculative decoding that sets a new state of the art in multi-sequence generation latency and that demonstrates superior GPU utilization as well as quality of generations within a time budget. For example, for a 7.8B-size model on a single A100 GPU and with a batch size of 8, each sequence is generated at an average speed of 5.8ms per token, the overall throughput being 1.1K tokens per second. These results represent state-of-the-art latency and a 2.15X speed-up over optimized regular decoding. Within a time budget that regular decoding does not finish, our system is able to generate sequences with HumanEval Pass@First of 43% and Pass@All of 61%, far exceeding what's feasible with single-sequence speculative decoding. Our peak GPU utilization during decoding reaches as high as 15.8%, more than 3X the highest of that of regular decoding and around 10X of single-sequence speculative decoding.