diff --git a/vllm/model_executor/layers/sampler.py b/vllm/model_executor/layers/sampler.py index 42a6a0e6..5b5f7db6 100644 --- a/vllm/model_executor/layers/sampler.py +++ b/vllm/model_executor/layers/sampler.py @@ -416,6 +416,33 @@ def _apply_top_k_top_p( p: torch.Tensor, k: torch.Tensor, ) -> torch.Tensor: + # CCCL insight from tuning_topk.cuh: radix select (used by torch.topk) + # is O(N × bits_per_pass) vs full sort O(N log N). For vocab=152064: + # topk ≈ 11 radix passes, sort ≈ 17 passes. 1.5x fewer kernel cycles. + # + # Fast path: when ALL sequences use top_p=1.0 (no nucleus sampling), + # we only need top-k selection, not full sort + cumsum. + # This skips: sort (152K elements) + softmax + cumsum + scatter + # and replaces with: topk (much cheaper) + scatter. + all_top_p_disabled = (p >= 1.0 - 1e-6).all() + if all_top_p_disabled: + # Pure top-k path: use torch.topk instead of full sort + # For k values, take the minimum k across all sequences + max_k = k.max().item() + if max_k > 0 and max_k < logits.size(1): + # Get top-k values and indices + topk_vals, topk_idx = torch.topk(logits, int(max_k), dim=-1) + # Mask out everything below top-k threshold per sequence + # topk_vals[:, -1] is the k-th largest value for each seq + actual_k_mask = torch.arange(int(max_k), device=k.device).unsqueeze(0) < k.unsqueeze(1) + topk_vals.masked_fill_(~actual_k_mask, -float("inf")) + # Get per-sequence threshold (smallest value kept) + threshold = topk_vals.min(dim=-1, keepdim=True).values + # Apply threshold to original logits + logits = logits.masked_fill(logits < threshold, -float("inf")) + return logits + + # Full path: sort + top-k + top-p (cumsum) logits_sort, logits_idx = logits.sort(dim=-1, descending=False) # Apply top-k.