[None][feat] Support layer and buffer reuse policies for attention and SSM

cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/AGENTS.md

@@ -49,7 +49,7 @@ them before changing `batch.*`, `sparseCache.*` or the index conversion.
 - Prefer composition when it expresses ownership or containment. Use
   inheritance only for a real IS-A relationship or required virtual dispatch,
   such as `CommittedPage : Page` and `EventManager : EventSink`.
-- Preserve strongly typed indices. `LayerGroupId` is a public lifecycle
+- Preserve strongly typed indices. `CacheGroupId` (and legacy `LayerGroupId`) is a public lifecycle
   semantic, `PoolGroupIndex` is a storage-layout index, `BlockOrdinal` is a
   sequence position, and `SlotId`/page indices address storage. Do not convert
   or compare them implicitly.
@@ -93,9 +93,41 @@ must eventually be `close()`d. Its normal flow is:
    for later requests; `stopCommitting()` finalizes this process.
 5. Suspend or close the request, returning pages to holding/eviction ownership.
 
-`LifeCycleRegistry` maps model layers to lifecycle groups. Attention lifecycles
+`LifeCycleRegistry` maps layer buffers to lifecycle groups. Attention lifecycles
 may have sliding-window and sink-token rules; SSM lifecycles represent a
-recurrent-state checkpoint. A pool-group index is a storage-layout index and is
+recurrent-state checkpoint. Both layer configs expose `reusePolicy`: `REQUIRED`
+constrains the prefix match, `OPTIONAL` offers state at the required endpoint for
+explicit selection on first resume, and `PRIVATE` never publishes reusable pages.
+A buffer may override its layer's default through `BufferConfig::reusePolicy`.
+This applies to attention and SSM buffers, so one model layer may span several
+lifecycle groups. Use `getBufferGroupId(layerId, role)` and `bufferGrouping()`
+(`get_buffer_group_id` and `buffer_grouping` in Python) for exact membership.
+These APIs return or index by `CacheGroupId`, a public alias of `LifeCycleId`.
+`LayerGroupId` remains a lifecycle alias for legacy APIs.
+The legacy `getLayerGroupId(layerId)` and `layerGrouping()` are deprecated in
+documentation only: uniform layers retain their results, while a split layer
+fails a `TLLM_CHECK_WITH_INFO` (for the queried layer or the entire partition,
+respectively), exposed as `RuntimeError` in Python. No deprecation warnings are emitted.
+
+Page conversion reads `BufferAttr::scratchPagesPerBlock`, counting only the
+layer's buffers in that coalesced slot. Scratch capacity uses the maximum exact
+layer share across pools within each lifecycle. Layer counts are temporary
+layout calculations, not persistent `(layer, lifecycle)` attributes.
+
+Execution backends that share a page or state index across coupled buffers must
+validate that they use the same policy; independent policies require independent
+index tables. In particular, the default attention view couples K/V and scales,
+and Mamba kernels couple recurrent, convolution, PLE, and replay state. Native
+KVCM supports independent policies for callers that provide those indices.
+The sparse fetch path configured by `sparseCache` still uses one group per layer;
+its setup uses the checked layer lookup and rejects split layers.
+
+SSM lifecycles are deduplicated by policy, so a manager can have up to three SSM
+groups; iterate `ssmLifeCycleIds()` for working state and use
+`requiredSsmLifeCycleId()` only for prefix pruning.
+Optional SSM checkpoints require an exact endpoint match, including partial
+blocks. Fresh or unselected recurrent state must be initialized or reconstructed
+by the caller before inference. A pool-group index is a storage-layout index and is
 not interchangeable with a layer ID or lifecycle ID.
 
 Sparsity is a property of a pool, not of a lifecycle: a layer's buffers may be

cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.cpp

