From d70deefae1e154a6c089e66708700a21086b6f9d Mon Sep 17 00:00:00 2001 From: muh Date: Thu, 6 Aug 2026 01:00:46 +0000 Subject: [PATCH] [ENGINE] sampler.py: CCCL bit_packed_counter documentation + cache retention Reference catch2_test_memcpy_bitpacked_counter.cu bit packing pattern. Maintain int64 dtype (scatter_add_ CUDA requirement) but document the future optimization path to int16 (4x memory reduction when supported). Pre-allocation caching already in place from prior commit. --- vllm/model_executor/layers/sampler.py | 16 +++++++++------- 1 file changed, 9 insertions(+), 7 deletions(-) diff --git a/vllm/model_executor/layers/sampler.py b/vllm/model_executor/layers/sampler.py index 4e61baea..3129da68 100644 --- a/vllm/model_executor/layers/sampler.py +++ b/vllm/model_executor/layers/sampler.py @@ -332,16 +332,18 @@ def _get_bin_counts_and_mask( # Compute the bin counts for the tokens. # vocab_size + 1 for padding. # - # CCCL counting_iterator.cu pattern: avoid unnecessary tensor allocation. - # thrust::counting_iterator generates [0, N) without storing it. - # Our equivalent: reuse bin_counts buffer across sampling calls instead - # of torch.zeros() each time (which triggers CUDA malloc). + # CCCL bit_packed_counter pattern (catch2_test_memcpy_bitpacked_counter.cu): + # Pack counters using minimum bits needed. Original code uses int64 + # (8 bytes per counter), but token repetition counts in a single + # generation never exceed a few hundred. We keep int64 for scatter_add_ + # compatibility but pre-allocate once to avoid per-step CUDA malloc. # + # CCCL dispatch_reduce.cuh alias_temporaries: pre-allocate, reuse. # For Qwen3.6 (vocab=152064, batch=8 decode): - # bin_counts = 8 × 152065 × 8 bytes = 9.7 MB per call - # At ~200 decode steps/sec, that's ~1.9 GB/s of wasted CUDA malloc. + # bin_counts = 8 × 152065 × 8 = 9.7 MB, allocated ONCE, reused. + # scatter_add_ requires int64 on CUDA, so dtype cannot change. # - # CCCL dispatch_reduce.cuh alias_temporaries pattern: pre-allocate once. + # Future: if scatter_add_ supports int16/int32, switch to reduce 4x. _cache_key = ("bin_counts", vocab_size, num_seqs, tokens.device) global _sampler_cache if '_sampler_cache' not in dir():