@@ -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