필사 모드: LLM 推理优化完全指南：KV Cache、Speculative Decoding、Continuous Batching

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
import time

def measure_prefill_decode_time(model, tokenizer, prompt: str, max_new_tokens: int = 100):
    """测量 Prefill 与 Decode 阶段耗时"""

    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
    input_len = inputs["input_ids"].size(1)

    # 测量 Prefill
    torch.cuda.synchronize()
    prefill_start = time.perf_counter()

    with torch.no_grad():
        # 直到第一个 token（Prefill + 第一次 Decode）
        first_output = model.generate(
            **inputs,
            max_new_tokens=1,
            do_sample=False
        )

    torch.cuda.synchronize()
    ttft = time.perf_counter() - prefill_start

    # 测量整体生成耗时
    torch.cuda.synchronize()
    total_start = time.perf_counter()

    with torch.no_grad():
        full_output = model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            do_sample=False
        )

    torch.cuda.synchronize()
    total_time = time.perf_counter() - total_start

    output_tokens = full_output.size(1) - input_len
    decode_time = total_time - ttft
    tpot = decode_time / max(output_tokens - 1, 1)

    print(f"输入 token 数: {input_len}")
    print(f"生成 token 数: {output_tokens}")
    print(f"TTFT（首 token 延迟）: {ttft * 1000:.1f} ms")
    print(f"TPOT（每 token 耗时）: {tpot * 1000:.1f} ms")
    print(f"吞吐量: {output_tokens / total_time:.1f} token/秒")

    return ttft, tpot

def analyze_memory_bandwidth():
    """LLM 推理的内存带宽分析"""

    # 示例：Llama-2-7B 配置
    model_params = {
        "num_layers": 32,
        "hidden_size": 4096,
        "num_heads": 32,
        "head_dim": 128,
        "vocab_size": 32000,
    }

    dtype_bytes = 2  # FP16：2 字节

    # 权重内存（只加载一次）
    # 每个 Transformer 层的权重
    attn_weight = 4 * model_params["hidden_size"] ** 2  # Q、K、V、O 投影
    ffn_weight = 8 * model_params["hidden_size"] ** 2  # Up、Gate、Down 投影（SwiGLU）
    layer_weight = (attn_weight + ffn_weight) * dtype_bytes

    total_weight_bytes = layer_weight * model_params["num_layers"]
    total_weight_gb = total_weight_bytes / 1e9

    print(f"模型权重: {total_weight_gb:.2f} GB")

    # KV Cache 内存（与序列长度成正比）
    seq_len = 2048
    kv_cache_per_token = (
        2 *  # K 和 V
        model_params["num_layers"] *
        model_params["num_heads"] *
        model_params["head_dim"] *
        dtype_bytes
    )

    kv_cache_total = kv_cache_per_token * seq_len / 1e6
    print(f"KV Cache（{seq_len} token）: {kv_cache_total:.2f} MB")
    print(f"每 token 的 KV Cache: {kv_cache_per_token} bytes")

    # A100 内存带宽：2 TB/s
    memory_bandwidth_tbs = 2.0  # TB/s

    # Decode 阶段：每生成一个 token 都要把权重读一遍
    # 批大小越小，内存读取相对计算量的占比就越高
    batch_size = 1
    flops_per_token = 2 * total_weight_bytes  # 近似 FLOPs

    # A100 FP16：312 TFLOPS
    compute_throughput = 312e12  # FLOPS

    # 以内存带宽为界的吞吐量
    memory_bound_tps = memory_bandwidth_tbs * 1e12 / total_weight_bytes
    # 以算力为界的吞吐量
    compute_bound_tps = compute_throughput / flops_per_token

    print(f"\n批大小为 {batch_size} 时:")
    print(f"内存瓶颈吞吐量: {memory_bound_tps:.1f} token/秒")
    print(f"算力瓶颈吞吐量: {compute_bound_tps:.1f} token/秒")
    print(f"实际瓶颈: {'内存' if memory_bound_tps < compute_bound_tps else '算力'}")

analyze_memory_bandwidth()

def estimate_inference_cost(
    model_size_b: float,
    tokens_per_request: int,
    requests_per_day: int,
    gpu_cost_per_hour: float = 3.0  # A100 每小时价格（美元）
):
    """推理成本估算"""

    # 吞吐量估算（经验数值）
    # 7B 模型：约 100 tok/s，70B 模型：约 20 tok/s（以 A100 为基准）
    throughput_tps = 100 / (model_size_b / 7) ** 0.6

    total_tokens_per_day = tokens_per_request * requests_per_day
    seconds_needed = total_tokens_per_day / throughput_tps
    hours_needed = seconds_needed / 3600

    # GPU 数量（考虑并行处理）
    # 这里假设用 1 块 GPU 处理
    daily_cost = hours_needed * gpu_cost_per_hour

    cost_per_1k_tokens = daily_cost / (total_tokens_per_day / 1000)

    print(f"模型: {model_size_b}B 参数")
    print(f"每日请求: {requests_per_day:,} 次")
    print(f"每请求 token 数: {tokens_per_request}")
    print(f"每日总 token 数: {total_tokens_per_day:,}")
    print(f"预计吞吐量: {throughput_tps:.1f} token/秒")
    print(f"所需 GPU 时间: {hours_needed:.2f} 小时")
    print(f"每日成本: ${daily_cost:.2f}")
    print(f"每 1K token 成本: ${cost_per_1k_tokens:.4f}")

# 示例
estimate_inference_cost(
    model_size_b=7.0,
    tokens_per_request=500,
    requests_per_day=10000
)

import torch
import torch.nn as nn
import math

class MultiHeadAttentionWithKVCache(nn.Module):
    """支持 KV Cache 的多头注意力"""

    def __init__(self, d_model: int, num_heads: int, max_seq_len: int = 4096):
        super().__init__()
        self.d_model = d_model
        self.num_heads = num_heads
        self.d_head = d_model // num_heads

        self.W_q = nn.Linear(d_model, d_model, bias=False)
        self.W_k = nn.Linear(d_model, d_model, bias=False)
        self.W_v = nn.Linear(d_model, d_model, bias=False)
        self.W_o = nn.Linear(d_model, d_model, bias=False)

        # 初始化 KV Cache
        self.register_buffer(
            'k_cache',
            torch.zeros(1, max_seq_len, num_heads, self.d_head)
        )
        self.register_buffer(
            'v_cache',
            torch.zeros(1, max_seq_len, num_heads, self.d_head)
        )
        self.cache_pos = 0

    def forward(
        self,
        x: torch.Tensor,
        use_cache: bool = True,
        position: int = None
    ):
        batch_size, seq_len, _ = x.shape

        # 计算 Q、K、V
        q = self.W_q(x).reshape(batch_size, seq_len, self.num_heads, self.d_head)
        k = self.W_k(x).reshape(batch_size, seq_len, self.num_heads, self.d_head)
        v = self.W_v(x).reshape(batch_size, seq_len, self.num_heads, self.d_head)

        if use_cache:
            # 把当前 K、V 存入缓存
            start_pos = self.cache_pos if position is None else position
            self.k_cache[:, start_pos:start_pos + seq_len] = k
            self.v_cache[:, start_pos:start_pos + seq_len] = v

            if position is None:
                self.cache_pos += seq_len

            # 使用缓存中的全部 K、V
            total_len = self.cache_pos if position is None else start_pos + seq_len
            k = self.k_cache[:, :total_len]
            v = self.v_cache[:, :total_len]

        # 计算注意力
        scale = math.sqrt(self.d_head)

        # 转换为 [batch, num_heads, seq_len, d_head] 形状
        q = q.transpose(1, 2)  # [B, H, S, D]
        k = k.transpose(1, 2)  # [B, H, T, D]（T：缓存的总长度）
        v = v.transpose(1, 2)

        # 注意力分数: [B, H, S, T]
        attn_scores = torch.matmul(q, k.transpose(-2, -1)) / scale

        # Softmax
        attn_weights = torch.softmax(attn_scores, dim=-1)

        # 加权求和: [B, H, S, D]
        output = torch.matmul(attn_weights, v)

        # 转换回原始形状
        output = output.transpose(1, 2).reshape(batch_size, seq_len, self.d_model)
        output = self.W_o(output)

        return output

    def clear_cache(self):
        """清空缓存"""
        self.k_cache.zero_()
        self.v_cache.zero_()
        self.cache_pos = 0

