From e52b5e2ee29e1d04d53323077f5c2769f5e92f9d Mon Sep 17 00:00:00 2001
From: Xiao Wang <24860335+xwang233@users.noreply.github.com>
Date: Wed, 21 Oct 2026 21:55:47 -0700
Subject: [PATCH 1/5] [TRTLLM-17096][feat] runtime hooks: model opt-in engine
 fast paths

Add RuntimeHooks, a per-model declaration of engine-side fast paths the
model provides itself: a step-input stager, a speculative sampler store
writer, a decode prologue, executor events, and running the overlap loop
on the execution stream. Every hook defaults to off, and a model without
a declaration runs the native path unchanged.

Also add two extension points a model can use: FMHA libraries a backend
puts ahead of the enabled ones, and the DSA indexer's reused top-k copy
as an overridable method.

Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com>
---
 .../_torch/attention/backends/fmha/manager.py |   5 +-
 .../attention/backends/sparse/dsa/indexer.py  | 136 ++-
 .../attention/backends/sparse/dsa/metadata.py |  98 +-
 .../_torch/attention/backends/trtllm.py       |  61 +-
 tensorrt_llm/_torch/pyexecutor/_util.py       |   8 +-
 .../_torch/pyexecutor/cuda_graph_runner.py    |  14 +-
 .../pyexecutor/engine/runners/common.py       |   8 +
 .../engine/runners/decoder/graphs.py          |   7 +
 .../engine/runners/decoder/inputs.py          |  98 +-
 .../engine/runners/decoder/runner.py          | 154 ++-
 .../kv_cache/kv_cache_manager_v2.py           |  13 +
 tensorrt_llm/_torch/pyexecutor/py_executor.py |  65 +-
 .../_torch/pyexecutor/runtime_hooks.py        | 221 +++++
 tensorrt_llm/_torch/speculative/eagle3.py     |   9 +-
 tensorrt_llm/_torch/speculative/interface.py  |  10 +-
 .../_torch/speculative/spec_sampler_base.py   |  82 +-
 .../attention/sparse/dsa/test_dsa_indexer.py  | 100 ++
 .../_torch/attention/test_fmha_manager.py     |  57 ++
 .../_torch/executor/test_gpu_keepalive.py     |   1 +
 .../_torch/executor/test_py_executor.py       |   3 +
 .../_torch/executor/test_runtime_hooks.py     | 924 ++++++++++++++++++
 21 files changed, 1942 insertions(+), 132 deletions(-)
 create mode 100644 tensorrt_llm/_torch/pyexecutor/runtime_hooks.py
 create mode 100644 tests/unittest/_torch/executor/test_runtime_hooks.py

diff --git a/tensorrt_llm/_torch/attention/backends/fmha/manager.py b/tensorrt_llm/_torch/attention/backends/fmha/manager.py
index c9afff66e1d8..a4770ad3dbba 100644
--- a/tensorrt_llm/_torch/attention/backends/fmha/manager.py
+++ b/tensorrt_llm/_torch/attention/backends/fmha/manager.py
@@ -308,7 +308,10 @@ def __init__(self, attn: TrtllmAttention) -> None:
         # normal forwarding never pays for an environment lookup.
         self._cache_sanity_check_enabled = _is_fmha_cache_sanity_check_enabled()
         self.fmha_libs: list[Fmha] = []
-        for fmha_cls in get_enabled_fmha_lib_classes():
+        for fmha_cls in (
+            *getattr(attn, "prepended_fmha_lib_classes", ()),
+            *get_enabled_fmha_lib_classes(),
+        ):
             if fmha_cls.is_available(attn):
                 self.fmha_libs.append(fmha_cls(attn))
 
