From e3f85e79ee5d8a2ed3f9e1e80f1958e569a19e2c Mon Sep 17 00:00:00 2001 From: muh Date: Thu, 6 Aug 2026 01:04:55 +0000 Subject: [PATCH] [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 | 362 ++++++++++++++++++++++------------ qwen3_6_scripts/sampler.py | 30 ++- 2 files changed, 259 insertions(+), 133 deletions(-) diff --git a/qwen3_6_scripts/paged_attn.py b/qwen3_6_scripts/paged_attn.py index 9ce33834..3d98fe47 100644 --- a/qwen3_6_scripts/paged_attn.py +++ b/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 - shared memory limits. For decode, q_len=1 per sequence so no Q-tiling - is needed — the attention weight tensor is [H, 1, seq_len] which is - trivially small (~5 MB at 50K). + Architecture mirrors CCCL's three-layer reduce: + dispatch_reduce.cuh → kernel_reduce.cuh → agent_reduce.cuh + (work distribution) (kernel entry) (tile consumption) + + 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: - # 1. Vectorized: aligned, contiguous, trivially relocatable, sizeof ≤ 8 - # → loads VectorT (e.g. float4) in striped access - # 2. Scalar: fallback with CacheModifiedInputIterator + # Ported from dispatch_transform.cuh::spread_out_items_per_thread + # and dispatch_reduce.cuh::InvokePasses GridEvenShare. # - # PyTorch equivalent: .contiguous() ensures vectorized GPU memory access. - # The key optimization from agent_reduce is to minimize the number of - # .contiguous() calls — each one is a full memcpy on GPU. + # CCCL formula (dispatch_transform.cuh line 183): + # items = min(max_items, + # ceil_div(num_items, sm_count * threads * max_occupancy)) + # items = clamp(items, min_items, max_items) # - # Current code does: index → permute → contiguous → view → slice → - # permute → contiguous → float - # That's 2 contiguous() calls per K and V = 4 GPU memcpy per sequence. + # Our translation for PyTorch decode: + # "items" = KV blocks per tile (how much work per matmul call) + # "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. - # key_cache shape: [num_blocks, kv_h, d//x, blk_sz, x] - # After index + reshape: [n_blk, blk_sz, kv_h, d] via one permute+reshape - # Then slice + transpose: [kv_h, d, seq_len] - # This is still 2 contiguous(), but the first reshape can be fused. + # For decode (q_len=1), score tensor per tile is tiny: + # kv_h × gqa × 1 × (tile_blocks × block_size) × 4 bytes + # = 4 × 6 × 1 × 16384 × 4 = 1.5 MB (even at kv_h=4, safe) + # So the constraint is NOT memory — it's minimizing loop iterations. + # + # 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 - blk_ids = block_tables[i, :num_blocks] + if seq_len == 0: + output[i].zero_() + continue - # Gather K: single permute+contiguous → view → slice → transpose - # key_cache[blk_ids]: [n, kv_h, d//x, blk_sz, x] - 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 + num_blocks_i = (seq_len + block_size - 1) // block_size + blk_ids = block_tables[i, :num_blocks_i] - # Gather V: same pattern - 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 reshaped once: [kv_h, gqa, 1, d] fp32 — tiny for decode 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] - attn_w = torch.matmul( - q_grouped * scale, # [kv_h, gqa, 1, d] - k_t.unsqueeze(1)) # [kv_h, 1, d, seq_len] - attn_w = torch.softmax(attn_w, dim=-1) + # Online softmax accumulators (CCCL summary_stats_data pattern) + # accumulator = {m (running max), l (running sum_exp), o (running output)} + m = torch.full((num_kv_heads, gqa_ratio, 1), + float('-inf'), dtype=torch.float32, device=dev) + 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] - out_i = torch.matmul(attn_w, v_t.unsqueeze(1)) - output[i] = out_i.view(num_heads, head_dim).to(orig_dtype) + # Tile over KV blocks — CCCL spread_out_items_per_thread pattern + # Adaptive: tile_blocks = ceil(num_blocks / target_tiles) + # 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 - and (max_num_partitions == 1 or num_seqs * num_heads > 512)) - # CCCL segmented_reduce three-tier: don't force V1 for all lengths. - # V2 partitioned execution is better for medium-length sequences - # (8K-32K) where V1's single-CTA approach underutilizes 16 SMs. - # But V2 native may fail on BI-V100 — catch and fallback to V1. + # CCCL dispatch_reduce.cuh two-path dispatch architecture: + # single-tile: num_items ≤ threads × items → one CTA, zero temp buffer + # multi-tile: GridEvenShare partitions across sm_count × occupancy CTAs + # + # Paged attention equivalent: + # 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). - # Try V2 native; if it fails, fallback to V1. - try: - assert _PARTITION_SIZE % block_size == 0 + # Run PagedAttention V2. + 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. - # Score tensor = [kv_h, gqa, q_len, tile_sz] × 4 bytes - # Solve: kv_h × gqa × q_len × tile_sz × 4 ≤ budget + # CCCL spread_out_items_per_thread adaptive tile sizing. + # + # 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 - # Round down to block_size boundary - 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)) + mem_max_tokens = _SMEM_BUDGET_BYTES // score_row_bytes + mem_max_tokens = (mem_max_tokens // block_size) * block_size 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 - # instead of allocating m_blk, m_new, corr each iteration. - torch.amax(s, dim=-1, out=_m_blk) - torch.maximum(m, _m_blk, out=_m_new) - exp_s = s - _m_new.unsqueeze(-1) + # exp_s = s - new_max (in-place exp after del s) + m_blk = s.amax(dim=-1) + m_new = torch.maximum(m, m_blk) + exp_s = s - m_new.unsqueeze(-1) del s exp_s.exp_() - torch.exp(m - _m_new, out=_corr) - m.copy_(_m_new) - l.mul_(_corr).add_(exp_s.sum(dim=-1)) - o.mul_(_corr.unsqueeze(-1)).add_( + corr = torch.exp(m - m_new) + m.copy_(m_new) + del m_blk, 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 + 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. - # CCCL CachingDeviceAllocator: same buffers as Phase 1. - torch.amax(s, dim=-1, out=_m_blk) - torch.maximum(m, _m_blk, out=_m_new) - exp_s = s - _m_new.unsqueeze(-1) + # Online softmax update (identical to context phase). + m_blk = s.amax(dim=-1) + m_new = torch.maximum(m, m_blk) + exp_s = s - m_new.unsqueeze(-1) del s exp_s.exp_() - torch.exp(m - _m_new, out=_corr) - m.copy_(_m_new) - l.mul_(_corr).add_(exp_s.sum(dim=-1)) - o.mul_(_corr.unsqueeze(-1)).add_( + corr = torch.exp(m - m_new) + m.copy_(m_new) + del m_blk, 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 + del exp_s, v_t, corr # -------------------------------------------------------------- # Finalize: normalize running output by normalization factor. diff --git a/qwen3_6_scripts/sampler.py b/qwen3_6_scripts/sampler.py index 5b5f7db6..3129da68 100644 --- a/qwen3_6_scripts/sampler.py +++ b/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, - device=tokens.device) + # + # 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 = 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