[TRTLLM-12669][feat] Eagle3 one-model draft sampling honors target sampling params

The Eagle3 one-model draft loop was hardcoded to greedy argmax, even
when the target sampler used non-greedy (temperature/top_k/top_p)
parameters. This made rejection sampling math degenerate: with draft
forced to argmax, p_draft is a one-hot on the argmax token, so
p_target / p_draft is zero everywhere else and acceptance is biased.

Route the linear draft loop through a new `_draft_sampler_advanced`
on `SpecWorkerBase` that reuses the per-request sampling tensors
already populated by `populate_sampling_params_for_one_model`
(request_temperatures / request_top_ks / request_top_ps) and the
shared seed/offset used by the target sampler. When the batch is
all-greedy (`is_all_greedy_sample=True`) it short-circuits to
`_draft_sampler_greedy`, matching the target sampler's argmax
fast-path so existing CUDA-graph variants are unaffected.

Other spec modes (Draft-Target, MTP, Pard) are untouched and keep
their current argmax draft sampling.

Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>

Author: zhaoyangwang-nvidia

tensorrt_llm/_torch/speculative/eagle3.py

@@ -717,7 +717,8 @@ def _forward_linear_draft_loop(self, inputs, attn_metadata, spec_metadata,
                 if spec_metadata.use_rejection_sampling:
                     draft_logits_list.append(logits.clone())
 
-                new_draft_token = self.draft_decoder(logits, draft_model)
+                new_draft_token = self.draft_decoder(logits, draft_model,
+                                                     spec_metadata, batch_size)
                 next_draft_tokens.append(new_draft_token)
                 # update inputs
                 hidden_states = hidden_states_to_save[gather_ids]
@@ -794,27 +795,27 @@ def draft_decoder(
         self,
         logits: torch.Tensor,
         draft_model: nn.Module,
+        spec_metadata: Optional[Eagle3OneModelSpecMetadata] = None,
+        batch_size: Optional[int] = None,
     ):
         '''
-        Sampling draft tokens with support for non-greedy sampling.
+        Sample draft tokens. When spec_metadata + batch_size are provided, use
+        the target's per-request sampling params (temperature/top_k/top_p);
+        otherwise fall back to argmax.
 
         Args:
-            logits: torch.Tensor
-                [num_tokens, vocab_size]
-                Logits produced by the draft model.
-            draft_model: nn.Module
-                The draft model.
-
-        Returns:
-            draft_tokens: torch.Tensor
-                [batch_size * max_draft_len]
-                Draft token ids. Flattened.
+            logits: [batch_size, vocab_size] - Draft model logits.
+            draft_model: The draft model.
+            spec_metadata: Carries per-request sampling param tensors. When
+                None, sampling is forced greedy.
+            batch_size: Active requests, used to slice per-request tensors.
         '''
 
         d2t = getattr(draft_model.model, "d2t", None)
-        draft_tokens = self._draft_sampler_greedy(logits, d2t)
-
-        return draft_tokens
+        if spec_metadata is not None and batch_size is not None:
+            return self._draft_sampler_advanced(logits, spec_metadata,
+                                                batch_size, d2t)
+        return self._draft_sampler_greedy(logits, d2t)
 
     def prepare_1st_drafter_inputs(
         self,

tensorrt_llm/_torch/speculative/interface.py

@@ -1141,6 +1141,62 @@ def _draft_sampler_greedy(self, logits: torch.Tensor, d2t=None):
 
         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,
+    ):
+        """
+        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)
+        """
+        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.use_flashinfer:
+            top_ks = top_ks.clamp(min=1, max=logits.shape[-1] - 1)
+            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,
+            use_flashinfer=self.use_flashinfer,
+            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 _compute_and_store_draft_probs(
         self,
         draft_logits_list: List[torch.Tensor],