diff --git a/tensorrt_llm/_torch/attention/backends/sparse/dsa/indexer.py b/tensorrt_llm/_torch/attention/backends/sparse/dsa/indexer.py
index a4d3c25b7bce..9e1ff5b6f902 100644
--- a/tensorrt_llm/_torch/attention/backends/sparse/dsa/indexer.py
+++ b/tensorrt_llm/_torch/attention/backends/sparse/dsa/indexer.py
@@ -1240,6 +1240,18 @@ def recompute_slot_mappings(
         metadata.host_slot_mapping_fp8[:total_new_kv_tokens] = fp8_flat_indices
         metadata.host_slot_mapping_scale[:total_new_kv_tokens] = scale_flat_indices
 
+        stager = getattr(metadata, "step_input_stager", None)
+        if stager is not None:
+            # Staged write; see the window banner in DecoderRunner._forward_decoder.
+            stager.copy(
+                metadata.slot_mapping_fp8[:total_new_kv_tokens],
+                metadata.host_slot_mapping_fp8[:total_new_kv_tokens],
+            )
+            stager.copy(
+                metadata.slot_mapping_scale[:total_new_kv_tokens],
+                metadata.host_slot_mapping_scale[:total_new_kv_tokens],
+            )
+            return
         metadata.slot_mapping_fp8[:total_new_kv_tokens].copy_(
             metadata.host_slot_mapping_fp8[:total_new_kv_tokens], non_blocking=True
         )
@@ -1514,7 +1526,13 @@ def prepare(metadata: DSAtrtllmAttentionMetadata) -> None:
         if num_generations > 0:
             # Prepare schedule metadata for fp8_paged_mqa_logits
             # This is a preprocessing step that computes scheduling information for the kernel
-            Indexer.prepare_scheduler_metadata(metadata)
+            # Staged write; see the window banner in DecoderRunner._forward_decoder.
+            stager = getattr(metadata, "step_input_stager", None)
+            if stager is None:
+                Indexer.prepare_scheduler_metadata(metadata)
+            elif not stager.scheduler_metadata(metadata):
+                # Reads staged KV lengths: run right after the staged launch.
+                stager.defer(lambda: Indexer.prepare_scheduler_metadata(metadata))
 
     def _update_k_cache(
         self, k_fp8: torch.Tensor, k_scale: torch.Tensor, metadata: DSAtrtllmAttentionMetadata
@@ -1987,9 +2005,24 @@ def sparse_attn_indexer(
         )
 
         if has_decode and not metadata.skip_indexer_for_gen_reqs and reuse_topk:
-            topk_indices_buffer[token_offset : token_offset + num_gen_tokens, :] = (
-                metadata.shared_topk_indices[:num_generations, :]
+            self._copy_reused_topk(
+                metadata,
+                topk_indices_buffer,
+                token_offset,
+                num_gen_tokens,
+                num_generations,
+                has_prefill,
+                gvr_prior_indices,
             )
+        elif (
+            has_decode
+            and not metadata.skip_indexer_for_gen_reqs
+            and self._fill_decode_topk(
+                metadata, q_fp8, weights, topk_indices_buffer, token_offset, num_gen_tokens
+            )
+        ):
+            # A subclass filled the generation rows.
+            pass
         elif has_decode and not metadata.skip_indexer_for_gen_reqs:
             # Get decode lengths per request (from seq_lens) for validation
             gen_seq_lens = metadata.seq_lens[num_contexts : num_contexts + num_generations]
@@ -2268,30 +2301,83 @@ def sparse_attn_indexer(
                 prev_topk_dst.copy_(last_mtp_topk)
 
         if self.mtp_index_share and metadata.in_mtp_draft_loop and not reuse_topk:
-            rows = None
-            if has_decode:
-                next_n = num_gen_tokens // num_generations
-                gen_topk = topk_indices_buffer[token_offset : token_offset + num_gen_tokens]
-                rows = self._mtp_last_accepted_rows(
-                    gen_topk, metadata, num_contexts, num_generations, next_n
-                )
-            if has_prefill:
-                ctx_last = (
-                    torch.cumsum(metadata.seq_lens_cuda[:num_contexts].to(torch.long), dim=0) - 1
-                )
-                ctx_rows = topk_indices_buffer[ctx_last]
-                rows = ctx_rows if rows is None else torch.cat([ctx_rows, rows])
-            if rows is not None:
-                shared_topk_indices = metadata.shared_topk_indices
-                if shared_topk_indices is None:
-                    metadata.shared_topk_indices = rows.contiguous()
-                else:
-                    row_start = num_contexts if is_generation else 0
-                    shared_topk_indices[
-                        row_start : row_start + rows.shape[0], : rows.shape[1]
-                    ].copy_(rows)
+            self._store_mtp_shared_rows(
+                topk_indices_buffer,
+                metadata,
+                num_contexts,
+                num_generations,
+                num_gen_tokens,
+                token_offset,
+                has_prefill,
+                has_decode,
+                is_generation,
+            )
         return topk_indices_buffer
 
+    def _copy_reused_topk(
+        self,
+        metadata: DSAtrtllmAttentionMetadata,
+        topk_indices_buffer: torch.Tensor,
+        token_offset: int,
+        num_gen_tokens: int,
+        num_generations: int,
+        has_prefill: bool,
+        gvr_prior_indices: Optional[torch.Tensor],
+    ) -> None:
+        """Copy the shared TopK rows an MTP draft step reuses into its generation rows."""
+        topk_indices_buffer[token_offset : token_offset + num_gen_tokens, :] = (
+            metadata.shared_topk_indices[:num_generations, :]
+        )
+
+    def _fill_decode_topk(
+        self,
+        metadata: DSAtrtllmAttentionMetadata,
+        q_fp8: torch.Tensor,
+        weights: torch.Tensor,
+        topk_indices_buffer: torch.Tensor,
+        token_offset: int,
+        num_gen_tokens: int,
+    ) -> bool:
+        """Fill the generation TopK rows ahead of the decode branch; True if it filled them.
+
+        This class leaves them to the decode branch of ``sparse_attn_indexer``.
+        """
+        return False
+
+    def _store_mtp_shared_rows(
+        self,
+        topk_indices_buffer: torch.Tensor,
+        metadata: DSAtrtllmAttentionMetadata,
+        num_contexts: int,
+        num_generations: int,
+        num_gen_tokens: int,
+        token_offset: int,
+        has_prefill: bool,
+        has_decode: bool,
+        is_generation: Optional[bool],
+    ) -> None:
+        """Store each request's last accepted TopK row for the next MTP draft step to reuse."""
+        rows = None
+        if has_decode:
+            next_n = num_gen_tokens // num_generations
+            gen_topk = topk_indices_buffer[token_offset : token_offset + num_gen_tokens]
+            rows = self._mtp_last_accepted_rows(
+                gen_topk, metadata, num_contexts, num_generations, next_n
+            )
+        if has_prefill:
+            ctx_last = torch.cumsum(metadata.seq_lens_cuda[:num_contexts].to(torch.long), dim=0) - 1
+            ctx_rows = topk_indices_buffer[ctx_last]
+            rows = ctx_rows if rows is None else torch.cat([ctx_rows, rows])
+        if rows is not None:
+            shared_topk_indices = metadata.shared_topk_indices
+            if shared_topk_indices is None:
+                metadata.shared_topk_indices = rows.contiguous()
+            else:
+                row_start = num_contexts if is_generation else 0
+                shared_topk_indices[row_start : row_start + rows.shape[0], : rows.shape[1]].copy_(
+                    rows
+                )
+
     def _mtp_last_accepted_rows(
         self,
         gen_topk: torch.Tensor,
diff --git a/tensorrt_llm/_torch/attention/backends/sparse/dsa/metadata.py b/tensorrt_llm/_torch/attention/backends/sparse/dsa/metadata.py
index e0773630dcb3..5c9c4686375c 100644
--- a/tensorrt_llm/_torch/attention/backends/sparse/dsa/metadata.py
+++ b/tensorrt_llm/_torch/attention/backends/sparse/dsa/metadata.py
@@ -1486,6 +1486,15 @@ def _update_indexer_k_cache_block_offsets(self) -> torch.Tensor:
         page_indices = pool_indices * cache_manager.indexer_k_cache_page_scale
         num_blocks = page_indices.shape[1]
         self.host_indexer_k_cache_block_offsets[: self.num_seqs, :num_blocks].copy_(page_indices)
+        if self.step_input_stager is not None:
+            # Staged write; see the window banner in DecoderRunner._forward_decoder.
+            self.step_input_stager.copy(
+                self.indexer_k_cache_block_offsets[: self.num_seqs],
+                self.host_indexer_k_cache_block_offsets[: self.num_seqs],
+                clamp_min=0,
+            )
+            self.step_input_stager.clamp_min(self.indexer_k_cache_block_offsets[self.num_seqs :], 0)
+            return pool_indices
         self.indexer_k_cache_block_offsets[: self.num_seqs].copy_(
             self.host_indexer_k_cache_block_offsets[: self.num_seqs], non_blocking=True
         )
@@ -1702,9 +1711,16 @@ def prepare_dense_topk_indices(self, kv_lens, device=False):  # device=False mea
                     kv_lens[self.num_contexts : self.num_seqs],
                     self.num_tokens - self.num_ctx_tokens,
                 )
-                self.topk_indices_buffer[gen_range, :].copy_(
-                    self.host_topk_indices_buffer[gen_range, :], non_blocking=True
-                )
+                if self.step_input_stager is not None:
+                    # Staged write; see the window banner in DecoderRunner._forward_decoder.
+                    self.step_input_stager.copy(
+                        self.topk_indices_buffer[gen_range, :],
+                        self.host_topk_indices_buffer[gen_range, :],
+                    )
+                else:
+                    self.topk_indices_buffer[gen_range, :].copy_(
+                        self.host_topk_indices_buffer[gen_range, :], non_blocking=True
+                    )
 
     def prepare_for_spec_decode(self, kv_lens: torch.Tensor):
         # The DeepGEMM paged-MQA kernel (fp8_paged_mqa_logits) runs a native
@@ -1804,9 +1820,16 @@ def prepare_for_spec_decode(self, kv_lens: torch.Tensor):
                 num_tokens = self.num_generations * expand_factor
                 gen_kv_lens_expanded = gen_kv_lens.repeat_interleave(expand_factor)
                 self.kv_lens_expanded_host[:num_tokens].copy_(gen_kv_lens_expanded)
-                self.kv_lens_expanded_cuda[:num_tokens].copy_(
-                    self.kv_lens_expanded_host[:num_tokens], non_blocking=True
-                )
+                if self.step_input_stager is not None:
+                    # Staged write; see the window banner in DecoderRunner._forward_decoder.
+                    self.step_input_stager.copy(
+                        self.kv_lens_expanded_cuda[:num_tokens],
+                        self.kv_lens_expanded_host[:num_tokens],
+                    )
+                else:
+                    self.kv_lens_expanded_cuda[:num_tokens].copy_(
+                        self.kv_lens_expanded_host[:num_tokens], non_blocking=True
+                    )
                 self._refresh_expanded_block_table(expand_factor)
         else:
             self.dsl_expand_factor = 1
@@ -2036,6 +2059,15 @@ def _refresh_expanded_block_table(self, repeat_factor: Optional[int] = None):
         self.host_block_table_expanded[:num_tokens, :max_len].copy_(
             expanded_blocks, non_blocking=True
         )
+        if self.step_input_stager is not None:
+            # Staged write; see the window banner in DecoderRunner._forward_decoder.
+            self.step_input_stager.copy(
+                self.block_table_expanded[:num_tokens],
+                self.host_block_table_expanded[:num_tokens],
+                clamp_min=0,
+            )
+            self.step_input_stager.clamp_min(self.block_table_expanded[num_tokens:], 0)
+            return
         self.block_table_expanded[:num_tokens].copy_(
             self.host_block_table_expanded[:num_tokens], non_blocking=True
         )
@@ -2076,9 +2108,15 @@ def prepare_for_indexer_k_cache(self):
             pad_cols.unsqueeze(0) >= num_blocks_per_seq[: self.num_seqs].unsqueeze(1), -1
         )
         # Copy to GPU
-        self.block_table[: self.num_seqs, :max_blocks_used].copy_(
-            host_block_table, non_blocking=True
-        )
+        if self.step_input_stager is not None:
+            # Staged write; see the window banner in DecoderRunner._forward_decoder.
+            self.step_input_stager.copy(
+                self.block_table[: self.num_seqs, :max_blocks_used], host_block_table
+            )
+        else:
+            self.block_table[: self.num_seqs, :max_blocks_used].copy_(
+                host_block_table, non_blocking=True
+            )
 
     def prepare_for_skip_indexer(self, kv_lens: torch.Tensor):
         num_extra_kv_tokens = self.kv_cache_params.num_extra_kv_tokens
@@ -2144,9 +2182,17 @@ def prepare_for_mla_rope_append(self, cached_token_lens: torch.Tensor, kv_lens:
                 dtype=torch.int64,
                 out=self.host_gen_cached_token_indptr[1 : self.num_generations + 1],
             )
-            self.gen_cached_token_indptr[: self.num_generations + 1].copy_(
-                self.host_gen_cached_token_indptr[: self.num_generations + 1], non_blocking=True
-            )
+            if self.step_input_stager is not None:
+                # Staged write; see the window banner in DecoderRunner._forward_decoder.
+                self.step_input_stager.copy(
+                    self.gen_cached_token_indptr[: self.num_generations + 1],
+                    self.host_gen_cached_token_indptr[: self.num_generations + 1],
+                )
+            else:
+                self.gen_cached_token_indptr[: self.num_generations + 1].copy_(
+                    self.host_gen_cached_token_indptr[: self.num_generations + 1],
+                    non_blocking=True,
+                )
             # generation kv indptr
             torch.cumsum(
                 kv_lens[self.num_contexts : self.num_seqs],
@@ -2154,9 +2200,16 @@ def prepare_for_mla_rope_append(self, cached_token_lens: torch.Tensor, kv_lens:
                 dtype=torch.int64,
                 out=self.host_gen_kv_indptr[1 : self.num_generations + 1],
             )
-            self.gen_kv_indptr[: self.num_generations + 1].copy_(
-                self.host_gen_kv_indptr[: self.num_generations + 1], non_blocking=True
-            )
+            if self.step_input_stager is not None:
+                # Staged write; see the window banner in DecoderRunner._forward_decoder.
+                self.step_input_stager.copy(
+                    self.gen_kv_indptr[: self.num_generations + 1],
+                    self.host_gen_kv_indptr[: self.num_generations + 1],
+                )
+            else:
+                self.gen_kv_indptr[: self.num_generations + 1].copy_(
+                    self.host_gen_kv_indptr[: self.num_generations + 1], non_blocking=True
+                )
         else:
             self.max_gen_seq_len = 0
 
@@ -2168,7 +2221,14 @@ def prepare_for_indices_conversion(self):
             self.seq_lens,
             dim=0,
         )
-        self.req_idx_per_token[: self.num_tokens].copy_(
-            self.host_req_idx_per_token[: self.num_tokens],
-            non_blocking=True,
-        )
+        if self.step_input_stager is not None:
+            # Staged write; see the window banner in DecoderRunner._forward_decoder.
+            self.step_input_stager.copy(
+                self.req_idx_per_token[: self.num_tokens],
+                self.host_req_idx_per_token[: self.num_tokens],
+            )
+        else:
+            self.req_idx_per_token[: self.num_tokens].copy_(
+                self.host_req_idx_per_token[: self.num_tokens],
+                non_blocking=True,
+            )
diff --git a/tensorrt_llm/_torch/attention/backends/trtllm.py b/tensorrt_llm/_torch/attention/backends/trtllm.py
index 1d7f183e6e85..64b29917880b 100644
--- a/tensorrt_llm/_torch/attention/backends/trtllm.py
+++ b/tensorrt_llm/_torch/attention/backends/trtllm.py
@@ -28,6 +28,7 @@
     from tensorrt_llm.mapping import Mapping
 
     from ...model_config import ModelConfig
+    from ...pyexecutor.runtime_hooks import StepInputStager
     from ...speculative.interface import SpecMetadata
     from ...speculative.spec_tree_manager import SpecTreeManager
     from .sparse.dsa.metadata import TokenMajorGenView
@@ -123,6 +124,8 @@ def generate_spec_decoding_packed_mask(max_num_requests: int,
 
 @dataclass(kw_only=True)
 class TrtllmAttentionMetadata(AttentionMetadata):
+    # A step-input stager while it records this metadata's writes, else None.
+    step_input_stager: ClassVar[Optional["StepInputStager"]] = None
     workspace: Optional[torch.Tensor] = None
     cuda_graph_workspace: Optional[torch.Tensor] = None
     workspace_reclaimable: bool = field(default=True, init=False)
@@ -1006,8 +1009,14 @@ def prepare(self) -> None:
             device='cpu',
         )
         self.prompt_lens_cpu[:self.num_seqs].copy_(prompt_lens)
-        self.prompt_lens_cuda[:self.num_seqs].copy_(
-            self.prompt_lens_cpu[:self.num_seqs], non_blocking=True)
+        stager = self.step_input_stager
+        if stager is not None:
+            # Staged write; see the window banner in DecoderRunner._forward_decoder.
+            stager.copy(self.prompt_lens_cuda[:self.num_seqs],
+                        self.prompt_lens_cpu[:self.num_seqs])
+        else:
+            self.prompt_lens_cuda[:self.num_seqs].copy_(
+                self.prompt_lens_cpu[:self.num_seqs], non_blocking=True)
 
         # number of tokens in the kv cache for each sequence in the batch
         cached_token_lens = torch.tensor(
@@ -1060,9 +1069,14 @@ def prepare(self) -> None:
         # the sequence length including the cached tokens and the input tokens.
         self.kv_lens[:self.num_seqs].copy_(
             kv_lens + self.kv_cache_params.num_extra_kv_tokens)
-        self.kv_lens_cuda[:self.num_seqs].copy_(maybe_pin_memory(
-            kv_lens[:self.num_seqs]),
-                                                non_blocking=True)
+        if stager is not None:
+            # Staged write; see the window banner in DecoderRunner._forward_decoder.
+            stager.copy(self.kv_lens_cuda[:self.num_seqs],
+                        kv_lens[:self.num_seqs])
+        else:
+            self.kv_lens_cuda[:self.num_seqs].copy_(maybe_pin_memory(
+                kv_lens[:self.num_seqs]),
+                                                    non_blocking=True)
         # total kv lens for context requests and generation requests, without extra tokens
         self.host_total_kv_lens[0] = kv_lens[:self.num_contexts].sum().item()
         self.host_total_kv_lens[1] = kv_lens[self.num_contexts:self.
@@ -1429,23 +1443,42 @@ def prepare_context_mla_with_cached_kv(self,
                      dtype=torch.int64,
                      out=self.host_ctx_cached_token_indptr[1:self.num_contexts +
                                                            1])
-        self.ctx_cached_token_indptr[:self.num_contexts + 1].copy_(
-            self.host_ctx_cached_token_indptr[:self.num_contexts + 1],
-            non_blocking=True)
+        stager = self.step_input_stager
+        if stager is not None:
+            # Staged write; see the window banner in DecoderRunner._forward_decoder.
+            stager.copy(
+                self.ctx_cached_token_indptr[:self.num_contexts + 1],
+                self.host_ctx_cached_token_indptr[:self.num_contexts + 1])
+        else:
+            self.ctx_cached_token_indptr[:self.num_contexts + 1].copy_(
+                self.host_ctx_cached_token_indptr[:self.num_contexts + 1],
+                non_blocking=True)
         torch.cumsum(
             self.seq_lens[:self.num_contexts],
             dim=0,
             dtype=torch.int64,
             out=self.host_ctx_uncached_token_indptr[1:self.num_contexts + 1])
-        self.ctx_uncached_token_indptr[:self.num_contexts + 1].copy_(
-            self.host_ctx_uncached_token_indptr[:self.num_contexts + 1],
-            non_blocking=True)
+        if stager is not None:
+            # Staged write; see the window banner in DecoderRunner._forward_decoder.
+            stager.copy(
+                self.ctx_uncached_token_indptr[:self.num_contexts + 1],
+                self.host_ctx_uncached_token_indptr[:self.num_contexts + 1])
+        else:
+            self.ctx_uncached_token_indptr[:self.num_contexts + 1].copy_(
+                self.host_ctx_uncached_token_indptr[:self.num_contexts + 1],
+                non_blocking=True)
         torch.cumsum(kv_lens[:self.num_contexts],
                      dim=0,
                      dtype=torch.int64,
                      out=self.host_ctx_kv_indptr[1:self.num_contexts + 1])
-        self.ctx_kv_indptr[:self.num_contexts + 1].copy_(
-            self.host_ctx_kv_indptr[:self.num_contexts + 1], non_blocking=True)
+        if stager is not None:
+            # Staged write; see the window banner in DecoderRunner._forward_decoder.
+            stager.copy(self.ctx_kv_indptr[:self.num_contexts + 1],
+                        self.host_ctx_kv_indptr[:self.num_contexts + 1])
+        else:
+            self.ctx_kv_indptr[:self.num_contexts + 1].copy_(
+                self.host_ctx_kv_indptr[:self.num_contexts + 1],
+                non_blocking=True)
 
     def spec_decoding_param_prepare_for_blackwell(self) -> None:
         """
@@ -1654,6 +1687,8 @@ def is_sm_version_trtllm_gen_kernel(self, sm):
 class TrtllmAttention(AttentionBackend[TrtllmAttentionMetadata]):
 
     Metadata = TrtllmAttentionMetadata
+    # FMHA library classes this instance tries before the registry's libraries.
+    prepended_fmha_lib_classes: tuple = ()
 
     @staticmethod
     def is_sm_version_trtllm_gen_kernel(sm):
diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py
index 4d3b4af850e8..2090179caa6c 100644
--- a/tensorrt_llm/_torch/pyexecutor/_util.py
+++ b/tensorrt_llm/_torch/pyexecutor/_util.py
@@ -53,6 +53,7 @@
 from ..speculative import (draft_prompt_lookahead, get_num_extra_kv_tokens,
                            get_num_spec_layers, get_spec_decoder,
                            should_use_separate_draft_kv_cache)
+from ..speculative.spec_sampler_base import SpecSampler
 from ..utils import is_gdn_replay_enabled
 from . import config_utils
 from .config_utils import (MambaKVCacheParams, _is_sliding_attention_layer,
@@ -77,6 +78,7 @@
 from .resource_manager import (KVCacheCompressionManager, KVCacheManager,
                                PeftCacheManager, ResourceManager,
                                ResourceManagerType)
+from .runtime_hooks import runtime_hooks_of
 from .sampler import EarlyStopSampler, EarlyStopWithMMResult, TorchSampler
 from .sampler.penalties import PenaltyStore
 from .scheduler import (BindCapacityScheduler, BindMicroBatchScheduler,
@@ -4587,7 +4589,11 @@ def instantiate_sampler(
     )
     if engine.spec_config is not None and engine.spec_config.spec_dec_mode.has_spec_decoder(
     ):
-        return get_spec_decoder(sampler_args, engine.spec_config)
+        sampler = get_spec_decoder(sampler_args, engine.spec_config)
+        store_writer = runtime_hooks_of(engine.model).sampler_store_writer
+        if store_writer is not None and isinstance(sampler, SpecSampler):
+            sampler.store_writer = store_writer
+        return sampler
 
     if mm_encoder_only:
         # NOTE: handle model outputs specially for mm encoder executor/engine.
diff --git a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py
index e0d5abd61b35..118499f61732 100644
--- a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py
+++ b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py
@@ -160,6 +160,10 @@ class CUDAGraphRunner:
     multiple graphs, keyed by batch shape and execution-path specializations.
     """
     WARMUP_STEPS = 1
+    # Tokens of input_ids / position_ids that a step-input stager already
+    # wrote into the static inputs for the next replay (0: the replay copies
+    # them).
+    staged_static_input_tokens = 0
 
     def __init__(self, config: CUDAGraphRunnerConfig):
         self.config = config
@@ -804,7 +808,12 @@ def replay(
                 f"replay() got {seqlen} tokens for key {key}, but the graph "
                 f"was captured for {expected_num_tokens} tokens. A shorter "
                 "input_ids leaves the tail of the static input buffer stale.")
-        static_tensors["input_ids"][:seqlen].copy_(input_ids)
+        # A step-input stager already wrote input_ids / position_ids[:seqlen]
+        # into the static buffers.
+        staged = self.staged_static_input_tokens == seqlen
+        self.staged_static_input_tokens = 0
+        if not staged:
+            static_tensors["input_ids"][:seqlen].copy_(input_ids)
 
         position_ids = current_inputs["position_ids"]
         if self.config.use_mrope:
@@ -851,7 +860,8 @@ def replay(
                     f"for key {key}, but expected {expected_position_ids_shape}. "
                     "torch.Tensor.copy_() silently broadcasts mismatched shapes, "
                     "which would corrupt the static input buffer.")
-            static_tensors["position_ids"][:, :seqlen].copy_(position_ids)
+            if not staged:
+                static_tensors["position_ids"][:, :seqlen].copy_(position_ids)
 
         num_encoder_tokens = key.num_encoder_tokens
         if num_encoder_tokens:
diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/common.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/common.py
index 98dce6b41550..cb24f68deef7 100644
--- a/tensorrt_llm/_torch/pyexecutor/engine/runners/common.py
+++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/common.py
@@ -22,6 +22,7 @@
 from .interface import ScheduledInputs
 
 if TYPE_CHECKING:
+    from tensorrt_llm._torch.pyexecutor.runtime_hooks import StepInputStager
     from tensorrt_llm._torch.pyexecutor.sampler.sampler import SampleStateTensors
 
 
@@ -137,11 +138,13 @@ def pad_input_tokens(
     enable_attention_dp: bool,
     prefill_cuda_graph_backend: PrefillCudaGraphBackend,
     prefill_cuda_graph_num_tokens: list[int],
+    step_input_stager: "StepInputStager | None" = None,
 ) -> int:
     """Pad the batch to the token count its prefill graph or attention-DP peers need.
 
     Zeroes the padded rows of ``input_ids`` and ``position_ids``, whose last
     dimension is the token one, and returns the token count the model runs on.
+    With ``step_input_stager`` the zeroing is staged on it.
     """
     padded_num_tokens, can_run_prefill_cuda_graph, attn_all_rank_num_tokens = get_padding_params(
         num_tokens,
@@ -158,6 +161,11 @@ def pad_input_tokens(
         attn_metadata.all_rank_num_tokens = attn_all_rank_num_tokens
     if attn_metadata.padded_num_tokens is None:
         return num_tokens
+    if step_input_stager is not None:
+        # Staged write; see the window banner in DecoderRunner._forward_decoder.
+        step_input_stager.zero(input_ids[num_tokens:padded_num_tokens])
+        step_input_stager.zero(position_ids[..., num_tokens:padded_num_tokens])
+        return padded_num_tokens
     input_ids[num_tokens:padded_num_tokens].fill_(0)
     position_ids[..., num_tokens:padded_num_tokens].fill_(0)
     return padded_num_tokens
diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/graphs.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/graphs.py
index 4046fd5ca8ea..1dc3fca7e80f 100644
--- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/graphs.py
+++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/graphs.py
@@ -487,11 +487,14 @@ def run(
         enable_spec_decode: bool,
         gather_context_logits: bool,
         position_id_offsets: torch.Tensor | None,
+        before_replay: Callable[[], None] | None = None,
     ) -> Any:
         """Run the model by replaying the graph for ``key``, capturing it first if needed.
 
         Without a key, record the model body while a breakable prefill graph
         captures, run through one when it covers the tokens, otherwise eagerly.
+        ``before_replay`` runs right before the replay of an already captured
+        graph, after the draft-cache metadata refresh.
         """
         with with_shared_pool(self.cuda_graph_runner.get_graph_pool()):
             if self.captures_model_body(key):
@@ -509,6 +512,7 @@ def run(
                 enable_spec_decode=enable_spec_decode,
                 gather_context_logits=gather_context_logits,
                 position_id_offsets=position_id_offsets,
+                before_replay=before_replay,
             )
 
     def _run(
@@ -524,6 +528,7 @@ def _run(
         enable_spec_decode: bool,
         gather_context_logits: bool,
         position_id_offsets: torch.Tensor | None,
+        before_replay: Callable[[], None] | None = None,
     ) -> Any:
         breakable_runner = self.breakable_cuda_graph_runner
         if key is None:
@@ -574,6 +579,8 @@ def capture_model(model_inputs: dict[str, Any]) -> Any:
                     attn_metadata, draft_kv_cache_manager
                 )
                 try:
+                    if before_replay is not None:
+                        before_replay()
                     with MoeLoadBalancerIterContext(moe_load_balancer):
                         outputs = self.cuda_graph_runner.replay(
                             key, inputs, position_id_offsets=position_id_offsets
diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/inputs.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/inputs.py
index ac57b3ef45eb..0aa1d1c35c76 100644
--- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/inputs.py
+++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/inputs.py
@@ -25,6 +25,7 @@
     ResourceManager,
     ResourceManagerType,
 )
+from tensorrt_llm._torch.pyexecutor.runtime_hooks import DecodePrologue, StepInputStager
 from tensorrt_llm._torch.pyexecutor.sampler import SampleStateTensors
 from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests
 from tensorrt_llm._torch.pyexecutor.scheduler.scheduler import RemoteTailPhase
@@ -311,6 +312,11 @@ class DecoderInputPreparer:
     path overrides ``_new_extra_inputs`` and ``_prepare_inputs_fast_path``.
     """
 
+    # The model's decode prologue (RuntimeHooks.decode_prologue), set by the runner.
+    decode_prologue: DecodePrologue | None = None
+    # The model's step-input stager while it records this preparer's writes, else None.
+    step_input_stager: StepInputStager | None = None
+
     def __init__(
         self,
         *,
@@ -518,11 +524,7 @@ def pad_tokens(
                 mapping=self.mapping,
                 dist=self.dist,
             )
-        return pad_input_tokens(
-            attn_metadata,
-            num_tokens,
-            num_contexts,
-            attn_all_rank_num_tokens,
+        padding = dict(
             input_ids=self._buffers.input_ids_cuda,
             position_ids=self._position_ids_buffer(),
             dist=self.dist,
@@ -534,6 +536,12 @@ def pad_tokens(
             ),
             prefill_cuda_graph_num_tokens=self._config.prefill_cuda_graph_num_tokens,
         )
+        if self.step_input_stager is not None:
+            # Staged write; see the window banner in DecoderRunner._forward_decoder.
+            padding["step_input_stager"] = self.step_input_stager
+        return pad_input_tokens(
+            attn_metadata, num_tokens, num_contexts, attn_all_rank_num_tokens, **padding
+        )
 
     def model_inputs(
         self,
@@ -836,10 +844,18 @@ def prepare_inputs(
         self._copy_position_ids_to_device(host)
 
         if enable_spec_decode:
-            self._buffers.gather_ids_cuda[: len(host.lists.gather_ids)].copy_(
-                torch.tensor(host.lists.gather_ids, dtype=torch.int, pin_memory=prefer_pinned()),
-                non_blocking=True,
+            gather_ids = torch.tensor(
+                host.lists.gather_ids, dtype=torch.int, pin_memory=prefer_pinned()
             )
+            if self.step_input_stager is not None:
+                # Staged write; see the window banner in DecoderRunner._forward_decoder.
+                self.step_input_stager.copy(
+                    self._buffers.gather_ids_cuda[: len(host.lists.gather_ids)], gather_ids
+                )
+            else:
+                self._buffers.gather_ids_cuda[: len(host.lists.gather_ids)].copy_(
+                    gather_ids, non_blocking=True
+                )
 
         self._prepare_attn_metadata(
             host,
@@ -1705,24 +1721,52 @@ def previous_seq_slots_device() -> torch.Tensor:
         assert total_num_tokens <= self._config.max_num_tokens, (
             f"total_num_tokens ({total_num_tokens}) should be less than or equal to max_num_tokens ({self._config.max_num_tokens})"  # noqa: E501
         )
+        stager = self.step_input_stager
         # if exist requests that do not have previous batch, copy input_ids and draft_tokens
         if num_tokens > 0:
             input_ids = torch.tensor(input_ids, dtype=torch.int, pin_memory=prefer_pinned())
-            self._buffers.input_ids_cuda[:num_tokens].copy_(input_ids, non_blocking=True)
+            if stager is not None:
+                # Staged write; see the window banner in DecoderRunner._forward_decoder.
+                stager.copy(self._buffers.input_ids_cuda[:num_tokens], input_ids)
+            else:
+                self._buffers.input_ids_cuda[:num_tokens].copy_(input_ids, non_blocking=True)
 
         if num_draft_tokens > 0:
             draft_tokens = torch.tensor(draft_tokens, dtype=torch.int, pin_memory=prefer_pinned())
-            self._buffers.draft_tokens_cuda[: len(draft_tokens)].copy_(
-                draft_tokens, non_blocking=True
-            )
+            if stager is not None:
+                # Staged write; see the window banner in DecoderRunner._forward_decoder.
+                stager.copy(self._buffers.draft_tokens_cuda[: len(draft_tokens)], draft_tokens)
+            else:
+                self._buffers.draft_tokens_cuda[: len(draft_tokens)].copy_(
+                    draft_tokens, non_blocking=True
+                )
         if self._config.is_spec_decode and len(num_accepted_draft_tokens) > 0:
             num_accepted_draft_tokens = torch.tensor(
                 num_accepted_draft_tokens, dtype=torch.int, pin_memory=prefer_pinned()
             )
-            self._buffers.num_accepted_draft_tokens_cuda[: len(num_accepted_draft_tokens)].copy_(
-                num_accepted_draft_tokens, non_blocking=True
+            accepted = self._buffers.num_accepted_draft_tokens_cuda[
+                : len(num_accepted_draft_tokens)
+            ]
+            if stager is not None:
+                # Staged write; see the window banner in DecoderRunner._forward_decoder.
+                stager.copy(accepted, num_accepted_draft_tokens)
+            else:
+                accepted.copy_(num_accepted_draft_tokens, non_blocking=True)
+        if next_draft_tokens_device is not None and stager is not None:
+            # Staged write; see the window banner in DecoderRunner._forward_decoder.
+            stager.overlap_inputs(
+                self._buffers,
+                previous_batch_indices,
+                previous_pos_indices,
+                num_tokens,
+                num_draft_tokens,
+                self.get_runtime_tokens_per_gen_step(runtime_draft_len),
+                len(extend_requests) - len(extend_dummy_requests),
+                new_tokens_device,
+                next_draft_tokens_device,
+                new_tokens_lens_device,
             )
-        if next_draft_tokens_device is not None:
+        elif next_draft_tokens_device is not None:
             # Initialize these two values to zeros
             self._buffers.previous_pos_id_offsets_cuda.zero_()
             self._buffers.previous_kv_lens_offsets_cuda.zero_()
@@ -1822,6 +1866,10 @@ def previous_seq_slots_device() -> torch.Tensor:
         # cover the previous-batch rows; padding-request rows are covered by
         # neither, so without this the embedding would gather stale contents
         # of the uninitialized buffer, which may exceed the vocabulary size.
+        if stager is not None and host.lists.unwritten_input_id_rows:
+            # A native write to input_ids_cuda inside the window needs a flush
+            # first; see the window banner in DecoderRunner._forward_decoder.
+            stager.flush()
         _zero_fill_row_spans(self._buffers.input_ids_cuda, host.lists.unwritten_input_id_rows)
 
         host.num_tokens = num_tokens
@@ -1870,6 +1918,11 @@ def _copy_position_ids_to_device(self, host: HostInputs) -> None:
                 self._buffers.mrope_position_ids_cuda[:, :, start_idx:end_idx].copy_(
                     segment[:, :, : end_idx - start_idx], non_blocking=True
                 )
+        elif self.step_input_stager is not None:
+            # Staged write; see the window banner in DecoderRunner._forward_decoder.
+            self.step_input_stager.position_ids(
+                self._buffers.position_ids_cuda[:total_num_tokens], host_position_ids
+            )
         else:
             self._buffers.position_ids_cuda[:total_num_tokens].copy_(
                 host_position_ids, non_blocking=True
@@ -2178,6 +2231,11 @@ def record_input_copy(self) -> None:
         self._prepare_inputs_event = torch.cuda.Event()
         self._prepare_inputs_event.record()
 
+    @property
+    def input_copy_event(self) -> torch.cuda.Event | None:
+        """The event ``record_input_copy`` recorded last."""
+        return self._prepare_inputs_event
+
     def wait_for_input_copy(self) -> None:
         if self._prepare_inputs_event is not None:
             self._prepare_inputs_event.synchronize()
@@ -2259,6 +2317,16 @@ def preprocess_inputs(
         Make some changes to the device inputs and avoid blocking the async data transfer
         """
         attn_meta = inputs.get("attn_metadata")
+        # A model's decode prologue may replace the device work below with one
+        # launch; it decides on the host before any device work and returns
+        # False to keep the code below.
+        prologue = self.decode_prologue
+        if prologue is not None and prologue(
+            self, inputs, enable_spec_decode=enable_spec_decode, runtime_draft_len=runtime_draft_len
+        ):
+            if self.guided_decoder is not None:
+                self.guided_decoder.token_event.record()
+            return inputs
         # Invalidate per-forward-pass caches so they are recomputed (and captured) on every _forward_step.
         if attn_meta is not None:
             attn_meta.on_update_kv_lens()
diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py
index 2480fecb6ce2..060f30b6b6a5 100644
--- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py
+++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py
@@ -32,6 +32,11 @@
 )
 from tensorrt_llm._torch.pyexecutor.guided_decoder import CapturableGuidedDecoder
 from tensorrt_llm._torch.pyexecutor.resource_manager import ResourceManager, ResourceManagerType
+from tensorrt_llm._torch.pyexecutor.runtime_hooks import (
+    DEFAULT_RUNTIME_HOOKS,
+    RuntimeHooks,
+    runtime_hooks_of,
+)
 from tensorrt_llm._torch.pyexecutor.sampler.sampler_common import SampleType
 from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests
 from tensorrt_llm._torch.pyexecutor.warmup_timer import _WarmupTimer
@@ -123,6 +128,9 @@ def __getattr__(self, name: str) -> Any:
 class DecoderRunner(ScheduledModelRunner):
     """Run scheduled decoder generation with its own buffers, graphs and metadata."""
 
+    # Step-loop behaviour the model opted into (RuntimeHooks); all off unless it declares otherwise.
+    runtime_hooks: RuntimeHooks = DEFAULT_RUNTIME_HOOKS
+
     # Optional execution paths that a subclass can turn off.
     _context_warmups_supported = True
     _context_graph_promotion_supported = True
@@ -195,6 +203,8 @@ def __init__(
         self._dummy_requests = self._make_dummy_requests()
         self._init_decoder_state()
         self._inputs = self._make_input_preparer(mm_item_scheduler)
+        self.runtime_hooks = runtime_hooks_of(self.model)
+        self._inputs.decode_prologue = self.runtime_hooks.decode_prologue
         self._graphs = self._make_cuda_graphs()
         self._initialize_moe_workspace_reclaimer()
         self._warmup = DecoderWarmup(
@@ -612,34 +622,100 @@ def _forward_decoder(
                     execution_requests.generation_requests,
                 )
 
-            with self._cuda_graph_metadata_prepare_scope(
-                attn_metadata,
-                key,
-                execution_requests,
-                execution_promoted_context_ids,
-                enable_spec_decode=enable_spec_decode,
-            ):
-                prepared = self._inputs.prepare_inputs(
+            # ==== Step-input window opens here ====
+            # When the model declares a step-input stager and begin_if_eligible
+            # returns it, every device write that prepare_inputs (and the
+            # metadata and KV-cache-manager code it calls) routes through
+            # ``step_input_stager`` (call sites marked "Staged write") is only
+            # recorded; close_input_window below, right before the graph
+            # replay, issues them all. Until then, do not read or write those
+            # buffers with any other device op: a read sees stale bytes and a
+            # write is overwritten. Route new writes through the stager or
+            # call stager.flush() first. Exception: input_ids_cuda and
+            # position_ids_cuda are also staged into the graph's static
+            # inputs, which the replay then does not refresh from them; to
+            # change them here, call stager.flush() and then write them
+            # natively. Nothing checks this contract in a production run.
+            # Contract: StepInputStager in pyexecutor/runtime_hooks.py.
+            stager = None
+            stager_source = self.runtime_hooks.step_input_stager
+            if stager_source is not None:
+                graph_replay = (
+                    can_run_graph
+                    and not self.cuda_graph_runner.needs_capture(key)
+                    and not self.cuda_graph_runner.is_warmup_only
+                )
+                stager = stager_source.begin_if_eligible(
+                    self._inputs,
                     execution_requests,
-                    kv_cache_manager,
                     attn_metadata,
                     spec_metadata,
                     new_tensors_device,
-                    cache_indirection_buffer,
-                    resource_manager,
-                    can_run_graph,
-                    execution_promoted_context_ids,
-                    use_lora_graph=use_lora_graph,
-                    enable_spec_decode=enable_spec_decode,
-                    runtime_draft_len=runtime_draft_len,
+                    graph_replay=graph_replay,
+                    static_inputs=(
+                        self.cuda_graph_runner.shared_static_tensors if graph_replay else None
+                    ),
                     is_dummy=is_dummy,
-                    gather_context_logits=gather_context_logits,
+                    enable_spec_decode=enable_spec_decode,
+                    promoted=bool(execution_promoted_context_ids),
                 )
+            try:
+                # The window spans the scope: the scope is a no-op whenever
+                # speculative decoding is on, which every stager requires.
+                with self._cuda_graph_metadata_prepare_scope(
+                    attn_metadata,
+                    key,
+                    execution_requests,
+                    execution_promoted_context_ids,
+                    enable_spec_decode=enable_spec_decode,
+                ):
+                    prepared = self._inputs.prepare_inputs(
+                        execution_requests,
+                        kv_cache_manager,
+                        attn_metadata,
+                        spec_metadata,
+                        new_tensors_device,
+                        cache_indirection_buffer,
+                        resource_manager,
+                        can_run_graph,
+                        execution_promoted_context_ids,
+                        use_lora_graph=use_lora_graph,
+                        enable_spec_decode=enable_spec_decode,
+                        runtime_draft_len=runtime_draft_len,
+                        is_dummy=is_dummy,
+                        gather_context_logits=gather_context_logits,
+                    )
+            except BaseException:
+                if stager is not None:
+                    stager.discard()
+                raise
             inputs = prepared.kwargs
             runtime_draft_len = prepared.runtime_draft_len
             if execution_promoted_context_ids:
                 self._inputs.record_promoted_iter_states(scheduled_requests)
-            self._inputs.record_input_copy()
+            close_input_window = None
+            window_open = stager is not None
+            if stager is None:
+                self._inputs.record_input_copy()
+            else:
+
+                def close_input_window() -> None:
+                    nonlocal window_open
+                    window_open = False
+                    # ==== Step-input window closes here ====
+                    # end() issues every staged write; nothing between
+                    # begin_if_eligible above and this call may have touched
+                    # the staged buffers with other device ops. The input-copy
+                    # event must follow the launch (host page-table reads
+                    # happen at execution time). A nonzero
+                    # static_inputs_tokens() makes the replay skip its
+                    # input_ids / position_ids copy.
+                    stager.end(record_event=False)
+                    self._inputs.record_input_copy()
+                    stager.set_completion_event(self._inputs.input_copy_event)
+                    self.cuda_graph_runner.staged_static_input_tokens = (
+                        stager.static_inputs_tokens()
+                    )
 
             def run_model(model_inputs: dict[str, Any]) -> Any:
                 with MoeLoadBalancerIterContext(moe_load_balancer):
@@ -661,24 +737,30 @@ def postprocess_inputs(model_inputs: dict[str, Any]) -> None:
 
             graph_key = key if can_run_graph else None
             captures_model_body = self._graphs.captures_model_body(graph_key)
-            outputs = self._graphs.run(
-                graph_key,
-                inputs,
-                run_model,
-                postprocess_inputs,
-                attn_metadata=attn_metadata,
-                draft_kv_cache_manager=draft_kv_cache_manager,
-                moe_load_balancer=moe_load_balancer,
-                enable_spec_decode=enable_spec_decode,
-                gather_context_logits=gather_context_logits,
-                position_id_offsets=(
-                    self._buffers.previous_pos_id_offsets_cuda
-                    if enable_spec_decode
-                    and not self._config.disable_overlap_scheduler
-                    and attn_metadata.kv_cache_manager is not None
-                    else None
-                ),
-            )
+            try:
+                outputs = self._graphs.run(
+                    graph_key,
+                    inputs,
+                    run_model,
+                    postprocess_inputs,
+                    attn_metadata=attn_metadata,
+                    draft_kv_cache_manager=draft_kv_cache_manager,
+                    moe_load_balancer=moe_load_balancer,
+                    enable_spec_decode=enable_spec_decode,
+                    gather_context_logits=gather_context_logits,
+                    position_id_offsets=(
+                        self._buffers.previous_pos_id_offsets_cuda
+                        if enable_spec_decode
+                        and not self._config.disable_overlap_scheduler
+                        and attn_metadata.kv_cache_manager is not None
+                        else None
+                    ),
+                    before_replay=close_input_window,
+                )
+            except BaseException:
+                if window_open:
+                    stager.discard()
+                raise
             if captures_model_body:
                 return outputs
 
diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py
index 07a5e93717aa..c825736a372c 100644
--- a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py
+++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py
@@ -134,6 +134,7 @@
 
 if TYPE_CHECKING:
     from tensorrt_llm._torch.attention.backends.interface import AttentionMetadata
+    from tensorrt_llm._torch.pyexecutor.runtime_hooks import StepInputStager
     from tensorrt_llm.llmapi.llm_args import DecodingBaseConfig
 
 KV_CACHE_ITERATION_STATS_DELTA_FIELDS = _KV_CACHE_ITERATION_STATS_DELTA_FIELDS
@@ -1249,6 +1250,8 @@ def _spec_recompute_claim_limit(
 class KVCacheManagerV2(BaseResourceManager):
     # Sparse managers attach a metadata batch during initialization.
     sparse_metadata_batch: Batch | None = None
+    # A step-input stager while it records this manager's writes, else None.
+    step_input_stager: Optional["StepInputStager"] = None
     # Filled lazily by _cold_pool_group_membership(); the grouping is fixed after construction.
     # Declared on the class so it is present even when an instance is built without running __init__.
     _cold_pool_group_membership_cache: Optional[tuple[tuple[int, frozenset[int]], ...]] = None
@@ -6754,6 +6757,16 @@ def copy_batch_block_offsets(
             )
             return
 
+        if self.step_input_stager is not None:
+            # Staged write; see the window banner in DecoderRunner._forward_decoder.
+            self.step_input_stager.block_offsets(
+                self.host_kv_cache_block_offsets,
+                dst_tensor,
+                copy_idx.tolist(),
+                self._index_scale_ints,
+                self.kv_offset.tolist(),
+            )
+            return
         copy_batch_block_offsets_to_device(
             self.host_kv_cache_block_offsets,
             dst_tensor,
diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py
index 9d326ab0c6e4..f06c73f28213 100644
--- a/tensorrt_llm/_torch/pyexecutor/py_executor.py
+++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py
@@ -119,6 +119,7 @@
                             get_from_waiting_queue, merge_requests)
 from .resource_manager import (NoFreeSlotsError, ResourceManager,
                                ResourceManagerType, request_context)
+from .runtime_hooks import ExecutorEvents, runtime_hooks_of
 from .sampler import (AsyncWorkerMixin, Sampler, SamplerEvent, SampleState,
                       SampleStateTensors)
 from .scheduler import (RequestScheduler, ScheduledRequests,
@@ -409,6 +410,16 @@ class PendingEncoderStep:
 
 
 class PyExecutor:
+    # True while the overlap loop runs with the execution stream current
+    # (_overlap_loop_stream_scope).
+    _loop_on_execution_stream = False
+    # Whether the model opted into running the overlap loop on the execution
+    # stream (RuntimeHooks.loop_on_execution_stream); read once at construction.
+    _execution_stream_loop_armed = False
+    # The model's executor events (RuntimeHooks.executor_events); read once at
+    # construction.
+    _executor_events: Optional[ExecutorEvents] = None
+
     # Minimum number of async micro batches for async PP execution.
     # This is a trade-off between memory usage and performance.
     # If the number of micro batches is too small, the executor will spend too much time in synchronization.
@@ -526,6 +537,9 @@ def __init__(
         # related modules
         self.resource_manager = resource_manager
         self.model_engine = model_engine
+        hooks = runtime_hooks_of(model_engine.model)
+        self._execution_stream_loop_armed = hooks.loop_on_execution_stream
+        self._executor_events = hooks.executor_events
         self._enable_adp_dummy_fixes = getattr(model_engine,
                                                "_enable_adp_dummy_fixes", False)
         self._enable_scheduler_aware_adp_dummy = getattr(
@@ -920,6 +934,10 @@ def __init__(
         # Ensure the default stream waits for execution_stream to complete
         # before subsequent operations.
         torch.cuda.current_stream().wait_stream(self.execution_stream)
+        if self._executor_events is not None:
+            # The model's actions after warmup and graph capture, on every rank
+            # (some are collective, so every rank reaches this point).
+            self._executor_events.after_warmup()
         self.encoder_launch_executor = (ThreadPoolExecutor(
             max_workers=1, thread_name_prefix="encoder-launch")
                                         if self.is_encoder_decoder else None)
@@ -3025,6 +3043,9 @@ def _executor_loop_cleanup(self):
                     f"Error waiting on PP send handles during cleanup: "
                     f"{traceback.format_exc()}")
 
+        if self._executor_events is not None and self._event_loop_completed:
+            self._executor_events.at_shutdown()
+
     def _pp_schedule_and_propagate(self, microbatch_id: int):
         """The first PP rank schedules the requests and propagates the result to all other PP ranks."""
 
@@ -5725,11 +5746,44 @@ def _wait_for_model_engine_input_copy(self):
         if wait_for_input_copy is not None:
             wait_for_input_copy()
 
+    @contextmanager
+    def _overlap_loop_stream_scope(self):
+        """Run the overlap loop with the execution stream current, when the model opted in.
+
+        _forward_step makes the execution stream wait for the current stream
+        before every forward and the current stream wait for the execution
+        stream after it. On the thread's default stream, which in steady decode
+        carries only these waits and timing events, step N+1's inputs wait for
+        decode graph N through two cross-stream event hops. With the execution
+        stream current the loop's stream work is queued in the same order on one
+        stream and _forward_step skips both waits.
+
+        Not taken with a two-model drafter, a guided decoder, a KV connector,
+        an encoder-decoder model or multimodal encoder item scheduling, whose
+        side streams synchronize with the current stream.
+        """
+        if not (self._execution_stream_loop_armed and self.drafter is None and
+                self.draft_model_engine is None and self.guided_decoder is None
+                and self.kv_connector_manager is None
+                and not self.is_encoder_decoder
+                and not self._mm_encoder_item_scheduling_enabled):
+            yield
+            return
+        with torch.cuda.stream(self.execution_stream):
+            self._loop_on_execution_stream = True
+            try:
+                yield
+            finally:
+                self._loop_on_execution_stream = False
+
     def _executor_loop_overlap(self):
         torch.cuda.set_device(self.device_id)
         # ensure the context is created, otherwise, some MPI calls will fail.
         CUASSERT(cudart.cudaSetDevice(self.device_id))
-        with self._profiler() as profile_step, self.hang_detector:
+        if self._executor_events is not None:
+            self._executor_events.before_overlap_loop()
+        with self._overlap_loop_stream_scope(), self._profiler() as profile_step, \
+                self.hang_detector:
             iter_start_time = time.time()
             iter_stats = None
             target_inputs = None
@@ -8484,8 +8538,10 @@ def forward(scheduled_requests, resource_manager, new_tensors_device,
             cache_indirection_buffer = self.sampler.get_cache_indirection()
 
             # Run model forward on the execution stream for proper synchronization
-            # with KVCacheTransferManager's onboard/offload operations.
-            self.execution_stream.wait_stream(torch.cuda.current_stream())
+            # with KVCacheTransferManager's onboard/offload operations. With the
+            # loop on the execution stream both waits are program order.
+            if not self._loop_on_execution_stream:
+                self.execution_stream.wait_stream(torch.cuda.current_stream())
             self._attach_encoder_output_to_execution_stream(scheduled_requests)
             with torch.cuda.stream(self.execution_stream):
                 outputs = forward(scheduled_requests, self.resource_manager,
@@ -8499,7 +8555,8 @@ def forward(scheduled_requests, resource_manager, new_tensors_device,
 
             # Ensure the default stream waits for execution_stream to complete
             # before downstream operations use the outputs.
-            torch.cuda.current_stream().wait_stream(self.execution_stream)
+            if not self._loop_on_execution_stream:
+                torch.cuda.current_stream().wait_stream(self.execution_stream)
 
             self._kv_connector_wait_for_save()
             if route_capture is not None:
diff --git a/tensorrt_llm/_torch/pyexecutor/runtime_hooks.py b/tensorrt_llm/_torch/pyexecutor/runtime_hooks.py
new file mode 100644
index 000000000000..24ffbccfbb8b
--- /dev/null
+++ b/tensorrt_llm/_torch/pyexecutor/runtime_hooks.py
@@ -0,0 +1,221 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Step-loop behaviour a model may opt into, read through ``runtime_hooks_of``.
+
+A model declares its hooks as ``runtime_hooks`` (an instance attribute, a class
+attribute or a property). Every field defaults to off, and a model that declares
+nothing gets ``RuntimeHooks()``: the step loop, the samplers and the attention
+metadata then run exactly as they do without this module. Every engine call site
+of a hook runs its original statements when the hook is off.
+
+The hooks change *where* and *how* the step loop moves data, never what a
+forward computes. The protocols below state only the methods an engine call
+site needs; the objects that implement them are owned by the model.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Protocol, Sequence
+
+if TYPE_CHECKING:
+    import torch
+
+
+class StepInputStager(Protocol):
+    """Records a decode step's input preparation and issues it as one launch.
+
+    ``begin_if_eligible`` opens a deferred-write window on a step it supports
+    and returns the stager, or returns None and the engine prepares the step
+    natively. While the window is open, the engine's call sites route their
+    device writes to the staging methods instead of issuing them. ``end``
+    issues every staged write in staging order, then the deferred native
+    operations; ``discard`` drops them after a failed preparation.
+
+    While its window is open the stager sets itself as ``step_input_stager`` on
+    the decoder's input preparer, the attention metadata, the speculative
+    metadata and the KV cache managers, and resets that attribute to None at
+    ``end`` and ``discard``. Every call site reads the attribute and runs its
+    original statement when it is None.
+    """
+
+    def begin_if_eligible(
+        self,
+        preparer: Any,
+        requests: Any,
+        attn_metadata: Any,
+        spec_metadata: Any,
+        new_tensors_device: Any,
+        *,
+        graph_replay: bool,
+        static_inputs: Optional[Dict[str, "torch.Tensor"]],
+        is_dummy: bool,
+        enable_spec_decode: bool,
+        promoted: bool,
+    ) -> Optional["StepInputStager"]:
+        """Open the window for this step, or return None to keep the native preparation.
+
+        ``preparer`` is the ``DecoderInputPreparer`` whose writes the window
+        records. ``static_inputs`` are the static input tensors of the CUDA
+        graph the step replays (None when it replays none).
+        """
+
+    def copy(
+        self, dst: "torch.Tensor", src: "torch.Tensor", clamp_min: Optional[int] = None
+    ) -> None:
+        """Stage ``dst.copy_(src, non_blocking=True)`` for a host ``src``."""
+
+    def zero(self, dst: "torch.Tensor") -> None:
+        """Stage ``dst.zero_()``."""
+
+    def clamp_min(self, dst: "torch.Tensor", lo: int) -> None:
+        """Stage ``dst.clamp_(min=lo)``."""
+
+    def block_offsets(
+        self,
+        host_table: "torch.Tensor",
+        dst: "torch.Tensor",
+        copy_index: Sequence[int],
+        index_scales: Sequence[int],
+        kv_offsets: Sequence[int],
+    ) -> None:
+        """Stage the KV cache manager's batch copy of block offsets to the device."""
+
+    def position_ids(self, dst: "torch.Tensor", host_position_ids: "torch.Tensor") -> None:
+        """Stage the position-id upload and the CUDA graph's static copy of it."""
+
+    def overlap_inputs(
+        self,
+        buffers: Any,
+        previous_batch_indices: Sequence[int],
+        previous_pos_indices: Sequence[int],
+        num_tokens: int,
+        num_draft_tokens: int,
+        tokens_per_gen_step: int,
+        num_extend_requests: int,
+        new_tokens: "torch.Tensor",
+        next_draft_tokens: "torch.Tensor",
+        new_tokens_lens: "torch.Tensor",
+    ) -> None:
+        """Stage the overlap scheduler's gathers of the previous step's outputs into ``buffers``."""
+
+    def scheduler_metadata(self, metadata: Any) -> bool:
+        """Stage the sparse-attention indexer's scheduler metadata.
+
+        Returns False, having staged nothing, when the caller must defer the
+        native computation instead.
+        """
+
+    def defer(self, fn: Callable[[], None]) -> None:
+        """Run ``fn`` (native device work that reads staged values) after the launch."""
+
+    def flush(self) -> None:
+        """Issue every write staged so far, before the window closes."""
+
+    def end(self, record_event: bool = True) -> None:
+        """Close the window and issue the staged preparation."""
+
+    def discard(self) -> None:
+        """Close the window after a failed preparation; the staged writes are dropped."""
+
+    def set_completion_event(self, event: "torch.cuda.Event") -> None:
+        """Event the caller recorded right after ``end(record_event=False)``."""
+
+    def static_inputs_tokens(self) -> int:
+        """Tokens of the CUDA graph's static inputs the last launch wrote (0: none)."""
+
+
+class SamplerStoreWriter(Protocol):
+    """Stores a speculative sampler step's outputs and copies the store to the host in one launch."""
+
+    def supported(
+        self,
+        store: Any,
+        new_tokens: "torch.Tensor",
+        next_new_tokens: "torch.Tensor",
+        new_tokens_lens: "torch.Tensor",
+        next_draft_tokens: "torch.Tensor",
+    ) -> bool:
+        """Whether this step's output tensors fit the writer; False keeps the native stores."""
+
+    def store_and_copy_out(
+        self,
+        store: Any,
+        slots: Sequence[int],
+        new_tokens: "torch.Tensor",
+        next_new_tokens: "torch.Tensor",
+        new_tokens_lens: "torch.Tensor",
+        next_draft_tokens: "torch.Tensor",
+        host_new_tokens: "torch.Tensor",
+        host_new_tokens_lens: "torch.Tensor",
+        host_next_draft_tokens: "torch.Tensor",
+        record_event: bool = True,
+    ) -> None:
+        """Scatter the outputs into the store rows ``slots`` and copy the store to the host tensors."""
+
+    def set_completion_event(self, event: "torch.cuda.Event") -> None:
+        """Event the caller recorded right after ``store_and_copy_out(record_event=False)``."""
+
+
+class DecodePrologue(Protocol):
+    """Replaces the decode step's input prologue with one launch.
+
+    Called first in the step's input preprocessing, inside CUDA graph capture
+    as well. It decides on the host before any device work; returning False
+    keeps the native prologue, and returning True means every engine-visible
+    tensor the native prologue writes holds the value it would have written.
+    """
+
+    def __call__(
+        self, preparer: Any, inputs: Any, *, enable_spec_decode: bool, runtime_draft_len: int
+    ) -> bool:
+        """``preparer`` is the ``DecoderInputPreparer`` running the step's preprocessing."""
+
+
+class ExecutorEvents(Protocol):
+    """Points in the executor's life a model may act at."""
+
+    def after_warmup(self) -> None:
+        """After warmup and CUDA graph capture, before the first request is served."""
+
+    def before_overlap_loop(self) -> None:
+        """Right before the overlap executor loop starts, on every rank; must not raise."""
+
+    def at_shutdown(self) -> None:
+        """When the executor loop has completed, during cleanup; must not raise."""
+
+
+@dataclass(frozen=True)
+class RuntimeHooks:
+    """The step-loop hooks a model opts into. All off by default."""
+
+    # Run the overlap executor loop with the execution stream current, so the
+    # loop's own stream work and the forward are ordered on one stream.
+    loop_on_execution_stream: bool = False
+    # Decode input preparation recorded and issued as one launch (see StepInputStager).
+    step_input_stager: Optional[StepInputStager] = None
+    # The speculative sampler's store update and host copy in one launch (see SamplerStoreWriter).
+    sampler_store_writer: Optional[SamplerStoreWriter] = None
+    # The decode step's input prologue in one launch (see DecodePrologue).
+    decode_prologue: Optional[DecodePrologue] = None
+    # Actions after warmup, before the overlap loop and at shutdown (see ExecutorEvents).
+    executor_events: Optional[ExecutorEvents] = None
+
+
+DEFAULT_RUNTIME_HOOKS = RuntimeHooks()
+
+
+def runtime_hooks_of(model: object) -> RuntimeHooks:
+    """The hooks ``model`` declares, or the all-off default when it declares none.
+
+    Declaring hooks is optional for every model class, so absence is the common
+    case and means the default.
+    """
+    hooks = getattr(model, "runtime_hooks", None)
+    if hooks is None:
+        return DEFAULT_RUNTIME_HOOKS
+    if not isinstance(hooks, RuntimeHooks):
+        raise TypeError(
+            f"{type(model).__name__}.runtime_hooks must be a RuntimeHooks, got {type(hooks).__name__}"
+        )
+    return hooks
diff --git a/tensorrt_llm/_torch/speculative/eagle3.py b/tensorrt_llm/_torch/speculative/eagle3.py
index 15993e5c26ba..0328e29746cc 100644
--- a/tensorrt_llm/_torch/speculative/eagle3.py
+++ b/tensorrt_llm/_torch/speculative/eagle3.py
@@ -348,8 +348,13 @@ def prepare(self):
                                      dtype=torch.int,
                                      device='cpu',
                                      pin_memory=prefer_pinned())
-        self.batch_indices_cuda[:num_seqs].copy_(batch_indices,
-                                                 non_blocking=True)
+        if self.step_input_stager is not None:
+            # Staged write; see the window banner in DecoderRunner._forward_decoder.
+            self.step_input_stager.copy(self.batch_indices_cuda[:num_seqs],
+                                        batch_indices)
+        else:
+            self.batch_indices_cuda[:num_seqs].copy_(batch_indices,
+                                                     non_blocking=True)
         # Snapshot the real (unpadded) scheduled token count before the decrement
         # below. maybe_capture_hidden_states runs during the target forward and
         # must bound the capture by this count; self.num_tokens is about to be
diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py
index 8edccc3e3f8a..a8bdaf4d4589 100644
--- a/tensorrt_llm/_torch/speculative/interface.py
+++ b/tensorrt_llm/_torch/speculative/interface.py
@@ -19,7 +19,7 @@
 from contextlib import contextmanager
 from dataclasses import dataclass, field
 from enum import IntEnum
-from typing import TYPE_CHECKING, List, Optional, Protocol, Type
+from typing import TYPE_CHECKING, ClassVar, List, Optional, Protocol, Type
 
 import torch
 from packaging.version import Version
@@ -37,6 +37,7 @@
 if TYPE_CHECKING:
     from ..pyexecutor.guided_decoder import CapturableGuidedDecoder
     from ..pyexecutor.llm_request import LlmRequest
+    from ..pyexecutor.runtime_hooks import StepInputStager
 
 if IS_FLASHINFER_AVAILABLE:
     import flashinfer
@@ -459,6 +460,8 @@ class SpecMetadata:
     """
     Metadata for speculative decoding.
     """
+    # A step-input stager while it records this metadata's writes, else None.
+    step_input_stager: ClassVar[Optional["StepInputStager"]] = None
     # The max number of requests in a single batch.
     max_num_requests: int
     # The number of draft layers. (Also the number of draft tokens for the linear tree.)
@@ -881,6 +884,11 @@ def _populate_request_rng_state(
                                        device='cuda')
 
         def _upload(dst: torch.Tensor, values: list[int]) -> None:
+            if self.step_input_stager is not None:
+                # Staged write; see the window banner in DecoderRunner._forward_decoder.
+                self.step_input_stager.copy(
+                    dst[:len(values)], torch.tensor(values, dtype=torch.int64))
+                return
             dst[:len(values)].copy_(torch.tensor(values,
                                                  dtype=torch.int64,
                                                  pin_memory=prefer_pinned()),
diff --git a/tensorrt_llm/_torch/speculative/spec_sampler_base.py b/tensorrt_llm/_torch/speculative/spec_sampler_base.py
index 5fb5b185dc0b..dc40cca3a7c5 100644
--- a/tensorrt_llm/_torch/speculative/spec_sampler_base.py
+++ b/tensorrt_llm/_torch/speculative/spec_sampler_base.py
@@ -27,9 +27,11 @@
 
 import torch
 
+from ..._utils import prefer_pinned
 from ...sampling_params import SamplingParams
 from ..pyexecutor.llm_request import LlmRequest, LlmRequestState, get_draft_token_length
 from ..pyexecutor.resource_manager import BaseResourceManager
+from ..pyexecutor.runtime_hooks import SamplerStoreWriter
 from ..pyexecutor.sampler import (
     DEFAULT_BEAM_IDX,
     AsyncWorkerMixin,
@@ -167,6 +169,10 @@ class SpecSampler(Sampler[SampleStateSpec], AsyncWorkerMixin):
 
     SampleState = SampleStateSpec
 
+    # The model's one-launch store writer, bound once by instantiate_sampler; None keeps
+    # the native stores and host copies.
+    store_writer: Optional[SamplerStoreWriter] = None
+
     def is_generation_model(self) -> bool:
         return True
 
@@ -442,8 +448,11 @@ def sample_async(
             for r in sampling_requests
         ]
 
-        slots = torch.as_tensor([r.py_seq_slot for r in sampling_requests], dtype=torch.long)
-        slots = slots.to(device="cuda", non_blocking=True)
+        store_writer = self.store_writer
+        slots = None
+        if store_writer is None:
+            slots = torch.as_tensor([r.py_seq_slot for r in sampling_requests], dtype=torch.long)
+            slots = slots.to(device="cuda", non_blocking=True)
 
         o_new_tokens = outputs["new_tokens"][num_skip : num_skip + num_sampling_requests]
         o_new_tokens_lens = outputs["new_tokens_lens"][num_skip : num_skip + num_sampling_requests]
@@ -514,11 +523,28 @@ def sample_async(
         elif o_next_new_tokens.shape[1] > next_new_tokens_width:
             o_next_new_tokens = o_next_new_tokens[:, :next_new_tokens_width]
 
-        # Use index_copy_ for efficient copying (slots are unique)
-        self.store.new_tokens.squeeze(-1).T.index_copy_(0, slots, o_new_tokens)
-        self.store.next_new_tokens.squeeze(-1).T.index_copy_(0, slots, o_next_new_tokens)
-        self.store.new_tokens_lens.index_copy_(0, slots, o_new_tokens_lens)
-        self.store.next_draft_tokens.index_copy_(0, slots, o_next_draft_tokens)
+        # The writer stores the four outputs and copies the three store tensors
+        # only; a step that publishes verify_lens keeps the native path.
+        staged_store = (
+            store_writer is not None
+            and o_verify_lens is None
+            and prefer_pinned()
+            and not self._async_worker_active()
+            and store_writer.supported(
+                self.store, o_new_tokens, o_next_new_tokens, o_new_tokens_lens, o_next_draft_tokens
+            )
+        )
+        if not staged_store:
+            if slots is None:
+                slots = torch.as_tensor(
+                    [r.py_seq_slot for r in sampling_requests], dtype=torch.long
+                )
+                slots = slots.to(device="cuda", non_blocking=True)
+            # Use index_copy_ for efficient copying (slots are unique)
+            self.store.new_tokens.squeeze(-1).T.index_copy_(0, slots, o_new_tokens)
+            self.store.next_new_tokens.squeeze(-1).T.index_copy_(0, slots, o_next_new_tokens)
+            self.store.new_tokens_lens.index_copy_(0, slots, o_new_tokens_lens)
+            self.store.next_draft_tokens.index_copy_(0, slots, o_next_draft_tokens)
 
         # Create sample state with async D2H copy
         device_tensors = SampleStateTensorsSpec(
@@ -527,13 +553,43 @@ def sample_async(
             next_draft_tokens=self.store.next_draft_tokens,
         )
 
-        host_tensors = SampleStateTensorsSpec(
-            new_tokens=self._copy_to_host(self.store.new_tokens),
-            new_tokens_lens=self._copy_to_host(self.store.new_tokens_lens),
-            next_draft_tokens=self._copy_to_host(self.store.next_draft_tokens),
-            verify_lens=(self._copy_to_host(o_verify_lens) if o_verify_lens is not None else None),
-        )
+        if staged_store:
+            # staged_store implies prefer_pinned(), so these buffers are pinned.
+            host_tensors = SampleStateTensorsSpec(
+                new_tokens=torch.empty_like(
+                    self.store.new_tokens, device="cpu", pin_memory=prefer_pinned()
+                ),
+                new_tokens_lens=torch.empty_like(
+                    self.store.new_tokens_lens, device="cpu", pin_memory=prefer_pinned()
+                ),
+                next_draft_tokens=torch.empty_like(
+                    self.store.next_draft_tokens, device="cpu", pin_memory=prefer_pinned()
+                ),
+            )
+            store_writer.store_and_copy_out(
+                self.store,
+                [r.py_seq_slot for r in sampling_requests],
+                o_new_tokens,
+                o_next_new_tokens,
+                o_new_tokens_lens,
+                o_next_draft_tokens,
+                host_tensors.new_tokens,
+                host_tensors.new_tokens_lens,
+                host_tensors.next_draft_tokens,
+                record_event=False,
+            )
+        else:
+            host_tensors = SampleStateTensorsSpec(
+                new_tokens=self._copy_to_host(self.store.new_tokens),
+                new_tokens_lens=self._copy_to_host(self.store.new_tokens_lens),
+                next_draft_tokens=self._copy_to_host(self.store.next_draft_tokens),
+                verify_lens=(
+                    self._copy_to_host(o_verify_lens) if o_verify_lens is not None else None
+                ),
+            )
         sampler_event = self._record_sampler_event()
+        if staged_store:
+            store_writer.set_completion_event(sampler_event.cuda_event)
 
         # Add dummy draft tokens to context requests for KV cache preparation
         for request in finished_context_requests:
diff --git a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py
index 1aad3d7d79c4..3fbef1289654 100644
--- a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py
+++ b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py
@@ -5327,3 +5327,103 @@ def _run_indexer():
     assert "indexer_topk_out_buffer" in metadata.cuda_graph_buffers.buffers, (
         "indexer topk-output buffer must be drawn from the cuda_graph_buffers arena"
     )
+
+
+# --- The subclass seams of Indexer.sparse_attn_indexer ----------------------------------------
+
+
+def _seam_indexer():
+    indexer = Indexer.__new__(Indexer)
+    torch.nn.Module.__init__(indexer)
+    return indexer
+
+
+def test_reused_topk_copy_is_the_row_copy():
+    gen = torch.Generator().manual_seed(0)
+    shared = torch.randint(0, 1000, (4, 16), generator=gen, dtype=torch.int32)
+    buffer = torch.full((12, 16), -1, dtype=torch.int32)
+    expected = buffer.clone()
+    expected[2:5, :] = shared[:3, :]
+
+    _seam_indexer()._copy_reused_topk(
+        SimpleNamespace(shared_topk_indices=shared), buffer, 2, 3, 3, False, None
+    )
+
+    assert torch.equal(buffer, expected)
+
+
+def test_fill_decode_topk_leaves_the_rows_to_the_decode_branch():
+    buffer = torch.full((6, 8), -1, dtype=torch.int32)
+
+    filled = _seam_indexer()._fill_decode_topk(SimpleNamespace(), None, None, buffer, 0, 6)
+
+    assert filled is False
+    assert bool((buffer == -1).all())
+
+
+def _stored_rows_inline(
+    indexer,
+    topk_indices_buffer,
+    metadata,
+    num_contexts,
+    num_generations,
+    num_gen_tokens,
+    token_offset,
+    has_prefill,
+    has_decode,
+    is_generation,
+):
+    """The statements ``sparse_attn_indexer`` ran in place of ``_store_mtp_shared_rows``."""
+    rows = None
+    if has_decode:
+        next_n = num_gen_tokens // num_generations
+        gen_topk = topk_indices_buffer[token_offset : token_offset + num_gen_tokens]
+        rows = indexer._mtp_last_accepted_rows(
+            gen_topk, metadata, num_contexts, num_generations, next_n
+        )
+    if has_prefill:
+        ctx_last = torch.cumsum(metadata.seq_lens_cuda[:num_contexts].to(torch.long), dim=0) - 1
+        ctx_rows = topk_indices_buffer[ctx_last]
+        rows = ctx_rows if rows is None else torch.cat([ctx_rows, rows])
+    if rows is not None:
+        shared_topk_indices = metadata.shared_topk_indices
+        if shared_topk_indices is None:
+            metadata.shared_topk_indices = rows.contiguous()
+        else:
+            row_start = num_contexts if is_generation else 0
+            shared_topk_indices[row_start : row_start + rows.shape[0], : rows.shape[1]].copy_(rows)
+
+
+@pytest.mark.parametrize("has_prefill, has_decode", [(False, True), (True, True), (True, False)])
+@pytest.mark.parametrize("is_generation", [True, False])
+@pytest.mark.parametrize("preset", [False, True], ids=["unset", "preset"])
+@pytest.mark.parametrize("accepted", [None, (0, 0, 1, 4, 2)], ids=["no_counts", "counts"])
+def test_mtp_shared_rows_are_the_inline_statements(
+    has_prefill, has_decode, is_generation, preset, accepted
+):
+    num_contexts, num_generations, next_n = 2, 3, 4
+    seq_lens = torch.tensor([5, 3] + [next_n] * num_generations, dtype=torch.int32)
+    token_offset = int(seq_lens[:num_contexts].sum())
+    num_gen_tokens = num_generations * next_n
+    gen = torch.Generator().manual_seed(1)
+    buffer = torch.randint(
+        0, 1000, (token_offset + num_gen_tokens, 6), generator=gen, dtype=torch.int32
+    )
+
+    def metadata():
+        return SimpleNamespace(
+            seq_lens_cuda=seq_lens.clone(),
+            mtp_num_accepted=None if accepted is None else torch.tensor(accepted),
+            shared_topk_indices=(torch.full((8, 8), -7, dtype=torch.int32) if preset else None),
+        )
+
+    args = (num_contexts, num_generations, num_gen_tokens, token_offset)
+    flags = (has_prefill, has_decode, is_generation)
+    expected, got = metadata(), metadata()
+    indexer = _seam_indexer()
+    _stored_rows_inline(indexer, buffer.clone(), expected, *args, *flags)
+    seam_buffer = buffer.clone()
+    indexer._store_mtp_shared_rows(seam_buffer, got, *args, *flags)
+
+    assert torch.equal(seam_buffer, buffer)
+    assert torch.equal(got.shared_topk_indices, expected.shared_topk_indices)
diff --git a/tests/unittest/_torch/attention/test_fmha_manager.py b/tests/unittest/_torch/attention/test_fmha_manager.py
index 7cfb59eee093..397bfd93a4b5 100644
--- a/tests/unittest/_torch/attention/test_fmha_manager.py
+++ b/tests/unittest/_torch/attention/test_fmha_manager.py
@@ -974,3 +974,60 @@ def __init__(self, attn: TrtllmAttention) -> None:
     assert isinstance(attn._fmha_manager.fmha_libs[0], _NewFmha)
     assert quant_states_during_construction == [True]
     assert events == [("support", "old", None)]
+
+
+def _listed_fmha(name: str, available: bool, checked: list) -> type:
+    """An FMHA library class that records its availability check."""
+
+    def is_available(cls, attn) -> bool:
+        checked.append((name, attn))
+        return available
+
+    return type(
+        name,
+        (),
+        {"is_available": classmethod(is_available), "__init__": lambda self, attn: None},
+    )
+
+
+@pytest.mark.parametrize("declared", [False, True], ids=["undeclared", "empty"])
+def test_manager_without_prepended_libraries_builds_the_registry_list(declared: bool) -> None:
+    checked: list = []
+    registry = [_listed_fmha("First", True, checked), _listed_fmha("Second", True, checked)]
+    attn = FakeAttention()
+    if declared:
+        attn.prepended_fmha_lib_classes = ()
+
+    with patch.object(fmha_manager, "get_enabled_fmha_lib_classes", return_value=registry):
+        manager = FmhaManager(attn)
+
+    assert [type(lib) for lib in manager.fmha_libs] == registry
+    assert TrtllmAttention.prepended_fmha_lib_classes == ()
+
+
+def test_manager_tries_a_prepended_library_first() -> None:
+    checked: list = []
+    registry = [_listed_fmha("First", True, checked), _listed_fmha("Second", True, checked)]
+    prepended = _listed_fmha("Prepended", True, checked)
+    attn = FakeAttention()
+    attn.prepended_fmha_lib_classes = (prepended,)
+
+    with patch.object(fmha_manager, "get_enabled_fmha_lib_classes", return_value=registry):
+        manager = FmhaManager(attn)
+
+    assert [type(lib) for lib in manager.fmha_libs] == [prepended, *registry]
+    assert [name for name, _ in checked] == ["Prepended", "First", "Second"]
+    assert all(checked_attn is attn for _, checked_attn in checked)
+
+
+def test_manager_skips_an_unavailable_prepended_library() -> None:
+    checked: list = []
+    registry = [_listed_fmha("First", True, checked)]
+    attn = FakeAttention()
+    attn.prepended_fmha_lib_classes = (_listed_fmha("Prepended", False, checked),)
+
+    with patch.object(fmha_manager, "get_enabled_fmha_lib_classes", return_value=registry):
+        manager = FmhaManager(attn)
+
+    assert [type(lib) for lib in manager.fmha_libs] == registry
+    assert [name for name, _ in checked] == ["Prepended", "First"]
diff --git a/tests/unittest/_torch/executor/test_gpu_keepalive.py b/tests/unittest/_torch/executor/test_gpu_keepalive.py
index 5c3318c00690..51d9b63b277e 100644
--- a/tests/unittest/_torch/executor/test_gpu_keepalive.py
+++ b/tests/unittest/_torch/executor/test_gpu_keepalive.py
@@ -315,6 +315,7 @@ class Stub:
     # notified first, then the keepalive is closed, then the PP handles are awaited.
     class CleanupStub:
         _executor_loop_cleanup = pe.PyExecutor._executor_loop_cleanup
+        _executor_events = None
 
         def __init__(self, keepalive):
             self.events = []
diff --git a/tests/unittest/_torch/executor/test_py_executor.py b/tests/unittest/_torch/executor/test_py_executor.py
index 20c64eaad57d..613e6b7ddd4e 100644
--- a/tests/unittest/_torch/executor/test_py_executor.py
+++ b/tests/unittest/_torch/executor/test_py_executor.py
@@ -102,6 +102,7 @@ def test_forward_step_carries_context_logits_request_to_runner(request_flags, ex
         iter_counter=0,
         sampler=Mock(get_cache_indirection=Mock(return_value=cache_indirection)),
         execution_stream=Mock(),
+        _loop_on_execution_stream=False,
         _iter_adp_dummy_ctx_tokens=0,
         _iter_adp_dummy_gen_tokens=0,
         _compute_adp_dummy_tokens=PyExecutor._compute_adp_dummy_tokens,
@@ -2245,6 +2246,8 @@ class _CleanupStub:
     """Stub for _executor_loop_cleanup: records the order in which the
     shutdown notification and PP-handle wait happen."""
 
+    _executor_events = None
+
     def __init__(self, pp_handles_raise=False):
         self.response_lock = threading.Lock()
         self.response_cv = threading.Condition(self.response_lock)
diff --git a/tests/unittest/_torch/executor/test_runtime_hooks.py b/tests/unittest/_torch/executor/test_runtime_hooks.py
new file mode 100644
index 000000000000..3993cfa7b42a
--- /dev/null
+++ b/tests/unittest/_torch/executor/test_runtime_hooks.py
@@ -0,0 +1,924 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""How the engine reads the step-loop hooks a model declares (``runtime_hooks_of``).
+
+The engine's call sites run their original statements when a hook is off: every
+object a hook attaches to declares it off at class level, every call a call
+site makes on a hook is a method of the hook's protocol, every read of a hook in
+the engine files is decided by its all-off value toward the original statements
+(checked on the source), and call sites of each hook run on stand-ins issue
+their original work with the hook off and hand it to the hook when it is on.
+"""
+
+import ast
+import contextlib
+import dataclasses
+import inspect
+import types
+from pathlib import Path
+
+import pytest
+import torch
+
+import tensorrt_llm
+from tensorrt_llm._torch.pyexecutor.runtime_hooks import (
+    DEFAULT_RUNTIME_HOOKS,
+    ExecutorEvents,
+    RuntimeHooks,
+    SamplerStoreWriter,
+    StepInputStager,
+    runtime_hooks_of,
+)
+
+pytestmark = pytest.mark.cpu_only
+
+
+class _PlainModel(torch.nn.Module):
+    """A model that declares no hooks, like every model outside the opt-in."""
+
+
+class _ClassHooksModel(torch.nn.Module):
+    runtime_hooks = RuntimeHooks(loop_on_execution_stream=True)
+
+
+class _Events:
+    def after_warmup(self) -> None:
+        pass
+
+    def before_overlap_loop(self) -> None:
+        pass
+
+    def at_shutdown(self) -> None:
+        pass
+
+
+@pytest.mark.parametrize("model", [_PlainModel(), object()], ids=["module", "object"])
+def test_a_model_without_hooks_gets_the_all_off_default(model):
+    hooks = runtime_hooks_of(model)
+
+    assert hooks is DEFAULT_RUNTIME_HOOKS
+    assert hooks == RuntimeHooks()
+    assert not hooks.loop_on_execution_stream
+    assert hooks.step_input_stager is None
+    assert hooks.sampler_store_writer is None
+    assert hooks.decode_prologue is None
+    assert hooks.executor_events is None
+
+
+def test_every_field_defaults_to_off():
+    for field in dataclasses.fields(RuntimeHooks):
+        assert getattr(DEFAULT_RUNTIME_HOOKS, field.name) in (False, None), field.name
+
+
+def test_hooks_declared_as_none_mean_the_default():
+    model = _PlainModel()
+    model.runtime_hooks = None
+
+    assert runtime_hooks_of(model) is DEFAULT_RUNTIME_HOOKS
+
+
+def test_class_attribute_hooks_are_returned_as_declared():
+    assert runtime_hooks_of(_ClassHooksModel()) is _ClassHooksModel.runtime_hooks
+
+
+def test_instance_attribute_hooks_are_returned_as_declared():
+    hooks = RuntimeHooks(executor_events=_Events())
+    model = _PlainModel()
+    model.runtime_hooks = hooks
+
+    assert runtime_hooks_of(model) is hooks
+
+
+def test_instance_attribute_overrides_class_attribute():
+    hooks = RuntimeHooks()
+    model = _ClassHooksModel()
+    model.runtime_hooks = hooks
+
+    assert runtime_hooks_of(model) is hooks
+
+
+@pytest.mark.parametrize(
+    "declared",
+    [True, {"loop_on_execution_stream": True}, "RuntimeHooks()"],
+    ids=["bool", "dict", "str"],
+)
+def test_hooks_of_the_wrong_type_raise(declared):
+    model = _PlainModel()
+    model.runtime_hooks = declared
+
+    with pytest.raises(TypeError, match=r"_PlainModel\.runtime_hooks must be a RuntimeHooks"):
+        runtime_hooks_of(model)
+
+
+def test_hooks_are_immutable_once_declared():
+    hooks = RuntimeHooks()
+
+    with pytest.raises(dataclasses.FrozenInstanceError):
+        hooks.loop_on_execution_stream = True
+
+
+def test_engine_classes_declare_the_stager_off():
+    """Every object a hook attaches to declares the attribute off at class level."""
+    from tensorrt_llm._torch.attention.backends.sparse.dsa.metadata import (
+        DSAtrtllmAttentionMetadata,
+    )
+    from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata
+    from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import CUDAGraphRunner
+    from tensorrt_llm._torch.pyexecutor.engine.runners.decoder.inputs import DecoderInputPreparer
+    from tensorrt_llm._torch.pyexecutor.engine.runners.decoder.runner import DecoderRunner
+    from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2
+    from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor
+    from tensorrt_llm._torch.speculative.eagle3 import Eagle3OneModelSpecMetadata
+    from tensorrt_llm._torch.speculative.interface import SpecMetadata
+    from tensorrt_llm._torch.speculative.spec_sampler_base import SpecSampler
+
+    for cls in (
+        TrtllmAttentionMetadata,
+        DSAtrtllmAttentionMetadata,
+        SpecMetadata,
+        Eagle3OneModelSpecMetadata,
+    ):
+        assert cls.step_input_stager is None, cls
+        # A class attribute, not a dataclass field: the constructors do not take it.
+        assert "step_input_stager" not in {f.name for f in dataclasses.fields(cls)}, cls
+    assert DecoderInputPreparer.step_input_stager is None
+    assert DecoderInputPreparer.decode_prologue is None
+    assert DecoderRunner.runtime_hooks is DEFAULT_RUNTIME_HOOKS
+    assert KVCacheManagerV2.step_input_stager is None
+    assert CUDAGraphRunner.staged_static_input_tokens == 0
+    assert SpecSampler.store_writer is None
+    assert PyExecutor._execution_stream_loop_armed is False
+    assert PyExecutor._loop_on_execution_stream is False
+    assert PyExecutor._executor_events is None
+
+
+# Per hook protocol: the local names and attribute names a call site holds the hook in, and
+# the engine files with call sites.
+_HOOK_CONSUMERS = {
+    StepInputStager: (
+        ("stager", "stager_source", "step_input_stager"),
+        ("step_input_stager",),
+        (
+            "pyexecutor/engine/runners/decoder/runner.py",
+            "pyexecutor/engine/runners/decoder/inputs.py",
+            "pyexecutor/engine/runners/common.py",
+            "pyexecutor/kv_cache/kv_cache_manager_v2.py",
+            "attention/backends/trtllm.py",
+            "attention/backends/sparse/dsa/metadata.py",
+            "attention/backends/sparse/dsa/indexer.py",
+            "speculative/eagle3.py",
+            "speculative/interface.py",
+        ),
+    ),
+    SamplerStoreWriter: (
+        ("store_writer",),
+        ("store_writer",),
+        ("speculative/spec_sampler_base.py",),
+    ),
+    ExecutorEvents: ((), ("_executor_events",), ("pyexecutor/py_executor.py",)),
+}
+
+
+def _protocol_methods(protocol) -> set:
+    return {
+        name
+        for name, value in vars(protocol).items()
+        if inspect.isfunction(value) and not name.startswith("_")
+    }
+
+
+def _hook_calls(tree: ast.AST, local_names, attr_names):
+    """``(method, line)`` of every call on a hook held in one of the locals or attributes."""
+    for node in ast.walk(tree):
+        if not (isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute)):
+            continue
+        receiver = node.func.value
+        if (isinstance(receiver, ast.Name) and receiver.id in local_names) or (
+            isinstance(receiver, ast.Attribute) and receiver.attr in attr_names
+        ):
+            yield node.func.attr, node.lineno
+
+
+@pytest.mark.parametrize("protocol", list(_HOOK_CONSUMERS), ids=lambda p: p.__name__)
+def test_every_hook_call_site_uses_its_protocol(protocol):
+    """A call site calls only its hook's protocol methods, and the protocol has no unused method."""
+    local_names, attr_names, consumers = _HOOK_CONSUMERS[protocol]
+    root = Path(tensorrt_llm.__file__).resolve().parent / "_torch"
+    methods = _protocol_methods(protocol)
+    used, unknown = set(), []
+    for rel in consumers:
+        tree = ast.parse((root / rel).read_text())
+        calls = list(_hook_calls(tree, local_names, attr_names))
+        assert calls, f"{rel} has no {protocol.__name__} call site"
+        for method, line in calls:
+            used.add(method)
+            if method not in methods:
+                unknown.append(f"{rel}:{line}: {method}")
+    assert not unknown, f"calls outside the {protocol.__name__} protocol:\n" + "\n".join(unknown)
+    assert used == methods
+
+
+class _StoreWriterModel(torch.nn.Module):
+    def __init__(self, writer) -> None:
+        super().__init__()
+        self.runtime_hooks = RuntimeHooks(sampler_store_writer=writer)
+
+
+def _instantiate_sampler(monkeypatch, model, sampler):
+    """``instantiate_sampler`` for a speculative engine whose spec decoder is ``sampler``."""
+    from tensorrt_llm._torch.pyexecutor import _util
+
+    monkeypatch.setattr(_util, "confidential_compute_enabled", lambda: False)
+    monkeypatch.setattr(_util, "resolve_max_num_sequences", lambda *args, **kwargs: 8)
+    monkeypatch.setattr(_util, "create_torch_sampler_args", lambda **kwargs: "sampler args")
+    monkeypatch.setattr(_util, "get_spec_decoder", lambda args, spec_config: sampler)
+    engine = types.SimpleNamespace(
+        max_seq_len=64,
+        spec_config=types.SimpleNamespace(
+            spec_dec_mode=types.SimpleNamespace(has_spec_decoder=lambda: True)
+        ),
+        model=model,
+    )
+    llm_args = types.SimpleNamespace(
+        sampler_force_async_worker=False,
+        disable_overlap_scheduler=False,
+        enable_speculative_beam_history_d2h=False,
+    )
+    return _util.instantiate_sampler(
+        engine,
+        llm_args,
+        None,
+        max_batch_size=8,
+        max_beam_width=1,
+        mm_encoder_only=False,
+        speculative_config=None,
+    )
+
+
+def test_instantiate_sampler_binds_no_writer_without_hooks(monkeypatch):
+    from tensorrt_llm._torch.speculative.spec_sampler_base import SpecSampler
+
+    sampler = SpecSampler.__new__(SpecSampler)
+
+    assert _instantiate_sampler(monkeypatch, _PlainModel(), sampler) is sampler
+    assert "store_writer" not in vars(sampler)
+    assert sampler.store_writer is None
+
+
+def test_instantiate_sampler_binds_the_declared_writer(monkeypatch):
+    from tensorrt_llm._torch.speculative.spec_sampler_base import SpecSampler
+
+    writer = object()
+    sampler = SpecSampler.__new__(SpecSampler)
+
+    assert _instantiate_sampler(monkeypatch, _StoreWriterModel(writer), sampler) is sampler
+    assert sampler.store_writer is writer
+
+
+def test_executor_events_protocol():
+    assert _protocol_methods(ExecutorEvents) == {
+        "after_warmup",
+        "before_overlap_loop",
+        "at_shutdown",
+    }
+    events = _Events()
+    assert RuntimeHooks(executor_events=events).executor_events is events
+
+
+class _StreamRecorder:
+    """Stands in for ``torch.cuda.stream``: records the streams it is entered with."""
+
+    def __init__(self) -> None:
+        self.entered = []
+
+    @contextlib.contextmanager
+    def __call__(self, stream):
+        self.entered.append(stream)
+        yield
+
+
+def _scope_executor(armed: bool, **overrides):
+    """A ``PyExecutor`` with only the attributes ``_overlap_loop_stream_scope`` reads."""
+    from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor
+
+    executor = PyExecutor.__new__(PyExecutor)
+    executor._execution_stream_loop_armed = armed
+    executor.execution_stream = object()
+    attrs = dict(
+        drafter=None,
+        draft_model_engine=None,
+        guided_decoder=None,
+        kv_connector_manager=None,
+        is_encoder_decoder=False,
+        _mm_encoder_item_scheduling_enabled=False,
+    )
+    attrs.update(overrides)
+    for name, value in attrs.items():
+        setattr(executor, name, value)
+    return executor
+
+
+def test_overlap_loop_scope_is_a_no_op_when_not_armed(monkeypatch):
+    executor = _scope_executor(armed=False)
+    recorder = _StreamRecorder()
+    monkeypatch.setattr(torch.cuda, "stream", recorder)
+
+    with executor._overlap_loop_stream_scope():
+        assert executor._loop_on_execution_stream is False
+    assert recorder.entered == []
+    assert "_loop_on_execution_stream" not in vars(executor)
+
+
+def test_overlap_loop_scope_engages_when_armed(monkeypatch):
+    executor = _scope_executor(armed=True)
+    recorder = _StreamRecorder()
+    monkeypatch.setattr(torch.cuda, "stream", recorder)
+
+    with executor._overlap_loop_stream_scope():
+        assert executor._loop_on_execution_stream is True
+    assert recorder.entered == [executor.execution_stream]
+    assert executor._loop_on_execution_stream is False
+
+
+def test_overlap_loop_scope_resets_the_flag_when_the_loop_raises(monkeypatch):
+    executor = _scope_executor(armed=True)
+    monkeypatch.setattr(torch.cuda, "stream", _StreamRecorder())
+
+    with pytest.raises(RuntimeError, match="loop failed"):
+        with executor._overlap_loop_stream_scope():
+            raise RuntimeError("loop failed")
+    assert executor._loop_on_execution_stream is False
+
+
+@pytest.mark.parametrize(
+    "name, value",
+    [
+        ("drafter", object()),
+        ("draft_model_engine", object()),
+        ("guided_decoder", object()),
+        ("kv_connector_manager", object()),
+        ("is_encoder_decoder", True),
+        ("_mm_encoder_item_scheduling_enabled", True),
+    ],
+)
+def test_overlap_loop_scope_stays_off_with_side_streams(monkeypatch, name, value):
+    executor = _scope_executor(armed=True, **{name: value})
+    recorder = _StreamRecorder()
+    monkeypatch.setattr(torch.cuda, "stream", recorder)
+
+    with executor._overlap_loop_stream_scope():
+        assert executor._loop_on_execution_stream is False
+    assert recorder.entered == []
+
+
+def test_instantiate_sampler_leaves_other_spec_decoders_alone(monkeypatch):
+    sampler = types.SimpleNamespace()
+
+    assert _instantiate_sampler(monkeypatch, _StoreWriterModel(object()), sampler) is sampler
+    assert "store_writer" not in vars(sampler)
+
+
+# ---------------------------------------------------------------------------------------------
+# With every hook off, every call site runs its original statements.
+#
+# The check below reads the engine files that consume a hook and evaluates each read of a hook
+# at its all-off value. An ``if`` that reads a hook must be decided by that value; the branch it
+# takes is checked like any other code and the other branch is not run, so it may use the hook
+# freely. Outside such tests, live code may only bind a hook value to a local (or to an
+# attribute that declares the hook) at its all-off value. Any other read is reported with its
+# file and line.
+# ---------------------------------------------------------------------------------------------
+
+_SEAM_CONSUMERS = (
+    "pyexecutor/engine/runners/decoder/runner.py",
+    "pyexecutor/engine/runners/decoder/inputs.py",
+    "pyexecutor/cuda_graph_runner.py",
+    "pyexecutor/kv_cache/kv_cache_manager_v2.py",
+    "pyexecutor/py_executor.py",
+    "pyexecutor/_util.py",
+    "attention/backends/trtllm.py",
+    "attention/backends/sparse/dsa/metadata.py",
+    "attention/backends/sparse/dsa/indexer.py",
+    "speculative/eagle3.py",
+    "speculative/interface.py",
+    "speculative/spec_sampler_base.py",
+)
+
+# Attributes that hold a hook, and their value when every hook is off.
+_SEAM_ATTRS = {
+    "step_input_stager": None,
+    "store_writer": None,
+    "_executor_events": None,
+    "_loop_on_execution_stream": False,
+    "_execution_stream_loop_armed": False,
+    "staged_static_input_tokens": 0,
+    "runtime_hooks": DEFAULT_RUNTIME_HOOKS,
+    "decode_prologue": None,
+}
+# Locals the call sites hold a hook (or a decision on one) in.
+_SEAM_LOCALS = ("stager", "stager_source", "store_writer", "prologue", "staged", "staged_store")
+_HOOK_FIELDS = {field.name for field in dataclasses.fields(RuntimeHooks)}
+
+
+class _Unknown:
+    """A value the all-off evaluation cannot decide."""
+
+
+_UNKNOWN = _Unknown()
+
+
+class _AllOffWalk:
+    """Walks one engine file with every hook at its all-off value."""
+
+    def __init__(self, rel: str) -> None:
+        self.rel = rel
+        self.problems = []
+        self.decided = 0
+
+    def report(self, node: ast.AST, what: str) -> None:
+        self.problems.append(f"{self.rel}:{node.lineno}: {what}")
+
+    # -- reads of a hook ------------------------------------------------------------------
+    @staticmethod
+    def _is_seam_call(node: ast.AST) -> bool:
+        if not (isinstance(node, ast.Call) and isinstance(node.func, ast.Name)):
+            return False
+        if node.func.id == "runtime_hooks_of":
+            return True
+        return (
+            node.func.id == "getattr"
+            and len(node.args) >= 2
+            and isinstance(node.args[1], ast.Constant)
+            and node.args[1].value == "step_input_stager"
+        )
+
+    def reads_seam(self, node: ast.AST, env: dict) -> bool:
+        for sub in ast.walk(node):
+            if isinstance(sub, ast.Attribute) and sub.attr in _SEAM_ATTRS:
+                return True
+            if isinstance(sub, ast.Name) and sub.id in env:
+                return True
+            if self._is_seam_call(sub):
+                return True
+        return False
+
+    def value(self, node: ast.AST, env: dict):
+        """The all-off value of ``node``, or ``_UNKNOWN``."""
+        if isinstance(node, ast.Constant):
+            return node.value
+        if isinstance(node, ast.Name):
+            return env.get(node.id, _UNKNOWN)
+        if isinstance(node, ast.Attribute):
+            receiver = self.value(node.value, env)
+            if receiver is DEFAULT_RUNTIME_HOOKS and node.attr in _HOOK_FIELDS:
+                return getattr(DEFAULT_RUNTIME_HOOKS, node.attr)
+            if node.attr in _SEAM_ATTRS:
+                return _SEAM_ATTRS[node.attr]
+            return _UNKNOWN
+        if self._is_seam_call(node):
+            return DEFAULT_RUNTIME_HOOKS if node.func.id == "runtime_hooks_of" else None
+        if isinstance(node, ast.UnaryOp) and isinstance(node.op, ast.Not):
+            operand = self.value(node.operand, env)
+            return _UNKNOWN if operand is _UNKNOWN else not operand
+        if isinstance(node, ast.BoolOp):
+            values = [self.value(v, env) for v in node.values]
+            known = [v for v in values if v is not _UNKNOWN]
+            if isinstance(node.op, ast.And):
+                if any(not v for v in known):
+                    return False
+                return True if len(known) == len(values) else _UNKNOWN
+            if any(known):
+                return True
+            return False if len(known) == len(values) else _UNKNOWN
+        if isinstance(node, ast.Compare) and len(node.ops) == 1:
+            left, right = self.value(node.left, env), self.value(node.comparators[0], env)
+            op = node.ops[0]
+            if isinstance(op, (ast.Is, ast.IsNot)) and left is not _UNKNOWN:
+                if isinstance(node.comparators[0], ast.Constant) and right is None:
+                    return (left is None) == isinstance(op, ast.Is)
+            if (
+                isinstance(op, ast.Eq)
+                and isinstance(node.left, ast.Attribute)
+                and node.left.attr == "staged_static_input_tokens"
+            ):
+                # No tokens were staged; a graph replay has at least one token.
+                return False
+        return _UNKNOWN
+
+    # -- statements -------------------------------------------------------------------------
+    def block(self, body, env: dict) -> bool:
+        """Walk live statements; True when the block always ends in return or raise."""
+        for stmt in body:
+            if self.stmt(stmt, env):
+                return True
+        return False
+
+    def bind(self, stmt, targets, value, env: dict) -> None:
+        reads = self.reads_seam(value, env)
+        names = [t.id for t in targets if isinstance(t, ast.Name)]
+        attrs = [t for t in targets if isinstance(t, ast.Attribute) and t.attr in _SEAM_ATTRS]
+        if not (reads or attrs or any(name in env or name in _SEAM_LOCALS for name in names)):
+            return
+        if len(names) + len(attrs) != len(targets):
+            self.report(stmt, "a hook value is bound to something other than a local")
+            return
+        off = self.value(value, env)
+        if off is _UNKNOWN:
+            if reads or attrs:
+                self.report(stmt, "a hook value with every hook off is not decided")
+            for name in names:
+                env.pop(name, None)
+            return
+        for name in names:
+            env[name] = off
+        for attr in attrs:
+            if off != _SEAM_ATTRS[attr.attr]:
+                self.report(stmt, f"{attr.attr} is set to a value other than its all-off value")
+
+    def stmt(self, stmt: ast.stmt, env: dict) -> bool:
+        if isinstance(stmt, (ast.FunctionDef, ast.AsyncFunctionDef)):
+            self.block(stmt.body, dict(env))
+            return False
+        if isinstance(stmt, ast.ClassDef):
+            self.block(stmt.body, {})
+            return False
+        if isinstance(stmt, (ast.If, ast.While)):
+            if not self.reads_seam(stmt.test, env):
+                body_ends = self.block(stmt.body, env)
+                else_ends = self.block(stmt.orelse, env)
+                return isinstance(stmt, ast.If) and body_ends and else_ends
+            taken = self.value(stmt.test, env)
+            if taken is _UNKNOWN:
+                self.report(stmt, "a test of a hook is not decided with every hook off")
+                return False
+            self.decided += 1
+            if isinstance(stmt, ast.While) and taken:
+                self.report(stmt, "a loop on a hook runs with every hook off")
+                return False
+            return self.block(stmt.body if taken else stmt.orelse, env)
+        if isinstance(stmt, ast.Assign):
+            self.bind(stmt, stmt.targets, stmt.value, env)
+            return False
+        if isinstance(stmt, ast.AnnAssign):
+            if stmt.value is not None:
+                self.bind(stmt, [stmt.target], stmt.value, env)
+            return False
+        if isinstance(stmt, (ast.With, ast.AsyncWith)):
+            for item in stmt.items:
+                if self.reads_seam(item.context_expr, env):
+                    self.report(stmt, "a hook is entered as a context")
+            return self.block(stmt.body, env)
+        if isinstance(stmt, ast.Try):
+            ends = self.block(stmt.body, env)
+            for handler in stmt.handlers:
+                self.block(handler.body, env)
+            self.block(stmt.orelse, env)
+            return self.block(stmt.finalbody, env) or (ends and not stmt.handlers)
+        if isinstance(stmt, (ast.For, ast.AsyncFor)):
+            if self.reads_seam(stmt.iter, env):
+                self.report(stmt, "a loop iterates over a hook")
+            self.block(stmt.body, env)
+            self.block(stmt.orelse, env)
+            return False
+        if self.reads_seam(stmt, env):
+            self.report(stmt, "a hook is used outside a test decided with every hook off")
+        return isinstance(stmt, (ast.Return, ast.Raise))
+
+
+@pytest.mark.parametrize("rel", _SEAM_CONSUMERS)
+def test_every_seam_read_has_a_native_branch(rel):
+    """With every hook off, each call site takes the branch with its original statements."""
+    root = Path(tensorrt_llm.__file__).resolve().parent / "_torch"
+    walk = _AllOffWalk(rel)
+    walk.block(ast.parse((root / rel).read_text()).body, {})
+    assert not walk.problems, "\n".join(walk.problems)
+    assert walk.decided, f"{rel} has no call site of a hook"
+
+
+def test_the_seam_check_reports_undecided_reads():
+    source = """
+def step(self, x):
+    stager = self.step_input_stager
+    if stager is not None:
+        stager.copy(x)
+    else:
+        x.copy_()
+    stager.flush()
+    if x or stager:
+        pass
+    self.step_input_stager = x
+"""
+    walk = _AllOffWalk("sample.py")
+    walk.block(ast.parse(source).body, {})
+    assert walk.decided == 1
+    assert walk.problems == [
+        "sample.py:8: a hook is used outside a test decided with every hook off",
+        "sample.py:9: a test of a hook is not decided with every hook off",
+        "sample.py:11: a hook value with every hook off is not decided",
+    ]
+
+
+def test_the_decode_prologue_call_site_uses_its_protocol():
+    """The input preparer calls the prologue with the protocol's positional and keyword parameters."""
+    from tensorrt_llm._torch.pyexecutor.runtime_hooks import DecodePrologue
+
+    signature = inspect.signature(DecodePrologue.__call__)
+    positional = [
+        p.name
+        for p in signature.parameters.values()
+        if p.kind is inspect.Parameter.POSITIONAL_OR_KEYWORD and p.name != "self"
+    ]
+    keywords = {
+        p.name for p in signature.parameters.values() if p.kind is inspect.Parameter.KEYWORD_ONLY
+    }
+    root = Path(tensorrt_llm.__file__).resolve().parent / "_torch"
+    tree = ast.parse((root / "pyexecutor/engine/runners/decoder/inputs.py").read_text())
+    calls = [
+        node
+        for node in ast.walk(tree)
+        if isinstance(node, ast.Call)
+        and isinstance(node.func, ast.Name)
+        and node.func.id == "prologue"
+    ]
+    assert len(calls) == 1
+    assert len(calls[0].args) == len(positional)
+    assert {keyword.arg for keyword in calls[0].keywords} == keywords
+
+
+# -- the decode prologue ------------------------------------------------------------------------
+
+
+class _AttnMetadata:
+    def __init__(self) -> None:
+        self.kv_lens_updates = 0
+
+    def on_update_kv_lens(self) -> None:
+        self.kv_lens_updates += 1
+
+
+def _prologue_preparer(prologue):
+    """A ``DecoderInputPreparer`` with only what ``preprocess_inputs`` reads without spec decoding."""
+    from tensorrt_llm._torch.pyexecutor.engine.runners.decoder.inputs import DecoderInputPreparer
+
+    preparer = DecoderInputPreparer.__new__(DecoderInputPreparer)
+    if prologue is not None:
+        preparer.decode_prologue = prologue
+    preparer.guided_decoder = None
+    preparer._config = types.SimpleNamespace(disable_overlap_scheduler=True)
+    preparer.mapping = types.SimpleNamespace(has_cp_helix=lambda: False)
+    return preparer
+
+
+def test_preprocess_inputs_runs_the_native_prologue_without_a_hook():
+    preparer = _prologue_preparer(None)
+    inputs = {"attn_metadata": _AttnMetadata()}
+
+    assert preparer.decode_prologue is None
+    assert preparer.preprocess_inputs(inputs, enable_spec_decode=False, runtime_draft_len=0) is (
+        inputs
+    )
+    assert inputs["attn_metadata"].kv_lens_updates == 1
+
+
+@pytest.mark.parametrize("engaged", [False, True], ids=["declined", "engaged"])
+def test_preprocess_inputs_keeps_the_native_prologue_when_the_hook_declines(engaged):
+    calls = []
+
+    def prologue(preparer, inputs, *, enable_spec_decode, runtime_draft_len):
+        calls.append((preparer, inputs, enable_spec_decode, runtime_draft_len))
+        return engaged
+
+    preparer = _prologue_preparer(prologue)
+    inputs = {"attn_metadata": _AttnMetadata()}
+
+    assert preparer.preprocess_inputs(inputs, enable_spec_decode=False, runtime_draft_len=3) is (
+        inputs
+    )
+    assert calls == [(preparer, inputs, False, 3)]
+    assert inputs["attn_metadata"].kv_lens_updates == (0 if engaged else 1)
+
+
+# -- the execution-stream loop --------------------------------------------------------------------
+
+
+class _Stream:
+    def __init__(self, name: str, log: list) -> None:
+        self.name, self.log = name, log
+
+    def wait_stream(self, other: "_Stream") -> None:
+        self.log.append(f"{self.name} waits for {other.name}")
+
+
+def _forward_step_executor(monkeypatch, log: list):
+    """A ``PyExecutor`` with only what ``_forward_step`` reads, logging its stream work."""
+    from tensorrt_llm._torch.pyexecutor import py_executor
+    from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor
+
+    current = _Stream("current", log)
+
+    @contextlib.contextmanager
+    def stream(s):
+        log.append(f"enter {s.name}")
+        yield
+        log.append(f"leave {s.name}")
+
+    monkeypatch.setattr(torch.cuda, "current_stream", lambda: current)
+    monkeypatch.setattr(torch.cuda, "stream", stream)
+    monkeypatch.setattr(py_executor, "nvtx_range", lambda message: (lambda fn: fn))
+
+    executor = PyExecutor.__new__(PyExecutor)
+    executor.iter_counter = 0
+    executor.execution_stream = _Stream("execution", log)
+    executor.resource_manager = object()
+    executor.sampler = types.SimpleNamespace(get_cache_indirection=lambda: None)
+    executor.model_engine = types.SimpleNamespace(
+        forward=lambda *args, **kwargs: log.append("forward") or "outputs"
+    )
+    executor._iter_adp_dummy_ctx_tokens = 0
+    executor._iter_adp_dummy_gen_tokens = 0
+    executor._compute_adp_dummy_tokens = lambda requests: (0, 0)
+    executor._maybe_record_hang_diagnostic_phase = lambda *args, **kwargs: None
+    executor._attach_encoder_output_to_execution_stream = lambda requests: None
+    executor._mark_cross_kv_projection_consumed = lambda requests: None
+    executor._kv_connector_wait_for_save = lambda: None
+    return executor
+
+
+_NO_REQUESTS = types.SimpleNamespace(
+    context_requests=[], num_context_requests=0, num_generation_requests=0
+)
+
+
+def test_forward_step_orders_the_streams_off_the_execution_stream_loop(monkeypatch):
+    log = []
+    executor = _forward_step_executor(monkeypatch, log)
+
+    assert executor._forward_step(_NO_REQUESTS) == "outputs"
+    assert log == [
+        "execution waits for current",
+        "enter execution",
+        "forward",
+        "leave execution",
+        "current waits for execution",
+    ]
+
+
+def test_forward_step_skips_the_waits_on_the_execution_stream_loop(monkeypatch):
+    log = []
+    executor = _forward_step_executor(monkeypatch, log)
+    executor._loop_on_execution_stream = True
+
+    assert executor._forward_step(_NO_REQUESTS) == "outputs"
+    assert log == ["enter execution", "forward", "leave execution"]
+
+
+# -- executor events --------------------------------------------------------------------------
+
+
+def _cleanup_executor(events, completed: bool):
+    """A ``PyExecutor`` with only what ``_executor_loop_cleanup`` reads, without pipeline stages."""
+    import threading
+
+    from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor
+
+    executor = PyExecutor.__new__(PyExecutor)
+    executor.response_cv = threading.Condition()
+    executor.shutdown_event = threading.Event()
+    executor.num_micro_batches = 0
+    executor._event_loop_completed = completed
+    if events is not None:
+        executor._executor_events = events
+    return executor
+
+
+class _RecordingEvents(_Events):
+    def __init__(self) -> None:
+        self.calls = []
+
+    def at_shutdown(self) -> None:
+        self.calls.append("at_shutdown")
+
+
+@pytest.mark.parametrize("completed", [False, True])
+def test_executor_loop_cleanup_without_events(completed):
+    executor = _cleanup_executor(None, completed)
+
+    executor._executor_loop_cleanup()
+    assert executor.is_shutdown and executor.shutdown_event.is_set()
+
+
+@pytest.mark.parametrize("completed", [False, True])
+def test_executor_loop_cleanup_calls_at_shutdown_after_a_completed_loop(completed):
+    events = _RecordingEvents()
+    executor = _cleanup_executor(events, completed)
+
+    executor._executor_loop_cleanup()
+    assert executor.is_shutdown
+    assert events.calls == (["at_shutdown"] if completed else [])
+
+
+# -- the step-input stager ------------------------------------------------------------------
+
+
+class _RecordingStager:
+    def __init__(self) -> None:
+        self.calls = []
+
+    def __getattr__(self, name):
+        return lambda *args, **kwargs: self.calls.append((name, args, kwargs))
+
+
+def _block_offsets_manager(monkeypatch, log: list):
+    """A ``KVCacheManagerV2`` with only what ``copy_batch_block_offsets`` reads (one page table)."""
+    from tensorrt_llm._torch.pyexecutor.kv_cache import kv_cache_manager_v2
+    from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2
+
+    monkeypatch.setattr(
+        kv_cache_manager_v2,
+        "copy_batch_block_offsets_to_device",
+        lambda *args: log.append(("native", args)),
+    )
+    manager = KVCacheManagerV2.__new__(KVCacheManagerV2)
+    manager._publish_sparse_metadata = lambda: None
+    manager.max_copy_beam_width = 1
+    manager.kv_cache_type = None
+    manager._use_per_layer_page_tables = False
+    manager.index_mapper = types.SimpleNamespace(
+        get_copy_index=lambda request_ids, num_contexts, beam_width, cross: torch.tensor([1, 0])
+    )
+    manager.host_kv_cache_block_offsets = torch.zeros(2, 4, dtype=torch.int32)
+    manager.index_scales = torch.ones(1, dtype=torch.int32)
+    manager._index_scale_ints = [1]
+    manager.kv_offset = torch.zeros(1, dtype=torch.int32)
+    manager._stream = types.SimpleNamespace(cuda_stream=7)
+    return manager
+
+
+def test_kv_cache_block_offsets_copy_natively_without_a_stager(monkeypatch):
+    log = []
+    manager = _block_offsets_manager(monkeypatch, log)
+    dst = torch.empty(2, 4, dtype=torch.int32)
+
+    manager.copy_batch_block_offsets(dst, [3, 4], 1, 0, 2)
+    assert len(log) == 1
+    kind, args = log[0]
+    assert kind == "native"
+    assert args[0] is manager.host_kv_cache_block_offsets and args[1] is dst
+    assert args[2].tolist() == [1, 0]
+    assert args[3] is manager.index_scales and args[4] is manager.kv_offset and args[5] == 7
+
+
+def test_kv_cache_block_offsets_are_staged_with_a_stager(monkeypatch):
+    log = []
+    manager = _block_offsets_manager(monkeypatch, log)
+    manager.step_input_stager = _RecordingStager()
+    dst = torch.empty(2, 4, dtype=torch.int32)
+
+    manager.copy_batch_block_offsets(dst, [3, 4], 1, 0, 2)
+    assert log == []
+    ((name, args, _),) = manager.step_input_stager.calls
+    assert name == "block_offsets"
+    assert args[0] is manager.host_kv_cache_block_offsets and args[1] is dst
+    assert args[2:] == ([1, 0], [1], [0])
+
+
+def _slot_mapping_inputs(tokens: int):
+    metadata = types.SimpleNamespace(
+        host_indexer_k_cache_block_offsets=torch.tensor([[0, 1]], dtype=torch.int32),
+        host_slot_mapping_fp8=torch.zeros(tokens, dtype=torch.int64),
+        host_slot_mapping_scale=torch.zeros(tokens, dtype=torch.int64),
+        slot_mapping_fp8=torch.full((tokens,), -1, dtype=torch.int64),
+        slot_mapping_scale=torch.full((tokens,), -1, dtype=torch.int64),
+    )
+    params = types.SimpleNamespace(
+        batch_size=1,
+        tokens_per_block=4,
+        head_dim=8,
+        new_kv_tokens=torch.tensor([tokens]),
+        data_bytes_per_token=8,
+        cached_kv_tokens=torch.tensor([1]),
+        quant_block_size=8,
+    )
+    return metadata, params
+
+
+def test_indexer_slot_mappings_copy_natively_without_a_stager():
+    from tensorrt_llm._torch.attention.backends.sparse.dsa.indexer import Indexer
+
+    metadata, params = _slot_mapping_inputs(3)
+
+    Indexer.recompute_slot_mappings(metadata, params)
+    assert torch.equal(metadata.slot_mapping_fp8, metadata.host_slot_mapping_fp8)
+    assert torch.equal(metadata.slot_mapping_scale, metadata.host_slot_mapping_scale)
+
+
+def test_indexer_slot_mappings_are_staged_with_a_stager():
+    from tensorrt_llm._torch.attention.backends.sparse.dsa.indexer import Indexer
+
+    metadata, params = _slot_mapping_inputs(3)
+    metadata.step_input_stager = _RecordingStager()
+
+    Indexer.recompute_slot_mappings(metadata, params)
+    assert [name for name, _, _ in metadata.step_input_stager.calls] == ["copy", "copy"]
+    assert bool((metadata.slot_mapping_fp8 == -1).all())
+    assert bool((metadata.slot_mapping_scale == -1).all())

From 37013e6a17fa932c75ba1c2bfd627031a6b4530b Mon Sep 17 00:00:00 2001
From: Xiao Wang <24860335+xwang233@users.noreply.github.com>
Date: Thu, 22 Oct 2026 09:34:08 -0700
Subject: [PATCH 2/5] [TRTLLM-17096][fix] runtime hooks: run
 ExecutorEvents.after_warmup inside the distributed warmup guard

after_warmup may run collectives, so a rank-local failure there must arm
the rank-crash watchdog like a warmup failure does.

Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com>
---
 tensorrt_llm/_torch/pyexecutor/py_executor.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py
index f06c73f28213..2bcac7d79198 100644
--- a/tensorrt_llm/_torch/pyexecutor/py_executor.py
+++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py
@@ -936,8 +936,10 @@ def __init__(
         torch.cuda.current_stream().wait_stream(self.execution_stream)
         if self._executor_events is not None:
             # The model's actions after warmup and graph capture, on every rank
-            # (some are collective, so every rank reaches this point).
-            self._executor_events.after_warmup()
+            # (some are collective, so every rank reaches this point, and a
+            # rank-local failure arms the rank-crash watchdog).
+            with _distributed_warmup_guard(self.dist, self.dist.mapping):
+                self._executor_events.after_warmup()
         self.encoder_launch_executor = (ThreadPoolExecutor(
             max_workers=1, thread_name_prefix="encoder-launch")
                                         if self.is_encoder_decoder else None)

From 247c83ec9c05cfc6bc4bf6b72f54b421f3412111 Mon Sep 17 00:00:00 2001
From: Xiao Wang <24860335+xwang233@users.noreply.github.com>
Date: Thu, 22 Oct 2026 11:51:34 -0700
Subject: [PATCH 3/5] [TRTLLM-17096][fix] runtime hooks: stage step inputs only
 on the replay of a captured graph

The step-input window closes right before the replay of an already captured
CUDA graph. An eager step, an initial capture or a warmup would run the model
before the staged writes are issued, so the decoder runner now offers the
window only on such a replay, and StepInputStager documents the restriction.

Add behavioral tests with reference hooks: the runner window lifecycle
(staged writes issued before the completion event and the replay, discard on
a failure while the window is open), steps that are not a replay never open a
window, before_replay runs only before the replay of a captured graph, and
the speculative sampler's store writer against the native stores on
nonconsecutive slots, with its native fallbacks.

Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com>
---
 .../engine/runners/decoder/runner.py          |  31 +-
 .../_torch/pyexecutor/runtime_hooks.py        |   8 +-
 .../_torch/executor/test_runtime_hooks.py     | 474 ++++++++++++++++++
 3 files changed, 498 insertions(+), 15 deletions(-)

diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py
index 060f30b6b6a5..e31f20da8e40 100644
--- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py
+++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py
@@ -640,25 +640,28 @@ def _forward_decoder(
             stager = None
             stager_source = self.runtime_hooks.step_input_stager
             if stager_source is not None:
+                # Only the replay of an already captured graph closes the window
+                # (before_replay); an eager step, an initial capture or a warmup
+                # would run the model before the staged writes are issued, so
+                # those steps always prepare natively.
                 graph_replay = (
                     can_run_graph
                     and not self.cuda_graph_runner.needs_capture(key)
                     and not self.cuda_graph_runner.is_warmup_only
                 )
-                stager = stager_source.begin_if_eligible(
-                    self._inputs,
-                    execution_requests,
-                    attn_metadata,
-                    spec_metadata,
-                    new_tensors_device,
-                    graph_replay=graph_replay,
-                    static_inputs=(
-                        self.cuda_graph_runner.shared_static_tensors if graph_replay else None
-                    ),
-                    is_dummy=is_dummy,
-                    enable_spec_decode=enable_spec_decode,
-                    promoted=bool(execution_promoted_context_ids),
-                )
+                if graph_replay:
+                    stager = stager_source.begin_if_eligible(
+                        self._inputs,
+                        execution_requests,
+                        attn_metadata,
+                        spec_metadata,
+                        new_tensors_device,
+                        graph_replay=graph_replay,
+                        static_inputs=self.cuda_graph_runner.shared_static_tensors,
+                        is_dummy=is_dummy,
+                        enable_spec_decode=enable_spec_decode,
+                        promoted=bool(execution_promoted_context_ids),
+                    )
             try:
                 # The window spans the scope: the scope is a no-op whenever
                 # speculative decoding is on, which every stager requires.
diff --git a/tensorrt_llm/_torch/pyexecutor/runtime_hooks.py b/tensorrt_llm/_torch/pyexecutor/runtime_hooks.py
index 24ffbccfbb8b..72987309443f 100644
--- a/tensorrt_llm/_torch/pyexecutor/runtime_hooks.py
+++ b/tensorrt_llm/_torch/pyexecutor/runtime_hooks.py
@@ -55,9 +55,15 @@ def begin_if_eligible(
     ) -> Optional["StepInputStager"]:
         """Open the window for this step, or return None to keep the native preparation.
 
+        The engine calls this only for the replay of an already captured CUDA
+        graph, so ``graph_replay`` is always True: the window closes right
+        before the replay, and an eager step, an initial capture or a warmup
+        would run the model before the staged writes are issued. Those steps
+        always prepare natively.
+
         ``preparer`` is the ``DecoderInputPreparer`` whose writes the window
         records. ``static_inputs`` are the static input tensors of the CUDA
-        graph the step replays (None when it replays none).
+        graph the step replays.
         """
 
     def copy(
diff --git a/tests/unittest/_torch/executor/test_runtime_hooks.py b/tests/unittest/_torch/executor/test_runtime_hooks.py
index 3993cfa7b42a..95745dd901db 100644
--- a/tests/unittest/_torch/executor/test_runtime_hooks.py
+++ b/tests/unittest/_torch/executor/test_runtime_hooks.py
@@ -922,3 +922,477 @@ def test_indexer_slot_mappings_are_staged_with_a_stager():
     assert [name for name, _, _ in metadata.step_input_stager.calls] == ["copy", "copy"]
     assert bool((metadata.slot_mapping_fp8 == -1).all())
     assert bool((metadata.slot_mapping_scale == -1).all())
+
+
+# -- the step-input window in the decoder runner -------------------------------------------------
+
+
+class _ReferenceStager:
+    """A step-input stager that accepts every step it is offered and logs the window's lifecycle."""
+
+    def __init__(self, log: list, accept: bool = True, tokens: int = 3) -> None:
+        self.log, self.accept, self.tokens = log, accept, tokens
+
+    def begin_if_eligible(
+        self,
+        preparer,
+        requests,
+        attn_metadata,
+        spec_metadata,
+        new_tensors_device,
+        *,
+        graph_replay,
+        static_inputs,
+        is_dummy,
+        enable_spec_decode,
+        promoted,
+    ):
+        self.log.append(("begin", graph_replay, static_inputs))
+        return self if self.accept else None
+
+    def end(self, record_event: bool = True) -> None:
+        self.log.append(("end", record_event))
+
+    def discard(self) -> None:
+        self.log.append("discard")
+
+    def set_completion_event(self, event) -> None:
+        self.log.append(("completion_event", event))
+
+    def static_inputs_tokens(self) -> int:
+        return self.tokens
+
+
+class _WindowGraphs:
+    """The ``DecoderCudaGraphs`` calls ``_forward_decoder`` makes, with one selected graph key.
+
+    ``run`` follows ``DecoderCudaGraphs._run``: ``before_replay`` runs only right before the
+    replay of an already captured graph, and the replay consumes the staged static inputs.
+    """
+
+    def __init__(self, log: list, key, *, needs_capture=False, warmup_only=False, fail=None):
+        self.log, self.key, self.fail = log, key, fail
+        self.cuda_graph_runner = types.SimpleNamespace(
+            needs_capture=lambda key: needs_capture,
+            is_warmup_only=warmup_only,
+            shared_static_tensors={"input_ids": torch.zeros(4, dtype=torch.int64)},
+            staged_static_input_tokens=0,
+        )
+
+    def attention_dp_execution_batch(self, requests):
+        return requests
+
+    def promote_final_contexts(self, requests, kv_cache_manager, **kwargs):
+        return requests, frozenset()
+
+    def pad_batch(self, requests, resource_manager, runtime_draft_len):
+        return contextlib.nullcontext(requests)
+
+    def select(self, requests, resource_manager, **kwargs):
+        return types.SimpleNamespace(
+            key=self.key,
+            use_lora_graph=False,
+            attn_metadata=kwargs["attn_metadata"],
+            spec_metadata=None,
+        )
+
+    def captures_model_body(self, key) -> bool:
+        return False
+
+    def run(self, key, inputs, run_model, postprocess_inputs, *, before_replay=None, **kwargs):
+        runner = self.cuda_graph_runner
+        if self.fail == "before_replay":
+            raise RuntimeError("graph run failed")
+        if key is None:
+            self.log.append("eager")
+        elif runner.needs_capture(key) or runner.is_warmup_only:
+            self.log.append("capture")
+        else:
+            if before_replay is not None:
+                before_replay()
+            self.log.append(("replay", runner.staged_static_input_tokens))
+            runner.staged_static_input_tokens = 0
+            if self.fail == "replay":
+                raise RuntimeError("replay failed")
+        return "outputs"
+
+
+class _WindowPreparer:
+    """The ``DecoderInputPreparer`` calls ``_forward_decoder`` makes."""
+
+    def __init__(self, log: list, fail: bool = False) -> None:
+        self.log, self.fail = log, fail
+        self.input_copy_event = "input-copy-event"
+
+    def prepare_inputs(self, *args, **kwargs):
+        self.log.append("prepare")
+        if self.fail:
+            raise RuntimeError("preparation failed")
+        return types.SimpleNamespace(kwargs={}, runtime_draft_len=0, gather_ids=None)
+
+    def record_input_copy(self) -> None:
+        self.log.append("record_input_copy")
+
+    def postprocess_inputs(self, *args, **kwargs) -> None:
+        pass
+
+
+def _window_runner(monkeypatch, graphs, preparer, stager):
+    """A ``DecoderRunner`` with only what ``_forward_decoder`` reads (no speculation, no CED)."""
+    from tensorrt_llm._torch.pyexecutor import ced_replay
+    from tensorrt_llm._torch.pyexecutor.engine.runners.decoder import runner as runner_module
+    from tensorrt_llm._torch.pyexecutor.engine.runners.decoder.runner import DecoderRunner
+
+    monkeypatch.setattr(runner_module, "get_draft_kv_cache_manager", lambda *args: None)
+    monkeypatch.setattr(ced_replay, "uses_private_decoder_cache", lambda *args: False)
+    attn_metadata = types.SimpleNamespace()
+    runner = DecoderRunner.__new__(DecoderRunner)
+    runner._config = types.SimpleNamespace(spec_config=None, use_mrope=False, is_spec_decode=False)
+    runner._metadata = types.SimpleNamespace(
+        setup_attention=lambda *args: attn_metadata,
+        attn_metadata=attn_metadata,
+        spec_metadata=None,
+    )
+    runner._graphs = graphs
+    runner._inputs = preparer
+    runner._stage_in_graph_sampling = None
+    runner._cuda_graph_metadata_prepare_scope = lambda *args, **kwargs: contextlib.nullcontext()
+    runner._execute_logit_post_processors = lambda *args: None
+    runner.forward_pass_callable = None
+    runner.moe_load_balancer = None
+    runner.model = types.SimpleNamespace()
+    runner.kv_cache_manager_key = "kv_cache_manager"
+    runner.runtime_hooks = RuntimeHooks(step_input_stager=stager)
+    return runner
+
+
+def _decode_step(runner):
+    forward_inputs = types.SimpleNamespace(
+        batch=types.SimpleNamespace(num_context_requests=0),
+        new_tensors_device=None,
+        cache_indirection_buffer=None,
+        gather_context_logits=False,
+        enable_spec_decode=False,
+        runtime_draft_len=0,
+    )
+    resource_manager = types.SimpleNamespace(get_resource_manager=lambda key: object())
+    return runner._forward_decoder(forward_inputs, resource_manager, is_dummy=False)
+
+
+def _names(log: list) -> list:
+    return [entry if isinstance(entry, str) else entry[0] for entry in log]
+
+
+_GRAPH_KEY = (1, 0, False)
+_STAGED_STEP = ["begin", "prepare", "end", "record_input_copy", "completion_event", "replay"]
+
+
+def test_a_replay_issues_the_staged_preparation_before_the_replay(monkeypatch):
+    log = []
+    graphs = _WindowGraphs(log, _GRAPH_KEY)
+    runner = _window_runner(monkeypatch, graphs, _WindowPreparer(log), _ReferenceStager(log))
+
+    assert _decode_step(runner) == "outputs"
+    assert log == [
+        ("begin", True, graphs.cuda_graph_runner.shared_static_tensors),
+        "prepare",
+        ("end", False),
+        "record_input_copy",
+        ("completion_event", "input-copy-event"),
+        ("replay", 3),
+    ]
+
+
+@pytest.mark.parametrize(
+    "key, needs_capture, warmup_only, ran",
+    [
+        (None, False, False, "eager"),
+        (_GRAPH_KEY, True, False, "capture"),
+        (_GRAPH_KEY, False, True, "capture"),
+    ],
+    ids=["eager", "initial_capture", "warmup_only"],
+)
+def test_a_step_that_is_not_a_replay_never_opens_a_window(
+    monkeypatch, key, needs_capture, warmup_only, ran
+):
+    log = []
+    graphs = _WindowGraphs(log, key, needs_capture=needs_capture, warmup_only=warmup_only)
+    runner = _window_runner(monkeypatch, graphs, _WindowPreparer(log), _ReferenceStager(log))
+
+    _decode_step(runner)
+    assert log == ["prepare", "record_input_copy", ran]
+    assert graphs.cuda_graph_runner.staged_static_input_tokens == 0
+
+
+def test_a_declined_replay_prepares_natively(monkeypatch):
+    log = []
+    graphs = _WindowGraphs(log, _GRAPH_KEY)
+    stager = _ReferenceStager(log, accept=False)
+    runner = _window_runner(monkeypatch, graphs, _WindowPreparer(log), stager)
+
+    _decode_step(runner)
+    assert _names(log) == ["begin", "prepare", "record_input_copy", "replay"]
+    assert log[-1] == ("replay", 0)
+
+
+def test_a_failed_preparation_discards_the_window(monkeypatch):
+    log = []
+    graphs = _WindowGraphs(log, _GRAPH_KEY)
+    runner = _window_runner(
+        monkeypatch, graphs, _WindowPreparer(log, fail=True), _ReferenceStager(log)
+    )
+
+    with pytest.raises(RuntimeError, match="preparation failed"):
+        _decode_step(runner)
+    assert _names(log) == ["begin", "prepare", "discard"]
+
+
+@pytest.mark.parametrize(
+    "fail, expected",
+    [
+        ("before_replay", ["begin", "prepare", "discard"]),
+        ("replay", _STAGED_STEP),
+    ],
+    ids=["window_open", "window_closed"],
+)
+def test_a_failed_run_discards_only_a_window_still_open(monkeypatch, fail, expected):
+    log = []
+    graphs = _WindowGraphs(log, _GRAPH_KEY, fail=fail)
+    runner = _window_runner(monkeypatch, graphs, _WindowPreparer(log), _ReferenceStager(log))
+
+    with pytest.raises(RuntimeError):
+        _decode_step(runner)
+    assert _names(log) == expected
+
+
+def test_replay_eager_replay_opens_one_window_per_replay(monkeypatch):
+    log = []
+    graphs = _WindowGraphs(log, _GRAPH_KEY)
+    runner = _window_runner(monkeypatch, graphs, _WindowPreparer(log), _ReferenceStager(log))
+
+    _decode_step(runner)
+    graphs.key = None
+    _decode_step(runner)
+    graphs.key = _GRAPH_KEY
+    _decode_step(runner)
+    assert _names(log) == [
+        *_STAGED_STEP,
+        "prepare",
+        "record_input_copy",
+        "eager",
+        *_STAGED_STEP,
+    ]
+    assert [entry for entry in log if entry[0] == "replay"] == [("replay", 3), ("replay", 3)]
+
+
+def _replay_graphs(monkeypatch, log: list, *, needs_capture: bool, warmup_only: bool):
+    """A ``DecoderCudaGraphs`` with only what ``_run`` reads, logging capture and replay."""
+    from tensorrt_llm._torch.pyexecutor.engine.runners.decoder import graphs as graphs_module
+    from tensorrt_llm._torch.pyexecutor.engine.runners.decoder.graphs import DecoderCudaGraphs
+
+    monkeypatch.setattr(graphs_module, "prepare_attn_metadata_for_draft_replay", lambda *args: None)
+    monkeypatch.setattr(
+        graphs_module, "restore_attn_metadata_after_draft_replay", lambda *args: None
+    )
+    monkeypatch.setattr(
+        graphs_module, "MoeLoadBalancerIterContext", lambda balancer: contextlib.nullcontext()
+    )
+    graphs = DecoderCudaGraphs.__new__(DecoderCudaGraphs)
+    graphs.breakable_cuda_graph_runner = None
+    graphs._torch_compile_bypass = contextlib.nullcontext
+    graphs.cuda_graph_runner = types.SimpleNamespace(
+        needs_capture=lambda key: needs_capture,
+        is_warmup_only=warmup_only,
+        capture=lambda key, fn, inputs, **kwargs: log.append("capture"),
+        replay=lambda key, inputs, **kwargs: log.append("replay"),
+    )
+    return graphs
+
+
+@pytest.mark.parametrize(
+    "key, needs_capture, warmup_only, expected",
+    [
+        (_GRAPH_KEY, False, False, ["before_replay", "replay"]),
+        (_GRAPH_KEY, True, False, ["capture", "replay"]),
+        (_GRAPH_KEY, True, True, ["capture"]),
+        (None, False, False, ["eager"]),
+    ],
+    ids=["replay", "capture_then_replay", "warmup_only", "eager"],
+)
+def test_before_replay_runs_only_before_the_replay_of_a_captured_graph(
+    monkeypatch, key, needs_capture, warmup_only, expected
+):
+    log = []
+    graphs = _replay_graphs(monkeypatch, log, needs_capture=needs_capture, warmup_only=warmup_only)
+
+    graphs._run(
+        key,
+        {"input_ids": torch.zeros(2, dtype=torch.int64)},
+        lambda inputs: log.append("eager"),
+        lambda inputs: None,
+        attn_metadata=None,
+        draft_kv_cache_manager=None,
+        moe_load_balancer=None,
+        enable_spec_decode=True,
+        gather_context_logits=False,
+        position_id_offsets=None,
+        before_replay=lambda: log.append("before_replay"),
+    )
+    assert log == expected
+
+
+# -- the speculative sampler's store writer -----------------------------------------------------
+
+
+class _ReferenceStoreWriter:
+    """Stores like the native ``index_copy_`` path, then copies the store to the host tensors."""
+
+    def __init__(self, supported: bool = True) -> None:
+        self._supported = supported
+        self.calls = []
+
+    def supported(self, store, new_tokens, next_new_tokens, new_tokens_lens, next_draft_tokens):
+        self.calls.append("supported")
+        return self._supported
+
+    def store_and_copy_out(
+        self,
+        store,
+        slots,
+        new_tokens,
+        next_new_tokens,
+        new_tokens_lens,
+        next_draft_tokens,
+        host_new_tokens,
+        host_new_tokens_lens,
+        host_next_draft_tokens,
+        record_event=True,
+    ) -> None:
+        self.calls.append(("store_and_copy_out", list(slots), record_event))
+        index = torch.tensor(slots, dtype=torch.long)
+        store.new_tokens.squeeze(-1).T.index_copy_(0, index, new_tokens)
+        store.next_new_tokens.squeeze(-1).T.index_copy_(0, index, next_new_tokens)
+        store.new_tokens_lens.index_copy_(0, index, new_tokens_lens)
+        store.next_draft_tokens.index_copy_(0, index, next_draft_tokens)
+        host_new_tokens.copy_(store.new_tokens)
+        host_new_tokens_lens.copy_(store.new_tokens_lens)
+        host_next_draft_tokens.copy_(store.next_draft_tokens)
+
+    def set_completion_event(self, event) -> None:
+        self.calls.append(("completion_event", event))
+
+
+def _cpu_sampling(monkeypatch, *, pinned: bool) -> None:
+    """Run ``SpecSampler.sample_async`` on CPU: device transfers stay on the host."""
+    from tensorrt_llm._torch.speculative import spec_sampler_base
+
+    original_to = torch.Tensor.to
+    original_empty_like = torch.empty_like
+
+    def to_cpu(tensor, *args, **kwargs):
+        if kwargs.get("device") == "cuda":
+            kwargs = {**kwargs, "device": "cpu"}
+        return original_to(tensor, *args, **kwargs)
+
+    def empty_like(tensor, *args, pin_memory=False, **kwargs):
+        # Pinning needs a CUDA context; the writer's host buffers are plain host tensors here.
+        return original_empty_like(tensor, *args, **kwargs)
+
+    monkeypatch.setattr(torch.Tensor, "to", to_cpu)
+    monkeypatch.setattr(torch, "empty_like", empty_like)
+    monkeypatch.setattr(spec_sampler_base, "prefer_pinned", lambda: pinned)
+    monkeypatch.setattr(spec_sampler_base, "handle_stop_criteria", lambda *args, **kwargs: False)
+
+
+def _store_sampler(writer, *, async_worker: bool = False):
+    """A ``SpecSampler`` with a six-slot store, draft length 3 and the given store writer."""
+    from tensorrt_llm._torch.speculative.spec_sampler_base import SpecSampler
+
+    sampler = SpecSampler.__new__(SpecSampler)
+    sampler.draft_len = 3
+    sampler.max_accepted_path_len = 4
+    sampler.max_seq_len = 128
+    sampler.store = types.SimpleNamespace(
+        new_tokens=torch.zeros(4, 6, 1, dtype=torch.int32),
+        next_new_tokens=torch.zeros(4, 6, 1, dtype=torch.int32),
+        next_draft_tokens=torch.zeros(6, 3, dtype=torch.int32),
+        new_tokens_lens=torch.zeros(6, dtype=torch.int32),
+    )
+    sampler.store_writer = writer
+    sampler._async_worker_active = lambda: async_worker
+    sampler._copy_to_host = lambda tensor: tensor.clone()
+    sampler._record_sampler_event = lambda: types.SimpleNamespace(cuda_event="sampler-event")
+    return sampler
+
+
+def _sample_generation_step(sampler, *, verify_lens: bool = False):
+    """Three generation requests in nonconsecutive slots 4, 1 and 3."""
+    from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestState
+
+    def request(request_id: int, slot: int):
+        return types.SimpleNamespace(
+            py_request_id=request_id,
+            py_seq_slot=slot,
+            state=LlmRequestState.GENERATION_IN_PROGRESS,
+            py_draft_tokens=[10, 11, 12],
+            py_draft_tokens_effective_len=None,
+            py_verify_len=None,
+            py_decoding_iter=0,
+        )
+
+    scheduled = types.SimpleNamespace(
+        context_requests_chunking=[],
+        context_requests_last_chunk=[],
+        generation_requests=[request(7, 4), request(8, 1), request(9, 3)],
+    )
+    outputs = dict(
+        new_tokens=torch.arange(12, dtype=torch.int32).reshape(3, 4) + 100,
+        new_tokens_lens=torch.tensor([1, 3, 2], dtype=torch.int32),
+        next_draft_tokens=torch.arange(9, dtype=torch.int32).reshape(3, 3) + 200,
+        next_new_tokens=torch.arange(12, dtype=torch.int32).reshape(3, 4) + 300,
+    )
+    if verify_lens:
+        outputs["verify_lens"] = torch.tensor([2, 3, 1], dtype=torch.int32)
+        outputs["verify_lens_in_output_order"] = True
+    return sampler.sample_async(scheduled, outputs, [])
+
+
+def _assert_same_stores(sampler, state, native, native_state) -> None:
+    for name in ("new_tokens", "next_new_tokens", "next_draft_tokens", "new_tokens_lens"):
+        assert torch.equal(getattr(sampler.store, name), getattr(native.store, name)), name
+    for name in ("new_tokens", "new_tokens_lens", "next_draft_tokens"):
+        assert torch.equal(getattr(state.host, name), getattr(native_state.host, name)), name
+
+
+def test_the_store_writer_matches_the_native_stores_on_nonconsecutive_slots(monkeypatch):
+    _cpu_sampling(monkeypatch, pinned=True)
+    native = _store_sampler(None)
+    writer = _ReferenceStoreWriter()
+    sampler = _store_sampler(writer)
+
+    native_state = _sample_generation_step(native)
+    state = _sample_generation_step(sampler)
+    assert writer.calls == [
+        "supported",
+        ("store_and_copy_out", [4, 1, 3], False),
+        ("completion_event", "sampler-event"),
+    ]
+    assert native.store.new_tokens_lens.tolist() == [0, 3, 0, 2, 1, 0]
+    _assert_same_stores(sampler, state, native, native_state)
+    assert state.host.verify_lens is None
+    assert state.device.new_tokens is sampler.store.next_new_tokens
+
+
+@pytest.mark.parametrize("case", ["declined", "verify_lens", "async_worker", "not_pinned"])
+def test_the_store_writer_falls_back_to_the_native_stores(monkeypatch, case):
+    _cpu_sampling(monkeypatch, pinned=case != "not_pinned")
+    native = _store_sampler(None)
+    writer = _ReferenceStoreWriter(supported=case != "declined")
+    sampler = _store_sampler(writer, async_worker=case == "async_worker")
+
+    native_state = _sample_generation_step(native, verify_lens=case == "verify_lens")
+    state = _sample_generation_step(sampler, verify_lens=case == "verify_lens")
+    assert writer.calls == (["supported"] if case == "declined" else [])
+    _assert_same_stores(sampler, state, native, native_state)
+    if case == "verify_lens":
+        assert state.host.verify_lens.tolist() == [2, 3, 1]

From ce1d3f397fe4613a3207827a74b81fb2485c0a46 Mon Sep 17 00:00:00 2001
From: Xiao Wang <24860335+xwang233@users.noreply.github.com>
Date: Mon, 26 Oct 2026 21:49:27 -0700
Subject: [PATCH 4/5] [TRTLLM-17096][chore] runtime hooks: drop the step-input
 stager and the sampler store writer

Main's one-model speculative decoding work now batches the overlap input
gathers and moves the sampler's host copies off the critical path natively,
so these two hooks no longer fit the code they replaced. Remove them and
their call sites; they can be ported onto the new native paths later.

Keep the decode prologue, the executor events, the execution-stream loop and
the two extension points (prepended FMHA libraries, the DSA indexer's
overridable top-k copy).

Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com>
---
 .../attention/backends/sparse/dsa/indexer.py  |  20 +-
 .../attention/backends/sparse/dsa/metadata.py |  98 +--
 .../_torch/attention/backends/trtllm.py       |  59 +-
 tensorrt_llm/_torch/pyexecutor/_util.py       |   8 +-
 .../_torch/pyexecutor/cuda_graph_runner.py    |  14 +-
 .../pyexecutor/engine/runners/common.py       |   8 -
 .../engine/runners/decoder/graphs.py          |   7 -
 .../engine/runners/decoder/inputs.py          |  86 +-
 .../engine/runners/decoder/runner.py          | 153 +---
 .../kv_cache/kv_cache_manager_v2.py           |  13 -
 .../_torch/pyexecutor/runtime_hooks.py        | 154 +---
 tensorrt_llm/_torch/speculative/eagle3.py     |   9 +-
 tensorrt_llm/_torch/speculative/interface.py  |  10 +-
 .../_torch/speculative/spec_sampler_base.py   |  82 +-
 .../_torch/executor/test_runtime_hooks.py     | 754 +-----------------
 15 files changed, 129 insertions(+), 1346 deletions(-)

diff --git a/tensorrt_llm/_torch/attention/backends/sparse/dsa/indexer.py b/tensorrt_llm/_torch/attention/backends/sparse/dsa/indexer.py
index 9e1ff5b6f902..d69ea5be0128 100644
--- a/tensorrt_llm/_torch/attention/backends/sparse/dsa/indexer.py
+++ b/tensorrt_llm/_torch/attention/backends/sparse/dsa/indexer.py
@@ -1240,18 +1240,6 @@ def recompute_slot_mappings(
         metadata.host_slot_mapping_fp8[:total_new_kv_tokens] = fp8_flat_indices
         metadata.host_slot_mapping_scale[:total_new_kv_tokens] = scale_flat_indices
 
-        stager = getattr(metadata, "step_input_stager", None)
-        if stager is not None:
-            # Staged write; see the window banner in DecoderRunner._forward_decoder.
-            stager.copy(
-                metadata.slot_mapping_fp8[:total_new_kv_tokens],
-                metadata.host_slot_mapping_fp8[:total_new_kv_tokens],
-            )
-            stager.copy(
-                metadata.slot_mapping_scale[:total_new_kv_tokens],
-                metadata.host_slot_mapping_scale[:total_new_kv_tokens],
-            )
-            return
         metadata.slot_mapping_fp8[:total_new_kv_tokens].copy_(
             metadata.host_slot_mapping_fp8[:total_new_kv_tokens], non_blocking=True
         )
@@ -1526,13 +1514,7 @@ def prepare(metadata: DSAtrtllmAttentionMetadata) -> None:
         if num_generations > 0:
             # Prepare schedule metadata for fp8_paged_mqa_logits
             # This is a preprocessing step that computes scheduling information for the kernel
-            # Staged write; see the window banner in DecoderRunner._forward_decoder.
-            stager = getattr(metadata, "step_input_stager", None)
-            if stager is None:
-                Indexer.prepare_scheduler_metadata(metadata)
-            elif not stager.scheduler_metadata(metadata):
-                # Reads staged KV lengths: run right after the staged launch.
-                stager.defer(lambda: Indexer.prepare_scheduler_metadata(metadata))
+            Indexer.prepare_scheduler_metadata(metadata)
 
     def _update_k_cache(
         self, k_fp8: torch.Tensor, k_scale: torch.Tensor, metadata: DSAtrtllmAttentionMetadata
diff --git a/tensorrt_llm/_torch/attention/backends/sparse/dsa/metadata.py b/tensorrt_llm/_torch/attention/backends/sparse/dsa/metadata.py
index 5c9c4686375c..e0773630dcb3 100644
--- a/tensorrt_llm/_torch/attention/backends/sparse/dsa/metadata.py
+++ b/tensorrt_llm/_torch/attention/backends/sparse/dsa/metadata.py
@@ -1486,15 +1486,6 @@ def _update_indexer_k_cache_block_offsets(self) -> torch.Tensor:
         page_indices = pool_indices * cache_manager.indexer_k_cache_page_scale
         num_blocks = page_indices.shape[1]
         self.host_indexer_k_cache_block_offsets[: self.num_seqs, :num_blocks].copy_(page_indices)
-        if self.step_input_stager is not None:
-            # Staged write; see the window banner in DecoderRunner._forward_decoder.
-            self.step_input_stager.copy(
-                self.indexer_k_cache_block_offsets[: self.num_seqs],
-                self.host_indexer_k_cache_block_offsets[: self.num_seqs],
-                clamp_min=0,
-            )
-            self.step_input_stager.clamp_min(self.indexer_k_cache_block_offsets[self.num_seqs :], 0)
-            return pool_indices
         self.indexer_k_cache_block_offsets[: self.num_seqs].copy_(
             self.host_indexer_k_cache_block_offsets[: self.num_seqs], non_blocking=True
         )
@@ -1711,16 +1702,9 @@ def prepare_dense_topk_indices(self, kv_lens, device=False):  # device=False mea
                     kv_lens[self.num_contexts : self.num_seqs],
                     self.num_tokens - self.num_ctx_tokens,
                 )
-                if self.step_input_stager is not None:
-                    # Staged write; see the window banner in DecoderRunner._forward_decoder.
-                    self.step_input_stager.copy(
-                        self.topk_indices_buffer[gen_range, :],
-                        self.host_topk_indices_buffer[gen_range, :],
-                    )
-                else:
-                    self.topk_indices_buffer[gen_range, :].copy_(
-                        self.host_topk_indices_buffer[gen_range, :], non_blocking=True
-                    )
+                self.topk_indices_buffer[gen_range, :].copy_(
+                    self.host_topk_indices_buffer[gen_range, :], non_blocking=True
+                )
 
     def prepare_for_spec_decode(self, kv_lens: torch.Tensor):
         # The DeepGEMM paged-MQA kernel (fp8_paged_mqa_logits) runs a native
@@ -1820,16 +1804,9 @@ def prepare_for_spec_decode(self, kv_lens: torch.Tensor):
                 num_tokens = self.num_generations * expand_factor
                 gen_kv_lens_expanded = gen_kv_lens.repeat_interleave(expand_factor)
                 self.kv_lens_expanded_host[:num_tokens].copy_(gen_kv_lens_expanded)
-                if self.step_input_stager is not None:
-                    # Staged write; see the window banner in DecoderRunner._forward_decoder.
-                    self.step_input_stager.copy(
-                        self.kv_lens_expanded_cuda[:num_tokens],
-                        self.kv_lens_expanded_host[:num_tokens],
-                    )
-                else:
-                    self.kv_lens_expanded_cuda[:num_tokens].copy_(
-                        self.kv_lens_expanded_host[:num_tokens], non_blocking=True
-                    )
+                self.kv_lens_expanded_cuda[:num_tokens].copy_(
+                    self.kv_lens_expanded_host[:num_tokens], non_blocking=True
+                )
                 self._refresh_expanded_block_table(expand_factor)
         else:
             self.dsl_expand_factor = 1
@@ -2059,15 +2036,6 @@ def _refresh_expanded_block_table(self, repeat_factor: Optional[int] = None):
         self.host_block_table_expanded[:num_tokens, :max_len].copy_(
             expanded_blocks, non_blocking=True
         )
-        if self.step_input_stager is not None:
-            # Staged write; see the window banner in DecoderRunner._forward_decoder.
-            self.step_input_stager.copy(
-                self.block_table_expanded[:num_tokens],
-                self.host_block_table_expanded[:num_tokens],
-                clamp_min=0,
-            )
-            self.step_input_stager.clamp_min(self.block_table_expanded[num_tokens:], 0)
-            return
         self.block_table_expanded[:num_tokens].copy_(
             self.host_block_table_expanded[:num_tokens], non_blocking=True
         )
@@ -2108,15 +2076,9 @@ def prepare_for_indexer_k_cache(self):
             pad_cols.unsqueeze(0) >= num_blocks_per_seq[: self.num_seqs].unsqueeze(1), -1
         )
         # Copy to GPU
-        if self.step_input_stager is not None:
-            # Staged write; see the window banner in DecoderRunner._forward_decoder.
-            self.step_input_stager.copy(
-                self.block_table[: self.num_seqs, :max_blocks_used], host_block_table
-            )
-        else:
-            self.block_table[: self.num_seqs, :max_blocks_used].copy_(
-                host_block_table, non_blocking=True
-            )
+        self.block_table[: self.num_seqs, :max_blocks_used].copy_(
+            host_block_table, non_blocking=True
+        )
 
     def prepare_for_skip_indexer(self, kv_lens: torch.Tensor):
         num_extra_kv_tokens = self.kv_cache_params.num_extra_kv_tokens
@@ -2182,17 +2144,9 @@ def prepare_for_mla_rope_append(self, cached_token_lens: torch.Tensor, kv_lens:
                 dtype=torch.int64,
                 out=self.host_gen_cached_token_indptr[1 : self.num_generations + 1],
             )
-            if self.step_input_stager is not None:
-                # Staged write; see the window banner in DecoderRunner._forward_decoder.
-                self.step_input_stager.copy(
-                    self.gen_cached_token_indptr[: self.num_generations + 1],
-                    self.host_gen_cached_token_indptr[: self.num_generations + 1],
-                )
-            else:
-                self.gen_cached_token_indptr[: self.num_generations + 1].copy_(
-                    self.host_gen_cached_token_indptr[: self.num_generations + 1],
-                    non_blocking=True,
-                )
+            self.gen_cached_token_indptr[: self.num_generations + 1].copy_(
+                self.host_gen_cached_token_indptr[: self.num_generations + 1], non_blocking=True
+            )
             # generation kv indptr
             torch.cumsum(
                 kv_lens[self.num_contexts : self.num_seqs],
@@ -2200,16 +2154,9 @@ def prepare_for_mla_rope_append(self, cached_token_lens: torch.Tensor, kv_lens:
                 dtype=torch.int64,
                 out=self.host_gen_kv_indptr[1 : self.num_generations + 1],
             )
-            if self.step_input_stager is not None:
-                # Staged write; see the window banner in DecoderRunner._forward_decoder.
-                self.step_input_stager.copy(
-                    self.gen_kv_indptr[: self.num_generations + 1],
-                    self.host_gen_kv_indptr[: self.num_generations + 1],
-                )
-            else:
-                self.gen_kv_indptr[: self.num_generations + 1].copy_(
-                    self.host_gen_kv_indptr[: self.num_generations + 1], non_blocking=True
-                )
+            self.gen_kv_indptr[: self.num_generations + 1].copy_(
+                self.host_gen_kv_indptr[: self.num_generations + 1], non_blocking=True
+            )
         else:
             self.max_gen_seq_len = 0
 
@@ -2221,14 +2168,7 @@ def prepare_for_indices_conversion(self):
             self.seq_lens,
             dim=0,
         )
-        if self.step_input_stager is not None:
-            # Staged write; see the window banner in DecoderRunner._forward_decoder.
-            self.step_input_stager.copy(
-                self.req_idx_per_token[: self.num_tokens],
-                self.host_req_idx_per_token[: self.num_tokens],
-            )
-        else:
-            self.req_idx_per_token[: self.num_tokens].copy_(
-                self.host_req_idx_per_token[: self.num_tokens],
-                non_blocking=True,
-            )
+        self.req_idx_per_token[: self.num_tokens].copy_(
+            self.host_req_idx_per_token[: self.num_tokens],
+            non_blocking=True,
+        )
diff --git a/tensorrt_llm/_torch/attention/backends/trtllm.py b/tensorrt_llm/_torch/attention/backends/trtllm.py
index 64b29917880b..c3d757fcd071 100644
--- a/tensorrt_llm/_torch/attention/backends/trtllm.py
+++ b/tensorrt_llm/_torch/attention/backends/trtllm.py
@@ -28,7 +28,6 @@
     from tensorrt_llm.mapping import Mapping
 
     from ...model_config import ModelConfig
-    from ...pyexecutor.runtime_hooks import StepInputStager
     from ...speculative.interface import SpecMetadata
     from ...speculative.spec_tree_manager import SpecTreeManager
     from .sparse.dsa.metadata import TokenMajorGenView
@@ -124,8 +123,6 @@ def generate_spec_decoding_packed_mask(max_num_requests: int,
 
 @dataclass(kw_only=True)
 class TrtllmAttentionMetadata(AttentionMetadata):
-    # A step-input stager while it records this metadata's writes, else None.
-    step_input_stager: ClassVar[Optional["StepInputStager"]] = None
     workspace: Optional[torch.Tensor] = None
     cuda_graph_workspace: Optional[torch.Tensor] = None
     workspace_reclaimable: bool = field(default=True, init=False)
@@ -1009,14 +1006,8 @@ def prepare(self) -> None:
             device='cpu',
         )
         self.prompt_lens_cpu[:self.num_seqs].copy_(prompt_lens)
-        stager = self.step_input_stager
-        if stager is not None:
-            # Staged write; see the window banner in DecoderRunner._forward_decoder.
-            stager.copy(self.prompt_lens_cuda[:self.num_seqs],
-                        self.prompt_lens_cpu[:self.num_seqs])
-        else:
-            self.prompt_lens_cuda[:self.num_seqs].copy_(
-                self.prompt_lens_cpu[:self.num_seqs], non_blocking=True)
+        self.prompt_lens_cuda[:self.num_seqs].copy_(
+            self.prompt_lens_cpu[:self.num_seqs], non_blocking=True)
 
         # number of tokens in the kv cache for each sequence in the batch
         cached_token_lens = torch.tensor(
@@ -1069,14 +1060,9 @@ def prepare(self) -> None:
         # the sequence length including the cached tokens and the input tokens.
         self.kv_lens[:self.num_seqs].copy_(
             kv_lens + self.kv_cache_params.num_extra_kv_tokens)
-        if stager is not None:
-            # Staged write; see the window banner in DecoderRunner._forward_decoder.
-            stager.copy(self.kv_lens_cuda[:self.num_seqs],
-                        kv_lens[:self.num_seqs])
-        else:
-            self.kv_lens_cuda[:self.num_seqs].copy_(maybe_pin_memory(
-                kv_lens[:self.num_seqs]),
-                                                    non_blocking=True)
+        self.kv_lens_cuda[:self.num_seqs].copy_(maybe_pin_memory(
+            kv_lens[:self.num_seqs]),
+                                                non_blocking=True)
         # total kv lens for context requests and generation requests, without extra tokens
         self.host_total_kv_lens[0] = kv_lens[:self.num_contexts].sum().item()
         self.host_total_kv_lens[1] = kv_lens[self.num_contexts:self.
@@ -1443,42 +1429,23 @@ def prepare_context_mla_with_cached_kv(self,
                      dtype=torch.int64,
                      out=self.host_ctx_cached_token_indptr[1:self.num_contexts +
                                                            1])
-        stager = self.step_input_stager
-        if stager is not None:
-            # Staged write; see the window banner in DecoderRunner._forward_decoder.
-            stager.copy(
-                self.ctx_cached_token_indptr[:self.num_contexts + 1],
-                self.host_ctx_cached_token_indptr[:self.num_contexts + 1])
-        else:
-            self.ctx_cached_token_indptr[:self.num_contexts + 1].copy_(
-                self.host_ctx_cached_token_indptr[:self.num_contexts + 1],
-                non_blocking=True)
+        self.ctx_cached_token_indptr[:self.num_contexts + 1].copy_(
+            self.host_ctx_cached_token_indptr[:self.num_contexts + 1],
+            non_blocking=True)
         torch.cumsum(
             self.seq_lens[:self.num_contexts],
             dim=0,
             dtype=torch.int64,
             out=self.host_ctx_uncached_token_indptr[1:self.num_contexts + 1])
-        if stager is not None:
-            # Staged write; see the window banner in DecoderRunner._forward_decoder.
-            stager.copy(
-                self.ctx_uncached_token_indptr[:self.num_contexts + 1],
-                self.host_ctx_uncached_token_indptr[:self.num_contexts + 1])
-        else:
-            self.ctx_uncached_token_indptr[:self.num_contexts + 1].copy_(
-                self.host_ctx_uncached_token_indptr[:self.num_contexts + 1],
-                non_blocking=True)
+        self.ctx_uncached_token_indptr[:self.num_contexts + 1].copy_(
+            self.host_ctx_uncached_token_indptr[:self.num_contexts + 1],
+            non_blocking=True)
         torch.cumsum(kv_lens[:self.num_contexts],
                      dim=0,
                      dtype=torch.int64,
                      out=self.host_ctx_kv_indptr[1:self.num_contexts + 1])
-        if stager is not None:
-            # Staged write; see the window banner in DecoderRunner._forward_decoder.
-            stager.copy(self.ctx_kv_indptr[:self.num_contexts + 1],
-                        self.host_ctx_kv_indptr[:self.num_contexts + 1])
-        else:
-            self.ctx_kv_indptr[:self.num_contexts + 1].copy_(
-                self.host_ctx_kv_indptr[:self.num_contexts + 1],
-                non_blocking=True)
+        self.ctx_kv_indptr[:self.num_contexts + 1].copy_(
+            self.host_ctx_kv_indptr[:self.num_contexts + 1], non_blocking=True)
 
     def spec_decoding_param_prepare_for_blackwell(self) -> None:
         """
diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py
index 2090179caa6c..4d3b4af850e8 100644
--- a/tensorrt_llm/_torch/pyexecutor/_util.py
+++ b/tensorrt_llm/_torch/pyexecutor/_util.py
@@ -53,7 +53,6 @@
 from ..speculative import (draft_prompt_lookahead, get_num_extra_kv_tokens,
                            get_num_spec_layers, get_spec_decoder,
                            should_use_separate_draft_kv_cache)
-from ..speculative.spec_sampler_base import SpecSampler
 from ..utils import is_gdn_replay_enabled
 from . import config_utils
 from .config_utils import (MambaKVCacheParams, _is_sliding_attention_layer,
@@ -78,7 +77,6 @@
 from .resource_manager import (KVCacheCompressionManager, KVCacheManager,
                                PeftCacheManager, ResourceManager,
                                ResourceManagerType)
-from .runtime_hooks import runtime_hooks_of
 from .sampler import EarlyStopSampler, EarlyStopWithMMResult, TorchSampler
 from .sampler.penalties import PenaltyStore
 from .scheduler import (BindCapacityScheduler, BindMicroBatchScheduler,
@@ -4589,11 +4587,7 @@ def instantiate_sampler(
     )
     if engine.spec_config is not None and engine.spec_config.spec_dec_mode.has_spec_decoder(
     ):
-        sampler = get_spec_decoder(sampler_args, engine.spec_config)
-        store_writer = runtime_hooks_of(engine.model).sampler_store_writer
-        if store_writer is not None and isinstance(sampler, SpecSampler):
-            sampler.store_writer = store_writer
-        return sampler
+        return get_spec_decoder(sampler_args, engine.spec_config)
 
     if mm_encoder_only:
         # NOTE: handle model outputs specially for mm encoder executor/engine.
diff --git a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py
index 118499f61732..e0d5abd61b35 100644
--- a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py
+++ b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py
@@ -160,10 +160,6 @@ class CUDAGraphRunner:
     multiple graphs, keyed by batch shape and execution-path specializations.
     """
     WARMUP_STEPS = 1
-    # Tokens of input_ids / position_ids that a step-input stager already
-    # wrote into the static inputs for the next replay (0: the replay copies
-    # them).
-    staged_static_input_tokens = 0
 
     def __init__(self, config: CUDAGraphRunnerConfig):
         self.config = config
@@ -808,12 +804,7 @@ def replay(
                 f"replay() got {seqlen} tokens for key {key}, but the graph "
                 f"was captured for {expected_num_tokens} tokens. A shorter "
                 "input_ids leaves the tail of the static input buffer stale.")
-        # A step-input stager already wrote input_ids / position_ids[:seqlen]
-        # into the static buffers.
-        staged = self.staged_static_input_tokens == seqlen
-        self.staged_static_input_tokens = 0
-        if not staged:
-            static_tensors["input_ids"][:seqlen].copy_(input_ids)
+        static_tensors["input_ids"][:seqlen].copy_(input_ids)
 
         position_ids = current_inputs["position_ids"]
         if self.config.use_mrope:
@@ -860,8 +851,7 @@ def replay(
                     f"for key {key}, but expected {expected_position_ids_shape}. "
                     "torch.Tensor.copy_() silently broadcasts mismatched shapes, "
                     "which would corrupt the static input buffer.")
-            if not staged:
-                static_tensors["position_ids"][:, :seqlen].copy_(position_ids)
+            static_tensors["position_ids"][:, :seqlen].copy_(position_ids)
 
         num_encoder_tokens = key.num_encoder_tokens
         if num_encoder_tokens:
diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/common.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/common.py
index cb24f68deef7..98dce6b41550 100644
--- a/tensorrt_llm/_torch/pyexecutor/engine/runners/common.py
+++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/common.py
@@ -22,7 +22,6 @@
 from .interface import ScheduledInputs
 
 if TYPE_CHECKING:
-    from tensorrt_llm._torch.pyexecutor.runtime_hooks import StepInputStager
     from tensorrt_llm._torch.pyexecutor.sampler.sampler import SampleStateTensors
 
 
@@ -138,13 +137,11 @@ def pad_input_tokens(
     enable_attention_dp: bool,
     prefill_cuda_graph_backend: PrefillCudaGraphBackend,
     prefill_cuda_graph_num_tokens: list[int],
-    step_input_stager: "StepInputStager | None" = None,
 ) -> int:
     """Pad the batch to the token count its prefill graph or attention-DP peers need.
 
     Zeroes the padded rows of ``input_ids`` and ``position_ids``, whose last
     dimension is the token one, and returns the token count the model runs on.
-    With ``step_input_stager`` the zeroing is staged on it.
     """
     padded_num_tokens, can_run_prefill_cuda_graph, attn_all_rank_num_tokens = get_padding_params(
         num_tokens,
@@ -161,11 +158,6 @@ def pad_input_tokens(
         attn_metadata.all_rank_num_tokens = attn_all_rank_num_tokens
     if attn_metadata.padded_num_tokens is None:
         return num_tokens
-    if step_input_stager is not None:
-        # Staged write; see the window banner in DecoderRunner._forward_decoder.
-        step_input_stager.zero(input_ids[num_tokens:padded_num_tokens])
-        step_input_stager.zero(position_ids[..., num_tokens:padded_num_tokens])
-        return padded_num_tokens
     input_ids[num_tokens:padded_num_tokens].fill_(0)
     position_ids[..., num_tokens:padded_num_tokens].fill_(0)
     return padded_num_tokens
diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/graphs.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/graphs.py
index 1dc3fca7e80f..4046fd5ca8ea 100644
--- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/graphs.py
+++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/graphs.py
@@ -487,14 +487,11 @@ def run(
         enable_spec_decode: bool,
         gather_context_logits: bool,
         position_id_offsets: torch.Tensor | None,
-        before_replay: Callable[[], None] | None = None,
     ) -> Any:
         """Run the model by replaying the graph for ``key``, capturing it first if needed.
 
         Without a key, record the model body while a breakable prefill graph
         captures, run through one when it covers the tokens, otherwise eagerly.
-        ``before_replay`` runs right before the replay of an already captured
-        graph, after the draft-cache metadata refresh.
         """
         with with_shared_pool(self.cuda_graph_runner.get_graph_pool()):
             if self.captures_model_body(key):
@@ -512,7 +509,6 @@ def run(
                 enable_spec_decode=enable_spec_decode,
                 gather_context_logits=gather_context_logits,
                 position_id_offsets=position_id_offsets,
-                before_replay=before_replay,
             )
 
     def _run(
@@ -528,7 +524,6 @@ def _run(
         enable_spec_decode: bool,
         gather_context_logits: bool,
         position_id_offsets: torch.Tensor | None,
-        before_replay: Callable[[], None] | None = None,
     ) -> Any:
         breakable_runner = self.breakable_cuda_graph_runner
         if key is None:
@@ -579,8 +574,6 @@ def capture_model(model_inputs: dict[str, Any]) -> Any:
                     attn_metadata, draft_kv_cache_manager
                 )
                 try:
-                    if before_replay is not None:
-                        before_replay()
                     with MoeLoadBalancerIterContext(moe_load_balancer):
                         outputs = self.cuda_graph_runner.replay(
                             key, inputs, position_id_offsets=position_id_offsets
diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/inputs.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/inputs.py
index 0aa1d1c35c76..694fa2d122da 100644
--- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/inputs.py
+++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/inputs.py
@@ -25,7 +25,7 @@
     ResourceManager,
     ResourceManagerType,
 )
-from tensorrt_llm._torch.pyexecutor.runtime_hooks import DecodePrologue, StepInputStager
+from tensorrt_llm._torch.pyexecutor.runtime_hooks import DecodePrologue
 from tensorrt_llm._torch.pyexecutor.sampler import SampleStateTensors
 from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests
 from tensorrt_llm._torch.pyexecutor.scheduler.scheduler import RemoteTailPhase
@@ -314,8 +314,6 @@ class DecoderInputPreparer:
 
     # The model's decode prologue (RuntimeHooks.decode_prologue), set by the runner.
     decode_prologue: DecodePrologue | None = None
-    # The model's step-input stager while it records this preparer's writes, else None.
-    step_input_stager: StepInputStager | None = None
 
     def __init__(
         self,
@@ -524,7 +522,11 @@ def pad_tokens(
                 mapping=self.mapping,
                 dist=self.dist,
             )
-        padding = dict(
+        return pad_input_tokens(
+            attn_metadata,
+            num_tokens,
+            num_contexts,
+            attn_all_rank_num_tokens,
             input_ids=self._buffers.input_ids_cuda,
             position_ids=self._position_ids_buffer(),
             dist=self.dist,
@@ -536,12 +538,6 @@ def pad_tokens(
             ),
             prefill_cuda_graph_num_tokens=self._config.prefill_cuda_graph_num_tokens,
         )
-        if self.step_input_stager is not None:
-            # Staged write; see the window banner in DecoderRunner._forward_decoder.
-            padding["step_input_stager"] = self.step_input_stager
-        return pad_input_tokens(
-            attn_metadata, num_tokens, num_contexts, attn_all_rank_num_tokens, **padding
-        )
 
     def model_inputs(
         self,
@@ -844,18 +840,10 @@ def prepare_inputs(
         self._copy_position_ids_to_device(host)
 
         if enable_spec_decode:
-            gather_ids = torch.tensor(
-                host.lists.gather_ids, dtype=torch.int, pin_memory=prefer_pinned()
+            self._buffers.gather_ids_cuda[: len(host.lists.gather_ids)].copy_(
+                torch.tensor(host.lists.gather_ids, dtype=torch.int, pin_memory=prefer_pinned()),
+                non_blocking=True,
             )
-            if self.step_input_stager is not None:
-                # Staged write; see the window banner in DecoderRunner._forward_decoder.
-                self.step_input_stager.copy(
-                    self._buffers.gather_ids_cuda[: len(host.lists.gather_ids)], gather_ids
-                )
-            else:
-                self._buffers.gather_ids_cuda[: len(host.lists.gather_ids)].copy_(
-                    gather_ids, non_blocking=True
-                )
 
         self._prepare_attn_metadata(
             host,
@@ -1721,52 +1709,24 @@ def previous_seq_slots_device() -> torch.Tensor:
         assert total_num_tokens <= self._config.max_num_tokens, (
             f"total_num_tokens ({total_num_tokens}) should be less than or equal to max_num_tokens ({self._config.max_num_tokens})"  # noqa: E501
         )
-        stager = self.step_input_stager
         # if exist requests that do not have previous batch, copy input_ids and draft_tokens
         if num_tokens > 0:
             input_ids = torch.tensor(input_ids, dtype=torch.int, pin_memory=prefer_pinned())
-            if stager is not None:
-                # Staged write; see the window banner in DecoderRunner._forward_decoder.
-                stager.copy(self._buffers.input_ids_cuda[:num_tokens], input_ids)
-            else:
-                self._buffers.input_ids_cuda[:num_tokens].copy_(input_ids, non_blocking=True)
+            self._buffers.input_ids_cuda[:num_tokens].copy_(input_ids, non_blocking=True)
 
         if num_draft_tokens > 0:
             draft_tokens = torch.tensor(draft_tokens, dtype=torch.int, pin_memory=prefer_pinned())
-            if stager is not None:
-                # Staged write; see the window banner in DecoderRunner._forward_decoder.
-                stager.copy(self._buffers.draft_tokens_cuda[: len(draft_tokens)], draft_tokens)
-            else:
-                self._buffers.draft_tokens_cuda[: len(draft_tokens)].copy_(
-                    draft_tokens, non_blocking=True
-                )
+            self._buffers.draft_tokens_cuda[: len(draft_tokens)].copy_(
+                draft_tokens, non_blocking=True
+            )
         if self._config.is_spec_decode and len(num_accepted_draft_tokens) > 0:
             num_accepted_draft_tokens = torch.tensor(
                 num_accepted_draft_tokens, dtype=torch.int, pin_memory=prefer_pinned()
             )
-            accepted = self._buffers.num_accepted_draft_tokens_cuda[
-                : len(num_accepted_draft_tokens)
-            ]
-            if stager is not None:
-                # Staged write; see the window banner in DecoderRunner._forward_decoder.
-                stager.copy(accepted, num_accepted_draft_tokens)
-            else:
-                accepted.copy_(num_accepted_draft_tokens, non_blocking=True)
-        if next_draft_tokens_device is not None and stager is not None:
-            # Staged write; see the window banner in DecoderRunner._forward_decoder.
-            stager.overlap_inputs(
-                self._buffers,
-                previous_batch_indices,
-                previous_pos_indices,
-                num_tokens,
-                num_draft_tokens,
-                self.get_runtime_tokens_per_gen_step(runtime_draft_len),
-                len(extend_requests) - len(extend_dummy_requests),
-                new_tokens_device,
-                next_draft_tokens_device,
-                new_tokens_lens_device,
+            self._buffers.num_accepted_draft_tokens_cuda[: len(num_accepted_draft_tokens)].copy_(
+                num_accepted_draft_tokens, non_blocking=True
             )
-        elif next_draft_tokens_device is not None:
+        if next_draft_tokens_device is not None:
             # Initialize these two values to zeros
             self._buffers.previous_pos_id_offsets_cuda.zero_()
             self._buffers.previous_kv_lens_offsets_cuda.zero_()
@@ -1866,10 +1826,6 @@ def previous_seq_slots_device() -> torch.Tensor:
         # cover the previous-batch rows; padding-request rows are covered by
         # neither, so without this the embedding would gather stale contents
         # of the uninitialized buffer, which may exceed the vocabulary size.
-        if stager is not None and host.lists.unwritten_input_id_rows:
-            # A native write to input_ids_cuda inside the window needs a flush
-            # first; see the window banner in DecoderRunner._forward_decoder.
-            stager.flush()
         _zero_fill_row_spans(self._buffers.input_ids_cuda, host.lists.unwritten_input_id_rows)
 
         host.num_tokens = num_tokens
@@ -1918,11 +1874,6 @@ def _copy_position_ids_to_device(self, host: HostInputs) -> None:
                 self._buffers.mrope_position_ids_cuda[:, :, start_idx:end_idx].copy_(
                     segment[:, :, : end_idx - start_idx], non_blocking=True
                 )
-        elif self.step_input_stager is not None:
-            # Staged write; see the window banner in DecoderRunner._forward_decoder.
-            self.step_input_stager.position_ids(
-                self._buffers.position_ids_cuda[:total_num_tokens], host_position_ids
-            )
         else:
             self._buffers.position_ids_cuda[:total_num_tokens].copy_(
                 host_position_ids, non_blocking=True
@@ -2231,11 +2182,6 @@ def record_input_copy(self) -> None:
         self._prepare_inputs_event = torch.cuda.Event()
         self._prepare_inputs_event.record()
 
-    @property
-    def input_copy_event(self) -> torch.cuda.Event | None:
-        """The event ``record_input_copy`` recorded last."""
-        return self._prepare_inputs_event
-
     def wait_for_input_copy(self) -> None:
         if self._prepare_inputs_event is not None:
             self._prepare_inputs_event.synchronize()
diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py
index e31f20da8e40..baacf2684be7 100644
--- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py
+++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py
@@ -622,103 +622,34 @@ def _forward_decoder(
                     execution_requests.generation_requests,
                 )
 
-            # ==== Step-input window opens here ====
-            # When the model declares a step-input stager and begin_if_eligible
-            # returns it, every device write that prepare_inputs (and the
-            # metadata and KV-cache-manager code it calls) routes through
-            # ``step_input_stager`` (call sites marked "Staged write") is only
-            # recorded; close_input_window below, right before the graph
-            # replay, issues them all. Until then, do not read or write those
-            # buffers with any other device op: a read sees stale bytes and a
-            # write is overwritten. Route new writes through the stager or
-            # call stager.flush() first. Exception: input_ids_cuda and
-            # position_ids_cuda are also staged into the graph's static
-            # inputs, which the replay then does not refresh from them; to
-            # change them here, call stager.flush() and then write them
-            # natively. Nothing checks this contract in a production run.
-            # Contract: StepInputStager in pyexecutor/runtime_hooks.py.
-            stager = None
-            stager_source = self.runtime_hooks.step_input_stager
-            if stager_source is not None:
-                # Only the replay of an already captured graph closes the window
-                # (before_replay); an eager step, an initial capture or a warmup
-                # would run the model before the staged writes are issued, so
-                # those steps always prepare natively.
-                graph_replay = (
-                    can_run_graph
-                    and not self.cuda_graph_runner.needs_capture(key)
-                    and not self.cuda_graph_runner.is_warmup_only
-                )
-                if graph_replay:
-                    stager = stager_source.begin_if_eligible(
-                        self._inputs,
-                        execution_requests,
-                        attn_metadata,
-                        spec_metadata,
-                        new_tensors_device,
-                        graph_replay=graph_replay,
-                        static_inputs=self.cuda_graph_runner.shared_static_tensors,
-                        is_dummy=is_dummy,
-                        enable_spec_decode=enable_spec_decode,
-                        promoted=bool(execution_promoted_context_ids),
-                    )
-            try:
-                # The window spans the scope: the scope is a no-op whenever
-                # speculative decoding is on, which every stager requires.
-                with self._cuda_graph_metadata_prepare_scope(
-                    attn_metadata,
-                    key,
+            with self._cuda_graph_metadata_prepare_scope(
+                attn_metadata,
+                key,
+                execution_requests,
+                execution_promoted_context_ids,
+                enable_spec_decode=enable_spec_decode,
+            ):
+                prepared = self._inputs.prepare_inputs(
                     execution_requests,
+                    kv_cache_manager,
+                    attn_metadata,
+                    spec_metadata,
+                    new_tensors_device,
+                    cache_indirection_buffer,
+                    resource_manager,
+                    can_run_graph,
                     execution_promoted_context_ids,
+                    use_lora_graph=use_lora_graph,
                     enable_spec_decode=enable_spec_decode,
-                ):
-                    prepared = self._inputs.prepare_inputs(
-                        execution_requests,
-                        kv_cache_manager,
-                        attn_metadata,
-                        spec_metadata,
-                        new_tensors_device,
-                        cache_indirection_buffer,
-                        resource_manager,
-                        can_run_graph,
-                        execution_promoted_context_ids,
-                        use_lora_graph=use_lora_graph,
-                        enable_spec_decode=enable_spec_decode,
-                        runtime_draft_len=runtime_draft_len,
-                        is_dummy=is_dummy,
-                        gather_context_logits=gather_context_logits,
-                    )
-            except BaseException:
-                if stager is not None:
-                    stager.discard()
-                raise
+                    runtime_draft_len=runtime_draft_len,
+                    is_dummy=is_dummy,
+                    gather_context_logits=gather_context_logits,
+                )
             inputs = prepared.kwargs
             runtime_draft_len = prepared.runtime_draft_len
             if execution_promoted_context_ids:
                 self._inputs.record_promoted_iter_states(scheduled_requests)
-            close_input_window = None
-            window_open = stager is not None
-            if stager is None:
-                self._inputs.record_input_copy()
-            else:
-
-                def close_input_window() -> None:
-                    nonlocal window_open
-                    window_open = False
-                    # ==== Step-input window closes here ====
-                    # end() issues every staged write; nothing between
-                    # begin_if_eligible above and this call may have touched
-                    # the staged buffers with other device ops. The input-copy
-                    # event must follow the launch (host page-table reads
-                    # happen at execution time). A nonzero
-                    # static_inputs_tokens() makes the replay skip its
-                    # input_ids / position_ids copy.
-                    stager.end(record_event=False)
-                    self._inputs.record_input_copy()
-                    stager.set_completion_event(self._inputs.input_copy_event)
-                    self.cuda_graph_runner.staged_static_input_tokens = (
-                        stager.static_inputs_tokens()
-                    )
+            self._inputs.record_input_copy()
 
             def run_model(model_inputs: dict[str, Any]) -> Any:
                 with MoeLoadBalancerIterContext(moe_load_balancer):
@@ -740,30 +671,24 @@ def postprocess_inputs(model_inputs: dict[str, Any]) -> None:
 
             graph_key = key if can_run_graph else None
             captures_model_body = self._graphs.captures_model_body(graph_key)
-            try:
-                outputs = self._graphs.run(
-                    graph_key,
-                    inputs,
-                    run_model,
-                    postprocess_inputs,
-                    attn_metadata=attn_metadata,
-                    draft_kv_cache_manager=draft_kv_cache_manager,
-                    moe_load_balancer=moe_load_balancer,
-                    enable_spec_decode=enable_spec_decode,
-                    gather_context_logits=gather_context_logits,
-                    position_id_offsets=(
-                        self._buffers.previous_pos_id_offsets_cuda
-                        if enable_spec_decode
-                        and not self._config.disable_overlap_scheduler
-                        and attn_metadata.kv_cache_manager is not None
-                        else None
-                    ),
-                    before_replay=close_input_window,
-                )
-            except BaseException:
-                if window_open:
-                    stager.discard()
-                raise
+            outputs = self._graphs.run(
+                graph_key,
+                inputs,
+                run_model,
+                postprocess_inputs,
+                attn_metadata=attn_metadata,
+                draft_kv_cache_manager=draft_kv_cache_manager,
+                moe_load_balancer=moe_load_balancer,
+                enable_spec_decode=enable_spec_decode,
+                gather_context_logits=gather_context_logits,
+                position_id_offsets=(
+                    self._buffers.previous_pos_id_offsets_cuda
+                    if enable_spec_decode
+                    and not self._config.disable_overlap_scheduler
+                    and attn_metadata.kv_cache_manager is not None
+                    else None
+                ),
+            )
             if captures_model_body:
                 return outputs
 
diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py
index c825736a372c..07a5e93717aa 100644
--- a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py
+++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py
@@ -134,7 +134,6 @@
 
 if TYPE_CHECKING:
     from tensorrt_llm._torch.attention.backends.interface import AttentionMetadata
-    from tensorrt_llm._torch.pyexecutor.runtime_hooks import StepInputStager
     from tensorrt_llm.llmapi.llm_args import DecodingBaseConfig
 
 KV_CACHE_ITERATION_STATS_DELTA_FIELDS = _KV_CACHE_ITERATION_STATS_DELTA_FIELDS
@@ -1250,8 +1249,6 @@ def _spec_recompute_claim_limit(
 class KVCacheManagerV2(BaseResourceManager):
     # Sparse managers attach a metadata batch during initialization.
     sparse_metadata_batch: Batch | None = None
-    # A step-input stager while it records this manager's writes, else None.
-    step_input_stager: Optional["StepInputStager"] = None
     # Filled lazily by _cold_pool_group_membership(); the grouping is fixed after construction.
     # Declared on the class so it is present even when an instance is built without running __init__.
     _cold_pool_group_membership_cache: Optional[tuple[tuple[int, frozenset[int]], ...]] = None
@@ -6757,16 +6754,6 @@ def copy_batch_block_offsets(
             )
             return
 
-        if self.step_input_stager is not None:
-            # Staged write; see the window banner in DecoderRunner._forward_decoder.
-            self.step_input_stager.block_offsets(
-                self.host_kv_cache_block_offsets,
-                dst_tensor,
-                copy_idx.tolist(),
-                self._index_scale_ints,
-                self.kv_offset.tolist(),
-            )
-            return
         copy_batch_block_offsets_to_device(
             self.host_kv_cache_block_offsets,
             dst_tensor,
diff --git a/tensorrt_llm/_torch/pyexecutor/runtime_hooks.py b/tensorrt_llm/_torch/pyexecutor/runtime_hooks.py
index 72987309443f..772845919949 100644
--- a/tensorrt_llm/_torch/pyexecutor/runtime_hooks.py
+++ b/tensorrt_llm/_torch/pyexecutor/runtime_hooks.py
@@ -4,8 +4,8 @@
 
 A model declares its hooks as ``runtime_hooks`` (an instance attribute, a class
 attribute or a property). Every field defaults to off, and a model that declares
-nothing gets ``RuntimeHooks()``: the step loop, the samplers and the attention
-metadata then run exactly as they do without this module. Every engine call site
+nothing gets ``RuntimeHooks()``: the step loop then runs exactly as it does
+without this module. Every engine call site
 of a hook runs its original statements when the hook is off.
 
 The hooks change *where* and *how* the step loop moves data, never what a
@@ -16,151 +16,7 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Protocol, Sequence
-
-if TYPE_CHECKING:
-    import torch
-
-
-class StepInputStager(Protocol):
-    """Records a decode step's input preparation and issues it as one launch.
-
-    ``begin_if_eligible`` opens a deferred-write window on a step it supports
-    and returns the stager, or returns None and the engine prepares the step
-    natively. While the window is open, the engine's call sites route their
-    device writes to the staging methods instead of issuing them. ``end``
-    issues every staged write in staging order, then the deferred native
-    operations; ``discard`` drops them after a failed preparation.
-
-    While its window is open the stager sets itself as ``step_input_stager`` on
-    the decoder's input preparer, the attention metadata, the speculative
-    metadata and the KV cache managers, and resets that attribute to None at
-    ``end`` and ``discard``. Every call site reads the attribute and runs its
-    original statement when it is None.
-    """
-
-    def begin_if_eligible(
-        self,
-        preparer: Any,
-        requests: Any,
-        attn_metadata: Any,
-        spec_metadata: Any,
-        new_tensors_device: Any,
-        *,
-        graph_replay: bool,
-        static_inputs: Optional[Dict[str, "torch.Tensor"]],
-        is_dummy: bool,
-        enable_spec_decode: bool,
-        promoted: bool,
-    ) -> Optional["StepInputStager"]:
-        """Open the window for this step, or return None to keep the native preparation.
-
-        The engine calls this only for the replay of an already captured CUDA
-        graph, so ``graph_replay`` is always True: the window closes right
-        before the replay, and an eager step, an initial capture or a warmup
-        would run the model before the staged writes are issued. Those steps
-        always prepare natively.
-
-        ``preparer`` is the ``DecoderInputPreparer`` whose writes the window
-        records. ``static_inputs`` are the static input tensors of the CUDA
-        graph the step replays.
-        """
-
-    def copy(
-        self, dst: "torch.Tensor", src: "torch.Tensor", clamp_min: Optional[int] = None
-    ) -> None:
-        """Stage ``dst.copy_(src, non_blocking=True)`` for a host ``src``."""
-
-    def zero(self, dst: "torch.Tensor") -> None:
-        """Stage ``dst.zero_()``."""
-
-    def clamp_min(self, dst: "torch.Tensor", lo: int) -> None:
-        """Stage ``dst.clamp_(min=lo)``."""
-
-    def block_offsets(
-        self,
-        host_table: "torch.Tensor",
-        dst: "torch.Tensor",
-        copy_index: Sequence[int],
-        index_scales: Sequence[int],
-        kv_offsets: Sequence[int],
-    ) -> None:
-        """Stage the KV cache manager's batch copy of block offsets to the device."""
-
-    def position_ids(self, dst: "torch.Tensor", host_position_ids: "torch.Tensor") -> None:
-        """Stage the position-id upload and the CUDA graph's static copy of it."""
-
-    def overlap_inputs(
-        self,
-        buffers: Any,
-        previous_batch_indices: Sequence[int],
-        previous_pos_indices: Sequence[int],
-        num_tokens: int,
-        num_draft_tokens: int,
-        tokens_per_gen_step: int,
-        num_extend_requests: int,
-        new_tokens: "torch.Tensor",
-        next_draft_tokens: "torch.Tensor",
-        new_tokens_lens: "torch.Tensor",
-    ) -> None:
-        """Stage the overlap scheduler's gathers of the previous step's outputs into ``buffers``."""
-
-    def scheduler_metadata(self, metadata: Any) -> bool:
-        """Stage the sparse-attention indexer's scheduler metadata.
-
-        Returns False, having staged nothing, when the caller must defer the
-        native computation instead.
-        """
-
-    def defer(self, fn: Callable[[], None]) -> None:
-        """Run ``fn`` (native device work that reads staged values) after the launch."""
-
-    def flush(self) -> None:
-        """Issue every write staged so far, before the window closes."""
-
-    def end(self, record_event: bool = True) -> None:
-        """Close the window and issue the staged preparation."""
-
-    def discard(self) -> None:
-        """Close the window after a failed preparation; the staged writes are dropped."""
-
-    def set_completion_event(self, event: "torch.cuda.Event") -> None:
-        """Event the caller recorded right after ``end(record_event=False)``."""
-
-    def static_inputs_tokens(self) -> int:
-        """Tokens of the CUDA graph's static inputs the last launch wrote (0: none)."""
-
-
-class SamplerStoreWriter(Protocol):
-    """Stores a speculative sampler step's outputs and copies the store to the host in one launch."""
-
-    def supported(
-        self,
-        store: Any,
-        new_tokens: "torch.Tensor",
-        next_new_tokens: "torch.Tensor",
-        new_tokens_lens: "torch.Tensor",
-        next_draft_tokens: "torch.Tensor",
-    ) -> bool:
-        """Whether this step's output tensors fit the writer; False keeps the native stores."""
-
-    def store_and_copy_out(
-        self,
-        store: Any,
-        slots: Sequence[int],
-        new_tokens: "torch.Tensor",
-        next_new_tokens: "torch.Tensor",
-        new_tokens_lens: "torch.Tensor",
-        next_draft_tokens: "torch.Tensor",
-        host_new_tokens: "torch.Tensor",
-        host_new_tokens_lens: "torch.Tensor",
-        host_next_draft_tokens: "torch.Tensor",
-        record_event: bool = True,
-    ) -> None:
-        """Scatter the outputs into the store rows ``slots`` and copy the store to the host tensors."""
-
-    def set_completion_event(self, event: "torch.cuda.Event") -> None:
-        """Event the caller recorded right after ``store_and_copy_out(record_event=False)``."""
+from typing import Any, Optional, Protocol
 
 
 class DecodePrologue(Protocol):
@@ -198,10 +54,6 @@ class RuntimeHooks:
     # Run the overlap executor loop with the execution stream current, so the
     # loop's own stream work and the forward are ordered on one stream.
     loop_on_execution_stream: bool = False
-    # Decode input preparation recorded and issued as one launch (see StepInputStager).
-    step_input_stager: Optional[StepInputStager] = None
-    # The speculative sampler's store update and host copy in one launch (see SamplerStoreWriter).
-    sampler_store_writer: Optional[SamplerStoreWriter] = None
     # The decode step's input prologue in one launch (see DecodePrologue).
     decode_prologue: Optional[DecodePrologue] = None
     # Actions after warmup, before the overlap loop and at shutdown (see ExecutorEvents).
diff --git a/tensorrt_llm/_torch/speculative/eagle3.py b/tensorrt_llm/_torch/speculative/eagle3.py
index 0328e29746cc..15993e5c26ba 100644
--- a/tensorrt_llm/_torch/speculative/eagle3.py
+++ b/tensorrt_llm/_torch/speculative/eagle3.py
@@ -348,13 +348,8 @@ def prepare(self):
                                      dtype=torch.int,
                                      device='cpu',
                                      pin_memory=prefer_pinned())
-        if self.step_input_stager is not None:
-            # Staged write; see the window banner in DecoderRunner._forward_decoder.
-            self.step_input_stager.copy(self.batch_indices_cuda[:num_seqs],
-                                        batch_indices)
-        else:
-            self.batch_indices_cuda[:num_seqs].copy_(batch_indices,
-                                                     non_blocking=True)
+        self.batch_indices_cuda[:num_seqs].copy_(batch_indices,
+                                                 non_blocking=True)
         # Snapshot the real (unpadded) scheduled token count before the decrement
         # below. maybe_capture_hidden_states runs during the target forward and
         # must bound the capture by this count; self.num_tokens is about to be
diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py
index a8bdaf4d4589..8edccc3e3f8a 100644
--- a/tensorrt_llm/_torch/speculative/interface.py
+++ b/tensorrt_llm/_torch/speculative/interface.py
@@ -19,7 +19,7 @@
 from contextlib import contextmanager
 from dataclasses import dataclass, field
 from enum import IntEnum
-from typing import TYPE_CHECKING, ClassVar, List, Optional, Protocol, Type
+from typing import TYPE_CHECKING, List, Optional, Protocol, Type
 
 import torch
 from packaging.version import Version
@@ -37,7 +37,6 @@
 if TYPE_CHECKING:
     from ..pyexecutor.guided_decoder import CapturableGuidedDecoder
     from ..pyexecutor.llm_request import LlmRequest
-    from ..pyexecutor.runtime_hooks import StepInputStager
 
 if IS_FLASHINFER_AVAILABLE:
     import flashinfer
@@ -460,8 +459,6 @@ class SpecMetadata:
     """
     Metadata for speculative decoding.
     """
-    # A step-input stager while it records this metadata's writes, else None.
-    step_input_stager: ClassVar[Optional["StepInputStager"]] = None
     # The max number of requests in a single batch.
     max_num_requests: int
     # The number of draft layers. (Also the number of draft tokens for the linear tree.)
@@ -884,11 +881,6 @@ def _populate_request_rng_state(
                                        device='cuda')
 
         def _upload(dst: torch.Tensor, values: list[int]) -> None:
-            if self.step_input_stager is not None:
-                # Staged write; see the window banner in DecoderRunner._forward_decoder.
-                self.step_input_stager.copy(
-                    dst[:len(values)], torch.tensor(values, dtype=torch.int64))
-                return
             dst[:len(values)].copy_(torch.tensor(values,
                                                  dtype=torch.int64,
                                                  pin_memory=prefer_pinned()),
diff --git a/tensorrt_llm/_torch/speculative/spec_sampler_base.py b/tensorrt_llm/_torch/speculative/spec_sampler_base.py
index dc40cca3a7c5..5fb5b185dc0b 100644
--- a/tensorrt_llm/_torch/speculative/spec_sampler_base.py
+++ b/tensorrt_llm/_torch/speculative/spec_sampler_base.py
@@ -27,11 +27,9 @@
 
 import torch
 
-from ..._utils import prefer_pinned
 from ...sampling_params import SamplingParams
 from ..pyexecutor.llm_request import LlmRequest, LlmRequestState, get_draft_token_length
 from ..pyexecutor.resource_manager import BaseResourceManager
-from ..pyexecutor.runtime_hooks import SamplerStoreWriter
 from ..pyexecutor.sampler import (
     DEFAULT_BEAM_IDX,
     AsyncWorkerMixin,
@@ -169,10 +167,6 @@ class SpecSampler(Sampler[SampleStateSpec], AsyncWorkerMixin):
 
     SampleState = SampleStateSpec
 
-    # The model's one-launch store writer, bound once by instantiate_sampler; None keeps
-    # the native stores and host copies.
-    store_writer: Optional[SamplerStoreWriter] = None
-
     def is_generation_model(self) -> bool:
         return True
 
@@ -448,11 +442,8 @@ def sample_async(
             for r in sampling_requests
         ]
 
-        store_writer = self.store_writer
-        slots = None
-        if store_writer is None:
-            slots = torch.as_tensor([r.py_seq_slot for r in sampling_requests], dtype=torch.long)
-            slots = slots.to(device="cuda", non_blocking=True)
+        slots = torch.as_tensor([r.py_seq_slot for r in sampling_requests], dtype=torch.long)
+        slots = slots.to(device="cuda", non_blocking=True)
 
         o_new_tokens = outputs["new_tokens"][num_skip : num_skip + num_sampling_requests]
         o_new_tokens_lens = outputs["new_tokens_lens"][num_skip : num_skip + num_sampling_requests]
@@ -523,28 +514,11 @@ def sample_async(
         elif o_next_new_tokens.shape[1] > next_new_tokens_width:
             o_next_new_tokens = o_next_new_tokens[:, :next_new_tokens_width]
 
-        # The writer stores the four outputs and copies the three store tensors
-        # only; a step that publishes verify_lens keeps the native path.
-        staged_store = (
-            store_writer is not None
-            and o_verify_lens is None
-            and prefer_pinned()
-            and not self._async_worker_active()
-            and store_writer.supported(
-                self.store, o_new_tokens, o_next_new_tokens, o_new_tokens_lens, o_next_draft_tokens
-            )
-        )
-        if not staged_store:
-            if slots is None:
-                slots = torch.as_tensor(
-                    [r.py_seq_slot for r in sampling_requests], dtype=torch.long
-                )
-                slots = slots.to(device="cuda", non_blocking=True)
-            # Use index_copy_ for efficient copying (slots are unique)
-            self.store.new_tokens.squeeze(-1).T.index_copy_(0, slots, o_new_tokens)
-            self.store.next_new_tokens.squeeze(-1).T.index_copy_(0, slots, o_next_new_tokens)
-            self.store.new_tokens_lens.index_copy_(0, slots, o_new_tokens_lens)
-            self.store.next_draft_tokens.index_copy_(0, slots, o_next_draft_tokens)
+        # Use index_copy_ for efficient copying (slots are unique)
+        self.store.new_tokens.squeeze(-1).T.index_copy_(0, slots, o_new_tokens)
+        self.store.next_new_tokens.squeeze(-1).T.index_copy_(0, slots, o_next_new_tokens)
+        self.store.new_tokens_lens.index_copy_(0, slots, o_new_tokens_lens)
+        self.store.next_draft_tokens.index_copy_(0, slots, o_next_draft_tokens)
 
         # Create sample state with async D2H copy
         device_tensors = SampleStateTensorsSpec(
@@ -553,43 +527,13 @@ def sample_async(
             next_draft_tokens=self.store.next_draft_tokens,
         )
 
-        if staged_store:
-            # staged_store implies prefer_pinned(), so these buffers are pinned.
-            host_tensors = SampleStateTensorsSpec(
-                new_tokens=torch.empty_like(
-                    self.store.new_tokens, device="cpu", pin_memory=prefer_pinned()
-                ),
-                new_tokens_lens=torch.empty_like(
-                    self.store.new_tokens_lens, device="cpu", pin_memory=prefer_pinned()
-                ),
-                next_draft_tokens=torch.empty_like(
-                    self.store.next_draft_tokens, device="cpu", pin_memory=prefer_pinned()
-                ),
-            )
-            store_writer.store_and_copy_out(
-                self.store,
-                [r.py_seq_slot for r in sampling_requests],
-                o_new_tokens,
-                o_next_new_tokens,
-                o_new_tokens_lens,
-                o_next_draft_tokens,
-                host_tensors.new_tokens,
-                host_tensors.new_tokens_lens,
-                host_tensors.next_draft_tokens,
-                record_event=False,
-            )
-        else:
-            host_tensors = SampleStateTensorsSpec(
-                new_tokens=self._copy_to_host(self.store.new_tokens),
-                new_tokens_lens=self._copy_to_host(self.store.new_tokens_lens),
-                next_draft_tokens=self._copy_to_host(self.store.next_draft_tokens),
-                verify_lens=(
-                    self._copy_to_host(o_verify_lens) if o_verify_lens is not None else None
-                ),
-            )
+        host_tensors = SampleStateTensorsSpec(
+            new_tokens=self._copy_to_host(self.store.new_tokens),
+            new_tokens_lens=self._copy_to_host(self.store.new_tokens_lens),
+            next_draft_tokens=self._copy_to_host(self.store.next_draft_tokens),
+            verify_lens=(self._copy_to_host(o_verify_lens) if o_verify_lens is not None else None),
+        )
         sampler_event = self._record_sampler_event()
-        if staged_store:
-            store_writer.set_completion_event(sampler_event.cuda_event)
 
         # Add dummy draft tokens to context requests for KV cache preparation
         for request in finished_context_requests:
diff --git a/tests/unittest/_torch/executor/test_runtime_hooks.py b/tests/unittest/_torch/executor/test_runtime_hooks.py
index 95745dd901db..f3a79f0e1180 100644
--- a/tests/unittest/_torch/executor/test_runtime_hooks.py
+++ b/tests/unittest/_torch/executor/test_runtime_hooks.py
@@ -3,11 +3,11 @@
 """How the engine reads the step-loop hooks a model declares (``runtime_hooks_of``).
 
 The engine's call sites run their original statements when a hook is off: every
-object a hook attaches to declares it off at class level, every call a call
-site makes on a hook is a method of the hook's protocol, every read of a hook in
-the engine files is decided by its all-off value toward the original statements
-(checked on the source), and call sites of each hook run on stand-ins issue
-their original work with the hook off and hand it to the hook when it is on.
+object that holds a hook declares it off at class level, every call a call site
+makes on a hook is a method of the hook's protocol, every read of a hook in the
+engine files is decided by its all-off value toward the original statements
+(checked on the source), and call sites of each hook run on stand-ins with the
+hook off and on.
 """
 
 import ast
@@ -25,8 +25,6 @@
     DEFAULT_RUNTIME_HOOKS,
     ExecutorEvents,
     RuntimeHooks,
-    SamplerStoreWriter,
-    StepInputStager,
     runtime_hooks_of,
 )
 
@@ -59,8 +57,6 @@ def test_a_model_without_hooks_gets_the_all_off_default(model):
     assert hooks is DEFAULT_RUNTIME_HOOKS
     assert hooks == RuntimeHooks()
     assert not hooks.loop_on_execution_stream
-    assert hooks.step_input_stager is None
-    assert hooks.sampler_store_writer is None
     assert hooks.decode_prologue is None
     assert hooks.executor_events is None
 
@@ -117,36 +113,14 @@ def test_hooks_are_immutable_once_declared():
         hooks.loop_on_execution_stream = True
 
 
-def test_engine_classes_declare_the_stager_off():
-    """Every object a hook attaches to declares the attribute off at class level."""
-    from tensorrt_llm._torch.attention.backends.sparse.dsa.metadata import (
-        DSAtrtllmAttentionMetadata,
-    )
-    from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata
-    from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import CUDAGraphRunner
+def test_engine_classes_declare_the_hooks_off():
+    """Every object that holds a hook declares it off at class level."""
     from tensorrt_llm._torch.pyexecutor.engine.runners.decoder.inputs import DecoderInputPreparer
     from tensorrt_llm._torch.pyexecutor.engine.runners.decoder.runner import DecoderRunner
-    from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2
     from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor
-    from tensorrt_llm._torch.speculative.eagle3 import Eagle3OneModelSpecMetadata
-    from tensorrt_llm._torch.speculative.interface import SpecMetadata
-    from tensorrt_llm._torch.speculative.spec_sampler_base import SpecSampler
-
-    for cls in (
-        TrtllmAttentionMetadata,
-        DSAtrtllmAttentionMetadata,
-        SpecMetadata,
-        Eagle3OneModelSpecMetadata,
-    ):
-        assert cls.step_input_stager is None, cls
-        # A class attribute, not a dataclass field: the constructors do not take it.
-        assert "step_input_stager" not in {f.name for f in dataclasses.fields(cls)}, cls
-    assert DecoderInputPreparer.step_input_stager is None
+
     assert DecoderInputPreparer.decode_prologue is None
     assert DecoderRunner.runtime_hooks is DEFAULT_RUNTIME_HOOKS
-    assert KVCacheManagerV2.step_input_stager is None
-    assert CUDAGraphRunner.staged_static_input_tokens == 0
-    assert SpecSampler.store_writer is None
     assert PyExecutor._execution_stream_loop_armed is False
     assert PyExecutor._loop_on_execution_stream is False
     assert PyExecutor._executor_events is None
@@ -155,26 +129,6 @@ def test_engine_classes_declare_the_stager_off():
 # Per hook protocol: the local names and attribute names a call site holds the hook in, and
 # the engine files with call sites.
 _HOOK_CONSUMERS = {
-    StepInputStager: (
-        ("stager", "stager_source", "step_input_stager"),
-        ("step_input_stager",),
-        (
-            "pyexecutor/engine/runners/decoder/runner.py",
-            "pyexecutor/engine/runners/decoder/inputs.py",
-            "pyexecutor/engine/runners/common.py",
-            "pyexecutor/kv_cache/kv_cache_manager_v2.py",
-            "attention/backends/trtllm.py",
-            "attention/backends/sparse/dsa/metadata.py",
-            "attention/backends/sparse/dsa/indexer.py",
-            "speculative/eagle3.py",
-            "speculative/interface.py",
-        ),
-    ),
-    SamplerStoreWriter: (
-        ("store_writer",),
-        ("store_writer",),
-        ("speculative/spec_sampler_base.py",),
-    ),
     ExecutorEvents: ((), ("_executor_events",), ("pyexecutor/py_executor.py",)),
 }
 
@@ -218,63 +172,6 @@ def test_every_hook_call_site_uses_its_protocol(protocol):
     assert used == methods
 
 
-class _StoreWriterModel(torch.nn.Module):
-    def __init__(self, writer) -> None:
-        super().__init__()
-        self.runtime_hooks = RuntimeHooks(sampler_store_writer=writer)
-
-
-def _instantiate_sampler(monkeypatch, model, sampler):
-    """``instantiate_sampler`` for a speculative engine whose spec decoder is ``sampler``."""
-    from tensorrt_llm._torch.pyexecutor import _util
-
-    monkeypatch.setattr(_util, "confidential_compute_enabled", lambda: False)
-    monkeypatch.setattr(_util, "resolve_max_num_sequences", lambda *args, **kwargs: 8)
-    monkeypatch.setattr(_util, "create_torch_sampler_args", lambda **kwargs: "sampler args")
-    monkeypatch.setattr(_util, "get_spec_decoder", lambda args, spec_config: sampler)
-    engine = types.SimpleNamespace(
-        max_seq_len=64,
-        spec_config=types.SimpleNamespace(
-            spec_dec_mode=types.SimpleNamespace(has_spec_decoder=lambda: True)
-        ),
-        model=model,
-    )
-    llm_args = types.SimpleNamespace(
-        sampler_force_async_worker=False,
-        disable_overlap_scheduler=False,
-        enable_speculative_beam_history_d2h=False,
-    )
-    return _util.instantiate_sampler(
-        engine,
-        llm_args,
-        None,
-        max_batch_size=8,
-        max_beam_width=1,
-        mm_encoder_only=False,
-        speculative_config=None,
-    )
-
-
-def test_instantiate_sampler_binds_no_writer_without_hooks(monkeypatch):
-    from tensorrt_llm._torch.speculative.spec_sampler_base import SpecSampler
-
-    sampler = SpecSampler.__new__(SpecSampler)
-
-    assert _instantiate_sampler(monkeypatch, _PlainModel(), sampler) is sampler
-    assert "store_writer" not in vars(sampler)
-    assert sampler.store_writer is None
-
-
-def test_instantiate_sampler_binds_the_declared_writer(monkeypatch):
-    from tensorrt_llm._torch.speculative.spec_sampler_base import SpecSampler
-
-    writer = object()
-    sampler = SpecSampler.__new__(SpecSampler)
-
-    assert _instantiate_sampler(monkeypatch, _StoreWriterModel(writer), sampler) is sampler
-    assert sampler.store_writer is writer
-
-
 def test_executor_events_protocol():
     assert _protocol_methods(ExecutorEvents) == {
         "after_warmup",
@@ -371,13 +268,6 @@ def test_overlap_loop_scope_stays_off_with_side_streams(monkeypatch, name, value
     assert recorder.entered == []
 
 
-def test_instantiate_sampler_leaves_other_spec_decoders_alone(monkeypatch):
-    sampler = types.SimpleNamespace()
-
-    assert _instantiate_sampler(monkeypatch, _StoreWriterModel(object()), sampler) is sampler
-    assert "store_writer" not in vars(sampler)
-
-
 # ---------------------------------------------------------------------------------------------
 # With every hook off, every call site runs its original statements.
 #
@@ -390,33 +280,20 @@ def test_instantiate_sampler_leaves_other_spec_decoders_alone(monkeypatch):
 # ---------------------------------------------------------------------------------------------
 
 _SEAM_CONSUMERS = (
-    "pyexecutor/engine/runners/decoder/runner.py",
     "pyexecutor/engine/runners/decoder/inputs.py",
-    "pyexecutor/cuda_graph_runner.py",
-    "pyexecutor/kv_cache/kv_cache_manager_v2.py",
     "pyexecutor/py_executor.py",
-    "pyexecutor/_util.py",
-    "attention/backends/trtllm.py",
-    "attention/backends/sparse/dsa/metadata.py",
-    "attention/backends/sparse/dsa/indexer.py",
-    "speculative/eagle3.py",
-    "speculative/interface.py",
-    "speculative/spec_sampler_base.py",
 )
 
 # Attributes that hold a hook, and their value when every hook is off.
 _SEAM_ATTRS = {
-    "step_input_stager": None,
-    "store_writer": None,
     "_executor_events": None,
     "_loop_on_execution_stream": False,
     "_execution_stream_loop_armed": False,
-    "staged_static_input_tokens": 0,
     "runtime_hooks": DEFAULT_RUNTIME_HOOKS,
     "decode_prologue": None,
 }
 # Locals the call sites hold a hook (or a decision on one) in.
-_SEAM_LOCALS = ("stager", "stager_source", "store_writer", "prologue", "staged", "staged_store")
+_SEAM_LOCALS = ("prologue",)
 _HOOK_FIELDS = {field.name for field in dataclasses.fields(RuntimeHooks)}
 
 
@@ -441,15 +318,10 @@ def report(self, node: ast.AST, what: str) -> None:
     # -- reads of a hook ------------------------------------------------------------------
     @staticmethod
     def _is_seam_call(node: ast.AST) -> bool:
-        if not (isinstance(node, ast.Call) and isinstance(node.func, ast.Name)):
-            return False
-        if node.func.id == "runtime_hooks_of":
-            return True
         return (
-            node.func.id == "getattr"
-            and len(node.args) >= 2
-            and isinstance(node.args[1], ast.Constant)
-            and node.args[1].value == "step_input_stager"
+            isinstance(node, ast.Call)
+            and isinstance(node.func, ast.Name)
+            and node.func.id == "runtime_hooks_of"
         )
 
     def reads_seam(self, node: ast.AST, env: dict) -> bool:
@@ -476,7 +348,7 @@ def value(self, node: ast.AST, env: dict):
                 return _SEAM_ATTRS[node.attr]
             return _UNKNOWN
         if self._is_seam_call(node):
-            return DEFAULT_RUNTIME_HOOKS if node.func.id == "runtime_hooks_of" else None
+            return DEFAULT_RUNTIME_HOOKS
         if isinstance(node, ast.UnaryOp) and isinstance(node.op, ast.Not):
             operand = self.value(node.operand, env)
             return _UNKNOWN if operand is _UNKNOWN else not operand
@@ -496,13 +368,6 @@ def value(self, node: ast.AST, env: dict):
             if isinstance(op, (ast.Is, ast.IsNot)) and left is not _UNKNOWN:
                 if isinstance(node.comparators[0], ast.Constant) and right is None:
                     return (left is None) == isinstance(op, ast.Is)
-            if (
-                isinstance(op, ast.Eq)
-                and isinstance(node.left, ast.Attribute)
-                and node.left.attr == "staged_static_input_tokens"
-            ):
-                # No tokens were staged; a graph replay has at least one token.
-                return False
         return _UNKNOWN
 
     # -- statements -------------------------------------------------------------------------
@@ -598,15 +463,15 @@ def test_every_seam_read_has_a_native_branch(rel):
 def test_the_seam_check_reports_undecided_reads():
     source = """
 def step(self, x):
-    stager = self.step_input_stager
-    if stager is not None:
-        stager.copy(x)
+    events = self._executor_events
+    if events is not None:
+        events.after_warmup()
     else:
         x.copy_()
-    stager.flush()
-    if x or stager:
+    events.at_shutdown()
+    if x or events:
         pass
-    self.step_input_stager = x
+    self._executor_events = x
 """
     walk = _AllOffWalk("sample.py")
     walk.block(ast.parse(source).body, {})
@@ -815,584 +680,3 @@ def test_executor_loop_cleanup_calls_at_shutdown_after_a_completed_loop(complete
     executor._executor_loop_cleanup()
     assert executor.is_shutdown
     assert events.calls == (["at_shutdown"] if completed else [])
-
-
-# -- the step-input stager ------------------------------------------------------------------
-
-
-class _RecordingStager:
-    def __init__(self) -> None:
-        self.calls = []
-
-    def __getattr__(self, name):
-        return lambda *args, **kwargs: self.calls.append((name, args, kwargs))
-
-
-def _block_offsets_manager(monkeypatch, log: list):
-    """A ``KVCacheManagerV2`` with only what ``copy_batch_block_offsets`` reads (one page table)."""
-    from tensorrt_llm._torch.pyexecutor.kv_cache import kv_cache_manager_v2
-    from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2
-
-    monkeypatch.setattr(
-        kv_cache_manager_v2,
-        "copy_batch_block_offsets_to_device",
-        lambda *args: log.append(("native", args)),
-    )
-    manager = KVCacheManagerV2.__new__(KVCacheManagerV2)
-    manager._publish_sparse_metadata = lambda: None
-    manager.max_copy_beam_width = 1
-    manager.kv_cache_type = None
-    manager._use_per_layer_page_tables = False
-    manager.index_mapper = types.SimpleNamespace(
-        get_copy_index=lambda request_ids, num_contexts, beam_width, cross: torch.tensor([1, 0])
-    )
-    manager.host_kv_cache_block_offsets = torch.zeros(2, 4, dtype=torch.int32)
-    manager.index_scales = torch.ones(1, dtype=torch.int32)
-    manager._index_scale_ints = [1]
-    manager.kv_offset = torch.zeros(1, dtype=torch.int32)
-    manager._stream = types.SimpleNamespace(cuda_stream=7)
-    return manager
-
-
-def test_kv_cache_block_offsets_copy_natively_without_a_stager(monkeypatch):
-    log = []
-    manager = _block_offsets_manager(monkeypatch, log)
-    dst = torch.empty(2, 4, dtype=torch.int32)
-
-    manager.copy_batch_block_offsets(dst, [3, 4], 1, 0, 2)
-    assert len(log) == 1
-    kind, args = log[0]
-    assert kind == "native"
-    assert args[0] is manager.host_kv_cache_block_offsets and args[1] is dst
-    assert args[2].tolist() == [1, 0]
-    assert args[3] is manager.index_scales and args[4] is manager.kv_offset and args[5] == 7
-
-
-def test_kv_cache_block_offsets_are_staged_with_a_stager(monkeypatch):
-    log = []
-    manager = _block_offsets_manager(monkeypatch, log)
-    manager.step_input_stager = _RecordingStager()
-    dst = torch.empty(2, 4, dtype=torch.int32)
-
-    manager.copy_batch_block_offsets(dst, [3, 4], 1, 0, 2)
-    assert log == []
-    ((name, args, _),) = manager.step_input_stager.calls
-    assert name == "block_offsets"
-    assert args[0] is manager.host_kv_cache_block_offsets and args[1] is dst
-    assert args[2:] == ([1, 0], [1], [0])
-
-
-def _slot_mapping_inputs(tokens: int):
-    metadata = types.SimpleNamespace(
-        host_indexer_k_cache_block_offsets=torch.tensor([[0, 1]], dtype=torch.int32),
-        host_slot_mapping_fp8=torch.zeros(tokens, dtype=torch.int64),
-        host_slot_mapping_scale=torch.zeros(tokens, dtype=torch.int64),
-        slot_mapping_fp8=torch.full((tokens,), -1, dtype=torch.int64),
-        slot_mapping_scale=torch.full((tokens,), -1, dtype=torch.int64),
-    )
-    params = types.SimpleNamespace(
-        batch_size=1,
-        tokens_per_block=4,
-        head_dim=8,
-        new_kv_tokens=torch.tensor([tokens]),
-        data_bytes_per_token=8,
-        cached_kv_tokens=torch.tensor([1]),
-        quant_block_size=8,
-    )
-    return metadata, params
-
-
-def test_indexer_slot_mappings_copy_natively_without_a_stager():
-    from tensorrt_llm._torch.attention.backends.sparse.dsa.indexer import Indexer
-
-    metadata, params = _slot_mapping_inputs(3)
-
-    Indexer.recompute_slot_mappings(metadata, params)
-    assert torch.equal(metadata.slot_mapping_fp8, metadata.host_slot_mapping_fp8)
-    assert torch.equal(metadata.slot_mapping_scale, metadata.host_slot_mapping_scale)
-
-
-def test_indexer_slot_mappings_are_staged_with_a_stager():
-    from tensorrt_llm._torch.attention.backends.sparse.dsa.indexer import Indexer
-
-    metadata, params = _slot_mapping_inputs(3)
-    metadata.step_input_stager = _RecordingStager()
-
-    Indexer.recompute_slot_mappings(metadata, params)
-    assert [name for name, _, _ in metadata.step_input_stager.calls] == ["copy", "copy"]
-    assert bool((metadata.slot_mapping_fp8 == -1).all())
-    assert bool((metadata.slot_mapping_scale == -1).all())
-
-
-# -- the step-input window in the decoder runner -------------------------------------------------
-
-
-class _ReferenceStager:
-    """A step-input stager that accepts every step it is offered and logs the window's lifecycle."""
-
-    def __init__(self, log: list, accept: bool = True, tokens: int = 3) -> None:
-        self.log, self.accept, self.tokens = log, accept, tokens
-
-    def begin_if_eligible(
-        self,
-        preparer,
-        requests,
-        attn_metadata,
-        spec_metadata,
-        new_tensors_device,
-        *,
-        graph_replay,
-        static_inputs,
-        is_dummy,
-        enable_spec_decode,
-        promoted,
-    ):
-        self.log.append(("begin", graph_replay, static_inputs))
-        return self if self.accept else None
-
-    def end(self, record_event: bool = True) -> None:
-        self.log.append(("end", record_event))
-
-    def discard(self) -> None:
-        self.log.append("discard")
-
-    def set_completion_event(self, event) -> None:
-        self.log.append(("completion_event", event))
-
-    def static_inputs_tokens(self) -> int:
-        return self.tokens
-
-
-class _WindowGraphs:
-    """The ``DecoderCudaGraphs`` calls ``_forward_decoder`` makes, with one selected graph key.
-
-    ``run`` follows ``DecoderCudaGraphs._run``: ``before_replay`` runs only right before the
-    replay of an already captured graph, and the replay consumes the staged static inputs.
-    """
-
-    def __init__(self, log: list, key, *, needs_capture=False, warmup_only=False, fail=None):
-        self.log, self.key, self.fail = log, key, fail
-        self.cuda_graph_runner = types.SimpleNamespace(
-            needs_capture=lambda key: needs_capture,
-            is_warmup_only=warmup_only,
-            shared_static_tensors={"input_ids": torch.zeros(4, dtype=torch.int64)},
-            staged_static_input_tokens=0,
-        )
-
-    def attention_dp_execution_batch(self, requests):
-        return requests
-
-    def promote_final_contexts(self, requests, kv_cache_manager, **kwargs):
-        return requests, frozenset()
-
-    def pad_batch(self, requests, resource_manager, runtime_draft_len):
-        return contextlib.nullcontext(requests)
-
-    def select(self, requests, resource_manager, **kwargs):
-        return types.SimpleNamespace(
-            key=self.key,
-            use_lora_graph=False,
-            attn_metadata=kwargs["attn_metadata"],
-            spec_metadata=None,
-        )
-
-    def captures_model_body(self, key) -> bool:
-        return False
-
-    def run(self, key, inputs, run_model, postprocess_inputs, *, before_replay=None, **kwargs):
-        runner = self.cuda_graph_runner
-        if self.fail == "before_replay":
-            raise RuntimeError("graph run failed")
-        if key is None:
-            self.log.append("eager")
-        elif runner.needs_capture(key) or runner.is_warmup_only:
-            self.log.append("capture")
-        else:
-            if before_replay is not None:
-                before_replay()
-            self.log.append(("replay", runner.staged_static_input_tokens))
-            runner.staged_static_input_tokens = 0
-            if self.fail == "replay":
-                raise RuntimeError("replay failed")
-        return "outputs"
-
-
-class _WindowPreparer:
-    """The ``DecoderInputPreparer`` calls ``_forward_decoder`` makes."""
-
-    def __init__(self, log: list, fail: bool = False) -> None:
-        self.log, self.fail = log, fail
-        self.input_copy_event = "input-copy-event"
-
-    def prepare_inputs(self, *args, **kwargs):
-        self.log.append("prepare")
-        if self.fail:
-            raise RuntimeError("preparation failed")
-        return types.SimpleNamespace(kwargs={}, runtime_draft_len=0, gather_ids=None)
-
-    def record_input_copy(self) -> None:
-        self.log.append("record_input_copy")
-
-    def postprocess_inputs(self, *args, **kwargs) -> None:
-        pass
-
-
-def _window_runner(monkeypatch, graphs, preparer, stager):
-    """A ``DecoderRunner`` with only what ``_forward_decoder`` reads (no speculation, no CED)."""
-    from tensorrt_llm._torch.pyexecutor import ced_replay
-    from tensorrt_llm._torch.pyexecutor.engine.runners.decoder import runner as runner_module
-    from tensorrt_llm._torch.pyexecutor.engine.runners.decoder.runner import DecoderRunner
-
-    monkeypatch.setattr(runner_module, "get_draft_kv_cache_manager", lambda *args: None)
-    monkeypatch.setattr(ced_replay, "uses_private_decoder_cache", lambda *args: False)
-    attn_metadata = types.SimpleNamespace()
-    runner = DecoderRunner.__new__(DecoderRunner)
-    runner._config = types.SimpleNamespace(spec_config=None, use_mrope=False, is_spec_decode=False)
-    runner._metadata = types.SimpleNamespace(
-        setup_attention=lambda *args: attn_metadata,
-        attn_metadata=attn_metadata,
-        spec_metadata=None,
-    )
-    runner._graphs = graphs
-    runner._inputs = preparer
-    runner._stage_in_graph_sampling = None
-    runner._cuda_graph_metadata_prepare_scope = lambda *args, **kwargs: contextlib.nullcontext()
-    runner._execute_logit_post_processors = lambda *args: None
-    runner.forward_pass_callable = None
-    runner.moe_load_balancer = None
-    runner.model = types.SimpleNamespace()
-    runner.kv_cache_manager_key = "kv_cache_manager"
-    runner.runtime_hooks = RuntimeHooks(step_input_stager=stager)
-    return runner
-
-
-def _decode_step(runner):
-    forward_inputs = types.SimpleNamespace(
-        batch=types.SimpleNamespace(num_context_requests=0),
-        new_tensors_device=None,
-        cache_indirection_buffer=None,
-        gather_context_logits=False,
-        enable_spec_decode=False,
-        runtime_draft_len=0,
-    )
-    resource_manager = types.SimpleNamespace(get_resource_manager=lambda key: object())
-    return runner._forward_decoder(forward_inputs, resource_manager, is_dummy=False)
-
-
-def _names(log: list) -> list:
-    return [entry if isinstance(entry, str) else entry[0] for entry in log]
-
-
-_GRAPH_KEY = (1, 0, False)
-_STAGED_STEP = ["begin", "prepare", "end", "record_input_copy", "completion_event", "replay"]
-
-
-def test_a_replay_issues_the_staged_preparation_before_the_replay(monkeypatch):
-    log = []
-    graphs = _WindowGraphs(log, _GRAPH_KEY)
-    runner = _window_runner(monkeypatch, graphs, _WindowPreparer(log), _ReferenceStager(log))
-
-    assert _decode_step(runner) == "outputs"
-    assert log == [
-        ("begin", True, graphs.cuda_graph_runner.shared_static_tensors),
-        "prepare",
-        ("end", False),
-        "record_input_copy",
-        ("completion_event", "input-copy-event"),
-        ("replay", 3),
-    ]
-
-
-@pytest.mark.parametrize(
-    "key, needs_capture, warmup_only, ran",
-    [
-        (None, False, False, "eager"),
-        (_GRAPH_KEY, True, False, "capture"),
-        (_GRAPH_KEY, False, True, "capture"),
-    ],
-    ids=["eager", "initial_capture", "warmup_only"],
-)
-def test_a_step_that_is_not_a_replay_never_opens_a_window(
-    monkeypatch, key, needs_capture, warmup_only, ran
-):
-    log = []
-    graphs = _WindowGraphs(log, key, needs_capture=needs_capture, warmup_only=warmup_only)
-    runner = _window_runner(monkeypatch, graphs, _WindowPreparer(log), _ReferenceStager(log))
-
-    _decode_step(runner)
-    assert log == ["prepare", "record_input_copy", ran]
-    assert graphs.cuda_graph_runner.staged_static_input_tokens == 0
-
-
-def test_a_declined_replay_prepares_natively(monkeypatch):
-    log = []
-    graphs = _WindowGraphs(log, _GRAPH_KEY)
-    stager = _ReferenceStager(log, accept=False)
-    runner = _window_runner(monkeypatch, graphs, _WindowPreparer(log), stager)
-
-    _decode_step(runner)
-    assert _names(log) == ["begin", "prepare", "record_input_copy", "replay"]
-    assert log[-1] == ("replay", 0)
-
-
-def test_a_failed_preparation_discards_the_window(monkeypatch):
-    log = []
-    graphs = _WindowGraphs(log, _GRAPH_KEY)
-    runner = _window_runner(
-        monkeypatch, graphs, _WindowPreparer(log, fail=True), _ReferenceStager(log)
-    )
-
-    with pytest.raises(RuntimeError, match="preparation failed"):
-        _decode_step(runner)
-    assert _names(log) == ["begin", "prepare", "discard"]
-
-
-@pytest.mark.parametrize(
-    "fail, expected",
-    [
-        ("before_replay", ["begin", "prepare", "discard"]),
-        ("replay", _STAGED_STEP),
-    ],
-    ids=["window_open", "window_closed"],
-)
-def test_a_failed_run_discards_only_a_window_still_open(monkeypatch, fail, expected):
-    log = []
-    graphs = _WindowGraphs(log, _GRAPH_KEY, fail=fail)
-    runner = _window_runner(monkeypatch, graphs, _WindowPreparer(log), _ReferenceStager(log))
-
-    with pytest.raises(RuntimeError):
-        _decode_step(runner)
-    assert _names(log) == expected
-
-
-def test_replay_eager_replay_opens_one_window_per_replay(monkeypatch):
-    log = []
-    graphs = _WindowGraphs(log, _GRAPH_KEY)
-    runner = _window_runner(monkeypatch, graphs, _WindowPreparer(log), _ReferenceStager(log))
-
-    _decode_step(runner)
-    graphs.key = None
-    _decode_step(runner)
-    graphs.key = _GRAPH_KEY
-    _decode_step(runner)
-    assert _names(log) == [
-        *_STAGED_STEP,
-        "prepare",
-        "record_input_copy",
-        "eager",
-        *_STAGED_STEP,
-    ]
-    assert [entry for entry in log if entry[0] == "replay"] == [("replay", 3), ("replay", 3)]
-
-
-def _replay_graphs(monkeypatch, log: list, *, needs_capture: bool, warmup_only: bool):
-    """A ``DecoderCudaGraphs`` with only what ``_run`` reads, logging capture and replay."""
-    from tensorrt_llm._torch.pyexecutor.engine.runners.decoder import graphs as graphs_module
-    from tensorrt_llm._torch.pyexecutor.engine.runners.decoder.graphs import DecoderCudaGraphs
-
-    monkeypatch.setattr(graphs_module, "prepare_attn_metadata_for_draft_replay", lambda *args: None)
-    monkeypatch.setattr(
-        graphs_module, "restore_attn_metadata_after_draft_replay", lambda *args: None
-    )
-    monkeypatch.setattr(
-        graphs_module, "MoeLoadBalancerIterContext", lambda balancer: contextlib.nullcontext()
-    )
-    graphs = DecoderCudaGraphs.__new__(DecoderCudaGraphs)
-    graphs.breakable_cuda_graph_runner = None
-    graphs._torch_compile_bypass = contextlib.nullcontext
-    graphs.cuda_graph_runner = types.SimpleNamespace(
-        needs_capture=lambda key: needs_capture,
-        is_warmup_only=warmup_only,
-        capture=lambda key, fn, inputs, **kwargs: log.append("capture"),
-        replay=lambda key, inputs, **kwargs: log.append("replay"),
-    )
-    return graphs
-
-
-@pytest.mark.parametrize(
-    "key, needs_capture, warmup_only, expected",
-    [
-        (_GRAPH_KEY, False, False, ["before_replay", "replay"]),
-        (_GRAPH_KEY, True, False, ["capture", "replay"]),
-        (_GRAPH_KEY, True, True, ["capture"]),
-        (None, False, False, ["eager"]),
-    ],
-    ids=["replay", "capture_then_replay", "warmup_only", "eager"],
-)
-def test_before_replay_runs_only_before_the_replay_of_a_captured_graph(
-    monkeypatch, key, needs_capture, warmup_only, expected
-):
-    log = []
-    graphs = _replay_graphs(monkeypatch, log, needs_capture=needs_capture, warmup_only=warmup_only)
-
-    graphs._run(
-        key,
-        {"input_ids": torch.zeros(2, dtype=torch.int64)},
-        lambda inputs: log.append("eager"),
-        lambda inputs: None,
-        attn_metadata=None,
-        draft_kv_cache_manager=None,
-        moe_load_balancer=None,
-        enable_spec_decode=True,
-        gather_context_logits=False,
-        position_id_offsets=None,
-        before_replay=lambda: log.append("before_replay"),
-    )
-    assert log == expected
-
-
-# -- the speculative sampler's store writer -----------------------------------------------------
-
-
-class _ReferenceStoreWriter:
-    """Stores like the native ``index_copy_`` path, then copies the store to the host tensors."""
-
-    def __init__(self, supported: bool = True) -> None:
-        self._supported = supported
-        self.calls = []
-
-    def supported(self, store, new_tokens, next_new_tokens, new_tokens_lens, next_draft_tokens):
-        self.calls.append("supported")
-        return self._supported
-
-    def store_and_copy_out(
-        self,
-        store,
-        slots,
-        new_tokens,
-        next_new_tokens,
-        new_tokens_lens,
-        next_draft_tokens,
-        host_new_tokens,
-        host_new_tokens_lens,
-        host_next_draft_tokens,
-        record_event=True,
-    ) -> None:
-        self.calls.append(("store_and_copy_out", list(slots), record_event))
-        index = torch.tensor(slots, dtype=torch.long)
-        store.new_tokens.squeeze(-1).T.index_copy_(0, index, new_tokens)
-        store.next_new_tokens.squeeze(-1).T.index_copy_(0, index, next_new_tokens)
-        store.new_tokens_lens.index_copy_(0, index, new_tokens_lens)
-        store.next_draft_tokens.index_copy_(0, index, next_draft_tokens)
-        host_new_tokens.copy_(store.new_tokens)
-        host_new_tokens_lens.copy_(store.new_tokens_lens)
-        host_next_draft_tokens.copy_(store.next_draft_tokens)
-
-    def set_completion_event(self, event) -> None:
-        self.calls.append(("completion_event", event))
-
-
-def _cpu_sampling(monkeypatch, *, pinned: bool) -> None:
-    """Run ``SpecSampler.sample_async`` on CPU: device transfers stay on the host."""
-    from tensorrt_llm._torch.speculative import spec_sampler_base
-
-    original_to = torch.Tensor.to
-    original_empty_like = torch.empty_like
-
-    def to_cpu(tensor, *args, **kwargs):
-        if kwargs.get("device") == "cuda":
-            kwargs = {**kwargs, "device": "cpu"}
-        return original_to(tensor, *args, **kwargs)
-
-    def empty_like(tensor, *args, pin_memory=False, **kwargs):
-        # Pinning needs a CUDA context; the writer's host buffers are plain host tensors here.
-        return original_empty_like(tensor, *args, **kwargs)
-
-    monkeypatch.setattr(torch.Tensor, "to", to_cpu)
-    monkeypatch.setattr(torch, "empty_like", empty_like)
-    monkeypatch.setattr(spec_sampler_base, "prefer_pinned", lambda: pinned)
-    monkeypatch.setattr(spec_sampler_base, "handle_stop_criteria", lambda *args, **kwargs: False)
-
-
-def _store_sampler(writer, *, async_worker: bool = False):
-    """A ``SpecSampler`` with a six-slot store, draft length 3 and the given store writer."""
-    from tensorrt_llm._torch.speculative.spec_sampler_base import SpecSampler
-
-    sampler = SpecSampler.__new__(SpecSampler)
-    sampler.draft_len = 3
-    sampler.max_accepted_path_len = 4
-    sampler.max_seq_len = 128
-    sampler.store = types.SimpleNamespace(
-        new_tokens=torch.zeros(4, 6, 1, dtype=torch.int32),
-        next_new_tokens=torch.zeros(4, 6, 1, dtype=torch.int32),
-        next_draft_tokens=torch.zeros(6, 3, dtype=torch.int32),
-        new_tokens_lens=torch.zeros(6, dtype=torch.int32),
-    )
-    sampler.store_writer = writer
-    sampler._async_worker_active = lambda: async_worker
-    sampler._copy_to_host = lambda tensor: tensor.clone()
-    sampler._record_sampler_event = lambda: types.SimpleNamespace(cuda_event="sampler-event")
-    return sampler
-
-
-def _sample_generation_step(sampler, *, verify_lens: bool = False):
-    """Three generation requests in nonconsecutive slots 4, 1 and 3."""
-    from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestState
-
-    def request(request_id: int, slot: int):
-        return types.SimpleNamespace(
-            py_request_id=request_id,
-            py_seq_slot=slot,
-            state=LlmRequestState.GENERATION_IN_PROGRESS,
-            py_draft_tokens=[10, 11, 12],
-            py_draft_tokens_effective_len=None,
-            py_verify_len=None,
-            py_decoding_iter=0,
-        )
-
-    scheduled = types.SimpleNamespace(
-        context_requests_chunking=[],
-        context_requests_last_chunk=[],
-        generation_requests=[request(7, 4), request(8, 1), request(9, 3)],
-    )
-    outputs = dict(
-        new_tokens=torch.arange(12, dtype=torch.int32).reshape(3, 4) + 100,
-        new_tokens_lens=torch.tensor([1, 3, 2], dtype=torch.int32),
-        next_draft_tokens=torch.arange(9, dtype=torch.int32).reshape(3, 3) + 200,
-        next_new_tokens=torch.arange(12, dtype=torch.int32).reshape(3, 4) + 300,
-    )
-    if verify_lens:
-        outputs["verify_lens"] = torch.tensor([2, 3, 1], dtype=torch.int32)
-        outputs["verify_lens_in_output_order"] = True
-    return sampler.sample_async(scheduled, outputs, [])
-
-
-def _assert_same_stores(sampler, state, native, native_state) -> None:
-    for name in ("new_tokens", "next_new_tokens", "next_draft_tokens", "new_tokens_lens"):
-        assert torch.equal(getattr(sampler.store, name), getattr(native.store, name)), name
-    for name in ("new_tokens", "new_tokens_lens", "next_draft_tokens"):
-        assert torch.equal(getattr(state.host, name), getattr(native_state.host, name)), name
-
-
-def test_the_store_writer_matches_the_native_stores_on_nonconsecutive_slots(monkeypatch):
-    _cpu_sampling(monkeypatch, pinned=True)
-    native = _store_sampler(None)
-    writer = _ReferenceStoreWriter()
-    sampler = _store_sampler(writer)
-
-    native_state = _sample_generation_step(native)
-    state = _sample_generation_step(sampler)
-    assert writer.calls == [
-        "supported",
-        ("store_and_copy_out", [4, 1, 3], False),
-        ("completion_event", "sampler-event"),
-    ]
-    assert native.store.new_tokens_lens.tolist() == [0, 3, 0, 2, 1, 0]
-    _assert_same_stores(sampler, state, native, native_state)
-    assert state.host.verify_lens is None
-    assert state.device.new_tokens is sampler.store.next_new_tokens
-
-
-@pytest.mark.parametrize("case", ["declined", "verify_lens", "async_worker", "not_pinned"])
-def test_the_store_writer_falls_back_to_the_native_stores(monkeypatch, case):
-    _cpu_sampling(monkeypatch, pinned=case != "not_pinned")
-    native = _store_sampler(None)
-    writer = _ReferenceStoreWriter(supported=case != "declined")
-    sampler = _store_sampler(writer, async_worker=case == "async_worker")
-
-    native_state = _sample_generation_step(native, verify_lens=case == "verify_lens")
-    state = _sample_generation_step(sampler, verify_lens=case == "verify_lens")
-    assert writer.calls == (["supported"] if case == "declined" else [])
-    _assert_same_stores(sampler, state, native, native_state)
-    if case == "verify_lens":
-        assert state.host.verify_lens.tolist() == [2, 3, 1]

From 6b322651209aac397cd6c172f8c9c21d721c8b43 Mon Sep 17 00:00:00 2001
From: Xiao Wang <24860335+xwang233@users.noreply.github.com>
Date: Tue, 27 Oct 2026 05:41:53 -0700
Subject: [PATCH 5/5] [TRTLLM-17096][fix] runtime hooks: read the model from
 DecoderRunner._model

main renamed the decoder runner's model attribute to _model, so reading
self.model failed at runner construction for every model.

Update two tests to main's code: the GVR prior test calls
sparse_attn_indexer on a SimpleNamespace, which now needs the indexer
seams bound; the forward-step stub needs route_capture,
enable_spec_decode and runtime_draft_len.

Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com>
---
 .../_torch/pyexecutor/engine/runners/decoder/runner.py       | 2 +-
 .../_torch/attention/sparse/dsa/test_indexer_gvr_prior.py    | 3 +++
 tests/unittest/_torch/executor/test_runtime_hooks.py         | 5 ++++-
 3 files changed, 8 insertions(+), 2 deletions(-)

diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py
index 07041d032d56..0f7586cf2c27 100644
--- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py
+++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py
@@ -211,7 +211,7 @@ def __init__(
         self._dummy_requests = self._make_dummy_requests()
         self._init_decoder_state()
         self._inputs = self._make_input_preparer(mm_item_scheduler)
-        self.runtime_hooks = runtime_hooks_of(self.model)
+        self.runtime_hooks = runtime_hooks_of(self._model)
         self._inputs.decode_prologue = self.runtime_hooks.decode_prologue
         self._graphs = self._make_cuda_graphs()
         self._initialize_moe_workspace_reclaimer()
diff --git a/tests/unittest/_torch/attention/sparse/dsa/test_indexer_gvr_prior.py b/tests/unittest/_torch/attention/sparse/dsa/test_indexer_gvr_prior.py
index 623a33a5e536..04aa0c8da075 100644
--- a/tests/unittest/_torch/attention/sparse/dsa/test_indexer_gvr_prior.py
+++ b/tests/unittest/_torch/attention/sparse/dsa/test_indexer_gvr_prior.py
@@ -177,6 +177,9 @@ def prefill(
         aux_stream=None,
     )
     indexer.forward_prepared = MethodType(Indexer.forward_prepared, indexer)
+    indexer._copy_reused_topk = MethodType(Indexer._copy_reused_topk, indexer)
+    indexer._fill_decode_topk = MethodType(Indexer._fill_decode_topk, indexer)
+    indexer._store_mtp_shared_rows = MethodType(Indexer._store_mtp_shared_rows, indexer)
     with (
         patch.object(
             top_k, "update_gvr_prior_from_prefill", wraps=top_k.update_gvr_prior_from_prefill
diff --git a/tests/unittest/_torch/executor/test_runtime_hooks.py b/tests/unittest/_torch/executor/test_runtime_hooks.py
index f3a79f0e1180..d2459e6804de 100644
--- a/tests/unittest/_torch/executor/test_runtime_hooks.py
+++ b/tests/unittest/_torch/executor/test_runtime_hooks.py
@@ -593,11 +593,14 @@ def stream(s):
 
     executor = PyExecutor.__new__(PyExecutor)
     executor.iter_counter = 0
+    executor.enable_spec_decode = False
+    executor.runtime_draft_len = 0
     executor.execution_stream = _Stream("execution", log)
     executor.resource_manager = object()
     executor.sampler = types.SimpleNamespace(get_cache_indirection=lambda: None)
     executor.model_engine = types.SimpleNamespace(
-        forward=lambda *args, **kwargs: log.append("forward") or "outputs"
+        forward=lambda *args, **kwargs: log.append("forward") or "outputs",
+        route_capture=None,
     )
     executor._iter_adp_dummy_ctx_tokens = 0
     executor._iter_adp_dummy_gen_tokens = 0