IceCache: Memory-efficient KV-cache Management for Long-Sequence LLMs
IceCache: Memory-efficient KV-cache Management for Long-Sequence LLMs IceCache: Memory-Efficient KV-cache Management for Long-Sequence LLMs Статья развивает область retrieval-based KV-cache-методов. Основная их идея такава: 1) на префилле: сгружаем KV-кеш на CPU; за счёт этого экономим GPU-память; 2) на декодинге: загружаем обратно на GPU только подмножество наиболее релевантных KV, считаем аттеншен по части токенов; за счёт этого эффективно снижаем seqlen, экономим compute. Существующие методы загружают кеш последовательными чанками. Проблема такого подхода в том, что релевантные токены разбросаны по памяти, и при загрузке их чанков на GPU переносится много лишних токенов. IceCache решает это через изменение layout’а KV-cache: страница в памяти определённому токену назначается не по его позиции, а на основе косинусной близости к опорным токенам. Таким образом, страницы становятся «семантическими», релевантные токены лежат более компактно, занимают меньше чанков, благодаря чему можно уменьшить трансфер CPU —> GPU, а также сократить эффективный seqlen при вычислении аттеншена. Реализация В основе метода — кастомная структура данных DCI-tree для задачи приближенного поиска k-ближайших соседей. С её помощью токены распределяются по страницам памяти, а также на шаге декода выбираются наиболее релевантные токены. На префилле параллельно с вычислением аттеншена происходит оффлоадинг KV-кеша на CPU, а затем индексирование ключей на CPU при помощи DCI-tree. На декодинге для данного query при помощи DCI-tree определяются страницы с релевантными токенами. Те из них, которые не использовались на предыдущем шаге декодинга, дозагружаются из CPU на GPU. Первые токены в последовательности, attention-синки, всё время находятся в GPU-памяти. То же самое происходит с «хвостом» декодируемых токенов. Как только набирается окно из N декодированных токенов, они асинхронно сгружаются на CPU и индексируются DCI-tree. Авторы реализовали DCI-tree на C, а также написали CUDA-кернел, эффективно копирующий страницы из CPU-памяти в нужные страницы PagedAttention. Код доступен на GitHub. Метрики Avg. accuracy: 1) до ~99% качества полного KV-кеша при бюджете 256 токенов; 2) при бюджете 64 — сопоставим или лучше бейзлайнов с x4 большим кешом. Latency на llama3.1-8B с 36k seqlen: 1) time to second token: 5,9с., на уровне OmniKV; 2) time per output token: 0,11с., против 0,05с. у OmniKV. Потенциальные проблемы Наиболее слабое место метода — необходимость синхронизации GPU и CPU перед вызовом аттеншена на каждом шаге декодинга. Авторы прямо указывают, что половину времени декодинга занимает поиск по DCI-tree, исполняемый на CPU. Вероятнее всего, для практического применения метода нужно будет заменить DCI-tree на структуру данных, в которой алгоритм поиска соседей адаптирован под GPU. При этом обновление дерева по-прежнему может асинхронно выполняться на CPU. Разбор подготовил ❣ Михаил Катунькин #YaICLR26 Душный NLP Summary материалов: - Душный NLP: Канал «Душный NLP» публикует разборы свежих статей от NLP-специалистов Яндекса. Содержит подробные и полезные материалы. Имеет более 6 411 подписчиков. Источник: оригинальный пост