[None][feat] DSv4 prep: attention op plumbing

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)

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.
@@ -1719,14 +1717,13 @@ def _run(
                 sparse_attn_offsets=forward_args.sparse.sparse_attn_offsets,
                 sparse_attn_indices_block_size=forward_args.sparse.
                 sparse_attn_indices_block_size,
-
-                # --- 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,
+                sparse_mla_topk_lens=forward_args.sparse.sparse_mla_topk_lens,
+                compressed_kv_cache_pool_ptr=forward_args.sparse.
+                compressed_kv_cache_pool_ptr,
                 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