TensorRT-LLM 源码完全拆解（十五）：显存与缓冲管理——BufferManager 如何把 KV 池、workspace 与 CUDA Graph 缓冲装进同一份预算

来源一：用户给的四个旋钮

按优先级由强到弱，初始化时可以影响预算的配置项有四类：

四者的关系可以用一句话概括：前两个是”我最多能花多少”（硬顶），第三个是”我愿意花剩下的百分之多少”（软顶），第四个决定”必须留下多少给别人”（预留）。 最终预算 = min(软顶, 硬顶) − 预留，其中的换算全部发生在 _util.py 的估算流程里。

注意这里有个容易踩的坑：max_gpu_total_bytes 在 KVCacheManagerV2 下是唯一的容量控制手段，max_tokens 退化成用户侧的一个约束提示；而在 V1 下，容量仍然由 max_tokens 控制。源码注释把这个分歧写得很直白：

1
# NOTE:
2
# For KVCacheManager, KvCacheCreator currently controls capacity using two parameters in KVCacheConfig:
3
#   • max_tokens
4
#   • max_gpu_total_bytes
5
# For KVCacheManagerV2, KvCacheCreator controls capacity using max_gpu_total_bytes only.
6
# This leaves max_tokens as a user-defined constraint.

读到这里应该形成的直觉是：“配置里那个数字”未必就是”实际决定预算的数字”。要判断谁在说话，必须先确定用的是哪一版管理器。

来源二：profiling 实测的峰值

软顶与硬顶只给出”可用总额”，真正决定”KV 池能拿走多少”的是一次实测。configure_kv_cache_capacity() 的流程是：

1. 模型权重加载完成，torch.cuda.empty_cache() 清掉缓存分配器的空洞，记录当前已用量；
2. 用一组哑请求（dummy requests）真的跑一遍前向，让所有临时缓冲真实地分配出来；
3. 读 torch.cuda.memory_stats()["allocated_bytes.all.peak"] 拿到 torch 侧峰值，再用 torch.cuda.mem_get_info() 拿设备级总量，两者相减得到 torch 之外的占用（NCCL、CUDA Graph）；
4. 减去这些，剩下的才是 KV 池能拿的钱。

第 3 步里那两个数字的差被单独记了一笔，名字叫 extra_cost：

1
activation_bytes = torch_peak_memory - model_bytes
2
extra_cost = max(total_used_bytes - torch_used_bytes, 0)
3
peak_memory = torch_peak_memory + extra_cost

activation_bytes 是”torch 分配器内部、前向过程中涨上去的部分”，也就是激活与 workspace；extra_cost 是”torch 分配器之外、前向过程中涨上去的部分”，注释直接点明它的典型来源是 NCCL 和 CUDA graphs。把这两类分开记账是有意义的：前者可以被缓存分配器复用、可以在前向结束后缩回去；后者（尤其是 CUDA Graph 池）一旦捕获就常驻，而且它的峰值未必出现在这次 profiling 的那一步。

这就引出了下一个问题：profiling 只能测到”它跑的那一步”，而那一步不一定是最费显存的一步。

来源三：显式预留（reserve）

正因为 profiling 会低估，TensorRT-LLM 在测完之后追加了一串显式预留。这些预留是本文的核心内容之一，因为它们正好就是”KV 池之外的显存”的清单。

源码里出现的预留共有五类，各自独立加成到 peak_memory 或从 KV 预算里减掉：

前两类直接加在峰值上（意味着它们挤占的是”KV 池之外”的钱），后两类直接从 KV 预算里扣（意味着它们挤占的是 KV 池自己的钱）。这个区别不是随意的——它反映了这两组 workspace 的分配时机：前者在 KV 池建立之前就已经存在（或必然存在），后者是在 KV 池建立之后、前向过程中才出现，所以只能从 KV 池已经拿走的额度里再借回来。

下面逐类看。这一段是理解”预算到底怎么算”的关键，值得慢读。

1
extra_seats = max_num_sequences - num_sequences_without_headroom
2
if extra_seats <= 0 or not vocab_size:
3
    return 0
4
counts_bytes = extra_seats * max_beam_width * vocab_size * torch.int32.itemsize
5
presence_prefix_bytes = extra_seats * vocab_size * torch.bool.itemsize
6
return counts_bytes + presence_prefix_bytes

1
@classmethod
2
def runtime_workspace_bytes_per_token(cls, model_config, mapping) -> int:
3
    """Per-token bytes to reserve for a workspace this backend stages whose size scales with a
4
    runtime quantity the KV-cache estimator does not drive to its serving maximum while profiling
5
    (e.g. ``total_kv_len``, inflated by KV-cache reuse). Default ``0`` -- correct for every backend
6
    that does not stage such context-MLA data.
7

8
    A non-zero rate is reserved from the KV budget by the estimator, and the scheduler caps the
9
    driving sum at what that reserve covers, so the declared buffer stays within its reservation
10
    (buffers the backend does not declare are still unaccounted for). Keep the rate identical to the
11
    runtime allocation's per-token cost or use a documented conservative upper bound.
12
    """
13

14
@classmethod
15
def runtime_workspace_is_chunked_prefill_bounded(cls, model_config) -> bool:
16
    """Whether chunked prefill bounds the declared workspace by the current chunk.
17

18
    Backends that still stage data proportional to the complete attended
19
    prefix must return ``False`` so KV-cache sizing reserves headroom for
20
    that prefix.
21
    """
22
    return True
23

24
@classmethod
25
def runtime_workspace_fixed_bytes(cls, model_config, mapping, *,
26
                                  max_batch_size, max_num_tokens, max_seq_len,
27
                                  enable_cuda_graph) -> int:
28
    """Capacity-sized workspace not bounded by summed attended KV length.
29

30
    The estimator reserves the full capacity alongside the profiled peak,
31
    covering arenas first allocated after the peak prefill step.
32
    """
33
    return 0

1
@classmethod
2
def runtime_workspace_bytes_per_token(cls, model_config, mapping) -> int:
3
    """Context-MLA workspace sized by summed attended KV length.
4

5
    Dense fp8 context-MLA stages expanded K/V. NVFP4 sparse MLA stages
6
    selected latent rows in fp8 plus selection/scan workspace. Cached
7
    tokens can grow both beyond the fresh-prefill profiling floor.
8
    """
