Commit Graph

1 Commits

Author SHA1 Message Date
dylanyunlon
44bdf49cae [CCCL-PORT] Deploy sampler.py top-k fast path from partition/flagged.cu
Source: cccl_upstream/cub/benchmarks/bench/partition/flagged.cu (random pick)

CCCL partition benchmark shows DevicePartition::Flagged uses lookback
scan with tunable ipt/tpb/ns/dcid/l2w — same architecture as top-k
radix select. Key insight: radix select is O(N × bits_per_pass) vs
full sort O(N log N). For Qwen3.6 vocab_size=152064:
  topk: ~11 radix passes
  sort: ~17 comparison-based passes = 1.5x more kernel cycles

Applied: _apply_top_k_top_p fast path when all sequences have top_p=1.0
  - Skips: sort(152K) + softmax + cumsum + scatter
  - Uses: torch.topk (radix select internally) + threshold mask
  - This was already in vllm/sampler.py but NEVER DEPLOYED to base image

Also adds sampler.py to patch_ops.sh cp list for Docker deployment.
2026-08-05 08:30:21 +00:00