Перенос KV-кэша между моделями Qwen3 вместо повторного префилла

Представьте кодового агента, который работает уже двадцать минут. Он прочитал репозиторий, открыл документацию, вызвал несколько инструментов и накопил около 30K токенов контекста на Qwen3-14B. И тут он доходит до шага, где более сильная модель действительно бы пригодилась.

Отправить следующий ход в Qwen3-32B — простая часть. Сохранить уже проделанную работу — сложная. Модели 32B придётся заново пережевать весь накопленный контекст, прежде чем она выдаст хоть один токен. Этот проход называется префиллом, а его результат — KV-кэш модели.

Недавний исследовательский препринт NVIDIA задаёт острый вопрос. Маленькая модель уже построила KV-кэш для этого разговора, так нельзя ли преобразовать это состояние в кэш большой модели, вместо того чтобы заставлять её строить его заново из сырого текста? В описанном эксперименте с Qwen3 14B → 32B маппер работает 277,6 мс на 32 768 токенах против 6975,3 мс на повторный префилл целевой моделью. Это примерно 25x для одного этапа. Результат многообещающий, но гораздо более узкий, чем «инференс LLM в 25 раз быстрее».

Почему смена модели выбрасывает работу

KV-кэш — это состояние внимания, которое авторегрессионный трансформер строит, пока читает промпт. Когда оно есть, модели не нужно заново считать внимание по всему префиксу для каждого генерируемого токена.

Внутри одной модели такое переиспользование работает естественно. Системы инференса уже кэшируют префиксы и перекидывают KV-состояние между совместимыми воркерами; например, NVIDIA TensorRT-LLM поддерживает переиспользование KV-кэша и обмен KV для дезагрегированного сервинга.

Ломается всё при смене модели.

Qwen3-14B
   │
   │  30K-token conversation
   ▼
14B KV cache
   │
   ✕ cannot simply hand this to 32B
   │
Qwen3-32B normally re-prefills the text

У моделей разного размера разное число слоёв и разные внутренние представления. Общий токенизатор или одно семейство не делает их кэши взаимозаменяемыми. NVIDIA рассматривает это как задачу отображения представлений.

14B KV cache
     ↓
model-pair mapper
     ↓
32B-compatible KV cache
     ↓
32B starts decoding

Исследование сознательно ограничено парами внутри одного семейства с совпадающими KV: у источника и цели должны совпадать число KV-голов и размерность головы. Проверенные пары взяты из Qwen3, Llama 3.1 и Ministral 3. Перенос между семействами, скажем из Qwen в Llama, остаётся открытым вопросом.

Линейного моста, как ни странно, хватает

Интересно не то, что маппер вообще можно обучить. Интересно, что довольно простой маппер в замкнутой форме хорошо работает на нескольких парах моделей.

Анализ линейной структуры в статье показывает, что для Qwen3 14B → 32B один лучший слой источника объясняет около 56% дисперсии ключей цели и 32% дисперсии значений. Если объединить набор слоёв источника, эти цифры вырастают примерно до 79% и 65%.

Рабочий маппер опирается на три идеи.

Выбрать полезные слои источника. Слой цели не просто копирует слой источника на той же глубине. Метод находит слои источника, чьи KV-состояния лучше всего предсказывают данный слой цели, и конкатенирует их признаки.

Подгонять ключи и значения отдельно для каждой головы цели. Сам мост — гребневая регрессия, которая решается независимо для каждого слоя и головы цели: одна подгонка для ключей, другая для значений. Упрощённая версия математики выглядит так:

import numpy as np

def fit_ridge(x, y, lam=0.01):
    """Fit Y ≈ XW + b. Conceptual demo, not NVIDIA's implementation."""
    x_mean = x.mean(axis=0, keepdims=True)
    y_mean = y.mean(axis=0, keepdims=True)
    xc = x - x_mean
    yc = y - y_mean
    eye = np.eye(xc.shape[1], dtype=xc.dtype)
    w = np.linalg.solve(xc.T @ xc + lam * eye, xc.T @ yc)
    b = (y_mean - x_mean @ w).squeeze(0)
    return w, b

def map_cache(x, w, b):
    return x @ w + b

В настоящем маппере x гораздо больше, поскольку он конкатенирует KV-признаки всех выбранных слоёв источника, и у каждой головы цели своя подгонка. Для калибровки используется 500 последовательностей FineWeb-Edu по 1024 токена, которые прореживаются примерно до 128K наблюдений на уровне токенов на каждую голову цели.

Убрать позицию перед отображением ключей. Эту деталь стоит запомнить больше всего. RoPE поворачивает ключи в зависимости от позиции токена, поэтому если отображать повёрнутые ключи напрямую, есть риск выучить распределение позиций в калибровочных данных из 1024 токенов, а не содержание. Метод сначала отменяет поворот RoPE модели-источника, отображает представление содержания без позиции и затем применяет RoPE целевой модели.

source key
    ↓
strip source RoPE
    ↓
map content representation
    ↓
apply target RoPE
    ↓
target key

Благодаря этому маппер продолжает работать на позициях, которых не видел при калибровке, включая тест задержки на 32K.

Что на самом деле покрывает главная цифра

Вот задержки для Qwen3 14B → 32B на узле 8xH100:

