Random CCCL pick: cub/test/catch2_test_device_topk_env_api.cu (290 lines, full)
CCCL DeviceTopK uses cuda::execution::output_ordering::unsorted —
top-k results are NOT sorted by default. The test sorts results
AFTER retrieval only for verification, not during the algorithm.
Our sampler's torch.topk(logits, k) defaults to sorted=True, which
adds an unnecessary final sort step after the radix selection.
For sampling, we only need the THRESHOLD value (min of top-k set)
to mask logits below it — the ordering within top-k is irrelevant.
Change: torch.topk(..., sorted=False) in the top-k fast path.
This skips the O(k log k) sort of the selected elements.
For Qwen3.6 with top_k=20, k=20 sort is cheap, but it's free
to eliminate and matches CCCL's unsorted-by-default design.
CCCL also teaches: determinism::not_guaranteed is acceptable for
top-k in sampling contexts (temperature > 0 = inherent randomness).
Base file modified: qwen3_6_scripts/sampler.py (deployed via patch_ops.sh)
Random CCCL pick: cub/cub/device/dispatch/dispatch_topk.cuh (480 lines, full read)
CCCL's DeviceTopK uses DoubleBuffer<key_in_t> to ping-pong between two
pre-allocated buffers across radix passes, achieving zero allocation in
the hot loop. Our sampler.py's _apply_top_k_top_p was allocating 2 new
tensors (logits_sort + logits_idx, each vocab_size×4B = 600KB) on every
single decode step via torch.sort().
Change: cache sort output tensors keyed on (batch, vocab, device) and
reuse them via torch.sort(..., out=(cached_sort, cached_idx)). This
eliminates 1.2MB of GPU allocation per decode step.
For competition max_num_seqs=1, vocab=152064:
Before: 2 × 152064 × 4B = 1.2MB allocated per step
After: 0 bytes allocated per step (reuse cached buffers)
At 395 tokens/sec target: saves 474MB/sec of allocator pressure.
BI-V100 has no async CUDA allocator, so this is synchronous overhead.
CCCL architecture insight used:
dispatch_topk.cuh line ~430: DoubleBuffer<key_in_t> key_bufs(alloc[3], alloc[2])
for pass: key_bufs.Current() → read, key_bufs.Alternate() → write, swap
Base file modified: qwen3_6_scripts/sampler.py (deployed via patch_ops.sh)
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.
