TensorRT-LLM 源码完全拆解（四）：KVCacheManager 与块管理——两级池、块表构造与前缀复用树如何记物理账

驱逐策略：LRU 队列与两级队列的一致性

BaseEvictionPolicy 的接口只有五个：getFreeBlock(cacheLevel, wantPlaceholder)、claimBlock(block[, priority, durationMs])、releaseBlock(block[, toFront])、getNumFreeBlocks(cacheLevel)。主实现是 LRUEvictionPolicy，为每个缓存层级（primary/secondary）各维护一条空闲队列——注意队列里放的是「有空闲内存的块」，不是「空闲内存」，回收块就是把它归还到对应层级的队列。

优先级如何进入 LRU？claimBlock 的重载带 priority/durationMs 参数，块被认领时记录保留信息；releaseBlock 归还时把块放到队列的尾部（或按 toFront 放头部）。块被驱逐的先后因此由「队列位置 + 优先级 + 过期时间」共同决定，而 secondaryOffloadMinPriority 这个构造参数给出了一条明确的边界：优先级低于该阈值的块不值得换出——直接销毁，换个新块进来更划算。这个阈值把「换出的固定成本」与「复用概率」放在一起权衡，是工程上非常务实的折中。

队列的完整性有专门的自检：WindowBlockManager::verifyQueueIntegrity() 与 BlockManager::verifyQueueIntegrity()。在调试构建下，它会验证「块的 mFreeBlockIterator 指向的队列位置与块自身的状态一致」「队列里没有重复块」「认领过的块不在队列里」。资源管理器最怕的不是分配失败，而是账实不符——某个块既不在任何队列、也没被任何请求引用，就永远漏出了账本。自检函数就是「账实相符」的断言层。

换入：块回来了，数据在 host

换入路径是 onboardBlock（BlockManager::onboardBlock → WindowBlockManager::onboardBlock）。当一个被换出的块在前缀树上被匹配命中时，它需要先回到主池才能被请求使用：分配一个主池目标块、mTransferManager->onboard(source, target, ...) 把 host 数据拷回显存、交换池索引，最后发一条 cacheLevelUpdated(secondary -> primary) 事件。换入同样要成对地认领两个层级的块，syncTransferManagerWithBufferManager 负责让传输管理器与缓冲管理器对池的布局保持同一视图。

为什么换入是按需的而不是预取？ 源码路径里换入只在「复用确实发生」时触发：前缀树匹配到块、且该块当前在次池。这样做的理由是预取需要预测「哪些前缀即将被用到」，而框架把预测责任交给了上层——调度器的前缀感知调度（第三篇的 beneficialToSkip）只是调整准入顺序（让可能命中的请求先进来），真正触发换入的是执行路径上的匹配。分离「调度排序」与「数据搬运」，各自复杂度都更低。

图：Mooncake 论文的以 KVCache 为中心的分离式架构（来源：arXiv:2407.00079）。TensorRT-LLM 的两级池是这套思想在单机框架内的最小实现，而跨实例的 KV 传输则由 cache_transmission 与 P/D 分离模块完成。

键：BlockKey 比「一串 token ID」复杂得多

BlockKey 的字段揭示了「什么条件下两块 KV 可以互相顶替」：

1
struct BlockKey {
2
    bool usesExtraIds = false;
3
    std::optional<LoraTaskIdType> loraTaskId = std::nullopt;
4
    VecUniqueTokens uniqueTokens;                       // 去重后的 token 序列（含 extra id）
5
    std::vector<MmKey> extraKeys;                       // 多模态内容哈希 + 块内偏移
6
    std::optional<std::string> cacheSalt = std::nullopt; // 缓存盐，隔离不同命名空间
7
};

• uniqueTokens 不是原始 token，而是 UniqueToken{token, extraId} 的序列。extra id 解决的正是p-tuning / prompt embedding 场景：不同请求可能带着相同的假 token ID（数值大于词表）却对应不同的 embedding 表。若只比 ID，两条用不同 embedding 的请求会错误复用彼此的 KV——这是正确性问题，不是性能问题，官方文档专门给了 extra id 的用法示例。
• loraTaskId：不同 LoRA adapter 的同一段文本，KV 完全不同（站内 SGLang 多 LoRA 服务一篇讲过适配器如何在批内共存），所以 LoRA 任务号必须进键。
• cacheSalt：多租户场景的命名空间隔离——同一段系统提示词，在租户 A 与租户 B 之间不应互相看见。盐值像密码学里的盐一样把哈希空间切开。
• extraKeys（多模态）：图像等内容的哈希与块内偏移，同样是为了让「token 序列相同但图像不同」的请求不误命中。注意 supportsPartialMatching() 对多模态键直接返回 false——键里含内容哈希时，前缀截断没有意义，干脆禁用部分匹配。

