[None][feat] Introduce the GDN MTP replay V2 design

tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py

@@ -35,6 +35,60 @@
 _FUSED_GDN_REPLAY_WORK_ITEMS_MAX_BATCH_SIZE = 256
 
 
+@triton.jit
+def _snapshot_gdn_replay_work_items_kernel(
+    state_indices,
+    history_len,
+    work_items,
+    n_writes_output,
+    num_decodes,
+    work_item_width: tl.constexpr,
+    position_field: tl.constexpr,
+    cache_slot_field: tl.constexpr,
+    history_len_field: tl.constexpr,
+    cache_buf_idx_field: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """Snapshot single-buffer history lengths before the layer replay loop."""
+    offsets = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+    active = offsets < num_decodes
+    slots = tl.load(state_indices + offsets, mask=active, other=0)
+    lengths = tl.load(history_len + slots, mask=active, other=0)
+    output_base = work_items + offsets * work_item_width
+    tl.store(output_base + position_field, offsets, mask=active)
+    tl.store(output_base + cache_slot_field, slots, mask=active)
+    tl.store(output_base + history_len_field, lengths, mask=active)
+    tl.store(output_base + cache_buf_idx_field, 0, mask=active)
+    if tl.program_id(0) == 0:
+        tl.store(n_writes_output, 0)
+
+
+def _build_gdn_replay_work_items(
+    state_indices: torch.Tensor,
+    history_len: torch.Tensor,
+    work_items: torch.Tensor,
+    n_writes: torch.Tensor,
+) -> None:
+    """Build the immutable inputs consumed by the split GDN MTP commit."""
+    num_decodes = state_indices.shape[0]
+    block_size = 256
+    _snapshot_gdn_replay_work_items_kernel[(triton.cdiv(
+        num_decodes, block_size), )](
+            state_indices,
+            history_len,
+            work_items,
+            n_writes,
+            num_decodes,
+            work_item_width=REPLAY_WORK_ITEM_WIDTH,
+            position_field=REPLAY_WORK_POSITION_IN_DECODE_BATCH,
+            cache_slot_field=REPLAY_WORK_CACHE_SLOT,
+            history_len_field=REPLAY_WORK_PNAT,
+            cache_buf_idx_field=REPLAY_WORK_CACHE_BUF_IDX,
+            BLOCK_SIZE=block_size,
+            num_warps=4,
+        )
+
+
 @triton.jit
 def _prepare_gdn_replay_work_items_kernel(
     state_indices,
@@ -408,15 +462,14 @@ def _prepare_replay_work_items(self, kv_cache_manager, batch_size: int,
         self.replay_num_decodes = num_decodes
         if num_decodes == 0:
             return
+        use_gdn_mtp_replay = getattr(kv_cache_manager, "use_gdn_mtp_replay",
+                                     False)
         use_gdn_all_layer_commit = getattr(
             kv_cache_manager, "use_gdn_cached_replay_all_layer_commit", False)
         if use_gdn_all_layer_commit:
             from tensorrt_llm._torch.modules.fla.cached_replay import \
                 CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE
 
-            # The fused small-batch GDN kernel commits its checkpoint in-layer
-            # and indexes cache metadata directly. Work items are only consumed
-            # by the partitioned replay + all-layer commit path.
             if num_decodes < CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE:
                 return
 
@@ -436,6 +489,18 @@ def _prepare_replay_work_items(self, kv_cache_manager, batch_size: int,
         replay_step_width = replay_metadata.replay_step_width
         replay_history_size = replay_metadata.replay_history_size
 
+        if use_gdn_mtp_replay:
+            _build_gdn_replay_work_items(
+                self.state_indices[num_contexts:batch_size],
+                replay_metadata.history_len,
+                self.replay_work_items,
+                self.replay_n_writes,
+            )
+            return
+
+        if cache_buf_idx is None:
+            raise RuntimeError("Mamba2 replay requires cache buffer indices")
+
         if (use_gdn_all_layer_commit
                 and num_decodes <= _FUSED_GDN_REPLAY_WORK_ITEMS_MAX_BATCH_SIZE):
             build_work_items = _build_replay_work_items_triton

tensorrt_llm/_torch/pyexecutor/_util.py

@@ -45,6 +45,7 @@
 from ..attention.backends import get_sparse_attn_kv_cache_manager
 from ..attention.backends.sparse.params import \
     get_indexer_mqa_logits_workspace_bytes
+from ..cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE
 from ..disaggregation.kv_cache_transceiver import (AttentionTypeCpp,
                                                    create_kv_cache_transceiver)
 from ..hostfunc import set_low_latency_dispatch
@@ -3747,19 +3748,15 @@ def _create_kv_cache_manager(
             ))
         num_mamba_layers = (0 if is_draft and mamba_params.num_draft_layers > 0
                             else mamba_params.num_mamba_layers)
-        # Replay state update for GDN MTP: mirrors the nemotron_hybrid gating
-        # above, minus the Mamba2-specific stochastic-rounding/Philox gate.
-        # The GDN replay kernel does a plain cast on checkpoint commit, so
-        # quantized SSM cache dtypes stay on the legacy path.
         sm = get_sm_version()
-        use_replay = spec_config is not None and sm >= 80
+        use_replay = spec_config is not None
         if spec_config is None:
             logger.info(
                 "GDN replay kernel requires speculative decoding; using "
                 "non-replay path")
         elif spec_config.tokens_per_gen_step > 8:
-            logger.info("GDN cached replay supports at most 8 tokens per "
-                        "generation step; using non-replay path")
+            logger.info("GDN replay supports at most 8 tokens per generation "
+                        "step; using non-replay path")
             use_replay = False
 
         # Tree attention: replay assumes a linear token sequence.
@@ -3769,36 +3766,53 @@ def _create_kv_cache_manager(
                         "using legacy MTP path")
             use_replay = False
 
-        if mamba_params.mamba_ssm_cache_dtype not in (torch.float32,
-                                                      torch.bfloat16,
-                                                      torch.float16):
-            logger.info(
-                "GDN replay kernel does not support quantized SSM cache "
-                f"dtype {mamba_params.mamba_ssm_cache_dtype}; using legacy "
-                "MTP path")
-            use_replay = False
+        # Experimental modeling targets retain their cached-U replay ABI.
+        use_legacy_gdn_replay = (getattr(model_config, "modeling_v2_target",
+                                         None) is not None)
+        if not use_legacy_gdn_replay:
+            # GDN replay is specialized for BF16 Qwen GDN on datacenter
+            # Blackwell and uses Cache Manager V2 for its split all-layer commit.
+            if not is_sm_100f(sm):
+                logger.info("GDN MTP replay currently requires datacenter "
+                            "Blackwell; using non-replay path")
+                use_replay = False
+            if not IS_CUTLASS_DSL_AVAILABLE:
+                logger.info("GDN MTP replay requires CUTLASS DSL; using "
+                            "non-replay path")
+                use_replay = False
+            if (mamba_params.mamba_ssm_cache_dtype != torch.bfloat16
+                    or mamba_params.dtype != torch.bfloat16):
+                logger.info("GDN MTP replay requires BF16 activations and "
+                            "recurrent states; using non-replay path")
+                use_replay = False
+            if (mamba_params.state_size != 128 or mamba_params.head_dim != 128
+                    or mamba_params.conv_kernel != 4
+                    or mamba_params.n_groups <= 0
+                    or mamba_params.num_heads % mamba_params.n_groups != 0):
+                logger.info(
+                    "GDN MTP replay requires Conv4, K=V=128, and an "
+                    "integral value/key head ratio; using non-replay path")
+                use_replay = False
 
         # Replay is enabled by default for eligible GDN MTP workloads.
         if not is_gdn_replay_enabled():
             use_replay = False
 
-        # GDN replay supports the contiguous C++ V1 state pool and the indirect
-        # per-layer state views exposed by V2. Mixed/Python does not expose an
-        # all-layer commit, so keep that manager but use non-replay MTP.
-        replay_manager_types = (CppMambaHybridCacheManager,
-                                MambaHybridCacheManagerV2)
         if use_replay and not issubclass(kv_cache_manager_cls,
-                                         replay_manager_types):
-            logger.info("GDN replay requires C++ V1 or V2 Mamba cache manager; "
-                        f"{kv_cache_manager_cls.__name__} was selected, so the "
-                        "non-replay MTP path will be used")
+                                         MambaHybridCacheManagerV2):
+            logger.info("GDN replay requires MambaHybridCacheManagerV2; "
+                        f"{kv_cache_manager_cls.__name__} was selected, so "
+                        "the non-replay MTP path will be used")
             use_replay = False
         logger.info("GDN replay state update: " +
                     ("ENABLED" if use_replay else "DISABLED"))
 
         mamba_manager_extra_kwargs = dict(manager_extra_kwargs)
         mamba_manager_extra_kwargs.update(
             _mamba_conv_layout_kwargs(kv_cache_manager_cls, "qwen3_next"))
+        if issubclass(kv_cache_manager_cls, MambaHybridCacheManagerV2):
+            mamba_manager_extra_kwargs["use_gdn_mtp_replay"] = (
+                not use_legacy_gdn_replay)
         if getattr(sparse_attention_config, "algorithm", None) == "qsa":
             # Resolve the side-cache shape from the same checkpoint geometry
             # used to construct the QSA index projection.