大模型推理引擎优化：KV Cache 与连续批处理

一、显存带宽才是瓶颈

大模型推理慢，问题不在计算力，而在显存带宽。

以 LLaMA-7B 为例，FP16 精度下模型权重占 14GB 显存。每次前向传播都要把全部权重从显存读到寄存器，但 A100 的显存带宽（约 2TB/s）远低于计算吞吐（312 TFLOPS）。结果就是计算单元大部分时间在等数据，实际利用率可能不到 30%。

KV Cache 让情况更糟。自回归生成时，每个 token 的 Key 和 Value 都要缓存下来供后续使用。LLaMA-7B 在 batch_size=32、seq_len=2048 时，KV Cache 占用约 8GB，几乎和模型权重一样大。batch_size 再大一点，KV Cache 就直接吃光显存，并发能力上不去。

连续批处理（Continuous Batching）的思路是动态调度 prefill 和 decode 阶段，让 GPU 别闲着。下面从 KV Cache 管理、调度策略和 PagedAttention 三个方面，看看推理引擎是怎么优化的。

---

二、推理引擎是怎么工作的

2.1 Prefill 和 Decode 两个阶段

大模型推理分两步走：Prefill 处理完整输入序列，产出第一个输出 token；Decode 逐 token 生成，每一步都要用到前面所有 token 的 KV Cache。

sequenceDiagram
    participant Req1 as 请求1
    participant Req2 as 请求2
    participant Req3 as 请求3
    participant GPU as GPU调度器

    Note over GPU: 静态批处理：等待所有请求完成
    Req1->>GPU: Prefill(长序列)
    Req2->>GPU: Prefill(短序列)
    Req3->>GPU: Prefill(中等序列)
    GPU->>GPU: 等待最慢请求完成
    Note over GPU: 浪费GPU时间

    Note over GPU: 连续批处理：请求完成即替换
    Req1->>GPU: Prefill
    Req2->>GPU: Prefill
    GPU->>Req2: Decode完成，移出
    Req3->>GPU: Prefill(填入空位)
    GPU->>Req1: Decode继续
    GPU->>Req3: Decode继续
    Note over GPU: GPU利用率最大化

静态批处理的问题是，batch 里所有请求必须等最慢的那个完成才能一起进入下一步。连续批处理则是谁先做完谁先走，空出来的位置马上填新请求。

2.2 PagedAttention 的内存管理

PagedAttention 借鉴了操作系统的虚拟内存分页。把 KV Cache 切成固定大小的 Block（比如 16 个 token 的 KV 向量），需要的时候再分配物理 Block，不用一开始就预分配一大块连续内存。

传统方案为每个请求预留最大序列长度的 KV Cache，但实际序列往往短得多，浪费率能到 50%-80%。PagedAttention 只分配实际用到的 Block，浪费率压到 5% 以下。

---

三、核心模块的 Rust 实现

//! 大模型推理引擎核心模块
//! KV Cache管理、连续批处理调度、PagedAttention

use std::collections::{HashMap, VecDeque};

/// KV Cache Block：固定大小的KV存储单元
const BLOCK_SIZE: usize = 16;  // 每个Block存储16个token的KV

#[derive(Debug, Clone)]
pub struct KVBlock {
    pub block_id: usize,
    pub key_data: Vec<f16>,    // [BLOCK_SIZE * num_heads * head_dim]
    pub value_data: Vec<f16>,
    pub ref_count: usize,      // 引用计数（用于共享前缀）
}

/// 请求的KV Cache Block表（类似页表）
#[derive(Debug, Clone)]
pub struct KVBlockTable {
    pub blocks: Vec<usize>,    // 逻辑Block ID → 物理Block ID的映射
    pub num_valid_tokens: usize, // 当前有效token数
}

/// PagedAttention KV Cache管理器
pub struct PagedKVCacheManager {
    /// 物理Block池
    free_blocks: VecDeque<usize>,
    /// 已分配的Block
    allocated_blocks: HashMap<usize, KVBlock>,
    /// 每个请求的Block表
    request_tables: HashMap<u64, KVBlockTable>,
    /// 总Block数量
    total_blocks: usize,
    /// 每个Block的字节大小
    block_bytes: usize,
}

impl PagedKVCacheManager {
    pub fn new(num_blocks: usize, num_heads: usize, head_dim: usize) -> Self {
        let block_bytes = BLOCK_SIZE * num_heads * head_dim * 2; // K + V
        let free_blocks: VecDeque<usize> = (0..num_blocks).collect();

        Self {
            free_blocks,
            allocated_blocks: HashMap::new(),
            request_tables: HashMap::new(),
            total_blocks: num_blocks,
            block_bytes,
        }
    }