numMatchingTokens 的实现是一个教科书式的「先比约束字段、再比 token」：

1
int BlockKey::numMatchingTokens(BlockKey const& other) const noexcept {
2
    SizeType32 numMatched{0};
3
    if (usesExtraIds == other.usesExtraIds && loraTaskId == other.loraTaskId
4
        && extraKeys == other.extraKeys && cacheSalt == other.cacheSalt) {
5
        auto [matchEnd, otherMatchEnd] = std::mismatch(
6
            uniqueTokens.begin(), uniqueTokens.end(), other.uniqueTokens.begin(), other.uniqueTokens.end());
7
        numMatched = std::distance(uniqueTokens.begin(), matchEnd);
8
    }
9
    return numMatched;
10
}

约束字段不一致时直接返回 0——「要么整体可比、要么不可比」，避免了「按 token 部分匹配但盐不同」这种半吊子状态。

树：一棵树服务所有窗口仓

查找结构是 UnifiedBlockTree，注释说得很直白：它是「统一树」，单实例服务所有 WindowBlockManager，通过把窗口大小当作 value key 实现。也就是说树上每个节点（对应一段 token 前缀）可以挂多个值：窗口 4096 的仓挂一个块、窗口 32768 的仓挂另一个块，互不干扰。

为什么合一棵树而不是每仓一棵？三个理由在源码注释与实现里都能读到：

1. 前缀是同一段 token 序列，按 token 前缀建树是自然键；不同窗口仓对同一前缀各存各的块，只是值不同。
2. 线性注意力仓（Mamba）用特殊的 value key：kRecurrentStates = -1（见 radixBlockTree.h）。负值刻意区别于所有合法窗口（>= 1），于是循环状态的快照块能和全注意力块共存于同一棵前缀链上——混合模型（如 Mamba + 全注意力层交错）的复用语义因此保持完整。
3. 锁只有一把：UnifiedBlockTree 持有一个 std::recursive_mutex，所有树操作（插入、lookup、detachFromLookupNode）都在其保护下。递归锁是刻意选择的：BlockManager::addSequenceBatch 的两阶段认领本身持有树锁，而它调用的 getFreeBlock / detachFromLookupNode 也要拿锁——用普通互斥量会自锁死。

块的树侧接口把「双向布线」写得很清楚：attachToLookupNode(node, windowSize) 把块挂到节点的某个窗口槽位，若块原本挂在别处，会先清除旧槽位并向上级联剪枝；detachFromLookupNode() 反向操作，空节点由自底向上、遇到第一个非空祖先即停的剪枝逻辑回收。根块（mCachedBlocksRoot）通过 setAsRoot(rootNode, windowSize) 挂在根节点上，供直接子块的 getPrevBlock() 查到。

站内 RadixAttention 完全拆解与 SGLang RadixCache 内部讲过 SGLang 的前缀树实现，可以对照阅读：SGLang 的树节点直接持有 token 片段与锁引用；TensorRT-LLM 的树是「节点 + 多窗口槽位 + 块句柄」的三层结构，且键是带盐与多模态哈希的 BlockKey。两边的相似处在于都把「满块」定义为可共享单位，差别在于 TensorRT-LLM 把多租户、多模态、LoRA 这些隔离维度直接编进了键。

查：findReusableBlockMatches 的三道闸门

匹配入口是 WindowBlockManager::findReusableBlockMatches(blockKeys, enablePartialReuse, copyOnPartialReuse, maxMatchedTokens)。它沿树自根向下逐块走，每步都有闸门：

闸门一：token 预算。 candidateMatchedTokens + numMatchedTokens > maxMatchedTokens 就停止——调用方传来的上限是「本请求当前可复用的最大 token 数」（预算口径细节见后文）。为什么不匹配了再说？因为匹配结果影响准入判断，必须受限匹配。

闸门二：SWA 的安全前缀截断。 这是这段代码里最微妙的部分。滑窗注意力下，一个块即使前缀匹配上，也可能因为窗口已滑过而不再覆盖当前需要的 K/V。代码用 updateSafePrefix 与 latestMissingAnchorEndToken 跟踪「最后一个缺块的锚点」：

