From a6fabea27478f051801718ea420c493b125c2618 Mon Sep 17 00:00:00 2001
From: William Zhang <133824995+2ez4bz@users.noreply.github.com>
Date: Tue, 4 Aug 2026 12:38:21 -0700
Subject: [PATCH 1/3] [None][fix] Fix one-model MTP KV cache accounting

* Why?

Overlapped one-model MTP could overstate the cached length or expose too
few FlashInfer pages during generation. This could access an invalid page
and cause a CUDA illegal memory access.

Draft-token accounting could also underreserve scheduler capacity or
produce out-of-range positions near the sequence limit.

* What?

Keep Python and C++ draft-token state synchronized and disable drafting
when the target positions could exceed the sequence limit.

Use the prompt KV boundary for the first generation step when dynamic
length correction is unavailable. For FlashInfer, expose all reserved
generation pages while maintaining the logical KV lengths independently.

Signed-off-by: William Zhang <133824995+2ez4bz@users.noreply.github.com>
---
 .../_torch/attention_backend/flashinfer.py    | 211 ++++++++++++++--
 .../_torch/auto_deploy/shim/ad_executor.py    |   9 +-
 tensorrt_llm/_torch/metadata.py               |   5 +
 .../_torch/pyexecutor/model_engine.py         |  59 ++++-
 tensorrt_llm/_torch/pyexecutor/py_executor.py |  41 +++
 .../_torch/pyexecutor/resource_manager.py     |   6 +
 .../_torch/pyexecutor/sampler/sampler.py      |   6 +
 .../attention/test_flashinfer_attention.py    | 236 ++++++++++++++++--
 .../_torch/executor/test_py_executor.py       | 116 ++++++++-
 .../executor/test_pytorch_model_engine.py     | 162 ++++++++++++
 .../_torch/sampler/test_torch_sampler.py      |  45 ++++
 .../singlegpu/shim/test_create_ad_executor.py |  21 +-
 12 files changed, 858 insertions(+), 59 deletions(-)

diff --git a/tensorrt_llm/_torch/attention_backend/flashinfer.py b/tensorrt_llm/_torch/attention_backend/flashinfer.py
index d49ab084a128..dd8d1374f4da 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,22 @@ 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,
+                         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])
 
-        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 +1009,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 +1039,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 +1527,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 +1611,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 +1656,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 +2014,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 +2046,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..d44419c09a73 100644
--- a/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py
+++ b/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py
@@ -1163,6 +1163,13 @@ def create_autodeploy_executor(
 
     This is the entrypoint API to the _autodeploy backend.
     """
+    spec_config = ad_config.speculative_config
+    if spec_config is not None and spec_config.spec_dec_mode.is_mtp_eagle_one_model():
+        raise NotImplementedError(
+            "AutoDeploy does not support MTP Eagle one-model speculative decoding because its "
+            "engine does not provide the draft-length state required by PyExecutor."
+        )
+
     # initialize process groups
     world_size = mpi_world_size()
     rank = mpi_rank()
@@ -1193,8 +1200,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."
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 4906ad52f5e5..760937be5936 100644
--- a/tensorrt_llm/_torch/pyexecutor/model_engine.py
+++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py
@@ -3347,6 +3347,14 @@ 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."""
+        return (self.enable_spec_decode and not self._disable_overlap_scheduler
+                and getattr(spec_config, '_use_shared_kv_cache', False)
+                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
@@ -3398,6 +3406,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()
@@ -3445,6 +3462,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:
@@ -4154,7 +4181,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()
 
@@ -5143,11 +5173,25 @@ 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)
+                # The first generation batch overlaps the context sampler, so there is a previous
+                # token tensor but no previous speculative target forward in the KV cache.
+                # Backends without dynamic KV lengths cannot apply the runtime acceptance-length
+                # correction in _preprocess_inputs and must start at the prompt boundary.
+                # py_decoding_iter is a proxy for "the previous forward for this request was its
+                # last context chunk": _update_requests lags _prepare_and_schedule_batch by exactly
+                # one iteration, so the sampler's first increment has not landed yet at this point
+                # and only here. This branch already assumes that same batch-to-batch continuity
+                # (previous_batch_idx indexes the immediately preceding batch's device tensors).
+                if (request.py_decoding_iter == 0
+                        and not hasattr(attn_metadata, "kv_lens_cuda")
+                        and not hasattr(attn_metadata,
+                                        "apply_spec_decode_kv_lens_offsets")):
+                    cached_token_num = request.max_beam_num_tokens
                 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)