def calculate_kv_cache_memory(
    model_config: dict,
    batch_size: int,
    seq_len: int,
    dtype_bytes: int = 2  # FP16
) -> dict:
    """计算 KV Cache 内存占用"""

    num_layers = model_config["num_layers"]
    num_kv_heads = model_config.get("num_kv_heads", model_config["num_heads"])
    head_dim = model_config["head_dim"]

    # KV Cache 大小：2（K+V）* layers * kv_heads * head_dim * seq_len * dtype
    kv_cache_bytes = (
        2 *           # K 和 V
        num_layers *
        num_kv_heads *
        head_dim *
        seq_len *
        batch_size *
        dtype_bytes
    )

    return {
        "kv_cache_bytes": kv_cache_bytes,
        "kv_cache_mb": kv_cache_bytes / 1e6,
        "kv_cache_gb": kv_cache_bytes / 1e9,
        "per_token_bytes": kv_cache_bytes // seq_len,
    }

# 各模型的 KV Cache 对比
models = {
    "Llama-2-7B": {
        "num_layers": 32, "num_heads": 32,
        "num_kv_heads": 32, "head_dim": 128
    },
    "Llama-2-13B": {
        "num_layers": 40, "num_heads": 40,
        "num_kv_heads": 40, "head_dim": 128
    },
    "Llama-2-70B (GQA)": {
        "num_layers": 80, "num_heads": 64,
        "num_kv_heads": 8, "head_dim": 128  # GQA：8 个 KV 头
    },
    "Mistral-7B (GQA)": {
        "num_layers": 32, "num_heads": 32,
        "num_kv_heads": 8, "head_dim": 128  # GQA：8 个 KV 头
    },
}

print("KV Cache 内存占用（batch=1, seq=4096）")
print("=" * 70)
for name, config in models.items():
    result = calculate_kv_cache_memory(config, batch_size=1, seq_len=4096)
    print(f"{name:<25} {result['kv_cache_gb']:.2f} GB  "
          f"(每 token {result['per_token_bytes']:,} bytes)")

import torch
import torch.nn as nn
import math

class GroupedQueryAttention(nn.Module):
    """Grouped Query Attention（GQA）实现"""

    def __init__(
        self,
        d_model: int,
        num_q_heads: int,
        num_kv_heads: int,
    ):
        super().__init__()
        assert num_q_heads % num_kv_heads == 0

        self.d_model = d_model
        self.num_q_heads = num_q_heads
        self.num_kv_heads = num_kv_heads
        self.num_groups = num_q_heads // num_kv_heads
        self.d_head = d_model // num_q_heads

        self.W_q = nn.Linear(d_model, num_q_heads * self.d_head, bias=False)
        self.W_k = nn.Linear(d_model, num_kv_heads * self.d_head, bias=False)
        self.W_v = nn.Linear(d_model, num_kv_heads * self.d_head, bias=False)
        self.W_o = nn.Linear(d_model, d_model, bias=False)

    def forward(self, x: torch.Tensor, kv_cache=None):
        batch_size, seq_len, _ = x.shape

        # Q: [B, S, num_q_heads * d_head]
        q = self.W_q(x).reshape(batch_size, seq_len, self.num_q_heads, self.d_head)
        k = self.W_k(x).reshape(batch_size, seq_len, self.num_kv_heads, self.d_head)
        v = self.W_v(x).reshape(batch_size, seq_len, self.num_kv_heads, self.d_head)

        # 更新 KV Cache
        if kv_cache is not None:
            k = torch.cat([kv_cache["k"], k], dim=1)
            v = torch.cat([kv_cache["v"], v], dim=1)
        new_kv_cache = {"k": k, "v": v}

        total_len = k.size(1)

        # 转换为 [B, num_heads, S, d_head] 形状
        q = q.transpose(1, 2)  # [B, Q_heads, S, d_head]
        k = k.transpose(1, 2)  # [B, KV_heads, T, d_head]
        v = v.transpose(1, 2)

        # GQA：将 KV 头重复到 Q 头的数量
        # k: [B, KV_heads, T, d_head] -> [B, Q_heads, T, d_head]
        k = k.repeat_interleave(self.num_groups, dim=1)
        v = v.repeat_interleave(self.num_groups, dim=1)

        # 计算注意力
        scale = math.sqrt(self.d_head)
        attn_scores = torch.matmul(q, k.transpose(-2, -1)) / scale
        attn_weights = torch.softmax(attn_scores, dim=-1)
        output = torch.matmul(attn_weights, v)

        output = output.transpose(1, 2).reshape(batch_size, seq_len, self.d_model)
        return self.W_o(output), new_kv_cache

# MHA vs GQA vs MQA 内存对比
def compare_attention_variants():
    """不同注意力变体的 KV Cache 内存对比"""

    # 以 70B 模型为例（Llama-2-70B）
    num_layers = 80
    d_head = 128
    seq_len = 4096
    batch_size = 1
    dtype_bytes = 2  # FP16

    variants = {
        "MHA (32 KV heads)": 32,
        "GQA (8 KV heads)": 8,
        "MQA (1 KV head)": 1,
    }

    print("70B 模型不同注意力变体的 KV Cache 对比")
    print(f"(seq_len={seq_len}, batch={batch_size})")
    print("=" * 55)

    for name, num_kv_heads in variants.items():
        kv_bytes = 2 * num_layers * num_kv_heads * d_head * seq_len * batch_size * dtype_bytes
        kv_gb = kv_bytes / 1e9
        print(f"{name:<25} {kv_gb:.2f} GB")

compare_attention_variants()

