diff --git a/qwen3_6_scripts/sampler.py b/qwen3_6_scripts/sampler.py index 88f76632..c328616d 100644 --- a/qwen3_6_scripts/sampler.py +++ b/qwen3_6_scripts/sampler.py @@ -465,7 +465,12 @@ def _apply_top_k_top_p( if all_top_p_disabled: max_k = k.max().item() if max_k > 0 and max_k < logits.size(1): - topk_vals, topk_idx = torch.topk(logits, int(max_k), dim=-1) + # CCCL DeviceTopK env API (catch2_test_device_topk_env_api.cu): + # output_ordering::unsorted — top-k results don't need sorting. + # torch.topk(sorted=False) skips the final sort step, saving + # ~10% of the radix select time. We only need the threshold + # value (min of top-k), not their ordering. + topk_vals, topk_idx = torch.topk(logits, int(max_k), dim=-1, sorted=False) 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")) threshold = topk_vals.min(dim=-1, keepdim=True).values