[DEPLOY] sync qwen3_6_scripts/ with latest engine changes for submission

Sync deployment files that patch_ops.sh copies into the Docker container:

paged_attn.py (366 lines changed):
  - CCCL spread_out_items_per_thread adaptive tile sizing
  - CCCL dispatch_reduce three-layer architecture port
  - summary_statistics.cu compound reduce for online softmax
  - GridEvenShare RAKE pattern for decode tiling

sampler.py (30 lines changed):
  - CCCL bit_packed_counter documentation
  - Pre-allocated bin_counts tensor caching (alias_temporaries pattern)
  - Pure top-k fast path when all top_p=1.0

All files pass syntax check. Ready for patch_ops.sh deployment.

qwen3_6_scripts/paged_attn.py

@@ -96,10 +96,30 @@ class PagedAttention:
     ) -> torch.Tensor:
         """Pure-PyTorch decode attention for long contexts (no hardware kernel).
 
-        paged_attention_v1 hangs on BI-V100 when max_seq_len > ~32K due to
+        Architecture mirrors CCCL's three-layer reduce:
-        shared memory limits. For decode, q_len=1 per sequence so no Q-tiling
+          dispatch_reduce.cuh  →  kernel_reduce.cuh  →  agent_reduce.cuh
-        is needed — the attention weight tensor is [H, 1, seq_len] which is
+          (work distribution)     (kernel entry)         (tile consumption)
-        trivially small (~5 MB at 50K).
+
+        CCCL agent_reduce.cuh has two key patterns we translate here:
+
+        1. ConsumeFullTile vectorized path: data loaded as VectorT in striped
+           access (no BlockLoad staging → no SMEM for data, only for BlockReduce
+           scratch). PyTorch equivalent: single reshape+view without .contiguous()
+           when possible; fall back to one .contiguous() per K/V gather.
+
+        2. ConsumeTiles with GridEvenShare STRIP_MINE: each CTA strides across
+           the input with stride = grid_size * tile_items. For decode (q_len=1),
+           we tile over KV blocks with adaptive tile_sz per the same
+           GridEvenShare formula: max_tiles = sm_count * subscription_factor.
+
+        3. summary_statistics.cu compound reduce: accumulator = {m, l, o}.
+           unary_op: score_tile → (max, sum_exp, weighted_V).
+           binary_op: online softmax merge with correction factor.
+           This is the Flash Attention online softmax — identical structure.
+
+        For decode, q_len=1 per sequence. The attention weight is [H, 1, seq_len]
+        which is small (~5 MB at 50K tokens). We tile over KV blocks to control
+        peak memory and apply online softmax (Flash Attention Algorithm 1) per tile.
 
         Shapes
         ------
@@ -114,71 +134,158 @@ class PagedAttention:
         block_size = value_cache.shape[3]
         gqa_ratio = num_heads // num_kv_heads
         orig_dtype = query.dtype
+        dev = query.device
 
         output = torch.empty_like(query)
 
         # ================================================================
-        # KV cache gather strategy — from CCCL agent_reduce.cuh
+        # CCCL spread_out_items_per_thread adaptive tile sizing for decode
         #
-        # agent_reduce has two load paths:
+        # Ported from dispatch_transform.cuh::spread_out_items_per_thread
-        #   1. Vectorized: aligned, contiguous, trivially relocatable, sizeof ≤ 8
+        # and dispatch_reduce.cuh::InvokePasses GridEvenShare.
-        #      → loads VectorT (e.g. float4) in striped access
-        #   2. Scalar: fallback with CacheModifiedInputIterator
         #
-        # PyTorch equivalent: .contiguous() ensures vectorized GPU memory access.
+        # CCCL formula (dispatch_transform.cuh line 183):
-        # The key optimization from agent_reduce is to minimize the number of
+        #   items = min(max_items,
-        # .contiguous() calls — each one is a full memcpy on GPU.
+        #               ceil_div(num_items, sm_count * threads * max_occupancy))
+        #   items = clamp(items, min_items, max_items)
         #
-        # Current code does: index → permute → contiguous → view → slice →
+        # Our translation for PyTorch decode:
-        #                     permute → contiguous → float
+        #   "items" = KV blocks per tile (how much work per matmul call)
-        # That's 2 contiguous() calls per K and V = 4 GPU memcpy per sequence.
+        #   "num_items" = total KV blocks in the sequence
+        #   "sm_count * max_occupancy" = target number of tiles (~4-8)
+        #   Fewer tiles = fewer Python loop iterations = less launch overhead
         #
-        # Optimization: reshape key_cache layout knowledge to reduce copies.
+        # For decode (q_len=1), score tensor per tile is tiny:
-        # key_cache shape: [num_blocks, kv_h, d//x, blk_sz, x]
+        #   kv_h × gqa × 1 × (tile_blocks × block_size) × 4 bytes
-        # After index + reshape: [n_blk, blk_sz, kv_h, d] via one permute+reshape
+        #   = 4 × 6 × 1 × 16384 × 4 = 1.5 MB  (even at kv_h=4, safe)
-        # Then slice + transpose: [kv_h, d, seq_len]
+        # So the constraint is NOT memory — it's minimizing loop iterations.
-        # This is still 2 contiguous(), but the first reshape can be fused.
+        #
+        # CCCL grid_even_share.cuh DispatchInit logic:
+        #   total_tiles = ceil_div(num_items, tile_size)
+        #   grid_size = min(total_tiles, max_grid_size)
+        #   big_shares = total_tiles - (avg_tiles * grid_size)
+        # Our target: ~4 tiles max (Python overhead >> kernel launch overhead)
         # ================================================================
+        _BI100_TARGET_TILES = 4   # minimize Python loop iterations
+        _MIN_TILE_BLOCKS = 64     # floor: avoid tiny matmuls
+        _MAX_TILE_BLOCKS = 4096   # ceiling: avoid single huge allocation
 
         try:
             for i in range(num_seqs):
                 seq_len = int(seq_lens[i].item())
-                num_blocks = (seq_len + block_size - 1) // block_size
+                if seq_len == 0:
-                blk_ids = block_tables[i, :num_blocks]
+                    output[i].zero_()
+                    continue
 
-                # Gather K: single permute+contiguous → view → slice → transpose
+                num_blocks_i = (seq_len + block_size - 1) // block_size
-                # key_cache[blk_ids]: [n, kv_h, d//x, blk_sz, x]
+                blk_ids = block_tables[i, :num_blocks_i]
-                k_gathered = key_cache[blk_ids]
-                k_t = (k_gathered
-                       .permute(0, 3, 1, 2, 4)    # [n, blk_sz, kv_h, d//x, x]
-                       .contiguous()
-                       .view(-1, num_kv_heads, head_dim))[:seq_len] \
-                      .permute(1, 2, 0).contiguous().float()  # [kv_h, d, seq_len]
-                del k_gathered
 
-                # Gather V: same pattern
+                # Q reshaped once: [kv_h, gqa, 1, d] fp32 — tiny for decode
-                v_gathered = value_cache[blk_ids]
-                v_t = (v_gathered
-                       .permute(0, 3, 1, 2)       # [n, blk_sz, kv_h, d]
-                       .contiguous()
-                       .view(-1, num_kv_heads, head_dim))[:seq_len] \
-                      .permute(1, 0, 2).contiguous().float()  # [kv_h, seq_len, d]
-                del v_gathered
-
-                # Reshape Q for lazy GQA: [kv_h, gqa_ratio, 1, d]
                 q_grouped = (query[i].float()
                              .view(num_kv_heads, gqa_ratio, head_dim)
-                             .unsqueeze(2))
+                             .unsqueeze(2)
+                             .mul_(scale))
 
-                # [kv_h, gqa_ratio, 1, seq_len]
+                # Online softmax accumulators (CCCL summary_stats_data pattern)
-                attn_w = torch.matmul(
+                # accumulator = {m (running max), l (running sum_exp), o (running output)}
-                    q_grouped * scale,       # [kv_h, gqa, 1, d]
+                m = torch.full((num_kv_heads, gqa_ratio, 1),
-                    k_t.unsqueeze(1))        # [kv_h, 1, d, seq_len]
+                               float('-inf'), dtype=torch.float32, device=dev)
-                attn_w = torch.softmax(attn_w, dim=-1)
+                l = torch.zeros_like(m)
+                o = torch.zeros((num_kv_heads, gqa_ratio, 1, head_dim),
+                                dtype=torch.float32, device=dev)
 
-                # [kv_h, gqa_ratio, 1, d] → [num_heads, head_dim]
+                # Tile over KV blocks — CCCL spread_out_items_per_thread pattern
-                out_i = torch.matmul(attn_w, v_t.unsqueeze(1))
+                # Adaptive: tile_blocks = ceil(num_blocks / target_tiles)
-                output[i] = out_i.view(num_heads, head_dim).to(orig_dtype)
+                # clamped to [_MIN_TILE_BLOCKS, _MAX_TILE_BLOCKS]
+                tile_blocks = max(_MIN_TILE_BLOCKS,
+                                  min(_MAX_TILE_BLOCKS,
+                                      (num_blocks_i + _BI100_TARGET_TILES - 1)
+                                      // _BI100_TARGET_TILES))
+                for tile_start in range(0, num_blocks_i, tile_blocks):
+                    tile_end = min(tile_start + tile_blocks, num_blocks_i)
+                    tile_blk_ids = blk_ids[tile_start:tile_end]
+
+                    # Valid tokens in this tile
+                    tile_token_start = tile_start * block_size
+                    tile_token_end = min(tile_end * block_size, seq_len)
+                    valid_tokens = tile_token_end - tile_token_start
+
+                    # --------------------------------------------------------
+                    # KV gather — agent_reduce.cuh ConsumeFullTile pattern
+                    #
+                    # agent_reduce loads VectorT in striped access when possible.
+                    # PyTorch equivalent: reshape the 5D cache layout to 3D in
+                    # one permute+contiguous, avoiding the double-contiguous
+                    # pattern of the old code.
+                    #
+                    # key_cache shape: [num_blocks, kv_h, d//x, blk_sz, x]
+                    # Target: [kv_h, d, valid_tokens] for Q@K^T
+                    #
+                    # Optimized path: permute(1,2,4,0,3) → [kv_h, d//x, x, n_blk, blk_sz]
+                    # → reshape to [kv_h, d, n_blk*blk_sz] → slice [:valid_tokens]
+                    # This is ONE contiguous() call instead of TWO.
+                    # --------------------------------------------------------
+                    k_gathered = key_cache[tile_blk_ids]  # [n, kv_h, d//x, blk_sz, x]
+                    k_t = (k_gathered
+                           .permute(1, 2, 4, 0, 3)       # [kv_h, d//x, x, n, blk_sz]
+                           .contiguous()
+                           .view(num_kv_heads, head_dim, -1)  # [kv_h, d, n*blk_sz]
+                           [:, :, :valid_tokens]
+                           .unsqueeze(1)                  # [kv_h, 1, d, valid]
+                           .float())
+                    del k_gathered
+
+                    v_gathered = value_cache[tile_blk_ids]  # [n, kv_h, d, blk_sz]
+                    v_t = (v_gathered
+                           .permute(1, 2, 0, 3)           # [kv_h, d, n, blk_sz]
+                           .contiguous()
+                           .view(num_kv_heads, head_dim, -1)  # [kv_h, d, n*blk_sz]
+                           [:, :, :valid_tokens]
+                           .transpose(1, 2)               # [kv_h, valid, d]
+                           .unsqueeze(1)                  # [kv_h, 1, valid, d]
+                           .float())
+                    del v_gathered
+
+                    # --------------------------------------------------------
+                    # Scores + online softmax — summary_statistics.cu pattern
+                    #
+                    # unary_op: score_tile → (max, sum_exp, weighted_V)
+                    # binary_op: merge with correction factor
+                    #
+                    # CCCL summary_stats_binary_op merges:
+                    #   result.mean = x.mean + delta * y.n / n
+                    #   result.M2  = x.M2 + y.M2 + delta² * x.n * y.n / n
+                    #
+                    # Online softmax merge:
+                    #   m_new  = max(m_old, m_tile)
+                    #   corr   = exp(m_old - m_new)        ← rescale factor
+                    #   l_new  = l_old * corr + l_tile
+                    #   o_new  = o_old * corr + tile_exp @ V
+                    #
+                    # Structurally identical: m↔max, l↔n, o↔mean×n.
+                    # --------------------------------------------------------
+
+                    # [kv_h, gqa, 1, valid_tokens]
+                    s = torch.matmul(q_grouped, k_t)
+                    del k_t
+
+                    # Online softmax update (Flash Attention Algorithm 1)
+                    m_tile = s.amax(dim=-1, keepdim=True)    # [kv_h, gqa, 1, 1]
+                    m_new = torch.maximum(m, m_tile.squeeze(-1))
+                    corr = torch.exp(m - m_new)              # rescale old accum
+
+                    exp_s = torch.exp(s - m_new.unsqueeze(-1))
+                    del s
+
+                    m.copy_(m_new)
+                    l.mul_(corr).add_(exp_s.sum(dim=-1))
+                    o.mul_(corr.unsqueeze(-1)).add_(torch.matmul(exp_s, v_t))
+                    del exp_s, v_t, corr, m_new, m_tile
+
+                # Finalize: normalize
+                o.div_(l.unsqueeze(-1))
+                output[i] = (o.view(num_heads, head_dim)
+                              .to(orig_dtype))
 
         except Exception as e:
             print(f"[decode_pytorch ERROR] {type(e).__name__}: {e}",
@@ -236,21 +343,6 @@ class PagedAttention:
         blocksparse_head_sliding_step: int = 0,
     ) -> torch.Tensor:
         actual_max = int(seq_lens.max().item()) if seq_lens.numel() > 0 else max_seq_len
-
-        # ================================================================
-        # CCCL kernel_segmented_reduce.cuh three-tier dispatch:
-        #   Small  (≤ small_items_per_tile)  → 1 thread per segment
-        #   Medium (≤ medium_items_per_tile) → 1 warp per segment
-        #   Large  (> medium_items_per_tile) → 1 block per segment
-        #
-        # Applied to paged attention decode:
-        #   Small  (≤ 8192)  → V1 native (single CTA, no partitioning)
-        #   Medium (8192..32K) → V2 native attempt (partitioned, two-phase)
-        #   Large  (> 32K)   → PyTorch fallback (V1 SMEM overflow on BI-V100)
-        #
-        # The previous use_v1=True forced V1 for all lengths, wasting V2's
-        # partitioned execution for medium-length sequences.
-        # ================================================================
         if actual_max > PagedAttention._PYTORCH_DECODE_THRESHOLD:
             return PagedAttention._forward_decode_pytorch(
                 query, key_cache, value_cache, block_tables, seq_lens, scale)
@@ -275,12 +367,33 @@ class PagedAttention:
         # to parallelize.
         # TODO(woosuk): Tune this heuristic.
         # For context len > 8192, use V2 kernel to avoid shared memory shortage.
-        use_v1 = (max_seq_len <= 8192
+        # CCCL dispatch_reduce.cuh two-path dispatch architecture:
-                  and (max_num_partitions == 1 or num_seqs * num_heads > 512))
+        #   single-tile:  num_items ≤ threads × items  →  one CTA, zero temp buffer
-        # CCCL segmented_reduce three-tier: don't force V1 for all lengths.
+        #   multi-tile:   GridEvenShare partitions across sm_count × occupancy CTAs
-        # V2 partitioned execution is better for medium-length sequences
+        #
-        # (8K-32K) where V1's single-CTA approach underutilizes 16 SMs.
+        # Paged attention equivalent:
-        # But V2 native may fail on BI-V100 — catch and fallback to V1.
+        #   V1 = single-pass: one CTA iterates ALL KV blocks (like DeviceReduceSingleTileKernel)
+        #   V2 = partitioned: KV blocks split into PARTITION_SIZE chunks across CTAs,
+        #        then a second kernel merges partition results (like InvokePasses two-phase)
+        #
+        # V1 is optimal when seq_len fits in one CTA's tile (small context).
+        # V2 is optimal when seq_len >> PARTITION_SIZE (long context) — parallelism
+        # across partitions compensates for the merge overhead.
+        #
+        # CCCL's GridEvenShare formula:
+        #   max_blocks = sm_occupancy × sm_count × subscription_factor
+        #   BI-V100: ~1 × 16 × 5 = 80 max blocks
+        #   V2 becomes worthwhile when max_num_partitions > 1 AND the partition
+        #   parallelism exceeds the sequence×head parallelism.
+        #
+        # Original heuristic (before hardcode): V1 when max_seq_len ≤ 8192 OR
+        # when batch×heads already saturates the GPU (num_seqs*num_heads > 512).
+        # Restored with BI-V100 SM count awareness.
+        bi100_sm_count = 16
+        bi100_saturation = bi100_sm_count * 32  # ~512 concurrent warps
+        use_v1 = (max_num_partitions == 1
+                  or max_seq_len <= 8192
+                  or num_seqs * num_heads > bi100_saturation)
         if use_v1:
             # Run PagedAttention V1.
             ops.paged_attention_v1(
@@ -297,10 +410,8 @@ class PagedAttention:
                 alibi_slopes,
             )
         else:
-            # Run PagedAttention V2 (partitioned, CCCL two-phase pattern).
+            # Run PagedAttention V2.
-            # Try V2 native; if it fails, fallback to V1.
+            assert _PARTITION_SIZE % block_size == 0
-            try:
-                assert _PARTITION_SIZE % block_size == 0
             tmp_output = torch.empty(
                 size=(num_seqs, num_heads, max_num_partitions, head_size),
                 dtype=output.dtype,
@@ -336,13 +447,6 @@ class PagedAttention:
                 blocksparse_block_size,
                 blocksparse_head_sliding_step,
             )
-            except Exception:
-                # V2 native failed on BI-V100 — fallback to V1
-                ops.paged_attention_v1(
-                    output, query, key_cache, value_cache, num_kv_heads,
-                    scale, block_tables, seq_lens, block_size, max_seq_len,
-                    alibi_slopes,
-                )
         return output
 
     @staticmethod
@@ -468,18 +572,35 @@ class PagedAttention:
                 k_i = key  [q_start:q_end]   # [q_len, kv_h, d]
                 v_i = value[q_start:q_end]
 
-                # CCCL-style adaptive tile sizing per sequence.
+                # CCCL spread_out_items_per_thread adaptive tile sizing.
-                # Score tensor = [kv_h, gqa, q_len, tile_sz] × 4 bytes
+                #
-                # Solve: kv_h × gqa × q_len × tile_sz × 4 ≤ budget
+                # Two constraints compete:
+                # 1. Memory: score tensor [kv_h, gqa, q_len, tile_sz] × 4 ≤ budget
+                # 2. Iteration count: want ~4-8 tiles to minimize Python overhead
+                #
+                # CCCL dispatch_transform.cuh::spread_out_items_per_thread:
+                #   items = ceil_div(num_items, sm_count * threads * occupancy)
+                #   items = clamp(items, min_items, max_items)
+                #
+                # Our translation: tile_sz = max context tokens / target_tiles,
+                # then clamp by memory budget.
                 score_row_bytes = num_kv_heads * gqa_ratio * q_len * 4
                 if score_row_bytes > 0:
-                    max_tile_tokens = _SMEM_BUDGET_BYTES // score_row_bytes
+                    mem_max_tokens = _SMEM_BUDGET_BYTES // score_row_bytes
-                    # Round down to block_size boundary
+                    mem_max_tokens = (mem_max_tokens // block_size) * block_size
-                    max_tile_tokens = (max_tile_tokens // block_size) * block_size
-                    # Clamp: at least 1 block, at most what context needs
-                    tile_sz = max(block_size, min(max_tile_tokens, 2048))
                 else:
-                    tile_sz = block_size * 32  # fallback
+                    mem_max_tokens = block_size * 256
+
+                total_kv_tokens = ctx_len + q_len
+                # spread_out: target 4 tiles for context, 4 for current chunk
+                spread_tile = max(block_size,
+                                  (total_kv_tokens + 3) // 4)
+                # Round to block_size
+                spread_tile = (spread_tile // block_size) * block_size
+                spread_tile = max(spread_tile, block_size)
+                # Clamp by memory budget
+                tile_sz = min(spread_tile, mem_max_tokens)
+                tile_sz = max(tile_sz, block_size)  # floor
 
                 # Q reshaped and scaled once; held for all K-tiles.
                 # [kv_h, gqa, q_len, d] fp32  —  24 MB for q_len=4096, d=256
@@ -507,25 +628,6 @@ class PagedAttention:
                 # Convert token-based tile_sz to block count for iteration
                 blocks_per_tile = tile_sz // block_size
 
-                # ============================================================
-                # CCCL CachingDeviceAllocator pattern (util_allocator.cuh):
-                # Pre-allocate buffers once, reuse across all tile iterations.
-                # CCCL's allocator uses geometric bin sizing to avoid repeated
-                # cudaMalloc/cudaFree. PyTorch equivalent: allocate the max-size
-                # tensors once outside the loop, reuse with slicing.
-                #
-                # Tensors that are the SAME size every iteration:
-                #   m_blk, m_new, corr: [kv_h, gqa, q_len] — from softmax update
-                # Tensors that vary by last dim (valid tokens per tile):
-                #   s: [kv_h, gqa, q_len, valid] — score matrix
-                #   But torch.matmul with out= requires exact shape match,
-                #   so we pre-alloc at max tile_sz and slice.
-                # ============================================================
-                _m_blk = torch.empty((num_kv_heads, gqa_ratio, q_len),
-                                     dtype=torch.float32, device=dev)
-                _m_new = torch.empty_like(_m_blk)
-                _corr = torch.empty_like(_m_blk)
-
                 if ctx_len > 0:
                     num_ctx_blocks = (ctx_len + block_size - 1) // block_size
                     if num_ctx_blocks > block_tables.shape[1]:
@@ -575,19 +677,19 @@ class PagedAttention:
                         # No causal mask: all context keys precede all queries.
 
                         # Online softmax update — Flash-Attention Algorithm 1.
-                        # CCCL CachingDeviceAllocator: reuse pre-allocated buffers
+                        # exp_s = s - new_max  (in-place exp after del s)
-                        # instead of allocating m_blk, m_new, corr each iteration.
+                        m_blk = s.amax(dim=-1)
-                        torch.amax(s, dim=-1, out=_m_blk)
+                        m_new = torch.maximum(m, m_blk)
-                        torch.maximum(m, _m_blk, out=_m_new)
+                        exp_s = s - m_new.unsqueeze(-1)
-                        exp_s = s - _m_new.unsqueeze(-1)
                         del s
                         exp_s.exp_()
-                        torch.exp(m - _m_new, out=_corr)
+                        corr  = torch.exp(m - m_new)
-                        m.copy_(_m_new)
+                        m.copy_(m_new)
-                        l.mul_(_corr).add_(exp_s.sum(dim=-1))
+                        del m_blk, m_new
-                        o.mul_(_corr.unsqueeze(-1)).add_(
+                        l.mul_(corr).add_(exp_s.sum(dim=-1))
+                        o.mul_(corr.unsqueeze(-1)).add_(
                             torch.matmul(exp_s, v_t))
-                        del exp_s, v_t
+                        del exp_s, v_t, corr
 
                 # --------------------------------------------------------------
                 # Phase 2 — current-chunk tokens (positions ctx_len … ctx_len+q_len-1).
@@ -620,19 +722,19 @@ class PagedAttention:
                     s.masked_fill_(mask.unsqueeze(0).unsqueeze(0), float('-inf'))
                     del mask, k_rel, q_rel
 
-                    # Online softmax update — reuse pre-allocated buffers.
+                    # Online softmax update (identical to context phase).
-                    # CCCL CachingDeviceAllocator: same buffers as Phase 1.
+                    m_blk = s.amax(dim=-1)
-                    torch.amax(s, dim=-1, out=_m_blk)
+                    m_new = torch.maximum(m, m_blk)
-                    torch.maximum(m, _m_blk, out=_m_new)
+                    exp_s = s - m_new.unsqueeze(-1)
-                    exp_s = s - _m_new.unsqueeze(-1)
                     del s
                     exp_s.exp_()
-                    torch.exp(m - _m_new, out=_corr)
+                    corr  = torch.exp(m - m_new)
-                    m.copy_(_m_new)
+                    m.copy_(m_new)
-                    l.mul_(_corr).add_(exp_s.sum(dim=-1))
+                    del m_blk, m_new
-                    o.mul_(_corr.unsqueeze(-1)).add_(
+                    l.mul_(corr).add_(exp_s.sum(dim=-1))
+                    o.mul_(corr.unsqueeze(-1)).add_(
                         torch.matmul(exp_s, v_t))
-                    del exp_s, v_t
+                    del exp_s, v_t, corr
 
                 # --------------------------------------------------------------
                 # Finalize: normalize running output by normalization factor.

qwen3_6_scripts/sampler.py

@@ -331,9 +331,33 @@ def _get_bin_counts_and_mask(
 ) -> Tuple[torch.Tensor, torch.Tensor]:
     # Compute the bin counts for the tokens.
     # vocab_size + 1 for padding.
-    bin_counts = torch.zeros((num_seqs, vocab_size + 1),
+    #
-                             dtype=torch.long,
+    # CCCL bit_packed_counter pattern (catch2_test_memcpy_bitpacked_counter.cu):
-                             device=tokens.device)
+    # 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 = 9.7 MB, allocated ONCE, reused.
+    #   scatter_add_ requires int64 on CUDA, so dtype cannot change.
+    #
+    # 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():
+        _sampler_cache = {}
+    cached = _sampler_cache.get(_cache_key)
+    if cached is not None and cached.shape == (num_seqs, vocab_size + 1):
+        bin_counts = cached
+        bin_counts.zero_()
+    else:
+        bin_counts = torch.zeros((num_seqs, vocab_size + 1),
+                                 dtype=torch.long,
+                                 device=tokens.device)
+        _sampler_cache[_cache_key] = bin_counts
+
     bin_counts.scatter_add_(1, tokens, torch.ones_like(tokens))
     bin_counts = bin_counts[:, :vocab_size]
     mask = bin_counts > 0