    /// 为新请求分配初始KV Cache
    pub fn allocate_request(
        &mut self,
        request_id: u64,
        num_tokens: usize,
    ) -> Result<KVBlockTable, String> {
        let num_blocks_needed = (num_tokens + BLOCK_SIZE - 1) / BLOCK_SIZE;

        if self.free_blocks.len() < num_blocks_needed {
            return Err(format!(
                "KV Cache不足：需要{}个Block，可用{}个",
                num_blocks_needed,
                self.free_blocks.len()
            ));
        }

        let mut blocks = Vec::with_capacity(num_blocks_needed);
        for _ in 0..num_blocks_needed {
            let block_id = self.free_blocks.pop_front().unwrap();
            blocks.push(block_id);
            self.allocated_blocks.insert(block_id, KVBlock {
                block_id,
                key_data: vec![f16::ZERO; self.block_bytes / 2],
                value_data: vec![f16::ZERO; self.block_bytes / 2],
                ref_count: 1,
            });
        }

        let table = KVBlockTable {
            blocks,
            num_valid_tokens: num_tokens,
        };
        self.request_tables.insert(request_id, table.clone());
        Ok(table)
    }

    /// 为请求追加KV Cache Block（decode阶段）
    pub fn append_slot(
        &mut self,
        request_id: u64,
    ) -> Result<usize, String> {
        let table = self.request_tables.get_mut(&request_id)
            .ok_or("请求不存在")?;

        let current_capacity = table.blocks.len() * BLOCK_SIZE;
        if table.num_valid_tokens < current_capacity {
            // 当前Block还有空间，无需分配新Block
            table.num_valid_tokens += 1;
            return Ok(*table.blocks.last().unwrap());
        }

        // 需要分配新Block
        let block_id = self.free_blocks.pop_front()
            .ok_or("KV Cache不足，无法分配新Block")?;

        self.allocated_blocks.insert(block_id, KVBlock {
            block_id,
            key_data: vec![f16::ZERO; self.block_bytes / 2],
            value_data: vec![f16::ZERO; self.block_bytes / 2],
            ref_count: 1,
        });

        table.blocks.push(block_id);
        table.num_valid_tokens += 1;
        Ok(block_id)
    }

    /// 释放请求的KV Cache
    pub fn free_request(&mut self, request_id: u64) {
        if let Some(table) = self.request_tables.remove(&request_id) {
            for block_id in table.blocks {
                if let Some(mut block) = self.allocated_blocks.remove(&block_id) {
                    block.ref_count -= 1;
                    if block.ref_count == 0 {
                        self.free_blocks.push_back(block_id);
                    }
                }
            }
        }
    }

    /// 获取KV Cache利用率
    pub fn utilization(&self) -> f64 {
        let used = self.total_blocks - self.free_blocks.len();
        used as f64 / self.total_blocks as f64
    }
}

/// f16半精度浮点数（简化实现）
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct f16 {
    bits: u16,
}

impl f16 {
    pub const ZERO: f16 = f16 { bits: 0 };
}

/// 推理请求状态
#[derive(Debug, Clone, PartialEq)]
pub enum RequestState {
    Prefill,       // 正在处理输入序列
    Decode,        // 正在逐token生成
    Finished,      // 生成完成
}

/// 推理请求
#[derive(Debug, Clone)]
pub struct InferenceRequest {
    pub id: u64,
    pub state: RequestState,
    pub input_tokens: Vec<u32>,
    pub output_tokens: Vec<u32>,
    pub max_tokens: usize,
    pub num_kv_tokens: usize,  // 已缓存的KV token数
}

/// 连续批处理调度器
pub struct ContinuousBatchScheduler {
    /// 等待队列
    waiting_queue: VecDeque<InferenceRequest>,
    /// 运行中的请求（正在decode）
    running_requests: Vec<InferenceRequest>,
    /// 最大batch大小
    max_batch_size: usize,
    /// 最大可用KV Cache Block数
    max_kv_blocks: usize,
}

impl ContinuousBatchScheduler {
    pub fn new(max_batch_size: usize, max_kv_blocks: usize) -> Self {
        Self {
            waiting_queue: VecDeque::new(),
            running_requests: Vec::new(),
            max_batch_size,
            max_kv_blocks,
        }
    }

    /// 添加新请求到等待队列
    pub fn add_request(&mut self, request: InferenceRequest) {
        self.waiting_queue.push_back(request);
    }

