[TRTLLM-12669][perf] bypass rejection sampling when batch is all-greedy

For all-greedy batches, rejection sampling acceptance is equivalent to
argmax. Skip it (and the associated draft-prob computation) to save the
rejection kernel and a softmax pass. Affects:
- _can_use_rejection_sampling in interface.py and eagle3_dynamic_tree.py
- _compute_and_store_draft_probs call in eagle3.py

Also revert the use_rejection_sampling exception in the populate fast-path
since the rejection path is no longer reachable when all-greedy.

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

Author: zhaoyangwang-nvidia

tensorrt_llm/_torch/speculative/eagle3.py

@@ -758,7 +758,10 @@ def _forward_linear_draft_loop(self, inputs, attn_metadata, spec_metadata,
                 gen_draft_tokens)
             next_draft_tokens[num_contexts:] = gen_draft_tokens
 
-        if spec_metadata.use_rejection_sampling and draft_logits_list:
+        # Skip when the whole batch is greedy: _can_use_rejection_sampling will
+        # bypass the rejection path anyway, so computing draft probs is wasted.
+        if (spec_metadata.use_rejection_sampling and draft_logits_list
+                and not spec_metadata.is_all_greedy_sample):
             d2t_param = getattr(draft_model.model, "d2t", None)
             spec_metadata.d2t = d2t_param.data if d2t_param is not None else None
             self._compute_and_store_draft_probs(draft_logits_list,

tensorrt_llm/_torch/speculative/eagle3_dynamic_tree.py

@@ -950,7 +950,13 @@ def _can_use_rejection_sampling(self, spec_metadata) -> bool:
         Returns:
             True if rejection sampling is enabled and the draft logit buffer is allocated
         """
-        return spec_metadata.use_rejection_sampling and self._draft_depth_logits_cat is not None
+        # Skip rejection sampling when the whole batch is greedy: argmax is
+        # equivalent and avoids the rejection kernel cost.
+        return (
+            spec_metadata.use_rejection_sampling
+            and self._draft_depth_logits_cat is not None
+            and not spec_metadata.is_all_greedy_sample
+        )
 
     def _finalize_dynamic_tree_verify_outputs(
         self,

tensorrt_llm/_torch/speculative/interface.py

@@ -653,9 +653,10 @@ def _normalize_request_sampling_params(
                                              dtype=torch.float32,
                                              device='cuda')
 
-        # All-greedy: argmax branch won't read these buffers, skip H->D copies.
-        # Rejection sampling always reads them, so it cannot use this fast-path.
-        if self.is_all_greedy_sample and not self.use_rejection_sampling:
+        # All-greedy: sampler takes the argmax branch (and rejection sampling
+        # is also bypassed for all-greedy), so the buffers are never read.
+        # Skip the H->D copies.
+        if self.is_all_greedy_sample:
             return
 
         # Phase 2: build per-token / per-request lists and copy to GPU.
@@ -1030,8 +1031,11 @@ def _accept_draft_tokens(self, logits, draft_tokens, num_contexts,
 
     def _can_use_rejection_sampling(self, spec_metadata: SpecMetadata,
                                     num_contexts: int) -> bool:
+        # Skip rejection sampling when the whole batch is greedy: the
+        # accepted result is identical to argmax and the base path is cheaper.
         return (spec_metadata.use_rejection_sampling
-                and spec_metadata.draft_probs_valid and num_contexts == 0)
+                and spec_metadata.draft_probs_valid and num_contexts == 0
+                and not spec_metadata.is_all_greedy_sample)
 
     def _sample_and_accept_draft_tokens_rejection(
         self,