9
    config = model_config.pretrained_config
10
    if not is_mla(config):
11
        return 0
12
    quant_config = model_config.quant_config
13
    quant_mode = getattr(quant_config, "quant_mode", None)
14
    sparse_algorithm = getattr(model_config.sparse_attention_config, "algorithm", None)
15
    nvfp4_dsa_context = (quant_mode is not None
16
                         and getattr(quant_mode, "has_fp4_kv_cache", lambda: False)()
17
                         and sparse_algorithm == "dsa"
18
                         and get_sm_version() >= 100)
19
    if nvfp4_dsa_context:
20
        # The output row has one fp8 value per latent element. The gather
21
        # also allocates three int32 arrays, CUB scan storage, and aligned
22
        # padding; 64 B/token is a conservative bound for that auxiliary
23
        # workspace.
24
        nvfp4_gather_aux_bytes_per_token = 64
25
        return int(config.kv_lora_rank + config.qk_rope_head_dim
26
                   + nvfp4_gather_aux_bytes_per_token)
27
    ...
28
    fp8_context_mla = (quant_config is not None
29
                       and quant_mode.has_fp8_kv_cache()
30
                       and get_sm_version() in (90, 100, 103, 107, 120))
31
    if not fp8_context_mla:
32
        return 0
33
    # Attention-DP runs the full head set per rank; otherwise heads shard across TP
34
    # (mirror mNumAttnHeads).
35
    attn_tp = 1 if mapping.enable_attention_dp else mapping.tp_size
36
    num_attn_heads = config.num_attention_heads // attn_tp
37
    # The buffer is skipped only where AttentionOp::useSparseMLA() holds, which needs all three of:
38
    #   * DSA / DeepSeek-V4 -- only these lower to the absorption path that reads K/V from the paged
39
    #     cache. Skip-softmax passes no sparse indices to C++, and its ignore-list can exclude a layer,
40
    #     so those layers still run dense MLA. The workspace is shared, so one dense layer forces the
41
    #     reserve.
42
    #   * SM 100 / 103 / 107 -- mUseTllmGen is `sm >= 100 && sm != 120 && sm != 121`.
43
    #   * short-seq MHA fallback off -- it routes short contexts back through the dense path.
44
    # Match the runtime predicate, not just "a sparse config exists": over-reserving costs KV pool,
45
    # under-reserving OOMs mid-forward.
46
    sparse_mla = (sparse_algorithm in ("dsa", "deepseek_v4")
47
                  and get_sm_version() in (100, 103, 107))
48
    ...

1
worst_case = min(max_batch_size, max_num_tokens) * max_seq_len
2
override = kv_cache_config.fp8_context_mla_kv_len_cap
3
if override is None:
4
    return worst_case
5
return min(max(int(override), max_seq_len), worst_case)

1
def get_mla_context_workspace_reserve(budget_bytes, k_bytes_per_token,
2
                                      w_bytes_per_token, kv_len_cap):
3
    """Reserve ``w * kv_len_cap`` (the worst-case summed attended KV), clamped to the per-token split
4
    ``budget * w / (k + w)`` so a memory-constrained node shares the budget at a common token count rather
5
    than starving the KV pool. The admission cap is ``reserve / w == min(kv_len_cap, budget / (k + w))``
6
    tokens; ..."""
7
    if not (budget_bytes > 0 and k_bytes_per_token > 0 and w_bytes_per_token > 0
8
            and kv_len_cap and kv_len_cap > 0):
9
        return 0, None
10
    reserve = min(
11
        w_bytes_per_token * kv_len_cap,
12
        budget_bytes * w_bytes_per_token / (k_bytes_per_token + w_bytes_per_token))
13
    return reserve, int(reserve / w_bytes_per_token)

来源三（续）：CacheCost——把每 token 成本写成仿射函数

_util.py 里有一个看起来不起眼、但把整套算术串起来的数据结构：

1
class CacheCost:
2
    """Affine KV cache budget: ``bytes = slope * tokens + intercept``.
3

4
    The legacy proportional case is just ``intercept = 0``.
5
    ...
6
    """
7
    slope: int
8
    intercept: int = 0
9
    bytes_per_request: Optional[int] = None
10
    quota_per_request: Optional[int] = None
11

12
    def tokens_for_budget(self, budget: int) -> int:
13
        """Memory budget -> max tokens. Clamps a negative result to 0."""
14
        if self.slope <= 0:
15
            return 0