@@ -490,7 +490,7 @@ std::vector<SharedPtr<Block>> Block::clearStaleBlocksAfterPageUnlink(
 
     // If this is a full-attention block or a sink block: evict subtree.
     // Mirrors Python: pages = remove_subtree(self)
-    if (alc && alc->requiresPrefixCoverage()
+    if (alc && requiresPrefixCoverage(*alc)
         && (!alc->windowSize.has_value() || block.ordinal() < BlockOrdinal{alc->numSinkBlocks}))
     {
         pruneStart = block.prev;
@@ -899,7 +899,7 @@ std::vector<BlockRadixTree::MatchResult> BlockRadixTree::pruneMatch(
         bool shortened = false;
         for (auto [lcId, attn] : attnLcs)
         {
-            if (!attn->requiresPrefixCoverage())
+            if (!requiresPrefixCoverage(*attn))
             {
                 continue;
             }
@@ -973,7 +973,7 @@ BlockRadixTree::ReuseMatch BlockRadixTree::match(
     auto rawMatched = matchTokenPath(reuseScope, tokens, knownNoDigest, enablePartialMatch);
     // Content-divergence depth is measured before page or recurrent-snapshot pruning.
     int const numReusableTokensBeforePruning = numMatchedTokens(rawMatched, mTokensPerBlock);
-    auto const ssmLcId = mLifeCycles.ssmLifeCycleId();
+    auto const ssmLcId = mLifeCycles.requiredSsmLifeCycleId();
     // Page requirements depend on the final endpoint. Back off before pruning
     // so SWA coverage and recurrent snapshots are validated at that endpoint.
     if (backoff > 0)

cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/config.h

@@ -130,8 +130,21 @@ inline std::optional<size_t> cacheTierMaxQuota(CacheTierConfig const& cfg)
     return host != nullptr ? host->maxQuota : std::nullopt;
 }
 
+//! Controls whether a lifecycle participates in prefix matching and page reuse.
+enum class ReusePolicy
+{
+    //! Commit reusable pages and require live-window coverage when matching a prefix.
+    REQUIRED = 0,
+    //! Keep request-local working pages; never publish or reuse them through the radix tree.
+    PRIVATE = 1,
+    //! Commit reusable pages without constraining the REQUIRED prefix match. The caller selects
+    //! available groups at resume; missing or unselected groups use fresh pages for recomputation.
+    //! These pages may be evicted independently without removing the REQUIRED prefix.
+    OPTIONAL = 2,
+};
+
 // ---------------------------------------------------------------------------
-// Buffer configuration (one KV buffer inside an attention layer).
+// Buffer configuration (one buffer inside an attention or SSM layer).
 // ---------------------------------------------------------------------------
 
 // ---------------------------------------------------------------------------
@@ -193,6 +206,15 @@ struct BufferConfig
     //! densely packed, so a padded page is not expressible.
     int sparseUnitChunks = 1;
 
+    //! Reuse policy of this buffer; nullopt inherits the layer's policy. Overrides can
+    //! place buffers of one attention or SSM layer in different lifecycle groups.
+    std::optional<ReusePolicy> reusePolicy;
+
+    [[nodiscard]] ReusePolicy effectiveReusePolicy(ReusePolicy layerPolicy) const noexcept
+    {
+        return reusePolicy.value_or(layerPolicy);
+    }
+
     void validate() const;
 
     //! Where this buffer's selection units sit inside a page, given how many the layer
@@ -292,19 +314,6 @@ inline void validateNoDuplicateBufferRoles(std::vector<BufferConfig> const& buff
 
 } // namespace detail
 
-//! Controls whether an attention lifecycle participates in prefix matching and page reuse.
-enum class AttentionReusePolicy
-{
-    //! Commit reusable pages and require live-window coverage when matching a prefix.
-    REQUIRED = 0,
-    //! Keep request-local working pages; never publish or reuse them through the radix tree.
-    PRIVATE = 1,
-    //! Commit reusable pages without constraining the REQUIRED prefix match. The caller selects
-    //! available groups at resume; missing or unselected groups use fresh pages for recomputation.
-    //! These pages may be evicted independently without removing the REQUIRED prefix.
-    OPTIONAL = 2,
-};
-
 struct AttentionLayerConfig
 {
     static constexpr LayerType type = LayerType::ATTENTION;
@@ -319,7 +328,8 @@ struct AttentionLayerConfig
     std::optional<int> numSinkTokens;
 
     // Private working pages never constrain prefix matching or enter the tree.
-    AttentionReusePolicy reusePolicy = AttentionReusePolicy::REQUIRED;
+    // Every buffer of an attention layer uses this policy.
+    ReusePolicy reusePolicy = ReusePolicy::REQUIRED;
     //! The sparse fetch path for this layer's sparse buffers, if it has one.
     std::optional<SparseCacheConfig> sparseCache;
 
@@ -341,6 +351,8 @@ struct SsmLayerConfig
 
     LayerId layerId = 0;
     std::vector<BufferConfig> buffers;
+    // Default policy of the layer's buffers; BufferConfig::reusePolicy may override it per buffer.
+    ReusePolicy reusePolicy = ReusePolicy::REQUIRED;
 
     void validate() const
     {
@@ -362,6 +374,11 @@ struct SsmLayerConfig
 
 using LayerConfig = std::variant<AttentionLayerConfig, SsmLayerConfig>;
 
+inline std::vector<BufferConfig> const& layerBuffers(LayerConfig const& layer)
+{
+    return std::visit([](auto const& cfg) -> std::vector<BufferConfig> const& { return cfg.buffers; }, layer);
+}
+
 // ---------------------------------------------------------------------------
 // KVCacheDesc — describes one KV cache request's capacity and history length.
 // Mirrors _config.py::KVCacheDesc.