class MultiHeadLatentAttention(nn.Module):
    """
    DeepSeek MLA - 将 KV Cache 压缩为低维潜在向量

    核心思路：
    - 不再以高维形式存储 KV，而是存储低维潜在向量 c_kv
    - 从 c_kv 还原出 K、V（上投影）
    - KV Cache 大小：num_layers * kv_lora_rank * seq_len
      （对比传统方式：2 * num_layers * num_kv_heads * d_head * seq_len）
    """

    def __init__(
        self,
        d_model: int = 5120,
        num_heads: int = 128,
        kv_lora_rank: int = 512,  # 低维潜在维度
        qk_nope_head_dim: int = 128,
        qk_rope_head_dim: int = 64,
        v_head_dim: int = 128,
    ):
        super().__init__()

        self.d_model = d_model
        self.num_heads = num_heads
        self.kv_lora_rank = kv_lora_rank
        self.qk_nope_head_dim = qk_nope_head_dim
        self.qk_rope_head_dim = qk_rope_head_dim
        self.v_head_dim = v_head_dim

        # Q 投影（LoRA 风格）
        self.q_a_proj = nn.Linear(d_model, 1536, bias=False)  # 下投影
        self.q_b_proj = nn.Linear(1536, num_heads * (qk_nope_head_dim + qk_rope_head_dim), bias=False)

        # KV 下投影：d_model -> kv_lora_rank
        # KV Cache 里只存储这一部分！
        self.kv_a_proj = nn.Linear(
            d_model,
            kv_lora_rank + qk_rope_head_dim,
            bias=False
        )

        # KV 上投影：kv_lora_rank -> K, V
        self.kv_b_proj = nn.Linear(
            kv_lora_rank,
            num_heads * (qk_nope_head_dim + v_head_dim),
            bias=False
        )

        self.o_proj = nn.Linear(num_heads * v_head_dim, d_model, bias=False)

    def forward(self, x: torch.Tensor, compressed_kv_cache=None):
        """
        Args:
            x: [batch, seq, d_model]
            compressed_kv_cache: [batch, cache_len, kv_lora_rank + rope_dim]
        """
        batch_size, seq_len, _ = x.shape

        # 计算 Q
        q = self.q_b_proj(self.q_a_proj(x))

        # KV 压缩（缓存中只存储这个结果）
        kv_compressed = self.kv_a_proj(x)  # [B, S, kv_lora_rank + rope_dim]

        # 更新 KV Cache
        if compressed_kv_cache is not None:
            kv_compressed_total = torch.cat([compressed_kv_cache, kv_compressed], dim=1)
        else:
            kv_compressed_total = kv_compressed

        # 从缓存的压缩 KV 中还原真正的 K、V（上投影）
        kv_content = kv_compressed_total[:, :, :self.kv_lora_rank]  # 排除 rope 部分
        kv_full = self.kv_b_proj(kv_content)  # [B, T, num_heads * (nope + v_dim)]

        # 最终的注意力计算（省略）

        return None, kv_compressed

# KV Cache 大小对比（以 DeepSeek-V2 为基准）
def compare_mla_vs_mha():
    """MLA vs MHA 的 KV Cache 对比"""

    seq_len = 4096
    dtype_bytes = 2  # BF16
    num_layers = 60  # DeepSeek-V2

    # MHA（传统方式）
    num_heads = 128
    head_dim = 128
    mha_kv_gb = 2 * num_layers * num_heads * head_dim * seq_len * dtype_bytes / 1e9

    # MLA（DeepSeek-V2）
    kv_lora_rank = 512
    rope_dim = 64
    mla_kv_gb = (kv_lora_rank + rope_dim) * num_layers * seq_len * dtype_bytes / 1e9

    print(f"MHA KV Cache: {mha_kv_gb:.2f} GB")
    print(f"MLA KV Cache: {mla_kv_gb:.2f} GB")
    print(f"节省比例: {mha_kv_gb / mla_kv_gb:.1f}x")

compare_mla_vs_mha()

请求 1: [PROMPT=200 tokens] [KV_CACHE=预留最多 2048-200=1848 tokens] → 内部碎片化
请求 2: [PROMPT=100 tokens] [KV_CACHE=预留 1948 tokens]
请求 3: 因内存不足而等待（外部碎片化）

from dataclasses import dataclass, field
from typing import Dict, List, Optional, Set
import torch

@dataclass
class PhysicalBlock:
    """物理内存块"""
    block_id: int
    block_size: int  # 块中可存储的 token 数（例如 16）
    device: str = "cuda"
    ref_count: int = 0  # 引用计数（用于 CoW）

    def __post_init__(self):
        # 实际分配 KV 张量
        # [2, num_layers, block_size, num_heads, head_dim]
        pass


@dataclass
class LogicalBlock:
    """逻辑块（与请求相映射）"""
    physical_block_id: int
    num_filled: int = 0  # 当前已填充的 token 数

class PagedKVCacheManager:
    """PagedAttention 的 KV Cache 管理器"""

    def __init__(
        self,
        num_physical_blocks: int,
        block_size: int,
        num_layers: int,
        num_kv_heads: int,
        head_dim: int,
        device: str = "cuda"
    ):
        self.block_size = block_size
        self.num_layers = num_layers
        self.num_kv_heads = num_kv_heads
        self.head_dim = head_dim
        self.device = device

        # 初始化物理块池
        self.free_blocks: List[int] = list(range(num_physical_blocks))
        self.all_blocks: Dict[int, PhysicalBlock] = {
            i: PhysicalBlock(block_id=i, block_size=block_size)
            for i in range(num_physical_blocks)
        }

        # 按请求划分的逻辑块表
        self.block_tables: Dict[int, List[LogicalBlock]] = {}

        # 实际的 KV Cache 张量
        # [num_blocks, 2, num_layers, block_size, num_kv_heads, head_dim]
        self.kv_cache = torch.zeros(
            num_physical_blocks, 2, num_layers, block_size, num_kv_heads, head_dim,
            dtype=torch.float16,
            device=device
        )

    def allocate_blocks_for_request(self, request_id: int, num_tokens: int):
        """为请求分配所需的块"""
        num_blocks_needed = (num_tokens + self.block_size - 1) // self.block_size

        if len(self.free_blocks) < num_blocks_needed:
            raise RuntimeError(f"内存不足: 需要 {num_blocks_needed} 个块，当前可用 {len(self.free_blocks)} 个")

        logical_blocks = []
        for i in range(num_blocks_needed):
            physical_id = self.free_blocks.pop(0)
            self.all_blocks[physical_id].ref_count = 1
            logical_blocks.append(
                LogicalBlock(physical_block_id=physical_id)
            )

        self.block_tables[request_id] = logical_blocks
        print(f"请求 {request_id}: 分配 {num_blocks_needed} 个块，"
              f"剩余块数: {len(self.free_blocks)}")

    def append_token(self, request_id: int, layer: int, token_pos: int, k: torch.Tensor, v: torch.Tensor):
        """把新 token 的 KV 追加到缓存中"""
        block_idx = token_pos // self.block_size
        token_in_block = token_pos % self.block_size

        logical_block = self.block_tables[request_id][block_idx]
        physical_id = logical_block.physical_block_id

        # 存入 KV Cache
        self.kv_cache[physical_id, 0, layer, token_in_block] = k  # K
        self.kv_cache[physical_id, 1, layer, token_in_block] = v  # V
        logical_block.num_filled = token_in_block + 1

    def get_physical_block_ids(self, request_id: int) -> List[int]:
        """返回请求对应的物理块 ID 列表"""
        return [lb.physical_block_id for lb in self.block_tables[request_id]]

    def free_request(self, request_id: int):
        """请求完成后释放对应的块"""
        if request_id in self.block_tables:
            for logical_block in self.block_tables[request_id]:
                phys_id = logical_block.physical_block_id
                self.all_blocks[phys_id].ref_count -= 1

                if self.all_blocks[phys_id].ref_count == 0:
                    self.free_blocks.append(phys_id)

            del self.block_tables[request_id]

    def copy_on_write(self, src_request_id: int, dst_request_id: int):
        """用于 Prefix Caching 的 Copy-on-Write"""
        src_blocks = self.block_tables[src_request_id]
        dst_blocks = []

        for logical_block in src_blocks:
            phys_id = logical_block.physical_block_id
            # 增加引用计数（实际只在写入时才复制）
            self.all_blocks[phys_id].ref_count += 1
            dst_blocks.append(
                LogicalBlock(physical_block_id=phys_id, num_filled=logical_block.num_filled)
            )

        self.block_tables[dst_request_id] = dst_blocks