1
auto updateSafePrefix = [&]() {
2
    if (!mIsSWA || latestMissingAnchorEndToken == 0
3
        || candidateMatchedTokens >= latestMissingAnchorEndToken + mWindowSize) {
4
        safeMatchCount = candidateMatches.size();
5
        safeMatchedTokens = candidateMatchedTokens;
6
    }
7
};

含义是：只有当匹配跨度已经超过「缺失锚点 + 一个窗口」时，整段才算安全可复用；否则把匹配列表截断到上一个安全点（candidateMatches.resize(safeMatchCount)）。直觉是：窗口大小为 W 的块只保证覆盖它自己往后 W 个位置；若中间有一段没有块（traversal-only 匹配，只沿树走但值槽为空），那么「缺失点到当前点」的距离必须至少有一个完整窗口，才有块真正覆盖当前需要的范围。

闸门三：部分匹配与写时复制。 若开启 enablePartialReuse，精确匹配失败后还会找 findPartiallyMatchingNodes——键是某已缓存块的前缀。部分命中时，块里只有前若干 token 有效，不能直接共享（后段数据属于另一条请求的上下文），于是要么复制（copyOnPartialReuse 为真），要么仅当块「无引用且是叶子」时才敢原地转用。两个开关的默认值与取舍在后文「工程取舍」里再展开。

匹配的输出是 ReuseMatchResult：匹配列表（每项含块指针、匹配 token 数、是否部分匹配、是否仅遍历即 traversal-only）、总匹配 token 数，以及 firstNewBlock——第一个没在树里找到的满块键。这个字段是给调度器用的：第三篇讲过 PrefixReuseSummary.firstNewBlock 决定「是否推迟这条请求」；此处正是它的产地。

认领：claimMatchingBlocks 的批内协调

匹配只是「找到」，认领（claim）才是「占住」。claimMatchingBlocks 要处理三个真实世界的麻烦：

麻烦一：beam search 的共享。 beam 宽度大于 1 时，同一请求的多条 beam 共享上下文块：isShareLastContextBlock = (mCacheType == kCROSS || inputLength % mTokensPerBlock == 0)，据此计算 numSharedContextBlocks——共享的前缀块只认领一次，最后一块在 beam 之间按需分配。CacheType::kCROSS（编码器-解码器跨注意力）的 KV 一次分配、永不增长，所以它天然可全共享。

麻烦二：同批多请求抢同一个块。 一个微批里可能有五条请求都命中同一段系统提示词的前缀块。这些块是「空闲队列里的叶子块」，如果每条请求各复制一份，复用收益瞬间归零。源码用 PartialClaimTracker 做批内仲裁：

1
// 命中同一个叶子块的多个请求里，只有当「没有更强者」的请求才原地复用，
2
// 其余的标记 needsCopy；若后续出现完整匹配，先前的复用承诺被降级为复制。
3
if (tIt != tracker.map.end()) {
4
    if (tIt->second.fullyMatched) { claimed.needsCopy = true; }        // 前者已完整匹配 → 我必须复制
5
    else {                                                            // 前者只是打算复用 → 让它改复制，我复用
6
        claimResults[tIt->second.requestIdx].claimedBlocks[tIt->second.claimedIdx].needsCopy = true;
7
        claimed.needsCopy = false;
8
        tIt->second.requestIdx = requestIdx;
9
        tIt->second.claimedIdx = result.claimedBlocks.size();
10
    }
11
} else {
12
    claimed.needsCopy = false;                                        // 第一个部分匹配的请求 → 原地复用
13
    tracker.map[blockId] = {requestIdx, result.claimedBlocks.size(), /*fullyMatched=*/false};
14
}

这段逻辑可以概括成一句话：「一个块只能有一个原地复用者，其余人复制」；部分匹配的复用者可以被后面出现的完整匹配者顶掉（完整匹配更强），复制者则要承担 shouldReleaseCopySource 的责任——若源块是被「无引用非叶」这类边缘状态借来的，最后一个拷贝者负责在复制完成后把它释放回账本。

麻烦三：认领即保护。 所有被匹配到的块立即 mEvictionPolicy->claimBlock(...)——从空闲队列摘除，防止驱逐策略在认领期间把它换出或销毁。这个「先占位再干活」的模式和内核里的 cp.async 预取一样：先预定资源，再执行操作，避免中间态被并发路径截胡。

