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

cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/AGENTS.md

@@ -65,14 +65,34 @@ 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. 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.
-All buffers within a layer share that policy. 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.
+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.
+The legacy `getLayerGroupId(layerId)` and `layerGrouping()` are deprecated in
+documentation only: uniform layers retain their results, while a split layer
+raises an ambiguity error (for the queried layer or the entire partition,
+respectively). 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.
+
+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

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).
 // ---------------------------------------------------------------------------
 
 struct BufferConfig
@@ -148,6 +161,15 @@ struct BufferConfig
     //! Buffers may coalesce only when their group and expanded size match.
     //! This does not change their lifecycle or shared slot allocation.
     int coalescingGroup = 0;
+
+    //! 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);
+    }
 };
 
 // ---------------------------------------------------------------------------
@@ -178,19 +200,6 @@ inline void validateNoDuplicateBufferRoles(std::vector<BufferConfig> const& buff
 
 } // namespace detail
 
-//! Controls whether a layer 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,
-};
-
 struct AttentionLayerConfig
 {
     static constexpr LayerType type = LayerType::ATTENTION;
@@ -205,6 +214,7 @@ struct AttentionLayerConfig
     std::optional<int> numSinkTokens;
 
     // Private working pages never constrain prefix matching or enter the tree.
+    // Every buffer of an attention layer uses this policy.
     ReusePolicy reusePolicy = ReusePolicy::REQUIRED;
 
     [[nodiscard]] std::optional<int> windowSize() const noexcept
@@ -235,6 +245,7 @@ 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
@@ -256,6 +267,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.

cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.cpp

@@ -30,6 +30,7 @@
 #include <map>
 #include <numeric>
 #include <stdexcept>
+#include <string>
 
 namespace tensorrt_llm::batch_manager::kv_cache_manager_v2
 {
@@ -72,10 +73,18 @@ std::vector<int> PageIndexConverter::operator()(
             // Scratch block: slot IDs come from ScratchDesc, not base_indices.
             int blockPos = ordinal - scratch->range.beg;
             int totalOffset = blockPos * scratchPages;
+            bool const crossesSlots = scale % scratchPages != 0;
+            if (crossesSlots)
+            {
+                // Unequal layer shares can put a block across two slots. Pack
+                // the layer's consecutive roles before resolving the slot ID,
+                // so a trailing role cannot wrap over an earlier block's data.
+                totalOffset += appliedLayerOffset % scratchPages;
+            }
             int slotIdx = totalOffset / scale;
             int slotId = scratch->slotIds[static_cast<size_t>(slotIdx)];
             int offset = totalOffset % scale;
-            index = slotId * scale + (offset + appliedLayerOffset) % scale;
+            index = slotId * scale + (crossesSlots ? offset : (offset + appliedLayerOffset) % scale);
         }
         else if (baseIndices[toSizeT(ordinal)] == kBadPageIndex.value())
         {
@@ -333,11 +342,9 @@ bool KvCacheManager::isSparse(LayerId layerId, DataRole role) const
 PageIndexConverter KvCacheManager::getPageIndexConverter(LayerId layerId, DataRole role) const
 {
     auto const& attr = mStorage->getBufferAttr(layerId, role);
-    auto const& layerAttr = mStorage->getLayerAttr(layerId);
     int scale = mStorage->mSlotToPageIndices.at(attr.lifeCycleId).at(attr.poolIndex);
     int offset = exactDiv(static_cast<int>(attr.offset), static_cast<int>(attr.size));
-    int scratchPages = layerAttr.slotUtil.at(attr.poolIndex);
-    return PageIndexConverter{scale, attr.expansion, offset, scratchPages};
+    return PageIndexConverter{scale, attr.expansion, offset, attr.scratchPagesPerBlock};
 }
 
 std::optional<bool> KvCacheManager::supportsIndexMode(PageIndexMode mode) const
@@ -461,16 +468,49 @@ std::vector<LayerId> KvCacheManager::layerIds() const
 
 LayerGroupId KvCacheManager::getLayerGroupId(LayerId layerId) const
 {
-    return mStorage->layerToLifeCycleIds().at(layerId);
+    auto const& lifeCycles = mStorage->layerToLifeCycleIds().at(layerId);
+    if (lifeCycles.size() != 1)
+    {
+        throw std::invalid_argument("Layer " + std::to_string(layerId) + " has buffers in "
+            + std::to_string(lifeCycles.size())
+            + " layer groups because its buffers use different reuse policies; look up its layer group per "
+              "buffer with get_buffer_group_id(layer_id, role) / getBufferGroupId(layerId, role)");
+    }
+    return lifeCycles.front();
+}
+
+LayerGroupId KvCacheManager::getBufferGroupId(LayerId layerId, DataRole const& role) const
+{
+    return mStorage->getBufferAttr(layerId, role).lifeCycleId;
 }
 
 TypedVec<LayerGroupId, std::vector<LayerId>> KvCacheManager::layerGrouping() const
 {
     LifeCycleId numLc = mLifeCycles.size();
     TypedVec<LayerGroupId, std::vector<LayerId>> result(numLc);
-    for (auto const& [lid, lc] : mStorage->layerToLifeCycleIds())
+    for (auto const& [lid, lifeCycles] : mStorage->layerToLifeCycleIds())
+    {
+        result.at(getLayerGroupId(lid)).push_back(lid);
+    }
+    return result;
+}
+
+TypedVec<LayerGroupId, std::vector<BufferId>> KvCacheManager::bufferGrouping() const
+{
+    TypedVec<LayerGroupId, std::vector<BufferId>> result(mLifeCycles.size());
+    std::unordered_map<LayerId, std::vector<BufferId>> buffersByLayer;
+    for (auto const& buffer : allBufferIds())
     {
-        result.at(lc).push_back(lid);
+        buffersByLayer[buffer.layerId].push_back(buffer);
+    }
+    // Keep layer order consistent with layerGrouping() for consumers that pair
+    // group membership with a declared list of global model layer IDs.
+    for (auto const& [layer, lifeCycles] : mStorage->layerToLifeCycleIds())
+    {
+        for (auto const& buffer : buffersByLayer.at(layer))
+        {
+            result.at(getBufferGroupId(buffer.layerId, buffer.role)).push_back(buffer);
+        }
     }
     return result;
 }

cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.h

@@ -229,15 +229,23 @@ class KvCacheManager : public std::enable_shared_from_this<KvCacheManager>
     int numLayers() const noexcept;
 
     std::vector<LayerId> layerIds() const;
+
+    //! Deprecated: use getBufferGroupId(). Returns the shared group of a uniform layer.
+    //! Throws std::invalid_argument if the queried layer spans several lifecycle groups.
     LayerGroupId getLayerGroupId(LayerId layerId) const;
 
-    // Layer grouping: layers with the same lifecycle share pool allocation.
-    // NOTE: the iteration order of the layer lists (and of the groups) is NOT
-    // part of the API contract and may differ across backends/runs. Do not rely
-    // on it to infer buffer/pool memory order — query poolGroupDescs()
-    // (PoolGroupDesc::pools[i].baseAddress + coalescedBuffers) for that.
+    //! Group of one buffer; throws std::out_of_range for an unknown buffer.
+    LayerGroupId getBufferGroupId(LayerId layerId, DataRole const& role) const;
+
+    //! Deprecated: use bufferGrouping(). Partitions uniform layers by lifecycle.
+    //! Throws std::invalid_argument if any layer spans several groups. Member order
+    //! is unspecified; use poolGroupDescs() to inspect physical buffer ordering.
     TypedVec<LayerGroupId, std::vector<LayerId>> layerGrouping() const;
 
+    //! Buffer membership indexed by lifecycle group ID. Each buffer appears once.
+    //! Member order is unspecified and does not describe physical memory ordering.
+    TypedVec<LayerGroupId, std::vector<BufferId>> bufferGrouping() const;
+
     // Iterator over all buffer identifiers. Mirrors Python's all_buffer_ids property.
     std::vector<BufferId> allBufferIds() const;
 

cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/lifeCycleRegistry.cpp

@@ -27,22 +27,21 @@ namespace tensorrt_llm::batch_manager::kv_cache_manager_v2
 // makeLifeCycle — factory dispatching on LayerConfig variant.
 // ---------------------------------------------------------------------------
 
-LifeCycle makeLifeCycle(LayerConfig const& layer, int tokensPerBlock)
+LifeCycle makeLifeCycle(LayerConfig const& layer, BufferConfig const& buffer, int tokensPerBlock)
 {
     return std::visit(
         [&](auto const& cfg) -> LifeCycle
         {
-            cfg.validate();
+            auto const policy = buffer.effectiveReusePolicy(cfg.reusePolicy);
             using T = std::decay_t<decltype(cfg)>;
             if constexpr (std::is_same_v<T, SsmLayerConfig>)
             {
-                return SsmLifeCycle{cfg.reusePolicy};
+                return SsmLifeCycle{policy};
             }
             else
             {
-                bool const isSparse = !cfg.buffers.empty() && cfg.buffers.front().isSparse;
                 return AttnLifeCycle::make(
-                    cfg.slidingWindowSize, cfg.numSinkTokens, tokensPerBlock, isSparse, cfg.reusePolicy);
+                    cfg.slidingWindowSize, cfg.numSinkTokens, tokensPerBlock, buffer.isSparse, policy);
             }
         },
         layer);
@@ -54,9 +53,8 @@ LifeCycle makeLifeCycle(LayerConfig const& layer, int tokensPerBlock)
 
 LifeCycleRegistry::LifeCycleRegistry(KVCacheManagerConfig const& config)
 {
-    for (auto const& layer : config.layers)
+    auto const registerLifeCycle = [this](LifeCycle const& lc)
     {
-        LifeCycle lc = makeLifeCycle(layer, config.tokensPerBlock);
         if (mLifeCycleIdMap.find(lc) == mLifeCycleIdMap.end())
         {
             check();
@@ -66,6 +64,14 @@ LifeCycleRegistry::LifeCycleRegistry(KVCacheManagerConfig const& config)
             if (std::holds_alternative<SsmLifeCycle>(lc))
                 mSsmLifeCycleIds.push_back(id);
         }
+    };
+    for (auto const& layer : config.layers)
+    {
+        std::visit([](auto const& cfg) { cfg.validate(); }, layer);
+        for (auto const& buffer : layerBuffers(layer))
+        {
+            registerLifeCycle(makeLifeCycle(layer, buffer, config.tokensPerBlock));
+        }
     }
     check();
 }

cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/lifeCycleRegistry.h

@@ -108,9 +108,9 @@ struct AttnLifeCycle
 };
 
 // ---------------------------------------------------------------------------
-// SsmLifeCycle — lifecycle for SSM (State Space Model) layers.
+// SsmLifeCycle — lifecycle for SSM (State Space Model) buffers.
 // All blocks before the last full block are stale (recurrent state).
-// SSM layers with the same reuse policy share a lifecycle.
+// SSM buffers with the same effective reuse policy share a lifecycle.
 // ---------------------------------------------------------------------------
 struct SsmLifeCycle
 {
@@ -178,8 +178,8 @@ inline HalfOpenRange<BlockOrdinal> computeScratchRange(
     return intersect(capStale, inputRange);
 }
 
-// Factory: create a LifeCycle from a LayerConfig variant.
-LifeCycle makeLifeCycle(LayerConfig const& layer, int tokensPerBlock);
+// Factory: create the LifeCycle that holds one buffer of a layer, using the buffer's effective reuse policy.
+LifeCycle makeLifeCycle(LayerConfig const& layer, BufferConfig const& buffer, int tokensPerBlock);
 
 // Integer id assigned to each unique LifeCycle.
 using LifeCycleId = StrongIndex<int, struct LifeCycleIdTag>;