Skip to content

Commit 8f2bd8f

Browse files
[TRTLLM-12669][fix] disable fast-path when use_rejection_sampling=True
The rejection sampling path always reads spec_metadata.temperatures, regardless of greedy/non-greedy. Skipping phase 2 there would leave stale buffer values and break greedy semantics in the rejection acceptance step. Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
1 parent 752ad68 commit 8f2bd8f

1 file changed

Lines changed: 2 additions & 1 deletion

File tree

tensorrt_llm/_torch/speculative/interface.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -654,7 +654,8 @@ def _normalize_request_sampling_params(
654654
device='cuda')
655655

656656
# All-greedy: argmax branch won't read these buffers, skip H->D copies.
657-
if self.is_all_greedy_sample:
657+
# Rejection sampling always reads them, so it cannot use this fast-path.
658+
if self.is_all_greedy_sample and not self.use_rejection_sampling:
658659
return
659660

660661
# Phase 2: build per-token / per-request lists and copy to GPU.

0 commit comments

Comments
 (0)