• 64 токена: 14,0 мс на маппинг против 61,7 мс на повторный префилл, примерно в 4 раза быстрее.
• 8K токенов: 67,8 мс против 1154,8 мс, примерно в 17 раз быстрее.
• 32K токенов: 277,6 мс против 6975,3 мс, примерно в 25 раз быстрее.

Читайте эти цифры внимательно. Они сравнивают время работы маппера со временем, которое принимающей модели нужно на повторный префилл. Это не ускорение генерации в целом в 25 раз, потому что декодирование всё равно идёт после. В статье есть и оговорка о развёртывании: в реальном многопроцессном сервинге отображённый кэш, возможно, придётся пересылать в целевой процесс, а бенчмарк маппера не превращает это в цифру с учётом сети.

Качество на одних парах хорошее, на других плохое. Qwen3 14B → 32B сохраняет 97,6% средней самостоятельной точности целевой 32B по пяти бенчмаркам. Qwen3 8B → 32B сохраняет 87,5%, Llama 3.1 8B → 70B — 72,8%, Ministral 3B → 8B — 76,2%. А вот Ministral 3B → 14B и 8B → 14B с линейным маппером падают до средних 44,2% и 41,6%. В этом исследовании совпадение KV-размерностей — необходимое условие, а не обещание чистой передачи.

Где это помогает агентам

Долгоживущие агенты — как раз то место, где повторный префилл бьёт больнее всего.

Вот как могла бы выглядеть политика маршрутизации на основе этого метода:

"""Illustrative serving API. The paper does not provide this interface."""
state = agent.run(model="qwen3-14b", task=task)

if state.needs_stronger_reasoning:
    mapped_kv = kv_bridge.convert(
        source_model="qwen3-14b",
        target_model="qwen3-32b",
        kv_cache=state.kv_cache,
    )
    state = agent.resume(
        model="qwen3-32b",
        kv_cache=mapped_kv,
    )

Маршрутизация на большую модель — не новость, это работает уже сейчас. Новое в том, что вместе с разговором может переезжать и состояние вычислений. Исследовательский агент мог бы дёшево исследовать на маленькой модели семейства, передать работу большой на трудном участке рассуждений, а потом вернуться вниз для рутины. Отображения направленные, так что для 14B → 32B и 32B → 14B нужны отдельные мапперы, и статья обосновывает как повышение качества, так и снижение затрат.

В многоходовом эксперименте Qwen3 14B и 32B чередовались в 100 диалогах CoQA. К десятому ходу разрыв при переходе от малой к большой модели вырос примерно на 1,7 процентного пункта, а дрейф при переходе от большой к малой рос примерно на 0,33 пункта за ход. Это неплохое свидетельство в пользу переключения в разговоре, хотя сами исследователи предупреждают, что на гораздо более длинных сессиях линейный дрейф может накапливаться.

Ограничения, важные в продакшене

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

• Рамки совместимости узкие. Все описанные эксперименты остаются внутри семейств, используют плотные модели с полным вниманием и охватывают шесть пар с одинаковым числом KV-голов и размерностью головы. Между семействами или на несовпадающих KV-формах метод не проверяли.
• Каждую пару нужно проверять отдельно. Два маппинга Ministral с гребневой регрессией проваливаются, хотя их KV-архитектура совпадает. Нелинейный MLP возвращает удержание на HellaSwag примерно на 24–37 процентных пунктов на этих неудачных парах, что обнадёживает, но показывает, что простая линейная история работает не везде.
• Маппер не маленький. Оценённые мапперы содержат примерно от 1,01B до 3,36B параметров и занимают 4–12 ГБ на направление. Qwen3 14B → 32B — около 1,07B параметров и 4 ГБ. Подгонка требует примерно 47–87 минут на пару на одном узле 8xH100.

Поскольку отображения направленные, маршрутизатор на много моделей накапливает состояние для каждой пары. При среднем размере маппера из статьи это 39 ГБ для трёх моделей, 79 ГБ для четырёх и 131 ГБ для пяти. Мапперы можно держать в памяти хоста или в хранилище, а не в VRAM, но это всё равно инфраструктура, которую нужно обслуживать.

И ещё: официальной реализации к статье пока нет. В записи на arXiv есть статья и её исходники, но репозитория NVIDIA для этого маппера нет, так что фрагменты выше — пояснения, а не готовый код. Контекст со стороны сервинга есть в документации TensorRT-LLM по переиспользованию KV-кэша и в записи статьи на arXiv.

Что стоит запомнить

KV-кэш обычно считали состоянием, принадлежащим одной модели. Если перенос между моделями станет достаточно надёжным, это допущение уйдёт.

Для инфраструктуры агентов интересное будущее — не просто маршрутизатор, выбирающий модель под каждый запрос. Это маршрутизатор, который переносит наработанное состояние разговора между совместимыми моделями, не оплачивая полный контекст второй раз.

Ничто из этого не делает KV-кэш переносимым в общем случае. Стоимость декодирования не исчезает, маршрутизация между семействами не решена, а самые слабые пары показывают, почему важно проверять каждую пару. Но результат Qwen3 14B → 32B показывает, что работу, сохранённую в состоянии внимания одной модели, другая модель может переиспользовать дешевле, чем пересчитать её из текста.

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