@@ -5756,7 +5800,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'):
diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py
index 60d801d904c7..fa031768d867 100644
--- a/tensorrt_llm/_torch/pyexecutor/py_executor.py
+++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py
@@ -3429,6 +3429,38 @@ def _handle_dynamic_draft_len(self,
                 if spec_config is not None and spec_config.is_linear_tree else
                 self.model_engine.max_total_draft_tokens)
 
+        if self._one_model_mtp_batch_needs_zero_draft(scheduled_batch):
+            # The target input width is batch-wide, so one unsafe request must
+            # disable drafting for every generation request in this batch.
+            for request in scheduled_batch.generation_requests:
+                request.py_draft_tokens = []
+            self.model_engine.runtime_draft_len = 0
+
+    def _one_model_mtp_batch_needs_zero_draft(
+            self, scheduled_batch: ScheduledRequests) -> bool:
+        """Return whether drafting could produce an out-of-range position."""
+        spec_config = self.model_engine.spec_config
+        runtime_draft_len = self.model_engine.runtime_draft_len
+        if (runtime_draft_len == 0 or spec_config is None
+                or not spec_config.spec_dec_mode.is_mtp_eagle_one_model()):
+            return False
+
+        # 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)
+
+        for request in scheduled_batch.generation_requests:
+            max_target_position = (request.max_beam_num_tokens - 1 +
+                                   max_pending_tokens + target_position_width -
+                                   1)
+            if max_target_position >= self.max_seq_len:
+                return True
+        return False
+
     @nvtx_range("_can_queue")
     def _can_queue(self, scheduled_batch):
 
