[None][feat] Support Gemma4 multi-head_dim pools and host-side slicing to provide local view to Triton kernels for SWA

cpp/tensorrt_llm/batch_manager/cacheFormatter.cpp

@@ -696,6 +696,11 @@ void CacheFormatter::unformat(tensorrt_llm::batch_manager::TransferSession& sess
     {
         NVTX3_SCOPED_RANGE(formatInputRecvBuffer);
 
+        // TODO(disagg-multi-dtype): pool 0's dtype is treated as canonical for the wire
+        // transport here.  Pools with differing dtypes are rejected up-front in
+        // CacheTransBufferManager's constructor (see cacheTransBuffer.cpp).  When
+        // per-pool dtype dispatch lands, this single dataType variable must be replaced
+        // with a per-pool lookup keyed by the source pool of each block.
         auto dataType = mCacheManager->getPrimaryPool(0)->getDataType();
         bool layerWise = common::getEnvDisaggLayerwise() && numKvPools == 1;
         if (layerWise)

cpp/tensorrt_llm/batch_manager/cacheTransBuffer.cpp

@@ -253,6 +253,28 @@ CacheTransBufferManager::CacheTransBufferManager(
 {
     // TODO: FP4 dataSize
     TLLM_CHECK(mCacheManager);
+    // TODO(disagg-multi-dtype): Per-pool dtype dispatch in formatter / transfer buffer
+    // not yet implemented.  Disagg currently picks pool 0's dtype as the canonical
+    // transport type (above), so any KV pool with a different dtype would be silently
+    // miscoerced on the wire.  Fail loudly until per-pool dispatch lands.  We restrict
+    // the comparison to KV pools (getNumPools(false, false)) since block-scale and
+    // indexer-K pools legitimately have their own dtypes and travel through their own
+    // code paths.
+    if (!transferIndexerKCache)
+    {
+        auto const numKvPools = mCacheManager->getBlockManager().getNumPools(
+            /*includeBlockScalePools=*/false, /*includeIndexerKCachePools=*/false);
+        auto const dtype0 = mCacheManager->getPrimaryPool(0)->getDataType();
+        for (SizeType32 i = 1; i < numKvPools; ++i)
+        {
+            auto const dtypeI = mCacheManager->getPrimaryPool(i)->getDataType();
+            TLLM_CHECK_WITH_INFO(dtypeI == dtype0,
+                "Disaggregated KV cache transfer does not yet support pools with differing dtypes "
+                "(pool 0 dtype=%d, pool %d dtype=%d). TODO(disagg-multi-dtype): per-pool dtype "
+                "dispatch in formatter.",
+                static_cast<int>(dtype0), i, static_cast<int>(dtypeI));
+        }
+    }
     TLLM_LOG_INFO("CacheTransBufferManager created for KV cache");
 }
 

cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp

@@ -346,6 +346,14 @@ void tb::kv_cache_manager::KVCacheManagerBindings::initBindings(nb::module_& m)
     nb::enum_<tbk::LinearAttentionMetadata::LinearCacheType>(m, "LinearCacheType")
         .value("RECURRENT_STATES", tbk::LinearAttentionMetadata::LinearCacheType::kRecurrentStates);
 
+    nb::class_<tbk::PoolConfiguration>(m, "PoolConfiguration")
+        .def(nb::init<>())
+        .def(nb::init<SizeType32, SizeType32, nvinfer1::DataType>(), nb::arg("window_size"), nb::arg("size_per_head"),
+            nb::arg("dtype"))
+        .def_rw("window_size", &tbk::PoolConfiguration::windowSize)
+        .def_rw("size_per_head", &tbk::PoolConfiguration::sizePerHead)
+        .def_rw("dtype", &tbk::PoolConfiguration::dtype);
+
     nb::class_<tbk::PrefixReuseSummary>(m, "PrefixReuseSummary")
         .def(nb::init<>())
         .def_ro("reusable_blocks_allocated", &tbk::PrefixReuseSummary::reusableBlocksAllocated)
@@ -415,6 +423,7 @@ void tb::kv_cache_manager::KVCacheManagerBindings::initBindings(nb::module_& m)
             nb::arg("world_config"), nb::arg("window_size_to_layers"), nb::arg("allotted_primary_mem_bytes"),
             nb::arg("allotted_secondary_mem_bytes"), nb::arg("extra_cost_memory"), nb::arg("kv_factor"),
             nb::arg("max_batch_size"), nb::arg("linear_attention_metadata") = std::nullopt,
+            nb::arg("pool_configurations") = std::vector<tbk::PoolConfiguration>{},
             nb::call_guard<nb::gil_scoped_release>())
         .def("allocate_pools", &BaseKVCacheManager::allocatePools, nb::call_guard<nb::gil_scoped_release>())
         .def("release_pools", &BaseKVCacheManager::releasePools, nb::call_guard<nb::gil_scoped_release>())
@@ -597,6 +606,18 @@ void tb::kv_cache_manager::KVCacheManagerBindings::initBindings(nb::module_& m)
         .def("analyze_prefix_reuse", &BaseKVCacheManager::analyzePrefixReuse, nb::arg("unique_tokens"),
             nb::arg("llm_request"), nb::call_guard<nb::gil_scoped_release>())
         .def("get_cache_block_ids", &BaseKVCacheManager::getCacheBlockIds, nb::call_guard<nb::gil_scoped_release>())
+        .def(
+            "get_num_front_blocks_removed",
+            [](BaseKVCacheManager const& self, tb::LlmRequest::RequestIdType requestId, SizeType32 windowSize)
+            {
+                auto const& seq = self.getSequence(requestId);
+                // Per-window query.  windowSize is required (no aggregation): the Python
+                // wrapper in resource_manager.py provides a "default to first window"
+                // convenience layer; the C++/nanobind boundary stays explicit so that
+                // every caller is forced to think about which pool's counter it wants.
+                return seq.getNumFrontBlocksRemoved(windowSize);
+            },
+            nb::arg("request_id"), nb::arg("window_size"), nb::call_guard<nb::gil_scoped_release>())
         .def("get_batch_cache_block_ids", &BaseKVCacheManager::getBatchCacheBlockIds,
             nb::call_guard<nb::gil_scoped_release>())
         .def("flush_iteration_events", &BaseKVCacheManager::flushIterationEvents,
@@ -631,7 +652,8 @@ void tb::kv_cache_manager::KVCacheManagerBindings::initBindings(nb::module_& m)
                  std::vector<SizeType32> const&, nvinfer1::DataType, SizeType32, int64_t, SizeType32, SizeType32, bool,
                  tbk::CacheType, std::optional<tensorrt_llm::executor::RetentionPriority>,
                  std::shared_ptr<tbk::KVCacheEventManager>, bool, bool, std::shared_ptr<tbc::KvCacheConnectorManager>,
-                 bool, SizeType32, SizeType32, bool, std::optional<tbk::LinearAttentionMetadata>>(),
+                 bool, SizeType32, SizeType32, bool, std::optional<tbk::LinearAttentionMetadata>,
+                 std::vector<tbk::PoolConfiguration> const&>(),
             nb::arg("num_kv_heads_per_layer"), nb::arg("size_per_head"), nb::arg("tokens_per_block"),
             nb::arg("blocks_per_window"), nb::arg("max_num_sequences"), nb::arg("max_beam_width"),
             nb::arg("max_attention_window_vec"), nb::arg("dtype"), nb::arg("sink_token_length"), nb::arg("stream"),
@@ -641,14 +663,21 @@ void tb::kv_cache_manager::KVCacheManagerBindings::initBindings(nb::module_& m)
             nb::arg("copy_on_partial_reuse") = true, nb::arg("kv_connector_manager") = nullptr,
             nb::arg("enable_indexer_k_cache") = false, nb::arg("indexer_k_cache_quant_block_size") = 128,
             nb::arg("indexer_k_cache_index_head_dim") = 0, nb::arg("indexer_k_cache_use_fp4") = false,
-            nb::arg("linear_attention_metadata").none() = std::nullopt, nb::call_guard<nb::gil_scoped_release>())
+            nb::arg("linear_attention_metadata").none() = std::nullopt,
+            nb::arg("pool_configurations") = std::vector<tbk::PoolConfiguration>{},
+            nb::call_guard<nb::gil_scoped_release>())
         .def(
             "scheduling_has_free_blocks",
             [](tbk::KVCacheManager& self, SizeType32 numRequired, SizeType32 windowSize)
             { return self.getBlockManager().schedulingHasFreeBlocks(numRequired, windowSize); },
             nb::arg("num_required"), nb::arg("window_size"), nb::call_guard<nb::gil_scoped_release>())
         .def_prop_ro(
             "is_variable_window", [](tbk::KVCacheManager& self) { return self.getBlockManager().isVariableWindow(); })
+        // Per-pool introspection: lets Python discover (windowSize, sizePerHead, dtype) per
+        // hosted pool so a single KVCacheManager can host mixed-shape pools without a
+        // Python-side wrapper duplicating the layer->pool routing.
+        .def_prop_ro("pool_configurations",
+            [](tbk::KVCacheManager& self) { return self.getBlockManager().getPoolConfigurations(); })
         .def("copy_linear_attention_block", &tbk::KVCacheManager::copyLinearAttentionBlock, nb::arg("llm_request"),
             nb::call_guard<nb::gil_scoped_release>())
         .def("copy_linear_attention_block_batch", &tbk::KVCacheManager::copyLinearAttentionBlockBatch,