Skip to content

Commit 752ad68

Browse files
[TRTLLM-12669][chore] trim verbose comments in populate_sampling_params
Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
1 parent 0720f38 commit 752ad68

1 file changed

Lines changed: 5 additions & 13 deletions

File tree

tensorrt_llm/_torch/speculative/interface.py

Lines changed: 5 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -581,9 +581,7 @@ def _normalize_request_sampling_params(
581581
is_greedy,
582582
)
583583

584-
# ===== Phase 1: scan sampling configs, compute flags only =====
585-
# Cache normalized params + num_tokens per request so phase 2 doesn't
586-
# have to re-run _normalize_request_sampling_params.
584+
# Phase 1: collect per-request flags and normalized values.
587585
per_request_normalized: list[tuple[float, int, float, int]] = []
588586
temperature_enabled = False
589587
top_k_enabled = False
@@ -625,9 +623,7 @@ def _normalize_request_sampling_params(
625623
self.skip_top_k = not top_k_enabled
626624
self.skip_top_p = not top_p_enabled
627625
self.has_greedy_requests = has_greedy_requests
628-
# All-greedy iff no sampling param is active across the whole batch.
629-
# When True, the sampler takes the argmax fast-path.
630-
# This boolean is part of the CUDA graph key so we capture both variants.
626+
# Used in the CUDA graph key to pick the argmax / advanced variant.
631627
self.is_all_greedy_sample = (self.skip_temperature and self.skip_top_k
632628
and self.skip_top_p)
633629

@@ -637,9 +633,7 @@ def _normalize_request_sampling_params(
637633

638634
if self.temperatures is None or self.temperatures.numel(
639635
) < required_flat_size:
640-
# The buffers are referenced by the captured CUDA graph (advanced
641-
# variant), so they must exist with stable addresses even when
642-
# the current batch is all-greedy and won't write into them.
636+
# Allocate once; the captured graph reads from these stable addresses.
643637
self.temperatures = torch.ones(required_flat_size,
644638
dtype=torch.float32,
645639
device='cuda')
@@ -659,13 +653,11 @@ def _normalize_request_sampling_params(
659653
dtype=torch.float32,
660654
device='cuda')
661655

662-
# ===== Fast-path: all-greedy batch skips list build + H->D copies =====
663-
# The sampler will take the argmax branch and never read these buffers,
664-
# so leaving their values stale is safe.
656+
# All-greedy: argmax branch won't read these buffers, skip H->D copies.
665657
if self.is_all_greedy_sample:
666658
return
667659

668-
# ===== Phase 2: build per-token / per-request lists and copy to GPU =====
660+
# Phase 2: build per-token / per-request lists and copy to GPU.
669661
temperatures: list[float] = []
670662
top_ks: list[int] = []
671663
top_ps: list[float] = []

0 commit comments

Comments
 (0)