Feat/pytorch vswa kvcachemanager

cpp/tensorrt_llm/pybind/batch_manager/kvCacheManager.cpp

@@ -326,7 +326,7 @@ void tb::kv_cache_manager::KVCacheManagerBindings::initBindings(py::module_& m)
         .def_static("calculate_max_num_blocks", &tbk::BaseKVCacheManager::calculateMaxNumBlocks, py::arg("config"),
             py::arg("is_cross_attention"), py::arg("dtype"), py::arg("model_config"), py::arg("world_config"),
             py::arg("window_size_to_layers"), py::arg("allotted_primary_mem_bytes"),
-            py::arg("allotted_secondary_mem_bytes"), py::arg("extra_cost_memory"), py::arg("kvFactor"))
+            py::arg("allotted_secondary_mem_bytes"), py::arg("extra_cost_memory"), py::arg("kv_factor"))
         .def("allocate_pools", &BaseKVCacheManager::allocatePools)
         .def("release_pools", &BaseKVCacheManager::releasePools)
         .def("start_scheduling", &BaseKVCacheManager::startScheduling)

tensorrt_llm/_torch/model_config.py

@@ -8,6 +8,7 @@
 
 from tensorrt_llm import logger
 from tensorrt_llm._utils import torch_dtype_to_binding
+from tensorrt_llm.bindings import LayerType as LayerTypeCpp
 from tensorrt_llm.functional import AllReduceStrategy
 from tensorrt_llm.logger import logger
 from tensorrt_llm.mapping import Mapping
@@ -270,11 +271,19 @@ def from_pretrained(cls,
         model_config._frozen = True
         return model_config
 
-    def get_bindings_model_config(self) -> "ModelConfigCpp":
+    def get_bindings_model_config(self,
+                                  tokens_per_block: Optional[int] = None
+                                  ) -> "ModelConfigCpp":
         """
         This method is used to construct the bindings config for the model.
         Currently it adheres to gptJsonConfig.cpp::createModelConfig, which assumes
         that an engine has been created.
+
+        Args:
+            tokens_per_block: The number of tokens per block. Please note that in PyTorch flow tokens_per_block is not available in the model config, instead it is defined in the executor config.
+
+        Returns:
+            The bindings model config.
         """
         # TODO smor- this isn't robust, and currently tested for LlamaConfig only
         # TODO smor- currently assuming no rnn layers, no MOE
@@ -293,6 +302,12 @@ def get_bindings_model_config(self) -> "ModelConfigCpp":
             hidden_size=hidden_size,
             data_type=torch_dtype_to_binding(
                 self.pretrained_config.torch_dtype))
+        if tokens_per_block is None:
+            logger.warning(
+                f"tokens_per_block is not set, using default value {model_config_cpp.tokens_per_block}"
+            )
+        else:
+            model_config_cpp.tokens_per_block = tokens_per_block
 
         mlp_hidden_size = None
         if self.pretrained_config.intermediate_size is not None:
@@ -323,6 +338,11 @@ def get_bindings_model_config(self) -> "ModelConfigCpp":
         model_config_cpp.mlp_hidden_size = mlp_hidden_size
         model_config_cpp.size_per_head = head_size
 
+        # NOTE: this method is not robust, for Gemma3ForCausalLM only
+        layer_types = self.get_layer_types()
+        if layer_types is not None:
+            model_config_cpp.layer_types = layer_types
+
         return model_config_cpp
 
     def _infer_nemotron_ffn_mult(self):
@@ -339,3 +359,18 @@ def _infer_nemotron_ffn_mult(self):
             biggest_ffn_mult, self.pretrained_config.hidden_size)
 
         return mlp_hidden_size
+
+    def get_layer_types(self) -> Optional[List[LayerTypeCpp]]:
+        """
+        This method is a hack to support the effort to switch to KvCacheManagerCpp.
+        Currently, it is only tested for Gemma3ForCausalLM. For other models, it will return None.
+        """
+        if self.pretrained_config.architectures[0] in ["Gemma3ForCausalLM"]:
+            logger.debug(
+                f"Setting layer types for {self.pretrained_config.architectures}"
+            )
+            return [
+                LayerTypeCpp.ATTENTION,
+            ] * self.pretrained_config.num_hidden_layers
+        else:
+            return None

tensorrt_llm/_torch/pyexecutor/_util.py

@@ -340,6 +340,13 @@ def _create_kv_cache_manager(
                 spec_config=spec_config,
             )
         else:
+            # NOTE: this is a workaround for VSWA to switch to calculate_max_num_blocks_from_cpp in KVCahceManager
+            is_vswa = executor_config.kv_cache_config.max_attention_window is not None and len(
+                set(executor_config.kv_cache_config.max_attention_window)) > 1
+            binding_model_config = model_engine.model.model_config.get_bindings_model_config(
+                tokens_per_block=executor_config.tokens_per_block
+            ) if is_vswa else None
+
             kv_cache_manager = KVCacheManager(
                 executor_config.kv_cache_config,
                 tensorrt_llm.bindings.internal.batch_manager.CacheType.SELF,
@@ -352,7 +359,8 @@ def _create_kv_cache_manager(
                 mapping=mapping,
                 dtype=kv_cache_dtype,
                 spec_config=spec_config,
-            )
+                max_num_tokens=executor_config.max_num_tokens,
+                model_config=binding_model_config)
         # KVCacheManager (Non-draft) modifies the max_seq_len field, update it to executor_config
         if model_engine.kv_cache_manager_key == ResourceManagerType.KV_CACHE_MANAGER:
             executor_config.max_seq_len = kv_cache_manager.max_seq_len

tensorrt_llm/_torch/pyexecutor/py_executor_creator.py

@@ -319,6 +319,13 @@ def create_py_executor(
 
     if executor_config.enable_chunked_context:
         chunk_unit_size = executor_config.tokens_per_block
+        max_attention_window = executor_config.kv_cache_config.max_attention_window
+        if max_attention_window and max_seq_len > min(max_attention_window):
+            # maxKvStepSizeInFmha = 256
+            chunk_unit_size = max(256, chunk_unit_size)
+            logger.info(
+                f"ChunkUnitSize is set to {chunk_unit_size} as sliding window attention is used."
+            )
         chunking_policy = (
             executor_config.scheduler_config.context_chunking_policy
             if executor_config.scheduler_config.context_chunking_policy