大模型推理的连续批处理技术：吞吐量翻倍的工程艺术

请求A: 输入50 token, 生成200 token  (总时长: 3秒)
请求B: 输入30 token, 生成20 token   (总时长: 0.5秒)
请求C: 输入100 token, 生成500 token (总时长: 7秒)

Step 1: batch = [A, B, C]  → 各生成第1个token
Step 2: batch = [A, B, C]  → 各生成第2个token
Step 3: batch = [A, B, C]  → B完成！B退出batch
Step 4: batch = [A, D, C]  → D是新请求，加入batch
Step 5: batch = [A, D, C]  → A完成！A退出
Step 6: batch = [E, D, C]  → E加入
...

传统KV Cache管理:
  请求A: [████████████████████____] 预分配20KB，用了12KB
  请求B: [████____________________] 预分配20KB，用了4KB
  浪费: 24KB空闲但不可用

PagedAttention:
  物理内存: [A_page0][B_page0][A_page1][C_page0][A_page2][B_page1]
  逻辑映射: A → [page0, page2, page5]  B → [page1, page4]
  
  按需分配，零碎片！

class BlockTable:
    def __init__(self, block_size=16, num_blocks=1024):
        self.block_size = block_size
        self.free_blocks = list(range(num_blocks))
        self.block_table = {}  # request_id → [block_indices]
    
    def allocate(self, request_id, num_tokens):
        """为请求分配KV Cache块"""
        blocks_needed = (num_tokens + self.block_size - 1) // self.block_size
        if len(self.free_blocks) < blocks_needed:
            raise OOMError("KV Cache内存不足")
        
        allocated = []
        for _ in range(blocks_needed):
            block_idx = self.free_blocks.pop(0)
            allocated.append(block_idx)
        
        self.block_table[request_id] = allocated
    
    def free(self, request_id):
        """释放请求的KV Cache块"""
        for block_idx in self.block_table[request_id]:
            self.free_blocks.append(block_idx)
        del self.block_table[request_id]

请求A: [system prompt 200 tokens][user query A]
请求B: [system prompt 200 tokens][user query B]

物理内存: [shared_system_block_0..12][A_query_block][B_query_block]
A的block table: [0,1,2,...,12, 13]
B的block table: [0,1,2,...,12, 14]

只有当A或B需要修改共享块时，才复制到新块（COW）

class ContinuousBatchScheduler:
    def __init__(self, max_batch_size, max_tokens):
        self.running = []      # 正在运行的请求
        self.waiting = []      # 等待队列
        self.max_batch_size = max_batch_size
        self.max_tokens = max_tokens  # 最大token数（输入+输出）
    
    def schedule(self):
        """每次iteration前调度"""
        # 1. 移出已完成的请求
        self.running = [r for r in self.running if not r.is_finished()]
        
        # 2. 尝试从等待队列加入新请求
        while self.waiting:
            req = self.waiting[0]
            if self.can_admit(req):
                self.running.append(req)
                self.waiting.pop(0)
                self.block_table.allocate(req.id, req.input_len)
            else:
                break  # 内存不够，停止
        
        # 3. 可能需要preempt（抢占）
        if not self.has_kv_cache_space():
            self.preempt_lowest_priority()
        
        return self.running
    
    def can_admit(self, req):
        """检查是否有足够资源接收新请求"""
        current_tokens = sum(r.total_tokens for r in self.running)
        new_tokens = req.input_len + req.max_output_len
        return (current_tokens + new_tokens <= self.max_tokens and
                len(self.running) < self.max_batch_size)

内存不足时:
  1. 选择优先级最低的请求
  2. 将其KV Cache换出到CPU内存
  3. 标记为"preempted"
  4. 释放物理块给其他请求

内存恢复时:
  1. 将preempted请求的KV Cache换回GPU
  2. 重新加入running队列

传统: [Prefill A (长输入)] → [Decode A, Decode B, Decode C]
      prefill时B和C的decode被阻塞

Chunked: [Prefill_A_chunk1, Decode B, Decode C]
         [Prefill_A_chunk2, Decode B, Decode C]
         [Decode A, Decode B, Decode C]
         
         decode不被阻塞！