KV Cache = 2 × 层数 × 头数 × 维度 × 序列长度 × 字节数
         = 2 × 40 × 40 × 128 × 1000 × 2 bytes (FP16)
         = 819,200,000 bytes
         ≈ 0.78 GB / 序列

没有 KV Cache：
  Token 1: 计算 K₁, V₁ → 与 Q₁ 做 Attention
  Token 2: 计算 K₁, K₂, V₁, V₂ → 与 Q₂ 做 Attention  ← 重复计算了 K₁, V₁
  Token 3: 计算 K₁, K₂, K₃, V₁, V₂, V₃ → 与 Q₃ 做 Attention  ← 重复计算了 K₁, V₁, K₂, V₂

有 KV Cache：
  Token 1: 计算 K₁, V₁ → 存入 Cache → 与 Q₁ 做 Attention
  Token 2: 计算 K₂, V₂ → 追加到 Cache → 与 Q₂ 做 Attention（只读 Cache 中的 K₁, V₁）
  Token 3: 计算 K₃, V₃ → 追加到 Cache → 与 Q₃ 做 Attention（只读 Cache 中的 K₁, V₁, K₂, V₂）

KV Cache 大小（字节） = 2 × L × H × D × S × B

其中：
  L = 层数（Transformer 层数）
  H = 每层的注意力头数
  D = 每个头的维度
  S = 序列长度（tokens）
  B = 每个参数的字节数（FP16 = 2）

示例：Llama-2-13B
  L = 40, H = 40, D = 128, S = 4096, B = 2 (FP16)
  KV Cache = 2 × 40 × 40 × 128 × 4096 × 2 = 3,355,443,200 bytes ≈ 3.13 GB

  10 个并发请求：10 × 3.13 GB = 31.3 GB（已占 A100-80GB 的 39%）
  加上模型权重 13B × 2 = 26 GB，总计 57.3 GB

请求 A: 实际生成 50 tokens → 预分配 4096 × 3.13GB/4096 ≈ 3.13 GB → 利用率 1.2%
请求 B: 实际生成 200 tokens → 预分配 4096 × 3.13GB/4096 ≈ 3.13 GB → 利用率 4.9%
请求 C: 实际生成 1000 tokens → 预分配 4096 × 3.13GB/4096 ≈ 3.13 GB → 利用率 24.4%

操作系统                    vLLM
─────────────────────────────────────────────
物理内存         →    GPU 显存（HBM）
页（Page）       →   KV Block（16 tokens）
页表（Page Table）→   Block Table（逻辑块 → 物理块映射）
虚拟内存         →   请求的虚拟 KV Cache 空间
换页（Swap）     →   显存不足时 Block 可以换出到 CPU 内存

一个 Block 的大小（字节）= 2 × L × H × D × 16 × B

对于 Llama-2-13B：
  Block Size = 2 × 40 × 40 × 128 × 16 × 2 = 13,107,200 bytes ≈ 12.5 MB

# 伪代码：Block Table 的映射关系
# 逻辑块: 请求的 KV Cache 按 16 token 间隔切分
# 物理块: GPU 显存中的实际存储位置

class BlockTable:
    """
    每个请求维护一个 Block Table，将逻辑块映射到物理块
    
    例如：请求 A 生成了 50 个 token，需要 ceil(50/16) = 4 个逻辑块
    Block Table = [物理块#3, 物理块#7, 物理块#12, 物理块#1]
    """
    def __init__(self, block_size=16):
        self.block_size = block_size
        self.mapping: dict[int, int] = {}  # 逻辑块 → 物理块
    
    def append(self, physical_block_id: int):
        """追加一个物理块"""
        logical_id = len(self.mapping)
        self.mapping[logical_id] = physical_block_id
    
    def get_physical(self, logical_id: int) -> int:
        """逻辑块 → 物理块转换"""
        return self.mapping[logical_id]
    
    @property
    def num_blocks(self) -> int:
        return len(self.mapping)

# 传统方式：连续内存访问
def traditional_attention(query, kv_cache):
    # kv_cache 是连续的内存块
    keys = kv_cache[:sequence_length]     # 直接切片
    values = kv_cache[:sequence_length]
    return scaled_dot_product_attention(query, keys, values)

