Куда девается память при инференсе: пейджинг KV cache и магия radix trie в Tailor

Почему ваш инференс-движок жрет память как не в себя

Вы запускаете модель. Контекст 4K токенов. VRAM заполняется за секунды. Каждый новый запрос - новая аллокация. Параллельные запросы? Забудьте. Память фрагментируется, производительность падает. Знакомая картина?

Классический KV cache - это монолит. Сплошной блок памяти под ключи и значения для каждого токена в контексте. Новый запрос - новый блок. Повторяющийся префикс (например, системный промпт) - все равно новый блок. Это расточительно. Глупо. И именно так работают большинство движков.

Два гениальных трюка, которые меняют правила игры

Репозиторий Tailor (не путать с портным) показывает, как можно сделать иначе. Результат: 1990 токенов в секунду на Llama 3.2 1B на RTX 4090. Цифра сама по себе ничего не говорит. Важно как.

Paged KV Cache: разбиваем монолит на страницы

Представьте оперативную память. Процессы не занимают сплошной блок. Они используют страницы. Операционная система управляет ими, выделяя и освобождая по мере необходимости. Почему бы не применить это к KV cache?

Paged KV cache делит пространство ключей и значений на фиксированные блоки - страницы. Например, по 256 токенов на страницу. Контекст из 1024 токенов займет 4 страницы. Новый запрос начинает заполнять свободные страницы. Когда контекст заканчивается (пользователь прерывает генерацию), страницы освобождаются и могут быть использованы другим запросом.

Но страницы - это только половина дела. Вторая половина - умное их использование.

Prefix Cache на Radix Trie: когда один промпт на всех

Сколько раз ваша система генерирует один и тот же системный промпт? "Ты - полезный ассистент..." Каждый запрос начинается с этих токенов. В классической схеме для каждого запроса эти токены занимают место в KV cache заново.

Prefix cache решает это радикально. Он кеширует вычисленные ключи и значения для общих префиксов запросов. Но не просто кеширует, а организует в radix trie (сжатое префиксное дерево).

Когда приходит новый запрос, движок проверяет: может быть, первые N токенов уже есть в кеше? Если да, он просто начинает генерацию с токена N+1, используя уже вычисленные ключи и значения. Это называется автоматическая инкрементальная декодизация.

Собираем движок по косточкам: от теории к коду

Хватит теории. Давайте посмотрим, как это реализовано в Tailor. Я разберу ключевые компоненты, а вы поймете, что ничего сверхсложного здесь нет. Только правильная архитектура.

1Структура данных для paged cache

Вот как выглядит базовая структура страницы в KV cache:

class KVCachePage:
    def __init__(self, page_id: int, capacity: int, dtype, device):
        self.page_id = page_id
        self.capacity = capacity  # токенов на страницу
        self.k_data = torch.zeros((capacity, hidden_size), dtype=dtype, device=device)
        self.v_data = torch.zeros((capacity, hidden_size), dtype=dtype, device=device)
        self.occupied = 0  # сколько токенов уже занято
        self.is_free = True
        self.next_page = None  # ссылка на следующую страницу

Просто, правда? Страница - это контейнер для K и V тензоров фиксированного размера. occupied показывает, сколько места использовано. is_free - свободна ли страница для нового запроса.

А вот как выглядит менеджер страниц (очень упрощенно):

class KVCacheManager:
    def __init__(self, total_pages: int, tokens_per_page: int):
        self.pages = [KVCachePage(i, tokens_per_page) for i in range(total_pages)]
        self.free_pages = list(range(total_pages))
        self.used_pages = {}  # page_id -> request_id
        
    def allocate_for_request(self, request_id: str, token_count: int):
        """Выделяем страницы под новый запрос"""
        pages_needed = (token_count + self.tokens_per_page - 1) // self.tokens_per_page
        if len(self.free_pages) < pages_needed:
            # Нет свободных страниц - нужно что-то освободить
            self._evict_oldest()
        
        allocated = []
        for _ in range(pages_needed):
            page_id = self.free_pages.pop()
            self.pages[page_id].is_free = False
            self.used_pages[page_id] = request_id
            allocated.append(page_id)
        
        return allocated  # возвращаем ID выделенных страниц

