diff --git a/tensorrt_llm/_torch/model_config.py b/tensorrt_llm/_torch/model_config.py
index fd15550fd6e1..7b445e991c21 100644
--- a/tensorrt_llm/_torch/model_config.py
+++ b/tensorrt_llm/_torch/model_config.py
@@ -141,6 +141,11 @@ class ModelConfig(Generic[TConfig]):
     skip_create_weights_in_init: bool = False
 
     spec_config: Optional["DecodingBaseConfig"] = None
+    # When False, the column-parallel LM head keeps its vocab-sharded output
+    # instead of all-gathering to full vocab. Used for one-model speculative
+    # draft models so greedy draft sampling can do a lighter TP gather. Defaults
+    # to True to preserve behavior for every non-draft model.
+    lm_head_gather_output: bool = True
     lora_config: Optional["LoraConfig"] = None
     sparse_attention_config: Optional["SparseAttentionConfig"] = None
 
diff --git a/tensorrt_llm/_torch/models/modeling_speculative.py b/tensorrt_llm/_torch/models/modeling_speculative.py
index 6e17258067a1..df685c4f2a27 100755
--- a/tensorrt_llm/_torch/models/modeling_speculative.py
+++ b/tensorrt_llm/_torch/models/modeling_speculative.py
@@ -777,7 +777,9 @@ def __init__(self, draft_config):
             draft_config.pretrained_config)
 
         # Remove spec_config to prevent recursive spec-dec initialization
-        draft_config_no_spec = replace(draft_config, spec_config=None)
+        draft_config_no_spec = replace(draft_config,
+                                       spec_config=None,
+                                       lm_head_gather_output=False)
 
         # Weights will be loaded later by ModelLoader.load_draft_weights()
         self.draft_model_full = DraftModelClass(draft_config_no_spec)
@@ -863,7 +865,9 @@ def __init__(self, draft_config):
                 pretrained_cfg.architectures = original_archs
 
         # Remove spec_config to prevent recursive spec-dec initialization
-        draft_config_no_spec = replace(draft_config, spec_config=None)
+        draft_config_no_spec = replace(draft_config,
+                                       spec_config=None,
+                                       lm_head_gather_output=False)
 
         # Weights will be loaded later by ModelLoader.load_draft_weights()
         self.draft_model_full = DraftModelClass(draft_config_no_spec)
@@ -1653,6 +1657,12 @@ def get_draft_model(model_config, draft_config, lm_head, model):
     elif spec_dec_mode.is_dflash():
         return DFlashForCausalLM(draft_config)
     elif spec_dec_mode.is_draft_target_one_model():
+        # Keep the draft LM head vocab-sharded so greedy draft sampling uses the
+        # lighter TP gather (see SpecWorkerBase.greedy_sample_draft_with_tp_gather).
+        was_frozen = draft_config._frozen
+        draft_config._frozen = False
+        draft_config.lm_head_gather_output = False
+        draft_config._frozen = was_frozen
         return AutoModelForCausalLM.from_config(draft_config)
     else:
         raise NotImplementedError(
@@ -1741,6 +1751,10 @@ def __init__(self, model: TModel, model_config: ModelConfig[TConfig]):
                 model_config.mapping,
                 use_separate_draft_kv_cache=self.use_separate_draft_kv_cache)
             if self.spec_worker is not None:
+                # Cache the static draft->target vocab map now that the draft
+                # model is loaded, so workers read self._d2t instead of probing
+                # draft_model.model.d2t on every forward.
+                self.spec_worker.set_draft_model(self.draft_model)
                 self.epilogue.append(self.spec_worker)
         self.layer_idx = -1
 