# 使用示例
manager = PagedKVCacheManager(
    num_physical_blocks=1000,
    block_size=16,
    num_layers=32,
    num_kv_heads=32,
    head_dim=128
)

# 处理 3 个请求
manager.allocate_blocks_for_request(request_id=1, num_tokens=200)
manager.allocate_blocks_for_request(request_id=2, num_tokens=500)
manager.allocate_blocks_for_request(request_id=3, num_tokens=100)

# 请求 1 完成后释放
manager.free_request(request_id=1)
print(f"\n请求 1 完成后的可用块数: {len(manager.free_blocks)}")

时间 t=0: [请求A: 500 token] [请求B: 100 token] [请求C: 300 token]
时间 t=1: 请求B 完成，但由于 A、C 仍在处理中，即便 GPU 有空闲也无法接入新请求
时间 t=2: 请求C 完成
时间 t=3: 请求A 完成 → 此时才能开始新一批！

from typing import List, Optional, Tuple
from dataclasses import dataclass
import asyncio
import torch
from queue import Queue
import threading

@dataclass
class Request:
    """推理请求"""
    request_id: str
    input_ids: List[int]
    max_new_tokens: int
    generated_ids: List[int] = None
    is_finished: bool = False

    def __post_init__(self):
        self.generated_ids = []

class ContinuousBatchingScheduler:
    """Continuous Batching 调度器"""

    def __init__(
        self,
        max_batch_size: int = 32,
        max_seq_len: int = 4096
    ):
        self.max_batch_size = max_batch_size
        self.max_seq_len = max_seq_len

        self.waiting_queue: List[Request] = []
        self.running_requests: List[Request] = []
        self.finished_requests: List[Request] = []

    def add_request(self, request: Request):
        """添加新请求"""
        self.waiting_queue.append(request)

    def _can_add_request(self, request: Request) -> bool:
        """检查是否可以把请求加入批次（内存检查）"""
        current_batch_size = len(self.running_requests) + 1
        if current_batch_size > self.max_batch_size:
            return False

        # KV Cache 内存检查（简化版）
        total_tokens = sum(
            len(r.input_ids) + len(r.generated_ids)
            for r in self.running_requests
        ) + len(request.input_ids)

        return total_tokens < self.max_seq_len * self.max_batch_size

    def schedule_iteration(self) -> Tuple[List[Request], List[str]]:
        """
        为一次 iteration 调度批次

        Returns:
            (待执行的请求, 已完成的请求 ID)
        """
        completed_ids = []

        # 处理已完成的请求
        still_running = []
        for req in self.running_requests:
            if req.is_finished:
                self.finished_requests.append(req)
                completed_ids.append(req.request_id)
            else:
                still_running.append(req)
        self.running_requests = still_running

        # 把等待中的请求加入批次（核心：立即填补空出的槽位）
        while self.waiting_queue and self._can_add_request(self.waiting_queue[0]):
            new_request = self.waiting_queue.pop(0)
            self.running_requests.append(new_request)
            print(f"将请求 {new_request.request_id} 加入批次 "
                  f"(当前批大小: {len(self.running_requests)})")

        return self.running_requests, completed_ids

    def simulate_one_step(self, model_forward_fn):
        """模拟一个步骤"""

        active_requests, completed = self.schedule_iteration()

        if not active_requests:
            return []

        # 准备当前批次的输入
        # Prefill 请求：只有 input_ids 的情况
        # Decode 请求：已有此前 KV Cache 的情况
        batch_input_ids = []
        for req in active_requests:
            if len(req.generated_ids) == 0:
                # Prefill
                batch_input_ids.append(req.input_ids)
            else:
                # Decode（只取最后生成的 token）
                batch_input_ids.append([req.generated_ids[-1]])

        # 执行模型（实际中会通过 PagedAttention 处理）
        outputs = model_forward_fn(batch_input_ids)

        # 处理下一个 token
        for req, next_token_id in zip(active_requests, outputs):
            req.generated_ids.append(next_token_id)

            # 检查终止条件
            if (next_token_id == 2 or  # EOS token
                    len(req.generated_ids) >= req.max_new_tokens):
                req.is_finished = True

        return completed

# 吞吐量对比模拟
def simulate_throughput_comparison():
    """静态批处理 vs Continuous Batching 吞吐量对比"""

    import random

    requests = [
        Request(
            request_id=str(i),
            input_ids=list(range(random.randint(50, 200))),
            max_new_tokens=random.randint(50, 500)
        )
        for i in range(20)
    ]

    # 静态批处理：等待所有请求都完成
    max_tokens_static = max(r.max_new_tokens for r in requests)
    total_iterations_static = max_tokens_static * 4  # 每批 4 个

    # Continuous Batching：一旦有请求完成就立刻加入新请求
    total_tokens = sum(r.max_new_tokens for r in requests)
    total_iterations_cb = total_tokens  # 粗略估算

    print(f"静态批处理预计 iteration 数: {total_iterations_static}")
    print(f"Continuous Batching 预计 iteration 数: {total_iterations_cb}")
    print(f"吞吐量提升: {total_iterations_static / total_iterations_cb:.2f}x")

传统方式: [大模型] → token1 → token2 → token3 → token4 → token5
投机方式: [小模型] → 并行生成 (token1, token2, token3, token4, token5)
          [大模型] → 一次性验证这 5 个 token（像 Prefill 一样并行！）
          只使用被接受的 token

import numpy as np
import torch
from typing import List, Tuple