16
        tokens = max((budget - self.intercept) // self.slope, 0)
17
        return tokens
18

19
    def bytes_for_tokens(self, tokens: int) -> int:
20
        """Token count -> memory bytes."""
21
        return self.slope * tokens + self.intercept

为什么需要 intercept？ 因为并非所有 KV 显存都随 token 数线性增长。举两个真实例子：

• Mamba / 线性注意力混合模型：除了按 token 增长的注意力 KV，还有一份循环状态，它是”每 batch 一个状态块”的固定开销，与序列长度无关。这份开销就是 intercept。
• MLA 的压缩缓存：每层的 latent 池之外，量化版本还要一份 block scale 池（形状相同、最后一维是 head_size / 16）。它同样按 token 增长，仍然算 slope，但乘的常数与主池不同。

再看两个可选字段的分工（文档字符串已经写明）：

• bytes_per_request：最坏情况下一条请求在 intercept 之外钉住的字节数，用来细化”target 与 draft 两个管理器之间如何分摊预算”；
• quota_per_request：分配器对单请求坚持要的配额（get_cache_quota_per_request），它给 draft 份额设下下界与上界。

而 __add__ 里那句注释回答了另一个问题——为什么相加时要丢掉这两个字段：

二者都不是仿射成本模型的一部分，所以算术（__add__）会丢掉它们。

也就是说，目标模型 + 草稿模型的合并成本只累加 slope 与 intercept；每请求的字节配额是”分配策略的元信息”，不是”容量公式的输入”。把这两类信息放进同一个结构、却让它们在算术上有不同命运，是一个很克制的设计——它避免了把”账”和”策略”混在一起。

CacheCost 还有两处工程细节值得记：

其一，TP/CP 缩放只改 slope，不改 intercept。 _get_draft_cache_cost() 里：

1
def scaled(cost: CacheCost) -> CacheCost:
2
    # Rescale only the slope: every other cost component (e.g. the
3
    # fixed intercept) is per-request rank-local bytes and stays
4
    # unscaled. dataclasses.replace keeps any component the rescale
5
    # does not touch, where rebuilding the CacheCost from named
6
    # fields would silently drop it.
7
    return dataclasses.replace(cost, slope=cost.slope * helix_cp_scale)

注释解释了为什么：intercept 是”每请求、本 rank 局部”的字节数（例如循环状态按 batch 分配，每个 rank 各自一份），不随 CP 切分而缩小；而 slope 是”每 token”的，在 helix CP 下每个 rank 只存每 cp_size 个页中的一个，所以要按 cp_size 放大成”per rank-LOCAL target token”的口径。

其二，用 dataclasses.replace 而不是重建。 注释直接点名了风险：重建会静默丢掉缩放没触及的字段（比如 bytes_per_request）。这类”复制对象时漏字段”的 bug 在成本模型里特别隐蔽——它不会报错，只会让预算少算一点，然后在某个长序列上 OOM。看到成本模型里有 dataclasses.replace，就应该确认它是在有意保留未触及字段。

三个来源合起来：一次完整的预算推导

把上面三段串起来，一次初始化里的预算推导顺序是：

1. 权重加载，empty_cache()，记录 model_bytes（torch 侧）与 total_used_bytes - model_bytes（torch 之外）；
2. profiling 前向：用哑请求跑一步，记录 torch_peak_memory、torch_used_bytes，并取 KV 管理器的 get_kv_cache_stats().allocated_bytes（估计阶段的临时 KV 也要加回来）；
3. 计算 activation_bytes = torch_peak_memory - model_bytes 与 extra_cost = total_used_bytes - torch_used_bytes；peak_memory = torch_peak_memory + extra_cost；
4. 加显式预留：多模态编码器、容量定尺 workspace、惩罚表 workspace；
5. 算候选预算：available_kv_mem = (total_gpu_memory - peak_memory + allocated_bytes) * fraction；
6. 按后端声明再从预算里扣：DeepSeek-V4 的 Indexer MQA logits workspace、context-MLA workspace reserve（同时得出准入上限 _fp8_ctx_mla_kv_len_cap）；
7. 与用户的硬顶取 min：max_tokens 换算成字节、以及 max_gpu_total_bytes，两者都只减不加；
8. 最终的 kv_cache_max_memory 写回 kv_cache_config.max_gpu_total_bytes，并被换算成 V1 需要的 max_tokens。

图：显存预算的切分在几何上就是”两条线性需求在总额约束下取交点”——这正是 CacheCost 与 get_mla_context_workspace_reserve 里那个 budget * w / (k + w) 分式的图像含义。roofline 的横轴是算术强度、纵轴是可达吞吐，而显存预算的版本横轴是 token 数、纵轴是字节数；两者判断”瓶颈在哪一侧”的方式是一样的。

第 5 步那个 (total_gpu_memory - peak_memory + allocated_bytes) * fraction 的形式值得停一下：为什么要加回 allocated_bytes？ _cal_max_memory 的注释回答了：

1
"""Calculate the max KV cache capacity.
2

3
NOTE: `allocated_bytes` is the total KV-cache memory that must be pre-allocated during the estimation
4
phase (for both the main and draft models) so the estimation run can complete successfully. When
5
computing `available_kv_mem`, add this amount back in.
6
"""

读法：peak_memory 里已经包含了为跑通 profiling 而临时分配的那份 KV 缓存，所以从”设备总量减峰值”得到的剩余量实际上被这份临时 KV 压低了一次。而这份临时 KV 是估计阶段的一次性产物，线上并不存在。所以要把它加回来，否则最终预算会系统性地偏小——小多少？恰好等于”为喂饱哑请求而必须准备的 KV 容量”，而这个量在长序列模型上可能是几十 GB 量级。这是一个典型的”测量工具自身占用被算进被测量”的问题，解法是在公式里显式补偿。

最后，profiling 阶段对 pool_ratio 与 avg_seq_len 的临时改写也值得一看：

1
# Estimation uses inferred pool sizing; the final manager uses the
2
# user-provided configuration.
3
self._kv_cache_config.pool_ratio = self._pool_ratio_in
4
self._kv_cache_config.avg_seq_len = self._avg_seq_len_in

而在 try_prepare_estimation() 里它们被临时改写：

1
# User-provided pool sizing can underprovision the temporary
2
# estimation cache and cause warmup to hang or fail. Override it
3
# for estimation, then restore it in configure_kv_cache_capacity().
4
self._kv_cache_config.pool_ratio = None
5
self._kv_cache_config.avg_seq_len = self._max_seq_len

估计期的池配置必须”就宽”、运行期的池配置必须”就用用户给的”。原因是估计期只是要跑通一次前向，池子配小了会让哑请求排不进去，warmup 直接挂住或失败；而运行期要尊重用户对池比例的意图（因为 VSWA 等场景下池比例决定各窗口怎么分预算）。同一组参数在两个阶段需要有不同取值，是”用真实执行来测量”这种方案固有的代价。

估算在哪些配置下会被关掉

还有一个反面清单不能漏：try_prepare_estimation() 会在四类配置下关闭估算，直接跳过 profiling：

关掉之后走的是”显式配额路径”。以 HELIX CP 为例，_configure_helix_kv_cache_capacity() 的注释说明了这里的口径差异：

显式配额直接透传；否则按比例定尺会设置 max_gpu_total_bytes（一个 rank-local 字节上限，管理器原样消费）。在这里设置 max_tokens 会超出比例目标：管理器会把这个旋钮按 1 / max_util_for_resume 放大。

这段注释的信息量很大：同一份”比例”，在两条路径下会被换算成不同的量，因此不能简单地在两条路径之间搬用同一个数字。更细的一点是 rank-local 与 global 的区别——注释里提到这个字节上限”管理器会在各 rank 之间取 min 同步、再换算成 global tokens”，也就是说各 rank 的空闲显存不均衡时，最终容量由最紧的那个 rank 决定。（跨 rank 同步这件事在 TensorRT-LLM 源码完全拆解（四） 里已经讲过它的记账侧后果，这里补的是”它为什么必须发生在字节层而不是 token 层”。）

路径一：虚拟内存分配器（如果启用了）

1
BufferManager::IBufferPtr BufferManager::gpu(std::size_t size, tensorrt_llm::DataType type) const
2
{
3
    if (auto vmAllocator = getVirtualMemoryAllocator())
4
    {
5
        return std::make_unique<VirtualAddressDeviceBuffer>(size, type, std::move(vmAllocator));
6
    }
7
    if (static_cast<bool>(mPool))
8
    {
9
        return std::make_unique<DeviceBuffer>(size, type, CudaAllocatorAsync{mStream, mPool});
10
    }
11
    // When memory pools are not supported, fallback to synchronous memory allocations.
12
    return gpuSync(size, type);
13
}

三条路径的优先级是：虚拟内存 > 内存池 > 同步 cudaMalloc。

虚拟内存路径基于 CUDA 的 VMM（Virtual Memory Management）API：先用 cuMemCreate 申请物理句柄，再用 cuMemAddressReserve 保留一段虚拟地址，最后 cuMemMap 把两者挂起来。访问器本身是线程安全的、可栈式切换的：

1
CudaVirtualMemoryAllocator getVirtualMemoryAllocator()
2
{
3
    std::shared_lock lock(sConfMutex);
4
    return CudaVirtualMemoryAllocator{sCurrentConf};
5
}
6

7
void pushVirtualMemoryAllocator(
8
    std::string const& tag, CudaVirtualMemoryAllocator::RestoreMode mode, std::shared_ptr<CudaStream> backStream)
9
{
10
    std::unique_lock lock(sConfMutex);
11
    sCurrentConf.swap(
12
        sConfStack.emplace_back(std::make_shared<AllocConf>(getVirtualMemoryManager(), tag, mode, backStream)));
13
}
14

15
void popVirtualMemoryAllocator() { ... }

为什么需要一个可 push/pop 的配置栈？ 因为”是否走虚拟内存”是一个进程级全局开关，但不同调用方希望在不同代码区域里得到不同行为（例如某个组件需要在它自己的作用域内拿到 VMM 分配，作用域结束恢复默认）。用一个受读写锁保护的栈，就得到了”作用域内的分配策略重载”——这是 C++ 里做全局策略局部覆盖的常见手法，代价是所有分配都要过一次锁。

顺带说明一句：VMM 这条路径与 TensorRT-LLM 源码完全拆解（四） 里讲过的”两级池、块句柄”是两个层面的事——那里讲的是”KV 页在池子里的逻辑记账”，这里讲的是”池子本身的物理内存是从哪儿来的”。同一套代码里既有 cudaMallocAsync 池，也可能有 VMM 映射，取决于配置。

路径二：primary 内存池 + 流序分配

1
BufferManager::BufferManager(CudaStreamPtr stream, bool trimPool)
2
    : mStream{std::move(stream)}
3
    , mTrimPool{trimPool}
4
{
5
    TLLM_CHECK_WITH_INFO(static_cast<bool>(mStream), "Undefined CUDA stream");
6
    mPool = CudaMemPool::getPrimaryPoolForDevice(mStream->getDevice());
7
}

它拿到的是设备的 “primary” 池，而不是自建一个。 primary 池是 CUDA 为每个设备提供的默认 cudaMemPool_t，被进程内所有使用者（包括 PyTorch 的缓存分配器）共享。这一点非常重要，因为它决定了后面那条”裁剪池”的操作到底在裁谁的内存：

1
//! \brief Hints the driver to trim the pool. Does not guarantee that the amount of reserved memory will actually
2
//! decrease, only guarantees that this amount after trimming will be larger than the provided size.
3
void memoryPoolTrimTo(std::size_t size);

注释里的两句限定词值得逐字读：“只是提示”（hints） 与 “只保证裁剪后保留量仍大于给定 size”。也就是说 memoryPoolTrimTo(0) 的正确理解不是”释放所有内存”，而是”告诉驱动：请尽量把保留量降到 0 附近，但我不保证一定成功，且不会低于你给的数”。而 ~BufferManager() 里那个 if (mTrimPool) memoryPoolTrimTo(0) 的含义是：析构一个”带裁剪语义”的 BufferManager，等价于向驱动提示”这块内存我不要了”。

CudaMemPool 还提供了三个不同步的统计量：

1
/// @brief Gets the amount of reserved memory in the memory pool stream, WITHOUT synchronizing.
2
[[nodiscard]] std::size_t memoryPoolReserved() const;
3

4
/// @brief Gets the amount of used memory in the memory pool, WITHOUT synchronizing.
5
[[nodiscard]] std::size_t memoryPoolUsed() const;
6

7
/// @brief Gets the amount of free memory in the memory pool, WITHOUT synchronizing.
8
[[nodiscard]] std::size_t memoryPoolFree() const;

reserved / used / free 三者构成池子的账本。在流序分配下，used 是”已分配但可能还没被释放点归还”的量，free 是”池子里可复用的量”。这三个数不需要同步是它们能出现在热路径上的前提——如果需要 cudaStreamSynchronize 才能读，就没法在调度循环里用了。

还有一个能力探测函数，它解释了整条回退链的边界条件：

1
/// @brief Returns a value indicating whether memory pools are supported on the device.
2
/// @details Memory pools depend on the presence of the UVM driver. On some systems, the UVM driver is explicitly
3
/// disabled.
4
static bool supportsMemoryPool(int deviceId);

内存池依赖 UVM 驱动的存在，而在某些系统上 UVM 驱动被显式禁用。 所以 gpu() 里那句 if (static_cast<bool>(mPool)) 会失败，落到路径三。在部分虚拟化或受限容器环境里，这就是”为什么显存明明够却报分配失败”的一个常见根因。

路径三：同步分配与其它内存类型

1
BufferManager::IBufferPtr BufferManager::gpuSync(std::size_t size, tensorrt_llm::DataType type)
2
{
3
    if (auto vmAllocator = getVirtualMemoryAllocator())
4
    {
5
        return std::make_unique<VirtualAddressDeviceBuffer>(size, type, std::move(vmAllocator));
6
    }
7
    return std::make_unique<StaticDeviceBuffer>(size, type, CudaAllocator{});
8
}

注意 gpuSync 也先查 VMM：如果启用了虚拟内存分配器，即使是”同步”入口也会走 VMM，只是它不再依赖 stream 的顺序语义。这也说明 VMM 路径在这套代码里被视为更优的默认。

BufferManager 提供的内存种类覆盖了完整的存储层次，这一点从工厂方法的名单就能看出来：

ipcNvls 一项与 TensorRT-LLM 源码完全拆解（十一）：并行拓扑与通信 里讲过的对称内存是同一件事的分配入口：它分配的是”一组 rank 都能以相同地址访问”的张量，因而必须由通信层与分配层共同保证地址一致。

从 Python 侧调用：PeftCacheManager 的例子

BufferManager 不只是 C++ 内部对象，Python 侧可以通过绑定直接构造它。PeftCacheManager 是最清楚的一个例子：

1
BufferManager = tensorrt_llm.bindings.internal.runtime.BufferManager
2
buffer_manager_stream = execution_stream.cuda_stream if execution_stream is not None else torch.cuda.current_stream(
3
).cuda_stream
4
buffer_manager = BufferManager(buffer_manager_stream, True)
5
logger.info(
6
    f"[PeftCacheManager] buffer_manager_stream: {buffer_manager_stream}"
7
)
8
self.impl = PeftCacheManagerCpp(config=peft_cache_manager_config,
9
                                model_config=model_config,
10
                                world_config=world_config,
11
                                buffer_manager=buffer_manager)

三个细节：

其一，第二个参数是 True，也就是启用裁剪语义。 LoRA adapter 缓存的分配是一次性、大块、可能长时间闲置的（一个 adapter 可能几小时不被用到），所以”析构时提示驱动归还”是有价值的。

其二，日志里显式打印了 buffer_manager_stream。 这不是无意义的调试信息——流错了会让流序释放点排到错误的流上，从而出现”内存明明释放了却仍被占用”或”提前归还导致踩内存”。把流值打进日志，是为了在出问题时能立刻确认这一条。

其三，PeftCacheManager 是”资源管理器”家族的成员。 这与 KV 管理器一样：它持有一份容量受限的显存（adapter 缓存），所以调度器必须把它作为准入条件之一。这一点在 TensorRT-LLM 源码完全拆解（二）：PyExecutor 装配线与单步闭环 里已经出现过——那里列出的准入条件里就有 peft_cache_manager.impl。这里补上了”它的显存从哪来”。

预算之外：设备侧账本的下一层

到这里为止的账都发生在设备侧。但线上系统里，KV 与权重可能有一部分并不住在 GPU 上。BufferManager 这一层（以及 pinned / pinnedPool 这些内存类型）在架构上为”往主机搬”留好了通道：host 侧的 KV 池（次级池）、前缀缓存卸载、adapter 预取都走这条路径。

图：把部分 KV 缓存放到 host 内存，会同时抬高前缀缓存命中率与作业级吞吐，但代价是”取回”要走 PCIe/CXL——也就是把显存预算问题转换成了带宽预算问题。这张图说明”显存不够”往往不是终点，而是”在哪一层存”的重新分配起点。

站内 KV Cache Offloading 完全拆解 与 SGLang 源码完全拆解（十六）：HiCache——三级 KV 缓存、分层前缀树与异步搬运 分别从框架无关视角与 SGLang 实现视角讲过这条路线；TensorRT-LLM 侧的对应物是次级池与 KV 连接器，属于设备预算之外的第二本账。看到”预算不够”时，先分清是要”多算一点预算”还是”换一层存”——这两条路的代价模型完全不同（前者是分配失败，后者是带宽与延迟）。

三个执行期钩子的语义

源码里对钩子有一段极简的注释，把每个钩子的时机讲清楚了：

1
def prepare_resources(self, scheduled_batch: ScheduledRequests):
2
    """确保资源管理器在 forward 之前按批分配资源。"""
3
def update_resources(self, scheduled_batch: ScheduledRequests):
4
    """在 forward 之后更新资源管理器的状态。"""
5
def free_resources(self, request: LlmRequest):
6
    """请求完成时释放资源。"""

而 KVCacheManager 对这三个钩子的实现（在上一篇里讲过细节）恰好演示了三种典型的”资源管理”问题：

prepare_resources 里的 fit_token_budget 是在”事后认账”。 调度器准入时按估计的可复用 token 折价计价，真实复用数要到 addSequence 才算出来；如果实际比估计少，前向就会多物化 token，撞上 total_num_tokens <= max_num_tokens 断言（GitHub #13318）。补救的办法是按量修剪批尾的 chunk，把被剪的请求变回分块状态。它解释了资源管理里一个反复出现的模式：

在”估算-实测”裂口面前，框架选择”延后计算”而不是”回滚分配”。 回滚分配意味着要把已经分配的块还回去、把已经加入的序列摘出来，会破坏注意力 DP 的投票与在途集合；延后计算只动”这一轮还要算多少 token”，不影响任何已建立的所有权关系。

update_resources 在 KV 管理器里是空实现，但注释留了扩展点：

1
def update_resources(self, scheduled_batch: ScheduledRequests):
2
    """KV 不需要更新：此函数当前不做操作。若 Python 侧支持 KV 复用，Radix 树管理在这里更新。"""

一个”故意空着但有明确用途”的钩子，比一个”删掉的钩子”更有信息量——它说明这套接口的设计者已经预见到”Python 侧实现一等公民复用管理器”这件事，并且提前把插入点留好了。

free_resources 有一个 pin_on_release 参数。 默认 False 表示正常回收；置 True 表示”钉住不回收”。它存在的理由是：有些场景下请求结束不代表这块 KV 没用了——例如它可能要被转移给另一个实例（P/D 分离场景），或者要被登记进复用树留给后来的请求。“释放”与”可复用”是两回事，这个参数就是把两者分开的开关。

资源管理器与显存预算的接口

资源管理器报容量给调度器，用的是 get_max_resource_count() 与 get_needed_resource_to_completion()。这一对的口径必须一致才有意义：如果前者按”池子总块数”报、后者按”到完成的净增量”报，比较它们就等于在比两个不同单位。

这套代码里最诚实的一处注释就出现在这里（出自 Python 侧）：

C++ 的 get_remaining_blocks_to_completion 与 Python 版本的语义不一致（C++ 排除已分配块、Python 包含），目前 Python 调度器用自己那版并显式注明 beam_width=1 “保持历史行为”。

账本跨语言实现的分裂被写进注释而不是被掩盖。 读这类代码时，看到”两版语义不一致且保留历史行为”的注释，应当把它当作”此处口径有坑”的路标：它意味着任何跨这两层的容量比较都可能是错的，而正确做法是统一到一侧。

三个字段决定它是什么

TrtllmAttentionMetadata 上与 workspace 有关的字段是：

1
workspace: Optional[torch.Tensor] = None
2
cuda_graph_workspace: Optional[torch.Tensor] = None
3
workspace_reclaimable: bool = field(default=True, init=False)

以及一个”按当前是否在图中切换”的访问器：

1
def effective_workspace(self) -> Optional[torch.Tensor]:
2
    """Attention-kernel workspace, switching to the CUDA-graph copy under capture."""
3
    return self.cuda_graph_workspace if self.is_cuda_graph else self.workspace

为什么需要两份 workspace？ 因为 CUDA Graph 捕获会把地址烘进图里：捕获期 kernel 读写的地址在重放时必须还是同一块。如果急切（eager）路径与图路径共用一块 workspace，那么”某个 eager 前向改变了这块张量的位置或大小”就会让已捕获的图指向失效地址。源码把这条风险直接写在注释里（在允许捕获的上下文管理器处）：

1
cuda_graph_workspace tensor, invalidating addresses baked into
2
previously captured graphs.

于是设计变成：急切路径用 workspace，图路径用 cuda_graph_workspace，两者互不干扰。这是”图内外必须双份”这条规则的又一处体现（前一节里 static_tensors 与 static_tensors_cpu 是同一规则的另一种形式）。

第三个字段 workspace_reclaimable 是”这块 workspace 能不能被回收”的开关，它的赋值逻辑解释了什么时候不能：

1
if metadata.is_cuda_graph or not fmha.supports_workspace_reclamation:
2
    # Conservatively disable reclamation for metadata used by graphs
3
    # or backends that can retain staged workspace state.
4
    metadata.workspace_reclaimable = False

两个情形都很直观：图里用过的不能回收（地址被烘进图），后端会”保留 staged 状态”的不能回收（比如把上次的索引/扫描结果留在缓冲里复用）。supports_workspace_reclamation 默认是 False（接口里的默认值），也就是说默认保守：只有后端明确声明”我不保留状态”，才允许回收。

grow-only：workspace 怎么长大

workspace 初始是空的：

1
if self.workspace is None:
2
    self.workspace = torch.empty((0, ), device='cuda', dtype=torch.int8)
3

4
if self.cuda_graph_workspace is None:
5
    self.cuda_graph_workspace = torch.empty((0, ), device='cuda', dtype=torch.int8)

之后由 FMHA 实现按需求 resize_。关键在于 WorkspaceShrinkPolicy 的注释一句话点明了它的形态：

1
# Retain storage; grow-only layer sizing accumulates this forward's peak.

“grow-only”：每个注意力层按自己的需求扩容，共享同一块 storage，于是这块 storage 的大小是”所有层里最费的那一层”，而不是”所有层相加”。这正是”workspace 共享”的全部收益所在——它在层间复用，所以只需要一份。

延迟回收：三次欠载前向之后才缩

但 grow-only 有个后果：一旦某个 shape 把 workspace 撑大，它就不会自己缩回去。长 prompt 的前向过后，即使后面全是短请求，那块显存仍然被占着。WorkspaceShrinkPolicy 就是为解决这件事而生的：

1
class WorkspaceShrinkPolicy:
2
    """Shrink after three underfilled forwards using their maximum demand."""
3

4
    def __init__(self, floor_bytes: int) -> None:
5
        if floor_bytes < 0:
6
            raise ValueError("Workspace floor must be nonnegative")
7
        self.floor_bytes = floor_bytes
8
        self.reset()
9

10
    def reset(self) -> None:
11
        self.remaining = 3
12
        self._window_max = 0
13

14
    def finish_forward(self, required_bytes: int, capacity_bytes: int, *, grew: bool) -> int | None:
15
        if required_bytes < 0 or capacity_bytes < max(required_bytes, self.floor_bytes):
16
            raise ValueError("Workspace capacity is below its requirement or warmup floor")
17
        if required_bytes == 0:
18
            return None
19
        if grew or capacity_bytes == self.floor_bytes or required_bytes == capacity_bytes:
20
            self.reset()
21
            return None
22
        self._window_max = max(self._window_max, required_bytes)
23
        self.remaining -= 1
24
        if self.remaining:
25
            return None
26
        target = max(self.floor_bytes, self._window_max)
27
        self.reset()
28
        return target

逐条读这个状态机：

• floor_bytes 是”warmup 建立的下限”，永远不缩到它以下。它是初始化时实测的容量，代表”这套配置下无论如何都要有的量”。
• remaining = 3：连续观察三次”欠载前向”（required_bytes < capacity_bytes）之后才动手，取这三次里的最大需求作为目标大小。
• 三种情况立即重置计数并放弃收缩：grew（这次涨过）、容量已经等于下限（无可缩）、需求恰好等于容量（正好用满）。“正好用满”也要重置，因为它说明当前容量是合适的，不该缩。
• 目标是 max(floor, window_max)：缩到”下限与近三次最大需求中较大的那个”，也就是留有余量、避免抖动。

为什么是三次？ 因为单次前向的需求有很强的随机性（batch 组成、序列长度分布、是否命中了缓存前缀）。按单次缩，会在长短请求交替时来回震荡：一缩就赶上大请求 → 又要涨回去。用三次滑窗的最大值作为目标，本质是”用一点滞后换掉抖动”，代价是收缩决策慢两步。 而”三次”这种小常数是经验值——它足够覆盖”长短交替”的常见周期，又不至于在负载真的降下来之后迟迟不还内存。

回收器：它到底怎么”缩”

真正执行收缩的是 EagerWorkspaceReclaimer，它用上下文管理器包住每次前向：

1
class EagerWorkspaceReclaimer:
2
    """Reclaim single-stream eager scratch, preserving the warmup floor."""
3

4
    def __init__(self, metadata: "TrtllmAttentionMetadata") -> None:
5
        if metadata.workspace is None:
6
            raise ValueError("Warmup must initialize eager workspace before reclamation")
7
        self.policy = WorkspaceShrinkPolicy(metadata.workspace.untyped_storage().nbytes())
8
        ...

注意 __init__ 的断言：“warmup 必须先初始化 eager workspace，回收器才能建立”。这是一个顺序约束，它的存在说明 floor 的语义依赖”初始化时那次实测”。

forward() 里有四个值得单独看的判断：

1
if self._active:
2
    raise RuntimeError("Overlapping forwards cannot share eager attention workspace")
3
if metadata.is_cuda_graph or torch.cuda.is_current_stream_capturing():
4
    yield
5
    return
6
...
7
if do_multi_stream() or (self._stream is not None and stream != self._stream):
8
    self._disabled = True
9
    logger.warning("Disabling eager workspace reclamation after a stream change")
10
    yield
11
    return

其一，“重叠前向不能共享 eager workspace”是一个直接的运行时断言。 它把”同一块 scratch 被两个前向同时用”这种最危险的状态直接封死。

其二，图路径与捕获期一律跳过。 这两者是同一件事的两种检测方式（元数据标记 + 当前是否在捕获）。

其三，流一变就永久禁用回收，并打一条 warning。这条非常关键：回收意味着”释放再重新分配”，而重新分配会改变地址。如果同一个 workspace 曾在流 A 上被使用、之后又在流 B 上被回收重建，那么在流 A 上还没执行完的工作会读到已释放的内存。“一次流切换即永久禁用”是最省事、也最安全的处理 —— 它不做任何推断（例如”这块内存是否已被流 A 用完了”），直接放弃优化。

其四，collect_demand 的逻辑：

1
collect_demand = capacity_before > self.policy.floor_bytes
2
if collect_demand:
3
    # Retain storage; grow-only layer sizing accumulates this forward's peak.
4
    metadata.workspace.resize_(0)

只有在”当前容量高于 floor”时才进入观察模式；而且在进入前先 resize_(0) —— 把张量的逻辑大小归零、保留 storage，这样这次前向里各层的 grow-only 扩容会重新”从零累加”，得到的 numel 恰好是”这次前向的真实需求”。这就是注释里那句 grow-only layer sizing accumulates this forward's peak 的含义：把残留容量清掉，才能测准这一次的需求。

而收缩本身的实现有一处值得注意的取舍：

1
if target is not None:
2
    device = metadata.workspace.device
3
    # resize_(smaller) retains storage. Release first, then
4
    # allocate; pure scratch has no payload to preserve.
5
    metadata.workspace = None
6
    metadata.workspace = torch.empty(target, dtype=torch.int8, device=device)

注释解释了为什么不直接 resize_(smaller)：resize_ 缩小只改逻辑大小、保留物理 storage，所以显存并不会真的还回去。 要真的缩，必须”先丢掉引用、再重新分配更小的一块”。代价是地址变了，所以前提是”纯粹是 scratch、没有需要保留的内容”。“纯 scratch”这个前提是被显式确认过的，而不是假设的——这句注释是这条代码敢直接换掉张量的唯一依据。

与预算的接缝：回收是”事后省”，预留是”事前算”

把这一节与上一节接起来，会看到一条完整的时间线：

“事前按最坏情况预留”与”事后按实测需求收缩”这两件事同时存在——前者保证不 OOM，后者保证长尾不浪费。它们并不矛盾：预留在初始化时把最坏情况的钱先圈出来（否则 KV 池会把钱花光），收缩在运行中把实际用不到的部分还回驱动。两者合起来的净效果是”上限由最坏情况决定、下限由实际负载决定”。

静态输入缓冲：为什么它们是”复用的”而不是”每图一份”

捕获前必须先准备好所有输入张量，并且这些张量在重放时会被就地改写。源码把这件事和它的约束写在一起：

1
after the static input buffers have been filled and outside any capture region, so
2
graph reads -- writes have to be in place, into the buffers capture recorded.

配套的运行时校验则在阻止”替换张量”这种写法：

1
f"{type(value).__name__}. Graph-visible buffers must be "
2
...
3
"Graph-visible buffers must be updated in place; "

“graph 可见的缓冲必须就地更新” 是一条硬约束，它把 CUDA Graph 的使用方式从”传参”变成了”改内存”。这也解释了为什么静态张量要按最大的支持规模开一次、然后靠切片取用：

1
self.shared_static_tensors = {
2
    "input_ids": torch.ones((max_total_tokens, ), device="cuda", dtype=torch.int32),
3
    "position_ids": torch.zeros((1, max_total_tokens), device="cuda", dtype=torch.int32),
4
    ...
5
}

1
def _create_shared_static_tensors(self) -> None:
2
    """Allocates static tensors sized for the largest supported num_tokens."""
3
    max_total_tokens = (
4
        self.config.max_num_tokens if self.is_encoder_decoder else min(
5
            self.max_supported_num_tokens, self.config.max_num_tokens))
6
    max_batch_size = self.max_supported_batch_size

max_supported_num_tokens / max_supported_batch_size 这两个”supported”上限与 config.max_num_tokens 取 min，说明静态张量的尺寸有一个独立于用户配置的上限。这很关键：如果静态张量真的按用户给的 max_num_tokens 开，那么”把 max_num_tokens 调大”这个纯调度侧的旋钮会顺带吃更多显存，而用户不会预期到这一点。把”支持的桶上限”与”配置的 token 上限”分开，避免了配置与显存之间出现一条隐藏耦合。

pinned host 影子缓冲：为什么 host 侧要 pageable

静态张量在设备侧有一份，在主机侧也有一份：

1
self.shared_static_tensors_cpu = {
2
    "input_ids": torch.ones((max_total_tokens, ), device="cpu", dtype=torch.int32,
3
                            pin_memory=prefer_pinned()),
4
    ...
5
    "seq_lens": torch.ones((max_batch_size, ), device="cpu", dtype=torch.int32,
6
                           pin_memory=prefer_pinned()),
7
}

对应地，解码侧的注释里有一句很值得看的话：

1
# prefer_pinned() is false: pageable host buffers are preferred, so the

也就是说 host 影子缓冲偏好”可分页”而不是”页锁定”。这与上一节 BufferManager 里 pinned / pinnedPool 的存在并不矛盾——它们是不同的使用场景：

• 需要异步传输的场合必须用 pinned，否则 H2D 会被驱动降级成同步拷贝；
• 只是”主机侧的一份镜像、用于就地拼装”的场合，pinned 反而不划算：页锁定内存是不可换出的，长期大量持有会压低整机可用内存。

而 seq_lens 那一份有一条特别说明，它是整套 CUDA Graph 设计里最精细的一处：

1
# Pinned static buffer for seq_lens. Each captured graph's attn_metadata._seq_lens
2
# is reseated (in maybe_get_cuda_graph) to a stable slice of this buffer, and the
3
# corresponding H2D copy into _seq_lens_cuda is captured inside the graph itself.

读法：seq_lens 每次重放都变（因为每条序列的长度在增长），如果每次都要 H2D 拷贝一个设备张量，那它就落在 graph 外面、每步都要一次 launch —— 这正好是 CUDA Graph 想消灭的开销。解法是让设备侧的 _seq_lens_cuda 也在 graph 里，而 H2D 拷贝本身被捕获进 graph。于是每步只需要纯 CPU 的 memcpy 到主机缓冲（注释明确说 per-replay seq_lens updates are pure CPU memcpys into that buffer），重放时 graph 内部的 H2D 节点会把它搬过去。

这就是”把拷贝也做成图的一部分”的典型手法：把”变化的数据”与”不变的结构”分离开，让变化的只有一次主机内存拷贝，不变的（包括拷贝动作本身）全部静态化。

额外的输入缓冲：用户声明的那些

还有一类静态张量来自”用户声明的额外模型输入”：

1
def _create_extra_input_buffers(self, max_total_tokens: int,
2
                                max_batch_size: int) -> None:
3
    """Allocate one static buffer per user-declared extra model input.
4

5
    The symbolic axis is sized at its bucket maximum. The pinned host
6
    shadow stages host-sourced inputs so their upload stays asynchronous.
7
    """

“符号轴按其桶最大值定尺”是这里的关键词：同一个输入张量在不同 batch 分桶下形状不同，而静态张量只能有一份，所以取所有桶的最大值。代价是”小桶下白白多占”，收益是”不必为每个桶各开一份”。在桶的数量可达十几个的情况下，这个取舍是明显划算的。

静态形状约束的三处体现

形状静态化带来的约束会在三个地方暴露出来，读源码时可以作为”是否理解了 CUDA Graph 代价”的检查点：

其一，输入必须 pad 到某个已捕获的形状。 源码在捕获键的构造里直接写明：

1
# [CUDA graph spec decode padding]
2
# We pad input IDs/position IDs to the maximum draft length (token per request).
3
# We're forced to do this because we cannot reallocate inputs over many graph runs.
4
num_tokens_for_capture = self._get_num_tokens_for_key(key)

“我们被迫这么做，因为在多次 graph 运行之间无法重新分配输入” —— 这句话把 padding 从”一种优化”降格成”一种必须”。投机解码下 token 数取决于 draft 长度，而 draft 长度可以动态变化；要把动态变化塞进静态形状，pad 是唯一出路。

其二，注意力元数据必须被”重座”（reseat）到稳定切片。 上面 seq_lens 的注释说的 reseated ... to a stable slice of this buffer 就是这件事：一个 graph 的 attn_metadata._seq_lens 必须指向它捕获时用的那段缓冲，不能在重放时换一个。

其三，多 rank 之间必须同意用哪张图。 注意力 DP 下的这段注释：

1
# The sampling tier joins the graph key, so it has to agree across the
2
# attention-DP ranks or they would replay different graphs -- the same
3
# reason can_run_cuda_graph and batch_size are all-gathered here. Ranks
4
# see different batches (an idle one resolves FULL), so fall back to
5
# FULL for the whole group unless every rank picked the same tier.

如果各 rank 重放不同的图，集合通信就会错配。 所以图键必须在 rank 间 all-gather 后达成一致，不一致时整组退回到最保守的档（FULL）。这是”分布式下的静态形状”比单卡麻烦得多的地方：不仅每张卡的形状要静态，还要所有卡选到同一张静态图。

图批次分桶选得好不好，直接反映在吞吐上

静态形状的代价不只是显存，还有”桶怎么选”。CUDA Graph 的捕获键里包含 batch size 与 token 数，因此可选的桶是一组离散值；批次落在哪个桶里，决定了这一段的执行效率。TensorRT-LLM 官方有一篇专门的博客讨论这件事，它给出的实测对照是：

图：同样的模型、同样的负载，只改 CUDA Graph 的批次分桶策略（+64 相对 x2），输出吞吐就有明显差异。这说明”图缓冲的静态形状”不是一个纯内存话题：它同时决定了 padding 的比例、重放命中率与调度器能塞进多少请求。

这篇博客的结论与本文的预算主线是同一件事的两面：分桶越密，padding 浪费越少、命中率越高，但需要捕获的图越多、每张图都要占一份常驻缓冲。所以”桶的粒度”本身就是一个用显存换吞吐的旋钮——它属于第二套预算的一部分，而不是纯粹的性能调优。源码里 filter_cuda_graph_batch_sizes 与 resolve_cuda_graph_batch_sizes 这类函数的存在，正是把”哪些桶值得捕获”变成了一条可计算的规则。

第二套预算意味着什么

把这一节收束成一句话：CUDA Graph 让”临时缓冲”变成了”常驻缓冲”，因此它必须被当成第二套预算来算。

第三行是关键：“不一定能被 profiling 测到” 正是 get_attention_workspace_fixed_bytes() 那句注释想表达的（“covering arenas first allocated after the peak prefill step”）。捕获发生在什么时候、捕获了哪些 shape，决定了哪块内存在 profiling 里是隐形的。 而源码给出的对策不是”想办法让 profiling 也捕获一遍”（那会让初始化成本翻倍、且捕获本身有副作用），而是让后端/组件按满容量显式申报。

这也解释了 why fixed_workspace_reserve 的名称里有 “fixed”：它与”attended KV 长度”无关，是一个按容量定尺的固定量，因此必须在峰值之外全额预留，而不是按比例分摊。“按满容量预留”这类开销不能参与”按比例共享预算”的协商 —— 因为协商的前提是”可以少给”，而定尺的 arena 少给就直接失败。