[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:
@@ -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
1344
qwen3_6_scripts/sampler.py
Normal file
File diff suppressed because it is too large
Load Diff
Reference in New Issue
Block a user