Перенесення 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 мс на повторний префіл, приблизно вчетверо швидше.
• 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 показує, що роботу, збережену в стані уваги однієї моделі, інша модель може перевикористати дешевше, ніж перерахувати її з тексту.

Якщо це перенесеться в справжні системи інференсу, каскади між моделями стануть набагато цікавішими для довгоживучих агентів. Питання буде вже не лише в тому, яка модель відповість на наступний хід, а й у тому, скільки обчислень попередньої моделі можна взяти з собою.