图：CacheBlend 给出的复用与选择性重算思路（来源：arXiv:2405.16444）。TensorRT-LLM 的复用建立在「严格前缀一致」之上（含盐、LoRA、多模态哈希的全部隔离维度），不涉及跨请求的近似复用；两条路线的精度-收益边界正好互为参照。

存货与回收：storeBlocks、releaseBlocks、releaseLastBlock

复用链的另一半是「把算好的块登记进树」。

• storeBlocks（WindowBlockManager::storeBlocks） 在请求推进到满块边界时调用：取该块的 BlockKey，把块挂到树上（attachToLookupNode），必要时置 mIsFull = true，并把 mPrevBlockInSeq 链续上；返回 (prepopulatedLen, blockIds) 供上层记账。只有满块才允许跨请求共享，原因在官方文档里：块内最后一个 token 的 KV 在「它自己」被算出来之前不完整，未满块的复用会导致缺 token。
• storeBlocksForReuse / storeNewBlock 是请求结束（或到达某个复用检查点）时对尾块的补登记：把最后一段已算 KV 的块也放进树，让后续请求可以少算这一段。注释口径与 getNeededBlocksOneStep 里那句「序列插入忽略最后一个 prompt token，因为它的 KV 无法恢复」呼应——最后一个 token 的 KV 永远不可复用，这是自回归结构决定的硬边界，不是实现懒省。
• releaseBlocks / releaseLastBlock / schedulingReleaseBlocks 负责反向操作：请求结束时若块仍被引用则只减引用；引用归零时把块归还空闲队列（releaseBlock），树上的登记按块的新状态刷新——若是叶子且无引用，它留在树上等复用（这就是「已缓存但空闲」的块）；若需要腾位则被驱逐时再摘树。

一个统计口径值得记住：KvCacheStats 里的 cacheHitRate = reusedBlocks / (reusedBlocks + missedBlocks)，并且区分了 allocNewBlocks（从未承载过 token 的全新块）与 allocTotalBlocks。这条口径是「复用收益」的可观测来源，也是事件系统对外报数的分母。

startScheduling / commit / schedulingRemoveSequence

KVCacheManager::startScheduling() 转发到每个仓的 startScheduling()，后者遍历空闲队列上每个块的 startScheduling()——把 mSchedulingRefCount 清零并重置调度期状态。这是「预测-提交」两阶段的分界点：调度器开始新一轮准入前，先把上一轮探测留下的临时占位全部清干净。

与之配套的是：

• schedulingRemoveSequence(requestId)：调度阶段的「移除序列」——不动树、不动引用，只释放调度期占位。第三篇讲过的 MAX_UTILIZATION 暂停-重试路径就走这条：先调度期移除，再决定是否重算。
• removeSequence(requestId)：运行期的真正移除——返回「最后一个块的降级块号」std::optional<KVCacheBlock::IdType>，供上层更新输出缓冲的映射（KV 被重算后，最后一块可能被释放重建）。
• commitAndGetBlockHashesForRequest(llmRequest, windowSize)：提交阶段。把本次分配涉及的块哈希化并返回给调用方（Python 侧对应 commit_and_get_block_hashes），让事件系统与外部观测拿到「这条请求现在持有哪些块的哈希」。这条接口也是第三篇所说「先预测、后提交」的提交端——预测阶段无副作用，提交阶段才落账。

analyzePrefixReuse：一份摘要服务三个消费者

KVCacheManager::analyzePrefixReuse 转发到 BlockManager → WindowBlockManager，对指定 token 序列做只读的树遍历，产出 PrefixReuseSummary：

1
struct PrefixReuseSummary {
2
    SizeType32 reusableBlocksAllocated{0};  // 已被引用（refCount > 0）的前缀块数：块预算用
3
    SizeType32 reusableBlocksAll{0};        // 树上有登记的前缀块总数（含空闲缓存块）：token 预算用
4
    std::optional<BlockKey> firstNewBlock{std::nullopt};  // 第一个未命中的满块键
5
};

三个字段对应三个消费者，第三篇已经讲过用法，这里补上「为什么分成两个计数」：

