Quản Lý KV-Cache Cho Sinh Văn Bản Tự Hồi Quy (Autoregressive)

Mục tiêu kiến thức

• Trong mô hình Transformer sinh văn bản (Decoder-only như GPT, LLaMA), tại mỗi bước sinh token mới t, việc tính lại toàn bộ Key (K) và Value (V) của các token cũ 1, …, t-1 sẽ có độ phức tạp O(T²) tính toán.
• KV Cache lưu trữ các tensor K, V của các bước trước đó trong bộ nhớ:
• Khi token mới đến, ta chỉ tính K_step, V_step của duy nhất token đó.
• Nối (concatenate) K_step, V_step vào cache theo chiều sequence length (dimension 2).
• Trả về tensor K, V đầy đủ để lớp Multi-Head Attention tính attention scores.
• Kỹ thuật Sliding Window / Max Context: Khi độ dài chuỗi vượt quá max_seq_len, chỉ giữ lại max_seq_len token mới nhất gần đây.

Yêu cầu

Xây dựng lớp KVCache:

class KVCache:
    def __init__(self, num_layers: int, max_seq_len: int | None = None):
        pass

    def update(self, key_step: torch.Tensor, value_step: torch.Tensor, layer_idx: int) -> tuple[torch.Tensor, torch.Tensor]:
        pass

    def reset(self) -> None:
        pass

    def get_seq_len(self, layer_idx: int) -> int:
        pass

• Định dạng tensor:
• key_step, value_step: Shape (batch_size, num_heads, seq_len_step, head_dim). Thường seq_len_step == 1 trong quá trình sinh token, hoặc > 1 ở bước Prefill.
• update(key_step, value_step, layer_idx):
• Nối tensor mới vào cache của tầng layer_idx dọc theo dim=2.
• Nếu max_seq_len được khai báo và seq_len > max_seq_len, cắt bỏ các token cũ ở đầu, chỉ giữ lại max_seq_len token cuối cùng [:, :, -self.max_seq_len:, :].
• Trả về tuple (full_key, full_value).

Ràng buộc

• Thời gian chạy tối đa: 2000ms.
• Giới hạn bộ nhớ: 512MB.
• Dữ liệu đầu vào hợp lệ theo đúng kiểu dữ liệu và miền giá trị được mô tả.