fix rebase breaks

Signed-off-by: yuhangh <[email protected]>

fix rebase bug.

Signed-off-by: Fanrong Li <[email protected]>

fix rebase bug.

Signed-off-by: Fanrong Li <[email protected]>

Author: heyuhhh

cpp/tensorrt_llm/common/attentionOp.h

@@ -494,10 +494,11 @@ class AttentionOp
             mPosShiftEnabled, mPagedContextFMHA, mFP8ContextFMHA, mFP8AttenOutput, mFP8ContextMLA, mFP8GenerationMLA,
             mChunkPrefillBufferBatchSize, mDenseContextFMHA, mHasFullAttentionMask, mIsSpecDecodingEnabled,
             mUseSpecDecoding, mIsSpecDecTree, mSpecDecodingIsGenerationLengthVariable, mSpecDecodingMaxGenerationLength,
-            mIsMLAEnabled, mIsGenerationMLA, mUseGenFlashMLA, mUseSparseAttention, mMLAParams.data(), mCpSize, mCpRank, mCpGroup,
-            mNumAttnHeads, mNumAttnKVHeads, mNumKVHeadsOrigin, mAttnTpSize, mAttnTpRank, mAttnCpSize, mAttnCpRank,
-            mUlyssesMQABroadcast, mEnableContextFMHA, mFMHAForceFP32Acc, mMultiBlockMode, mEnableXQA, mUseKVCache,
-            mSkipAttn, mFuseFp4Quant, mRuntimeSparseAttentionParams.data(), mNbMultiBlockSemaphores, mAttentionChunkSize.value_or(-1));
+            mIsMLAEnabled, mIsGenerationMLA, mUseGenFlashMLA, mUseSparseAttention, mMLAParams.data(), mCpSize, mCpRank,
+            mCpGroup, mNumAttnHeads, mNumAttnKVHeads, mNumKVHeadsOrigin, mAttnTpSize, mAttnTpRank, mAttnCpSize,
+            mAttnCpRank, mUlyssesMQABroadcast, mEnableContextFMHA, mFMHAForceFP32Acc, mMultiBlockMode, mEnableXQA,
+            mUseKVCache, mSkipAttn, mFuseFp4Quant, mRuntimeSparseAttentionParams.data(), mNbMultiBlockSemaphores,
+            mAttentionChunkSize.value_or(-1));
     };
 
 private:

cpp/tensorrt_llm/nanobind/thop/bindings.cpp

@@ -45,13 +45,13 @@ void initBindings(nb::module_& m)
         nb::arg("q_scaling"), nb::arg("rotary_embedding_int_params"), nb::arg("rotary_embedding_base"),
         nb::arg("rotary_embedding_scales"), nb::arg("rotary_embedding_max_position_info"),
         nb::arg("use_paged_context_fmha"), nb::arg("attention_input_type") = std::nullopt, nb::arg("is_mla_enable"),
-        nb::arg("chunked_prefill_buffer_batch_size") = std::nullopt, 
-        nb::arg("q_lora_rank") = std::nullopt, nb::arg("kv_lora_rank") = std::nullopt,
-        nb::arg("qk_nope_head_dim") = std::nullopt, nb::arg("qk_rope_head_dim") = std::nullopt,
-        nb::arg("v_head_dim") = std::nullopt, nb::arg("mrope_rotary_cos_sin") = std::nullopt,
-        nb::arg("mrope_position_deltas") = std::nullopt, nb::arg("attention_chunk_size") = std::nullopt,
-        nb::arg("softmax_stats_tensor") = std::nullopt, nb::arg("spec_decoding_bool_params"),
-        nb::arg("spec_decoding_tensor_params"), nb::arg("sparse_attention_params") = std::nullopt,
-        "Multi-head attention operation", nb::call_guard<nb::gil_scoped_release>());
+        nb::arg("chunked_prefill_buffer_batch_size") = std::nullopt, nb::arg("q_lora_rank") = std::nullopt,
+        nb::arg("kv_lora_rank") = std::nullopt, nb::arg("qk_nope_head_dim") = std::nullopt,
+        nb::arg("qk_rope_head_dim") = std::nullopt, nb::arg("v_head_dim") = std::nullopt,
+        nb::arg("mrope_rotary_cos_sin") = std::nullopt, nb::arg("mrope_position_deltas") = std::nullopt,
+        nb::arg("attention_chunk_size") = std::nullopt, nb::arg("softmax_stats_tensor") = std::nullopt,
+        nb::arg("spec_decoding_bool_params"), nb::arg("spec_decoding_tensor_params"),
+        nb::arg("sparse_attention_params") = std::nullopt, "Multi-head attention operation",
+        nb::call_guard<nb::gil_scoped_release>());
 }
 } // namespace tensorrt_llm::nanobind::thop

cpp/tensorrt_llm/pybind/thop/bindings.cpp

