diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index f4189f22ed92..5ae9ce95fbd0 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -1301,7 +1301,7 @@ def _create_cross_kv_cache_manager( max_seq_len=max_seq_len, max_batch_size=self._max_batch_size, spec_config=None, - sparse_attn_config=None, + sparse_attention_config=None, max_num_tokens=self._max_num_tokens, max_beam_width=1, kv_connector_manager=None, diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index e1551c88edd8..4e5cac179c02 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -5234,6 +5234,10 @@ def _prepare_tp_inputs_encoder( # Build a fresh, no-cache attention metadata for the encoder # pass. We do not reuse ``self.attn_metadata`` because that # object is bound to the decoder's KV-cache manager. + sparse_metadata_params = ( + self.sparse_attention_config.to_sparse_metadata_params( + pretrained_config=self.model.model_config.pretrained_config) + if self.sparse_attention_config is not None else None) encoder_attn_metadata = self.attn_backend.Metadata( max_num_requests=self.batch_size, max_num_tokens=self.max_num_tokens, @@ -5244,7 +5248,7 @@ def _prepare_tp_inputs_encoder( enable_flash_mla=self.model.model_config.enable_flash_mla, enable_context_mla_with_cached_kv=False, cache_indirection=None, - sparse_attention_config=self.sparse_attention_config, + sparse_metadata_params=sparse_metadata_params, num_heads_per_kv=1, ) assert isinstance(