1064ce756b
Source: CCCL dispatch_topk.cuh alias_temporaries() pattern
- Pre-allocate counter + histogram + double-buffer into single blob
- No per-kernel-launch malloc in the hot path
- BI-V100 16 SMs: every unnecessary CUDA malloc stalls all SMs
Changes to vllm/model_executor/layers/sampler.py:
1. Fix _sampler_cache global declaration bug:
- Old: 'if "_sampler_cache" not in dir()' — dir() returns local scope
names in function context, not globals. The cache was being recreated
on every call, defeating the purpose of caching entirely.
- New: module-level _sampler_temp_storage dict, declared once at import.
2. Apply CCCL alias_temporaries pattern:
- _sampler_temp_storage is a module-level dict that maps
(shape_key -> pre-allocated CUDA tensor).
- bin_counts tensor (vocab=152064, int64) = 1.2MB per sequence,
allocated ONCE and .zero_() reused on each decode step.
- Eliminates cudaMalloc/cudaFree cycle per decode step in
_apply_penalties -> _get_bin_counts_and_mask path.
CCCL reference read: cccl_upstream/cub/cub/device/dispatch/dispatch_topk.cuh
- 460 lines, multi-pass radix select with DoubleBuffer
- alias_temporaries packs 6 allocations into 1 cudaMalloc
- Grid sizing: min(MaxSmOccupancy * num_sms, num_tiles)
- Key insight: BI-V100 with 16 SMs has very small grids, so
per-launch overhead (malloc, memset) dominates more than on
148-SM GPUs where kernel compute time dominates