KV Cache

KV Cache (Key-Value Cache) là một kỹ thuật tối ưu hóa quan trọng được sử dụng chủ yếu trong Large Language Models (LLMs) và các kiến trúc dựa trên Transformer khác nhằm tăng tốc inference latency và giảm chi phí tính toán. Về bản chất, bộ nhớ đệm KV lưu trữ các ma trận Key và Value được tạo bởi attention mechanism cho các token trước đó trong một chuỗi. Bằng cách lưu các phép tính trung gian này, mô hình tránh việc phải tính toán lại trạng thái chú ý cho toàn bộ lịch sử hội thoại mỗi khi tạo một token mới. Quá trình này biến quy trình text generation từ một hoạt động có độ phức tạp bậc hai thành tuyến tính, giúp các tương tác thời gian thực với chatbot và AI agents khả thi hơn.

Cơ chế và Lợi ích#

Trong một mô hình Transformer tiêu chuẩn, việc tạo từ tiếp theo đòi hỏi phải chú ý đến tất cả các từ trước đó để hiểu ngữ cảnh. Nếu không có bộ nhớ đệm, mô hình sẽ cần tính toán lại các mối quan hệ toán học cho toàn bộ chuỗi ở mỗi bước. KV cache giải quyết vấn đề này bằng cách hoạt động như một ngân hàng bộ nhớ.

• Cải thiện tốc độ: Bằng cách truy xuất các key và value đã tính toán trước từ bộ nhớ, hệ thống tăng tốc đáng kể inference engine. Điều này rất cần thiết cho các ứng dụng đòi hỏi độ trễ thấp, chẳng hạn như real-time inference trong bot chăm sóc khách hàng.
• Hiệu quả tài nguyên: Mặc dù làm tăng mức sử dụng bộ nhớ (VRAM), phương pháp này lại giảm đáng kể lượng tính toán (FLOPs) cần thiết cho mỗi token. Sự đánh đổi này thường được quản lý thông qua các kỹ thuật như model quantization hoặc phân trang, tương tự cách hệ điều hành quản lý RAM.
• Mở rộng ngữ cảnh: Quản lý hiệu quả KV cache cho phép các mô hình xử lý context window lớn hơn, giúp chúng xử lý các tài liệu dài hoặc duy trì các cuộc hội thoại mạch lạc trong thời gian dài.

Các ứng dụng trong thực tế#

KV cache là thành phần cơ bản trong việc triển khai AI tạo sinh hiện đại, nhưng các nguyên tắc của nó cũng mở rộng sang cả computer vision (CV).

1. Chatbot tạo sinh: Các dịch vụ như ChatGPT hoặc Claude phụ thuộc rất nhiều vào KV caching. Khi người dùng đặt câu hỏi tiếp theo, mô hình không đọc lại toàn bộ lịch sử trò chuyện từ đầu. Thay vào đó, nó nối thêm đầu vào mới vào các trạng thái được lưu vào bộ nhớ đệm của lượt trước, cho phép phản hồi gần như tức thì.

2. Hiểu video: Trong các tác vụ video understanding, các mô hình xử lý khung hình theo tuần tự. Tương tự như các token văn bản, các đặc trưng trực quan từ các khung hình trước có thể được lưu vào bộ nhớ đệm để giúp mô hình theo dõi đối tượng hoặc nhận dạng hành động mà không cần xử lý lại toàn bộ lịch sử video. Điều này đặc biệt có liên quan đối với action recognition nơi ngữ cảnh thời gian đóng vai trò cốt yếu.

Quản lý bộ nhớ hiệu quả#

Khi các mô hình ngày càng lớn hơn, kích thước của KV cache có thể trở thành một nút thắt cổ chai, tiêu tốn hàng gigabyte bộ nhớ GPU. Các tiến bộ gần đây tập trung vào việc tối ưu hóa khả năng lưu trữ này.

• PagedAttention: Lấy cảm hứng từ bộ nhớ ảo trong hệ điều hành, PagedAttention (do vLLM giới thiệu) cho phép KV cache được lưu trữ trong các khối bộ nhớ không liên tục. Điều này làm giảm sự phân mảnh và cho phép kích thước batch lớn hơn trong quá trình model serving.
• Lượng tử hóa KV Cache: Để tiết kiệm không gian, các nhà phát triển thường áp dụng mixed precision hoặc lượng tử hóa int8 dành riêng cho các giá trị được lưu vào bộ nhớ đệm. Điều này làm giảm dung lượng bộ nhớ, cho phép các thiết bị edge AI bị giới hạn RAM chạy được các mô hình mạnh mẽ.
• Prompt Caching: Một kỹ thuật liên quan trong đó các trạng thái KV của một system prompt tĩnh (ví dụ: "Bạn là một trợ lý lập trình hữu ích") được tính toán một lần và sử dụng lại trên nhiều phiên người dùng khác nhau. Đây là tính năng cốt lõi để tối ưu hóa quy trình làm việc prompt engineering ở quy mô lớn.

Phân biệt các khái niệm liên quan#

Việc phân biệt KV Cache với các thuật ngữ tối ưu hóa và caching khác là rất hữu ích:

• KV Cache và Prompt Caching: KV Cache thường đề cập đến bộ nhớ động, theo từng token được sử dụng trong một luồng tạo sinh duy nhất. Prompt caching đề cập cụ thể đến việc lưu trữ trạng thái đã xử lý của một chỉ lệnh đầu vào cố định để sử dụng lại trên nhiều lệnh gọi suy luận độc lập.
• KV Cache và Embeddings: Embeddings là các biểu diễn véc-tơ của dữ liệu đầu vào (văn bản hoặc hình ảnh) nắm bắt ý nghĩa ngữ nghĩa. KV cache lưu trữ các activations (key và value) bắt nguồn từ các embedding này trong các lớp chú ý, dành riêng cho mục đích tạo chuỗi.
• KV Cache và Model Weights: Trọng số mô hình là các tham số tĩnh, đã học của mạng nơ-ron. KV cache bao gồm dữ liệu tạm thời, động được tạo ra trong quá trình chuyển tiếp (forward pass) của một chuỗi đầu vào cụ thể.

Ví dụ: Ngữ cảnh trong các mô hình thị giác#

Mặc dù KV caching nổi tiếng nhất trong NLP, khái niệm duy trì trạng thái cũng áp dụng cho các mô hình thị giác tiên tiến. Trong ví dụ dưới đây, chúng ta mô phỏng ý tưởng truyền trạng thái (ngữ cảnh) trong kịch bản theo dõi video bằng cách sử dụng Ultralytics YOLO26. Ở đây, bộ theo dõi duy trì danh tính của các đối tượng qua các khung hình, về mặt khái niệm tương tự cách bộ nhớ đệm duy trì ngữ cảnh qua các token.

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()}")

Các nhà phát triển muốn quản lý tập dữ liệu và triển khai các mô hình đã tối ưu hóa có thể sử dụng Ultralytics Platform, giúp đơn giản hóa quy trình từ gán nhãn dữ liệu đến model deployment hiệu quả. Đối với những ai quan tâm đến cơ chế sâu hơn của attention, các thư viện như PyTorch cung cấp các khối nền tảng nơi các cơ chế bộ nhớ đệm này được triển khai.