def speculative_decode_step(
    draft_model,
    target_model,
    input_ids: torch.Tensor,
    draft_steps: int = 4,
    temperature: float = 1.0
) -> Tuple[torch.Tensor, int, int]:
    """
    Speculative Decoding 的一个步骤

    Returns:
        (生成的 token, 被接受的 token 数, 草稿 token 数)
    """
    batch_size = input_ids.size(0)

    # 1. 用草稿模型生成候选 token
    draft_tokens = []
    draft_probs = []

    current_ids = input_ids.clone()

    for _ in range(draft_steps):
        with torch.no_grad():
            draft_output = draft_model(current_ids)
            draft_logits = draft_output.logits[:, -1, :]  # [B, vocab_size]

        # 计算草稿概率
        if temperature > 0:
            draft_prob = torch.softmax(draft_logits / temperature, dim=-1)
        else:
            draft_prob = torch.zeros_like(draft_logits)
            draft_prob.scatter_(1, draft_logits.argmax(dim=-1, keepdim=True), 1.0)

        # 采样草稿 token
        draft_token = torch.multinomial(draft_prob, num_samples=1)  # [B, 1]
        draft_tokens.append(draft_token)
        draft_probs.append(draft_prob)

        # 为下一步追加 token
        current_ids = torch.cat([current_ids, draft_token], dim=1)

    # 把草稿 token 合并为一个张量
    draft_sequence = torch.cat(draft_tokens, dim=1)  # [B, draft_steps]
    candidate_ids = torch.cat([input_ids, draft_sequence], dim=1)

    # 2. 用验证模型一次性验证草稿 token
    with torch.no_grad():
        target_output = target_model(candidate_ids)
        target_logits = target_output.logits[:, input_ids.size(1) - 1:-1, :]  # [B, draft_steps, vocab_size]

    # 验证模型给出的概率
    if temperature > 0:
        target_probs = torch.softmax(target_logits / temperature, dim=-1)
    else:
        target_probs = torch.zeros_like(target_logits)
        target_probs.scatter_(2, target_logits.argmax(dim=-1, keepdim=True), 1.0)

    # 3. 决定每个草稿 token 是接受还是拒绝
    accepted_tokens = []
    num_accepted = 0

    for step in range(draft_steps):
        token = draft_sequence[:, step]  # [B]

        # 计算接受概率: min(1, p_target / p_draft)
        p_draft = draft_probs[step].gather(1, token.unsqueeze(1)).squeeze(1)
        p_target = target_probs[:, step, :].gather(1, token.unsqueeze(1)).squeeze(1)

        acceptance_prob = torch.clamp(p_target / (p_draft + 1e-8), max=1.0)

        # 随机决定接受或拒绝
        random_val = torch.rand_like(acceptance_prob)
        accepted = random_val < acceptance_prob  # [B]

        if not accepted.all():
            # 在第一个被拒绝的 token 处停止
            break

        accepted_tokens.append(token)
        num_accepted += 1

    # 4. 最后一个 token：由验证模型生成（或从修正后的分布中采样）
    last_target_logits = target_output.logits[:, input_ids.size(1) + num_accepted - 1, :]

    if temperature > 0:
        last_prob = torch.softmax(last_target_logits / temperature, dim=-1)
    else:
        last_prob = torch.zeros_like(last_target_logits)
        last_prob.scatter_(1, last_target_logits.argmax(dim=-1, keepdim=True), 1.0)

    # 分布修正（发生拒绝时）
    if num_accepted < draft_steps:
        # 使用 max(0, p_target - p_draft)
        correction = torch.clamp(
            last_prob - draft_probs[num_accepted],
            min=0
        )
        correction = correction / (correction.sum(dim=-1, keepdim=True) + 1e-8)
        last_token = torch.multinomial(correction, num_samples=1)
    else:
        last_token = torch.multinomial(last_prob, num_samples=1)

    accepted_tokens.append(last_token.squeeze(1))

    final_tokens = torch.stack(accepted_tokens, dim=1)

    return final_tokens, num_accepted, draft_steps


def analyze_speedup(acceptance_rate: float, draft_steps: int = 4) -> dict:
    """按接受率分析加速效果"""

    # 计算期望值
    # E[accepted tokens] = sum_{k=0}^{K} alpha^k = (1 - alpha^{K+1}) / (1 - alpha)
    # 其中 alpha = acceptance_rate, K = draft_steps

    expected_accepted = sum(
        acceptance_rate ** k for k in range(draft_steps + 1)
    )

    # 实际加速比（包含草稿模型的开销）
    # 假设草稿模型是验证模型大小的 1/10
    draft_model_ratio = 0.1

    # 耗时：草稿 K 步 + 验证 1 步
    # 传统方式：K+1 步
    # 投机方式：K * draft_ratio + 1 步（验证）
    steps_with_speculative = draft_steps * draft_model_ratio + 1
    expected_tokens_with_speculative = expected_accepted

    speedup = expected_tokens_with_speculative / steps_with_speculative

    return {
        "acceptance_rate": acceptance_rate,
        "draft_steps": draft_steps,
        "expected_accepted_tokens": expected_accepted,
        "speedup": speedup
    }

# 按接受率输出加速效果
print("Speculative Decoding 按接受率的加速效果（草稿 K=4）")
print("=" * 55)
for alpha in [0.5, 0.6, 0.7, 0.8, 0.9, 0.95]:
    result = analyze_speedup(alpha, draft_steps=4)
    print(f"接受率 {alpha:.0%}: 期望接受 {result['expected_accepted_tokens']:.2f} token, "
          f"加速 {result['speedup']:.2f}x")

import torch
import torch.nn as nn

class MedusaHead(nn.Module):
    """
    Medusa: 在单个模型上附加多个草稿头

    每个头预测未来的一个 token：
    - Head 1: 预测 t+1
    - Head 2: 预测 t+2
    - Head N: 预测 t+N
    """

    def __init__(
        self,
        hidden_size: int,
        vocab_size: int,
        num_heads: int = 4,
        hidden_layers: int = 1
    ):
        super().__init__()
        self.num_heads = num_heads

        # 为每个未来 token 位置准备独立的头
        self.heads = nn.ModuleList([
            nn.Sequential(
                *[nn.Linear(hidden_size, hidden_size, bias=False),
                  nn.SiLU()] * hidden_layers,
                nn.Linear(hidden_size, vocab_size, bias=False)
            )
            for _ in range(num_heads)
        ])

    def forward(self, hidden_states: torch.Tensor):
        """
        Args:
            hidden_states: [batch, seq, hidden_size] - 基础模型最后一层的隐藏状态

        Returns:
            每个未来位置对应的 logits 列表 [batch, seq, vocab]
        """
        return [head(hidden_states) for head in self.heads]


