04-TensorRT-LLM 运行时引擎：动态批处理、KV 缓存与 CUDA Graph

4 minute read

Published: July 06, 2026

背景与目标

在现代大语言模型（LLM）推理服务中，请求的到达时间、序列长度、解码步数高度动态。传统的“一请求一批次”模式导致 GPU 计算资源利用率低下，因为解码阶段每步只产生单个 token，计算密度极低。TensorRT-LLM 的运行时引擎围绕连续批处理（continuous batching）、分页 KV 缓存（paged KV cache）和 CUDA Graph 重放三大支柱，将请求调度、内存管理、kernel 发射深度耦合，在毫秒级延迟约束下最大化吞吐。

运行时引擎的核心抽象包括：

调度器（Scheduler）：决定哪些请求进入下一个 micro-batch，平衡内存压力与计算负载。
KV 缓存管理器（KVCacheManager）：以固定大小 block 为粒度管理物理存储，支持前缀复用、滑动窗口临时页以及多 tier 卸载。
注意力后端：消费 block table 在 GPU 上计算分页注意力，配合 CUDA Graph 减少 kernel launch 开销。
投机解码（Speculative Decoding）：与主执行循环深度集成，利用 draft model 生成多 token 候选并通过树形结构验证。

以下将从物理块管理、调度策略、注意力 kernel 寻址以及 CUDA Graph 的动态重放展开，揭示整个运行时引擎的底层实现。

核心设计与实现

1. 分页 KV 缓存：块表映射与零拷贝传输

TensorRT-LLM 中，每个序列的 KV cache 被分割为固定大小的块（block），默认每块容纳 tokens_per_block 个 token。逻辑块到物理存储的映射通过块表（block table）完成，这是一个形状为 [num_pools, capacity, 2, max_blocks_per_seq] 的张量，其中 2 代表每层注意力所需的 K/V 两个池，max_blocks_per_seq 是单个序列可占用的最大块数。

KVCacheManagerV2 在 Python 层维护一份 host-pinned 的块表 host_kv_cache_block_offsets，其中每个请求的块表由 IndexMapper 分配。每次迭代，只有活跃请求的块表条目被拷贝到 GPU 端的 staging tensor，使用的是一段自定义 CUDA kernel：

# tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py: L1550-L1570
def copy_batch_block_offsets(self, dst_tensor, request_ids, ...):
    copy_batch_block_offsets_to_device(
        self.host_kv_cache_block_offsets,
        dst_tensor,
        copy_idx,
        self.index_scales,
        self.kv_offset,
        self._stream.cuda_stream
    )

copy_batch_block_offsets_to_device 是 C++ kernel 的 Python 封装，它在 GPU 上并行完成以下操作：

根据 copy_idx 从 host-pinned 源中选取活跃请求的块号。
应用 index_scales 对块号缩放（适配不同池的大小）。
加上 kv_offset 得到最终物理偏移，写入目标张量。

由于源内存是 host-pinned（页锁定），PCIe 传输速度可接近带宽上限，且 kernel 避免了每个块号单独 cudaMemcpy 的调用开销。这种设计将块表更新视为一种“批处理归约”，与调度器的选批逻辑自然衔接。

2. 前缀复用与 radix tree 去重

为了减少重复 prompt 前缀的 KV 计算与存储，KVCacheManagerV2 在创建 context 时，会先将 prompt token 序列与多模态内容哈希组合成 TokenIdExt，然后在一棵并发 radix tree 中查找前缀。匹配到的块直接用于新请求，无需复制 K/V 数据，因为物理块可以多对一映射。具体做法是：

# tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py: L410-L450
def _augment_tokens_for_block_reuse(self, tokens, req, start=0, end=None):
    ...
    result[offset:offset+length] = gen_multimodal_cache_key_tokens(
        vocab_size, digest, length, token_offset=item_token_offset
    )
    ...

gen_multimodal_cache_key_tokens 使用 Blake3 哈希对图像等多模态内容生成伪 token，确保相同图片在不同请求中产生相同的 token pattern。
radix tree 的节点携带块的引用计数，当请求完成时引用计数递减，归零后块被回收。
调度器中的 _beneficial_to_skip 方法利用前缀分析，如果某 pending 请求的第一个新块恰好被一个正在执行 chunked prefill 的请求贡献，则跳过该请求，待前者释放块后再调度，避免重复分配。这本质是一种 生产者-消费者流水线优化：让 prefill chunk 成为块的“生产者”，后续消费请求延迟等待。

# tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py: _beneficial_to_skip
def _beneficial_to_skip(self, req, newly_contributed_context_blocks, ...):
    ...
    if summary.first_new_block in newly_contributed_context_blocks:
        return True
    ...

3. 注意力 kernel 中的块寻址逻辑

在 GPU 上，注意力 kernel 必须通过块表将逻辑 token 位置映射到物理内存地址。以 FMHA V2 kernel 为例，其寻址逻辑如下（简化）：

// cpp/kernels/fmha_v2/src/fused_multihead_attention.cpp: L250-L270
// PagedAttention: compute physical address
int logical_block_idx = token_idx / tokens_per_block;
int intra_block_offset = token_idx % tokens_per_block;
int physical_block = block_offsets[batch_idx * max_blocks_per_seq + logical_block_idx];
int physical_idx = physical_block * tokens_per_block + intra_block_offset;
K_ptr = K_pool + physical_idx * head_size * stride_k;

block_offsets 正是前述 copy_batch_block_offsets 拷贝到 GPU 的那份张量。Kernel 先通过逻辑块索引读取物理块号，再结合块内偏移算出最终 K/V 的线性地址。由于每个 token 的寻址都依赖 block_offsets 表，这部分内存访问无法合并，属于 随机读取 模式。为了缓解，tokens_per_block 通常设为 16~64，使得同一个 warp 内的线程大概率落在同一物理块内，减少表查找次数。

对于分组查询注意力（GQA/MQA），kernel 进一步设置 h_q_per_kv 步长，使得多个 Q 头共享同一 K、V，从而将 block table 查找的代价均摊。

优化技术细节

4. 滑动窗口临时页复用（SWA Scratch Reuse）

某些模型（如 Mistral）采用滑动窗口注意力，只有最近 W 个 token 的 KV 缓存需要保留。为节省内存，KVCacheManagerV2 不分配独立物理块，而是维护一组 scratch pages 在多个请求间临时共用。其核心是 _copy_swa_block_offsets_with_scratch_compiled 这个 torch.compile 优化的 CUDA 函数：

# tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py: L300-L350
@maybe_compile(options={"max-autotune": True})
def _copy_swa_block_offsets_with_scratch_compiled(
    block_offsets, copy_idx, pool_ids, scales, layer_offsets, scratch_pages,
    block_positions, scratch_begs, scratch_ends, scratch_slots, num_contexts, output
):
    ...
    converted = torch.where(
        scratch_mask_by_pool,
        scratch_index,
        converted
    )
    output.copy_(converted.permute(0, 2, 1, 3))

其原理是：在向 GPU 拷贝块表时，判断某逻辑块位置是否落在窗口的 scratch 范围内（由 block_positions 与 scratch_begs/ends 比较）。若是，则将物理块号替换为对应的 scratch page 编号，并使用 slots 索引分配。这样，同一物理 scratch page 可以被不同请求的不同窗口复用，而语义上互不干扰，因为滑动窗口外的 token 不会被注意力 kernel 读取。

此优化的代价是增加一个额外的条件赋值 kernel 启动，但它与块表拷贝融合在一个编译函数中，并通过 torch.where 实现分支消除，几乎没有额外开销。

5. 连续批处理调度：容量感知与微批打包

调度器分两级运作：

容量调度（GuaranteedNoEvict / MaxUtilization policy）：遍历 pending 请求，检查其所需的 KV cache 块数（考虑前缀复用后）是否可分配，同时评估是否会因 chunked prefill 触发串行化（_beneficial_to_skip）。通过后请求被加入 active_requests。
微批打包：将活跃请求按 max_batch_size 和 max_num_tokens 裁剪为 micro-batch。对于超长 context，按 EQUAL_PROGRESS 或 FCFS 策略分块，每个 chunk 获得 token budget，形成多个连续 micro-batch，逐步消费长序列。

代码片段展示了 GuaranteedNoEvict 策略的循环：

# tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py: L280-L330
def schedule(self, scheduler, active_requests):
    ...
    for req in pending_requests:
        # 检查是否有足够的 block
        if not scheduler.can_allocate(req):
            continue
        # 前缀感知跳过
        if scheduler._beneficial_to_skip(req, ...):
            continue
        active_requests.append(req)
    ...

这里的 can_allocate 并非简单计数，而是调用 KVCacheManager 的 probe_reuse 方法，在 radix tree 中实时查找可共享块，从而估计剩余所需新块数量。这种紧密耦合让调度器能根据实际内存碎片程度做出精确决策。