• 块预算（capacity / 分配侧）只认 reusableBlocksAllocated。注释解释得直白：一个「空闲但有缓存」的块被复用时，仍要从空闲池里借出一个块（它只是从「空闲缓存态」转为「被引用态」），所以不能把它当成「不需要池容量」；只有已被引用的块（多条请求共享）才真正减少空闲池需求。记账如果混用两者，会与驱逐策略的空闲计数双重计算，导致超额准入。
• token 预算（compute / 计算侧）认 reusableBlocksAll。因为无论块当前是什么状态，命中前缀都会让引擎通过 prepopulatedPromptLen 跳过这段重算——省的是算力，不是容量。
• firstNewBlock 供调度器的跳过检查判断「要不要推迟这条请求」；std::nullopt 表示「满前缀全命中」或「还没有可探测的满块」。

预算侧还有一处极细的边界：maxRecoverableSharedBlocks = (promptInputLen - 1) / tokensPerBlock——减一是因为最后一个 prompt token 的 KV 不可恢复。当 prompt 恰好落在块边界上时，按全部 unique token 数去数可复用满块，会多认一个块；减一把这个 off-by-one 修掉。这类「差一」在分页系统里普遍存在（块内最后一个 token 的 KV 依赖前向计算），把它写进公式而不是靠调用方自觉，是源码值得学的地方。

图：Gemma 3 技术报告给出的滑窗（local）与全注意力（global）层配比对质量的影响（来源：arXiv:2503.19786）。VSWA 模型里不同层窗口不同，正是 BlockManager 按窗口分仓、并让各仓在共享树上各占槽位的根本原因。

setOffsets：两级翻译与三种布局

1
void WindowBlockManager::setOffsets(tk::KVCacheIndex* offsetsPtr, tensorrt_llm::Dims const& offsetsShape,
2
    SizeType32 beamIdx, SizeType32 blockIdx, KVCacheBlock::IdType blockId) const
3
{
4
    auto const& block = getBlockById(blockId);
5
    bool const isSwaPlaceholder = blockId == KVCacheBlock::kPlaceholderBlockId;
6
    for (poolIdx...) {
7
        const auto& pool = mPools.at(poolIdx);
8
        for (auto const xIdx : {kIdx, vIdx}) {          // K 与 V 各写一次
9
            auto const offsetIndex = flat_index(offsetsShape.d, poolIdx, beamIdx, xIdx, blockIdx);
10
            auto const fieldIdx = (mCacheType == CacheType::kSELFKONLY || isRecurrentState()) ? 0 : xIdx;
11
            auto const blockIndex = [&]() -> tk::KVCacheIndex {
12
                if (isSwaPlaceholder || block->isPlaceholder()) { return tk::KVCacheIndex::nullIndex; }
13
                if (pool.layerFirstLayout) {                 // {numLayers, numBlocks, kvFactor, blockSize}
14
                    return tk::KVCacheIndex{common::flat_index3(
15
                        layerIdx, block->getMemoryPoolBlockIndex(), fieldIdx, mNumPrimaryBlocks, mKVFactor)};
16
                }
17
                return tk::KVCacheIndex{common::flat_index3(
18
                    block->getMemoryPoolBlockIndex(), layerIdx, fieldIdx, pool.numLayers, mKVFactor)};
19
            }();
20
            offsetsPtr[offsetIndex] = blockIndex;
21
        }
22
    }
23
}

逐行读下来的要点：

1. 两级寻址。表里的条目不是「块号」，而是 KVCacheIndex——一个能把「池归属 + 池内槽位」编码进同一个整数的封装。内核拿到索引后，用高位选择池、低位寻址槽位；分离式部署时同一份索引结构还能跨实例解释。
2. 三种布局。layerFirstLayout 决定池内维度序：层优布局是 {numLayers, numBlocks, kvFactor, blockSize}，块优布局是 {numBlocks, numLayers, kvFactor, blockSize}。为什么两种都要？同一池被多层共享（同 KV head 配置的层放一个池），层优布局让「按层切图」的后端（如逐层执行的 CUDA Graph 捕获）拿到连续地址；块优布局让「按块寻址」的常规 FMHA 内核实现更直。这是一个纯性能取向的布局开关，不影响语义，值就藏在每一层对内核访存模式的实测里。
3. fieldIdx 对 kSELFKONLY 恒为 0。自注意力-仅-K 的缓存只有 K 没有 V（或 V 与 K 同址），写表时只写一份。
4. 占位块写空索引。SWA 占位块与线性注意力占位块都映射到 nullIndex，内核读到它就知道「这块没有可见 KV」。

copyBlockOffsets：把表搬进张量

