diff --git a/tensorrt_llm/_torch/attention_backend/interface.py b/tensorrt_llm/_torch/attention_backend/interface.py index 2856a1769fe8..df9735463b9d 100644 --- a/tensorrt_llm/_torch/attention_backend/interface.py +++ b/tensorrt_llm/_torch/attention_backend/interface.py @@ -733,6 +733,10 @@ class AttentionSparseArgs: sparse_attn_indices: Optional[torch.Tensor] = None sparse_attn_offsets: Optional[torch.Tensor] = None sparse_attn_indices_block_size: int = 0 + # DeepSeek-V4 sparse-MLA only: per-token compressed top-k lengths and the + # base pointer of the compressed KV cache pool (compress_ratio > 1). + sparse_mla_topk_lens: Optional[torch.Tensor] = None + compressed_kv_cache_pool_ptr: Optional[int] = None @dataclass(kw_only=True, slots=True) diff --git a/tensorrt_llm/_torch/attention_backend/trtllm.py b/tensorrt_llm/_torch/attention_backend/trtllm.py index bf1268549778..1d409a181530 100644 --- a/tensorrt_llm/_torch/attention_backend/trtllm.py +++ b/tensorrt_llm/_torch/attention_backend/trtllm.py @@ -22,9 +22,9 @@ get_model_extra_attrs) from .interface import (AttentionBackend, AttentionForwardArgs, AttentionInputType, AttentionMask, AttentionMetadata, - AttentionSparseArgs, KVCacheParams, MLAParams, - PositionalEmbeddingParams, PredefinedAttentionMask, - RopeParams, merge_attention_forward_args) + KVCacheParams, MLAParams, PositionalEmbeddingParams, + PredefinedAttentionMask, RopeParams, + merge_attention_forward_args) # Enable TRTLLM-Gen attention backend by default. Set # TRTLLM_ENABLE_TRTLLM_GEN_ATTENTION=0 to force the thop.attention path. @@ -44,10 +44,7 @@ # ``thop.attention`` kwargs hard-wired to a literal at the call site (no # rich object owns them). Sync test enforces both the kwarg name and the # literal value. -_THOP_LITERALS: dict = { - "sparse_mla_topk_lens": None, - "compressed_kv_cache_pool_ptr": None, -} +_THOP_LITERALS: dict = {} @functools.cache @@ -114,6 +111,7 @@ class TrtllmAttentionMetadata(AttentionMetadata): is_spec_dec_dynamic_tree: bool = False # parameters required for spec-dec mode + max_total_draft_tokens: Optional[int] = None spec_decoding_position_offsets: Optional[torch.Tensor] = None # C++ attention op requires a 2-D position_offsets tensor and reads # sizes()[1] as the generation length / packed-mask row stride. @@ -1631,6 +1629,8 @@ def _run( spec_decoding_bl_tree_mask_offset=metadata. spec_decoding_bl_tree_mask_offset, spec_decoding_bl_tree_mask=metadata.spec_decoding_bl_tree_mask, + spec_decoding_target_max_draft_tokens=metadata. + max_total_draft_tokens, spec_bl_tree_first_sparse_mask_offset_kv=metadata. spec_bl_tree_first_sparse_mask_offset_kv, num_sparse_topk=metadata.num_sparse_topk, @@ -1719,14 +1719,11 @@ def _run( sparse_attn_offsets=forward_args.sparse.sparse_attn_offsets, sparse_attn_indices_block_size=forward_args.sparse. sparse_attn_indices_block_size, + sparse_mla_topk_lens=forward_args.sparse.sparse_mla_topk_lens, + compressed_kv_cache_pool_ptr=forward_args.sparse. + compressed_kv_cache_pool_ptr, - # --- Literals intentionally None (see _THOP_LITERALS) --- - # ``sparse_mla_topk_lens`` and ``compressed_kv_cache_pool_ptr`` - # stay as literal ``None`` until DeepSeek V4 sparse-MLA lands. - sparse_mla_topk_lens=None, - compressed_kv_cache_pool_ptr=None, - spec_decoding_target_max_draft_tokens=getattr( - metadata, 'max_total_draft_tokens', None), + # --- Literals intentionally in _THOP_LITERALS --- ) if self.print_skip_softmax_stat: @@ -1811,14 +1808,13 @@ def forward( forward_args) at_idx, at_off = self.sparse_attn_predict(q, k, metadata, forward_args) - forward_args.sparse = AttentionSparseArgs( - sparse_kv_indices=kv_idx, - sparse_kv_offsets=kv_off, - sparse_attn_indices=at_idx, - sparse_attn_offsets=at_off, - sparse_attn_indices_block_size=self.sparse_attention_config. - get_indices_block_size(), - ) + sparse_args = forward_args.sparse + sparse_args.sparse_kv_indices = kv_idx + sparse_args.sparse_kv_offsets = kv_off + sparse_args.sparse_attn_indices = at_idx + sparse_args.sparse_attn_offsets = at_off + sparse_args.sparse_attn_indices_block_size = ( + self.sparse_attention_config.get_indices_block_size()) # Compute FlashMLA tile-scheduler metadata once per forward pass. # The flag is reset in prepare_flash_mla() and update_for_spec_dec() to trigger