# PagedAttention 方式：通过 Block Table 间接访问
def paged_attention(query, key_blocks, value_blocks, block_table):
    """
    key_blocks: [num_physical_blocks, block_size, num_heads, head_dim]
    block_table: [num_logical_blocks] → 物理块 ID 列表
    
    关键：每个逻辑块通过 block_table 映射到物理块，
    物理块可以分散在显存的任意位置。GPU Kernel 需要
    根据 block_table 来 gather 分散的 KV 数据。
    """
    gathered_keys = []
    gathered_values = []
    
    for logical_id in range(len(block_table)):
        physical_id = block_table[logical_id]
        gathered_keys.append(key_blocks[physical_id])
        gathered_values.append(value_blocks[physical_id])
    
    # 拼接为连续张量后计算 Attention
    keys = torch.cat(gathered_keys, dim=0)
    values = torch.cat(gathered_values, dim=0)
    return scaled_dot_product_attention(query, keys, values)

class BlockAllocator:
    """
    GPU 显存的 Block 分配器
    
    设计决策：为什么用固定大小的 Block（16 tokens）？
    - 太小（如 4 tokens）：Block Table 条目过多，映射开销大
    - 太大（如 64 tokens）：内部碎片严重（生成了 17 tokens 需要 2 个 Block，浪费 15 tokens 空间）
    - 16 tokens 是 vLLM 团队通过实验确定的最佳平衡点：
      大多数请求的额外生成长度是 16 的倍数，碎片率 < 5%
    
    设计决策：为什么用简单的 free list 而非伙伴系统？
    所有 Block 大小相同，没有"合并相邻空闲块"的需求。
    简单的 free list 分配和释放都是 O(1)，比伙伴系统的 O(log n) 更快。
    """
    
    def __init__(self, num_blocks: int):
        self.num_blocks = num_blocks
        self.free_blocks: list[int] = list(range(num_blocks))  # 空闲物理块列表
        self.allocated: dict[int, int] = {}  # 物理块 → 引用计数
    
    def allocate(self) -> int:
        """分配一个物理块"""
        if not self.free_blocks:
            raise MemoryError("No free blocks available")
        block_id = self.free_blocks.pop()
        self.allocated[block_id] = 1  # 引用计数 = 1
        return block_id
    
    def free(self, block_id: int):
        """释放一个物理块"""
        if block_id not in self.allocated:
            return
        self.allocated[block_id] -= 1
        if self.allocated[block_id] <= 0:
            del self.allocated[block_id]
            self.free_blocks.append(block_id)
    
    def increment_ref(self, block_id: int):
        """增加引用计数（用于 Copy-on-Write）"""
        if block_id in self.allocated:
            self.allocated[block_id] += 1
    
    @property
    def free_count(self) -> int:
        return len(self.free_blocks)
    
    @property
    def utilization(self) -> float:
        """显存利用率"""
        if self.num_blocks == 0:
            return 0.0
        return 1.0 - (self.free_count / self.num_blocks)

class PrefixCache:
    """
    前缀缓存管理器
    
    工作原理：
    1. 第一个请求计算 System Prompt 的 KV Cache，存入物理块 [P0, P1, P2]
    2. 第二个请求需要同样的 System Prompt，直接引用物理块 [P0, P1, P2]
       → 引用计数 +1，不需要重新计算
    3. 当第二个请求在 System Prompt 之后生成新 token 时，
       需要修改最后一个共享块 P2（因为追加了新 token）
       → Copy-on-Write：复制 P2 → P2'，新 token 写入 P2'
       → 第一个请求仍然使用 P2，第二个请求使用 P2'
    
    引用计数机制：
    - 物理块 [P0, P1]: ref_count = 2（两个请求共享）
    - 物理块 [P2]: ref_count = 1（第一个请求独占）
    - 物理块 [P2']: ref_count = 1（第二个请求的副本）
    """
    
    def __init__(self):
        # 前缀哈希 → 物理块列表
        self.cache: dict[str, list[int]] = {}
        self.allocator: BlockAllocator = None
    
    def match_prefix(self, prompt_tokens: tuple) -> tuple[list[int], int]:
        """查找匹配的前缀缓存，返回 (物理块列表, 匹配的 token 数)"""
        # 从最长前缀开始匹配
        for prefix_len in range(len(prompt_tokens), 0, -1):
            prefix_hash = str(hash(prompt_tokens[:prefix_len]))
            if prefix_hash in self.cache:
                blocks = self.cache[prefix_hash]
                # 增加引用计数
                for block_id in blocks:
                    self.allocator.increment_ref(block_id)
                return blocks, prefix_len
        
        return [], 0
    
    def store(self, prompt_tokens: tuple, blocks: list[int]):
        """存储前缀缓存"""
        prefix_hash = str(hash(prompt_tokens))
        if prefix_hash not in self.cache:
            self.cache[prefix_hash] = blocks
            for block_id in blocks:
                self.allocator.increment_ref(block_id)

