[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.
This commit is contained in:
dylanyunlon
2026-08-05 08:30:21 +00:00
parent 327f9fbf40
commit 44bdf49cae
2 changed files with 1352 additions and 0 deletions

View File

@@ -58,6 +58,14 @@ echo "[patch_ops] arg_utils.py → engine/"
cp ./logits_processor.py $V/model_executor/layers/logits_processor.py
echo "[patch_ops] logits_processor.py → model_executor/layers/"
# --- sampler.py: CCCL-ported top-k fast path for sampling --------------------
# When all sequences use top_p=1.0, skip full sort+cumsum and use torch.topk.
# CCCL partition/flagged.cu insight: radix select is O(N×bits_per_pass) vs
# full sort O(N log N). For vocab=152064: topk ~11 passes vs sort ~17 passes.
# FULL FILE REPLACEMENT.
cp ./sampler.py $V/model_executor/layers/sampler.py
echo "[patch_ops] sampler.py → model_executor/layers/"
# --- transformers: Qwen3_5 tokenizer / model files --------------------------
pip install transformers==4.55.3 -i https://pypi.tuna.tsinghua.edu.cn/simple
cp -r ./qwen3_5 /usr/local/lib/python3.10/site-packages/transformers/models/

1344
qwen3_6_scripts/sampler.py Normal file

File diff suppressed because it is too large Load Diff