2Radix Trie для prefix cache

Теперь самое интересное - префиксный кеш. Реализация trie:

class RadixTrieNode:
    def __init__(self, token_id=None):
        self.token_id = token_id
        self.children = {}  # token_id -> RadixTrieNode
        self.kv_cache_ref = None  # ссылка на вычисленные K/V
        self.is_terminal = False  # конец последовательности
        
class PrefixCache:
    def __init__(self):
        self.root = RadixTrieNode()
        
    def insert_sequence(self, token_ids: List[int], kv_data):
        """Вставляем последовательность токенов и связанные K/V"""
        node = self.root
        for token in token_ids:
            if token not in node.children:
                node.children[token] = RadixTrieNode(token)
            node = node.children[token]
        node.kv_cache_ref = kv_data
        node.is_terminal = True
        
    def find_longest_prefix(self, token_ids: List[int]):
        """Ищем самый длинный префикс, который уже есть в кеше"""
        node = self.root
        matched = []
        
        for token in token_ids:
            if token in node.children:
                node = node.children[token]
                matched.append(token)
            else:
                break
        
        if matched and node.is_terminal:
            return matched, node.kv_cache_ref
        return [], None

Когда приходит запрос "Ты - полезный ассистент. Напиши код...", движок сначала проверяет префиксный кеш. Находит "Ты - полезный ассистент." уже вычисленным. Берет готовые K/V. Начинает генерацию сразу с "Напиши код...". Экономия - десятки матричных умножений.

3Интеграция в forward pass

Вот где магия становится реальностью. Модифицированный forward pass внимания:

def attention_forward_with_paged_cache(query, key, value, cache_manager, prefix_cache, token_ids):
    """Внимание с paged KV cache и prefix cache"""
    
    # 1. Проверяем prefix cache
    cached_prefix, cached_kv = prefix_cache.find_longest_prefix(token_ids)
    if cached_prefix:
        # Уже есть вычисленные K/V для префикса
        cached_k, cached_v = cached_kv
        start_from = len(cached_prefix)
        # Берем только новые токены для вычислений
        new_token_ids = token_ids[start_from:]
    else:
        cached_k, cached_v = None, None
        new_token_ids = token_ids
    
    # 2. Вычисляем K/V для новых токенов
    new_k = compute_key(new_token_ids)
    new_v = compute_value(new_token_ids)
    
    # 3. Объединяем с кешированными (если есть)
    if cached_k is not None:
        k = torch.cat([cached_k, new_k], dim=0)
        v = torch.cat([cached_v, new_v], dim=0)
    else:
        k, v = new_k, new_v
    
    # 4. Выделяем страницы в paged cache для новых K/V
    page_ids = cache_manager.allocate_for_request(request_id, len(new_token_ids))
    
    # 5. Распределяем новые K/V по страницам
    offset = 0
    for page_id in page_ids:
        page = cache_manager.pages[page_id]
        tokens_to_fill = min(page.capacity - page.occupied, len(new_token_ids) - offset)
        
        page.k_data[page.occupied:page.occupied + tokens_to_fill] = \
            new_k[offset:offset + tokens_to_fill]
        page.v_data[page.occupied:page.occupied + tokens_to_fill] = \
            new_v[offset:offset + tokens_to_fill]
        
        page.occupied += tokens_to_fill
        offset += tokens_to_fill
    
    # 6. Вычисляем внимание как обычно
    scores = torch.matmul(query, k.transpose(-2, -1))
    weights = F.softmax(scores, dim=-1)
    output = torch.matmul(weights, v)
    
    return output

Обратите внимание на шаг 5: мы не просто записываем K/V в сплошной тензор. Мы распределяем их по страницам. Каждая страница знает, сколько в ней занято места. Это позволяет эффективно переиспользовать память.

Подводные камни, о которых молчат в рекламе

Красиво на бумаге. В реальности - десятки нюансов, которые могут все испортить.

Фрагментация все равно происходит