@@ -3838,6 +3870,15 @@ def _prepare_and_schedule_batch(self):
                         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` - and would also suppress the
+                # `current_num_draft_tokens == 0` signal that `_handle_dynamic_draft_len` uses to
+                # request a one-hot draft-probs placeholder under rejection sampling.
+                if not request.py_draft_tokens:
+                    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/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py
index 2ea7d89a6fbc..5d1c7949e538 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 ed58c84ac854..cc7f838bfcf2 100644
--- a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py
+++ b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py
@@ -4088,6 +4088,12 @@ def _process_requests(
         )
         can_use_stable_greedy_path = (
             bool(generation_requests)
+            # A speculative target can temporarily carry no drafts when one-model MTP
+            # disables drafting near the sequence limit. Keep it on the speculative/device
+            # finish-state path; explicit draft-model batches are genuinely single-step.
+            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..839450cc9927 100644
--- a/tests/unittest/_torch/attention/test_flashinfer_attention.py
+++ b/tests/unittest/_torch/attention/test_flashinfer_attention.py
@@ -2,6 +2,7 @@
 import unittest
 from collections import defaultdict
 from dataclasses import dataclass
+from types import SimpleNamespace
 from typing import List, Optional, Union
 from unittest import mock
 
@@ -13,7 +14,8 @@
                                                    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
@@ -64,31 +66,226 @@ class CUDAGraphTestScenario:
     dtype: torch.dtype
 
 
+def _create_kv_cache_manager() -> KVCacheManager:
+    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,
+    )
+
+
 class TestFlashInferAttention(unittest.TestCase):
 
+    def test_generation_page_table_uses_reserved_block_count(self):
+        manager = SimpleNamespace(kv_cache_map={
+            98: SimpleNamespace(num_blocks=4),
+            99: SimpleNamespace(num_blocks=325),
+        }, )
+
+        self.assertEqual(
+            flashinfer_backend._get_page_table_num_blocks(manager, [98, 99],
+                                                          [3, 324],
+                                                          num_contexts=1),
+            [3, 325],
+        )
+        v1_manager = SimpleNamespace(get_batch_cache_indices=mock.Mock(
+            return_value=[list(range(4)), list(range(325))]))
+        self.assertEqual(
+            flashinfer_backend._get_page_table_num_blocks(v1_manager, [98, 99],
+                                                          [3, 324],
+                                                          num_contexts=1),
+            [3, 325],
+        )
+        v1_manager.get_batch_cache_indices.assert_called_once_with([98, 99])
+
+    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:
+            metadata = FlashInferAttentionMetadata(
+                seq_lens=torch.tensor([2, 4, 4], dtype=torch.int32),
+                num_contexts=1,
+                kv_cache_params=KVCacheParams(use_cache=True),
+                max_num_requests=3,
+                max_num_tokens=10,
+                kv_cache_manager=kv_cache_manager,
+            )
+            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[:3].copy_(cached_token_lens)
+            metadata._paged_kv_last_page_len[:3].copy_(last_page_lens)
+            metadata._positions[:10].copy_(positions)
+            kv_lens_buffer = torch.tensor([35, 66],
+                                          dtype=torch.int32,
+                                          device="cuda")
+            metadata._plan_params_to_wrappers = {
+                object():
+                FlashInferWrappers(is_planned=True,
+                                   decode_wrapper=SimpleNamespace(
+                                       _kv_lens_buffer=kv_lens_buffer))
+            }
+            offsets = torch.tensor([-3, -1], dtype=torch.int32, device="cuda")
+
+            metadata.apply_spec_decode_kv_lens_offsets(offsets,
+                                                       num_generations=2,
+                                                       tokens_per_generation=4)
+
+            torch.testing.assert_close(
+                metadata._cached_token_lens[:3],
+                torch.tensor([10, 28, 61], dtype=torch.int32, device="cuda"))
+            torch.testing.assert_close(metadata._paged_kv_last_page_len[:3],
+                                       last_page_lens)
+            torch.testing.assert_close(
+                metadata._positions[:10],
+                torch.tensor([10, 11, 28, 29, 30, 31, 61, 62, 63, 64],
+                             dtype=torch.int32,
+                             device="cuda"))
+            torch.testing.assert_close(
+                kv_lens_buffer,
+                torch.tensor([32, 65], dtype=torch.int32, device="cuda"))
+
+            # 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_(
+                torch.tensor([35, 66], dtype=torch.int32, device="cuda"))
+            metadata._publish_decode_wrapper_kv_lens(
+                next(iter(
+                    metadata._plan_params_to_wrappers.values())).decode_wrapper)
+            torch.testing.assert_close(
+                kv_lens_buffer,
+                torch.tensor([32, 65], dtype=torch.int32, device="cuda"))
+
+            metadata.apply_spec_decode_kv_lens_offsets(
+                offsets,
+                num_generations=2,
+                tokens_per_generation=4,
+                restore=True,
+            )
+
+            torch.testing.assert_close(metadata._cached_token_lens[:3],
+                                       cached_token_lens)
+            torch.testing.assert_close(metadata._paged_kv_last_page_len[:3],
+                                       last_page_lens)
+            torch.testing.assert_close(metadata._positions[:10], positions)
+            torch.testing.assert_close(
+                kv_lens_buffer,
+                torch.tensor([35, 66], dtype=torch.int32, device="cuda"))
+
+            # A shared external-assistant view is Q-only: its cached length is
+            # the full accepted target prefix and must not include its query token.
+            metadata.seq_lens = torch.ones(2, dtype=torch.int32)
+            metadata.num_contexts = 0
+            metadata._is_shared_kv_draft_view = True
+            metadata._draft_kv_runtime_lens[:2].copy_(
+                torch.tensor([32, 65], dtype=torch.int32, device="cuda"))
+            metadata._cached_token_lens[:2].zero_()
+            metadata._paged_kv_last_page_len[:2].zero_()
+            kv_lens_buffer.copy_(
+                torch.tensor([35, 66], dtype=torch.int32, device="cuda"))
+            metadata._update_draft_kv_lengths()
+            torch.testing.assert_close(
+                kv_lens_buffer,
+                torch.tensor([32, 65], dtype=torch.int32, device="cuda"))
+
+    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):
+                self.plan_args = args
+
+        kv_cache_manager = _create_kv_cache_manager()
+        metadata = FlashInferAttentionMetadata(
+            seq_lens=torch.ones(2, dtype=torch.int32),
+            num_contexts=0,
+            kv_cache_params=KVCacheParams(use_cache=True),
+            max_num_requests=2,
+            max_num_tokens=2,
+            kv_cache_manager=kv_cache_manager,
+        )
+        metadata._uses_full_generation_page_table = True
+        metadata.num_generation_blocks = 5
+        metadata.num_context_blocks = 0
+        metadata.paged_kv_indptr_decode[:3].copy_(
+            torch.tensor([0, 2, 5], dtype=torch.int32, device="cuda"))
+        metadata._paged_kv_indices[:5].copy_(
+            torch.arange(5, dtype=torch.int32, device="cuda"))
+        metadata._paged_kv_last_page_len[:2].copy_(
+            torch.tensor([5, 5], dtype=torch.int32, device="cuda"))
+        metadata._qo_indptr[:3].copy_(
+            torch.tensor([0, 1, 2], dtype=torch.int32, device="cuda"))
+        metadata._cached_token_lens[:2].copy_(
+            torch.tensor([7, 13], dtype=torch.int32, device="cuda"))
+        logical_kv_lens = torch.tensor([8, 14],
+                                       dtype=torch.int32,
+                                       device="cuda")
+        metadata._logical_kv_lens[:2].copy_(logical_kv_lens)
+        metadata._positions[:2].copy_(logical_kv_lens)
+        offsets = torch.tensor([-3, -1], dtype=torch.int32, device="cuda")
+
+        metadata.apply_spec_decode_kv_lens_offsets(offsets,
+                                                   num_generations=2,
+                                                   tokens_per_generation=1)
+
+        torch.testing.assert_close(
+            metadata._logical_kv_lens[:2],
+            torch.tensor([5, 13], dtype=torch.int32, device="cuda"))
+
+        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 reservation-width table would derive [13, 21]; MLA must see the live logical lengths.
+        torch.testing.assert_close(
+            wrapper.plan_args[3],
+            torch.tensor([5, 13], dtype=torch.int32, device="cuda"))
+
+        metadata.apply_spec_decode_kv_lens_offsets(
+            offsets,
+            num_generations=2,
+            tokens_per_generation=1,
+            restore=True,
+        )
+        torch.testing.assert_close(metadata._logical_kv_lens[:2],
+                                   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 +346,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 161ac02eb995..8cb520eed721 100644
--- a/tests/unittest/_torch/executor/test_py_executor.py
+++ b/tests/unittest/_torch/executor/test_py_executor.py
@@ -43,6 +43,7 @@
     ScheduledRequests,
     SerializableSchedulerOutput,
 )
+from tensorrt_llm.llmapi.llm_args import MTPDecodingConfig
 from tensorrt_llm.runtime.kv_cache_manager_v2 import OutOfPagesError
 
 pytestmark = pytest.mark.cpu_only
@@ -2487,12 +2488,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
@@ -2569,6 +2575,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()
@@ -2578,7 +2586,103 @@ 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_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
+
+    @classmethod
+    def _make_runtime_draft_executor(cls, disable_overlap_scheduler: bool):
+        ex = object.__new__(PyExecutor)
+        spec_config = MTPDecodingConfig(
+            max_draft_len=cls.MAX_TOTAL_DRAFT_TOKENS,
+            mtp_eagle_one_model=True,
+        )
+        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.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,sequence_length",
+        [
+            (False, 10),
+            (True, 14),
+        ],
+    )
+    def test_one_model_mtp_uses_zero_draft_near_sequence_limit(
+        self, disable_overlap_scheduler, sequence_length
+    ):
+        ex = self._make_runtime_draft_executor(disable_overlap_scheduler)
+        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,sequence_length",
+        [
+            (False, 9),
+            (True, 13),
+        ],
+    )
+    def test_one_model_mtp_keeps_drafting_with_position_headroom(
+        self, disable_overlap_scheduler, sequence_length
+    ):
+        ex = self._make_runtime_draft_executor(disable_overlap_scheduler)
+        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]
diff --git a/tests/unittest/_torch/executor/test_pytorch_model_engine.py b/tests/unittest/_torch/executor/test_pytorch_model_engine.py
index ed0f3bd4d00b..fdd3782e2cb6 100644
--- a/tests/unittest/_torch/executor/test_pytorch_model_engine.py
+++ b/tests/unittest/_torch/executor/test_pytorch_model_engine.py
@@ -36,7 +36,9 @@
 # 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
@@ -1373,6 +1375,166 @@ 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_first_speculative_generation_uses_prompt_kv_boundary(self) -> None:
+        spec_config = SADecodingConfig(
+            max_draft_len=3,
+            draft_len_schedule={1: 3},
+        )
+        model_engine, kv_cache_manager = create_model_engine_and_kvcache(
+            spec_config=spec_config)
+        model_engine.runtime_draft_len = 3
+        resource_manager = ResourceManager(
+            {ResourceManagerType.KV_CACHE_MANAGER: kv_cache_manager})
+        attn_metadata = AttentionMetadata(max_num_requests=4,
+                                          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
+        # The final context batch populated this overlap slot, but it did not
+        # run a speculative target step.
+        generation.py_batch_idx = 0
+        generation.py_draft_tokens = [0, 0, 0]
+        self.assertEqual(generation.py_decoding_iter, 0)
+
+        graph_batch = ScheduledRequests()
+        graph_batch.generation_requests = [generation]
+        overlap_state = SampleStateTensorsSpec(
+            new_tokens=torch.zeros((4, 4, 1), dtype=torch.int32, device="cuda"),
+            new_tokens_lens=torch.ones(4, dtype=torch.int32, device="cuda"),
+            next_draft_tokens=torch.zeros((4, 3),
+                                          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,
+        )
+
+        prompt_len = generation.max_beam_num_tokens
+        self.assertEqual(
+            attn_metadata.kv_cache_params.num_cached_tokens_per_seq,
+            [prompt_len])
+        self.assertEqual(generation.cached_tokens, prompt_len)
+        model_engine._preprocess_inputs(inputs)
+        self.assertEqual(inputs["position_ids"][0, :4].cpu().tolist(),
+                         list(range(prompt_len, prompt_len + 4)))
+        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)
+        model_engine.runtime_draft_len = spec_config.max_draft_len
+        model_engine.previous_pos_id_offsets_cuda = torch.zeros(
+            8, dtype=torch.int32, device="cuda")
+        model_engine.previous_kv_lens_offsets_cuda = torch.tensor(
+            [-3, -1], dtype=torch.int32, device="cuda")
+
+        attn_metadata = FlashInferAttentionMetadata(
+            seq_lens=torch.tensor([2, 4, 4], dtype=torch.int32),
+            num_contexts=1,
+            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")
+        attn_metadata._cached_token_lens[:3].copy_(cached_token_lens)
+        attn_metadata._positions[:10].copy_(
+            torch.tensor([10, 11, 31, 32, 33, 34, 62, 63, 64, 65],
+                         dtype=torch.int32,
+                         device="cuda"))
+        inputs = {
+            "input_ids": torch.zeros(10, dtype=torch.int32, device="cuda"),
+            "position_ids": torch.zeros((1, 10),
+                                        dtype=torch.int32,
+                                        device="cuda"),
+            "attn_metadata": attn_metadata,
+        }
+
+        model_engine._preprocess_inputs(inputs)
+
+        torch.testing.assert_close(
+            attn_metadata._cached_token_lens[:3],
+            torch.tensor([10, 28, 61], dtype=torch.int32, device="cuda"))
+        torch.testing.assert_close(
+            attn_metadata._positions[:10],
+            torch.tensor([10, 11, 28, 29, 30, 31, 61, 62, 63, 64],
+                         dtype=torch.int32,
+                         device="cuda"))
+
+        model_engine._postprocess_inputs(inputs)
+
+        torch.testing.assert_close(attn_metadata._cached_token_lens[:3],
+                                   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)
+        model_engine.runtime_draft_len = spec_config.max_draft_len
+        model_engine.previous_pos_id_offsets_cuda = torch.zeros(
+            8, dtype=torch.int32, device="cuda")
+        model_engine.previous_kv_lens_offsets_cuda = torch.tensor(
+            [-3, -1], dtype=torch.int32, device="cuda")
+
+        attn_metadata = FlashInferAttentionMetadata(
+            seq_lens=torch.tensor([2, 4, 4], dtype=torch.int32),
+            num_contexts=3,
+            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 = 2
+        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[:3].copy_(cached_token_lens)
+        attn_metadata._positions[:10].copy_(positions)
+        inputs = {
+            "input_ids": torch.zeros(10, dtype=torch.int32, device="cuda"),
+            "position_ids": torch.zeros((1, 10),
+                                        dtype=torch.int32,
+                                        device="cuda"),
+            "attn_metadata": attn_metadata,
+        }
+
+        model_engine._preprocess_inputs(inputs)
+
+        torch.testing.assert_close(
+            attn_metadata._cached_token_lens[:3],
+            torch.tensor([10, 28, 61], dtype=torch.int32, device="cuda"))
+        torch.testing.assert_close(
+            attn_metadata._positions[:10],
+            torch.tensor([10, 11, 28, 29, 30, 31, 61, 62, 63, 64],
+                         dtype=torch.int32,
+                         device="cuda"))
+
+        model_engine._postprocess_inputs(inputs)
+
+        torch.testing.assert_close(attn_metadata._cached_token_lens[:3],
+                                   cached_token_lens)
+        torch.testing.assert_close(attn_metadata._positions[:10], 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 4a29f602bba6..31f5ca5b215b 100644
--- a/tests/unittest/_torch/sampler/test_torch_sampler.py
+++ b/tests/unittest/_torch/sampler/test_torch_sampler.py
@@ -753,6 +753,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 = []
@@ -812,6 +813,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..a134923bfeb0 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
 
@@ -133,6 +133,25 @@ def make_mock_engine(
     return mock_engine, kv_cache_manager
 
 
+def test_create_executor_rejects_mtp_eagle_one_model_before_runtime_setup():
+    ad_config = LlmArgs(
+        model="test-model",
+        speculative_config=MTPDecodingConfig(max_draft_len=3, mtp_eagle_one_model=True),
+        transforms={"compile_model": {"piecewise_enabled": False}},
+    )
+
+    with (
+        patch("tensorrt_llm._torch.auto_deploy.shim.ad_executor.mpi_world_size") as mpi_world_size,
+        pytest.raises(
+            NotImplementedError,
+            match="AutoDeploy does not support MTP Eagle one-model speculative decoding",
+        ),
+    ):
+        create_autodeploy_executor(ad_config)
+
+    mpi_world_size.assert_not_called()
+
+
 @contextmanager
 def _mock_ad_engine_build(mock_engine, *, vocab_size_padded: int = 1000):
     with (

From 94fb07b7f520c30ad30be8f51fddd9d641a6ddf9 Mon Sep 17 00:00:00 2001
From: William Zhang <133824995+2ez4bz@users.noreply.github.com>
Date: Wed, 12 Aug 2026 16:22:05 -0700
Subject: [PATCH 2/3] [None][fix] Guard one-model MTP draft positions

* Why?
One-model MTP can consume positions beyond the target verification
span near the maximum sequence length. The shared-KV Gemma4 assistant
could therefore index one rotary row past the model limit and poison
the CUDA context.

* What?
Include shared and non-shared drafter lookahead in the zero-draft
boundary check, with coverage for overlap and non-overlap scheduling.

Signed-off-by: William Zhang <133824995+2ez4bz@users.noreply.github.com>
---
 tensorrt_llm/_torch/pyexecutor/py_executor.py | 15 +++++++---
 .../_torch/executor/test_py_executor.py       | 29 ++++++++++++-------
 2 files changed, 29 insertions(+), 15 deletions(-)

diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py
index fa031768d867..bec9d953f849 100644
--- a/tensorrt_llm/_torch/pyexecutor/py_executor.py
+++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py
@@ -3452,12 +3452,19 @@ def _one_model_mtp_batch_needs_zero_draft(
                               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_target_position = (request.max_beam_num_tokens - 1 +
-                                   max_pending_tokens + target_position_width -
-                                   1)
-            if max_target_position >= self.max_seq_len:
+            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
 
diff --git a/tests/unittest/_torch/executor/test_py_executor.py b/tests/unittest/_torch/executor/test_py_executor.py
index 8cb520eed721..a6551bd0ab8f 100644
--- a/tests/unittest/_torch/executor/test_py_executor.py
+++ b/tests/unittest/_torch/executor/test_py_executor.py
@@ -2618,12 +2618,15 @@ def test_one_model_mtp_preserves_sampler_draft_tokens(self):
         assert gen.num_draft_tokens == self.MAX_TOTAL_DRAFT_TOKENS
 
     @classmethod
-    def _make_runtime_draft_executor(cls, disable_overlap_scheduler: bool):
+    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,
@@ -2652,16 +2655,18 @@ def _make_generation_batch(cls, *sequence_lengths: int):
         return batch
 
     @pytest.mark.parametrize(
-        "disable_overlap_scheduler,sequence_length",
+        "disable_overlap_scheduler,use_shared_target_kv,sequence_length",
         [
-            (False, 10),
-            (True, 14),
+            (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, sequence_length
+        self, disable_overlap_scheduler, use_shared_target_kv, sequence_length
     ):
-        ex = self._make_runtime_draft_executor(disable_overlap_scheduler)
+        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)
@@ -2670,16 +2675,18 @@ def test_one_model_mtp_uses_zero_draft_near_sequence_limit(
         assert all(request.py_draft_tokens == [] for request in batch.generation_requests)
 
     @pytest.mark.parametrize(
-        "disable_overlap_scheduler,sequence_length",
+        "disable_overlap_scheduler,use_shared_target_kv,sequence_length",
         [
-            (False, 9),
-            (True, 13),
+            (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, sequence_length
+        self, disable_overlap_scheduler, use_shared_target_kv, sequence_length
     ):
-        ex = self._make_runtime_draft_executor(disable_overlap_scheduler)
+        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)

From 973fd7dbd294bd3f4bd4b200b043485045675acc Mon Sep 17 00:00:00 2001
From: William Zhang <133824995+2ez4bz@users.noreply.github.com>
Date: Wed, 12 Aug 2026 22:32:20 -0700
Subject: [PATCH 3/3] Address review comments

Signed-off-by: William Zhang <133824995+2ez4bz@users.noreply.github.com>
---
 .../_torch/auto_deploy/shim/ad_executor.py    |   6 +-
 .../_torch/pyexecutor/model_engine.py         |  29 +---
 tensorrt_llm/_torch/pyexecutor/py_executor.py | 141 +++++++++---------
 .../_torch/pyexecutor/py_executor_creator.py  |  29 +++-
 .../_torch/executor/test_py_executor.py       |  95 +++++++++++-
 .../executor/test_pytorch_model_engine.py     |  48 +++---
 .../singlegpu/shim/test_create_ad_executor.py |  21 +--
 7 files changed, 235 insertions(+), 134 deletions(-)

diff --git a/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py b/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py
index d44419c09a73..e783796f758e 100644
--- a/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py
+++ b/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py
@@ -1164,11 +1164,6 @@ def create_autodeploy_executor(
     This is the entrypoint API to the _autodeploy backend.
     """
     spec_config = ad_config.speculative_config
-    if spec_config is not None and spec_config.spec_dec_mode.is_mtp_eagle_one_model():
-        raise NotImplementedError(
-            "AutoDeploy does not support MTP Eagle one-model speculative decoding because its "
-            "engine does not provide the draft-length state required by PyExecutor."
-        )
 
     # initialize process groups
     world_size = mpi_world_size()
@@ -1341,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/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py
index 760937be5936..3446366e93c8 100644
--- a/tensorrt_llm/_torch/pyexecutor/model_engine.py
+++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py
@@ -5080,9 +5080,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
@@ -5094,7 +5096,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
@@ -5175,20 +5177,6 @@ def append_cross_attention_state(request: LlmRequest,
 
                 cached_token_num = (past_seen_token_num +
                                     runtime_tokens_per_gen_step)
-                # The first generation batch overlaps the context sampler, so there is a previous
-                # token tensor but no previous speculative target forward in the KV cache.
-                # Backends without dynamic KV lengths cannot apply the runtime acceptance-length
-                # correction in _preprocess_inputs and must start at the prompt boundary.
-                # py_decoding_iter is a proxy for "the previous forward for this request was its
-                # last context chunk": _update_requests lags _prepare_and_schedule_batch by exactly
-                # one iteration, so the sampler's first increment has not landed yet at this point
-                # and only here. This branch already assumes that same batch-to-batch continuity
-                # (previous_batch_idx indexes the immediately preceding batch's device tensors).
-                if (request.py_decoding_iter == 0
-                        and not hasattr(attn_metadata, "kv_lens_cuda")
-                        and not hasattr(attn_metadata,
-                                        "apply_spec_decode_kv_lens_offsets")):
-                    cached_token_num = request.max_beam_num_tokens
                 num_cached_tokens_per_seq.append(
                     cached_token_num - request.py_num_compressed_tokens)
                 request.cached_tokens = cached_token_num
@@ -7109,9 +7097,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 bec9d953f849..9679b5ded105 100644
--- a/tensorrt_llm/_torch/pyexecutor/py_executor.py
+++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py
@@ -3359,91 +3359,92 @@ 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 self._one_model_mtp_batch_needs_zero_draft(scheduled_batch):
+        if needs_zero_draft:
             # The target input width is batch-wide, so one unsafe request must
-            # disable drafting for every generation request in this batch.
+            # disable drafting for every generation request on every ADP rank.
             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) -> bool:
+            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
-        runtime_draft_len = self.model_engine.runtime_draft_len
-        if (runtime_draft_len == 0 or spec_config is None
+        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.
@@ -3872,6 +3873,9 @@ 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,
@@ -3881,10 +3885,11 @@ def _prepare_and_schedule_batch(self):
                 # 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` - and would also suppress the