场景：10 个并发请求，共享 2000 tokens 的 System Prompt
模型：Llama-2-13B

传统方式：每个请求独立存储 System Prompt 的 KV Cache
  10 × (2000/4096) × 3.13 GB = 10 × 1.53 GB = 15.3 GB

PagedAttention 方式：共享前缀
  1 × 1.53 GB = 1.53 GB（唯一的 System Prompt KV Cache）
  + 10 × 各自生成部分的 KV Cache

节省：15.3 - 1.53 = 13.77 GB（节省 90%）

时间线（Static Batching）：

Batch 1: [请求A(100 tokens) | 请求B(50 tokens) | 请求C(200 tokens) | 请求D(30 tokens)]
         ↓ 所有请求必须在同一时刻完成 prefill 才能开始 decode
         ↓ 请求 D 在 30 tokens 时就结束了，但必须等待请求 C 生成完 200 tokens
         ↓ 请求 D 的 KV Cache 占着显存但不再使用
         
         时间 ─────────────────────────────────────────────────────────────────→
         请求A: ████████████████████████████████████████████████████████████████ (100 tokens)
         请求B: ██████████████████████████ (50 tokens, 完成后等待)
         请求C: ████████████████████████████████████████████████████████████████ (200 tokens)
         请求D: ████████████ (30 tokens, 完成后等待最久)
         
         GPU 利用率: 30% (请求 B 和 D 完成后，GPU 在等待 C 完成)

时间线（Continuous Batching）：

Step 1: [请求A(1/100) | 请求B(1/50) | 请求C(1/200) | 请求D(1/30)]
Step 2: [请求A(2/100) | 请求B(2/50) | 请求C(2/200) | 请求D(2/30)]
...
Step 30: [请求A(30/100) | 请求B(30/50) | 请求C(30/200) | 请求D(30/30)] ← D 完成
Step 31: [请求A(31/100) | 请求B(31/50) | 请求C(31/200) | 请求E(1/?) ]  ← 新请求 E 加入
...
Step 50: [请求A(50/100) | 请求B(50/50) | 请求C(50/200) | 请求E(20/?)]  ← B 完成
Step 51: [请求A(51/100) | 请求C(51/200) | 请求E(21/?)  | 请求F(1/?) ]  ← 新请求 F 加入
...
Step 100: [请求A(100/100) | 请求C(100/200) | 请求E(70/?) | 请求F(50/?)] ← A 完成
Step 101: [请求C(101/200) | 请求E(71/?)  | 请求F(51/?) | 请求G(1/?) ]  ← 新请求 G 加入

