长上下文优化：注意力稀疏化与 KV Cache 管理

KV bytes per token = 2 (K 和 V) × n_layers × n_kv_heads × head_dim × dtype_bytes

Llama-3-8B: 2 × 32 层 × 8 kv_heads × 128 head_dim × 2 (FP16)
          = 131072 bytes ≈ 128 KB / token
128k token → 128000 × 128KB ≈ 16 GB   ← 仅 KV，不含权重

def sliding_window_mask(seq_len: int, window: int):
    """生成滑动窗口注意力 mask。"""
    import torch
    idx = torch.arange(seq_len)
    # |i - j| <= window 的位置可见
    mask = (idx[None, :] - idx[:, None]).abs() <= window
    return torch.where(mask, 0.0, float("-inf"))

class StreamingKVCache:
    """注意力汇 + 滑动窗口的 KV 管理。"""
    def __init__(self, num_sinks: int = 4, window: int = 4096):
        self.num_sinks = num_sinks
        self.window = window
        self.keys, self.values = [], []

    def append(self, k, v):
        self.keys.append(k)
        self.values.append(v)
        total = len(self.keys)
        if total > self.num_sinks + self.window:
            # 保留前 num_sinks 个 + 最近 window 个，丢弃中间
            keep = self.keys[: self.num_sinks] + self.keys[-self.window :]
            self.keys = keep
            self.values = self.values[: self.num_sinks] + self.values[-self.window :]

class H2OEviction:
    def __init__(self, budget: int = 2048, recent: int = 512):
        self.budget = budget          # 总保留 KV 数
        self.recent = recent          # 最近窗口强制保留
        self.scores = None            # 每个位置的累计注意力

    def update_scores(self, attn_weights):
        # attn_weights: [heads, q_len, kv_len]，对 query 维度求和
        cur = attn_weights.sum(dim=-2).mean(dim=0)
        if self.scores is None:
            self.scores = cur
        else:
            self.scores[: cur.shape[0]] += cur

    def select(self, kv_len: int):
        keep = self.budget - self.recent
        # 最近 recent 个无条件保留
        recent_idx = list(range(kv_len - self.recent, kv_len))
        # 其余按分数取 top
        head_idx = self.scores[: kv_len - self.recent].topk(keep).indices.tolist()
        return sorted(set(head_idx + recent_idx))

无缓存: 生成 n 个 token 需要 O(n²) 次注意力计算
有缓存: 每步 O(n) 读取缓存，总计 O(n²) 读取但 O(n) 计算

MHA: n_kv_heads = n_heads        （32/32，KV 最大）
GQA: 1 < n_kv_heads < n_heads    （8/32，折中，主流）
MQA: n_kv_heads = 1              （1/32，KV 最小，精度略降）

逻辑 KV:  [t0 t1 ... t15][t16 ... t31][t32 ...]
物理块:   块 #7          块 #2          块 #19
块表:     [7, 2, 19, ...]   ← 非连续，按需分配

vllm serve Qwen/Qwen2.5-7B-Instruct \
  --max-model-len 32768 \
  --gpu-memory-utilization 0.90 \
  --block-size 16

# vLLM 自动启用，OpenAI 兼容接口无需改代码
resp = client.chat.completions.create(
    model="Qwen/Qwen2.5-7B-Instruct",
    messages=[
        {"role": "system", "content": LONG_SYSTEM_PROMPT},   # 长且固定 → 命中缓存
        {"role": "user", "content": user_input},
    ],
)

# vLLM 开启 KV INT8/FP8
# --kv-cache-dtype fp8  或  auto

标准:   QK^T → 写 HBM → softmax → 写 HBM → 乘 V   （多次 HBM 往返）
Flash:  分块载入 SRAM → 在线 softmax 累加 → 直接出结果（HBM 只读写 O(n)）

# 安装并验证
# pip install flash-attn --no-build-isolation
from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen2.5-7B-Instruct",
    attn_implementation="flash_attention_2",
    torch_dtype="bfloat16",
    device_map="auto",
)

# vLLM 开启分块 prefill，降低长 prompt 对在线请求的干扰
vllm serve Qwen/Qwen2.5-7B-Instruct --enable-chunked-prefill --max-num-batched-tokens 8192

async def maybe_summarize(history: list[dict], max_tokens: int = 8000):
    if count_tokens(history) <= max_tokens:
        return history

    # 保留最近 N 轮原文
    keep_recent = 4
    old, recent = history[:-keep_recent], history[-keep_recent:]

    summary = await llm.complete(
        "把以下对话压缩成要点摘要，保留事实、数字、约定与未决问题：\n"
        + format_turns(old)
    )
    return [{"role": "system", "content": f"[历史摘要] {summary}"}] + recent

RoPE 把位置 m 编码为旋转角度 θ_i · m，其中 θ_i = base^(-2i/d)，base 默认 10000

# transformers 中启用 YaRN
model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen2.5-7B-Instruct",
    rope_scaling={
        "type": "yarn",
        "factor": 4.0,
        "original_max_position_embeddings": 32768,
    },
)

def needle_test(model, doc: str, needle: str, question: str, positions=(0.1, 0.5, 0.9)):
    results = {}
    for p in positions:
        idx = int(len(doc) * p)
        injected = doc[:idx] + f"\n{needle}\n" + doc[idx:]
        answer = model.generate(f"{injected}\n\n{question}")
        results[f"pos_{int(p*100)}%"] = needle in answer
    return results

1. 开启 FlashAttention-2：零成本加速
2. 开启 Chunked Prefill：保护在线请求的 TTFT
3. 开启 Prefix Caching：多轮/RAG 场景必开
4. 设置合理的 max-model-len：不要盲目拉满，KV 显存按需
5. 显存吃紧时：KV FP8 + GQA 模型 + 上下文压缩
6. 需要精确回忆：禁用驱逐类稀疏，改用 RAG 外挂检索

def kv_cache_gb(
    n_layers: int,
    n_kv_heads: int,
    head_dim: int,
    max_len: int,
    batch: int = 1,
    dtype_bytes: int = 2,      # FP16=2, FP8/INT8=1
) -> float:
    per_token = 2 * n_layers * n_kv_heads * head_dim * dtype_bytes
    return per_token * max_len * batch / (1024 ** 3)

# Llama-3-8B，FP16，8k 上下文
print(round(kv_cache_gb(32, 8, 128, 8192), 2))      # 2.0 GB
# 同配置拉到 128k
print(round(kv_cache_gb(32, 8, 128, 131072), 2))    # 32.0 GB
# 换成 FP8 KV
print(round(kv_cache_gb(32, 8, 128, 131072, dtype_bytes=1), 2))  # 16.0 GB

# 短上下文高并发实例
vllm serve Qwen/Qwen2.5-7B-Instruct --max-model-len 8192 --gpu-memory-utilization 0.92

# 长文档实例（低并发，开 KV FP8 省显存）
vllm serve Qwen/Qwen2.5-7B-Instruct \
  --max-model-len 131072 --kv-cache-dtype fp8 --enable-chunked-prefill

方案 A（全量上下文）：100k 输入 token × $3/M = $0.30 / 次
方案 B（RAG top-5）：  4k 输入 token × $3/M  = $0.012 / 次
                     + 检索成本（可忽略）
                     差 25 倍

用户请求
  ├─ 短问题（< 8k） → 直接进上下文
  └─ 长文档 → 先检索定位相关段落（RAG）
               └─ 若需全局推理 → 送入长上下文实例
                   └─ 配合前缀缓存复用文档 KV