KVCacheManager::copyBlockOffsets(ITensor& output, SizeType32 outputSlotOffset, RequestIdType requestId) 把某请求的块列表逐块调用 setOffsets 写进设备张量的指定槽位；getBlockOffsetsOfBatch 遍历整批。getCacheBlockIds(requestId, windowSize) 则返回 vector<vector<vector<SizeType32>>>——per-beam 的块号列表再嵌套 per-pool，与 getBatchCacheBlockIds 一起构成「批-请求-beam-池」四维寻址。Python 侧对应 get_cache_indices，而 groupLayersByWindowSize 把层按窗口分组，保证每个仓只处理自己那些层的表项。

这里能看出这套设计与 vLLM 块表的异同：两者都是「逻辑位置→物理块」的翻译表，但 vLLM 的实现是单层映射（逻辑块号列表），TensorRT-LLM 的表同时携带池选择、层布局与 beam 维度——因为它的池结构本身就是多仓多池的。复杂度不是凭空来的，每一维都对应一个真实功能（VSWA / VGQA / beam search）。

跨 rank 一致性：账本不能各记各的

并行部署下每个 rank 都有自己的 KV 分片，但「总共能装多少块」必须一致——否则调度器在不同 rank 上会算出不同的准入结果，出现「rank 0 答应、rank 1 拒绝」的裂脑。两处保护：

• synchronizeBlockCountsAcrossRanks：在 calculateMaxNumBlocks 末尾对每个窗口的块数做 allreduce（取最小值，见 Python 侧 dist.allreduce(max_tokens, op=ReduceOp.MIN)），保证所有 rank 用同一套块数起步。
• isSortedVectorIdenticalAcrossAllRanks：检查各 rank 的窗口序列一致，不一致直接报错。

Python 侧的 calculate_max_num_blocks 还处理了一个工程细节：free_mem_fraction 默认 0.9，用户若同时设置了 max_tokens，取两者较小值并给出警告——「显存比例」与「token 上限」双配置的一致性由框架兜底，避免用户以为自己配的是上限、实际被另一个约束压掉。

calculateMaxNumBlocks：一个池的块数从哪来

BaseKVCacheManager::calculateMaxNumBlocks 的输入是一组显存预算（allottedPrimaryMemBytes / allottedSecondaryMemBytes）与窗口映射表，流程是：

1. 每个窗口算「每 token 缓存字节数」：cacheSizePerToken = kvFactor * Σ(numKvHeads) * sizePerHead，再乘 dtype 大小；池配置（PoolConfiguration{windowSize, sizePerHead, dtype}）可覆盖管理器级默认值——这就是「同一模型不同层 head_dim 不同」（如 Gemma4 的 SWA 层 256 维 + 全注意力层 512 维）能各自成池的原因。
2. 窗口间分预算：computeWindowSizeShares 默认按窗口数均分（注释坦承「更合理的做法是按窗口大小比例分配，但当前为了让各层性能一致，先均分」），可用环境变量 TRTLLM_WINDOW_SIZE_SHARES 覆盖，例如 0.4,0.6 给窗口 512 与 32768 分别 40% 与 60%。框架把「不均分」留成显式开关，而不是悄悄做启发式——默认保守、优化留给人。
3. 块数 = 预算 / （每 token 字节 × tokensPerBlock），向上取整。kv_cache_config.max_tokens 若非空且非 VSWA，作为额外上限被取 min（源码里有条注释直接提醒：在 VSWA 下设置 maxTokens「是个很奇怪的概念」，因为它限制的是每个窗口各自的 token 上限）。
4. 次池同理，来自 host_cache_size；若块复用被禁用而次池仍非零，会打印「禁用复用时分配次池内存没有意义」的警告。
5. 线性注意力/Mamba 混合模型特例：只有两种「窗口」（max_seq_len 与 kRecurrentStates）、且禁用复用时，循环状态按 max_batch_size 静态分配（每 batch 一个状态块）；否则按内存预算动态算。这解释了为什么混合模型的容量口径与纯 Transformer 不同。
6. 最后跨 rank 同步块数（上节）。

getNeededBlocksOneStep：两种预算各算一次

这是全篇最该逐行读的函数之一。它对一次前向「新需要多少块」的估算分三种请求状态：

状态一：上下文初始化 + 首个分块（或分离式的生成初始化）。 先算 prompt 覆盖长度：