Представьте: у вас 10 страниц по 256 токенов. Приходит запрос на 300 токенов. Ему нужно 2 страницы. Он занимает страницы 0 и 1. Потом приходит запрос на 200 токенов - занимает страницу 2. Первый запрос завершается, освобождает страницы 0 и 1. Теперь у вас свободные страницы 0, 1 и занятая 2. Приходит запрос на 500 токенов. Ему нужно 2 страницы. Страницы 0 и 1 свободны - отлично. Но они не последовательны в памяти? А кому какое дело! Это же отдельные тензоры.

Prefix cache не всегда выгоден

Системный промпт из 10 токенов - да, выгодно кешировать. А если префикс всего 2 токена? Вычисление занимает микросекунды. Поиск в trie тоже занимает время. Может оказаться, что поиск дольше вычисления. Tailor решает это пороговым значением: кешируем только если префикс длиннее N токенов (обычно 8-16).

Синхронизация - ад для параллельных запросов

Два запроса одновременно хотят добавить один и тот же префикс в кеш. Кто первый? Нужна блокировка. Но блокировка тормозит. Решение в Tailor: optimistic concurrency. Первый запрос вычисляет префикс и помещает в кеш. Второй проверяет - если уже есть, использует. Если нет (редкий случай race condition), вычисляет сам. Немного избыточных вычислений, зато нет блокировок.

Как Tailor добивается 1990 tok/s: неочевидные оптимизации

Цифра 1990 токенов в секунду для Llama 3.2 1B на RTX 4090 впечатляет. Но paged cache и prefix cache - не единственные причины.

• Квантование KV cache в int8. Ключи и значения не требуют FP16 точности. 8 бит достаточно. Это еще 50% экономии памяти. Об этом мы подробно говорили в статье про Binary KV cache.
• Пакетная обработка страниц. Вместо работы с каждой страницей отдельно, Tailor группирует операции. Все K тензоры со всех страниц обрабатываются одним ядром CUDA.
• Предиктивное выделение. Движок предсказывает, сколько страниц понадобится для типичного запроса, и держит их готовыми.
• Асинхронные операции с памятью. Пока GPU вычисляет, CPU готовит следующие страницы.

Но главное - отсутствие аллокаций во время инференса. В классическом движке каждый новый токен требует проверки: поместится ли в выделенную память? Если нет - переаллокация всего тензора с копированием. Это смерть производительности. В paged cache страницы выделены заранее. Новый токен просто записывается в следующую ячейку уже существующей страницы.

Стоит ли писать свой движок?

Если вам нужно:
1. Обрабатывать сотни параллельных запросов
2. Иметь гарантированное время ответа (no GC pauses)
3. Максимально использовать дорогую GPU память
4. Изучить внутренности LLM до винтика
Тогда да, стоит.

Если же вам просто нужна работающая модель для экспериментов - используйте vLLM, TGI или llama.cpp. Они уже реализуют похожие оптимизации.

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

Что будет дальше? Прогноз от того, кто видел эту кухню изнутри

Paged KV cache и prefix cache - не конечная точка. Уже видны следующие шаги:

1. Иерархический кеш: горячие префиксы в GPU памяти, теплые - в CPU RAM, холодные - на SSD. Как в современных СУБД.
2. Адаптивный размер страниц: не фиксированные 256 токенов, а динамические страницы под размер запроса.
3. Сжатие KV cache on-the-fly: редко используемые части контекста можно сжимать без потери качества.
4. Распределенный KV cache: для моделей с триллионом параметров кеш будет разбиваться между несколькими GPU.

Самое интересное: эти оптимизации сделают возможным то, что сегодня кажется фантастикой. Персональные LLM с контекстом в миллион токенов, работающие на смартфоне. Реальные диалоги с историей длиной в часы, а не в минуты. Многозадачные агенты, которые держат в памяти десятки параллельных процессов.

Но для этого нужно перестать думать о памяти как о безликом ресурсе. Нужно думать о ней как о структуре данных. И Tailor показывает, как это делать.

P.S. Если после этой статьи вам захотелось поковыряться в коде - отлично. Начните с форка Tailor. Замените фиксированные страницы на динамические. Добавьте сжатие. Экспериментируйте. Именно так появляются прорывные технологии - не чтением документации, а разбором на детали и сборкой заново.