超长上下文窗口工程实战：从KV Cache优化到注意力机制革新

# KV Cache内存占用计算
def calc_kv_cache_memory(
    num_layers: int,
    num_kv_heads: int,
    head_dim: int,
    seq_len: int,
    batch_size: int = 1,
    dtype_bytes: int = 2  # FP16
) -> float:
    # 计算KV Cache的内存占用（GB）
    # K和V各一份，所以乘以2
    bytes_per_token = 2 * num_kv_heads * head_dim * dtype_bytes * num_layers
    total_bytes = bytes_per_token * seq_len * batch_size
    return total_bytes / (1024 ** 3)

# Qwen3.8-27B参数（64层, 8个KV头, 每头128维）
kv_mem_32k = calc_kv_cache_memory(64, 8, 128, 32768)
kv_mem_128k = calc_kv_cache_memory(64, 8, 128, 131072)
kv_mem_1m = calc_kv_cache_memory(64, 8, 128, 1048576)

print(f"32K上下文 KV Cache: {kv_mem_32k:.2f} GB")
print(f"128K上下文 KV Cache: {kv_mem_128k:.2f} GB")
print(f"1M上下文 KV Cache: {kv_mem_1m:.2f} GB")

from collections import defaultdict
import torch

class PagedKVCache:
    def __init__(self, num_layers, num_kv_heads, head_dim,
                 page_size=16, num_pages=4096):
        self.page_size = page_size
        self.num_pages = num_pages
        # 预分配页表内存
        self.kv_pages = {
            layer: torch.zeros(
                num_pages, 2, page_size,
                num_kv_heads, head_dim,
                dtype=torch.float16
            )
            for layer in range(num_layers)
        }
        # 页表映射: seq_id -> [page_indices]
        self.page_table = defaultdict(list)
        # 空闲页列表
        self.free_pages = list(range(num_pages))

    def allocate(self, seq_id, num_tokens):
        # 为序列分配KV Cache页
        num_pages_needed = (num_tokens + self.page_size - 1) // self.page_size
        if len(self.free_pages) < num_pages_needed:
            raise MemoryError("KV Cache内存不足")
        pages = [self.free_pages.pop() for _ in range(num_pages_needed)]
        self.page_table[seq_id] = pages
        return pages

    def free(self, seq_id):
        # 释放序列的KV Cache页
        for page_idx in self.page_table[seq_id]:
            self.free_pages.append(page_idx)
        del self.page_table[seq_id]

import torch

def quantize_kv_cache_to_int8(kv_cache: torch.Tensor) -> tuple:
    # 将FP16 KV Cache量化为INT8
    abs_max = kv_cache.abs().amax(dim=-1, keepdim=True)
    scale = abs_max / 127.0
    scale = torch.clamp(scale, min=1e-8)
    quantized = (kv_cache / scale).round().clamp(-128, 127).to(torch.int8)
    return quantized, scale

def dequantize_kv_cache_from_int8(quantized: torch.Tensor,
                                   scale: torch.Tensor) -> torch.Tensor:
    # 将INT8 KV Cache反量化为FP16
    return quantized.to(torch.float16) * scale

class SlidingWindowKVCache:
    def __init__(self, window_size=8192, num_layers=64):
        self.window_size = window_size
        self.caches = [None] * num_layers

    def update(self, layer_idx, new_k, new_v):
        if self.caches[layer_idx] is None:
            self.caches[layer_idx] = (new_k, new_v)
        else:
            k, v = self.caches[layer_idx]
            # 拼接新KV
            k = torch.cat([k, new_k], dim=-2)
            v = torch.cat([v, new_v], dim=-2)
            # 滑动窗口截断
            if k.shape[-2] > self.window_size:
                k = k[..., -self.window_size:, :]
                v = v[..., -self.window_size:, :]
            self.caches[layer_idx] = (k, v)
        return self.caches[layer_idx]

import torch.nn as nn
import torch

class GatedDeltaNetAttention(nn.Module):
    def __init__(self, dim, num_heads, head_dim):
        super().__init__()
        self.num_heads = num_heads
        self.head_dim = head_dim
        self.q_proj = nn.Linear(dim, num_heads * head_dim)
        self.k_proj = nn.Linear(dim, num_heads * head_dim)
        self.v_proj = nn.Linear(dim, num_heads * head_dim)
        self.gate = nn.Linear(dim, num_heads)
        # Delta规则的状态矩阵
        self.register_buffer(
            'state', torch.zeros(num_heads, head_dim, head_dim)
        )

    def forward(self, x):
        B, T, C = x.shape
        q = self.q_proj(x).view(B, T, self.num_heads, self.head_dim)
        k = self.k_proj(x).view(B, T, self.num_heads, self.head_dim)
        v = self.v_proj(x).view(B, T, self.num_heads, self.head_dim)
        g = torch.sigmoid(self.gate(x))
        outputs = []
        for t in range(T):
            qt, kt, vt = q[:, t], k[:, t], v[:, t]
            gt = g[:, t]
            # Delta规则更新状态
            kv = torch.einsum('bhd,bhe->bde', vt, kt)
            kk = torch.einsum('bhd,bhe->bde', kt, kt)
            norm_sq = kt.norm(dim=-1, keepdim=True).unsqueeze(-1) + 1e-6
            self.state = (gt.unsqueeze(-1).unsqueeze(-1) * self.state
                          + (1 - gt.unsqueeze(-1).unsqueeze(-1))
                          * (kv - kk @ self.state / norm_sq))
            out = torch.einsum('bhd,bde->bhe', qt, self.state)
            outputs.append(out)
        return torch.stack(outputs, dim=1)

import torch

def flash_attention(q, k, v, block_size=128):
    # 分块注意力计算
    # q, k, v: (batch, heads, seq_len, head_dim)
    B, H, N, D = q.shape
    scale = D ** -0.5
    output = torch.zeros_like(q)

    for i in range(0, N, block_size):
        qi = q[:, :, i:i+block_size]
        max_score = torch.full(
            (B, H, min(block_size, N-i)), -1e9
        )
        sum_exp = torch.zeros(B, H, min(block_size, N-i))
        acc = torch.zeros_like(qi)

        for j in range(0, N, block_size):
            kj = k[:, :, j:j+block_size]
            vj = v[:, :, j:j+block_size]
            scores = torch.einsum('bhmd,bhnd->bhmn', qi, kj) * scale

            # 在线Softmax（关键创新）
            block_max = scores.amax(dim=-1, keepdim=True)
            new_max = torch.maximum(max_score, block_max)
            exp_diff = torch.exp(max_score - new_max)
            exp_scores = torch.exp(scores - new_max)
            sum_exp = sum_exp * exp_diff.squeeze(-1) + exp_scores.sum(dim=-1)
            acc = acc * exp_diff + torch.einsum(
                'bhmn,bhnd->bhmd', exp_scores, vj
            )
            max_score = new_max.squeeze(-1)

        output[:, :, i:i+block_size] = acc / sum_exp.unsqueeze(-1)
    return output