1
auto const maxDraftTokensToAdd = std::min(req.getNumDraftTokens(), req.mMaxNewTokens);
2
auto const promptCacheLen
3
    = std::min((isCrossKv() ? req.getEncoderOutputLen() : req.mPromptLen) + maxDraftTokensToAdd,
4
          windowSize + mChunkSize)
5
    + mSinkBubbleLength;

三个加项各有一句可解释的理由：

• maxDraftTokensToAdd：投机解码的草稿 token 也会写入 KV（站内 SGLang 投机解码一篇讲过草稿-验证链路），必须计入容量；取 min(numDraftTokens, mMaxNewTokens) 防止「草稿数超过剩余生成长度」这种配置下的虚高。
• min(..., windowSize + mChunkSize)：滑窗注意力下历史只保留一个窗口；分块预填充一次最多推进 mChunkSize。取两者之和是「本次前向可能新增的 KV 上界」。
• mSinkBubbleLength：attention sink 相关的气泡长度。实现是 sinkTokensInLastBlock = sinkTokenLen % tokensPerBlock ? tokensPerBlock - sinkTokensInLastBlock : 0——sink token 占不满一个块时，为它保留一个整块的气泡。注释写明这是为 StreamingLLM 式的 sink 保底（TODO 里还有「为 streamLLM 完善」的条目）。

然后按块拆分：numSharedBlocks（满块，可跨 beam 共享）+ numUnSharedBlocks（尾块 × beam 宽），得 numRequiredBlocks。最后——复用折扣，这是与第三篇的接缝处：

1
if (mEnableBlockReuse && !mBlockManager.isVariableWindow() && !isCrossKv()
2
    && !req.isDisaggGenerationInitState())
3
{
4
    auto const summary = cachedSummary.has_value() ? cachedSummary.value()
5
                                                   : analyzePrefixReuse(req.getUniqueTokens(0), req);
6
    auto const promptInputLen = std::min(req.mPromptLen, windowSize + mChunkSize);
7
    auto const maxRecoverableSharedBlocks = (promptInputLen - 1) / getTokensPerBlock();
8
    // 块预算：只有「已被引用」的可复用块才真正减少空闲池需求
9
    auto const reusableAllocatedBlocks
10
        = std::min({summary.reusableBlocksAllocated, numSharedBlocks, maxRecoverableSharedBlocks});
11
    numRequiredBlocks -= reusableAllocatedBlocks;
12
    // token 预算：所有命中前缀都省算力，记在请求对象上供微批调度器计价
13
    auto const reusableAllBlocks
14
        = std::min({summary.reusableBlocksAll, numSharedBlocks, maxRecoverableSharedBlocks});
15
    req.setEstimatedReusableTokens(reusableAllBlocks * getTokensPerBlock());
16
}

四个守卫条件把「折扣不适用的场景」全部排除：未开块复用、VSWA（多窗口下前缀语义复杂，第三篇讲过可变窗口下前缀感知整体失效）、跨 KV、分离式生成初始化（该状态只搬 KV 不算上下文）。然后两个计数各自受限、各自记账，最后 setEstimatedReusableTokens 把估算挂到请求上——微批调度器用 reuse_adjusted_compute 折价计价时读的就是它。

状态二：生成中。 算的是「接下来的若干 token 会不会跨块」：

1
auto const tokensPerStep = req.getNumDraftTokens() + 1;
2
auto const maxTokensToAdd = std::min((twoStepsLookAhead ? 2 : 1) * tokensPerStep, maxTokensToAddToKVCache);
3
auto const numNextTokens = numCurrTokens + maxTokensToAdd;
4
// 若越过该窗口的最大 token 数 → 返回 0（滑窗下不再增长）
5
auto const numRequiredBlocks = (ceilDiv(numNextTokens, tpb) - ceilDiv(numCurrTokens, tpb)) * beamWidth;

两个细节：tokensPerStep = draft + 1（草稿 + 被验证的目标 token）；twoStepsLookAhead 是给「两阶段预看」的选项——第三篇里 GUARANTEED_NO_EVICT 的到完成预留走一步展望，微批侧会按需做两步展望，避免「这一步刚好不跨块、下一步就跨块」的漏算。

状态三：其他状态返回 0。 不增长就不需要。

getRemainingBlocksToCompletion：到完成的预留口径