class ContinuousBatchingScheduler:
    """
    连续批处理调度器
    
    设计决策：为什么保持 Batch 大小恒定？
    如果 Batch 大小波动，GPU 的利用率会波动。
    理想情况下，每完成一个请求就立即补充一个新请求，保持 Batch 满载。
    
    设计决策：优先调度哪个请求？
    生产级 vLLM 使用 FCFS（先来先服务）等启发式策略，并配合抢占与换出机制。
    教学代码中简化为 Round-Robin，以说明公平调度的核心思想：
    1. Round-Robin 保证公平性，每个请求的等待时间有上界
    2. 在延迟敏感场景下，公平调度能避免低优先级请求饿死
    3. 实际部署时还会综合考虑请求长度、显存占用、SLO 等因素
    """
    
    def __init__(self, max_batch_size: int, block_allocator: BlockAllocator):
        self.max_batch_size = max_batch_size
        self.block_allocator = block_allocator
        self.running: list[Request] = []   # 正在处理的请求
        self.waiting: list[Request] = []   # 等待队列
        self.block_tables: dict[int, BlockTable] = {}  # request_id → BlockTable
    
    def add_request(self, request: Request):
        """添加新请求到等待队列"""
        self.waiting.append(request)
    
    def step(self) -> list[tuple[int, int]]:
        """
        执行一个推理 step，返回本 step 生成的 token
        
        返回值: [(request_id, token_id), ...]
        """
        # 1. 移除已完成的请求
        completed = [r for r in self.running if r.is_finished()]
        for req in completed:
            self.running.remove(req)
            # 释放 KV Cache 占用的物理块
            block_table = self.block_tables.pop(req.id, None)
            if block_table:
                for physical_id in block_table.mapping.values():
                    self.block_allocator.free(physical_id)
        
        # 2. 从等待队列中补充新请求（直到 Batch 满或显存不足）
        while len(self.running) < self.max_batch_size and self.waiting:
            next_req = self.waiting.pop(0)
            
            # 检查是否需要前缀缓存
            blocks, prefix_len = self._check_prefix_cache(next_req)
            
            # 检查显存是否足够
            remaining_tokens = next_req.max_tokens - prefix_len
            needed_blocks = (remaining_tokens + 15) // 16  # 向上取整
            
            if self.block_allocator.free_count >= needed_blocks:
                # 分配物理块
                block_table = self._allocate_blocks(next_req, needed_blocks, blocks)
                self.block_tables[next_req.id] = block_table
                self.running.append(next_req)
            else:
                # 显存不足，放回等待队列
                self.waiting.insert(0, next_req)
                break
        
        # 3. 执行推理（伪代码）
        if not self.running:
            return []
        
        # 构建 Batch 输入
        batch_inputs = self._build_batch_inputs()
        
        # 调用模型推理（使用 PagedAttention）
        next_tokens = self._model_forward(batch_inputs)
        
        # 4. 更新每个请求的状态
        results = []
        for i, req in enumerate(self.running):
            token_id = next_tokens[i]
            req.append_token(token_id)
            results.append((req.id, token_id))
            
            # 如果当前逻辑块满了，分配新物理块
            block_table = self.block_tables[req.id]
            if req.generated_tokens % 16 == 0 and req.generated_tokens > 0:
                try:
                    new_block = self.block_allocator.allocate()
                    block_table.append(new_block)
                except MemoryError:
                    # 显存不足 → 请求被抢占（preemption）
                    # 将 KV Cache 换出到 CPU 内存
                    self._preempt_request(req)
        
        return results
    
    def _allocate_blocks(self, request: Request, num_blocks: int, 
                         shared_blocks: list[int] = None) -> BlockTable:
        """为请求分配物理块"""
        bt = BlockTable()
        
        # 先添加共享的前缀块
        if shared_blocks:
            for block_id in shared_blocks:
                bt.append(block_id)
        
        # 分配新物理块
        for _ in range(num_blocks):
            block_id = self.block_allocator.allocate()
            bt.append(block_id)
        
        return bt
    
    def _preempt_request(self, request: Request):
        """
        请求抢占：将 KV Cache 换出到 CPU 内存
        
        设计决策：当显存不足时，为什么选择"换出"而非"拒绝"？
        拒绝新请求 → 用户体验差（请求失败）
        换出旧请求 → 用户体验降级（延迟增加），但请求不会失败
        
        换出策略：选择已生成最多 token 的请求（最不可能很快完成的长请求）
        理由：长请求占据大量显存，且换出损失的 CPU-GPU 传输时间
        相对于其总生成时间占比小。
        """
        # 选择要换出的请求（生成 token 最多的）
        victim = max(self.running, key=lambda r: r.generated_tokens)
        
        # 将其 KV Cache 从 GPU 显存复制到 CPU 内存
        bt = self.block_tables[victim.id]
        cpu_blocks = []
        for physical_id in bt.mapping.values():
            # 伪代码：cudaMemcpy D2H
            cpu_block = copy_to_cpu(physical_id)  
            cpu_blocks.append(cpu_block)
            self.block_allocator.free(physical_id)
        
        # 保存到请求对象
        victim.swapped_blocks = cpu_blocks
        victim.is_swapped = True
        self.running.remove(victim)