diff --git a/tensorrt_llm/_torch/models/modeling_utils.py b/tensorrt_llm/_torch/models/modeling_utils.py
index 5355bb7d3193..c45572f577f6 100755
--- a/tensorrt_llm/_torch/models/modeling_utils.py
+++ b/tensorrt_llm/_torch/models/modeling_utils.py
@@ -431,7 +431,7 @@ def __init__(self, model: TModel, *, config: ModelConfig[TConfig],
                 dtype=config.pretrained_config.torch_dtype,
                 mapping=config.mapping,
                 tensor_parallel_mode=TensorParallelMode.COLUMN,
-                gather_output=True,
+                gather_output=config.lm_head_gather_output,
                 reduce_output=False,
                 use_custom_cublas_mm=getattr(model, 'use_custom_cublas_mm',
                                              False),
diff --git a/tensorrt_llm/_torch/pyexecutor/llm_request.py b/tensorrt_llm/_torch/pyexecutor/llm_request.py
index 63bfb318e284..6d7522ec79a8 100644
--- a/tensorrt_llm/_torch/pyexecutor/llm_request.py
+++ b/tensorrt_llm/_torch/pyexecutor/llm_request.py
@@ -736,6 +736,10 @@ def __init__(
         self.py_target_probs = None
         self.py_last_draft_tokens = None
         self.py_num_accepted_draft_tokens = 0
+        # One-model rejection: set per-iteration by _handle_dynamic_draft_len when
+        # this gen request produced 0 real draft tokens, so _prepare_tp_inputs
+        # one-hots its stale draft_probs slot. Consumed (and cleared) there.
+        self.py_needs_onehot_draft_probs = False
         self.py_num_accepted_draft_tokens_indices = []
         self.py_rewind_draft_token_separate_adjustment = 0
         self.py_per_pos_drafted = [0] * MAX_SPEC_DECODE_POSITIONS
diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py
index 2ec62d5d3cae..00aed5522ae0 100644
--- a/tensorrt_llm/_torch/pyexecutor/model_engine.py
+++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py
@@ -3296,6 +3296,10 @@ def _prepare_tp_inputs(
         draft_tokens = []
         draft_lens = []
         gen_request_seq_slots = []  # per generation request
+        # One-model rejection: slots of gen requests that produced 0 real draft
+        # tokens this step (marked in _handle_dynamic_draft_len); their stale
+        # draft_probs rows are one-hot'd after spec_metadata.prepare().
+        padding_gen_slots = []
         multimodal_params_list = []
         mrope_position_ids = [
         ]  # (start_idx, end_idx, (3,1,L) mrope_pos_ids) per multimodal request
@@ -3541,6 +3545,10 @@ def append_cross_attention_state(request: LlmRequest,
             self.runtime_draft_len)
         runtime_draft_token_buffer_width = runtime_tokens_per_gen_step - 1
         for request in extend_requests:
+            if getattr(request, "py_needs_onehot_draft_probs", False):
+                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
             request_ids.append(request.py_request_id)
             request_accepted_path[
                 request.
@@ -4278,6 +4286,11 @@ def previous_seq_slots_device():
             spec_metadata.populate_sampling_params_for_one_model(
                 scheduled_requests.all_requests())
             spec_metadata.prepare()
+            # One-model rejection: one-hot the stale draft_probs rows of gen
+            # requests that produced no draft tokens this step, so the (possibly
+            # captured) rejection kernel reads a legal placeholder distribution.
+            spec_metadata.write_padding_onehot_draft_probs(
+                padding_gen_slots, self.runtime_draft_len)
             inputs['spec_metadata'] = spec_metadata
 
             if self.enable_attention_dp:
diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py
index bf20ed9e495e..ed45c421f290 100644
--- a/tensorrt_llm/_torch/pyexecutor/py_executor.py
+++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py
@@ -3194,8 +3194,17 @@ def _handle_dynamic_draft_len(self,
                 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 = (
diff --git a/tensorrt_llm/_torch/speculative/dflash.py b/tensorrt_llm/_torch/speculative/dflash.py
index efa31fa2cc1e..9c90628d8708 100644
--- a/tensorrt_llm/_torch/speculative/dflash.py
+++ b/tensorrt_llm/_torch/speculative/dflash.py
@@ -407,17 +407,8 @@ def forward(
 
         self._execute_guided_decoder_if_present(logits)
 
-        # Target now emits K+1 logits per gen request and the previous step
-        # stored K draft tokens per gen request (no filler padding).
-        if num_gens > 0:
-            draft_tokens = spec_metadata.draft_tokens.reshape(num_gens, K)
-        else:
-            draft_tokens = spec_metadata.draft_tokens.reshape(0, K)
-
-        logits_for_accept = logits
-
-        accepted_tokens, num_accepted_tokens = self._sample_and_accept_draft_tokens_base(
-            logits_for_accept, draft_tokens, num_contexts, batch_size, spec_metadata
+        accepted_tokens, num_accepted_tokens = self.sample_and_accept_draft_tokens(
+            logits, attn_metadata, spec_metadata
         )
 
         # Update GDN/Mamba recurrent states to the accepted token's state.
@@ -498,17 +489,22 @@ def forward(
                 vocab_size = gen_logits.shape[-1]
                 gen_logits = gen_logits.reshape(num_gens, K, vocab_size)
 
-                d2t = getattr(draft_model.model, "d2t", None)
-                gen_draft_tokens = torch.argmax(gen_logits, dim=-1, keepdim=False).long()
-
-                if d2t is not None:
-                    gen_draft_tokens = d2t[gen_draft_tokens] + gen_draft_tokens
-
-                gen_draft_tokens = gen_draft_tokens.type(torch.int32)
+                gen_draft_tokens = self.sample_draft_tokens(
+                    gen_logits,
+                    spec_metadata,
+                    batch_size,
+                    num_contexts=num_contexts,
+                )
 
         else:
             gen_draft_tokens = torch.empty((0, K), dtype=torch.int32, device="cuda")
 
+        # Context requests are not drafted by the block worker (zero placeholder
+        # token); fill their draft-prob slot rows with a one-hot placeholder so
+        # they are a legal distribution when they become gen requests next iter.
+        gen_vocab = vocab_size if num_gens > 0 else None
+        self.write_context_onehot_draft_probs(spec_metadata, num_contexts, num_gens, K, gen_vocab)
+
         if num_contexts > 0 and num_gens > 0:
             ctx_draft_tokens = torch.zeros((num_contexts, K), dtype=torch.int32, device="cuda")
             next_draft_tokens = torch.cat([ctx_draft_tokens, gen_draft_tokens], dim=0)
diff --git a/tensorrt_llm/_torch/speculative/draft_target.py b/tensorrt_llm/_torch/speculative/draft_target.py
index fe6dc7d1e4c7..6df31b2bd223 100644
--- a/tensorrt_llm/_torch/speculative/draft_target.py
+++ b/tensorrt_llm/_torch/speculative/draft_target.py
@@ -257,10 +257,14 @@ def forward(
                     hidden_states[gather_ids], draft_model.lm_head, attn_metadata, True
                 )
                 if self.guided_decoder is not None:
-                    d2t = getattr(draft_model.model, "d2t", None)
-                    self.guided_decoder.execute_draft_batch(logits, d2t, draft_step=i)
+                    self.guided_decoder.execute_draft_batch(logits, self._d2t, draft_step=i)
 
-                new_draft_token = self.draft_decoder(logits, draft_model)
+                new_draft_token = self.sample_draft_tokens(
+                    logits,
+                    spec_metadata,
+                    batch_size,
+                    draft_step=i,
+                )
                 next_draft_tokens.append(new_draft_token)
 
                 # Update inputs and metadata for next draft step
@@ -314,36 +318,6 @@ def forward(
             "next_new_tokens": next_new_tokens,
         }
 
-    def sample_and_accept_draft_tokens(
-        self,
-        logits: torch.Tensor,
-        attn_metadata: AttentionMetadata,
-        spec_metadata: DraftTargetOneModelSpecMetadata,
-    ):
-        batch_size = attn_metadata.num_seqs
-        num_contexts = attn_metadata.num_contexts
-        num_gens = batch_size - num_contexts
-        runtime_draft_len = spec_metadata.runtime_draft_len
-
-        if spec_metadata.draft_tokens is None:
-            draft_tokens = torch.zeros(
-                (num_gens, runtime_draft_len), dtype=torch.int, device=logits.device
-            )
-        else:
-            draft_tokens = spec_metadata.draft_tokens.reshape(num_gens, runtime_draft_len)
-
-        return self._sample_and_accept_draft_tokens_base(
-            logits, draft_tokens, num_contexts, batch_size, spec_metadata
-        )
-
-    def draft_decoder(
-        self,
-        logits: torch.Tensor,
-        draft_model: nn.Module,
-    ):
-        d2t = getattr(draft_model.model, "d2t", None)
-        return self._draft_sampler_greedy(logits, d2t)
-
     def prepare_1st_drafter_inputs(
         self,
         input_ids: torch.LongTensor,
diff --git a/tensorrt_llm/_torch/speculative/eagle3.py b/tensorrt_llm/_torch/speculative/eagle3.py
index 3978040724b1..d9b5be1ecc1e 100644
--- a/tensorrt_llm/_torch/speculative/eagle3.py
+++ b/tensorrt_llm/_torch/speculative/eagle3.py
@@ -10,7 +10,6 @@
 from tensorrt_llm.mapping import Mapping
 
 from ..attention_backend import AttentionMetadata
-from ..distributed.ops import allgather
 from ..model_config import ModelConfig
 from ..pyexecutor.llm_request import LlmRequest
 from ..pyexecutor.mamba_cache_manager import MambaHybridCacheManager
@@ -891,9 +890,8 @@ def _forward_linear_draft_loop(self, inputs, attn_metadata, spec_metadata,
                         self.guided_decoder.execute_draft_batch(logits,
                                                                 draft_step=i)
                     else:
-                        d2t = getattr(draft_model.model, "d2t", None)
                         self.guided_decoder.execute_draft_batch(logits,
-                                                                d2t,
+                                                                self._d2t,
                                                                 draft_step=i)
 
                 # When ADP+LM-head-TP pads logits to max_num_requests, the
@@ -905,13 +903,20 @@ def _forward_linear_draft_loop(self, inputs, attn_metadata, spec_metadata,
                 # would otherwise fail to broadcast in apply_temperature. This
                 # also keeps next_draft_tokens and the draft_probs buffer
                 # token_count-sized without a post-hoc trim.
+                mapping_lm_head_tp = None
                 if use_lm_head_tp_in_adp:
                     logits = logits[:token_count]
-                new_draft_token = self.draft_decoder(logits,
-                                                     draft_model,
-                                                     spec_metadata,
-                                                     batch_size,
-                                                     draft_step=i)
+                    # The MTP head built this per-forward mapping when producing
+                    # the vocab-sharded logits; the sampler needs it to gather.
+                    mapping_lm_head_tp = getattr(
+                        draft_model.mtp_layers[0].shared_head,
+                        "mapping_lm_head_tp", None)
+                new_draft_token = self.sample_draft_tokens(
+                    logits,
+                    spec_metadata,
+                    batch_size,
+                    draft_step=i,
+                    mapping_lm_head_tp=mapping_lm_head_tp)
                 next_draft_tokens.append(new_draft_token)
 
                 # Update hidden states for the next iteration.
@@ -977,20 +982,6 @@ def _forward_linear_draft_loop(self, inputs, attn_metadata, spec_metadata,
                 gen_draft_tokens)
             next_draft_tokens[num_contexts:] = gen_draft_tokens
 
-        # Probs were already scattered into the slot-indexed buffer by
-        # _draft_sampler_advanced_for_rejection on each draft step (non-greedy
-        # batches only). All-greedy batches skip storage — rejection sampling
-        # will be bypassed by _can_use_rejection_sampling. Finalize the validity
-        # flag and d2t for next-iter target-side verification.
-        if spec_metadata.use_rejection_sampling:
-            if not spec_metadata.is_all_greedy_sample:
-                d2t_param = getattr(getattr(draft_model, 'model', None), "d2t",
-                                    None)
-                spec_metadata.d2t = d2t_param.data if d2t_param is not None else None
-                spec_metadata.draft_probs_valid = True
-            else:
-                spec_metadata.draft_probs_valid = False
-
         return next_draft_tokens
 
     def _get_step_all_rank_num_tokens(self, spec_metadata, step_idx: int):
@@ -1043,63 +1034,6 @@ def _prepare_flash_mla_generation_layout(self, attn_metadata, num_contexts,
         attn_metadata.block_ids_per_seq[:batch_size, :].copy_(
             reorder_block_ids_per_seq, non_blocking=True)
 
-    @torch.compile(options={"max-autotune": True})
-    def _get_local_max_and_combined(self, logits, mapping_lm_tp=None):
-        local_max_values, local_argmax = torch.max(logits, dim=-1, keepdim=True)
-        vocab_per_rank = logits.shape[-1]
-        mapping_lm_tp = mapping_lm_tp if mapping_lm_tp is not None else \
-            self.model_config.mapping
-        max_index_per_rank = local_argmax.type(
-            torch.int32) + (mapping_lm_tp.tp_rank * vocab_per_rank)
-        max_index_per_rank_float = max_index_per_rank.float()
-        local_max_values_float32 = local_max_values.float()
-        combined = torch.stack(
-            [max_index_per_rank_float, local_max_values_float32],
-            dim=-1).flatten(-2)
-        return combined
-
-    @torch.compile(options={"max-autotune": True})
-    def _get_draft_tokens_from_gathered(self, gathered):
-        gathered_indices_float = gathered[..., 0::2]
-        gathered_values_float = gathered[..., 1::2]
-        max_indices = torch.argmax(gathered_values_float, dim=-1, keepdim=True)
-        draft_tokens = torch.gather(gathered_indices_float, -1,
-                                    max_indices).squeeze(-1).type(torch.int32)
-        return draft_tokens
-
-    def draft_sampler(
-        self,
-        logits: torch.Tensor,
-        mapping_lm_head_tp=None,
-    ):
-        """TP-aware greedy draft token sampler (MTP Eagle path).
-
-        Falls back to simple argmax when no tensor parallelism is active or
-        when only attention DP is enabled without LM-head TP.
-        """
-        if (self.model_config is not None
-                and hasattr(self.model_config, 'mapping')
-                and self.model_config.mapping.tp_size > 1
-                and not self.model_config.mapping.enable_attention_dp):
-            combined = self._get_local_max_and_combined(logits)
-            gathered = allgather(combined, self.model_config.mapping, dim=-1)
-            return self._get_draft_tokens_from_gathered(gathered)
-        elif (self.model_config is not None
-              and hasattr(self.model_config, 'mapping')
-              and self.model_config.mapping.tp_size > 1
-              and self.model_config.mapping.enable_lm_head_tp_in_adp):
-            combined = self._get_local_max_and_combined(logits,
-                                                        mapping_lm_head_tp)
-            gathered = allgather(combined, mapping_lm_head_tp, dim=-1)
-            batch_size = logits.shape[0]
-            local_batch_size = batch_size // mapping_lm_head_tp.tp_size
-            gathered = gathered.view(mapping_lm_head_tp.tp_size,
-                                     local_batch_size, -1)
-            sliced_gathered = gathered[mapping_lm_head_tp.tp_rank]
-            return self._get_draft_tokens_from_gathered(sliced_gathered)
-        else:
-            return self._draft_sampler_greedy(logits)
-
     @torch.compile(options={"max-autotune": True})
     def _topk_kernel(self, gen_logprobs, num_gens, mtp_num_modules,
                      spec_metadata):
@@ -1210,78 +1144,6 @@ def sample_and_accept_draft_tokens(
         return self._accept_draft_tokens(logits, draft_tokens, num_contexts,
                                          batch_size, spec_metadata)
 
-    def draft_decoder(
-        self,
-        logits: torch.Tensor,
-        draft_model: nn.Module,
-        spec_metadata: Eagle3OneModelSpecMetadata,
-        batch_size: int,
-        draft_step: Optional[int] = None,
-    ):
-        '''
-        Sample draft tokens using the target's per-request sampling params
-        (temperature/top_k/top_p).
-
-        When rejection sampling is enabled and draft_step is provided, take the
-        single-pass path that also scatters the draft prob distribution into the
-        slot-indexed buffer (avoids a redundant softmax later).
-
-        Args:
-            logits: [batch_size, vocab_size] - Draft model logits.
-            draft_model: The draft model.
-            spec_metadata: Carries per-request sampling param tensors.
-            batch_size: Active requests, used to slice per-request tensors.
-            draft_step: Current draft step index (0..max_draft_len-1). Required
-                for the rejection-sampling code path so probs are written to
-                the correct slice of spec_metadata.draft_probs.
-        '''
-
-        d2t = getattr(getattr(draft_model, 'model', None), "d2t", None)
-        # All-greedy fast path must stay TP-aware. When the draft LM head is
-        # tensor-parallel (tp_size>1 without attention DP, or LM-head-TP in
-        # ADP), the draft logits are sharded along the vocab dim. A plain
-        # per-rank argmax then picks a different token on each rank, which
-        # desyncs the speculative-decoding control flow across ranks and
-        # deadlocks the next collective (observed as a generation hang on
-        # MTP-Eagle + TP). draft_sampler() all-gathers the sharded logits
-        # before argmax (and falls back to a plain argmax when no TP gather is
-        # needed). Eagle3 (non-MTP) keeps its d2t-aware argmax.
-        if spec_metadata.is_all_greedy_sample:
-            # Only plain tensor parallelism (tp_size>1 without attention DP)
-            # shards the draft logits over the vocab dim and thus needs
-            # draft_sampler()'s all-gather argmax. The LM-head-TP-in-ADP case
-            # already produces full-vocab logits per rank (gathered upstream),
-            # and the no-TP / Eagle3 cases need nothing, so they take the plain
-            # d2t-aware argmax. (Routing ADP/LM-head-TP through draft_sampler
-            # without its mapping_lm_head_tp arg hits the None-mapping branch
-            # and crashes with 'NoneType has no attribute tp_group'.)
-            if (self.is_mtp_eagle and self.model_config is not None
-                    and hasattr(self.model_config, 'mapping')
-                    and self.model_config.mapping.tp_size > 1
-                    and not self.model_config.mapping.enable_attention_dp):
-                return self.draft_sampler(logits)
-            return self._draft_sampler_greedy(logits, d2t)
-        # Non-greedy (advanced) draft sampling has the same TP hazard as the
-        # greedy path: when the draft LM head is plain tensor-parallel
-        # (tp_size>1 without attention DP), each rank only holds a vocab shard
-        # of the draft logits. Random per-rank sampling then draws different
-        # tokens on different ranks, desyncing the spec-decode control flow and
-        # deadlocking the next collective. All-gather the shards into the full
-        # vocab first so every rank samples from the same distribution with the
-        # shared seed. (Greedy uses draft_sampler()'s lighter max+index gather;
-        # random sampling needs the full distribution. The LM-head-TP-in-ADP
-        # case is handled upstream and must not be gathered again here.)
-        if (self.is_mtp_eagle and self.model_config is not None
-                and hasattr(self.model_config, 'mapping')
-                and self.model_config.mapping.tp_size > 1
-                and not self.model_config.mapping.enable_attention_dp):
-            logits = allgather(logits, self.model_config.mapping, dim=-1)
-        if spec_metadata.use_rejection_sampling and draft_step is not None:
-            return self._draft_sampler_advanced_for_rejection(
-                logits, spec_metadata, batch_size, d2t, draft_step)
-        return self._draft_sampler_advanced(logits, spec_metadata, batch_size,
-                                            d2t)
-
     def prepare_1st_drafter_inputs(
         self,
         input_ids: torch.LongTensor,
@@ -1344,10 +1206,12 @@ class MTPEagleWorker(Eagle3OneModelWorker):
     def __init__(self,
                  spec_config,
                  model_config: Optional[ModelConfig] = None,
-                 use_separate_draft_kv_cache: bool = False):
+                 use_separate_draft_kv_cache: bool = False,
+                 *,
+                 mapping: Optional[Mapping] = None):
         super().__init__(
             spec_config,
-            mapping=None,
+            mapping=mapping,
             model_config=model_config,
             use_separate_draft_kv_cache=use_separate_draft_kv_cache)
         # Preserved for callers/tests that still expect this attribute.
diff --git a/tensorrt_llm/_torch/speculative/eagle3_dynamic_tree.py b/tensorrt_llm/_torch/speculative/eagle3_dynamic_tree.py
index 7d9e836f5b9a..7edbcc5e6756 100644
--- a/tensorrt_llm/_torch/speculative/eagle3_dynamic_tree.py
+++ b/tensorrt_llm/_torch/speculative/eagle3_dynamic_tree.py
@@ -585,7 +585,6 @@ def _forward_draft_loop(
     ):
         """Dynamic tree draft loop with growing context."""
         spec_tree_manager = self.spec_tree_manager
-        self._d2t = getattr(draft_model.model, "d2t", None)
 
         assert batch_size <= self._max_batch_size, (
             f"batch_size {batch_size} exceeds pre-allocated max_batch_size {self._max_batch_size}"
@@ -636,9 +635,7 @@ def _forward_draft_loop(
                 hidden_states[gather_ids], draft_model.lm_head, attn_metadata, True
             )
 
-            new_draft_tokens, new_draft_scores = self.sample(
-                logits, self.K, draft_model=draft_model
-            )
+            new_draft_tokens, new_draft_scores = self.sample(logits, self.K)
 
             previous_draft_scores = self.update_draft_tokens_and_scores(
                 cur_draft_idx=0,
@@ -698,9 +695,7 @@ def _forward_draft_loop(
                     selected_hs, draft_model.lm_head, attn_metadata, True
                 )
 
-                new_draft_tokens, new_draft_scores = self.sample(
-                    logits, self.K, draft_model=draft_model
-                )
+                new_draft_tokens, new_draft_scores = self.sample(logits, self.K)
 
                 # Reshape for update: [batch_size, K, K]
                 new_draft_tokens = new_draft_tokens.reshape(batch_size, self.K, self.K)
@@ -1015,14 +1010,12 @@ def _finalize_dynamic_tree_verify_outputs(
         ).to(torch.int32)
 
     @nvtx_range("eagle3_dyn.sample")
-    def sample(
-        self, logits: torch.Tensor, max_top_k: int, draft_model=None
-    ) -> tuple[torch.Tensor, torch.Tensor]:
+    def sample(self, logits: torch.Tensor, max_top_k: int) -> tuple[torch.Tensor, torch.Tensor]:
         """TopK sampling with softmax for dynamic tree."""
         topk_indices, topk_values = _sample_softmax_topk(logits, max_top_k)
-        # Apply draft-to-target vocab mapping if the draft model has it
-        if draft_model is not None and hasattr(draft_model.model, "d2t"):
-            d2t = draft_model.model.d2t.data
+        # Apply draft-to-target vocab mapping when draft/target vocabs differ.
+        if self._d2t is not None:
+            d2t = self._d2t.data
             topk_indices = topk_indices + d2t[topk_indices]
         return topk_indices, topk_values
 
diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py
index 6244873efc43..275298a200c6 100644
--- a/tensorrt_llm/_torch/speculative/interface.py
+++ b/tensorrt_llm/_torch/speculative/interface.py
@@ -521,7 +521,6 @@ class SpecMetadata:
     skip_temperature: bool = False
     skip_top_k: bool = False
     skip_top_p: bool = False
-    has_greedy_requests: bool = False
     # Pre-computed top_k_max scalar (CPU-side) to avoid CUDA-graph-incompatible
     # dynamic boolean tensor indexing inside verify_dynamic_tree_rejection_from_logits_out.
     top_k_max: int = 0
@@ -539,68 +538,120 @@ class SpecMetadata:
     # Slot-indexed buffers (draft_probs) must span this full range.
     # 0 falls back to max_num_requests.
     num_seq_slots: int = 0
+    # Draft-model vocab size. full_draft_probs is allocated only when it differs
+    # from vocab_size; 0 (unknown) or a value equal to vocab_size means shared
+    # vocab and skips the buffer.
+    draft_vocab_size: int = 0
     # Draft probabilities buffer for rejection sampling, indexed by py_seq_slot
     # so per-request data is stable across iterations regardless of batch
     # composition shifts (chunking ctx, gen completion, new ctx joining).
     # Shape: [num_seq_slots, max_draft_len, vocab_size].
     draft_probs: Optional[torch.Tensor] = None
     draft_probs_vocab_size: int = 0
-    # Whether draft_probs contains valid data.
-    draft_probs_valid: bool = False
+    # Scratch row index that dummy/padding requests (py_seq_slot is None) route
+    # to, captured when draft_probs is allocated so it tracks the buffer's real
+    # size. It must NOT be re-derived from max_num_requests at use time:
+    # create_cuda_graph_metadata shrinks max_num_requests to the graph bucket
+    # size while sharing the full-size buffer, so a bucket-relative index would
+    # collide with a real request's slot row and overwrite its draft probs.
+    dummy_slot_row: int = 0
     # Last dimension size of the draft logits/probs stored in draft_probs.
     draft_probs_last_dim: int = 0
     # Per-request slot ids (py_seq_slot) for the current batch, in batch order.
     # Used to scatter draft probs by slot at write time and gather them by slot
     # at the next iter's verify. Shape: [max_num_requests], dtype=long.
     batch_slot_ids: Optional[torch.Tensor] = None
-    # Draft-to-target vocab offset tensor.
-    d2t: Optional[torch.Tensor] = None
-    # Pre-allocated scratch for draft probs expanded to the target vocab size.
-    # Filled with zeros once at prepare(); each rejection iter only overwrites
-    # the positions selected by d2t (or [:draft_vocab] when there is no d2t),
-    # so the zeros outside those positions persist across iterations and we
-    # avoid a per-iter 64 MB zero-fill on the (max_num_requests, max_draft_len,
-    # vocab_size) tensor. Shape: [max_num_requests, max_draft_len, vocab_size].
+    # Draft probs expanded to the target vocab size. Zero-filled once at
+    # prepare(); each rejection iter overwrites only the d2t-selected positions
+    # (or [:draft_vocab] when there is no d2t).
+    # Shape: [max_num_requests, max_draft_len, vocab_size].
     full_draft_probs: Optional[torch.Tensor] = None
-    # Cached d2t-projected target vocab indices, computed once on first use
-    # (d2t is a model-static tensor). Replaces the per-iter
-    # arange + (source + d2t) % vocab_size kernel sequence inside the d2t
-    # padding step. Shape: [draft_vocab_size], dtype long.
+    # Cached d2t-projected target vocab indices, computed once on first use.
+    # Shape: [draft_vocab_size], dtype long.
     d2t_target_indices: Optional[torch.Tensor] = None
 
     def __post_init__(self):
         pass
 
-    def prepare(self):
+    def prepare_rejection_sampling_buffers(self):
         """
-        Hook to be called before the forward step of the model.
+        Allocate the slot-indexed buffers used by one-model rejection sampling.
+
+        Idempotent and gated on ``use_rejection_sampling``.
         """
-        if (self.use_rejection_sampling and self.draft_probs is None
-                and self.vocab_size > 0):
-            # 3D [slot, draft_step, vocab] so we can scatter/gather by slot id
-            # and avoid the brittle "batch position == buffer position" mapping.
-            slot_capacity = self.num_seq_slots or self.max_num_requests
+        if not self.use_rejection_sampling:
+            return
+
+        # Slot-indexed buffers span the full SeqSlotManager pool: py_seq_slot
+        # can range over [0, num_seq_slots), which under DeepSeek-V4 overlap
+        # exceeds max_num_requests. Fall back to max_num_requests when the pool
+        # size is unknown (0). One extra scratch row at index ``slot_capacity``
+        # absorbs CUDA-graph dummy/padding requests (``py_seq_slot is None``).
+        slot_capacity = self.num_seq_slots or self.max_num_requests
+        num_slot_rows = slot_capacity + 1
+
+        if self.draft_probs is None and self.vocab_size > 0:
+            # [slot, draft_step, vocab]: scatter/gather by stable slot id.
             self.draft_probs = torch.empty(
-                (slot_capacity, self.max_draft_len, self.vocab_size),
+                (num_slot_rows, self.max_draft_len, self.vocab_size),
                 dtype=torch.float32,
                 device='cuda')
             self.draft_probs_vocab_size = self.vocab_size
-        if (self.use_rejection_sampling and self.batch_slot_ids is None
-                and self.max_num_requests > 0):
+            # Dummy requests route to the extra scratch row (the buffer's last
+            # row at index ``slot_capacity``). Capture it here against the real
+            # allocation size so it stays correct on graph copies that shrink
+            # max_num_requests and under overlap where slot_capacity exceeds it.
+            self.dummy_slot_row = slot_capacity
+        if self.batch_slot_ids is None and self.max_num_requests > 0:
             self.batch_slot_ids = torch.empty((self.max_num_requests, ),
                                               dtype=torch.long,
                                               device='cuda')
-        if (self.use_rejection_sampling and self.full_draft_probs is None
-                and self.vocab_size > 0):
-            # Zero-fill once. Subsequent iters only overwrite the d2t-mapped
-            # positions (constant across iters since d2t is model-static), so
-            # untouched positions stay 0 forever — saves the per-iter 64 MB
-            # zero-fill in _sample_and_accept_draft_tokens_rejection.
+        # full_draft_probs (d2t-expanded) is read only when draft and target
+        # vocabularies differ; skip it otherwise. Zero-filled once.
+        if (self.full_draft_probs is None and self.vocab_size > 0
+                and self.draft_vocab_size not in (0, self.vocab_size)):
             self.full_draft_probs = torch.zeros(
-                (self.max_num_requests, self.max_draft_len, self.vocab_size),
+                (num_slot_rows, self.max_draft_len, self.vocab_size),
                 dtype=torch.float32,
                 device='cuda')
 
+    def write_padding_onehot_draft_probs(self, padding_slot_ids, draft_len):
+        """Write a one-hot draft-prob row (prob 1.0 at draft-vocab token id 0,
+        the placeholder token) into each padding gen request's stable slot row.
+
+        Padding requests are gen requests that entered this iteration with 0 real
+        draft tokens (e.g. a runtime_draft_len K->0->K dynamic-draft-len toggle).
+        Their slot's draft_probs row was never scattered by the draft sampler, so
+        the next iteration's (possibly CUDA-graph-captured) rejection kernel would
+        read a stale/uninitialized distribution. Writing a legal one-hot row makes
+        acceptance reject the placeholder and resample from the target (equivalent
+        to strict acceptance) for those rows. Written eagerly before graph replay
+        into the stable draft_probs buffer, so the replayed kernel reads it.
+
+        Idempotent w.r.t. context->gen transitions whose row was already one-hot'd
+        by write_context_onehot_draft_probs. The width matches the value already
+        published in draft_probs_last_dim (what acceptance reads), so it is NOT
+        overwritten here. No-op unless rejection is enabled and slots exist.
+        Static shapes -> CUDA-graph safe.
+        """
+        if (not padding_slot_ids
+                or not getattr(self, "use_rejection_sampling", False)
+                or self.draft_probs is None):
+            return
+        onehot_vocab = (self.draft_probs_last_dim if self.draft_probs_last_dim
+                        > 0 else self.draft_probs_vocab_size)
+        slots = torch.tensor(padding_slot_ids,
+                             dtype=torch.long,
+                             device=self.draft_probs.device)
+        self.draft_probs[slots, :draft_len, :onehot_vocab] = 0.0
+        self.draft_probs[slots, :draft_len, 0] = 1.0
+
+    def prepare(self):
+        """
+        Hook to be called before the forward step of the model.
+        """
+        self.prepare_rejection_sampling_buffers()
+
     def create_cuda_graph_metadata(self, max_batch_size: int):
         """
         Creates metadata for CUDA graph execution.
@@ -634,8 +685,8 @@ def _scan_one_model_sampling(
     ) -> tuple[list[tuple[float, int, float, int]], list[int]]:
         """Single source of truth for one-engine sampling-param detection.
 
-        Scans the batch's sampling configs and sets skip_*/has_greedy_requests/
-        is_all_greedy_sample (honoring the warmup capture override). Returns
+        Scans the batch's sampling configs and sets skip_*/is_all_greedy_sample
+        (honoring the warmup capture override). Returns
         ``(per_request_normalized, per_request_slot_ids)`` for buffer
         population. Does NOT allocate or fill GPU buffers, so it is safe to call
         before the CUDA graph key is built.
@@ -693,7 +744,7 @@ def _normalize_request_sampling_params(
         temperature_enabled = False
         top_k_enabled = False
         top_p_enabled = False
-        has_greedy_requests = False
+        has_non_greedy_requests = False
         per_request_slot_ids: list[int] = []
 
         for request in requests:
@@ -722,32 +773,32 @@ def _normalize_request_sampling_params(
             temperature_enabled |= use_temperature
             top_k_enabled |= use_top_k
             top_p_enabled |= use_top_p
-            has_greedy_requests |= is_greedy
+            has_non_greedy_requests |= not is_greedy
 
             per_request_normalized.append(
                 (temp_val, tk_val, tp_val, num_tokens))
-            # py_seq_slot is a stable per-request id used to scatter / gather
-            # draft probs across iterations. Dummies / unallocated slots fall
-            # back to 0 (any valid index is fine — the data at that slot will
-            # be overwritten on the next real iteration before being read).
+            # py_seq_slot is a stable per-request id used to scatter/gather draft
+            # probs across iterations. Dummy/padding requests (py_seq_slot is
+            # None) route to the scratch row captured at allocation time (the
+            # buffer's real last row), not max_num_requests, which a graph copy
+            # shrinks to the bucket size and would alias a real request's slot.
             per_request_slot_ids.append(
-                request.py_seq_slot if request.py_seq_slot is not None else 0)
+                request.py_seq_slot if request.
+                py_seq_slot is not None else self.dummy_slot_row)
 
         self.skip_temperature = not temperature_enabled
         self.skip_top_k = not top_k_enabled
         self.skip_top_p = not top_p_enabled
-        self.has_greedy_requests = has_greedy_requests
         # Used in the CUDA graph key to pick the argmax / advanced variant.
-        self.is_all_greedy_sample = (self.skip_temperature and self.skip_top_k
-                                     and self.skip_top_p)
-
-        # Warmup-time override (set via runtime attribute by the model engine):
-        # force the advanced-sampling code path so the CUDA graph for the
-        # (is_all_greedy_sample=False) key gets captured. Dummy warmup requests
-        # carry no sampling params, so the natural detection above always
-        # returns True; this branch substitutes synthetic non-greedy scalars
-        # into the per-request data and lets Phase 2 run normally to populate
-        # the GPU buffers used by the captured kernels.
+        # All-greedy iff EVERY request is greedy. Derived from per-request
+        # greediness, not from the skip_* filter flags (a non-greedy request may
+        # enable no filter, e.g. temperature=1.0 with top_k/top_p unset).
+        self.is_all_greedy_sample = not has_non_greedy_requests
+
+        # Warmup-time override: force the advanced-sampling path so the CUDA
+        # graph for the (is_all_greedy_sample=False) key gets captured. Dummy
+        # warmup requests carry no sampling params, so substitute synthetic
+        # non-greedy scalars to populate the GPU buffers.
         if getattr(self, '_force_non_greedy_for_capture', False):
             self.skip_temperature = False
             self.skip_top_k = False
@@ -765,13 +816,8 @@ def update_is_all_greedy_sample(self, requests: list["LlmRequest"]) -> None:
 
         Must be called BEFORE the CUDA graph key is built (the key includes
         ``is_all_greedy_sample`` to choose the argmax vs advanced-sampling graph
-        variant). ``populate_sampling_params_for_one_model`` runs later, inside
-        ``_prepare_inputs``, and re-derives the same flag while filling the GPU
-        sampling buffers. Computing the flag here first keeps the selected graph
-        consistent with the buffers ``populate`` fills; otherwise the key would
-        use the previous iteration's stale value and could replay the advanced
-        graph against unpopulated (greedy) buffers, which can hang/corrupt the
-        run (notably for MTP with num_nextn>=2).
+        variant), so the selected graph stays consistent with the buffers
+        ``populate_sampling_params_for_one_model`` fills later.
         """
         if not self.spec_dec_mode.use_one_engine():
             return
@@ -789,6 +835,11 @@ def populate_sampling_params_for_one_model(
         if not self.spec_dec_mode.use_one_engine():
             return
 
+        # Allocate the rejection buffers before copying py_seq_slot values into
+        # batch_slot_ids below; this runs earlier than prepare() in the
+        # model-engine flow. No-op unless use_rejection_sampling is set.
+        self.prepare_rejection_sampling_buffers()
+
         if self.temperatures is None:
             # Ensures determinism across ranks.
             torch.manual_seed(0)
@@ -917,6 +968,9 @@ def __init__(self, use_separate_draft_kv_cache: bool = False):
         self.seed: Optional[torch.Tensor] = None
         self.offset: Optional[torch.Tensor] = None
         self.use_separate_draft_kv_cache = use_separate_draft_kv_cache
+        # Static draft->target vocab offset map, cached once the draft model is
+        # loaded (see set_draft_model). None when draft and target share a vocab.
+        self._d2t: Optional[torch.Tensor] = None
         # Lazily-initialized state for the fractional synthetic acceptance
         # rate. The pool is a fixed-seed, rank-independent table of uniform
         # [0, 1) values; the counter is a device-side int64 advanced in-place
@@ -1025,6 +1079,13 @@ def set_guided_decoder(self,
         self.guided_decoder = guided_decoder
         return True
 
+    def set_draft_model(self, draft_model) -> None:
+        """Cache the static draft->target vocab offset map (``d2t``) once the
+        draft model is loaded. ``d2t`` is a model-static parameter present only
+        when the draft and target vocabularies differ; stays None otherwise.
+        """
+        self._d2t = getattr(getattr(draft_model, "model", None), "d2t", None)
+
     def _prepare_attn_metadata_for_spec_dec(self, attn_metadata):
         """
         Prepare attention metadata before speculative decoding draft token generation.
@@ -1248,26 +1309,219 @@ def _accept_draft_tokens(self, logits, draft_tokens, num_contexts,
             stored_vocab = (spec_metadata.draft_probs_last_dim
                             if spec_metadata.draft_probs_last_dim > 0 else
                             spec_metadata.draft_probs_vocab_size)
-            # Gather the slot rows for the gen subset. The buffer was filled
-            # at the previous draft step indexed by py_seq_slot, so each gen
-            # request reads back exactly its own probs, regardless of batch
-            # composition changes since then.
-            gen_slot_ids = spec_metadata.batch_slot_ids[num_contexts:batch_size]
-            draft_probs = spec_metadata.draft_probs[
-                gen_slot_ids, :draft_len, :stored_vocab]
-            return self._sample_and_accept_draft_tokens_rejection(
-                logits, draft_tokens, draft_probs, num_contexts, batch_size,
-                spec_metadata)
+            # Fail closed: run the rejection kernel only when every buffer is
+            # present and correctly shaped; otherwise fall back to strict
+            # acceptance.
+            if self._rejection_buffers_valid(draft_tokens, draft_len,
+                                             stored_vocab, num_contexts,
+                                             batch_size, logits, spec_metadata):
+                # Gather the gen subset's slot rows, filled at the previous draft
+                # step indexed by py_seq_slot.
+                gen_slot_ids = spec_metadata.batch_slot_ids[
+                    num_contexts:batch_size]
+                draft_probs = spec_metadata.draft_probs[
+                    gen_slot_ids, :draft_len, :stored_vocab]
+                return self._sample_and_accept_draft_tokens_rejection(
+                    logits, draft_tokens, draft_probs, num_contexts, batch_size,
+                    spec_metadata)
         return self._sample_and_accept_draft_tokens_base(
             logits, draft_tokens, num_contexts, batch_size, spec_metadata)
 
+    def _draft_logits_are_sharded(self, logits, spec_metadata):
+        """Whether the draft logits are vocab-sharded and need a TP gather.
+
+        Sharded when tp_size>1 and the logits' last dim is narrower than the
+        DRAFT head's own full vocab -- either plain TP (no attention DP) or the
+        ADP + LM-head-TP mode, both of which produce vocab-sharded draft logits.
+        Replicated full-vocab logits (borrowed/gathered LM head, plain attention
+        DP, or a single rank) are not sharded.
+
+        The reference width is the draft head's own full vocab (``draft_vocab_size``,
+        falling back to ``vocab_size`` when unknown/shared), NOT the target
+        ``vocab_size``: an Eagle3 reduced-vocab draft head produces full,
+        replicated ``[tokens, draft_vocab_size]`` logits that are narrower than the
+        target vocab; comparing against the target vocab would misclassify those as
+        sharded and gather identical copies (overflowing the d2t table).
+        """
+        mapping = self.mapping
+        if mapping is None or getattr(mapping, "tp_size", 1) <= 1:
+            return False
+        # Under attention DP every rank is data-parallel -- it owns a distinct
+        # set of requests -- so the draft logits must NOT be cross-rank gathered,
+        # whether they are replicated full-vocab (plain ADP) or vocab-sharded
+        # (ADP + LM-head TP). A per-rank argmax on the rank's own logits is the
+        # correct proposal (verified later); a gather would splice in a
+        # mismatched collective across ranks that hold different token counts,
+        # desyncing them into a hang / DeepEP launch failure. Only plain TP
+        # (tp>1 without ADP), where all ranks share the same tokens sharded over
+        # the vocab dim, needs the gather.
+        if getattr(mapping, "enable_attention_dp", False):
+            return False
+        draft_full_vocab = (getattr(spec_metadata, "draft_vocab_size", 0)
+                            or getattr(spec_metadata, "vocab_size", 0) or 0)
+        return bool(draft_full_vocab) and logits.shape[-1] < draft_full_vocab
+
+    def maybe_gather_sharded_draft_logits(self,
+                                          logits,
+                                          spec_metadata,
+                                          mapping_lm_head_tp=None):
+        """All-gather TP-sharded draft logits to full vocab before advanced sampling.
+
+        Advanced (non-greedy) draft sampling needs the full-vocab distribution.
+        Gathers shards only for a non-greedy batch when the logits are sharded
+        (see ``_draft_logits_are_sharded``); replicated full-vocab logits are
+        returned unchanged.
+
+        Plain TP gathers vocab shards over ``self.mapping``. Under ADP + LM-head
+        TP the worker has already trimmed the LM-head-TP padding rows so each rank
+        holds ``[token_count, vocab_shard]`` for its own tokens; a vocab-dim
+        all-gather over ``mapping_lm_head_tp`` restores full vocab (no token
+        re-slice is needed after the trim).
+        """
+        if (spec_metadata is None or spec_metadata.is_all_greedy_sample
+                or not self._draft_logits_are_sharded(logits, spec_metadata)):
+            return logits
+
+        # Only plain TP (no attention DP) reaches here -- ADP variants return
+        # early via _draft_logits_are_sharded, since their ranks are
+        # data-parallel and must not gather draft logits across ranks.
+        from ..distributed.ops import allgather
+        return allgather(logits, self.mapping, dim=-1)
+
+    def advanced_sample_draft(self,
+                              logits: torch.Tensor,
+                              spec_metadata: "SpecMetadata",
+                              batch_size: int,
+                              draft_step: Optional[int] = None):
+        """Per-step advanced (non-greedy) draft sampler for step workers
+        (MTP, DraftTarget).
+
+        With rejection enabled and a ``draft_step``, samples via
+        ``sampling_batch_spec_dec_one_model_for_rejection`` and scatters this
+        step's proposal distribution into the slot-indexed ``draft_probs``
+        buffer; otherwise uses ``sampling_batch_spec_dec_one_model`` (tokens
+        only). Returns tokens in draft-vocab space (the caller applies d2t).
+        Expects 2D ``[batch_size, vocab]`` logits (one row per request).
+        """
+        temperatures = spec_metadata.request_temperatures[:batch_size]
+        top_ks = spec_metadata.request_top_ks[:batch_size]
+        top_ps = spec_metadata.request_top_ps[:batch_size]
+
+        self._update_advance_draft_sampling_seed(logits.device)
+        if spec_metadata.use_rejection_sampling and draft_step is not None:
+            draft_tokens, probs = (
+                sampling_batch_spec_dec_one_model_for_rejection(
+                    logits,
+                    temperatures,
+                    top_ks,
+                    top_ps,
+                    seed=self.seed,
+                    offset=self.offset))
+            # Scatter probs into the slot-indexed buffer so each request's data
+            # lands at its stable py_seq_slot row regardless of batch shifts.
+            assert spec_metadata.batch_slot_ids is not None, (
+                "batch_slot_ids must be populated by "
+                "populate_sampling_params_for_one_model before draft probs "
+                "storage")
+            batch_slots = spec_metadata.batch_slot_ids[:batch_size]
+            vocab = probs.shape[-1]
+            spec_metadata.draft_probs[batch_slots, draft_step, :vocab] = probs
+            spec_metadata.draft_probs_last_dim = vocab
+        else:
+            draft_tokens = sampling_batch_spec_dec_one_model(logits,
+                                                             temperatures,
+                                                             top_ks,
+                                                             top_ps,
+                                                             seed=self.seed,
+                                                             offset=self.offset)
+
+        return draft_tokens.type(torch.int32)
+
+    def _reshape_draft_tokens_for_accept(self, spec_metadata, num_gens, device):
+        """Reshape the stored draft tokens to ``[num_gens, runtime_draft_len]``
+        for acceptance. Default assumes one draft token per step (DraftTarget,
+        DFlash); workers with a different buffer layout (e.g. PARD's 2K-1
+        entries) override this.
+        """
+        runtime_draft_len = spec_metadata.runtime_draft_len
+        if spec_metadata.draft_tokens is None:
+            return torch.zeros((num_gens, runtime_draft_len),
+                               dtype=torch.int,
+                               device=device)
+        return spec_metadata.draft_tokens.reshape(num_gens, runtime_draft_len)
+
+    def _reshape_logits_for_accept(self, logits, num_contexts, num_gens,
+                                   spec_metadata):
+        """Reshape target logits to the ``[num_contexts + num_gens*(K+1), vocab]``
+        layout expected by acceptance. Default is identity (target already emits
+        one logit per accepted position); workers that emit extra positions
+        (e.g. PARD's 2K per gen request) override this.
+        """
+        return logits
+
+    def sample_and_accept_draft_tokens(self, logits, attn_metadata,
+                                       spec_metadata):
+        """Sample the golden token and verify previously proposed draft tokens.
+
+        Default implementation for one-model workers whose acceptance differs
+        only in how draft tokens / target logits are reshaped (DraftTarget,
+        PARD, DFlash): unpack batch sizes, reshape via the overridable hooks,
+        then route through ``_accept_draft_tokens`` (strict or rejection). Workers
+        with a materially different acceptance path (MTP, Eagle3: relaxed /
+        THOP / extra ``input_ids``) override this method entirely.
+        """
+        batch_size = attn_metadata.num_seqs
+        num_contexts = attn_metadata.num_contexts
+        num_gens = batch_size - num_contexts
+        draft_tokens = self._reshape_draft_tokens_for_accept(
+            spec_metadata, num_gens, logits.device)
+        logits = self._reshape_logits_for_accept(logits, num_contexts, num_gens,
+                                                 spec_metadata)
+        return self._accept_draft_tokens(logits, draft_tokens, num_contexts,
+                                         batch_size, spec_metadata)
+
+    def _rejection_buffers_valid(self, draft_tokens, draft_len, stored_vocab,
+                                 num_contexts, batch_size, logits,
+                                 spec_metadata) -> bool:
+        """Fail-closed guard: return True only when the slot-indexed draft-prob
+        buffers exist and every shape the rejection path dereferences is valid;
+        otherwise the caller falls back to strict acceptance. Inspects only
+        host-side tensor shapes -- no ``.item()`` / value read on CUDA tensors --
+        so it stays CUDA-graph-capture safe.
+        """
+        draft_probs = spec_metadata.draft_probs
+        batch_slot_ids = spec_metadata.batch_slot_ids
+        if draft_probs is None or batch_slot_ids is None:
+            return False
+        if stored_vocab <= 0:
+            return False
+        num_gens = batch_size - num_contexts
+        # draft_probs must cover the slice [:, :draft_len, :stored_vocab].
+        if draft_probs.dim() != 3:
+            return False
+        if draft_probs.shape[1] < draft_len or draft_probs.shape[
+                2] < stored_vocab:
+            return False
+        if draft_tokens.dim() != 2 or draft_tokens.shape[0] != num_gens:
+            return False
+        # logits must cover context rows (1 each) + gen rows (draft_len + 1 each).
+        logits_rows = logits.shape[0] if logits.dim() > 1 else 1
+        if logits_rows < num_contexts + num_gens * (draft_len + 1):
+            return False
+        # Slot ids for the gen subset must exist (range safety is guaranteed by
+        # construction).
+        if batch_slot_ids.shape[0] < batch_size:
+            return False
+        if batch_slot_ids[num_contexts:batch_size].shape[0] != num_gens:
+            return False
+        return True
+
     def _can_use_rejection_sampling(self, spec_metadata: SpecMetadata) -> bool:
         # Skip rejection sampling when the whole batch is greedy: the accepted
         # result is identical to argmax and the base path is cheaper. Mixed
         # batches (context + gen) are handled via slot-indexed draft probs and
         # are split inside _sample_and_accept_draft_tokens_rejection.
         return (spec_metadata.use_rejection_sampling
-                and spec_metadata.draft_probs_valid
                 and not spec_metadata.is_all_greedy_sample)
 
     def _sample_and_accept_draft_tokens_rejection(
@@ -1343,16 +1597,17 @@ def _sample_and_accept_draft_tokens_rejection(
             assert draft_probs.shape[1] == runtime_draft_len, (
                 f"draft_probs draft length mismatch: {draft_probs.shape[1]} != "
                 f"{runtime_draft_len}")
-            d2t = getattr(spec_metadata, "d2t", None)
+            d2t = self._d2t.data if self._d2t is not None else None
             if draft_vocab_size != vocab_size:
-                # Use the pre-allocated buffer from spec_metadata.prepare()
-                # (zero-filled once at init; untouched positions stay 0).
-                # Falls back to per-iter allocation if the buffer is not
-                # configured, e.g. when use_rejection_sampling was off at
-                # prepare() time.
                 if spec_metadata.full_draft_probs is not None:
-                    full_draft_probs = spec_metadata.full_draft_probs[:num_gens]
+                    # Slice to runtime_draft_len so the max_draft_len buffer
+                    # never passes stale extra rows to the rejection kernel.
+                    full_draft_probs = spec_metadata.full_draft_probs[:
+                                                                      num_gens, :
+                                                                      runtime_draft_len]
                 else:
+                    # Buffer not pre-allocated (e.g. rejection off at prepare()):
+                    # fall back to a per-iter allocation.
                     full_draft_probs = torch.zeros(
                         (num_gens, runtime_draft_len, vocab_size),
                         dtype=torch.float32,
@@ -1419,142 +1674,255 @@ def _sample_and_accept_draft_tokens_rejection(
             spec_metadata=spec_metadata)
         return accepted_tokens, num_accepted_tokens
 
-    def _draft_sampler_greedy(self, logits: torch.Tensor, d2t=None):
+    def _update_advance_draft_sampling_seed(self, device):
+        """Increment the draft sampler's RNG seed for this draft-sampling call
+        (lazily initializing the seed/offset tensors on first use), so each call
+        samples with a fresh, deterministic seed."""
+        if self.seed is None:
+            self.seed = torch.tensor([0], dtype=torch.int64, device=device)
+            self.offset = torch.tensor([0], dtype=torch.int64, device=device)
+        self.seed += 1
+        self.seed %= (2**31)
+
+    def _draft_sampler_greedy(self, logits: torch.Tensor):
         """
         Simple greedy draft token sampling using argmax.
 
         Args:
             logits: [num_tokens, vocab_size] - Draft model logits
-            d2t: Optional dictionary offset tensor for vocab mapping
 
         Returns:
             draft_tokens: [num_tokens] - Sampled draft token ids (int32)
         """
         draft_tokens = greedy(logits, return_probs=False)[0]
 
-        # Apply d2t (offsets between draft and target model dictionaries)
-        if d2t is not None:
-            draft_tokens = d2t[draft_tokens] + draft_tokens
+        # Apply the cached draft->target vocab offset map.
+        if self._d2t is not None:
+            draft_tokens = self._d2t[draft_tokens] + draft_tokens
 
         return draft_tokens.type(torch.int32)
 
-    def _draft_sampler_advanced(
-        self,
-        logits: torch.Tensor,
-        spec_metadata: "SpecMetadata",
-        batch_size: int,
-        d2t: Optional[torch.Tensor] = None,
-    ):
+    def _get_local_max_and_combined(self, logits, mapping_lm_tp=None):
+        """Pack each rank's local (global_argmax_index, max_value) for a
+        distributed argmax over a vocab-sharded draft LM head.
         """
-        Draft token sampling using per-request sampling parameters from the
-        target's sampling config. Falls back to argmax when the batch is
-        all-greedy.
-
-        Args:
-            logits: [batch_size, vocab_size] - Draft model logits (one row per
-                request, since each draft step emits one token per request).
-            spec_metadata: Source of per-request temperatures / top_k / top_p
-                tensors populated by populate_sampling_params_for_one_model.
-            batch_size: Number of active requests in the batch.
-            d2t: Optional dictionary offset tensor for vocab mapping.
-
-        Returns:
-            draft_tokens: [batch_size] - Sampled draft token ids (int32)
+        local_max_values, local_argmax = torch.max(logits, dim=-1, keepdim=True)
+        vocab_per_rank = logits.shape[-1]
+        mapping_lm_tp = mapping_lm_tp if mapping_lm_tp is not None else self.mapping
+        max_index_per_rank = local_argmax.type(
+            torch.int32) + (mapping_lm_tp.tp_rank * vocab_per_rank)
+        max_index_per_rank_float = max_index_per_rank.float()
+        local_max_values_float32 = local_max_values.float()
+        # Interleaved layout: [idx0, val0, idx1, val1, ...] after all-gather.
+        combined = torch.stack(
+            [max_index_per_rank_float, local_max_values_float32],
+            dim=-1).flatten(-2)
+        return combined
+
+    @torch.compile(options={"max-autotune": True})
+    def _get_draft_tokens_from_gathered(self, gathered):
+        """Pick the global-argmax token id from the all-gathered per-rank
+        (index, value) pairs produced by ``_get_local_max_and_combined``.
         """
-        if spec_metadata.is_all_greedy_sample:
-            return self._draft_sampler_greedy(logits, d2t)
-
-        temperatures = spec_metadata.request_temperatures[:batch_size]
-        top_ks = spec_metadata.request_top_ks[:batch_size]
-        top_ps = spec_metadata.request_top_ps[:batch_size]
-
-        if self.seed is None:
-            self.seed = torch.tensor([0],
-                                     dtype=torch.int64,
-                                     device=logits.device)
-            self.offset = torch.tensor([0],
-                                       dtype=torch.int64,
-                                       device=logits.device)
-        self.seed += 1
-        self.seed %= (2**31)
-
-        draft_tokens = sampling_batch_spec_dec_one_model(logits,
-                                                         temperatures,
-                                                         top_ks,
-                                                         top_ps,
-                                                         seed=self.seed,
-                                                         offset=self.offset)
-
-        if d2t is not None:
-            draft_tokens = d2t[draft_tokens] + draft_tokens
-
-        return draft_tokens.type(torch.int32)
-
-    def _draft_sampler_advanced_for_rejection(
-        self,
-        logits: torch.Tensor,
-        spec_metadata: "SpecMetadata",
-        batch_size: int,
-        d2t: Optional[torch.Tensor] = None,
-        draft_step: int = 0,
-    ):
+        gathered_indices_float = gathered[..., 0::2]
+        gathered_values_float = gathered[..., 1::2]
+        max_indices = torch.argmax(gathered_values_float, dim=-1, keepdim=True)
+        draft_tokens = torch.gather(gathered_indices_float, -1,
+                                    max_indices).squeeze(-1).type(torch.int32)
+        return draft_tokens
+
+    def greedy_sample_draft_with_tp_gather(self,
+                                           logits: torch.Tensor,
+                                           spec_metadata=None,
+                                           mapping_lm_head_tp=None):
+        """Greedy draft-token sampling with a TP all-gather of the argmax.
+
+        When the draft LM head is vocab-sharded under plain tensor parallelism,
+        a per-rank argmax disagrees across ranks and desyncs speculative
+        decoding. Gather only each rank's local (index, value) and pick the
+        global argmax. Falls back to plain argmax when the logits are not
+        vocab-sharded (see ``_draft_logits_are_sharded``) -- e.g. a borrowed or
+        gathered full-vocab draft head. Returns tokens in draft-vocab space (the
+        caller applies d2t). Expects 2D ``[num_tokens, vocab_shard]`` logits.
         """
-        Rejection-sampling-aware variant of ``_draft_sampler_advanced``.
-
-        Single-pass compute + sample + scatter: computes the per-request prob
-        distribution once via TRT-LLM's fused ``compute_probs_from_logits``
-        (temp + top_k + top_p + softmax + greedy override in one CUDA kernel),
-        samples the draft token from that distribution, and scatters the same
-        probs into the slot-indexed ``spec_metadata.draft_probs`` buffer for
-        next-iter rejection verification. Replaces the previous two-stage path
-        (flashinfer fused sampling kernel + a redundant softmax pass to store
-        probs).
-
-        All-greedy batches take the cheaper argmax path —
-        ``_can_use_rejection_sampling`` will bypass rejection for those anyway.
+        mapping = self.mapping
+        sharded = self._draft_logits_are_sharded(logits, spec_metadata)
+        if (sharded and mapping is not None
+                and getattr(mapping, "tp_size", 1) > 1
+                and not mapping.enable_attention_dp):
+            from ..distributed.ops import allgather
+            combined = self._get_local_max_and_combined(logits)
+            gathered = allgather(combined, mapping, dim=-1)
+            return self._get_draft_tokens_from_gathered(gathered)
+        # No TP gather for attention-DP (incl. ADP + LM-head TP): each rank owns
+        # its own requests, so a per-rank argmax is the correct proposal and a
+        # cross-rank gather here would desync the ranks (see
+        # _draft_logits_are_sharded). Plain argmax; caller applies d2t.
+        return torch.argmax(logits, dim=-1).type(torch.int32)
+
+    def advanced_sample_draft_block(self, gen_logits: torch.Tensor,
+                                    spec_metadata: "SpecMetadata",
+                                    num_contexts: int, batch_size: int):
+        """Block counterpart of ``advanced_sample_draft`` for gen-only workers
+        (PARD, DFLASH): produces all ``K`` draft positions per gen request in one
+        forward from ``[num_gens, K, vocab]`` logits.
+
+        With rejection enabled, samples via
+        ``sampling_batch_spec_dec_one_model_for_rejection`` and scatters the K
+        proposal rows into ``draft_probs[gen_slot_ids, 0:K, :]``; otherwise uses
+        ``sampling_batch_spec_dec_one_model`` (tokens only). Only called for a
+        non-greedy batch (the all-greedy path is handled by the caller). Returns
+        ``[num_gens, K]`` int32 tokens in draft-vocab space (the caller applies
+        d2t); stored probs likewise stay in draft-vocab space.
         """
-        if spec_metadata.is_all_greedy_sample:
-            return self._draft_sampler_greedy(logits, d2t)
-
-        temperatures = spec_metadata.request_temperatures[:batch_size]
-        top_ks = spec_metadata.request_top_ks[:batch_size]
-        top_ps = spec_metadata.request_top_ps[:batch_size]
-
-        if self.seed is None:
-            self.seed = torch.tensor([0],
-                                     dtype=torch.int64,
-                                     device=logits.device)
-            self.offset = torch.tensor([0],
-                                       dtype=torch.int64,
-                                       device=logits.device)
-        self.seed += 1
-        self.seed %= (2**31)
-
-        draft_tokens, probs = sampling_batch_spec_dec_one_model_for_rejection(
-            logits,
-            temperatures,
-            top_ks,
-            top_ps,
-            seed=self.seed,
-            offset=self.offset,
-        )
+        num_gens, K, vocab = gen_logits.shape
+        if num_gens == 0:
+            return torch.empty((0, K),
+                               dtype=torch.int32,
+                               device=gen_logits.device)
+
+        # Take the gen slice and repeat each request's value K times to line up
+        # with the flattened [num_gens*K, vocab] logits (K rows per request).
+        temps = spec_metadata.request_temperatures[
+            num_contexts:batch_size].repeat_interleave(K)
+        top_ks = spec_metadata.request_top_ks[
+            num_contexts:batch_size].repeat_interleave(K)
+        top_ps = spec_metadata.request_top_ps[
+            num_contexts:batch_size].repeat_interleave(K)
+
+        self._update_advance_draft_sampling_seed(gen_logits.device)
+        flat_logits = gen_logits.reshape(num_gens * K, vocab)
+
+        if getattr(spec_metadata, "use_rejection_sampling", False):
+            flat_tokens, flat_probs = (
+                sampling_batch_spec_dec_one_model_for_rejection(
+                    flat_logits,
+                    temps,
+                    top_ks,
+                    top_ps,
+                    seed=self.seed,
+                    offset=self.offset))
+            # Scatter the K prob rows per gen request into its stable slot row.
+            if spec_metadata.draft_probs is not None:
+                assert spec_metadata.batch_slot_ids is not None, (
+                    "batch_slot_ids must be populated before block draft prob "
+                    "storage")
+                gen_slot_ids = spec_metadata.batch_slot_ids[
+                    num_contexts:batch_size]
+                probs = flat_probs.reshape(num_gens, K, vocab)
+                spec_metadata.draft_probs[gen_slot_ids, :K, :vocab] = probs
+                spec_metadata.draft_probs_last_dim = vocab
+        else:
+            flat_tokens = sampling_batch_spec_dec_one_model(flat_logits,
+                                                            temps,
+                                                            top_ks,
+                                                            top_ps,
+                                                            seed=self.seed,
+                                                            offset=self.offset)
+
+        return flat_tokens.reshape(num_gens, K).type(torch.int32)
+
+    def write_context_onehot_draft_probs(self,
+                                         spec_metadata,
+                                         num_contexts,
+                                         num_gens,
+                                         draft_len,
+                                         gen_vocab=None):
+        """Write a one-hot draft-prob distribution (prob 1.0 at draft-vocab token
+        id 0, the placeholder draft token) into each context request's stable
+        slot row, so the row is a legal distribution when the context request
+        becomes a generation request next iteration and rejection acceptance
+        reads ``draft_probs[slot, :draft_len, :stored_vocab]``.
+
+        Block workers (PARD/DFlash) do not draft context requests (they get a
+        zero placeholder token), leaving their slot rows unwritten. Mixed
+        iterations reuse the vocab width the gen scatter just set (``gen_vocab``);
+        pure-context iterations have no gen logits, so use the buffer's full
+        draft-vocab width and publish it via ``draft_probs_last_dim`` (acceptance
+        next iter reads this scalar). No-op unless rejection is enabled and the
+        slot-indexed buffers exist. Static shapes -> CUDA-graph safe.
+        """
+        if (num_contexts <= 0
+                or not getattr(spec_metadata, "use_rejection_sampling", False)
+                or spec_metadata.draft_probs is None):
+            return
+        ctx_slot_ids = spec_metadata.batch_slot_ids[:num_contexts]
+        if num_gens > 0:
+            onehot_vocab = gen_vocab  # matches the gen scatter's width
+        else:
+            onehot_vocab = spec_metadata.draft_probs_vocab_size
+            spec_metadata.draft_probs_last_dim = onehot_vocab
+        spec_metadata.draft_probs[ctx_slot_ids, :draft_len, :onehot_vocab] = 0.0
+        spec_metadata.draft_probs[ctx_slot_ids, :draft_len, 0] = 1.0
+
+    def sample_draft_tokens(self,
+                            logits,
+                            spec_metadata,
+                            batch_size,
+                            *,
+                            num_contexts=0,
+                            draft_step=None,
+                            mapping_lm_head_tp=None):
+        """Unified draft-token production entry for all one-model workers.
+
+        Branches by logits rank: 3D ``[num_gens, K, vocab]`` is the block form
+        (gen-only workers emitting all K positions in one forward, e.g.
+        PARD/DFlash); 2D ``[num_tokens, vocab]`` is the per-step form
+        (autoregressive workers called once per draft step, e.g. MTP/
+        DraftTarget). ``draft_step`` applies only to the step form and
+        ``num_contexts`` only to the block form. d2t is read from ``self._d2t``.
+
+        In a mixed (context + generation) batch the block form receives
+        gen-only logits (context requests draft from accepted tokens they do not
+        have yet), so ``num_contexts`` slices the full-batch per-request metadata
+        down to the gen segment. The step form receives full-batch logits
+        (context requests draft from target hidden states already available), so
+        no slicing is needed.
+        """
+        is_block = logits.dim() == 3
+        use_rejection = getattr(spec_metadata, "use_rejection_sampling", False)
+
+        # Draft tokens use argmax unless rejection sampling is engaged for a
+        # non-greedy batch. Rejection sampling is the only path that needs the
+        # draft's stochastic proposal distribution (stored in draft_probs); every
+        # other path (all-greedy, or non-greedy with strict/exact-match
+        # acceptance) accepts a draft token only when it equals the target's
+        # choice, and there argmax maximizes the acceptance rate (E[accept] =
+        # max_i p_i >= sum_i p_i^2 = E[accept] for a stochastic draft). This
+        # matches sglang/vLLM, which draft with argmax/top-k by default and apply
+        # sampling params only on the target/acceptance side.
+        advanced = use_rejection and not spec_metadata.is_all_greedy_sample
+
+        # All samplers below return tokens in draft-vocab space; d2t is applied
+        # once after the branch.
+        if not advanced:
+            # greedy_sample_draft_with_tp_gather expects 2D [tokens, vocab];
+            # flatten the block form and restore its shape (step form is 2D).
+            batch_shape = logits.shape[:-1]
+            tokens = self.greedy_sample_draft_with_tp_gather(
+                logits.reshape(-1, logits.shape[-1]), spec_metadata,
+                mapping_lm_head_tp)
+            tokens = tokens.reshape(batch_shape)
+        else:
+            # Advanced sampling gathers the vocab-sharded draft logits to full
+            # vocab, then samples (scattering this step's proposal distribution
+            # into draft_probs).
+            logits = self.maybe_gather_sharded_draft_logits(
+                logits, spec_metadata, mapping_lm_head_tp)
+            if is_block:
+                tokens = self.advanced_sample_draft_block(
+                    logits, spec_metadata, num_contexts, batch_size).long()
+            else:
+                tokens = self.advanced_sample_draft(logits,
+                                                    spec_metadata,
+                                                    batch_size,
+                                                    draft_step=draft_step)
 
-        # Scatter probs into the slot-indexed buffer (shaped
-        # [max_num_requests, max_draft_len, vocab_size]). Each request's data
-        # always lands at its stable py_seq_slot row regardless of batch
-        # composition shifts across iterations.
-        assert spec_metadata.batch_slot_ids is not None, (
-            "batch_slot_ids must be populated by "
-            "populate_sampling_params_for_one_model before draft probs storage")
-        batch_slots = spec_metadata.batch_slot_ids[:batch_size]
-        vocab = probs.shape[-1]
-        spec_metadata.draft_probs[batch_slots, draft_step, :vocab] = probs
-        spec_metadata.draft_probs_last_dim = vocab
-
-        if d2t is not None:
-            draft_tokens = d2t[draft_tokens] + draft_tokens
+        # Map draft-vocab token ids to target vocab (no-op for shared vocab).
+        if self._d2t is not None:
+            tokens = self._d2t[tokens] + tokens
 
-        return draft_tokens.type(torch.int32)
+        return tokens.type(torch.int32)
 
     def _execute_guided_decoder_if_present(self, logits):
         """Execute guided decoder on target model logits if available."""
diff --git a/tensorrt_llm/_torch/speculative/mtp.py b/tensorrt_llm/_torch/speculative/mtp.py
index 6afed0fc02e9..fcf0603a04c1 100644
--- a/tensorrt_llm/_torch/speculative/mtp.py
+++ b/tensorrt_llm/_torch/speculative/mtp.py
@@ -5,10 +5,8 @@
 import torch
 
 from tensorrt_llm._utils import prefer_pinned
-from tensorrt_llm.mapping import Mapping
 
 from ..attention_backend import AttentionMetadata
-from ..distributed.ops import allgather
 from ..pyexecutor.llm_request import LlmRequest
 from ..pyexecutor.resource_manager import BaseResourceManager, SlotManager
 from ..pyexecutor.sampler import TorchSampler
@@ -278,10 +276,19 @@ class MTPWorker(SpecWorkerBase):
     def __init__(self,
                  spec_config: "MTPDecodingConfig",
                  model_config=None,
-                 use_separate_draft_kv_cache: bool = False):
+                 use_separate_draft_kv_cache: bool = False,
+                 *,
+                 mapping=None):
         super().__init__(use_separate_draft_kv_cache)
         self.spec_config = spec_config
         self.model_config = model_config
+        # Use the mapping passed by get_spec_worker (the same Mapping that shards
+        # the draft LM head). model_config.mapping is unreliable here -- for the
+        # MTP worker it can be None -- so reading it caused greedy draft sampling
+        # to skip the TP argmax gather and desync the ranks into a hang under
+        # plain TP + overlap scheduler.
+        self.mapping = mapping if mapping is not None else getattr(
+            model_config, "mapping", None)
         self.is_thop = False
         self.sa_enhancer: Optional[SADraftEnhancer] = None
         if spec_config.sa_config is not None:
@@ -469,7 +476,10 @@ def forward(
                     self.guided_decoder.execute_draft_batch(logits,
                                                             draft_step=i)
 
-                new_draft_token = self.draft_sampler(logits)
+                new_draft_token = self.sample_draft_tokens(logits,
+                                                           spec_metadata,
+                                                           batch_size,
+                                                           draft_step=i)
                 next_draft_tokens.append(new_draft_token)
                 # shift input_ids and hidden_states
                 input_ids = draft_inputs["input_ids"]
@@ -841,6 +851,16 @@ def sample_and_accept_draft_tokens(
                 runtime_draft_len,
                 spec_metadata=spec_metadata)
 
+        # Rejection sampling acceptance. _can_use_rejection_sampling() requires
+        # use_rejection_sampling and a non-all-greedy batch; otherwise falls
+        # through to strict acceptance below. Context rows take the target's
+        # first sampled token; gen rows run the rejection kernel.
+        elif self._can_use_rejection_sampling(spec_metadata):
+            draft_tokens = spec_metadata.draft_tokens.reshape(
+                num_gens, runtime_draft_len)
+            accepted_tokens, num_accepted_tokens = self._accept_draft_tokens(
+                logits, draft_tokens, num_contexts, batch_size, spec_metadata)
+
         # Strict acceptance
         else:
             if self.is_thop:
@@ -1092,80 +1112,3 @@ def prepare_drafter_inputs(
             "hidden_states": return_hidden_states,
             "attn_metadata": attn_metadata,
         }
-
-    @torch.compile(options={"max-autotune": True})
-    def get_local_max_and_combined(self, logits, mapping_lm_tp=None):
-        local_max_values, local_argmax = torch.max(logits, dim=-1, keepdim=True)
-        # Adjust indices based on TP rank and size
-        vocab_per_rank = logits.shape[-1]
-        mapping_lm_tp = mapping_lm_tp if mapping_lm_tp is not None else self.model_config.mapping
-        max_index_per_rank = local_argmax.type(
-            torch.int32) + (mapping_lm_tp.tp_rank * vocab_per_rank)
-        # Use torch.stack and flatten instead of view+cat to avoid torch.compile issues
-        # Convert both to float32 to ensure consistent dtype
-        max_index_per_rank_float = max_index_per_rank.float()
-        local_max_values_float32 = local_max_values.float()
-
-        # Stack and flatten to get interleaved layout: [idx0, val0, idx1, val1, ...]
-        combined = torch.stack(
-            [max_index_per_rank_float, local_max_values_float32],
-            dim=-1).flatten(-2)
-        return combined
-
-    @torch.compile(options={"max-autotune": True})
-    def get_draft_tokens_from_gathered(self, gathered):
-        gathered_indices_float = gathered[..., 0::2]  # Even positions: indices
-        gathered_values_float = gathered[..., 1::2]  # Odd positions: values
-
-        # Find the rank with maximum value
-        max_indices = torch.argmax(gathered_values_float, dim=-1, keepdim=True)
-
-        # Get the corresponding token indices and convert back to int32
-        draft_tokens = torch.gather(gathered_indices_float, -1,
-                                    max_indices).squeeze(-1).type(torch.int32)
-        return draft_tokens
-
-    def draft_sampler(
-        self,
-        logits: torch.Tensor,
-        mapping_lm_head_tp: Mapping = None,
-    ):
-        '''
-        Sampling draft tokens.
-
-        Args:
-            logits: torch.Tensor
-                [num_tokens, vocab_size]
-                Logits produced by the draft model.
-
-        Returns:
-            draft_tokens: torch.Tensor
-                [batch_size * max_draft_len]
-                Draft token ids. Flattened.
-        '''
-        if (self.model_config is not None
-                and hasattr(self.model_config, 'mapping')
-                and self.model_config.mapping.tp_size
-                > 1) and not (self.model_config.mapping.enable_attention_dp):
-            combined = self.get_local_max_and_combined(logits)
-            gathered = allgather(combined, self.model_config.mapping, dim=-1)
-            draft_tokens = self.get_draft_tokens_from_gathered(gathered)
-        elif (self.model_config is not None
-              and hasattr(self.model_config, 'mapping')
-              and self.model_config.mapping.tp_size
-              > 1) and self.model_config.mapping.enable_lm_head_tp_in_adp:
-            # For ADP + LM head TP mode, we need to find the global argmax across all TP ranks
-            combined = self.get_local_max_and_combined(logits,
-                                                       mapping_lm_head_tp)
-            gathered = allgather(combined, mapping_lm_head_tp, dim=-1)
-            batch_size = logits.shape[0]
-            local_batch_size = batch_size // mapping_lm_head_tp.tp_size
-            gathered = gathered.view(mapping_lm_head_tp.tp_size,
-                                     local_batch_size, -1)
-            sliced_gathered = gathered[mapping_lm_head_tp.tp_rank]
-            draft_tokens = self.get_draft_tokens_from_gathered(sliced_gathered)
-        else:
-            # Simple argmax if no TP or no model config
-            draft_tokens = self._draft_sampler_greedy(logits)
-
-        return draft_tokens
diff --git a/tensorrt_llm/_torch/speculative/pard.py b/tensorrt_llm/_torch/speculative/pard.py
index 34a4509314cd..969add54136e 100644
--- a/tensorrt_llm/_torch/speculative/pard.py
+++ b/tensorrt_llm/_torch/speculative/pard.py
@@ -196,25 +196,8 @@ def forward(
 
         self._execute_guided_decoder_if_present(logits)
 
-        # draft_tokens buffer has (2K-1) entries per gen request; extract the K real drafts
-        if num_gens > 0:
-            draft_tokens = spec_metadata.draft_tokens[: num_gens * (2 * K - 1)]
-            draft_tokens = draft_tokens.reshape(num_gens, 2 * K - 1)[:, :K]
-        else:
-            draft_tokens = spec_metadata.draft_tokens.reshape(0, K)
-
-        # logits have 2K entries per gen request; extract K+1 for acceptance
-        if num_gens > 0:
-            ctx_logits = logits[:num_contexts]
-            vocab_size = logits.shape[-1]
-            gen_logits_2k = logits[num_contexts:].reshape(num_gens, 2 * K, vocab_size)
-            gen_logits_kp1 = gen_logits_2k[:, : K + 1, :].reshape(-1, vocab_size)
-            logits_for_accept = torch.cat([ctx_logits, gen_logits_kp1], dim=0)
-        else:
-            logits_for_accept = logits
-
-        accepted_tokens, num_accepted_tokens = self._sample_and_accept_draft_tokens_base(
-            logits_for_accept, draft_tokens, num_contexts, batch_size, spec_metadata
+        accepted_tokens, num_accepted_tokens = self.sample_and_accept_draft_tokens(
+            logits, attn_metadata, spec_metadata
         )
 
         # Pad accepted_tokens from (batch, K+1) to (batch, 2K) to match sampler buffer
@@ -285,14 +268,12 @@ def forward(
                 vocab_size = gen_logits.shape[-1]
                 gen_logits = gen_logits.reshape(num_gens, K, vocab_size)
 
-                # Use torch.argmax directly to avoid cute_argmax stride issues
-                d2t = getattr(draft_model.model, "d2t", None)
-                gen_draft_tokens = torch.argmax(gen_logits, dim=-1, keepdim=False).long()
-
-                if d2t is not None:
-                    gen_draft_tokens = d2t[gen_draft_tokens] + gen_draft_tokens
-
-                gen_draft_tokens = gen_draft_tokens.type(torch.int32)
+                gen_draft_tokens = self.sample_draft_tokens(
+                    gen_logits,
+                    spec_metadata,
+                    batch_size,
+                    num_contexts=num_contexts,
+                )
 
                 if self.sa_enhancer is not None and sa_manager is not None:
                     gen_draft_tokens = self.sa_enhancer.maybe_override_all_draft_tokens(
@@ -314,6 +295,12 @@ def forward(
         else:
             gen_draft_tokens = torch.empty((0, 2 * K - 1), dtype=torch.int32, device="cuda")
 
+        # Context requests are not drafted by the block worker (zero placeholder
+        # token); fill their draft-prob slot rows with a one-hot placeholder so
+        # they are a legal distribution when they become gen requests next iter.
+        gen_vocab = vocab_size if num_gens > 0 else None
+        self.write_context_onehot_draft_probs(spec_metadata, num_contexts, num_gens, K, gen_vocab)
+
         if num_contexts > 0 and num_gens > 0:
             ctx_draft_tokens = torch.zeros(
                 (num_contexts, 2 * K - 1), dtype=torch.int32, device="cuda"
@@ -347,23 +334,26 @@ def forward(
             "next_new_tokens": next_new_tokens,
         }
 
-    def draft_decoder(
-        self,
-        logits: torch.Tensor,
-        draft_model: nn.Module,
-    ):
-        """
-        Sample draft tokens using greedy decoding.
-
-        Args:
-            logits: [num_tokens, vocab_size] from the draft model.
-            draft_model: The draft model (used to read the d2t mapping).
-
-        Returns:
-            draft_tokens: [batch_size * max_draft_len] flattened token ids.
-        """
-        d2t = getattr(draft_model.model, "d2t", None)
-        return self._draft_sampler_greedy(logits, d2t)
+    def _reshape_draft_tokens_for_accept(self, spec_metadata, num_gens, device):
+        # draft_tokens buffer has (2K-1) entries per gen request; extract the K real drafts.
+        K = spec_metadata.runtime_draft_len
+        if num_gens > 0:
+            draft_tokens = spec_metadata.draft_tokens[: num_gens * (2 * K - 1)]
+            return draft_tokens.reshape(num_gens, 2 * K - 1)[:, :K]
+        # Context-only batch: return an empty view without reshaping the
+        # (possibly preallocated, non-empty) draft_tokens buffer.
+        return torch.empty((0, K), dtype=torch.int32, device=device)
+
+    def _reshape_logits_for_accept(self, logits, num_contexts, num_gens, spec_metadata):
+        # logits have 2K entries per gen request; extract K+1 for acceptance.
+        if num_gens == 0:
+            return logits
+        K = spec_metadata.runtime_draft_len
+        ctx_logits = logits[:num_contexts]
+        vocab_size = logits.shape[-1]
+        gen_logits_2k = logits[num_contexts:].reshape(num_gens, 2 * K, vocab_size)
+        gen_logits_kp1 = gen_logits_2k[:, : K + 1, :].reshape(-1, vocab_size)
+        return torch.cat([ctx_logits, gen_logits_kp1], dim=0)
 
     def prepare_1st_drafter_inputs(
         self,
diff --git a/tensorrt_llm/_torch/speculative/suffix_automaton.py b/tensorrt_llm/_torch/speculative/suffix_automaton.py
index 1eb7be0bcea5..12bb77208fe4 100644
--- a/tensorrt_llm/_torch/speculative/suffix_automaton.py
+++ b/tensorrt_llm/_torch/speculative/suffix_automaton.py
@@ -87,6 +87,17 @@ class SuffixAutomatonManager(BaseResourceManager):
     and are CUDA graph compatible.
 
     Used as the resource manager for both NGram and MTP+SA speculative decoding.
+
+    Rejection sampling is not supported for NGram or SA drafting. Both are
+    retrieval-based drafters: they propose draft tokens by matching against
+    previously seen context (n-gram lookup / suffix-automaton traversal) instead
+    of sampling from a model head, so they emit only token ids with no per-token
+    proposal distribution ``q(x)`` over the vocabulary. Rejection sampling's
+    acceptance test needs ``q`` to form the ratio ``min(1, p(x)/q(x))`` and the
+    residual ``(p - q)+`` correction on rejection; without ``q`` the target
+    distribution cannot be preserved. The gate lives in
+    ``TorchLlmArgs.validate_speculative_config`` (NGram falls outside the
+    supported-mode whitelist; SA is rejected via ``rs_sa_active``).
     """
 
     def __init__(
diff --git a/tensorrt_llm/_torch/speculative/utils.py b/tensorrt_llm/_torch/speculative/utils.py
index d7a0e8185da5..4f9c5af8846a 100644
--- a/tensorrt_llm/_torch/speculative/utils.py
+++ b/tensorrt_llm/_torch/speculative/utils.py
@@ -42,6 +42,37 @@ def _is_effective_dynamic_tree(spec_config) -> bool:
             and getattr(spec_config, 'dynamic_tree_max_topK', 0) > 1)
 
 
+def _get_draft_vocab_size(spec_config, target_vocab_size: int) -> int:
+    """Draft-model vocab size, used to decide whether rejection sampling needs
+    the d2t-expanded ``full_draft_probs`` buffer (only when it differs from the
+    target vocab).
+
+    Reads the draft model's ``config.json`` from ``spec_config.speculative_model``.
+    Eagle3 configs store the target vocab in ``vocab_size`` and the reduced head
+    width in ``draft_vocab_size``, so ``draft_vocab_size`` is read first, falling
+    back to ``vocab_size`` (or a nested ``text_config.vocab_size``). Returns
+    ``target_vocab_size`` (shared vocab, no buffer needed) when there is no
+    separate draft model or the config cannot be read.
+    """
+    draft_dir = getattr(spec_config, "speculative_model", None)
+    if not draft_dir:
+        return target_vocab_size
+    try:
+        import json
+        import os
+        cfg_path = os.path.join(str(draft_dir), "config.json")
+        if not os.path.isfile(cfg_path):
+            return target_vocab_size
+        with open(cfg_path) as f:
+            cfg = json.load(f)
+        vs = cfg.get("draft_vocab_size") or cfg.get("vocab_size")
+        if vs is None:
+            vs = (cfg.get("text_config") or {}).get("vocab_size")
+        return int(vs) if vs else target_vocab_size
+    except (OSError, ValueError, TypeError, AttributeError):
+        return target_vocab_size
+
+
 def get_spec_metadata(spec_config,
                       model_config,
                       max_num_requests,
@@ -57,6 +88,9 @@ def get_spec_metadata(spec_config,
     num_seq_slots = (num_seq_slots
                      if num_seq_slots is not None else max_num_requests)
     vocab_size = getattr(model_config, "vocab_size", 0)
+    # Draft-model vocab size, used to gate the d2t-expanded full_draft_probs
+    # buffer allocation (see SpecMetadata.prepare_rejection_sampling_buffers).
+    draft_vocab_size = _get_draft_vocab_size(spec_config, vocab_size)
     if spec_config.spec_dec_mode.is_mtp_eagle_one_model():
         # MTP Eagle one-model reuses Eagle3 one-model metadata for the
         # unified worker/sampler/slot_ids plumbing, but skips per-layer
@@ -76,6 +110,7 @@ def get_spec_metadata(spec_config,
             use_rejection_sampling=use_rejection_sampling,
             vocab_size=vocab_size,
             num_seq_slots=num_seq_slots,
+            draft_vocab_size=draft_vocab_size,
             spec_resource_manager=spec_resource_manager,
         )
     if spec_config.spec_dec_mode.is_mtp_vanilla():
@@ -86,6 +121,9 @@ def get_spec_metadata(spec_config,
             mtp_num_modules=spec_config.max_draft_len,
             max_num_requests=max_num_requests,
             mtp_hidden_states_manager=spec_resource_manager,
+            use_rejection_sampling=use_rejection_sampling,
+            vocab_size=vocab_size,
+            draft_vocab_size=draft_vocab_size,
         )
     if spec_config.spec_dec_mode.is_mtp_eagle():
         return Eagle3SpecMetadata(
@@ -134,6 +172,7 @@ def get_spec_metadata(spec_config,
             layers_to_capture=spec_config.eagle3_layers_to_capture,
             use_rejection_sampling=use_rejection_sampling,
             vocab_size=vocab_size,
+            draft_vocab_size=draft_vocab_size,
             spec_resource_manager=spec_resource_manager,
             use_dynamic_tree=_is_effective_dynamic_tree(spec_config),
             eagle_choices=spec_config.eagle_choices,
@@ -145,6 +184,9 @@ def get_spec_metadata(spec_config,
             spec_dec_mode=spec_config.spec_dec_mode,
             max_num_requests=max_num_requests,
             spec_resource_manager=spec_resource_manager,
+            use_rejection_sampling=use_rejection_sampling,
+            vocab_size=vocab_size,
+            draft_vocab_size=draft_vocab_size,
         )
     if spec_config.spec_dec_mode.is_dflash():
         target_layer_ids = getattr(spec_config, 'target_layer_ids', None)
@@ -157,6 +199,9 @@ def get_spec_metadata(spec_config,
             hidden_size=model_config.hidden_size,
             max_num_tokens=max_num_tokens,
             dtype=model_config.torch_dtype,
+            use_rejection_sampling=use_rejection_sampling,
+            vocab_size=vocab_size,
+            draft_vocab_size=draft_vocab_size,
         )
     if spec_config.spec_dec_mode.is_draft_target_one_model():
         return DraftTargetOneModelSpecMetadata(
@@ -165,6 +210,9 @@ def get_spec_metadata(spec_config,
             spec_dec_mode=spec_config.spec_dec_mode,
             max_num_requests=max_num_requests,
             max_num_tokens=max_num_tokens,
+            use_rejection_sampling=use_rejection_sampling,
+            vocab_size=vocab_size,
+            draft_vocab_size=draft_vocab_size,
         )
     if spec_config.spec_dec_mode.is_save_hidden_states():
         return SaveHiddenStatesSpecMetadata(
@@ -396,10 +444,15 @@ def get_spec_worker(spec_config,
                     use_separate_draft_kv_cache: bool = False):
     spec_dec_mode = spec_config.spec_dec_mode
     if spec_dec_mode.is_mtp_vanilla():
-        return MTPWorker(spec_config, model_config, use_separate_draft_kv_cache)
+        return MTPWorker(spec_config,
+                         model_config,
+                         use_separate_draft_kv_cache,
+                         mapping=mapping)
     if spec_dec_mode.is_mtp_eagle_one_model():
-        return MTPEagleWorker(spec_config, model_config,
-                              use_separate_draft_kv_cache)
+        return MTPEagleWorker(spec_config,
+                              model_config,
+                              use_separate_draft_kv_cache,
+                              mapping=mapping)
     if spec_dec_mode.is_eagle3_one_model():
         if _is_effective_dynamic_tree(spec_config):
             return Eagle3OneModelDynamicTreeWorker(spec_config, mapping,
diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py
index cc6a66e53a08..96bb64a02f3e 100644
--- a/tensorrt_llm/llmapi/llm_args.py
+++ b/tensorrt_llm/llmapi/llm_args.py
@@ -1690,14 +1690,17 @@ def validate_draft_len_schedule_and_sort(cls, v, info):
 
     @model_validator(mode='after')
     def validate_rejection_sampling_config(self):
-        """Disable rejection sampling when SA-enhanced configurations are
-        active, since SA may override the proposed draft tokens. This is a
-        silent fallback so the new default (True) does not break sa_config
-        users.
+        """Disable rejection sampling when SA-enhanced configurations are active.
+
+        Only silently disable a default-inherited value; an explicit
+        ``use_rejection_sampling=True`` is preserved so
+        ``TorchLlmArgs.validate_speculative_config`` can raise for the
+        unsupported SA combination.
         """
         if self.use_rejection_sampling and getattr(self, 'sa_config',
                                                    None) is not None:
-            self.use_rejection_sampling = False
+            if "use_rejection_sampling" not in self.model_fields_set:
+                self.use_rejection_sampling = False
         return self
 
     @model_validator(mode='before')
@@ -4963,17 +4966,127 @@ def validate_speculative_config(self):
                     exclude={"decoding_type"})
                 self.speculative_config = Eagle3DecodingConfig(**eagle_data)
 
-            if self.speculative_config.use_rejection_sampling and not isinstance(
-                    self.speculative_config, Eagle3DecodingConfig):
-                # Rejection sampling is only wired up for Eagle3 one-model paths.
-                # Silently fall back for other spec types so the new default
-                # (True) does not break them.
-                # TODO: extend rejection sampling to the remaining speculative
-                # decoding paths (MTP / DraftTarget / PARD / DFlash /
-                # SaveHiddenStates / SA) and unify the dispatch in SpecMetadata
-                # so new spec algorithms get rejection sampling for free; once
-                # all paths are covered this whitelist guard can be removed.
-                self.speculative_config.use_rejection_sampling = False
+            if self.speculative_config.use_rejection_sampling:
+                # Supported paths: Eagle3 one-model, MTP-Eagle one-model,
+                # vanilla MTP, PARD, DFlash, DraftTarget one-model. Classify by
+                # spec-dec mode; MTP-Eagle one-model shares the Eagle3 one-model
+                # worker/metadata/sampler, so it rides the same unified
+                # production/acceptance path.
+                #
+                # Every supported path runs a neural draft head that yields a
+                # per-token proposal distribution q(x). Rejection sampling needs
+                # q to form min(1, p/q) and the (p - q)+ residual correction, so
+                # retrieval-based drafters that emit only token ids with no q are
+                # excluded: NGram is simply absent from this whitelist, and SA
+                # (sa_config) is rejected below via rs_sa_active.
+                from tensorrt_llm._torch.speculative.interface import \
+                    SpeculativeDecodingMode as TorchSpeculativeDecodingMode
+
+                spec_mode = self.speculative_config.spec_dec_mode
+                is_supported = (isinstance(self.speculative_config,
+                                           Eagle3DecodingConfig)
+                                or spec_mode == TorchSpeculativeDecodingMode.MTP
+                                or spec_mode.is_mtp_eagle_one_model()
+                                or spec_mode.is_pard() or spec_mode.is_dflash()
+                                or spec_mode.is_draft_target_one_model())
+
+                # Combinations that break the proposal-distribution invariant.
+                # SA is gated for every method; relaxed / parallel / guided gates
+                # apply only to the newly wired methods (vanilla MTP, PARD,
+                # DFlash, DraftTarget one-model).
+                is_new_rejection_method = (
+                    spec_mode == TorchSpeculativeDecodingMode.MTP
+                    or spec_mode.is_pard() or spec_mode.is_dflash()
+                    or spec_mode.is_draft_target_one_model())
+                # Plain tensor parallelism is supported (the draft path
+                # all-gathers vocab-sharded draft logits before rejection, see
+                # SpecWorkerBase.maybe_gather_sharded_draft_logits).
+                # attention-DP and context parallelism remain gated.
+                rs_parallel_active = (self.context_parallel_size > 1
+                                      or self.enable_attention_dp)
+                rs_guided_active = self.guided_decoding_backend is not None
+                rs_sa_active = getattr(self.speculative_config, "sa_config",
+                                       None) is not None
+                rs_relaxed_active = getattr(
+                    self.speculative_config,
+                    "use_relaxed_acceptance_for_thinking", False)
+                # Eagle3 dynamic-tree rejection records at most kMaxTriedPerLevel
+                # (=32, cpp/.../speculativeDecoding/dynamicTreeKernels.cu) tried
+                # siblings per level; dynamic_tree_max_topK > 32 yields a wrong
+                # correction distribution, so cap supported topK at 32.
+                _DYNAMIC_TREE_REJECTION_MAX_TOPK = 32
+                rs_dynamic_tree_active = getattr(self.speculative_config,
+                                                 "use_dynamic_tree", False)
+                rs_dynamic_tree_topk_unsupported = (
+                    rs_dynamic_tree_active
+                    and (self.speculative_config.dynamic_tree_max_topK
+                         or 0) > _DYNAMIC_TREE_REJECTION_MAX_TOPK)
+                mtp_unsupported_combo = (
+                    rs_sa_active or rs_dynamic_tree_topk_unsupported
+                    or (is_new_rejection_method and
+                        (rs_relaxed_active or rs_parallel_active
+                         or rs_guided_active)))
+
+                if not is_supported or mtp_unsupported_combo:
+                    # Explicit opt-in raises; a default-inherited value is
+                    # silently disabled for backward compatibility.
+                    explicitly_set = "use_rejection_sampling" in \
+                        self.speculative_config.model_fields_set
+                    if explicitly_set:
+                        reasons = []
+                        if not is_supported:
+                            reasons.append(
+                                f"spec mode {spec_mode.name} is not a supported "
+                                "rejection-sampling path (supported: the Eagle3 "
+                                "one-model path, vanilla MTP, PARD, DFlash, and "
+                                "DraftTarget one-model)")
+                        if rs_sa_active:
+                            reasons.append("SA (sa_config) is active")
+                        if rs_dynamic_tree_topk_unsupported:
+                            reasons.append(
+                                "dynamic_tree_max_topK="
+                                f"{self.speculative_config.dynamic_tree_max_topK}"
+                                " exceeds the dynamic-tree rejection kernel limit "
+                                f"of {_DYNAMIC_TREE_REJECTION_MAX_TOPK} tried "
+                                "siblings per level")
+                        if is_new_rejection_method:
+                            if rs_relaxed_active:
+                                reasons.append(
+                                    "relaxed-thinking acceptance is enabled")
+                            if rs_parallel_active:
+                                reasons.append(
+                                    "tensor/context parallelism or attention-DP "
+                                    "is active (the draft path resolves only "
+                                    "the global argmax, not full distributions)"
+                                )
+                            if rs_guided_active:
+                                reasons.append("guided decoding is enabled")
+                        raise ValueError(
+                            "use_rejection_sampling=True is not supported for "
+                            f"this configuration: {'; '.join(reasons)}.")
+                    self.speculative_config.use_rejection_sampling = False
+                elif self.speculative_config.use_rejection_sampling:
+                    # The non-dynamic-tree one-model rejection path depends on
+                    # FlashInfer (sampling_from_probs /
+                    # chain_speculative_sampling); fail fast if it is missing.
+                    # The dynamic-tree path uses
+                    # verify_dynamic_tree_rejection_out /
+                    # compute_probs_from_logits (CUDA op with a PyTorch
+                    # fallback) and does not require FlashInfer.
+                    rs_dynamic_tree_active = getattr(self.speculative_config,
+                                                     "use_dynamic_tree", False)
+                    if not rs_dynamic_tree_active:
+                        from tensorrt_llm._torch.flashinfer_utils import \
+                            IS_FLASHINFER_AVAILABLE
+                        if not IS_FLASHINFER_AVAILABLE:
+                            raise ValueError(
+                                "use_rejection_sampling=True requires FlashInfer, "
+                                "which is not available in this environment. "
+                                "Install flashinfer-python or set "
+                                "use_rejection_sampling=False.")
+                # The enabled paths (Eagle3, vanilla MTP, PARD, DFlash,
+                # DraftTarget one-model) share the SpecMetadata slot-indexed
+                # dispatch.
 
             if isinstance(self.speculative_config, PARDDecodingConfig):
                 assert self.speculative_config.max_draft_len > 0, "PARD max_draft_len must be > 0"
diff --git a/tests/integration/test_lists/test-db/l0_h100.yml b/tests/integration/test_lists/test-db/l0_h100.yml
index 3fc1f846b163..df098b45454b 100644
--- a/tests/integration/test_lists/test-db/l0_h100.yml
+++ b/tests/integration/test_lists/test-db/l0_h100.yml
@@ -49,6 +49,7 @@ l0_h100:
   # function's leaves on Hopper, keep the rest of the sampler dir.
   - unittest/_torch/sampler -k "not test_speculative_d2h_parity_real_predictor"
   - unittest/_torch/speculative/test_eagle3.py
+  - unittest/_torch/speculative/test_rejection_buffers_guard.py
   - unittest/_torch/speculative/hw_agnostic
   - unittest/_torch/thop/parallel
   - unittest/_torch/thop/parallel_hw_agnostic
diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dflash.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dflash.py
index 2a3def69a639..bb920e27bfae 100644
--- a/tests/unittest/_torch/speculative/hw_agnostic/test_dflash.py
+++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dflash.py
@@ -104,5 +104,35 @@ def test_dflash_qwen3_5_4b(disable_overlap_scheduler: bool):
     _run_and_check(llm_config, min_avg_accepted=1.0)
 
 
+@pytest.mark.high_cuda_memory
+def test_dflash_qwen3_8b_rejection():
+    """DFlash with rejection sampling on: the block-capture rejection path
+    (draft-prob scatter -> fail-closed guard -> rejection acceptance) runs
+    end-to-end with non-greedy sampling and produces coherent output."""
+    total_mem_gb = torch.cuda.get_device_properties(0).total_memory / 1e9
+    if total_mem_gb < 35:
+        pytest.skip("Not enough memory to load target + draft model")
+
+    models_path = llm_models_root()
+    llm_config = _make_llm_config(
+        target_model_dir=f"{models_path}/Qwen3/Qwen3-8B",
+        dflash_model_dir=f"{models_path}/Qwen3-8B-DFlash-b16",
+        disable_overlap_scheduler=True,
+    )
+    llm_config["speculative_config"].use_rejection_sampling = True
+
+    llm = LLM(**llm_config)
+    # Non-greedy so rejection sampling actually engages (all-greedy bypasses it).
+    outputs = llm.generate(
+        PROMPTS, SamplingParams(max_tokens=32, temperature=0.8, top_p=0.95, top_k=50, seed=1234)
+    )
+    llm.shutdown()
+
+    assert len(outputs) == len(PROMPTS)
+    for out in outputs:
+        assert len(out.outputs[0].token_ids) > 0
+        assert out.outputs[0].text.strip()
+
+
 if __name__ == "__main__":
     unittest.main()
diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_draft_target.py b/tests/unittest/_torch/speculative/hw_agnostic/test_draft_target.py
index 4bda8318c08c..dd2b3e5cd7a0 100644
--- a/tests/unittest/_torch/speculative/hw_agnostic/test_draft_target.py
+++ b/tests/unittest/_torch/speculative/hw_agnostic/test_draft_target.py
@@ -66,5 +66,51 @@ def test_llama_draft_target(use_cuda_graph: bool, attn_backend: str):
         assert similar(text_spec, text_ref)
 
 
+@pytest.mark.high_cuda_memory
+def test_llama_draft_target_rejection():
+    """DraftTarget one-model with rejection sampling on: the rejection path
+    (draft-prob capture -> fail-closed guard -> rejection acceptance) runs
+    end-to-end with non-greedy sampling and produces coherent output."""
+    total_mem_gb = torch.cuda.get_device_properties(0).total_memory / 1e9
+    if total_mem_gb < 60:
+        pytest.skip("Not enough memory to load target model")
+
+    models_path = llm_models_root()
+    target_model_dir = f"{models_path}/llama-3.1-model/Llama-3.1-8B-Instruct"
+    draft_model_dir = f"{models_path}/llama-3.2-models/Llama-3.2-1B-Instruct"
+
+    spec_config = DraftTargetDecodingConfig(
+        max_draft_len=4,
+        speculative_model=draft_model_dir,
+        use_rejection_sampling=True,
+    )
+
+    llm = LLM(
+        model=target_model_dir,
+        backend="pytorch",
+        attn_backend="TRTLLM",
+        disable_overlap_scheduler=True,
+        max_batch_size=2,
+        kv_cache_config=KvCacheConfig(enable_block_reuse=False, max_tokens=8192),
+        max_num_tokens=2048,
+        speculative_config=spec_config,
+    )
+    prompts = [
+        "The capital of France is",
+        "The president of the United States is",
+    ]
+    # Non-greedy so rejection sampling actually engages (all-greedy bypasses it).
+    sampling_params = SamplingParams(
+        max_tokens=32, temperature=0.8, top_p=0.95, top_k=50, seed=1234
+    )
+    outputs = llm.generate(prompts, sampling_params)
+    llm.shutdown()
+
+    assert len(outputs) == len(prompts)
+    for out in outputs:
+        assert len(out.outputs[0].token_ids) > 0
+        assert out.outputs[0].text.strip()
+
+
 if __name__ == "__main__":
     unittest.main()
diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_mtp.py b/tests/unittest/_torch/speculative/hw_agnostic/test_mtp.py
index d987ed442b5d..a381dbd94907 100644
--- a/tests/unittest/_torch/speculative/hw_agnostic/test_mtp.py
+++ b/tests/unittest/_torch/speculative/hw_agnostic/test_mtp.py
@@ -1,13 +1,20 @@
+import os
+import sys
 import unittest
 
+import pytest
 import torch
 from parameterized import parameterized
 
 import tensorrt_llm
+from tensorrt_llm import LLM, SamplingParams
 from tensorrt_llm._torch.attention_backend import TrtllmAttentionMetadata
 from tensorrt_llm._torch.metadata import KVCacheParams
 from tensorrt_llm._torch.speculative.mtp import MTPHiddenStatesManager, MTPSpecMetadata, MTPWorker
-from tensorrt_llm.llmapi import MTPDecodingConfig
+from tensorrt_llm.llmapi import KvCacheConfig, MTPDecodingConfig
+
+sys.path.append(os.path.join(os.path.dirname(__file__), ".."))
+from utils.llm_data import llm_models_root
 
 
 def unittest_name_func(testcase_func, param_num, param):
@@ -1724,3 +1731,90 @@ def test_prepare_drafter_inputs(
 
         torch.testing.assert_close(draft_inputs["input_ids"], ref_input_ids)
         torch.testing.assert_close(draft_inputs["hidden_states"], ref_previous_hidden_states)
+
+
+@pytest.mark.high_cuda_memory
+def test_vanilla_mtp_rejection():
+    """Vanilla MTP with rejection sampling on: the per-step rejection path
+    (draft-prob scatter -> fail-closed guard -> rejection acceptance) runs
+    end-to-end with non-greedy sampling and produces coherent output."""
+    total_mem_gb = torch.cuda.get_device_properties(0).total_memory / 1e9
+    if total_mem_gb < 60:
+        pytest.skip("Not enough memory to load DeepSeek-V3-Lite")
+
+    model_dir = f"{llm_models_root()}/DeepSeek-V3-Lite/bf16"
+    spec_config = MTPDecodingConfig(
+        max_draft_len=1,
+        use_mtp_vanilla=True,
+        speculative_model=model_dir,
+        use_rejection_sampling=True,
+    )
+    llm = LLM(
+        model=model_dir,
+        backend="pytorch",
+        disable_overlap_scheduler=True,
+        max_batch_size=2,
+        kv_cache_config=KvCacheConfig(enable_block_reuse=False, max_tokens=8192),
+        max_num_tokens=2048,
+        speculative_config=spec_config,
+        trust_remote_code=True,
+    )
+    prompts = [
+        "The capital of France is",
+        "The president of the United States is",
+    ]
+    # Non-greedy so rejection sampling actually engages (all-greedy bypasses it).
+    sampling_params = SamplingParams(
+        max_tokens=32, temperature=0.8, top_p=0.95, top_k=50, seed=1234
+    )
+    outputs = llm.generate(prompts, sampling_params)
+    llm.shutdown()
+
+    assert len(outputs) == len(prompts)
+    for out in outputs:
+        assert len(out.outputs[0].token_ids) > 0
+        assert out.outputs[0].text.strip()
+
+
+@pytest.mark.high_cuda_memory
+def test_mtp_eagle_one_model_rejection():
+    """MTP-Eagle one-model with rejection sampling on: it shares the Eagle3
+    one-model worker/sampler, so the same per-step rejection path runs
+    end-to-end with non-greedy sampling and produces coherent output."""
+    total_mem_gb = torch.cuda.get_device_properties(0).total_memory / 1e9
+    if total_mem_gb < 60:
+        pytest.skip("Not enough memory to load DeepSeek-V3-Lite")
+
+    model_dir = f"{llm_models_root()}/DeepSeek-V3-Lite/bf16"
+    spec_config = MTPDecodingConfig(
+        max_draft_len=1,
+        use_mtp_vanilla=False,
+        mtp_eagle_one_model=True,
+        speculative_model=model_dir,
+        use_rejection_sampling=True,
+    )
+    llm = LLM(
+        model=model_dir,
+        backend="pytorch",
+        disable_overlap_scheduler=True,
+        max_batch_size=2,
+        kv_cache_config=KvCacheConfig(enable_block_reuse=False, max_tokens=8192),
+        max_num_tokens=2048,
+        speculative_config=spec_config,
+        trust_remote_code=True,
+    )
+    prompts = [
+        "The capital of France is",
+        "The president of the United States is",
+    ]
+    # Non-greedy so rejection sampling actually engages (all-greedy bypasses it).
+    sampling_params = SamplingParams(
+        max_tokens=32, temperature=0.8, top_p=0.95, top_k=50, seed=1234
+    )
+    outputs = llm.generate(prompts, sampling_params)
+    llm.shutdown()
+
+    assert len(outputs) == len(prompts)
+    for out in outputs:
+        assert len(out.outputs[0].token_ids) > 0
+        assert out.outputs[0].text.strip()
diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_pard.py b/tests/unittest/_torch/speculative/hw_agnostic/test_pard.py
index 543f269cec2e..6e2c805b54f9 100644
--- a/tests/unittest/_torch/speculative/hw_agnostic/test_pard.py
+++ b/tests/unittest/_torch/speculative/hw_agnostic/test_pard.py
@@ -12,12 +12,6 @@
 sys.path.append(os.path.join(os.path.dirname(__file__), ".."))
 
 
-@pytest.fixture(scope="function")
-def enforce_single_worker(monkeypatch):
-    monkeypatch.setenv("TLLM_WORKER_USE_SINGLE_PROCESS", "1")
-    yield
-
-
 @pytest.mark.parametrize("disable_overlap_scheduler", [True, False])
 def test_pard(disable_overlap_scheduler: bool):
     """Test PARD speculative decoding with CUDA graph support.
@@ -76,5 +70,50 @@ def test_pard(disable_overlap_scheduler: bool):
     llm_spec.shutdown()
 
 
+@pytest.mark.high_cuda_memory
+def test_pard_rejection():
+    """PARD with rejection sampling on: the block-capture rejection path
+    (draft-prob scatter -> fail-closed guard -> rejection acceptance) runs
+    end-to-end with non-greedy sampling and produces coherent output."""
+    total_mem_gb = torch.cuda.get_device_properties(0).total_memory / 1e9
+    if total_mem_gb < 35:
+        pytest.skip("Not enough memory to load target + draft model")
+
+    models_path = llm_models_root()
+    pard_model_dir = f"{models_path}/PARD-Llama-3.2-1B"
+    target_model_dir = f"{models_path}/llama-3.1-model/Llama-3.1-8B-Instruct"
+
+    spec_config = PARDDecodingConfig(
+        max_draft_len=4,
+        speculative_model=pard_model_dir,
+        use_rejection_sampling=True,
+    )
+    llm_spec = LLM(
+        model=target_model_dir,
+        attn_backend="TRTLLM",
+        disable_overlap_scheduler=True,
+        cuda_graph_config=CudaGraphConfig(batch_sizes=[1, 2], enable_padding=True),
+        max_batch_size=2,
+        kv_cache_config=KvCacheConfig(enable_block_reuse=False, max_tokens=2048),
+        max_seq_len=2048,
+        speculative_config=spec_config,
+    )
+    prompts = [
+        "The capital of France is",
+        "The president of the United States is",
+    ]
+    # Non-greedy so rejection sampling actually engages (all-greedy bypasses it).
+    sampling_params = SamplingParams(
+        max_tokens=32, temperature=0.8, top_p=0.95, top_k=50, seed=1234
+    )
+    outputs = llm_spec.generate(prompts, sampling_params)
+    llm_spec.shutdown()
+
+    assert len(outputs) == len(prompts)
+    for out in outputs:
+        assert len(out.outputs[0].token_ids) > 0
+        assert out.outputs[0].text.strip()
+
+
 if __name__ == "__main__":
     unittest.main()
diff --git a/tests/unittest/_torch/speculative/test_rejection_buffers_guard.py b/tests/unittest/_torch/speculative/test_rejection_buffers_guard.py
new file mode 100644
index 000000000000..e2cbe8752d7a
--- /dev/null
+++ b/tests/unittest/_torch/speculative/test_rejection_buffers_guard.py
@@ -0,0 +1,272 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""GPU unit tests for the one-model rejection slot-buffer allocation and the
+fail-closed acceptance guard.
+
+These exercise the pure buffer/guard logic with synthetic tensors and lightweight
+attribute stand-ins (``types.SimpleNamespace``), so they do not construct a full
+``SpecMetadata`` or run any model forward:
+
+- ``SpecMetadata.prepare_rejection_sampling_buffers`` only reads/writes ``self``
+  attributes (``use_rejection_sampling``, ``vocab_size``, ``max_num_requests``,
+  ``max_draft_len``, ``draft_probs``, ``batch_slot_ids``, ``full_draft_probs``).
+- ``SpecWorkerBase._rejection_buffers_valid`` only reads its arguments and
+  ``spec_metadata`` attributes (no ``self`` use), so it can be called unbound.
+
+Requires CUDA because the buffers are allocated on ``device='cuda'``.
+"""
+
+import types
+
+import pytest
+import torch
+
+from tensorrt_llm._torch.speculative.interface import SpecMetadata, SpecWorkerBase
+
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available(), reason="rejection buffers are CUDA tensors"
+)
+
+R, K, V = 8, 4, 32  # max_num_requests, max_draft_len, vocab_size
+
+
+def _alloc_meta(**over):
+    base = dict(
+        use_rejection_sampling=True,
+        draft_probs=None,
+        vocab_size=V,
+        draft_vocab_size=V,
+        max_num_requests=R,
+        max_draft_len=K,
+        num_seq_slots=0,
+        batch_slot_ids=None,
+        full_draft_probs=None,
+        draft_probs_vocab_size=0,
+    )
+    base.update(over)
+    return types.SimpleNamespace(**base)
+
+
+def test_prepare_buffers_allocates_when_enabled():
+    # Shared draft/target vocab: draft_probs + batch_slot_ids are allocated;
+    # full_draft_probs is skipped (only needed when the vocabs differ).
+    m = _alloc_meta()
+    SpecMetadata.prepare_rejection_sampling_buffers(m)
+    assert m.draft_probs is not None and tuple(m.draft_probs.shape) == (R + 1, K, V)
+    assert m.batch_slot_ids is not None and m.batch_slot_ids.shape[0] == R
+    assert m.batch_slot_ids.dtype == torch.long
+    assert m.full_draft_probs is None
+    assert m.draft_probs_vocab_size == V
+
+
+def test_prepare_buffers_allocates_full_draft_probs_on_vocab_mismatch():
+    # Distinct draft vocab: full_draft_probs (d2t-expanded) is allocated.
+    m = _alloc_meta(draft_vocab_size=V - 1)
+    SpecMetadata.prepare_rejection_sampling_buffers(m)
+    assert m.full_draft_probs is not None and tuple(m.full_draft_probs.shape) == (R + 1, K, V)
+
+
+def test_prepare_buffers_span_seq_slot_pool():
+    # Under DeepSeek-V4 overlap scheduling the SeqSlotManager pool
+    # (num_seq_slots) can exceed max_num_requests; py_seq_slot then indexes past
+    # max_num_requests. Slot-indexed buffers must span the full pool plus the
+    # dummy scratch row, and dummy_slot_row must land on that last row so a real
+    # slot never aliases the dummy row.
+    pool = 2 * R
+    m = _alloc_meta(num_seq_slots=pool)
+    SpecMetadata.prepare_rejection_sampling_buffers(m)
+    assert tuple(m.draft_probs.shape) == (pool + 1, K, V)
+    assert m.dummy_slot_row == pool
+
+
+def test_prepare_buffers_noop_when_disabled():
+    m = _alloc_meta(use_rejection_sampling=False)
+    SpecMetadata.prepare_rejection_sampling_buffers(m)
+    assert m.draft_probs is None
+    assert m.batch_slot_ids is None
+    assert m.full_draft_probs is None
+
+
+def _valid_state():
+    return types.SimpleNamespace(
+        draft_probs=torch.empty((R, K, V), device="cuda"),
+        batch_slot_ids=torch.arange(R, device="cuda", dtype=torch.long),
+    )
+
+
+# num_contexts=0, num_gens=4
+_NUM_CTX = 0
+_BATCH = 4
+_NUM_GENS = _BATCH - _NUM_CTX
+
+
+def _good_args():
+    draft_tokens = torch.zeros((_NUM_GENS, K), dtype=torch.int, device="cuda")
+    logits = torch.zeros((_NUM_CTX + _NUM_GENS * (K + 1), V), device="cuda")
+    return draft_tokens, K, V, _NUM_CTX, _BATCH, logits
+
+
+_DUMMY = object()  # _rejection_buffers_valid does not use self
+
+
+def _call(meta, draft_tokens, draft_len, stored_vocab, num_contexts, batch_size, logits):
+    return SpecWorkerBase._rejection_buffers_valid(
+        _DUMMY, draft_tokens, draft_len, stored_vocab, num_contexts, batch_size, logits, meta
+    )
+
+
+def test_guard_true_on_valid_state():
+    assert _call(_valid_state(), *_good_args()) is True
+
+
+# Each builder returns (meta, *_call args) for one malformed state that the
+# fail-closed guard must reject. Kept as a single parametrized test so the guard's
+# rejection branches read as one table.
+def _case_draft_probs_missing():
+    m = _valid_state()
+    m.draft_probs = None
+    return (m, *_good_args())
+
+
+def _case_batch_slot_ids_missing():
+    m = _valid_state()
+    m.batch_slot_ids = None
+    return (m, *_good_args())
+
+
+def _case_stored_vocab_nonpositive():
+    dt, dl, _sv, nc, bs, lg = _good_args()
+    return (_valid_state(), dt, dl, 0, nc, bs, lg)
+
+
+def _case_stored_vocab_exceeds_buffer():
+    dt, dl, _sv, nc, bs, lg = _good_args()
+    return (_valid_state(), dt, dl, V + 1, nc, bs, lg)
+
+
+def _case_draft_len_exceeds_buffer():
+    # draft_probs has K steps; ask for K+1.
+    draft_tokens = torch.zeros((_NUM_GENS, K + 1), dtype=torch.int, device="cuda")
+    logits = torch.zeros((_NUM_CTX + _NUM_GENS * (K + 2), V), device="cuda")
+    return (_valid_state(), draft_tokens, K + 1, V, _NUM_CTX, _BATCH, logits)
+
+
+def _case_draft_tokens_wrong_rows():
+    dt = torch.zeros((_NUM_GENS + 1, K), dtype=torch.int, device="cuda")
+    _, dl, sv, nc, bs, lg = _good_args()
+    return (_valid_state(), dt, dl, sv, nc, bs, lg)
+
+
+def _case_too_few_logits_rows():
+    dt, dl, sv, nc, bs, _lg = _good_args()
+    too_few = torch.zeros((nc + _NUM_GENS, V), device="cuda")  # missing +1 each
+    return (_valid_state(), dt, dl, sv, nc, bs, too_few)
+
+
+def _case_batch_slot_ids_too_short():
+    m = _valid_state()
+    m.batch_slot_ids = torch.arange(_BATCH - 1, device="cuda", dtype=torch.long)
+    return (m, *_good_args())
+
+
+@pytest.mark.parametrize(
+    "build_case",
+    [
+        _case_draft_probs_missing,
+        _case_batch_slot_ids_missing,
+        _case_stored_vocab_nonpositive,
+        _case_stored_vocab_exceeds_buffer,
+        _case_draft_len_exceeds_buffer,
+        _case_draft_tokens_wrong_rows,
+        _case_too_few_logits_rows,
+        _case_batch_slot_ids_too_short,
+    ],
+    ids=lambda f: f.__name__.removeprefix("_case_"),
+)
+def test_guard_false_on_malformed_state(build_case):
+    assert _call(*build_case()) is False
+
+
+# --------------------------------------------------------------------------
+# Fail-closed acceptance dispatch: prove _accept_draft_tokens() routes to
+# strict/base acceptance when the buffers are malformed, and routes to the
+# rejection method when the state is valid. The two acceptance methods are
+# stubbed to record which path ran; the real _can_use_rejection_sampling and
+# _rejection_buffers_valid are exercised.
+# --------------------------------------------------------------------------
+
+
+class _Worker(SpecWorkerBase):
+    @property
+    def max_draft_len(self) -> int:
+        return K
+
+
+def _dispatch_meta(**over):
+    base = dict(
+        use_rejection_sampling=True,
+        is_all_greedy_sample=False,
+        draft_probs_vocab_size=V,
+        draft_probs_last_dim=V,
+        batch_slot_ids=torch.arange(R, device="cuda", dtype=torch.long),
+        draft_probs=torch.zeros((R, K, V), device="cuda"),
+    )
+    base.update(over)
+    return types.SimpleNamespace(**base)
+
+
+def _make_worker():
+    w = _Worker()
+    calls = {"base": 0, "rejection": 0}
+
+    def _base(logits, draft_tokens, num_contexts, batch_size, spec_metadata):
+        calls["base"] += 1
+        return ("base", None)
+
+    def _rej(logits, draft_tokens, draft_probs, num_contexts, batch_size, spec_metadata):
+        calls["rejection"] += 1
+        return ("rejection", None)
+
+    w._sample_and_accept_draft_tokens_base = _base
+    w._sample_and_accept_draft_tokens_rejection = _rej
+    return w, calls
+
+
+def test_accept_dispatch_fails_closed_on_malformed_buffers():
+    w, calls = _make_worker()
+    # Buffers are malformed (draft_probs None) -> _rejection_buffers_valid must
+    # return False, so acceptance falls back to base and the rejection kernel is
+    # skipped.
+    meta = _dispatch_meta(draft_probs=None)
+    draft_tokens = torch.zeros((_NUM_GENS, K), dtype=torch.int, device="cuda")
+    logits = torch.zeros((_NUM_CTX + _NUM_GENS * (K + 1), V), device="cuda")
+    out = w._accept_draft_tokens(logits, draft_tokens, _NUM_CTX, _BATCH, meta)
+    assert out == ("base", None)
+    assert calls["base"] == 1 and calls["rejection"] == 0
+
+
+def test_accept_dispatch_routes_to_rejection_on_valid_state():
+    w, calls = _make_worker()
+    meta = _dispatch_meta()  # valid buffers
+    draft_tokens = torch.zeros((_NUM_GENS, K), dtype=torch.int, device="cuda")
+    logits = torch.zeros((_NUM_CTX + _NUM_GENS * (K + 1), V), device="cuda")
+    out = w._accept_draft_tokens(logits, draft_tokens, _NUM_CTX, _BATCH, meta)
+    assert out == ("rejection", None)
+    assert calls["rejection"] == 1 and calls["base"] == 0
+
+
+def test_accept_dispatch_base_when_all_greedy():
+    w, calls = _make_worker()
+    # All-greedy batch: rejection is bypassed regardless of buffers.
+    meta = _dispatch_meta(is_all_greedy_sample=True)
+    draft_tokens = torch.zeros((_NUM_GENS, K), dtype=torch.int, device="cuda")
+    logits = torch.zeros((_NUM_CTX + _NUM_GENS * (K + 1), V), device="cuda")
+    out = w._accept_draft_tokens(logits, draft_tokens, _NUM_CTX, _BATCH, meta)
+    assert out == ("base", None)
+    assert calls["base"] == 1 and calls["rejection"] == 0
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main([__file__, "-v"]))