getRemainingBlocksToCompletion(req, windowSize, cachedSummary) 与上面成对，回答「这条请求跑完还要多少块」。它对跨 KV、上下文初始化、生成中分别给出「已完成部分之外的余量」；GUARANTEED_NO_EVICT 调度策略（第三篇讲过）在准纳入口用它做整生命周期预留。Python 侧对此有一条诚实得罕见的 TODO 注释：C++ 的 get_remaining_blocks_to_completion 与 Python 版本的语义不一致（C++ 排除已分配块、Python 包含），目前 Python 调度器用自己那版并显式注明 beam_width=1 「保持历史行为」。账本跨语言实现的分裂被写进注释而不是被掩盖——读源码时看到这种注释，应当把它当作「此处口径有坑」的路标。

Python 侧 KVCacheManager：三个钩子与一条补救路径

PyTorch 后端的 resource_manager.py 里，KVCacheManager 继承 BaseResourceManager，兑现四个接口：

• prepare_resources(scheduled_batch)：每步前向之前按批分配。文档口径是「首次进入上下文的请求按整个上下文分配；生成中的请求按下一步分配；块式组织且块内有余量时可能不实际分配」。源码里交叉 KV 走独立路径（一次分配、永不增长），上下文请求按首块/尾块分流。
• update_resources：官方文档承认「KV 不需要更新，此函数当前不做操作」——但留下一条明确接口位：「若 Python 侧支持 KV 复用，Radix 树管理在这里更新」。这是为「Python 实现的一等公民复用管理器」预留的扩展点。
• free_resources(request, pin_on_release=False)：请求结束回收；pin_on_release 允许把块钉住不回收（供外部复用/传输），默认关闭。
• get_max_resource_count / get_needed_resource_to_completion：给调度器的容量口径，后者显式注明 beam 口径与 C++ 不一致。

此外就是本篇开头提到的 get_request_kv_block_budget：返回「到完成所需块数 vs 主池容量」的一对值，允许 PyExecutor 直接拒绝不可能完成的请求。它有一条极严的适用性守卫（kv_block_budget_applies）：仅当「恰好是基类 KVCacheManager」且「非 VSWA、非线性注意力、所有窗口等于 max_seq_len、且非 CROSS 或 CROSS 特批」才生效——注释解释了原因：稀疏/压缩 KV 管理器（RocketKV、DSA 缓存）保留的 token 少于全序列，按稠密口径估会误杀请求；Mamba 混合把容量分给额外的状态缓存，单个「需要/容量」对表达不了。「估不准就别估」被写成了类型判断，这段守卫本身就是「预算接口的适用边界」最好的教材。

最后是那条值得单独讲的补救路径 fit_token_budget：

微批调度器准入时按估计的可复用 token 折价计价（reuse_adjusted_compute 用 estimated_reusable_tokens）；真实复用数要到 addSequence 里才算出来（prepopulated_prompt_len）。若实际比估计少，前向就会多物化 token，撞上 total_num_tokens <= max_num_tokens 断言（GitHub #13318）。

补救发生在 prepare_resources 之后：此时块已分配、序列已加入，请求不能退出批——但也不必退出。函数从批尾部（先剪「尾块请求」，因为只有尾块请求享受了复用折扣与草稿 token 计价）开始裁剪上下文 chunk，把被剪的请求变回分块状态，多出来的 token 下一轮再算。源码注释把「为什么只能在此刻、只能剪不能踢」讲得很透：剪枝不动批成员，因此不会破坏注意力 DP 的投票、在途集合或其它管理器的每请求状态。估算-实测裂口靠「延后计算」而不是「回滚分配」来弥补——这是资源管理里非常典型的「以时间换一致性」。

顺带的产物：显存布局与索引器缓存

容量核算与块表构造之间还有两块值得点到的实现：

• FP4 KV 的块级缩放池：NVFP4 缓存会为每个普通池再建一个「块缩放因子池」，形状相同、最后一维是 head_size / 16，由 createBlockScalePools 建立、_merge_kv_cache_pool_pointers 在 Python 侧把紧凑缩放行对齐到物理池序。量化 KV 的元数据也走同一套「池-块」账本，这是量化推理类文章里常被忽略的工程细节。
• DSA 索引器 K 缓存池：createIndexerKCachePools 为 DeepSeek 稀疏注意力的闪电索引器单独建池（可 FP4、块大小默认 128），并有按层的掩码 indexerKCacheLayerMask。站内 DSA 稀疏注意力完全拆解讲过索引器在算法层做什么；在块管理层，它只是「又一个带特殊布局的池」。