KV Cache

KV Cache (Key-Value Cache) is a critical optimization technique used primarily in Large Language Models (LLMs) and other Transformer-based architectures to accelerate inference latency and reduce computational costs. At its core, the KV cache stores the Key and Value matrices generated by the attention mechanism for previous tokens in a sequence. By saving these intermediate calculations, the model avoids recomputing the attention states for the entire history of the conversation every time it generates a new token. This process transforms the text generation workflow from a quadratic complexity operation into a linear one, making real-time interactions with chatbots and AI agents feasible.

Механизм и преимущества#

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

• Повышение скорости: извлекая предварительно вычисленные ключи и значения из памяти, система значительно ускоряет работу inference engine. Это необходимо для приложений, требующих низкой задержки, таких как real-time inference в ботах поддержки.
• Эффективность использования ресурсов: хотя кэш увеличивает объем используемой памяти (VRAM), он значительно снижает объем вычислений (FLOPs), необходимых для одного токена. Этот компромисс часто решается с помощью таких методов, как model quantization или страничная организация памяти, аналогично тому, как операционные системы управляют RAM.
• Расширенный контекст: эффективное управление KV-кэшем позволяет моделям работать с большим context window, давая им возможность обрабатывать длинные документы или поддерживать связные беседы в течение долгого времени.

Реальные приложения#

KV-кэш — это фундаментальный компонент при развертывании современного генеративного ИИ, однако его принципы также применимы и к computer vision (CV).

1. Генеративные чат-боты: такие сервисы, как ChatGPT или Claude, активно используют кэширование KV. Когда ты задаешь уточняющий вопрос, модель не перечитывает всю историю чата с нуля. Вместо этого она добавляет новый ввод к закэшированным состояниям предыдущего шага, обеспечивая практически мгновенные ответы.

2. Понимание видео: в задачах video understanding модели обрабатывают кадры последовательно. Подобно текстовым токенам, визуальные признаки из прошлых кадров могут кэшироваться, помогая модели отслеживать объекты или распознавать действия без повторной обработки всей истории видео. Это особенно актуально для action recognition, где важен временной контекст.

Эффективное управление памятью#

По мере роста размеров моделей объем KV cache может стать узким местом, потребляя гигабайты памяти GPU. Последние достижения сосредоточены на оптимизации этого хранилища.

• PagedAttention: вдохновленный виртуальной памятью в операционных системах, механизм PagedAttention (представленный в vLLM) позволяет хранить KV-кэш в несмежных блоках памяти. Это снижает фрагментацию и позволяет использовать большие размеры батчей при model serving.
• Квантование KV-кэша: для экономии места разработчики часто применяют mixed precision или квантование int8 специально для закэшированных значений. Это снижает требования к памяти, позволяя edge AI устройствам с ограниченным объемом RAM запускать мощные модели.
• Кэширование промптов: смежный метод, при котором состояния KV для статического системного промпта (например, "You are a helpful coding assistant") вычисляются один раз и переиспользуются в разных пользовательских сеансах. Это ключевая функция для оптимизации рабочих процессов prompt engineering в больших масштабах.

Разграничение похожих концепций#

Полезно различать KV Cache и другие термины кэширования и оптимизации:

• KV Cache vs. Prompt Caching: KV-кэш обычно относится к динамической памяти токен за токеном, используемой во время одного потока генерации. Кэширование промптов конкретно означает хранение обработанного состояния фиксированной входной инструкции для повторного использования в нескольких независимых вызовах генерации.
• KV Cache vs. Embeddings: Эмбеддинги — это векторные представления входных данных (текста или изображений), которые отражают семантический смысл. KV-кэш сохраняет активации (ключи и значения), полученные из этих эмбеддингов в слоях внимания, конкретно для задачи генерации последовательности.
• KV Cache vs. Model Weights: Веса модели — это статические изученные параметры нейронной сети. KV-кэш состоит из динамических временных данных, создаваемых во время прямого прохода для конкретной входной последовательности.

Пример: Контекст в моделях зрения#

Хотя кэширование KV наиболее известтно в NLP, концепция сохранения состояния применима и к продвинутым моделям компьютерного зрения. В примере ниже мы имитируем передачу состояния (контекста) в сценарии трекинга видео с помощью Ultralytics YOLO26. Здесь трекер сохраняет идентичность объектов между кадрами, что концептуально похоже на то, как кэш сохраняет контекст между токенами.

from ultralytics import YOLO

# Load the Ultralytics YOLO26 model
model = YOLO("yolo26n.pt")

# Track objects in a video, maintaining identity state across frames
# The 'track' mode effectively caches object features to link detections
results = model.track(source="https://ultralytics.com/images/bus.jpg", show=False)

# Print the ID of the tracked objects
if results[0].boxes.id is not None:
    print(f"Tracked IDs: {results[0].boxes.id.numpy()}")

Разработчики, желающие управлять датасетами и развертывать оптимизированные модели, могут использовать Ultralytics Platform, которая упрощает весь пайплайн от аннотации данных до эффективного model deployment. Для тех, кто интересуется глубокими механизмами внимания, такие библиотеки, как PyTorch, предоставляют базовые блоки, на которых реализованы эти механизмы кэширования.