-                # `current_num_draft_tokens == 0` signal that `_handle_dynamic_draft_len` uses to
-                # request a one-hot draft-probs placeholder under rejection sampling.
+                # 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
 
diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
index 680d607efeb5..e36609b603e3 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
+
+    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 llm_args.attn_backend.upper() == "VANILLA":
+        reason = "VANILLA attention lacks dynamic speculative KV lengths"
+    elif (llm_args.attn_backend.upper() == "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,12 +477,7 @@ 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"
             and spec_config.spec_dec_mode.use_one_engine()
diff --git a/tests/unittest/_torch/executor/test_py_executor.py b/tests/unittest/_torch/executor/test_py_executor.py
index a6551bd0ab8f..98fe70098f3b 100644
--- a/tests/unittest/_torch/executor/test_py_executor.py
+++ b/tests/unittest/_torch/executor/test_py_executor.py
@@ -2525,7 +2525,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
@@ -2538,11 +2542,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 = []
 
@@ -2594,6 +2612,26 @@ def test_one_model_mtp_populates_draft_tokens_for_scheduling(self):
         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.
 
@@ -2634,6 +2672,8 @@ def _make_runtime_draft_executor(
         )
         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
 
@@ -2693,3 +2733,54 @@ def test_one_model_mtp_keeps_drafting_with_position_headroom(
 
         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 fdd3782e2cb6..9584b1a542f4 100644
--- a/tests/unittest/_torch/executor/test_pytorch_model_engine.py
+++ b/tests/unittest/_torch/executor/test_pytorch_model_engine.py
@@ -44,8 +44,8 @@
     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
 
 
@@ -1375,14 +1375,22 @@ 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_first_speculative_generation_uses_prompt_kv_boundary(self) -> None:
-        spec_config = SADecodingConfig(
+    def test_dynamic_tree_prepare_preserves_explicit_zero_draft(self) -> None:
+        max_total_draft_tokens = 12
+        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=3,
-            draft_len_schedule={1: 3},
+            mtp_eagle_one_model=True,
+            use_dynamic_tree=True,
+            dynamic_tree_max_topK=4,
         )
-        model_engine, kv_cache_manager = create_model_engine_and_kvcache(
-            spec_config=spec_config)
-        model_engine.runtime_draft_len = 3
+        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
+        model_engine.runtime_draft_len = 0
         resource_manager = ResourceManager(
             {ResourceManagerType.KV_CACHE_MANAGER: kv_cache_manager})
         attn_metadata = AttentionMetadata(max_num_requests=4,
@@ -1392,18 +1400,17 @@ def test_first_speculative_generation_uses_prompt_kv_boundary(self) -> None:
 
         generation = _create_request_with_tokens([50, 51, 52, 53, 54], 1)
         generation.py_seq_slot = 0
-        # The final context batch populated this overlap slot, but it did not
-        # run a speculative target step.
         generation.py_batch_idx = 0
-        generation.py_draft_tokens = [0, 0, 0]
-        self.assertEqual(generation.py_decoding_iter, 0)
+        generation.py_draft_tokens = []
 
         graph_batch = ScheduledRequests()
         graph_batch.generation_requests = [generation]
         overlap_state = SampleStateTensorsSpec(
-            new_tokens=torch.zeros((4, 4, 1), dtype=torch.int32, device="cuda"),
+            new_tokens=torch.zeros((max_total_draft_tokens + 1, 4, 1),
+                                   dtype=torch.int32,
+                                   device="cuda"),
             new_tokens_lens=torch.ones(4, dtype=torch.int32, device="cuda"),
-            next_draft_tokens=torch.zeros((4, 3),
+            next_draft_tokens=torch.zeros((4, max_total_draft_tokens),
                                           dtype=torch.int32,
                                           device="cuda"),
         )
@@ -1418,14 +1425,11 @@ def test_first_speculative_generation_uses_prompt_kv_boundary(self) -> None:
             resource_manager=resource_manager,
         )
 
-        prompt_len = generation.max_beam_num_tokens
-        self.assertEqual(
-            attn_metadata.kv_cache_params.num_cached_tokens_per_seq,
-            [prompt_len])
-        self.assertEqual(generation.cached_tokens, prompt_len)
-        model_engine._preprocess_inputs(inputs)
-        self.assertEqual(inputs["position_ids"][0, :4].cpu().tolist(),
-                         list(range(prompt_len, prompt_len + 4)))
+        self.assertEqual(model_engine.runtime_draft_len, 0)
+        self.assertEqual(model_engine.get_runtime_tokens_per_gen_step(0), 1)
+        self.assertEqual(attn_metadata.seq_lens.tolist(), [1])
+        self.assertEqual(spec_metadata.seq_lens, [1])
+        self.assertEqual(inputs["input_ids"].numel(), 1)
         kv_cache_manager.shutdown()
 
     def test_overlap_input_processing_applies_flashinfer_kv_offsets(
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 a134923bfeb0..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
@@ -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,23 +134,23 @@ def make_mock_engine(
     return mock_engine, kv_cache_manager
 
 
-def test_create_executor_rejects_mtp_eagle_one_model_before_runtime_setup():
+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 (
-        patch("tensorrt_llm._torch.auto_deploy.shim.ad_executor.mpi_world_size") as mpi_world_size,
-        pytest.raises(
-            NotImplementedError,
-            match="AutoDeploy does not support MTP Eagle one-model speculative decoding",
-        ),
-    ):
-        create_autodeploy_executor(ad_config)
+    with _mock_py_executor_creation(mock_engine) as py_executor_cls:
+        result = create_autodeploy_executor(ad_config)
 
-    mpi_world_size.assert_not_called()
+    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