class MedusaModel(nn.Module):
    """Medusa 完整模型"""

    def __init__(self, base_model, vocab_size: int, num_medusa_heads: int = 4):
        super().__init__()
        self.base_model = base_model
        hidden_size = base_model.config.hidden_size

        self.medusa_heads = MedusaHead(
            hidden_size=hidden_size,
            vocab_size=vocab_size,
            num_heads=num_medusa_heads
        )

    def forward(self, input_ids: torch.Tensor, use_medusa: bool = False):
        # 执行基础模型
        base_output = self.base_model(
            input_ids,
            output_hidden_states=True
        )

        base_logits = base_output.logits

        if not use_medusa:
            return base_logits, None

        # 用 Medusa 头预测未来 token
        last_hidden_state = base_output.hidden_states[-1]
        medusa_logits = self.medusa_heads(last_hidden_state)

        return base_logits, medusa_logits

    def generate_with_medusa(
        self,
        input_ids: torch.Tensor,
        max_new_tokens: int = 100,
        temperature: float = 1.0,
        medusa_choices: int = 16,  # 候选 token 数
        threshold: float = 0.09    # 接受阈值
    ):
        """使用 Medusa 加速生成"""

        current_ids = input_ids.clone()
        all_accepted = []

        while len(all_accepted) < max_new_tokens:
            # 用 Medusa 头预测候选 token
            base_logits, medusa_logits = self.forward(current_ids, use_medusa=True)

            # 在每个位置选出排名靠前的候选
            candidates = []
            base_probs = torch.softmax(base_logits[:, -1, :] / temperature, dim=-1)
            top_tokens = torch.topk(base_probs, medusa_choices)[1]

            for head_logits in medusa_logits:
                head_probs = torch.softmax(head_logits[:, -1, :] / temperature, dim=-1)
                candidates.append(torch.topk(head_probs, medusa_choices)[1])

            # 用树状注意力验证候选（简化版）
            # 实际实现会用树状掩码来做高效验证
            best_token = top_tokens[0, 0]
            all_accepted.append(best_token.item())

            current_ids = torch.cat([current_ids, best_token.unsqueeze(0).unsqueeze(0)], dim=1)

            if best_token.item() == 2:  # EOS
                break

        return all_accepted

标准 Attention 的内存操作:
1. 从 HBM 读取 Q、K            → 读取: O(N * d)
2. 计算 S = Q @ K.T            → 写入: O(N^2) ← 瓶颈！
3. 从 HBM 读取 S 做 Softmax    → 读取: O(N^2)
4. 存储 P = softmax(S)         → 写入: O(N^2)
5. 读取 P 计算 P @ V           → 读取: O(N^2)
6. 存储最终结果                → 写入: O(N * d)

HBM 访问总量: O(N^2)（与序列长度的平方成正比！）

import torch
import math

def flash_attention_v1(Q, K, V, block_size=64):
    """
    FlashAttention v1 的简化实现
    通过分块（tiling），避免把整个注意力矩阵存入 HBM

    核心: 用 Online Softmax 算法按块处理
    """
    batch_size, num_heads, seq_len, d_head = Q.shape

    scale = 1.0 / math.sqrt(d_head)
    Q = Q * scale

    # 初始化输出张量（保留在 SRAM 中）
    O = torch.zeros_like(Q)
    L = torch.zeros(batch_size, num_heads, seq_len, 1, device=Q.device)  # softmax 分母
    M = torch.full((batch_size, num_heads, seq_len, 1), float('-inf'), device=Q.device)  # 最大值

    num_blocks = (seq_len + block_size - 1) // block_size

    for j in range(num_blocks):
        # 加载 K、V 块（HBM → SRAM）
        k_start = j * block_size
        k_end = min((j + 1) * block_size, seq_len)
        K_j = K[:, :, k_start:k_end, :]
        V_j = V[:, :, k_start:k_end, :]

        for i in range(num_blocks):
            # 加载 Q 块
            q_start = i * block_size
            q_end = min((i + 1) * block_size, seq_len)
            Q_i = Q[:, :, q_start:q_end, :]
            O_i = O[:, :, q_start:q_end, :]
            L_i = L[:, :, q_start:q_end, :]
            M_i = M[:, :, q_start:q_end, :]

            # 计算注意力分数（在 SRAM 中）
            S_ij = torch.matmul(Q_i, K_j.transpose(-2, -1))  # [B, H, Br, Bc]

            # 更新 Online Softmax
            M_ij_new = torch.maximum(M_i, S_ij.max(dim=-1, keepdim=True)[0])
            P_ij = torch.exp(S_ij - M_ij_new)
            L_ij_new = torch.exp(M_i - M_ij_new) * L_i + P_ij.sum(dim=-1, keepdim=True)

            # 更新输出（重新缩放）
            O_i_new = (
                torch.exp(M_i - M_ij_new) * O_i +
                torch.matmul(P_ij, V_j)
            )

            # 把该块的结果写回 HBM
            O[:, :, q_start:q_end, :] = O_i_new
            L[:, :, q_start:q_end, :] = L_ij_new
            M[:, :, q_start:q_end, :] = M_ij_new

    # 最终归一化
    O = O / L

    return O


def compare_attention_implementations():
    """FlashAttention vs 标准注意力对比"""

    batch_size = 2
    num_heads = 32
    seq_len = 4096
    d_head = 128

    Q = torch.randn(batch_size, num_heads, seq_len, d_head, device='cuda', dtype=torch.float16)
    K = torch.randn_like(Q)
    V = torch.randn_like(Q)

    # PyTorch SDPA（内置 FlashAttention 2 实现）
    import torch.nn.functional as F

    with torch.backends.cuda.sdp_kernel(
        enable_flash=True,
        enable_math=False,
        enable_mem_efficient=False
    ):
        flash_output = F.scaled_dot_product_attention(Q, K, V)

    # 标准注意力
    scale = 1.0 / math.sqrt(d_head)
    attn_scores = torch.matmul(Q, K.transpose(-2, -1)) * scale
    attn_weights = torch.softmax(attn_scores, dim=-1)
    standard_output = torch.matmul(attn_weights, V)

    # 结果对比
    max_diff = (flash_output - standard_output).abs().max().item()
    print(f"FlashAttention vs 标准注意力的最大差异: {max_diff:.6f}")

    # 内存占用对比
    standard_attn_matrix_size = batch_size * num_heads * seq_len * seq_len * 2  # FP16
    print(f"标准注意力矩阵内存: {standard_attn_matrix_size / 1e9:.2f} GB")
    print(f"FlashAttention 矩阵内存: ~0 GB（分块处理，无需存储）")

import torch
import torch.nn.functional as F
from torch.nn.attention import SDPBackend, sdpa_kernel

def modern_attention(q, k, v, is_causal=True, dropout_p=0.0):
    """
    使用 PyTorch 2.0+ 的 scaled_dot_product_attention

    会自动选择 FlashAttention 2/3
    """

    # 自动选择后端（Flash、Memory-efficient、Math）
    output = F.scaled_dot_product_attention(
        q, k, v,
        attn_mask=None,
        dropout_p=dropout_p,
        is_causal=is_causal,  # 因果掩码
        scale=None  # 为 None 时使用 1/sqrt(d_head)
    )

    return output

# 强制指定某个后端
def attention_with_flash_backend(q, k, v):
    with sdpa_kernel(SDPBackend.FLASH_ATTENTION):
        return F.scaled_dot_product_attention(q, k, v, is_causal=True)

def attention_with_efficient_backend(q, k, v):
    with sdpa_kernel(SDPBackend.EFFICIENT_ATTENTION):
        return F.scaled_dot_product_attention(q, k, v, is_causal=True)