@@ -45,13 +45,13 @@ void initBindings(pybind11::module_& m)
         py::arg("q_scaling"), py::arg("rotary_embedding_int_params"), py::arg("rotary_embedding_base"),
         py::arg("rotary_embedding_scales"), py::arg("rotary_embedding_max_position_info"),
         py::arg("use_paged_context_fmha"), py::arg("attention_input_type") = std::nullopt, py::arg("is_mla_enable"),
-        py::arg("chunked_prefill_buffer_batch_size") = std::nullopt, 
-        py::arg("q_lora_rank") = std::nullopt, py::arg("kv_lora_rank") = std::nullopt,
-        py::arg("qk_nope_head_dim") = std::nullopt, py::arg("qk_rope_head_dim") = std::nullopt,
-        py::arg("v_head_dim") = std::nullopt, py::arg("mrope_rotary_cos_sin") = std::nullopt,
-        py::arg("mrope_position_deltas") = std::nullopt, py::arg("attention_chunk_size") = std::nullopt,
-        py::arg("softmax_stats_tensor") = std::nullopt, py::arg("spec_decoding_bool_params"),
-        py::arg("spec_decoding_tensor_params"), py::arg("sparse_attention_params") = std::nullopt,
-        "Multi-head attention operation", py::call_guard<py::gil_scoped_release>());
+        py::arg("chunked_prefill_buffer_batch_size") = std::nullopt, py::arg("q_lora_rank") = std::nullopt,
+        py::arg("kv_lora_rank") = std::nullopt, py::arg("qk_nope_head_dim") = std::nullopt,
+        py::arg("qk_rope_head_dim") = std::nullopt, py::arg("v_head_dim") = std::nullopt,
+        py::arg("mrope_rotary_cos_sin") = std::nullopt, py::arg("mrope_position_deltas") = std::nullopt,
+        py::arg("attention_chunk_size") = std::nullopt, py::arg("softmax_stats_tensor") = std::nullopt,
+        py::arg("spec_decoding_bool_params"), py::arg("spec_decoding_tensor_params"),
+        py::arg("sparse_attention_params") = std::nullopt, "Multi-head attention operation",
+        py::call_guard<py::gil_scoped_release>());
 }
 } // namespace tensorrt_llm::pybind::thop

tensorrt_llm/_torch/attention_backend/sparse/rocket.py

@@ -15,7 +15,7 @@
 from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager
 from tensorrt_llm._utils import get_size_in_bytes, next_power_of_two
 from tensorrt_llm.bindings import DataType
-from tensorrt_llm.bindings.executor import ExecutorConfig, KvCacheConfig
+from tensorrt_llm.bindings.executor import KvCacheConfig
 from tensorrt_llm.bindings.internal.batch_manager import \
     CacheType as CacheTypeCpp
 from tensorrt_llm.mapping import Mapping
@@ -843,9 +843,7 @@ def compute_page_count(self, token_count: int, tokens_per_page: int) -> int:
 
     @staticmethod
     def get_cache_size_per_token(model_config: ModelConfig,
-                                 executor_config: ExecutorConfig,
-                                 mapping: Mapping):
-        sparse_attn_config = executor_config.sparse_attention_config
+                                 tokens_per_block: int, mapping: Mapping):
         # get kv cache dtype bytes
         mem_per_token = 2
         quant_config = model_config.quant_config
@@ -875,7 +873,7 @@ def get_cache_size_per_token(model_config: ModelConfig,
 
         # K and V
         # 2 for K and V, 2 * kt_tokens_per_block / tokens_per_block for KT cache
-        tokens_per_block = executor_config.tokens_per_block
+        sparse_attn_config = model_config.sparse_attention_config
         kt_tokens_per_block = next_power_of_two(
             math.ceil(tokens_per_block / sparse_attn_config.page_size))
         kv_factor = 2 + 2 * kt_tokens_per_block / tokens_per_block

tensorrt_llm/_torch/pyexecutor/_util.py

@@ -10,7 +10,8 @@
 from tensorrt_llm._utils import str_dtype_to_binding, torch_dtype_to_str
 from tensorrt_llm.bindings.executor import DecodingMode
 from tensorrt_llm.llmapi.llm_args import (PeftCacheConfig, SamplerType,
-                                          SpeculativeConfig, SparseAttentionConfig)
+                                          SparseAttentionConfig,
+                                          SpeculativeConfig)
 from tensorrt_llm.logger import logger
 from tensorrt_llm.lora_helper import (LoraConfig,
                                       get_default_trtllm_modules_to_hf_modules)
@@ -40,10 +41,9 @@
 GB = 1 << 30
 
 
-def get_kv_cache_manager_cls(model_config: ModelConfig,
-                             executor_config: ExecutorConfig):
+def get_kv_cache_manager_cls(model_config: ModelConfig):
     config = model_config.pretrained_config
-    sparse_attn_config = executor_config.sparse_attention_config
+    sparse_attn_config = model_config.sparse_attention_config
     if is_mla(config):
         return KVCacheManager
     elif is_nemotron_hybrid(config):