6. CUDA Graph 动态重放与块表更新

解码阶段每步 kernel 调用模式高度固定，但 batch 组成可能变化（请求完成或有新请求加入）。TensorRT-LLM 使用 CUDA Graph 捕获每种 batch size 的计算图，并在重放前更新 输入张量地址 和 块表内容。

由于 CUDA Graph 要求所有 kernel 参数在捕获时确定，动态的 block table 通过以下方式处理：

分配一块足够大的“块表缓冲张量”，其大小对应 max_batch_size。
在每次重放前，调用 copy_batch_block_offsets 将新 block table 拷贝到该缓冲的具体偏移。
Graph 内部的注意力 kernel 始终从此缓冲读取块表，从而无需重新捕获 Graph。

这种策略类似于 动态图注入：将易变的块表数据隔离为外部更新的缓冲区，Graph 本身保持静态。当某些请求完成，其块表条目可能被设为无效块号（如 -1），kernel 中通过掩码跳过，但 Graph 形状不变。

投机解码（speculative decoding）进一步复杂化了此流程：draft 模型与 target 模型的验证 kernel 也需要共享块表。TensorRT-LLM 采用 DecoderState 中的 cacheIndirection 张量实现树形草稿的 KV cache 重定向，使得同一物理块的 K/V 可被多个草稿路径安全共享，避免深度拷贝。

性能与权衡

上述设计将内存管理、调度和 kernel 发射紧密耦合，其性能优势显著：

内存利用率：分页机制避免了外部碎片，前缀复用可节约大量重复前缀的存储，滑动窗口临时页进一步削减峰值显存。相比静态分配，同样 GPU 可承载 2~3 倍并发请求。
计算-通信重叠：host-pinned 块表零拷贝传输与 GPU kernel 异步执行，确保 CPU 更新块表时 GPU 仍在计算上一轮。
调度开销：radix tree 查找复杂度 O(L)，但通过缓存 PrefixReuseSummary 和并发无锁结构，调度器遍历 pending 请求的总开销控制在几十微秒以内。
CUDA Graph 收益：将解码阶段数十个小 kernel 合并为单次 launch，大幅降低 CPU launch 延迟，对于 batch size 较小时提升尤为明显。

取舍：

分页带来了额外的块表随机访问开销，当 tokens_per_block 过小时，查表代价可能吃掉收益。工程上选 16/32/64 为常用值。
前缀复用依赖准确的哈希，Blake3 计算增加预处理延迟，但分摊到大量共享 prefix 的场景（如多轮对话、系统 prompt）是可接受的。
CUDA Graph 要求固定 shape，限制了动态性的上限，例如 batch size 变化须重新捕获；实践中常用多个预捕获 size 覆盖典型值。
滑动窗口临时页共享增加了 block table 更新 kernel 的复杂度，且 scratch pages 数量固定，超额订阅时可能引起部分请求暂停，需结合容量调度谨慎配置。

总结

TensorRT-LLM 的运行时引擎在 物理块映射、前缀感知调度、零拷贝块表传输 以及 CUDA Graph 注入式更新 之间构建了一条精密流水线。它将请求的生命周期分解为逻辑块的分配、映射与回收，并通过调度器提前规划内存需求，使得 GPU 核持续运转而无须空闲等待。无论是超长 context 的分块预填充，还是投机解码的树状 KV 共享，底层都依赖同一套 paged memory 机制与块表映射范式。理解这些机制有助于深度优化大模型推理服务，在有限的 GPU 资源下压榨出每一点吞吐。

Twitter Facebook LinkedIn

bug404

04-TensorRT-LLM 运行时引擎：动态批处理、KV 缓存与 CUDA Graph

背景与目标

核心设计与实现

1. 分页 KV 缓存：块表映射与零拷贝传输

2. 前缀复用与 radix tree 去重

3. 注意力 kernel 中的块寻址逻辑

优化技术细节

4. 滑动窗口临时页复用（SWA Scratch Reuse）

5. 连续批处理调度：容量感知与微批打包

6. CUDA Graph 动态重放与块表更新

性能与权衡

相关论文

总结

llama.cpp 高级特性：推测解码与语法约束生成

1. 背景与目标

llama.cpp 服务端揭秘：连续批处理与请求调度

背景与目标

llama.cpp KV Cache 与注意力机制：环形缓冲到 Flash Attention

1. 背景与目标

llama.cpp 量化技术深度解析：从 8-bit 到 1.5-bit 与 GGUF 格式

1. 背景与目标