// PagedAttention CUDA Kernel 简化版
// 每个 thread block 处理一个请求的一个注意力头

__global__ void paged_attention_kernel(
    float* output,              // [num_requests, num_heads, head_dim]
    const float* query,         // [num_requests, num_heads, head_dim]
    const float* key_cache,     // [num_blocks, block_size, num_heads, head_dim]
    const float* value_cache,   // [num_blocks, block_size, num_heads, head_dim]
    const int* block_tables,    // [num_requests, max_num_blocks_per_req]
    const int* block_table_lens,// [num_requests] 每个请求的逻辑块数
    const int block_size,       // 16
    const float scale           // 1/sqrt(head_dim)
) {
    int req_id = blockIdx.x;    // 每个请求一个 thread block
    int head_id = blockIdx.y;   // 每个注意力头独立处理
    
    int num_blocks = block_table_lens[req_id];
    int seq_len = num_blocks * block_size;
    
    // 共享内存：用于存储中间结果
    extern __shared__ float shared_mem[];
    float* qk_scores = shared_mem;  // [seq_len]
    
    // 步骤 1：计算 Q * K^T（逐 block 计算）
    for (int block_idx = 0; block_idx < num_blocks; block_idx++) {
        int physical_block = block_tables[req_id * MAX_BLOCKS + block_idx];
        
        // 读取该 block 的 Key 向量
        // 每个线程处理 block 中的几个 token
        for (int t = threadIdx.x; t < block_size; t += blockDim.x) {
            float dot = 0.0f;
            for (int d = 0; d < HEAD_DIM; d++) {
                dot += query[req_id * NUM_HEADS * HEAD_DIM + 
                            head_id * HEAD_DIM + d] *
                       key_cache[physical_block * BLOCK_SIZE * NUM_HEADS * HEAD_DIM +
                                t * NUM_HEADS * HEAD_DIM +
                                head_id * HEAD_DIM + d];
            }
            qk_scores[block_idx * block_size + t] = dot * scale;
        }
    }
    __syncthreads();
    
    // 步骤 2：Softmax
    float max_score = -INFINITY;
    for (int i = 0; i < seq_len; i++) {
        max_score = fmaxf(max_score, qk_scores[i]);
    }
    
    float sum_exp = 0.0f;
    for (int i = 0; i < seq_len; i++) {
        qk_scores[i] = expf(qk_scores[i] - max_score);
        sum_exp += qk_scores[i];
    }
    for (int i = 0; i < seq_len; i++) {
        qk_scores[i] /= sum_exp;
    }
    __syncthreads();
    
    // 步骤 3：Attention * V（逐 block 计算）
    float result[HEAD_DIM] = {0.0f};
    for (int block_idx = 0; block_idx < num_blocks; block_idx++) {
        int physical_block = block_tables[req_id * MAX_BLOCKS + block_idx];
        
        for (int t = 0; t < block_size; t++) {
            float weight = qk_scores[block_idx * block_size + t];
            for (int d = 0; d < HEAD_DIM; d++) {
                result[d] += weight * 
                    value_cache[physical_block * BLOCK_SIZE * NUM_HEADS * HEAD_DIM +
                               t * NUM_HEADS * HEAD_DIM +
                               head_id * HEAD_DIM + d];
            }
        }
    }
    
    // 写入输出
    for (int d = 0; d < HEAD_DIM; d++) {
        output[req_id * NUM_HEADS * HEAD_DIM + 
               head_id * HEAD_DIM + d] = result[d];
    }
}