@@ -93,46 +93,7 @@ def __init__(
         self._max_seq_len = max_seq_len
         self._max_batch_size = max_batch_size
         self._kv_cache_manager_cls = get_kv_cache_manager_cls(
-            model_engine.model.model_config, executor_config)
-
-    @staticmethod
-    def _get_cache_size_per_token(model_config: ModelConfig,
-                                  mapping: Mapping) -> int:
-        mem_per_token = 2
-        quant_config = model_config.quant_config
-        if quant_config is not None and quant_config.quant_mode.has_fp8_kv_cache(
-        ):
-            mem_per_token = 1
-
-        config = model_config.pretrained_config
-
-        num_key_value_heads = getattr(config, 'num_key_value_heads',
-                                      config.num_attention_heads)
-        if isinstance(num_key_value_heads, Iterable):
-            num_key_value_heads = sum(num_key_value_heads) / len(
-                num_key_value_heads)
-
-        mla = is_mla(config)
-        tp_size = 1 if mapping.enable_attention_dp else mapping.tp_size
-
-        kv_factor = 2
-        if mla:
-            # MLA has kv_lora_rank and qk_rope_head_dim
-            head_dim = config.kv_lora_rank + config.qk_rope_head_dim
-            kv_factor = 1
-        else:
-            _head_dim = getattr(config, 'head_dim', None)
-            if not isinstance(_head_dim, int):
-                _head_dim = config.hidden_size // config.num_attention_heads
-            head_dim = _head_dim * num_key_value_heads // tp_size
-
-        # provide at least 1 layer to prevent division by zero cache size
-        num_attention_layers = max(
-            len(mapping.pp_layers(model_config.get_num_attention_layers())), 1)
-        mem_per_token *= num_attention_layers * head_dim
-        # K and V
-        mem_per_token *= kv_factor
-        return mem_per_token
+            model_engine.model.model_config)
 
     def _get_free_gpu_memory_fraction(self) -> float:
         fraction = self._kv_cache_config.free_gpu_memory_fraction
@@ -144,11 +105,11 @@ def _get_kv_size_per_token(self):
         model_config = self._model_engine.model.model_config
         mapping = self._mapping
         kv_size_per_token = self._kv_cache_manager_cls.get_cache_size_per_token(
-            model_config, self._executor_config, mapping)
+            model_config, self._tokens_per_block, mapping)
         if self._draft_model_engine is not None:
             draft_model_config = self._draft_model_engine.model.model_config
             kv_size_per_token += self._kv_cache_manager_cls.get_cache_size_per_token(
-                draft_model_config, self._executor_config, mapping)
+                draft_model_config, self._tokens_per_block, mapping)
         return kv_size_per_token
 
     def _cal_max_memory(self, peak_memory, total_gpu_memory, fraction,

tensorrt_llm/_torch/pyexecutor/resource_manager.py

@@ -10,7 +10,6 @@
 import tensorrt_llm
 import tensorrt_llm.bindings
 from tensorrt_llm.bindings.BuildInfo import ENABLE_MULTI_DEVICE
-from tensorrt_llm.bindings.executor import ExecutorConfig
 from tensorrt_llm.lora_helper import LoraConfig
 from tensorrt_llm.lora_manager import LoraManager, LoraModelConfig
 from tensorrt_llm.sampling_params import SamplingParams
@@ -279,11 +278,7 @@ def append_to_kv_heads_per_layer(num_kv_heads_per_layer: List[int],
             # Standard case: use original Python implementation
             self.blocks_in_primary_pool, self.blocks_in_secondary_pool = self.calculate_max_num_blocks(
                 kv_cache_config=kv_cache_config,
-                head_dim=head_dim,
-                tokens_per_block=tokens_per_block,
                 mapping=mapping,
-                dtype=dtype,
-                kv_factor=self.kv_factor,
             )
             blocks_per_window = {
                 self.max_attention_window_vec[0]:
@@ -549,8 +544,7 @@ def calculate_scaling_factor_size_bytes(
 
     @staticmethod
     def get_cache_size_per_token(model_config: ModelConfig,
-                                 executor_config: ExecutorConfig,
-                                 mapping: Mapping):
+                                 tokens_per_block: int, mapping: Mapping):
         # get kv cache dtype bytes
         mem_per_token = 2
         quant_config = model_config.quant_config
@@ -605,13 +599,9 @@ def get_cache_bytes_per_token(self):
                 scaling_factor_dtype=DataType.FP8)
         return cache_size_bytes_per_token
 
-    def calculate_max_num_blocks(self,
-                                 kv_cache_config: KvCacheConfigCpp,
-                                 head_dim: int,
-                                 tokens_per_block: int,
-                                 mapping: Mapping,
-                                 dtype: DataType,
-                                 kv_factor: int = 2):
+    def calculate_max_num_blocks(self, kv_cache_config: KvCacheConfigCpp,
+                                 mapping: Mapping):
+        tokens_per_block = self.tokens_per_block
         free_mem_fraction = (kv_cache_config.free_gpu_memory_fraction
                              if kv_cache_config.free_gpu_memory_fraction
                              is not None else 0.9)