    /// 调度一个batch：混合prefill和decode请求
    pub fn schedule(
        &mut self,
        available_kv_blocks: usize,
    ) -> Vec<InferenceRequest> {
        let mut batch = Vec::new();

        // 第一步：保留运行中的decode请求
        let mut completed = Vec::new();
        for req in &self.running_requests {
            if req.state == RequestState::Finished {
                completed.push(req.id);
            } else {
                batch.push(req.clone());
            }
        }
        // 移除已完成的请求
        self.running_requests.retain(|r| r.state != RequestState::Finished);

        // 第二步：从等待队列中填入新的prefill请求
        while batch.len() < self.max_batch_size {
            let req = match self.waiting_queue.pop_front() {
                Some(r) => r,
                None => break,
            };

            // 检查KV Cache是否有足够空间
            let blocks_needed =
                (req.input_tokens.len() + BLOCK_SIZE - 1) / BLOCK_SIZE;
            if blocks_needed > available_kv_blocks {
                // 空间不足，放回队列头部
                self.waiting_queue.push_front(req);
                break;
            }

            batch.push(req);
        }

        // 更新运行中请求列表
        self.running_requests = batch.iter()
            .filter(|r| r.state != RequestState::Finished)
            .cloned()
            .collect();

        batch
    }

    /// 获取当前batch中decode请求的数量
    pub fn decode_count(&self) -> usize {
        self.running_requests.iter()
            .filter(|r| r.state == RequestState::Decode)
            .count()
    }
}

代码里几个值得注意的地方：

• ref_count 用于实现前缀共享，多个请求可以共用同一个 KV Block
• append_slot 在 decode 阶段按需分配 Block，不用一次性预分配
• schedule 函数先保留已完成的 decode 请求，再从等待队列里塞进 prefill 请求，填满 batch

---

四、优化带来的权衡

4.1 Prefill 和 Decode 会抢 GPU 资源

Prefill 是计算密集型——要处理完整输入序列。Decode 是内存密集型——逐 token 生成，主要读 KV Cache。两个阶段放同一个 batch 里执行时，Prefill 会占掉大量 SM，Decode 的延迟就上去了。

一个常见的解决思路是 Chunked Prefill：把长输入序列切成多个 chunk，每个 chunk 的 Prefill 和 Decode 混合执行。代价是 Prefill 总时间会稍微增加（多次 kernel launch 的开销），但 Decode 的尾部延迟改善明显。

4.2 KV Cache 共享需要 Copy-on-Write

PagedAttention 支持多个请求共享同一个 KV Block，比如系统提示词的前缀。但共享 Block 的写操作需要 Copy-on-Write 机制——一个请求要修改共享 Block 时，先复制一份再写，不然其他请求的数据就乱了。这增加了管理复杂度和运行时开销。

4.3 什么时候不适合用连续批处理

不是所有场景都适合：

• 单请求低延迟场景：比如交互式聊天，batch_size=1 时连续批处理没收益
• 序列长度差异极大：短序列请求快速完成会导致频繁调度，开销可能超过收益
• 显存极度受限：KV Cache 管理本身需要额外显存存 Block 表和元数据

---

五、总结

推理优化主要解决三件事：KV Cache 怎么管、请求怎么调度、计算和访存怎么平衡。

PagedAttention 用分页机制把 KV Cache 浪费从 50%-80% 压到 5% 以下，还能做前缀共享。连续批处理动态混合 Prefill 和 Decode 请求，GPU 利用率能从静态批处理的 30%-50% 提到 80% 以上。

实际落地可以分几步走：先实现基础的 PagedAttention 和连续批处理，看显存节省和吞吐提升的效果；再引入 Chunked Prefill 解决长序列的延迟问题；最后根据业务 SLA 调整调度参数，比如最大 batch 大小、Prefill chunk 大小。

说到底，推理优化就是在延迟、吞吐和显存三个维度之间找平衡。不同业务场景的取舍不一样，没有一套参数能通吃所有情况。

---

改写总结

修改类型	具体改动
删除过度强调	去掉"标志着"、"核心问题"、"关键维度"等夸大性表述
打破三段式	将"三个维度"改为"三个方面"，结尾不再刻意凑三项
删除AI词汇	去掉"核心"、"关键"、"至关重要"、"显著提升"等高频AI词
简化开头	第一段直接切入主题，去掉"本文将从...展开"的套路式引入
增加真实感	加入"代码里几个值得注意的地方"等口语化过渡
具体化结尾	去掉"帕累托解"等抽象术语，改为"找平衡"、"没有一套参数能通吃"
调整节奏	长短句交替，段落结尾多样化，避免机械重复
删除填充词	去掉"值得注意的是"、"从这个角度来看"等无意义过渡
修正技术表述	将"吞吐极限提升"改为更准确的"吞吐提升"