diff --git a/tensorrt_llm/_torch/attention_backend/flashinfer.py b/tensorrt_llm/_torch/attention_backend/flashinfer.py
index d49ab084a128..b4d49af40534 100644
--- a/tensorrt_llm/_torch/attention_backend/flashinfer.py
+++ b/tensorrt_llm/_torch/attention_backend/flashinfer.py
@@ -20,7 +20,7 @@
 import weakref
 from dataclasses import dataclass, field
 from itertools import chain
-from typing import Any, Dict, NewType, Optional, TypeAlias, cast
+from typing import Any, Dict, List, NewType, Optional, TypeAlias, cast
 
 if sys.version_info[:2] >= (3, 12):
     from typing import override
@@ -89,6 +89,29 @@ def _slice_paged_kv_cache_heads(
     return paged_kv_cache[tuple(index)]
 
 
+def _get_page_table_num_blocks(kv_cache_manager, request_ids,
+                               logical_num_blocks: List[int],
+                               num_contexts: int) -> List[int]:
+    """Keep context rows logical and expose every reserved generation page."""
+    if hasattr(kv_cache_manager, "kv_cache_map"):
+        reserved_num_blocks = [
+            kv_cache_manager.kv_cache_map[req_id].num_blocks
+            for req_id in request_ids
+        ]
+    else:
+        # V1 keeps allocation state behind the C++ manager. Asking for the untrimmed page tables is
+        # its public equivalent of V2's num_blocks.
+        reserved_num_blocks = [
+            len(block_ids) for block_ids in
+            kv_cache_manager.get_batch_cache_indices(request_ids)
+        ]
+    return list(logical_num_blocks[:num_contexts]) + [
+        max(logical, reserved)
+        for logical, reserved in zip(logical_num_blocks[num_contexts:],
+                                     reserved_num_blocks[num_contexts:])
+    ]
+
+
 def _append_paged_kv_cache(
     append_key: torch.Tensor,
     append_value: torch.Tensor,
@@ -509,9 +532,14 @@ def _do_plan_mla_decode(self, plan_params: MLAPlanParams) -> None:
         qo_indptr = self._qo_indptr[num_ctx:num_ctx + num_gen +
                                     1] - self._qo_indptr[num_ctx]
 
-        num_pages_per_seq = kv_indptr[1:] - kv_indptr[:-1]
-        kv_len_arr = (num_pages_per_seq -
-                      1) * plan_params.page_size + kv_last_page
+        if self._uses_full_generation_page_table:
+            # Reservation-width page tables deliberately expose unused pages to overlap decode.
+            # MLA still needs the device-logical lengths, including any live acceptance rewind.
+            kv_len_arr = self._logical_kv_lens[num_ctx:num_ctx + num_gen]
+        else:
+            num_pages_per_seq = kv_indptr[1:] - kv_indptr[:-1]
+            kv_len_arr = (num_pages_per_seq -
+                          1) * plan_params.page_size + kv_last_page
 
         self._mla_decode_wrapper.plan(
             qo_indptr,
@@ -767,6 +795,98 @@ def update_shared_kv_draft_lengths(
             num_accepted_tokens[num_contexts:num_seqs])
         self._update_draft_kv_lengths()
 
+    def apply_spec_decode_kv_lens_offsets(
+        self,
+        offsets: torch.Tensor,
+        num_generations: int,
+        tokens_per_generation: int,
+        *,
+        num_chunked_contexts: int = 0,
+        restore: bool = False,
+    ) -> None:
+        """Apply overlap-scheduler KV rewinds to FlashInfer runtime state.
+
+        `prepare()` builds the target page-table metadata from the maximum speculative step width.
+        With the overlap scheduler, the actual number of accepted tokens is still device-resident,
+        so generation or extend-context rows can be shorter than that host-side upper bound.
+        Keep every device consumer of the runtime KV length in lockstep without synchronizing back
+        to the host.
+
+        Args:
+            offsets: Signed per-request corrections ordered like the trailing extend-context rows
+                when ``num_chunked_contexts`` is nonzero, otherwise like the generation rows.
+                Applying the rewind adds these values to the runtime lengths and positions.
+            num_generations: Number of generation rows in the metadata.
+            tokens_per_generation: Number of contiguous query positions that receive each request's
+                offset.
+            num_chunked_contexts: Number of trailing context rows representing extend requests.
+            restore: Whether to subtract the same corrections. Calling once with ``restore=False``
+                and once with ``restore=True`` exactly reverses the mutation.
+        """
+        if self._is_shared_kv_draft_view or self._is_separate_kv_draft_view:
+            raise RuntimeError(
+                "Speculative KV offsets must be applied to target metadata")
+        if num_chunked_contexts == 0 and num_generations == 0:
+            return
+
+        direction = -1 if restore else 1
+        num_contexts = self.num_contexts
+        if num_chunked_contexts > 0:
+            # Linear-tree speculative decoding packs extend requests at the tail of the context
+            # partition. Their overlap offsets remain live even though no rows are classified as
+            # generation, so adjust both the trailing request rows and their trailing query tokens.
+            row_slice = slice(num_contexts - num_chunked_contexts, num_contexts)
+            runtime_offsets = offsets[:num_chunked_contexts]
+            num_runtime_tokens = num_chunked_contexts * tokens_per_generation
+            token_slice = slice(self.num_ctx_tokens - num_runtime_tokens,
+                                self.num_ctx_tokens)
+        else:
+            row_slice = slice(num_contexts, num_contexts + num_generations)
+            runtime_offsets = offsets[:num_generations]
+            num_runtime_tokens = num_generations * tokens_per_generation
+            token_slice = slice(self.num_ctx_tokens,
+                                self.num_ctx_tokens + num_runtime_tokens)
+
+        self._cached_token_lens[row_slice].add_(runtime_offsets,
+                                                alpha=direction)
+        if self._uses_full_generation_page_table:
+            self._logical_kv_lens[row_slice].add_(runtime_offsets,
+                                                  alpha=direction)
+
+        # Keep the host-preallocated page structure intact. In particular, `last_page_len` must stay
+        # consistent with `paged_kv_indptr`: changing only one of them at a page boundary would
+        # describe a different KV length. Appends use the corrected explicit positions below, while
+        # trtllm-gen decode uses its corrected `_kv_lens_buffer`.
+
+        token_offsets = runtime_offsets.repeat_interleave(tokens_per_generation)
+        self._positions[token_slice].add_(token_offsets, alpha=direction)
+
+        # trtllm-gen decode wrappers own a separate persistent logical-length buffer, including
+        # under CUDA graphs. Publish the exact lengths rather than incrementing the existing values:
+        # mixed prefill / decode batches may have re-planned a wrapper from the structural
+        # page-table upper bound after overlap preprocessing.
+        for wrappers in self._plan_params_to_wrappers.values():
+            self._publish_decode_wrapper_kv_lens(wrappers.decode_wrapper)
+
+    def _publish_decode_wrapper_kv_lens(self, decode_wrapper) -> None:
+        """Publish device-logical lengths after a trtllm-gen decode plan."""
+        kv_lens_buffer = getattr(decode_wrapper, "_kv_lens_buffer", None)
+        if kv_lens_buffer is None or self.num_generations == 0:
+            return
+        start = self.num_contexts
+        end = start + self.num_generations
+        if self._is_shared_kv_draft_view:
+            # The external assistant is Q-only and does not append its query to the target KV cache.
+            # Its cached length is already the full accepted target prefix.
+            kv_lens_buffer[:self.num_generations].copy_(
+                self._cached_token_lens[start:end])
+        else:
+            torch.add(
+                self._cached_token_lens[start:end],
+                self.seq_lens_kv_cuda[start:end],
+                out=kv_lens_buffer[:self.num_generations],
+            )
+
     def _prepare_full_draft_page_table(self) -> None:
         """Expose every allocated draft page and use device KV lengths."""
         if self._uses_full_draft_page_table:
@@ -812,19 +932,16 @@ def _update_draft_kv_lengths(self) -> None:
                         self.page_size,
                         out=self._paged_kv_last_page_len[:num_seqs])
         self._paged_kv_last_page_len[:num_seqs].add_(1)
-        if self._is_shared_kv_draft_view:
-            return
+        if not self._is_shared_kv_draft_view:
+            self._cached_token_lens[:num_seqs].sub_(1)
+            self._positions[:num_seqs].copy_(self._cached_token_lens[:num_seqs])
+            torch.arange(num_seqs + 1, out=self._qo_indptr[:num_seqs + 1])
+            torch.arange(num_seqs, out=self._batch_indices[:num_seqs])
 
-        self._cached_token_lens[:num_seqs].sub_(1)
-        self._positions[:num_seqs].copy_(self._cached_token_lens[:num_seqs])
-        torch.arange(num_seqs + 1,
-                     dtype=torch.int32,
-                     device=self._qo_indptr.device,
-                     out=self._qo_indptr[:num_seqs + 1])
-        torch.arange(num_seqs,
-                     dtype=torch.int32,
-                     device=self._batch_indices.device,
-                     out=self._batch_indices[:num_seqs])
+        # CUDA-graph draft wrappers retain the plan's private logical-length buffer. Keep it
+        # synchronized with every accepted-prefix update.
+        for wrappers in self._plan_params_to_wrappers.values():
+            self._publish_decode_wrapper_kv_lens(wrappers.decode_wrapper)
 
     def update_for_spec_dec(self) -> None:
         if not self._is_separate_kv_draft_view:
@@ -886,6 +1003,13 @@ def _post_init_with_buffers(self, buffers) -> None:
         self._cached_token_lens = torch.empty((self.max_num_requests, ),
                                               dtype=torch.int,
                                               device='cuda')
+        self._logical_kv_lens = self.get_empty(
+            buffers,
+            (self.max_num_requests, ),
+            dtype=torch.int,
+            cache_name="_logical_kv_lens",
+            capture_graph=capture_graph,
+        )
         self._draft_kv_runtime_lens = self.get_empty(
             buffers,
             (self.max_num_requests, ),
@@ -909,6 +1033,7 @@ def _post_init_with_buffers(self, buffers) -> None:
         self._host_pool_indices: Dict[int, torch.Tensor] = {}
         self._host_paged_kv_indices: Optional[torch.Tensor] = None
         self._host_paged_kv_indptr_decode: Optional[torch.Tensor] = None
+        self._uses_full_generation_page_table = False
         self._max_num_blocks_per_seq = 0
 
         # VSWA (Variable Sliding Window Attention): models with per-layer
@@ -1396,9 +1521,25 @@ def _to_int32_tensor(arr: np.ndarray) -> torch.Tensor:
         # so a device-side computation would force a sync per step.
         kv_lens_host = np.asarray(num_cached_tokens_per_seq,
                                   dtype=np.int64) + self.seq_lens_kv.numpy()
-        num_blocks = (kv_lens_host + self.page_size - 1) // self.page_size
+        logical_num_blocks = ((kv_lens_host + self.page_size - 1) //
+                              self.page_size)
+        num_blocks = logical_num_blocks
+        use_full_generation_page_table = bool(
+            getattr(self.kv_cache_params, "use_full_generation_page_table",
+                    False))
+        self._uses_full_generation_page_table = use_full_generation_page_table
+        if use_full_generation_page_table:
+            assert self.request_ids is not None
+            num_blocks = np.asarray(
+                _get_page_table_num_blocks(
+                    self.kv_cache_manager,
+                    self.request_ids,
+                    logical_num_blocks.tolist(),
+                    self.num_contexts,
+                ),
+                dtype=np.int64,
+            )
         self.num_blocks = num_blocks.tolist()
-
         assert self.request_ids is not None
 
         # start and end indices of each sequence in the ragged key and value
@@ -1464,9 +1605,9 @@ def _to_int32_tensor(arr: np.ndarray) -> torch.Tensor:
             self._vswa_active_pool_id = primary_pool_id
 
         # number of tokens in the last cache block used by each sequence,
-        # derived on the host so no GPU arithmetic or sync is needed.
+        # derived from the logical (not reservation-width) page count.
         paged_kv_last_page_len = _to_int32_tensor(kv_lens_host -
-                                                  (num_blocks - 1) *
+                                                  (logical_num_blocks - 1) *
                                                   self.page_size)
         self._paged_kv_last_page_len[:paged_kv_last_page_len.size(0)].copy_(
             paged_kv_last_page_len, non_blocking=True)
@@ -1509,11 +1650,20 @@ def _to_int32_tensor(arr: np.ndarray) -> torch.Tensor:
 
         # For cross attention, num_tokens is 0 during decode, and we don't need to update kv cache.
         if self.num_tokens > 0:
+            if use_full_generation_page_table:
+                logical_kv_lens = _to_int32_tensor(kv_lens_host)
+                self._logical_kv_lens[:logical_kv_lens.numel()].copy_(
+                    logical_kv_lens, non_blocking=True)
+                position_kv_lens = self._logical_kv_lens[:self.num_seqs]
+            else:
+                position_kv_lens = flashinfer.get_seq_lens(
+                    self.paged_kv_indptr,
+                    self.paged_kv_last_page_len,
+                    self.page_size,
+                )
             batch_indices, positions = flashinfer.get_batch_indices_positions(
                 self.kv_indptr,
-                flashinfer.get_seq_lens(self.paged_kv_indptr,
-                                        self.paged_kv_last_page_len,
-                                        self.page_size),
+                position_kv_lens,
                 self.num_tokens,
             )
             self._batch_indices[:batch_indices.size(0)].copy_(batch_indices,
@@ -1858,6 +2008,11 @@ def prefill_plan():
 
         def decode_plan():
             assert decode_wrapper is not None
+            if (self._uses_full_generation_page_table
+                    and decode_wrapper._backend != "trtllm-gen"):
+                raise ValueError(
+                    "Reservation-width FlashInfer page tables require the "
+                    "trtllm-gen decode backend's independent KV lengths.")
             # Host int32 indptr (retained by prepare, which always runs
             # before plans): flashinfer moves it to the device itself, and
             # its indptr.cpu()/get_seq_lens calls stay free of D2H syncs.
@@ -1885,6 +2040,10 @@ def decode_plan():
                 o_data_type=o_dtype,
                 block_tables=block_tables,
             )
+            # plan() rebuilds trtllm-gen's private KV-length buffer from the structural page table.
+            # In a mixed overlap batch this plan can run after `_preprocess_inputs()` has applied
+            # the device acceptance rewind, so republish the corrected logical lengths immediately.
+            self._publish_decode_wrapper_kv_lens(decode_wrapper)
 
         # Must sync after append_paged_kv_cache and before plan.
         torch.cuda.current_stream().synchronize()
diff --git a/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py b/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py
index 1ed345a1914f..e783796f758e 100644
--- a/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py
+++ b/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py
@@ -1163,6 +1163,8 @@ def create_autodeploy_executor(
 
     This is the entrypoint API to the _autodeploy backend.
     """
+    spec_config = ad_config.speculative_config
+
     # initialize process groups
     world_size = mpi_world_size()
     rank = mpi_rank()
@@ -1193,8 +1195,6 @@ def create_autodeploy_executor(
         ad_config=ad_config, dist_config=dc, mapping=dist_mapping, dist=dist
     )
 
-    spec_config = ad_config.speculative_config
-
     if spec_config is not None and ad_config.guided_decoding_backend is not None:
         raise ValueError(
             "Guided decoding is not currently supported for speculative decoding in AutoDeploy."
@@ -1336,6 +1336,7 @@ def create_autodeploy_executor(
         max_beam_width=ad_config.max_beam_width,
         max_draft_len=max_draft_len,
         max_total_draft_tokens=max_total_draft_tokens,
+        max_seq_len=engine.cache_seq_interface.info.max_seq_len,
         guided_decoder=guided_decoder,
         kv_cache_transceiver=kv_cache_transceiver,
         resource_governor_queue=resource_governor_queue,
diff --git a/tensorrt_llm/_torch/metadata.py b/tensorrt_llm/_torch/metadata.py
index fd076cc4cab0..943a4eebf922 100644
--- a/tensorrt_llm/_torch/metadata.py
+++ b/tensorrt_llm/_torch/metadata.py
@@ -30,6 +30,11 @@ class KVCacheParams:
     # The number of extra kv for draft tokens
     num_extra_kv_tokens: Optional[int] = 0
 
+    # Stage every reserved generation page when device-resident speculative
+    # state can advance the runtime KV length past the host logical snapshot.
+    # The attention backend must carry the logical length independently.
+    use_full_generation_page_table: bool = False
+
 
 class CacheType(Enum):
     # Linear KV cache stores all the cached tokens of a sequence in a single page.
diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py
index 0b624cb6362d..afc8fdb8d108 100644
--- a/tensorrt_llm/_torch/pyexecutor/model_engine.py
+++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py
@@ -3354,6 +3354,17 @@ def get_max_num_sequences(self) -> int:
         num_batches = self.mapping.pp_size
         return num_batches * self.batch_size
 
+    def _should_use_full_generation_page_table(
+            self, spec_config: Optional[DecodingBaseConfig],
+            attn_metadata: AttentionMetadata) -> bool:
+        """Return whether overlap decode needs every reserved generation page."""
+        if spec_config is None:
+            return False
+        return (self.enable_spec_decode and not self._disable_overlap_scheduler
+                and spec_config._use_shared_kv_cache
+                # At time of writing, this is only the case for FlashInfer.
+                and hasattr(attn_metadata, 'apply_spec_decode_kv_lens_offsets'))
+
     def _preprocess_inputs(self, inputs: Dict[str, Any]):
         """
         Make some changes to the device inputs and avoid blocking the async data transfer
@@ -3405,6 +3416,15 @@ def _preprocess_inputs(self, inputs: Dict[str, Any]):
                                 previous_kv_lens_offsets_cuda[:num_gen_requests]
                             )
                     inputs['attn_metadata'].on_update_kv_lens()
+                elif hasattr(inputs['attn_metadata'],
+                             'apply_spec_decode_kv_lens_offsets'):
+                    inputs['attn_metadata'].apply_spec_decode_kv_lens_offsets(
+                        self.previous_kv_lens_offsets_cuda,
+                        num_gen_requests,
+                        self.get_runtime_tokens_per_gen_step(
+                            self.runtime_draft_len),
+                        num_chunked_contexts=num_chunked_ctx_requests,
+                    )
 
         if self.guided_decoder is not None:
             self.guided_decoder.token_event.record()
@@ -3452,6 +3472,16 @@ def _postprocess_inputs(self, inputs: Dict[str, Any]):
                                 self.
                                 previous_kv_lens_offsets_cuda[:num_gen_requests]
                             )
+                elif hasattr(inputs['attn_metadata'],
+                             'apply_spec_decode_kv_lens_offsets'):
+                    inputs['attn_metadata'].apply_spec_decode_kv_lens_offsets(
+                        self.previous_kv_lens_offsets_cuda,
+                        num_gen_requests,
+                        self.get_runtime_tokens_per_gen_step(
+                            self.runtime_draft_len),
+                        num_chunked_contexts=num_chunked_ctx_requests,
+                        restore=True,
+                    )
 
     def _get_all_rank_num_tokens(self, attn_metadata: AttentionMetadata):
         if self.enable_attention_dp:
@@ -4161,7 +4191,10 @@ def _prepare_incremental_update_metadata(
         attn_metadata.kv_cache_params = KVCacheParams(
             use_cache=True,
             num_cached_tokens_per_seq=num_cached_tokens_per_seq,
-            num_extra_kv_tokens=get_num_extra_kv_tokens(spec_config))
+            num_extra_kv_tokens=get_num_extra_kv_tokens(spec_config),
+            use_full_generation_page_table=(
+                self._should_use_full_generation_page_table(
+                    spec_config, attn_metadata)))
         attn_metadata.kv_cache_manager = kv_cache_manager
         attn_metadata.prepare()
 
@@ -5057,9 +5090,11 @@ def append_cross_attention_state(request: LlmRequest,
             assert spec_config.spec_dec_mode.support_overlap_scheduler(
             ), f"{spec_config.decoding_type} does not support overlap scheduler"
 
-        # For tree decoding, runtime_draft_len should match total tree
-        # tokens (not tree depth).  py_executor resets it every iteration.
-        if spec_config is not None and not spec_config.is_linear_tree:
+        # For active tree decoding, runtime_draft_len should match total tree
+        # tokens (not tree depth). Preserve an explicit zero selected by the
+        # executor for this iteration.
+        if (spec_config is not None and not spec_config.is_linear_tree
+                and self.runtime_draft_len != 0):
             self.runtime_draft_len = self.max_total_draft_tokens
 
         # will contain previous batch indices of generation requests
@@ -5071,7 +5106,7 @@ def append_cross_attention_state(request: LlmRequest,
         for request in extend_requests:
             is_promoted_context = (request.py_request_id
                                    in promoted_context_request_ids)
-            if getattr(request, "py_needs_onehot_draft_probs", False):
+            if request.py_needs_onehot_draft_probs:
                 if request.py_seq_slot is not None:
                     padding_gen_slots.append(request.py_seq_slot)
                 request.py_needs_onehot_draft_probs = False  # consume once
@@ -5150,11 +5185,11 @@ def append_cross_attention_state(request: LlmRequest,
                 previous_pos_indices.extend([previous_batch_idx] *
                                             runtime_tokens_per_gen_step)
 
+                cached_token_num = (past_seen_token_num +
+                                    runtime_tokens_per_gen_step)
                 num_cached_tokens_per_seq.append(
-                    past_seen_token_num + runtime_tokens_per_gen_step -
-                    request.py_num_compressed_tokens)
-                request.cached_tokens = (past_seen_token_num +
-                                         runtime_tokens_per_gen_step)
+                    cached_token_num - request.py_num_compressed_tokens)
+                request.cached_tokens = cached_token_num
                 if self.enable_spec_decode and spec_config.spec_dec_mode.extend_ctx(
                         self.attn_backend) and spec_config.is_linear_tree:
                     prompt_lengths.append(runtime_tokens_per_gen_step)
@@ -5804,7 +5839,10 @@ def previous_seq_slots_device():
         attn_metadata.kv_cache_params = KVCacheParams(
             use_cache=True,
             num_cached_tokens_per_seq=num_cached_tokens_per_seq,
-            num_extra_kv_tokens=get_num_extra_kv_tokens(spec_config))
+            num_extra_kv_tokens=get_num_extra_kv_tokens(spec_config),
+            use_full_generation_page_table=(
+                self._should_use_full_generation_page_table(
+                    spec_config, attn_metadata)))
         attn_metadata.kv_cache_manager = kv_cache_manager
 
         if hasattr(self.model.model_config.pretrained_config, 'chunk_size'):
@@ -7110,9 +7148,8 @@ def forward(self,
 
         graph_requests = scheduled_requests
         promoted_context_request_ids: frozenset[int] = frozenset()
-        # Non-linear tree input preparation expands runtime_draft_len to the
-        # total tree width after graph selection. Only linear-tree zero-draft
-        # iterations can therefore safely reuse a zero-draft graph.
+        # Keep zero-draft graph promotion conservative for non-linear trees;
+        # their metadata and capture shapes are based on the configured tree.
         can_promote_spec_decode = (not self.enable_spec_decode
                                    or (not self.is_draft_model
                                        and self.runtime_draft_len == 0
diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py
index bcd1363138d7..cf709bc28afb 100644
--- a/tensorrt_llm/_torch/pyexecutor/py_executor.py
+++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py
@@ -3370,75 +3370,115 @@ def _handle_dynamic_draft_len(self,
         if not hasattr(self.model_engine, 'max_draft_len'):
             return
 
+        spec_config = self.model_engine.spec_config
+        dynamic_draft_len_enabled = (
+            spec_config is not None
+            and spec_config.draft_len_schedule is not None
+            and spec_config.spec_dec_mode.support_dynamic_draft_len())
         if self.speculation_permanently_disabled:
-            for request in scheduled_batch.generation_requests:
-                request.py_draft_tokens = []
-            self.model_engine.runtime_draft_len = 0
-            return
-
-        if (self.model_engine.spec_config is not None
-                and self.model_engine.spec_config.draft_len_schedule is not None
-                and self.model_engine.spec_config.spec_dec_mode.
-                support_dynamic_draft_len()):
+            runtime_draft_len = 0
+        elif dynamic_draft_len_enabled:
             from tensorrt_llm._torch.speculative.utils import \
                 get_draft_len_for_batch_size
 
-            spec_dec_mode = self.model_engine.spec_config.spec_dec_mode
-
-            # 1. Resolve runtime draft length from schedule
             runtime_draft_len = get_draft_len_for_batch_size(
-                self.model_engine.spec_config.draft_len_schedule,
-                scheduled_batch.batch_size, self.model_engine.max_draft_len)
-            # 2. Pad or truncate draft tokens to the resolved length
-            DRAFT_BUFFER_PAD = 0  # Buffer sentinel, not PARD mask_token_id.
-            rejection_on = getattr(self.model_engine.spec_config,
-                                   "use_rejection_sampling", False)
-            for request in scheduled_batch.generation_requests:
-                current_num_draft_tokens = len(request.py_draft_tokens)
-                # One-model rejection: a gen request entering with 0 real draft
-                # tokens produced no draft-prob scatter for its slot last iter,
-                # so next iter's rejection kernel would read a stale draft_probs
-                # row. Mark it (pre-pad signal) so _prepare_tp_inputs writes a
-                # one-hot placeholder row after spec_metadata.prepare().
-                request.py_needs_onehot_draft_probs = (
-                    rejection_on and current_num_draft_tokens == 0)
-                if spec_dec_mode.is_pard():
-                    # special case: PARD carries 2K-1 draft tokens per request
-                    runtime_draft_token_buffer_width = (
-                        self.model_engine.spec_config.
-                        get_runtime_tokens_per_gen_step(runtime_draft_len) - 1)
-                    current_runtime_draft_len = (
-                        current_num_draft_tokens +
-                        1) // 2 if current_num_draft_tokens > 0 else 0
-                    real_draft_tokens = request.py_draft_tokens[:min(
-                        current_runtime_draft_len, runtime_draft_len)]
-                    real_draft_tokens.extend(
-                        [DRAFT_BUFFER_PAD] *
-                        (runtime_draft_len - len(real_draft_tokens)))
-                    request.py_draft_tokens = real_draft_tokens + [
-                        DRAFT_BUFFER_PAD
-                    ] * (runtime_draft_token_buffer_width -
-                         len(real_draft_tokens))
-                else:
-                    if current_num_draft_tokens < runtime_draft_len:
-                        padding_needed = (runtime_draft_len -
-                                          current_num_draft_tokens)
-                        request.py_draft_tokens.extend([DRAFT_BUFFER_PAD] *
-                                                       padding_needed)
-                    elif current_num_draft_tokens > runtime_draft_len:
-                        request.py_draft_tokens = request.py_draft_tokens[:
-                                                                          runtime_draft_len]
-
-            self.model_engine.runtime_draft_len = runtime_draft_len
+                spec_config.draft_len_schedule, scheduled_batch.batch_size,
+                self.model_engine.max_draft_len)
         else:
             # Linear-tree modes (incl. PARD) use logical K; tree decoding
             # (e.g. EAGLE3 dynamic tree) uses total tree tokens. Same
             # selection as _prepare_tp_inputs and _get_graphs_to_capture.
-            spec_config = self.model_engine.spec_config
-            self.model_engine.runtime_draft_len = (
-                self.model_engine.max_draft_len
-                if spec_config is not None and spec_config.is_linear_tree else
-                self.model_engine.max_total_draft_tokens)
+            runtime_draft_len = (self.model_engine.max_draft_len
+                                 if spec_config is not None
+                                 and spec_config.is_linear_tree else
+                                 self.model_engine.max_total_draft_tokens)
+
+        needs_zero_draft = self._one_model_mtp_batch_needs_zero_draft(
+            scheduled_batch, runtime_draft_len)
+        if (spec_config is not None
+                and spec_config.spec_dec_mode.is_mtp_eagle_one_model()
+                and self.enable_attention_dp):
+            needs_zero_draft = any(self.dist.tp_allgather(needs_zero_draft))
+
+        if needs_zero_draft or runtime_draft_len == 0:
+            # The target input width is batch-wide, so a committed zero must
+            # clear drafting inputs for every generation request.
+            for request in scheduled_batch.generation_requests:
+                request.py_draft_tokens = []
+            self.model_engine.runtime_draft_len = 0
+            return
+
+        self.model_engine.runtime_draft_len = runtime_draft_len
+        if not dynamic_draft_len_enabled:
+            return
+
+        draft_buffer_pad = 0  # Buffer sentinel, not PARD mask_token_id.
+        rejection_on = spec_config.use_rejection_sampling
+        spec_dec_mode = spec_config.spec_dec_mode
+        for request in scheduled_batch.generation_requests:
+            current_num_draft_tokens = len(request.py_draft_tokens)
+            # Preserve a pre-schedule zero-proposal signal across placeholder
+            # insertion until _prepare_tp_inputs one-hots the stale row.
+            request.py_needs_onehot_draft_probs |= (rejection_on
+                                                    and current_num_draft_tokens
+                                                    == 0)
+            if spec_dec_mode.is_pard():
+                # special case: PARD carries 2K-1 draft tokens per request
+                runtime_draft_token_buffer_width = (
+                    spec_config.get_runtime_tokens_per_gen_step(
+                        runtime_draft_len) - 1)
+                current_runtime_draft_len = ((current_num_draft_tokens + 1) //
+                                             2 if current_num_draft_tokens > 0
+                                             else 0)
+                real_draft_tokens = request.py_draft_tokens[:min(
+                    current_runtime_draft_len, runtime_draft_len)]
+                real_draft_tokens.extend(
+                    [draft_buffer_pad] *
+                    (runtime_draft_len - len(real_draft_tokens)))
+                request.py_draft_tokens = real_draft_tokens + [
+                    draft_buffer_pad
+                ] * (runtime_draft_token_buffer_width - len(real_draft_tokens))
+            elif current_num_draft_tokens < runtime_draft_len:
+                padding_needed = runtime_draft_len - current_num_draft_tokens
+                request.py_draft_tokens.extend([draft_buffer_pad] *
+                                               padding_needed)
+            elif current_num_draft_tokens > runtime_draft_len:
+                request.py_draft_tokens = request.py_draft_tokens[:
+                                                                  runtime_draft_len]
+
+    def _one_model_mtp_batch_needs_zero_draft(
+            self, scheduled_batch: ScheduledRequests,
+            runtime_draft_len: int) -> bool:
+        """Return whether drafting could produce an out-of-range position."""
+        spec_config = self.model_engine.spec_config
+        if (spec_config is None
+                or not spec_config.spec_dec_mode.is_mtp_eagle_one_model()):
+            return False
+        if runtime_draft_len == 0:
+            return True
+
+        # With overlap, the host request has not incorporated the previous iteration's accepted
+        # tokens yet. `_preprocess_inputs` adds that count to every target position on device.
+        # Use the maximum possible count because reading the exact value here would incur a sync.
+        max_pending_tokens = (0 if self.disable_overlap_scheduler else
+                              self.model_engine.max_draft_len + 1)
+        target_position_width = spec_config.get_runtime_tokens_per_gen_step(
+            runtime_draft_len)
+        # The one-model drafter can consume positions beyond the target verification span.
+        # A shared-KV assistant runs every Q-only draft step one position after the last accepted
+        # target token. The regular MTP-Eagle loop advances that position between its K draft
+        # forwards, reaching K - 1 positions beyond the target span.
+        draft_position_lookahead = (1 if getattr(
+            spec_config, '_use_shared_kv_cache', False) else max(
+                runtime_draft_len - 1, 0))
+
+        for request in scheduled_batch.generation_requests:
+            max_draft_position = (request.max_beam_num_tokens - 1 +
+                                  max_pending_tokens + target_position_width -
+                                  1 + draft_position_lookahead)
+            if max_draft_position >= self.max_seq_len:
+                return True
+        return False
 
     @nvtx_range("_can_queue")
     def _can_queue(self, scheduled_batch):
@@ -3844,11 +3884,24 @@ def _prepare_and_schedule_batch(self):
             # two-model normalization so scheduling reserves the correct token budget.
             # model_engine is guarded first so partially-constructed executors in unit tests
             # (which may not set model_engine) do not raise AttributeError.
+            rejection_on = (
+                self.model_engine.spec_config is not None
+                and self.model_engine.spec_config.use_rejection_sampling)
             for request in self.active_requests:
                 if request.state not in (
                         LlmRequestState.GENERATION_IN_PROGRESS,
                         LlmRequestState.DISAGG_GENERATION_INIT):
                     continue
+                # Only fill in a placeholder when the Python-side list is empty (e.g. a
+                # DISAGG_GENERATION_INIT request, which never gets a draft snapshot). Overwriting it
+                # unconditionally would clobber the real draft tokens the one-model spec sampler
+                # wrote at the end of the previous iteration - which the overlap-disabled path
+                # reads back in `_prepare_tp_inputs`. Capture the independent zero-proposal
+                # signal before inserting the scheduling placeholder so rejection sampling can
+                # one-hot the otherwise-stale draft-probability row.
+                if not request.py_draft_tokens:
+                    request.py_needs_onehot_draft_probs |= rejection_on
+                    request.py_draft_tokens = [0] * self.max_total_draft_tokens
                 request.draft_tokens = [0] * self.max_total_draft_tokens
 
         scheduled_batch, scheduler_fitting_disagg_gen_init_requests, num_fitting_reqs = self._schedule(
diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
index 680d607efeb5..1901ee1d66c4 100644
--- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
+++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
@@ -58,6 +58,28 @@
     for sm_version in _MLA_CHUNKED_PREFILL_SUPPORTED_SM_VERSIONS)
 
 
+def _disable_unsupported_speculative_overlap_scheduler(
+        llm_args: TorchLlmArgs,
+        spec_config: Optional[SpeculativeConfig]) -> None:
+    if llm_args.disable_overlap_scheduler or spec_config is None:
+        return
+
+    attn_backend = llm_args.attn_backend.upper()
+    spec_dec_mode = spec_config.spec_dec_mode
+    reason = None
+    if not spec_dec_mode.support_overlap_scheduler():
+        reason = f"speculation mode {spec_dec_mode.name}"
+    elif attn_backend == "VANILLA":
+        reason = "VANILLA attention lacks dynamic speculative KV lengths"
+    elif (attn_backend == "FLASHINFER" and not spec_dec_mode.use_one_engine()):
+        reason = ("FLASHINFER extend-context attention lacks dynamic "
+                  "speculative KV lengths")
+
+    if reason is not None:
+        logger.warning(f"Disable overlap scheduler for {reason}")
+        llm_args.disable_overlap_scheduler = True
+
+
 class _ExecutorMemoryMonitor:
     """Currently this focuses on tracking memory usage and related errors."""
 
@@ -455,14 +477,10 @@ def create_py_executor(
         from tensorrt_llm._torch.speculative import suggest_spec_config
         spec_config = suggest_spec_config(max_batch_size)
 
-    if not llm_args.disable_overlap_scheduler and spec_config is not None:
-        if not spec_config.spec_dec_mode.support_overlap_scheduler():
-            logger.warning(
-                f"Disable overlap scheduler for speculation mode {spec_config.spec_dec_mode.name}"
-            )
-            llm_args.disable_overlap_scheduler = True
+    _disable_unsupported_speculative_overlap_scheduler(llm_args, spec_config)
 
-    if (spec_config is not None and llm_args.attn_backend == "FLASHINFER"
+    if (spec_config is not None
+            and llm_args.attn_backend.upper() == "FLASHINFER"
             and spec_config.spec_dec_mode.use_one_engine()
             and not spec_config._use_shared_kv_cache):
         raise ValueError(
diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py
index 827bd5e534b7..c2e35908cef9 100644
--- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py
+++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py
@@ -164,6 +164,12 @@ def free_resources(self, request: LlmRequest):
     def shutdown(self):
         pass
 
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.shutdown()
+
 
 def get_pp_layers(
     num_layers: int,
diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py
index 2597174fb6d8..b2c6564b31aa 100644
--- a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py
+++ b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py
@@ -3715,6 +3715,14 @@ def _process_requests(
         )
         can_use_stable_greedy_path = (
             bool(generation_requests)
+            # A one-model MTP target request may temporarily have zero draft tokens when
+            # drafting is disabled near the sequence limit. It is still a speculative target
+            # and must use the device-side finish-state path. In contrast, requests marked
+            # `py_is_draft=True` belong to a separate draft model, whose sampler calls always
+            # produce exactly one token per request.
+            and (
+                self.max_tokens == 1 or all(request.py_is_draft for request in generation_requests)
+            )
             and self.max_beam_width == 1
             and scheduled_requests.num_context_requests == 0
             and len(generation_requests) <= raw_logits_cuda.shape[0]
diff --git a/tests/unittest/_torch/attention/test_flashinfer_attention.py b/tests/unittest/_torch/attention/test_flashinfer_attention.py
index fe00e3d69d07..104a8b4d31e2 100644
--- a/tests/unittest/_torch/attention/test_flashinfer_attention.py
+++ b/tests/unittest/_torch/attention/test_flashinfer_attention.py
@@ -1,3 +1,17 @@
+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 import random
 import unittest
 from collections import defaultdict
@@ -13,14 +27,17 @@
                                                    FlashInferAttentionMetadata)
 from tensorrt_llm._torch.attention_backend import \
     flashinfer as flashinfer_backend
-from tensorrt_llm._torch.attention_backend.flashinfer import PlanParams
+from tensorrt_llm._torch.attention_backend.flashinfer import (
+    FlashInferWrappers, MLAPlanParams, PlanParams)
 from tensorrt_llm._torch.attention_backend.interface import \
     PredefinedAttentionMask
 from tensorrt_llm._torch.metadata import KVCacheParams
+from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import KVCacheManagerV2
 from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager
 from tensorrt_llm._utils import prefer_pinned
 from tensorrt_llm.bindings.executor import KvCacheConfig
 from tensorrt_llm.functional import AttentionMaskType
+from tensorrt_llm.llmapi.llm_args import KvCacheConfig as LlmKvCacheConfig
 from tensorrt_llm.mapping import Mapping
 
 
@@ -64,31 +81,281 @@ class CUDAGraphTestScenario:
     dtype: torch.dtype
 
 
+@dataclass
+class FakeDecodeWrapper:
+    _kv_lens_buffer: torch.Tensor
+
+
+def _create_kv_cache_manager(max_tokens: int = 256,
+                             max_seq_len: int = 64) -> KVCacheManager:
+    return KVCacheManager(
+        KvCacheConfig(max_tokens=max_tokens),
+        tensorrt_llm.bindings.internal.batch_manager.CacheType.SELF,
+        num_layers=1,
+        num_kv_heads=1,
+        head_dim=128,
+        tokens_per_block=32,
+        max_seq_len=max_seq_len,
+        max_batch_size=2,
+        mapping=Mapping(world_size=1, tp_size=1, rank=0),
+        dtype=tensorrt_llm.bindings.DataType.BF16,
+    )
+
+
+def _create_kv_cache_manager_v2(max_tokens: int,
+                                max_seq_len: int) -> KVCacheManagerV2:
+    return KVCacheManagerV2(
+        LlmKvCacheConfig(max_tokens=max_tokens,
+                         max_util_for_resume=1.0,
+                         enable_block_reuse=False),
+        tensorrt_llm.bindings.internal.batch_manager.CacheType.SELF,
+        num_layers=1,
+        num_kv_heads=1,
+        head_dim=128,
+        tokens_per_block=32,
+        max_seq_len=max_seq_len,
+        max_batch_size=2,
+        mapping=Mapping(world_size=1, tp_size=1, rank=0),
+        dtype=tensorrt_llm.bindings.DataType.BF16,
+    )
+
+
 class TestFlashInferAttention(unittest.TestCase):
 
+    def test_generation_page_table_uses_reserved_block_count(self):
+        if not torch.cuda.is_available():
+            self.skipTest("CUDA is required for KV cache managers")
+
+        request_ids = [98, 99]
+        tokens_per_block = 32
+        reserved_block_counts = [4, 325]
+        token_counts = [
+            num_blocks * tokens_per_block - 1
+            for num_blocks in reserved_block_counts
+        ]
+        max_tokens = sum(token_counts) + 2
+        max_seq_len = max(token_counts) + 1
+
+        with _create_kv_cache_manager(
+                max_tokens=max_tokens,
+                max_seq_len=max_seq_len,
+        ) as v1_manager, _create_kv_cache_manager_v2(
+                max_tokens,
+                max_seq_len,
+        ) as v2_manager:
+            v1_manager.add_dummy_requests(request_ids,
+                                          token_counts,
+                                          is_gen=True)
+            v2_manager.add_dummy_requests(request_ids,
+                                          token_counts,
+                                          is_gen=True)
+
+            for manager in (v1_manager, v2_manager):
+                self.assertEqual(
+                    flashinfer_backend._get_page_table_num_blocks(
+                        manager,
+                        request_ids,
+                        [3, 324],
+                        num_contexts=1,
+                    ),
+                    [3, 325],
+                )
+
+    def test_spec_decode_kv_lens_offsets_update_logical_decode_state(self):
+        if not torch.cuda.is_available():
+            self.skipTest("CUDA is required for FlashInfer metadata")
+
+        kv_cache_manager = _create_kv_cache_manager()
+        with kv_cache_manager:
+            num_contexts = 1
+            num_generations = 2
+            tokens_per_generation = 4
+            seq_lens = [2] + [tokens_per_generation] * num_generations
+            num_seqs = len(seq_lens)
+            num_tokens = sum(seq_lens)
+            # Pack one context row, then two four-token speculative decode rows.
+            metadata = FlashInferAttentionMetadata(
+                seq_lens=torch.tensor(seq_lens, dtype=torch.int32),
+                num_contexts=num_contexts,
+                kv_cache_params=KVCacheParams(use_cache=True),
+                max_num_requests=num_seqs,
+                max_num_tokens=num_tokens,
+                kv_cache_manager=kv_cache_manager,
+            )
+            # Seed the logical state at the host-planned upper bound. The
+            # context row is a control and must remain untouched below.
+            cached_token_lens = torch.tensor([10, 31, 62],
+                                             dtype=torch.int32,
+                                             device="cuda")
+            last_page_lens = torch.tensor([12, 3, 2],
+                                          dtype=torch.int32,
+                                          device="cuda")
+            positions = torch.tensor([10, 11, 31, 32, 33, 34, 62, 63, 64, 65],
+                                     dtype=torch.int32,
+                                     device="cuda")
+            metadata._cached_token_lens[:num_seqs].copy_(cached_token_lens)
+            metadata._paged_kv_last_page_len[:num_seqs].copy_(last_page_lens)
+            metadata._positions[:num_tokens].copy_(positions)
+            planned_kv_lens = (cached_token_lens[num_contexts:] +
+                               tokens_per_generation)
+            kv_lens_buffer = planned_kv_lens.clone()
+            metadata._plan_params_to_wrappers = {
+                object():
+                FlashInferWrappers(
+                    is_planned=True,
+                    decode_wrapper=FakeDecodeWrapper(kv_lens_buffer))
+            }
+            # The decode requests accepted one and three of their four
+            # speculative tokens.
+            offsets = torch.tensor([-3, -1], dtype=torch.int32, device="cuda")
+            corrected_cached_token_lens = cached_token_lens.clone()
+            corrected_cached_token_lens[num_contexts:].add_(offsets)
+            corrected_positions = positions.clone()
+            corrected_positions[sum(seq_lens[:num_contexts]):].add_(
+                offsets.repeat_interleave(tokens_per_generation))
+            corrected_kv_lens = planned_kv_lens + offsets
+
+            metadata.apply_spec_decode_kv_lens_offsets(offsets, num_generations,
+                                                       tokens_per_generation)
+
+            # Rewind logical decode state while preserving structural
+            # page-table lengths.
+            torch.testing.assert_close(metadata._cached_token_lens[:num_seqs],
+                                       corrected_cached_token_lens)
+            torch.testing.assert_close(
+                metadata._paged_kv_last_page_len[:num_seqs], last_page_lens)
+            torch.testing.assert_close(metadata._positions[:num_tokens],
+                                       corrected_positions)
+            torch.testing.assert_close(kv_lens_buffer, corrected_kv_lens)
+
+            # A mixed prefill/decode forward can lazily plan after overlap
+            # preprocessing. Simulate plan() restoring the host upper bound and
+            # verify the post-plan publication restores device-logical lengths.
+            kv_lens_buffer.copy_(planned_kv_lens)
+            metadata._publish_decode_wrapper_kv_lens(
+                next(iter(
+                    metadata._plan_params_to_wrappers.values())).decode_wrapper)
+            torch.testing.assert_close(kv_lens_buffer, corrected_kv_lens)
+
+            # Restoring must exactly reverse the logical-state correction.
+            metadata.apply_spec_decode_kv_lens_offsets(
+                offsets,
+                num_generations,
+                tokens_per_generation,
+                restore=True,
+            )
+
+            torch.testing.assert_close(metadata._cached_token_lens[:num_seqs],
+                                       cached_token_lens)
+            torch.testing.assert_close(
+                metadata._paged_kv_last_page_len[:num_seqs], last_page_lens)
+            torch.testing.assert_close(metadata._positions[:num_tokens],
+                                       positions)
+            torch.testing.assert_close(kv_lens_buffer, planned_kv_lens)
+
+            # A shared external-assistant view is Q-only: its cached length is
+            # the full accepted target prefix and excludes its query token.
+            metadata.seq_lens = torch.ones(num_generations, dtype=torch.int32)
+            metadata.num_contexts = 0
+            metadata._is_shared_kv_draft_view = True
+            metadata._draft_kv_runtime_lens[:num_generations].copy_(
+                corrected_kv_lens)
+            metadata._cached_token_lens[:num_generations].zero_()
+            metadata._paged_kv_last_page_len[:num_generations].zero_()
+            kv_lens_buffer.copy_(planned_kv_lens)
+            metadata._update_draft_kv_lengths()
+            torch.testing.assert_close(kv_lens_buffer, corrected_kv_lens)
+
+    def test_mla_decode_uses_offset_logical_lengths_with_reserved_pages(self):
+        if not torch.cuda.is_available():
+            self.skipTest("CUDA is required for FlashInfer metadata")
+
+        class FakeMLADecodeWrapper:
+
+            def plan(self, *args, **kwargs):
+                # Capture the KV lengths passed to MLA during planning.
+                self.plan_args = args
+
+        num_generations = 2
+        tokens_per_generation = 1
+        num_generation_blocks = 5
+        kv_cache_manager = _create_kv_cache_manager()
+        metadata = FlashInferAttentionMetadata(
+            seq_lens=torch.ones(num_generations, dtype=torch.int32),
+            num_contexts=0,
+            kv_cache_params=KVCacheParams(use_cache=True),
+            max_num_requests=num_generations,
+            max_num_tokens=num_generations * tokens_per_generation,
+            kv_cache_manager=kv_cache_manager,
+        )
+        # Reserve two and three pages. At page size 8 with five entries in each
+        # final page, the structural upper-bound lengths are 13 and 21.
+        metadata._uses_full_generation_page_table = True
+        metadata.num_generation_blocks = num_generation_blocks
+        metadata.num_context_blocks = 0
+        metadata.paged_kv_indptr_decode[:num_generations + 1].copy_(
+            torch.tensor([0, 2, 5], dtype=torch.int32, device="cuda"))
+        metadata._paged_kv_indices[:num_generation_blocks].copy_(
+            torch.arange(num_generation_blocks,
+                         dtype=torch.int32,
+                         device="cuda"))
+        metadata._paged_kv_last_page_len[:num_generations].fill_(5)
+        metadata._qo_indptr[:num_generations + 1].copy_(
+            torch.arange(num_generations + 1, dtype=torch.int32, device="cuda"))
+        # Track shorter device-logical lengths independently of reserved pages.
+        logical_kv_lens = torch.tensor([8, 14],
+                                       dtype=torch.int32,
+                                       device="cuda")
+        metadata._cached_token_lens[:num_generations].copy_(
+            logical_kv_lens - tokens_per_generation)
+        metadata._logical_kv_lens[:num_generations].copy_(logical_kv_lens)
+        metadata._positions[:num_generations].copy_(logical_kv_lens)
+        # Device acceptance results rewind the logical lengths to 5 and 13.
+        offsets = torch.tensor([-3, -1], dtype=torch.int32, device="cuda")
+        corrected_logical_kv_lens = logical_kv_lens + offsets
+
+        metadata.apply_spec_decode_kv_lens_offsets(offsets, num_generations,
+                                                   tokens_per_generation)
+
+        torch.testing.assert_close(metadata._logical_kv_lens[:num_generations],
+                                   corrected_logical_kv_lens)
+
+        wrapper = FakeMLADecodeWrapper()
+        metadata._mla_decode_wrapper = wrapper
+        metadata._do_plan_mla_decode(
+            MLAPlanParams(
+                num_heads=2,
+                kv_lora_rank=4,
+                qk_rope_head_dim=2,
+                page_size=8,
+                q_dtype=torch.bfloat16,
+                kv_dtype=torch.bfloat16,
+            ))
+
+        # The page table derives [13, 21], but MLA needs the live lengths.
+        torch.testing.assert_close(wrapper.plan_args[3],
+                                   corrected_logical_kv_lens)
+
+        # Restore live lengths without rebuilding reservations.
+        metadata.apply_spec_decode_kv_lens_offsets(
+            offsets,
+            num_generations,
+            tokens_per_generation,
+            restore=True,
+        )
+        torch.testing.assert_close(metadata._logical_kv_lens[:num_generations],
+                                   logical_kv_lens)
+        kv_cache_manager.shutdown()
+
     def test_separate_kv_draft_metadata_uses_draft_manager(self):
         if not torch.cuda.is_available():
             self.skipTest("CUDA is required for FlashInfer metadata")
         if torch.cuda.get_device_capability() not in ((10, 0), (10, 3)):
             self.skipTest("FlashInfer trtllm-gen requires SM100 or SM103")
 
-        def create_manager():
-            return KVCacheManager(
-                KvCacheConfig(max_tokens=256),
-                tensorrt_llm.bindings.internal.batch_manager.CacheType.SELF,
-                num_layers=1,
-                num_kv_heads=1,
-                head_dim=128,
-                tokens_per_block=32,
-                max_seq_len=64,
-                max_batch_size=2,
-                mapping=Mapping(world_size=1, tp_size=1, rank=0),
-                dtype=tensorrt_llm.bindings.DataType.BF16,
-            )
-
-        target_manager = create_manager()
-        draft_manager = create_manager()
-        try:
+        target_manager = _create_kv_cache_manager()
+        draft_manager = _create_kv_cache_manager()
+        with target_manager, draft_manager:
             target_manager.add_dummy_requests([0, 1], [31, 45], is_gen=True)
             draft_manager.add_dummy_requests([0, 1], [31, 45],
                                              is_gen=True,
@@ -149,9 +416,6 @@ def create_manager():
             replan.assert_not_called()
             self.assertEqual(refresh_block_tables.call_count,
                              len(draft_metadata._plan_params_to_wrappers))
-        finally:
-            target_manager.shutdown()
-            draft_manager.shutdown()
 
     def test_ragged_no_kv_cuda_graph_uses_stable_indptr_aliases(self):
         if not torch.cuda.is_available():
diff --git a/tests/unittest/_torch/executor/test_py_executor.py b/tests/unittest/_torch/executor/test_py_executor.py
index 10953bbbdbc2..00adf7d9b0b0 100644
--- a/tests/unittest/_torch/executor/test_py_executor.py
+++ b/tests/unittest/_torch/executor/test_py_executor.py
@@ -37,6 +37,9 @@
     PyExecutor,
     _ADPForwardIntent,
 )
+from tensorrt_llm._torch.pyexecutor.py_executor_creator import (
+    _disable_unsupported_speculative_overlap_scheduler,
+)
 from tensorrt_llm._torch.pyexecutor.resource_manager import NoFreeSlotsError, ResourceManagerType
 from tensorrt_llm._torch.pyexecutor.scheduler import (
     FCFSWaitingQueue,
@@ -44,6 +47,12 @@
     ScheduledRequests,
     SerializableSchedulerOutput,
 )
+from tensorrt_llm.llmapi.llm_args import (
+    DraftTargetDecodingConfig,
+    MTPDecodingConfig,
+    SpeculativeConfig,
+    TorchLlmArgs,
+)
 from tensorrt_llm.runtime.kv_cache_manager_v2 import OutOfPagesError
 
 pytestmark = pytest.mark.cpu_only
@@ -2622,6 +2631,65 @@ def test_handles_error_on_single_rank(self):
         assert len(stub.handle_errors_calls) == 1
 
 
+class TestUnsupportedSpeculativeOverlapScheduler:
+    @staticmethod
+    def _make_args(attn_backend: str, spec_config: SpeculativeConfig) -> TorchLlmArgs:
+        return TorchLlmArgs(
+            model="dummy",
+            attn_backend=attn_backend,
+            speculative_config=spec_config,
+            disable_overlap_scheduler=False,
+        )
+
+    def test_vanilla_disables_speculative_overlap(self) -> None:
+        llm_args = self._make_args(
+            "VANILLA",
+            MTPDecodingConfig(max_draft_len=3, mtp_eagle_one_model=True),
+        )
+
+        _disable_unsupported_speculative_overlap_scheduler(llm_args, llm_args.speculative_config)
+
+        assert llm_args.disable_overlap_scheduler
+
+    def test_vanilla_with_overlap_already_disabled_is_unchanged(self) -> None:
+        llm_args = self._make_args(
+            "VANILLA",
+            MTPDecodingConfig(max_draft_len=3, mtp_eagle_one_model=True),
+        )
+        llm_args.disable_overlap_scheduler = True
+
+        _disable_unsupported_speculative_overlap_scheduler(llm_args, llm_args.speculative_config)
+
+        assert llm_args.disable_overlap_scheduler
+
+    @pytest.mark.parametrize("attn_backend", ["FLASHINFER", "flashinfer"])
+    def test_flashinfer_two_model_extend_context_disables_speculative_overlap(
+        self, attn_backend: str
+    ) -> None:
+        spec_config = DraftTargetDecodingConfig(
+            max_draft_len=3,
+            speculative_model="dummy",
+        )
+        # Public configs currently select the one-model DraftTarget path. Exercise
+        # the retained legacy two-model mode that uses FlashInfer extend-context.
+        spec_config._draft_target_one_model = False
+        llm_args = self._make_args(attn_backend, spec_config)
+
+        _disable_unsupported_speculative_overlap_scheduler(llm_args, llm_args.speculative_config)
+
+        assert llm_args.disable_overlap_scheduler
+
+    def test_flashinfer_one_engine_keeps_speculative_overlap(self) -> None:
+        llm_args = self._make_args(
+            "FLASHINFER",
+            MTPDecodingConfig(max_draft_len=3, mtp_eagle_one_model=True),
+        )
+
+        _disable_unsupported_speculative_overlap_scheduler(llm_args, llm_args.speculative_config)
+
+        assert not llm_args.disable_overlap_scheduler
+
+
 class TestOneModelMTPDraftTokenScheduling:
     """Regression tests for the one-model MTP over-scheduling bug (#16101).
 
@@ -2635,12 +2703,17 @@ class TestOneModelMTPDraftTokenScheduling:
     forward then builds a uniform ``1 + runtime_draft_len`` per gen request and
     overshoots ``max_num_tokens`` (``total_num_tokens > max_num_tokens``).
 
-    The fix populates ``request.draft_tokens = [0] * max_total_draft_tokens``
-    on every in-progress generation request so scheduling reserves the correct
-    token budget. This test drives ``_prepare_and_schedule_batch`` for a
-    one-model-MTP executor and asserts generation requests get
-    ``num_draft_tokens == max_total_draft_tokens`` while context requests are
-    left untouched.
+    The fix populates both the Python and C++ draft-token representations on
+    every in-progress generation request so both schedulers reserve the
+    correct token budget. This test drives ``_prepare_and_schedule_batch`` for
+    a one-model-MTP executor and asserts generation requests get the full
+    draft-token budget while context requests are left untouched.
+
+    The Python-side fill is placeholder-only: with the overlap scheduler
+    disabled, `_prepare_tp_inputs` sources a generation request's draft
+    tokens from `py_draft_tokens`, which the one-model spec sampler wrote at
+    the end of the previous iteration. Overwriting a populated list here would
+    feed zeros to the target model and collapse the acceptance rate.
 
     NOTE: Like ``test_fetch_called_once_even_in_benchmark_disagg`` in
     ``test_benchmark_disagg.py``, this uses ``object.__new__(PyExecutor)`` to
@@ -2667,7 +2740,11 @@ def _make_llm_request(request_id: int, state: LlmRequestState) -> LlmRequest:
         return req
 
     @classmethod
-    def _make_one_model_mtp_executor(cls, active_requests):
+    def _make_one_model_mtp_executor(
+        cls,
+        active_requests: list[LlmRequest],
+        use_rejection_sampling: bool = False,
+    ) -> PyExecutor:
         """Construct a partially-initialised one-model-MTP PyExecutor.
 
         drafter is None (one-model MTP has no separate drafter) and
@@ -2680,11 +2757,25 @@ def _make_one_model_mtp_executor(cls, active_requests):
         ex = object.__new__(PyExecutor)
         ex.drafter = None
         ex.max_total_draft_tokens = cls.MAX_TOTAL_DRAFT_TOKENS
-        ex.model_engine = Mock(is_spec_decode=True)
+        spec_config = MTPDecodingConfig(
+            max_draft_len=cls.MAX_TOTAL_DRAFT_TOKENS,
+            mtp_eagle_one_model=True,
+            use_rejection_sampling=use_rejection_sampling,
+        )
+        ex.model_engine = Mock(
+            is_spec_decode=True,
+            spec_config=spec_config,
+            max_draft_len=cls.MAX_TOTAL_DRAFT_TOKENS,
+            max_total_draft_tokens=cls.MAX_TOTAL_DRAFT_TOKENS,
+        )
         ex.kv_cache_transceiver = None
         ex.is_shutdown = False
         ex.enable_iter_perf_stats = False
         ex.enable_attention_dp = False
+        ex.disable_overlap_scheduler = False
+        ex.speculation_permanently_disabled = False
+        ex.max_seq_len = 64
+        ex.dist = Mock()
         ex.active_requests = active_requests
         ex.waiting_queue = []
 
@@ -2717,6 +2808,8 @@ def test_one_model_mtp_populates_draft_tokens_for_scheduling(self):
         # Precondition: no draft tokens reserved yet on either gen request.
         assert gen.num_draft_tokens == 0
         assert disagg_gen.num_draft_tokens == 0
+        assert gen.py_draft_tokens == []
+        assert disagg_gen.py_draft_tokens == []
 
         ex = self._make_one_model_mtp_executor([gen, disagg_gen, ctx])
         scheduled_batch, _ = ex._prepare_and_schedule_batch()
@@ -2726,7 +2819,199 @@ def test_one_model_mtp_populates_draft_tokens_for_scheduling(self):
         # full draft-token budget so the micro-batch scheduler reserves
         # beam + max_total_draft_tokens.
         assert gen.num_draft_tokens == self.MAX_TOTAL_DRAFT_TOKENS
+        assert gen.py_draft_tokens == [0] * self.MAX_TOTAL_DRAFT_TOKENS
         # Disaggregated case: decode-worker request awaiting KV also normalized.
         assert disagg_gen.num_draft_tokens == self.MAX_TOTAL_DRAFT_TOKENS
+        assert disagg_gen.py_draft_tokens == [0] * self.MAX_TOTAL_DRAFT_TOKENS
         # Context requests are not generation requests and must be left alone.
         assert ctx.num_draft_tokens == 0
+        assert ctx.py_draft_tokens == []
+
+    def test_one_model_mtp_preserves_zero_proposal_signal_for_rejection(self) -> None:
+        gen = self._make_llm_request(0, LlmRequestState.GENERATION_IN_PROGRESS)
+        ex = self._make_one_model_mtp_executor([gen], use_rejection_sampling=True)
+
+        ex._prepare_and_schedule_batch()
+
+        # Scheduling still sees the full reservation, while the independent
+        # signal records that no proposal probabilities were produced.
+        assert gen.num_draft_tokens == self.MAX_TOTAL_DRAFT_TOKENS
+        assert gen.py_draft_tokens == [0] * self.MAX_TOTAL_DRAFT_TOKENS
+        assert gen.py_needs_onehot_draft_probs
+
+        batch = ScheduledRequests()
+        batch.append_generation_request(gen)
+        ex._handle_dynamic_draft_len(batch)
+
+        assert ex.model_engine.runtime_draft_len == self.MAX_TOTAL_DRAFT_TOKENS
+        assert gen.py_needs_onehot_draft_probs
+        assert gen.py_draft_tokens == [0] * self.MAX_TOTAL_DRAFT_TOKENS
+
+    def test_one_model_mtp_preserves_sampler_draft_tokens(self):
+        """Normalization must not clobber real draft tokens.
+
+        With `disable_overlap_scheduler=True` the one-model spec sampler writes the next iteration's
+        draft tokens into `py_draft_tokens`, and `_prepare_tp_inputs` reads them straight back into
+        `input_ids` / `draft_tokens_cuda` (there is no previous-iteration device tensor to source
+        them from). Overwriting them with the zero placeholder leaves the target model verifying
+        token id 0, silently dropping the acceptance rate to chance.
+        Only the C++ count needs unconditional normalization.
+        """
+        sampler_drafts = [7, 8]
+        assert len(sampler_drafts) < self.MAX_TOTAL_DRAFT_TOKENS
+
+        gen = self._make_llm_request(0, LlmRequestState.GENERATION_IN_PROGRESS)
+        gen.py_draft_tokens = list(sampler_drafts)
+
+        ex = self._make_one_model_mtp_executor([gen])
+        ex._prepare_and_schedule_batch()
+
+        assert gen.py_draft_tokens == sampler_drafts
+        # The C++ count is still normalized for the micro-batch scheduler.
+        assert gen.num_draft_tokens == self.MAX_TOTAL_DRAFT_TOKENS
+
+    def test_permanent_disable_clears_non_mtp_eagle_draft_tokens(self) -> None:
+        ex = self._make_runtime_draft_executor(
+            disable_overlap_scheduler=False, use_shared_target_kv=False
+        )
+        ex.model_engine.spec_config = MTPDecodingConfig(
+            max_draft_len=self.MAX_TOTAL_DRAFT_TOKENS,
+            use_mtp_vanilla=True,
+        )
+        ex.speculation_permanently_disabled = True
+        batch = self._make_generation_batch(4)
+
+        ex._handle_dynamic_draft_len(batch)
+
+        assert ex.model_engine.runtime_draft_len == 0
+        assert batch.generation_requests[0].py_draft_tokens == []
+
+    @classmethod
+    def _make_runtime_draft_executor(
+        cls, disable_overlap_scheduler: bool, use_shared_target_kv: bool
+    ):
+        ex = object.__new__(PyExecutor)
+        spec_config = MTPDecodingConfig(
+            max_draft_len=cls.MAX_TOTAL_DRAFT_TOKENS,
+            mtp_eagle_one_model=True,
+        )
+        spec_config._use_shared_kv_cache = use_shared_target_kv
+        ex.model_engine = Mock(
+            spec_config=spec_config,
+            max_draft_len=cls.MAX_TOTAL_DRAFT_TOKENS,
+            max_total_draft_tokens=cls.MAX_TOTAL_DRAFT_TOKENS,
+        )
+        ex.disable_overlap_scheduler = disable_overlap_scheduler
+        ex.speculation_permanently_disabled = False
+        ex.enable_attention_dp = False
+        ex.dist = Mock()
+        ex.max_seq_len = 16
+        return ex
+
+    @classmethod
+    def _make_generation_batch(cls, *sequence_lengths: int):
+        batch = ScheduledRequests()
+        for request_id, sequence_length in enumerate(sequence_lengths):
+            request = LlmRequest(
+                request_id=request_id,
+                max_new_tokens=10,
+                input_tokens=list(range(sequence_length)),
+                sampling_config=SamplingConfig(1),
+                is_streaming=False,
+                draft_tokens=None,
+            )
+            request.state = LlmRequestState.GENERATION_IN_PROGRESS
+            request.py_draft_tokens = [7, 8, 9]
+            batch.append_generation_request(request)
+        return batch
+
+    @pytest.mark.parametrize(
+        "disable_overlap_scheduler,use_shared_target_kv,sequence_length",
+        [
+            (False, False, 8),
+            (True, False, 12),
+            (False, True, 9),
+            (True, True, 13),
+        ],
+    )
+    def test_one_model_mtp_uses_zero_draft_near_sequence_limit(
+        self, disable_overlap_scheduler, use_shared_target_kv, sequence_length
+    ):
+        ex = self._make_runtime_draft_executor(disable_overlap_scheduler, use_shared_target_kv)
+        batch = self._make_generation_batch(sequence_length, 4)
+
+        ex._handle_dynamic_draft_len(batch)
+
+        assert ex.model_engine.runtime_draft_len == 0
+        assert all(request.py_draft_tokens == [] for request in batch.generation_requests)
+
+    @pytest.mark.parametrize(
+        "disable_overlap_scheduler,use_shared_target_kv,sequence_length",
+        [
+            (False, False, 7),
+            (True, False, 11),
+            (False, True, 8),
+            (True, True, 12),
+        ],
+    )
+    def test_one_model_mtp_keeps_drafting_with_position_headroom(
+        self, disable_overlap_scheduler, use_shared_target_kv, sequence_length
+    ):
+        ex = self._make_runtime_draft_executor(disable_overlap_scheduler, use_shared_target_kv)
+        batch = self._make_generation_batch(sequence_length)
+
+        ex._handle_dynamic_draft_len(batch)
+
+        assert ex.model_engine.runtime_draft_len == self.MAX_TOTAL_DRAFT_TOKENS
+        assert batch.generation_requests[0].py_draft_tokens == [7, 8, 9]
+
+    @pytest.mark.parametrize("peer_needs_zero_draft", [False, True])
+    def test_one_model_mtp_synchronizes_zero_draft_across_attention_dp(
+        self, peer_needs_zero_draft: bool
+    ) -> None:
+        ex = self._make_runtime_draft_executor(
+            disable_overlap_scheduler=False, use_shared_target_kv=False
+        )
+        ex.enable_attention_dp = True
+        ex.dist.tp_allgather.return_value = [False, peer_needs_zero_draft]
+        batch = self._make_generation_batch(7)
+
+        ex._handle_dynamic_draft_len(batch)
+
+        ex.dist.tp_allgather.assert_called_once_with(False)
+        if peer_needs_zero_draft:
+            assert ex.model_engine.runtime_draft_len == 0
+            assert batch.generation_requests[0].py_draft_tokens == []
+        else:
+            assert ex.model_engine.runtime_draft_len == self.MAX_TOTAL_DRAFT_TOKENS
+            assert batch.generation_requests[0].py_draft_tokens == [7, 8, 9]
+
+    def test_one_model_mtp_zero_draft_is_collective_free_without_attention_dp(self) -> None:
+        ex = self._make_runtime_draft_executor(
+            disable_overlap_scheduler=False, use_shared_target_kv=False
+        )
+        batch = self._make_generation_batch(8)
+
+        ex._handle_dynamic_draft_len(batch)
+
+        assert ex.model_engine.runtime_draft_len == 0
+        ex.dist.tp_allgather.assert_not_called()
+
+    def test_one_model_mtp_synchronizes_locally_selected_zero_draft(self) -> None:
+        ex = self._make_runtime_draft_executor(
+            disable_overlap_scheduler=False, use_shared_target_kv=False
+        )
+        ex.model_engine.spec_config = MTPDecodingConfig(
+            max_draft_len=self.MAX_TOTAL_DRAFT_TOKENS,
+            mtp_eagle_one_model=True,
+            draft_len_schedule={1: self.MAX_TOTAL_DRAFT_TOKENS, 2: 0},
+        )
+        ex.enable_attention_dp = True
+        ex.dist.tp_allgather.return_value = [True, False]
+        batch = self._make_generation_batch(4, 4)
+
+        ex._handle_dynamic_draft_len(batch)
+
+        ex.dist.tp_allgather.assert_called_once_with(True)
+        assert ex.model_engine.runtime_draft_len == 0
+        assert batch.generation_requests[0].py_draft_tokens == []
diff --git a/tests/unittest/_torch/executor/test_pytorch_model_engine.py b/tests/unittest/_torch/executor/test_pytorch_model_engine.py
index 84be2e7db3f1..1b4e2aa928fc 100644
--- a/tests/unittest/_torch/executor/test_pytorch_model_engine.py
+++ b/tests/unittest/_torch/executor/test_pytorch_model_engine.py
@@ -36,14 +36,16 @@
 # isort: on
 from utils.util import skip_ray
 
+from tensorrt_llm._torch.attention_backend import FlashInferAttentionMetadata
 from tensorrt_llm._torch.attention_backend.interface import AttentionMetadata
+from tensorrt_llm._torch.metadata import KVCacheParams
 from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests
 from tensorrt_llm._torch.speculative.spec_sampler_base import \
     SampleStateTensorsSpec
 from tensorrt_llm.bindings.executor import KvCacheConfig
 from tensorrt_llm.inputs.registry import BaseMultimodalDummyInputsBuilder
-from tensorrt_llm.llmapi import (CudaGraphConfig, SADecodingConfig,
-                                 SamplingParams)
+from tensorrt_llm.llmapi import (CudaGraphConfig, MTPDecodingConfig,
+                                 SADecodingConfig, SamplingParams)
 from tensorrt_llm.mapping import CpType, Mapping
 
 
@@ -1373,6 +1375,230 @@ def set_attn_max_seq_len(self, max_seq_len: int) -> None:
                                  (encoder_batch_size, encoder_max_num_tokens))
                 self.assertEqual(encoder.max_seq_len, expected_max_seq_len)
 
+    def test_dynamic_tree_prepare_preserves_explicit_zero_draft(self) -> None:
+        max_draft_len = 3
+        dynamic_tree_max_top_k = 4
+        max_total_draft_tokens = max_draft_len * dynamic_tree_max_top_k
+        max_num_requests = 4
+        runtime_draft_len = 0
+        expected_tokens_per_gen_step = 1
+        # Allocate buffers for the dynamic tree's full 3-level, top-4 width.
+        allocation_config = SADecodingConfig(
+            max_draft_len=max_total_draft_tokens, )
+        model_engine, kv_cache_manager = create_model_engine_and_kvcache(
+            spec_config=allocation_config)
+        dynamic_tree_config = MTPDecodingConfig(
+            max_draft_len=max_draft_len,
+            mtp_eagle_one_model=True,
+            use_dynamic_tree=True,
+            dynamic_tree_max_topK=dynamic_tree_max_top_k,
+        )
+        model_engine.spec_config = dynamic_tree_config
+        model_engine.max_draft_len = dynamic_tree_config.max_draft_len
+        model_engine.max_total_draft_tokens = max_total_draft_tokens
+        # Zero explicitly disables drafting for this iteration. Tree-width
+        # normalization must not replace it with max_total_draft_tokens.
+        model_engine.runtime_draft_len = runtime_draft_len
+        resource_manager = ResourceManager(
+            {ResourceManagerType.KV_CACHE_MANAGER: kv_cache_manager})
+        attn_metadata = AttentionMetadata(max_num_requests=max_num_requests,
+                                          max_num_tokens=32,
+                                          kv_cache_manager=kv_cache_manager)
+        attn_metadata.is_cuda_graph = False
+
+        generation = _create_request_with_tokens([50, 51, 52, 53, 54], 1)
+        generation.py_seq_slot = 0
+        generation.py_batch_idx = 0
+        generation.py_draft_tokens = []
+
+        # Supply overlap buffers even though this request should consume only
+        # its single target token.
+        graph_batch = ScheduledRequests()
+        graph_batch.generation_requests = [generation]
+        overlap_state = SampleStateTensorsSpec(
+            new_tokens=torch.zeros(
+                (max_total_draft_tokens + 1, max_num_requests, 1),
+                dtype=torch.int32,
+                device="cuda"),
+            new_tokens_lens=torch.ones(max_num_requests,
+                                       dtype=torch.int32,
+                                       device="cuda"),
+            next_draft_tokens=torch.zeros(
+                (max_num_requests, max_total_draft_tokens),
+                dtype=torch.int32,
+                device="cuda"),
+        )
+        spec_metadata = Mock(_force_non_greedy_for_capture=False)
+
+        inputs, _ = model_engine._prepare_tp_inputs(
+            scheduled_requests=graph_batch,
+            kv_cache_manager=kv_cache_manager,
+            attn_metadata=attn_metadata,
+            spec_metadata=spec_metadata,
+            new_tensors_device=overlap_state,
+            resource_manager=resource_manager,
+        )
+
+        # Preparation must retain the no-draft shape in the engine, metadata,
+        # and inputs.
+        expected_seq_lens = [expected_tokens_per_gen_step]
+        self.assertEqual(model_engine.runtime_draft_len, runtime_draft_len)
+        self.assertEqual(
+            model_engine.get_runtime_tokens_per_gen_step(runtime_draft_len),
+            expected_tokens_per_gen_step)
+        self.assertEqual(attn_metadata.seq_lens.tolist(), expected_seq_lens)
+        self.assertEqual(spec_metadata.seq_lens, expected_seq_lens)
+        self.assertEqual(inputs["input_ids"].numel(),
+                         expected_tokens_per_gen_step)
+        kv_cache_manager.shutdown()
+
+    def test_overlap_input_processing_applies_flashinfer_kv_offsets(
+            self) -> None:
+        spec_config = SADecodingConfig(max_draft_len=3)
+        model_engine, kv_cache_manager = create_model_engine_and_kvcache(
+            spec_config=spec_config)
+        num_contexts = 1
+        num_generations = 2
+        runtime_draft_len = spec_config.max_draft_len
+        tokens_per_generation = spec_config.get_runtime_tokens_per_gen_step(
+            runtime_draft_len)
+        seq_lens = [2] + [tokens_per_generation] * num_generations
+        num_seqs = len(seq_lens)
+        num_tokens = sum(seq_lens)
+        model_engine.runtime_draft_len = runtime_draft_len
+        model_engine.previous_pos_id_offsets_cuda = torch.zeros(
+            num_generations * tokens_per_generation,
+            dtype=torch.int32,
+            device="cuda")
+        # Four tokens were planned per decode row, but the requests accepted
+        # only one and three.
+        offsets = torch.tensor([-3, -1], dtype=torch.int32, device="cuda")
+        model_engine.previous_kv_lens_offsets_cuda = offsets
+
+        # Pack one context row followed by the two speculative decode rows.
+        attn_metadata = FlashInferAttentionMetadata(
+            seq_lens=torch.tensor(seq_lens, dtype=torch.int32),
+            num_contexts=num_contexts,
+            kv_cache_params=KVCacheParams(use_cache=True),
+            max_num_requests=4,
+            max_num_tokens=32,
+            kv_cache_manager=kv_cache_manager,
+        )
+        attn_metadata.num_chunked_ctx_requests = 0
+        cached_token_lens = torch.tensor([10, 31, 62],
+                                         dtype=torch.int32,
+                                         device="cuda")
+        positions = torch.tensor([10, 11, 31, 32, 33, 34, 62, 63, 64, 65],
+                                 dtype=torch.int32,
+                                 device="cuda")
+        attn_metadata._cached_token_lens[:num_seqs].copy_(cached_token_lens)
+        attn_metadata._positions[:num_tokens].copy_(positions)
+        corrected_cached_token_lens = cached_token_lens.clone()
+        corrected_cached_token_lens[num_contexts:].add_(offsets)
+        corrected_positions = positions.clone()
+        corrected_positions[sum(seq_lens[:num_contexts]):].add_(
+            offsets.repeat_interleave(tokens_per_generation))
+        inputs = {
+            "input_ids":
+            torch.zeros(num_tokens, dtype=torch.int32, device="cuda"),
+            "position_ids":
+            torch.zeros((1, num_tokens), dtype=torch.int32, device="cuda"),
+            "attn_metadata":
+            attn_metadata,
+        }
+
+        model_engine._preprocess_inputs(inputs)
+
+        # Preprocessing rewinds logical KV lengths and positions only for the
+        # decode rows.
+        torch.testing.assert_close(attn_metadata._cached_token_lens[:num_seqs],
+                                   corrected_cached_token_lens)
+        torch.testing.assert_close(attn_metadata._positions[:num_tokens],
+                                   corrected_positions)
+
+        model_engine._postprocess_inputs(inputs)
+
+        # Postprocessing restores metadata reused by CUDA graph capture.
+        torch.testing.assert_close(attn_metadata._cached_token_lens[:num_seqs],
+                                   cached_token_lens)
+        kv_cache_manager.shutdown()
+
+    def test_overlap_input_processing_applies_flashinfer_extend_ctx_kv_offsets(
+            self) -> None:
+        spec_config = SADecodingConfig(max_draft_len=3)
+        model_engine, kv_cache_manager = create_model_engine_and_kvcache(
+            spec_config=spec_config)
+        num_contexts = 3
+        num_chunked_contexts = 2
+        runtime_draft_len = spec_config.max_draft_len
+        tokens_per_generation = spec_config.get_runtime_tokens_per_gen_step(
+            runtime_draft_len)
+        seq_lens = [2] + [tokens_per_generation] * num_chunked_contexts
+        num_tokens = sum(seq_lens)
+        model_engine.runtime_draft_len = runtime_draft_len
+        model_engine.previous_pos_id_offsets_cuda = torch.zeros(
+            num_chunked_contexts * tokens_per_generation,
+            dtype=torch.int32,
+            device="cuda")
+        # Four tokens were planned per extend row, but the requests accepted
+        # only one and three.
+        offsets = torch.tensor([-3, -1], dtype=torch.int32, device="cuda")
+        model_engine.previous_kv_lens_offsets_cuda = offsets
+
+        # All rows are contexts, but the trailing two are speculative requests
+        # packed by extend_ctx. They need offsets even with no generation rows.
+        attn_metadata = FlashInferAttentionMetadata(
+            seq_lens=torch.tensor(seq_lens, dtype=torch.int32),
+            num_contexts=num_contexts,
+            kv_cache_params=KVCacheParams(use_cache=True),
+            max_num_requests=4,
+            max_num_tokens=32,
+            kv_cache_manager=kv_cache_manager,
+        )
+        attn_metadata.num_chunked_ctx_requests = num_chunked_contexts
+        cached_token_lens = torch.tensor([10, 31, 62],
+                                         dtype=torch.int32,
+                                         device="cuda")
+        positions = torch.tensor([10, 11, 31, 32, 33, 34, 62, 63, 64, 65],
+                                 dtype=torch.int32,
+                                 device="cuda")
+        attn_metadata._cached_token_lens[:num_contexts].copy_(cached_token_lens)
+        attn_metadata._positions[:num_tokens].copy_(positions)
+        first_chunked_context = num_contexts - num_chunked_contexts
+        corrected_cached_token_lens = cached_token_lens.clone()
+        corrected_cached_token_lens[first_chunked_context:].add_(offsets)
+        corrected_positions = positions.clone()
+        num_chunked_context_tokens = (num_chunked_contexts *
+                                      tokens_per_generation)
+        corrected_positions[-num_chunked_context_tokens:].add_(
+            offsets.repeat_interleave(tokens_per_generation))
+        inputs = {
+            "input_ids":
+            torch.zeros(num_tokens, dtype=torch.int32, device="cuda"),
+            "position_ids":
+            torch.zeros((1, num_tokens), dtype=torch.int32, device="cuda"),
+            "attn_metadata":
+            attn_metadata,
+        }
+
+        model_engine._preprocess_inputs(inputs)
+
+        # The leading context stays unchanged; both trailing rows are rewound.
+        torch.testing.assert_close(
+            attn_metadata._cached_token_lens[:num_contexts],
+            corrected_cached_token_lens)
+        torch.testing.assert_close(attn_metadata._positions[:num_tokens],
+                                   corrected_positions)
+
+        model_engine._postprocess_inputs(inputs)
+
+        # Restore both the cached-length and position corrections.
+        torch.testing.assert_close(
+            attn_metadata._cached_token_lens[:num_contexts], cached_token_lens)
+        torch.testing.assert_close(attn_metadata._positions[:num_tokens],
+                                   positions)
+        kv_cache_manager.shutdown()
+
     def test_pad_generation_requests(self) -> None:
         model_engine, kv_cache_manager = create_model_engine_and_kvcache()
         resource_manager = ResourceManager(
diff --git a/tests/unittest/_torch/sampler/test_torch_sampler.py b/tests/unittest/_torch/sampler/test_torch_sampler.py
index b9dcd85d195a..ced29552e8af 100644
--- a/tests/unittest/_torch/sampler/test_torch_sampler.py
+++ b/tests/unittest/_torch/sampler/test_torch_sampler.py
@@ -759,6 +759,7 @@ def _uut(res=res):
 
 def test_stable_greedy_cache_key_includes_sequence_slots(monkeypatch: pytest.MonkeyPatch):
     sampler = object.__new__(TorchSampler)
+    sampler.max_tokens = 1
     sampler.max_beam_width = 1
     sampler._stable_greedy_request_ids = []
     sampler._stable_greedy_seq_slots = []
@@ -818,6 +819,50 @@ def copy_without_cuda(tensor: torch.Tensor, *args: Any, **kwargs: Any) -> torch.
         assert new_tokens[0, seq_slot, 0].item() == 2
 
 
+@force_ampere
+@pytest.mark.parametrize(("is_draft", "expected_stable"), [(False, False), (True, True)])
+def test_speculative_sampler_stable_greedy_requires_draft_batch(
+    is_draft: bool, expected_stable: bool
+):
+    sampler = TorchSampler(
+        TorchSampler.Args(
+            max_seq_len=16,
+            max_draft_len=3,
+            max_num_sequences=1,
+            max_beam_width=1,
+            max_total_draft_tokens=3,
+            disable_overlap_scheduler=False,
+        )
+    )
+    request = LlmRequest(
+        request_id=0,
+        max_new_tokens=4,
+        input_tokens=[1],
+        sampling_config=SamplingConfig(),
+        seq_slot=0,
+        is_streaming=False,
+        is_draft=is_draft,
+    )
+    admission = ScheduledRequests()
+    admission.context_requests_last_chunk = [request]
+    sampler.setup_sampler_step(admission)
+
+    scheduled_requests = ScheduledRequests()
+    scheduled_requests.generation_requests = [request]
+    logits = torch.tensor([[0.0, 1.0, 2.0]], device="cuda")
+
+    *_, new_tokens_host, single_step_greedy = sampler._process_requests(
+        scheduled_requests,
+        {"logits": logits},
+        sampler.store.new_tokens,
+        [0],
+    )
+    torch.cuda.synchronize()
+
+    assert single_step_greedy is expected_stable
+    assert new_tokens_host.reshape(-1)[0].item() == 2
+
+
 @force_ampere
 def test_greedy_no_repeat_ngram_uses_token_ban_path():
     sampler = TorchSampler(
diff --git a/tests/unittest/auto_deploy/singlegpu/shim/test_create_ad_executor.py b/tests/unittest/auto_deploy/singlegpu/shim/test_create_ad_executor.py
index afeef5471baf..be5f3a12e558 100644
--- a/tests/unittest/auto_deploy/singlegpu/shim/test_create_ad_executor.py
+++ b/tests/unittest/auto_deploy/singlegpu/shim/test_create_ad_executor.py
@@ -25,7 +25,7 @@
 from tensorrt_llm._torch.auto_deploy.llm_args import LlmArgs
 from tensorrt_llm._torch.auto_deploy.shim.ad_executor import create_autodeploy_executor
 from tensorrt_llm._torch.pyexecutor.kv_cache_transceiver import AttentionTypeCpp
-from tensorrt_llm.llmapi import CacheTransceiverConfig
+from tensorrt_llm.llmapi import CacheTransceiverConfig, MTPDecodingConfig
 
 pytestmark = pytest.mark.cpu_only
 
@@ -93,6 +93,7 @@ class MockPyExecutor:
     max_beam_width: int
     max_draft_len: int
     max_total_draft_tokens: int
+    max_seq_len: int
     guided_decoder: Any
     kv_cache_transceiver: Any = None
     resource_governor_queue: Any = None
@@ -133,6 +134,25 @@ def make_mock_engine(
     return mock_engine, kv_cache_manager
 
 
+def test_create_executor_supports_mtp_eagle_one_model_with_resolved_max_seq_len() -> None:
+    ad_config = LlmArgs(
+        model="test-model",
+        max_seq_len=128,
+        speculative_config=MTPDecodingConfig(max_draft_len=3, mtp_eagle_one_model=True),
+        transforms={"compile_model": {"piecewise_enabled": False}},
+    )
+    resolved_max_seq_len = 256
+    mock_engine, _ = make_mock_engine(max_seq_len=resolved_max_seq_len)
+
+    with _mock_py_executor_creation(mock_engine) as py_executor_cls:
+        result = create_autodeploy_executor(ad_config)
+
+    py_executor_cls.assert_called_once()
+    assert result.max_draft_len == 3
+    assert result.max_total_draft_tokens == 3
+    assert result.max_seq_len == resolved_max_seq_len
+
+
 @contextmanager
 def _mock_ad_engine_build(mock_engine, *, vocab_size_padded: int = 1000):
     with (