# FlashAttention 各版本特点
flash_versions = {
    "FlashAttention 1": {
        "paper": "arXiv:2205.14135",
        "key_innovation": "分块（Tiling）+ Online Softmax",
        "memory": "O(N)（无需存储注意力矩阵）",
        "speedup": "相比标准注意力提速 2-4x"
    },
    "FlashAttention 2": {
        "paper": "arXiv:2307.08691",
        "key_innovation": "工作划分优化，支持 FP16/BF16",
        "memory": "O(N)",
        "speedup": "相比标准注意力提速 5-9x（H100 上）"
    },
    "FlashAttention 3": {
        "paper": "arXiv:2407.08608",
        "key_innovation": "针对 H100 优化，支持 FP8，异步流水线",
        "memory": "O(N)",
        "speedup": "相比 FA2 提速 1.5-2x（H100 上）"
    },
}

for name, info in flash_versions.items():
    print(f"\n{name}")
    for k, v in info.items():
        print(f"  {k}: {v}")

import torch
import torch.distributed as dist

class TensorParallelLinear(torch.nn.Module):
    """
    Tensor Parallel 的 Linear 层
    列拆分方式（Column Parallel）
    """

    def __init__(
        self,
        in_features: int,
        out_features: int,
        world_size: int,
        rank: int
    ):
        super().__init__()
        self.world_size = world_size
        self.rank = rank

        # 每块 GPU 负责 out_features // world_size 个输出神经元
        self.local_out_features = out_features // world_size

        self.weight = torch.nn.Parameter(
            torch.randn(self.local_out_features, in_features) / (in_features ** 0.5)
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # 本地计算
        local_output = torch.nn.functional.linear(x, self.weight)

        # 用 All-gather 把所有 GPU 的输出合并起来
        # (在实际分布式环境中执行)
        # dist.all_gather(output_list, local_output)

        return local_output


def setup_tensor_parallel_llm(model_name: str, tp_size: int):
    """
    Tensor Parallel LLM 配置示例（vLLM 方式）

    vLLM 内部就是采用这种方式
    """

    from vllm import LLM, SamplingParams

    # vLLM 的 tensor_parallel_size 配置
    llm = LLM(
        model=model_name,
        tensor_parallel_size=tp_size,  # GPU 数量
        gpu_memory_utilization=0.9
    )

    return llm

from vllm import LLM, SamplingParams
from vllm.engine.arg_utils import AsyncEngineArgs
from vllm.engine.async_llm_engine import AsyncLLMEngine
import asyncio
import time

# vLLM 基本用法
def vllm_basic_usage():
    """vLLM 的基本用法"""

    llm = LLM(
        model="meta-llama/Llama-2-7b-hf",
        tensor_parallel_size=1,         # GPU 数量
        gpu_memory_utilization=0.90,    # GPU 内存使用率
        max_model_len=4096,             # 最大序列长度
        quantization=None,              # "awq", "gptq", "squeezellm"
        dtype="auto",                   # "float16", "bfloat16"
        max_num_seqs=256,               # 最大并发序列数
        enable_prefix_caching=True,     # 启用前缀缓存
        use_v2_block_manager=True,      # PagedAttention v2
    )

    sampling_params = SamplingParams(
        temperature=0.8,
        top_p=0.95,
        max_tokens=200,
        presence_penalty=0.0,
        frequency_penalty=0.0,
    )

    prompts = [
        "Explain quantum computing in simple terms",
        "What is the future of artificial intelligence?",
        "How does the human brain work?",
    ]

    # 批量推理
    outputs = llm.generate(prompts, sampling_params)

    for output in outputs:
        print(f"Prompt: {output.prompt[:50]}...")
        print(f"Output: {output.outputs[0].text[:100]}...")
        print(f"Tokens generated: {len(output.outputs[0].token_ids)}")
        print()

    return outputs


# vLLM 异步服务
async def vllm_async_server():
    """使用 vLLM 异步引擎"""

    engine_args = AsyncEngineArgs(
        model="meta-llama/Llama-2-7b-hf",
        tensor_parallel_size=1,
        gpu_memory_utilization=0.90,
        max_model_len=4096,
        enable_prefix_caching=True,
    )

    engine = AsyncLLMEngine.from_engine_args(engine_args)

    async def generate_stream(prompt: str, request_id: str):
        sampling_params = SamplingParams(
            temperature=0.8,
            max_tokens=200
        )

        full_text = ""
        async for output in engine.generate(prompt, sampling_params, request_id):
            if output.outputs:
                delta = output.outputs[0].text[len(full_text):]
                full_text = output.outputs[0].text

                if delta:
                    print(f"[{request_id}] {delta}", end="", flush=True)

            if output.finished:
                print(f"\n[{request_id}] 完成")

    # 同时处理多个请求
    await asyncio.gather(
        generate_stream("What is AI?", "req_1"),
        generate_stream("Explain machine learning", "req_2"),
        generate_stream("What is deep learning?", "req_3"),
    )

import subprocess
import json
import time
import requests

def benchmark_vllm_server(
    model: str,
    num_requests: int = 100,
    max_tokens: int = 100,
    concurrency: int = 10
):
    """vLLM 服务器基准测试"""

    results = {
        "total_requests": num_requests,
        "concurrency": concurrency,
        "latencies": [],
        "ttfts": [],
        "throughputs": []
    }

    import asyncio
    import aiohttp

    async def send_request(session, prompt, request_id):
        start = time.perf_counter()
        first_token_time = None

        payload = {
            "model": model,
            "prompt": prompt,
            "max_tokens": max_tokens,
            "stream": True,
            "temperature": 0.0
        }

        async with session.post(
            "http://localhost:8000/v1/completions",
            json=payload
        ) as response:
            async for line in response.content:
                if line.startswith(b"data: "):
                    data = line[6:].decode()
                    if data.strip() == "[DONE]":
                        break

                    if first_token_time is None:
                        first_token_time = time.perf_counter() - start

        end = time.perf_counter()
        return {
            "latency": end - start,
            "ttft": first_token_time,
        }

    async def run_benchmark():
        prompts = [
            f"Tell me about topic number {i}." for i in range(num_requests)
        ]

        start = time.perf_counter()

        async with aiohttp.ClientSession() as session:
            # 并发请求
            tasks = []
            for i, prompt in enumerate(prompts):
                if len(tasks) >= concurrency:
                    done, tasks = await asyncio.wait(
                        tasks, return_when=asyncio.FIRST_COMPLETED
                    )
                    for task in done:
                        result = await task
                        results["latencies"].append(result["latency"])
                        if result["ttft"]:
                            results["ttfts"].append(result["ttft"])

                tasks.add(asyncio.ensure_future(
                    send_request(session, prompt, i)
                ))

            # 处理剩余任务
            for coro in asyncio.as_completed(tasks):
                result = await coro
                results["latencies"].append(result["latency"])

        total_time = time.perf_counter() - start
        total_tokens = num_requests * max_tokens
        results["throughput"] = total_tokens / total_time

    asyncio.run(run_benchmark())

    # 计算统计量
    import statistics
    latencies = results["latencies"]

    return {
        "avg_latency_ms": statistics.mean(latencies) * 1000,
        "p50_latency_ms": statistics.median(latencies) * 1000,
        "p99_latency_ms": sorted(latencies)[int(len(latencies) * 0.99)] * 1000,
        "avg_ttft_ms": statistics.mean(results["ttfts"]) * 1000 if results["ttfts"] else 0,
        "throughput_tps": results.get("throughput", 0),
    }

from vllm import LLM, SamplingParams

def demonstrate_prefix_caching():
    """演示前缀缓存的效果"""

    llm = LLM(
        model="meta-llama/Llama-2-7b-hf",
        enable_prefix_caching=True,  # 启用前缀缓存
        max_model_len=4096,
    )

    # 较长的系统提示词（所有请求共用）
    system_prompt = """You are a helpful AI assistant with expertise in:
    - Python programming and software development
    - Machine learning and deep learning
    - Data science and statistics
    - Cloud computing and DevOps
    [... 长系统提示词 ...]""" * 10  # 1000+ token

    questions = [
        "How do I optimize a Python loop?",
        "What is gradient descent?",
        "Explain containerization.",
        "What is a neural network?",
    ]

    sampling_params = SamplingParams(temperature=0.7, max_tokens=100)

    # 第一批：没有缓存（冷启动）
    import time

    cold_prompts = [f"{system_prompt}\n\nQuestion: {q}" for q in questions]

    cold_start = time.time()
    llm.generate(cold_prompts, sampling_params)
    cold_time = time.time() - cold_start

    # 第二批：相同的系统提示词（缓存命中！）
    warm_start = time.time()
    llm.generate(cold_prompts, sampling_params)
    warm_time = time.time() - warm_start

    print(f"第一次（无缓存）: {cold_time:.2f} 秒")
    print(f"第二次（缓存命中）: {warm_time:.2f} 秒")
    print(f"速度提升: {cold_time / warm_time:.2f}x")


def radix_tree_prefix_cache():
    """基于 Radix Tree 的前缀缓存实现"""

    class RadixNode:
        def __init__(self):
            self.children: dict = {}
            self.kv_cache_block_id: int = None

    class RadixTreeCache:
        """用 Radix Tree 管理 token 序列，从而共享公共前缀的 KV 缓存"""

        def __init__(self):
            self.root = RadixNode()
            self.cache_hits = 0
            self.cache_misses = 0

        def insert(self, token_ids: list, block_id: int):
            """插入 token 序列及其对应的 KV Cache 块 ID"""
            node = self.root
            for token_id in token_ids:
                if token_id not in node.children:
                    node.children[token_id] = RadixNode()
                node = node.children[token_id]
            node.kv_cache_block_id = block_id

        def lookup(self, token_ids: list) -> tuple:
            """查找给定 token 序列的最长匹配前缀"""
            node = self.root
            matched_len = 0
            last_block_id = None

            for i, token_id in enumerate(token_ids):
                if token_id in node.children:
                    node = node.children[token_id]
                    matched_len = i + 1
                    if node.kv_cache_block_id is not None:
                        last_block_id = node.kv_cache_block_id
                else:
                    break

            if last_block_id is not None:
                self.cache_hits += 1
            else:
                self.cache_misses += 1

            return matched_len, last_block_id

        def get_hit_rate(self) -> float:
            total = self.cache_hits + self.cache_misses
            return self.cache_hits / total if total > 0 else 0.0

    return RadixTreeCache()

class LLMOptimizationChecklist:
    """LLM 推理优化清单"""

    optimizations = [
        {
            "category": "基础设置",
            "level": 1,
            "items": [
                {
                    "name": "使用 FP16/BF16",
                    "impact": "高",
                    "effort": "低",
                    "description": "从 FP32 换成 FP16，内存节省 2 倍，速度提升",
                    "code": """
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.bfloat16,  # 或 float16
    device_map="auto"
)"""
                },
                {
                    "name": "启用 Flash Attention 2",
                    "impact": "高",
                    "effort": "低",
                    "description": "注意力计算提速 2-4x，同时节省内存",
                    "code": """
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    attn_implementation="flash_attention_2",
    torch_dtype=torch.bfloat16
)"""
                },
            ]
        },
        {
            "category": "KV Cache 优化",
            "level": 2,
            "items": [
                {
                    "name": "选择 GQA/MQA 模型",
                    "impact": "高",
                    "effort": "中",
                    "description": "KV Cache 缩减 4-8 倍，可处理更大的批量"
                },
                {
                    "name": "前缀缓存",
                    "impact": "中",
                    "effort": "低",
                    "description": "复用公共系统提示词的 KV Cache"
                },
            ]
        },
        {
            "category": "批处理优化",
            "level": 3,
            "items": [
                {
                    "name": "Continuous Batching (vLLM)",
                    "impact": "极高",
                    "effort": "低",
                    "description": "吞吐量提升 2-5x",
                    "code": """
from vllm import LLM, SamplingParams

llm = LLM(
    model=model_name,
    gpu_memory_utilization=0.90,
    enable_prefix_caching=True,
)"""
                },
            ]
        },
        {
            "category": "模型优化",
            "level": 4,
            "items": [
                {
                    "name": "AWQ 4-bit 量化",
                    "impact": "高",
                    "effort": "中",
                    "description": "内存减少 4x，速度提升 1.5-2x",
                    "code": """
from awq import AutoAWQForCausalLM

model = AutoAWQForCausalLM.from_quantized(
    "model-awq-4bit",
    fuse_layers=True
)"""
                },
                {
                    "name": "Speculative Decoding",
                    "impact": "中",
                    "effort": "高",
                    "description": "速度提升 2-3x（需要合适的草稿模型）"
                },
            ]
        },
        {
            "category": "硬件优化",
            "level": 5,
            "items": [
                {
                    "name": "Tensor Parallelism",
                    "impact": "极高",
                    "effort": "中",
                    "description": "通过多 GPU 实现近乎线性的吞吐量提升"
                },
                {
                    "name": "CUDA 图捕获",
                    "impact": "中",
                    "effort": "高",
                    "description": "消除内核启动开销"
                },
            ]
        }
    ]

    @classmethod
    def print_checklist(cls):
        print("=" * 70)
        print("LLM 推理优化分阶段清单")
        print("=" * 70)

        for category in cls.optimizations:
            print(f"\n[级别 {category['level']}] {category['category']}")
            print("-" * 50)

            for item in category['items']:
                impact_emoji = {"极高": "★★★", "高": "★★", "中": "★", "低": "☆"}
                print(f"  ✓ {item['name']}")
                print(f"    效果: {impact_emoji.get(item['impact'], '?')} {item['impact']}")
                print(f"    说明: {item['description']}")

        print("\n推荐的优化顺序:")
        print("1. 切换到 BF16/FP16（立即见效，零成本）")
        print("2. Flash Attention 2（立即见效，只需安装包）")
        print("3. 用 vLLM 提供服务（最大化吞吐量）")
        print("4. AWQ/GPTQ 4-bit 量化（内存节省 4 倍）")
        print("5. Speculative Decoding（改善延迟）")
        print("6. 多 GPU Tensor Parallelism（扩展规模）")

LLMOptimizationChecklist.print_checklist()