[None][feat] Checkpointing variant of replay for MTP for mamba models

tensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.py

@@ -38,6 +38,7 @@
 from .._compat import KvCacheConfig, nvtx_range, prefer_pinned, str_dtype_to_torch
 from ..utils.logger import ad_logger
 from ..utils.node_utils import extract_op_args, get_op_schema
+from .mamba.replay_metadata import REPLAY_WORK_ITEM_WIDTH
 
 Constant = Union[int, float, str, None]
 
@@ -2204,10 +2205,9 @@ def from_base(
 
 
 class ReplayOldXHandler(SpeculativeOnly, StateResourceHandler):
-    """Per-layer old_x cache for the replay SSM kernel (single-buffered, bf16).
+    """Per-layer old_x cache for the replay SSM kernel (double-buffered, bf16).
 
-    Shape: (max_batch, T, num_heads, head_dim) — T is determined by the manager's
-    spec_config (max_draft_len + 1), not by this handler.  Acts as a type marker.
+    Shape: (max_batch, 2, replay_history_size, num_heads, head_dim).
     Routes to MambaHybridCacheManager via get_replay_old_x(layer_idx).
     """
 
@@ -2235,7 +2235,7 @@ def __eq__(self, other) -> bool:
 class ReplayOldBHandler(SpeculativeOnly, StateResourceHandler):
     """Per-layer old_B cache for the replay SSM kernel (double-buffered, bf16).
 
-    Shape: (max_batch, 2, T, n_groups, d_state) — T from manager.
+    Shape: (max_batch, 2, replay_history_size, n_groups, d_state).
     Routes to MambaHybridCacheManager via get_replay_old_B(layer_idx).
     """
 
@@ -2263,7 +2263,7 @@ def __eq__(self, other) -> bool:
 class ReplayOldDtHandler(SpeculativeOnly, StateResourceHandler):
     """Per-layer old_dt cache for the replay SSM kernel (double-buffered, fp32).
 
-    Shape: (max_batch, 2, num_heads, T) — T from manager.
+    Shape: (max_batch, 2, num_heads, replay_history_size).
     Routes to MambaHybridCacheManager via get_replay_old_dt(layer_idx).
     """
 
@@ -2285,7 +2285,7 @@ def __eq__(self, other) -> bool:
 class ReplayOldDAcumsumHandler(SpeculativeOnly, StateResourceHandler):
     """Per-layer old_dA_cumsum cache for the replay SSM kernel (double-buffered, fp32).
 
-    Shape: (max_batch, 2, num_heads, T) — T from manager.
+    Shape: (max_batch, 2, num_heads, replay_history_size).
     Routes to MambaHybridCacheManager via get_replay_old_dA_cumsum(layer_idx).
     """
 
@@ -2346,6 +2346,35 @@ def __eq__(self, other) -> bool:
         return isinstance(other, ReplayPrevNumAcceptedHandler)
 
 
+class ReplayWorkItemsHandler(ResourceHandler):
+    """Shared per-forward replay work items for the checkpoint replay SSM kernel.
+
+    Shape: (max_batch, REPLAY_WORK_ITEM_WIDTH) int32.  Each row is
+    (position_in_decode_batch, cache_slot, prev_num_accepted_tokens, cache_buf_idx).
+    """
+
+    def allocate(self, sequence_info) -> torch.Tensor:
+        return torch.empty(
+            sequence_info.max_num_state_slots,
+            REPLAY_WORK_ITEM_WIDTH,
+            device=sequence_info.device,
+            dtype=torch.int32,
+        )
+
+    def __eq__(self, other) -> bool:
+        return isinstance(other, ReplayWorkItemsHandler)
+
+
+class ReplayNWritesHandler(ResourceHandler):
+    """Shared single-element device tensor holding the replay write-count."""
+
+    def allocate(self, sequence_info) -> torch.Tensor:
+        return torch.empty(1, device=sequence_info.device, dtype=torch.int32)
+
+    def __eq__(self, other) -> bool:
+        return isinstance(other, ReplayNWritesHandler)
+
+
 class IntermediateConvStateHandler(SpeculativeOnly, StateResourceHandler):
     """Intermediate conv state cache descriptor for speculative decoding.
 

tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/flashinfer_backend_mamba.py

@@ -13,7 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import List, Optional
+from typing import List, Optional, cast
 
 import torch
 from flashinfer.mamba import selective_state_update as _flashinfer_ssm_update
@@ -31,11 +31,13 @@
     IntermediateSSMStateHandler,
     MHACallable,
     ReplayCacheBufIdxHandler,
+    ReplayNWritesHandler,
     ReplayOldBHandler,
     ReplayOldDAcumsumHandler,
     ReplayOldDtHandler,
     ReplayOldXHandler,
     ReplayPrevNumAcceptedHandler,
+    ReplayWorkItemsHandler,
     ResourceHandlerDict,
 )
 from .mamba_backend_common import (
@@ -64,7 +66,7 @@ def _fi_align(t: torch.Tensor) -> torch.Tensor:
         "ssm_state_cache",
         "intermediate_ssm_state_cache",
         # replay buffers: written in-place by the precompute and main kernels
-        # (double-buffered B/dt/dA_cumsum, single-buffered x); None in non-replay mode
+        # (double-buffered x/B/dt/dA_cumsum); None in non-replay mode
         "replay_old_x",
         "replay_old_b",
         "replay_old_dt",
@@ -95,14 +97,22 @@ def _flashinfer_cached_ssm(
     intermediate_ssm_state_cache: Optional[
         torch.Tensor
     ],  # [spec_state_size, max_draft_len+1, num_heads, head_dim, d_state]; None in replay mode
-    replay_old_x: Optional[torch.Tensor],  # [max_batch, T, nheads, head_dim]; None in non-replay
-    replay_old_b: Optional[torch.Tensor],  # [max_batch, 2, T, ngroups, dstate]; None in non-replay
-    replay_old_dt: Optional[torch.Tensor],  # [max_batch, 2, nheads, T] fp32; None in non-replay
+    replay_old_x: Optional[
+        torch.Tensor
+    ],  # [max_batch, 2, history, nheads, head_dim]; None in non-replay
+    replay_old_b: Optional[
+        torch.Tensor
+    ],  # [max_batch, 2, history, ngroups, dstate]; None in non-replay
+    replay_old_dt: Optional[
+        torch.Tensor
+    ],  # [max_batch, 2, nheads, history] fp32; None in non-replay
     replay_old_da_cumsum: Optional[
         torch.Tensor
-    ],  # [max_batch, 2, nheads, T] fp32; None in non-replay
+    ],  # [max_batch, 2, nheads, history] fp32; None in non-replay
     replay_cache_buf_idx: Optional[torch.Tensor],  # [max_batch] int32; None in non-replay
     replay_prev_num_accepted: Optional[torch.Tensor],  # [max_batch] int32; None in non-replay
+    replay_work_items: Optional[torch.Tensor],  # [max_batch, 4] int32; None in non-replay
+    replay_n_writes: Optional[torch.Tensor],  # [1] int32; None in non-replay
     # CONSTANTS
     time_step_limit: List[float],
     chunk_size: int,
@@ -198,6 +208,34 @@ def _flashinfer_cached_ssm(
 
         use_replay = batch_info.is_use_replay()
         if use_replay:
+            missing_replay_tensors = [
+                name
+                for name, tensor in (
+                    ("replay_old_x", replay_old_x),
+                    ("replay_old_b", replay_old_b),
+                    ("replay_old_dt", replay_old_dt),
+                    ("replay_old_da_cumsum", replay_old_da_cumsum),
+                    ("replay_cache_buf_idx", replay_cache_buf_idx),
+                    ("replay_prev_num_accepted", replay_prev_num_accepted),
+                    ("replay_work_items", replay_work_items),
+                    ("replay_n_writes", replay_n_writes),
+                )
+                if tensor is None
+            ]
+            if missing_replay_tensors:
+                raise RuntimeError(
+                    "flashinfer_cached_ssm replay path missing required tensors: "
+                    f"{', '.join(missing_replay_tensors)}"
+                )
+            replay_old_x = cast(torch.Tensor, replay_old_x)
+            replay_old_b = cast(torch.Tensor, replay_old_b)
+            replay_old_dt = cast(torch.Tensor, replay_old_dt)
+            replay_old_da_cumsum = cast(torch.Tensor, replay_old_da_cumsum)
+            replay_cache_buf_idx = cast(torch.Tensor, replay_cache_buf_idx)
+            replay_prev_num_accepted = cast(torch.Tensor, replay_prev_num_accepted)
+            replay_work_items = cast(torch.Tensor, replay_work_items)
+            replay_n_writes = cast(torch.Tensor, replay_n_writes)
+
             # Replay path: fast-forward SSM state via tl.dot on cached values.
             # State is updated in-place; no disable_state_update needed.
             # x_extend/B_extend/C_extend are non-contiguous views from the CUDA graph's
@@ -218,13 +256,20 @@ def _flashinfer_cached_ssm(
                 B_extend,
                 C_extend,
                 out=preallocated_ssm_out_e,
+                n_writes=replay_n_writes,
+                replay_work_items=replay_work_items[:num_extend],
                 D=D_full,
                 dt_bias=dt_bias_hp,
                 dt_softplus=True,
                 state_batch_indices=slot_idx_extend,
                 launch_with_pdl=True,  # PDL chain: triton_causal_conv extend → precompute → main
             )
         else:
+            if intermediate_ssm_state_cache is None:
+                raise RuntimeError(
+                    "flashinfer_cached_ssm non-replay extend branch requires "
+                    "intermediate_ssm_state_cache"
+                )
             if intermediate_ssm_state_cache.size(1) < tokens_per_extend:
                 raise RuntimeError(
                     "flashinfer_cached_ssm: intermediate_ssm_state_cache is too small "
@@ -342,14 +387,22 @@ def _flashinfer_cached_ssm_fake(
     intermediate_ssm_state_cache: Optional[
         torch.Tensor
     ],  # [spec_state_size, max_draft_len+1, num_heads, head_dim, d_state]; None in replay mode
-    replay_old_x: Optional[torch.Tensor],  # [max_batch, T, nheads, head_dim]; None in non-replay
-    replay_old_b: Optional[torch.Tensor],  # [max_batch, 2, T, ngroups, dstate]; None in non-replay
-    replay_old_dt: Optional[torch.Tensor],  # [max_batch, 2, nheads, T] fp32; None in non-replay
+    replay_old_x: Optional[
+        torch.Tensor
+    ],  # [max_batch, 2, history, nheads, head_dim]; None in non-replay
+    replay_old_b: Optional[
+        torch.Tensor
+    ],  # [max_batch, 2, history, ngroups, dstate]; None in non-replay
+    replay_old_dt: Optional[
+        torch.Tensor
+    ],  # [max_batch, 2, nheads, history] fp32; None in non-replay
     replay_old_da_cumsum: Optional[
         torch.Tensor
-    ],  # [max_batch, 2, nheads, T] fp32; None in non-replay
+    ],  # [max_batch, 2, nheads, history] fp32; None in non-replay
     replay_cache_buf_idx: Optional[torch.Tensor],  # [max_batch] int32; None in non-replay
     replay_prev_num_accepted: Optional[torch.Tensor],  # [max_batch] int32; None in non-replay
+    replay_work_items: Optional[torch.Tensor],  # [max_batch, 4] int32; None in non-replay
+    replay_n_writes: Optional[torch.Tensor],  # [1] int32; None in non-replay
     # CONSTANTS
     time_step_limit: List[float],
     chunk_size: int,
@@ -400,7 +453,7 @@ def get_cache_initializers(
 
         ssm_h = ret["ssm_state_cache"]
 
-        # All 7 optional caches are always registered positionally (None = unused in this mode).
+        # Optional replay/spec caches are registered positionally (None = unused in this mode).
         # intermediate_ssm_state_cache: real in non-replay, None in replay.
         # replay_old_*: real in replay mode (SM80+), None otherwise.
         if use_replay:
@@ -418,6 +471,8 @@ def get_cache_initializers(
             ret["replay_old_da_cumsum"] = ReplayOldDAcumsumHandler(num_heads=ssm_h.num_heads)
             ret["replay_cache_buf_idx"] = ReplayCacheBufIdxHandler()
             ret["replay_prev_num_accepted"] = ReplayPrevNumAcceptedHandler()
+            ret["replay_work_items"] = ReplayWorkItemsHandler()
+            ret["replay_n_writes"] = ReplayNWritesHandler()
         else:
             ret["intermediate_ssm_state_cache"] = IntermediateSSMStateHandler.from_base(ssm_h)
             ret["replay_old_x"] = None
@@ -426,4 +481,6 @@ def get_cache_initializers(
             ret["replay_old_da_cumsum"] = None
             ret["replay_cache_buf_idx"] = None
             ret["replay_prev_num_accepted"] = None
+            ret["replay_work_items"] = None
+            ret["replay_n_writes"] = None
         return ret

tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/replay_metadata.py

@@ -0,0 +1,22 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Replay metadata layout shared by AutoDeploy Mamba descriptors."""
+
+REPLAY_WORK_POSITION_IN_DECODE_BATCH = 0
+REPLAY_WORK_CACHE_SLOT = 1
+REPLAY_WORK_PNAT = 2
+REPLAY_WORK_CACHE_BUF_IDX = 3
+REPLAY_WORK_ITEM_WIDTH = 4

tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py

@@ -993,6 +993,7 @@ def _prepare_inputs(
             _ungathered_new_lens=new_tokens_lens,
             **extra_args,
         )
+        self.cache_seq_interface.prepare_replay_metadata()
 
         self.iter_states["num_ctx_requests"] = num_prefill
         self.iter_states["num_ctx_tokens"] = num_prefill_tokens