diff --git a/engine_cccl_patterns.py b/engine_cccl_patterns.py new file mode 100644 index 00000000..f1f3c427 --- /dev/null +++ b/engine_cccl_patterns.py @@ -0,0 +1,408 @@ +""" +engine_cccl_patterns.py — CCCL 系统级设计模式移植到 BI-V100 vllm 引擎 +========================================================================== + +从 CCCL 源码中提取的不是参数值,而是架构设计模式。 +每个模式引用具体的 CCCL 源文件和行号。 + +核心发现(来自完整 CCCL 源码阅读): + +1. Reduce vs Scan 的 SMEM 差异: + - agent_reduce.cuh: 数据直接 striped load 到寄存器,NOT SMEM staging + → SMEM 只给 BlockReduce 的 warp shuffle scratch + → items_per_thread 不受 SMEM 限制,只受 register pressure 限制 + → BI-V100 可以用 items=24 (CCCL SM100 只用 items=16) + - agent_scan.cuh: 数据先 BlockLoad 到 SMEM staging buffer + → SMEM = BlockLoad::TempStorage ∪ BlockStore::TempStorage ∪ (BlockScan + Prefix) + → items_per_thread 严格受 tpb * ipt * type_size ≤ 48KB 约束 + → BI-V100 和 SM100 共享这个约束 + +2. scan delay 在 BI-V100 上完全无效: + - single_pass_scan_operators.cuh 第 130 行: + if (gridDim.x < GridThreshold=500) { __threadfence_block(); } + else { __nanosleep(Delay); } + - BI-V100: 16 SMs × 2 CTAs/SM = 32 blocks << 500 + - 结论: 所有 delay 策略退化为 __threadfence_block() + - 意味着 dcid/ns/l2w 三个参数在 BI-V100 上无效,不需要调 + +3. dispatch_reduce.cuh 的 two-phase 模式 = paged_attention_v2: + - Phase 1: DeviceReduceKernel → 每个 CTA 算一个 tile partition + - GridEvenShare 均匀分配 → 对应 V2 的 partition 分配 + - StableReductionOrder=false → atomic 聚合 (BI-V100: 16 SM 低争用) + - StableReductionOrder=true → write to d_out[blockIdx.x] + Phase 2 + - Phase 2: DeviceReduceSingleTileKernel → 一个 CTA 归约所有 partition 结果 + - 对应 V2 的 cross-partition log-sum-exp merge + +4. agent_reduce.cuh 的向量化加载条件: + ATTEMPT_VECTORIZATION = (vec_size > 1) && (items % vec == 0) + && is_pointer && is_trivially_relocatable + && sizeof(InputT) <= 8 + - PyTorch 等价: 用 .view().reshape() 做 contiguous 后 torch.bmm (已实现) + - 不等价: scatter/gather 非连续内存 → 强制 scalar path + +5. cc_dispatch.cuh 的 policy 折叠: + - lowest_cc_resolver: 多个 CC 生成相同 policy → 共享 kernel 实例化 + - BI-V100 等价: 所有 Qwen3.6 配置 (bf16, head_dim=256, kv_heads=4) + → 预计算一套配置,不做运行时 dispatch +""" + +import torch +from dataclasses import dataclass +from typing import Dict, Optional, Tuple + +# ═══════════════════════════════════════════════════════════════ +# Hardware descriptor — mirrors muh/include/muh/hardware.cuh +# ═══════════════════════════════════════════════════════════════ + +@dataclass(frozen=True) +class HW: + """BI-V100 hardware profile, confirmed via ixsmi on Phanthy Cloud.""" + sm_count: int = 16 + smem_per_block: int = 49152 # 48 KiB + warp_size: int = 32 + max_threads: int = 1024 + hbm_bw_gbps: int = 900 + l2_bytes: int = 6 * 1024 * 1024 # 6 MiB + bw_per_sm_gbps: float = 900 / 16 # 56.25 GB/s ≈ B200 level + bytes_in_flight: int = 64 * 1024 # bench_bi100.py verified: bif=8 wins + max_concurrent_ctas: int = 32 # 16 SM × ~2 occupancy + + # CCCL single_pass_scan_operators.cuh GridThreshold + # All grids < 500 blocks → delay() becomes __threadfence_block() + scan_delay_threshold: int = 500 + +BI100 = HW() + + +# ═══════════════════════════════════════════════════════════════ +# Pattern 1: CCCL GridEvenShare work distribution +# Source: cub/grid/grid_even_share.cuh +# Used by: dispatch_reduce.cuh line ~200 +# ═══════════════════════════════════════════════════════════════ + +def grid_even_share( + num_items: int, + sm_count: int = BI100.sm_count, + sm_occupancy: int = 2, + subscription_factor: int = 5, # CCCL util_device.cuh default + tile_size: int = 512 * 24, # threads × items for reduce +) -> Dict: + """ + CCCL's GridEvenShare maps work to CTAs. + dispatch_reduce.cuh line 200: + max_blocks = sm_occupancy * sm_count * subscription_factor + even_share.DispatchInit(num_items, max_blocks, tile_size) + + Returns partition plan for paged_attention_v2. + """ + max_blocks = sm_occupancy * sm_count * subscription_factor + # GridEvenShare.DispatchInit: divide num_items into even tiles + num_tiles = (num_items + tile_size - 1) // tile_size + grid_size = min(num_tiles, max_blocks) + + # For BI-V100: max_blocks = 2 × 16 × 5 = 160 + # For 100K tokens with tile=12288: num_tiles=9, grid=9 + # For 100K tokens with partition=1024: num_tiles=98, grid=98 + + return { + "num_items": num_items, + "tile_size": tile_size, + "max_blocks": max_blocks, + "grid_size": grid_size, + "items_per_cta": (num_items + grid_size - 1) // grid_size if grid_size > 0 else num_items, + "single_tile": num_tiles <= 1, + } + + +# ═══════════════════════════════════════════════════════════════ +# Pattern 2: CCCL AgentReduce tile consumption +# Source: agent_reduce.cuh ConsumeFullTile (two paths) +# Key insight: reduce does NOT use BlockLoad SMEM staging +# ═══════════════════════════════════════════════════════════════ + +def reduce_tile_config( + accum_size: int, # sizeof(AccumT) in bytes + hw: HW = BI100, +) -> Dict: + """ + Compute optimal reduce tile config for BI-V100. + + CCCL agent_reduce.cuh insight: data goes to REGISTERS not SMEM. + The SMEM constraint that limits scan (tpb*ipt*type_size ≤ 48KB) + does NOT apply to reduce. Instead, register pressure is the limit: + - Each thread holds AccumT items[ITEMS_PER_THREAD] in registers + - BI-V100 has 64K registers/SM (255 per thread max) + - items=24 for float32 → 24 registers → acceptable + - Larger items → fewer CTAs possible → but 16 SMs only need ~32 CTAs anyway + + Vectorized load condition (agent_reduce.cuh line ~243): + ATTEMPT_VECTORIZATION = vec_size > 1 && items % vec == 0 + && is_pointer && is_trivially_relocatable && sizeof <= 8 + """ + # Register pressure limit + regs_per_item = accum_size // 4 # 1 reg = 4 bytes for float32 + if regs_per_item < 1: + regs_per_item = 1 + + # Target: ~40 registers per thread total (data + overhead) + # 255 max regs per thread, but high reg usage reduces occupancy + max_items_by_regs = min(64, 40 // regs_per_item) + + # CCCL SM100 reference values + cccl_items = {1: 32, 2: 24, 4: 16, 8: 16, 16: 16} + reference = cccl_items.get(accum_size, 16) + + # BI-V100 adjustment: 16 SMs → larger tiles to compensate + # Each CTA should process more data (fewer CTAs total) + # Scale: items = reference × (SM100_count / BI100_count)^0.3 + # = reference × (148/16)^0.3 ≈ reference × 2.2 + # But cap at register limit + bi100_items = min(max_items_by_regs, int(reference * 2.0)) + + # Threads: 512 for most types (CCCL SM100 default) + # Except float64 where CCCL uses 640 → BI-V100 uses 384 (12 warps, clean) + threads = 384 if accum_size >= 8 else 512 + + # Vectorization + if accum_size <= 8 and bi100_items % 2 == 0: + vec_size = 2 if accum_size >= 4 else 4 + else: + vec_size = 1 + + return { + "threads": threads, + "items": bi100_items, + "vec_size": vec_size, + "tile_size": threads * bi100_items, + "regs_per_thread": bi100_items * regs_per_item + 16, # +16 for overhead + "smem_limited": False, # reduce is NOT SMEM limited + "cccl_reference_items": reference, + } + + +# ═══════════════════════════════════════════════════════════════ +# Pattern 3: CCCL AgentScan tile with SMEM staging +# Source: agent_scan.cuh ConsumeTile +# Key insight: scan DOES use BlockLoad SMEM staging → strict SMEM limit +# ═══════════════════════════════════════════════════════════════ + +def scan_tile_config( + accum_size: int, + hw: HW = BI100, +) -> Dict: + """ + Compute optimal scan tile config for BI-V100. + + CCCL agent_scan.cuh: uses BlockLoad → data goes through SMEM staging. + _TempStorage is a union of: + - BlockLoadT::TempStorage (tpb * items * type_size) + - BlockStoreT::TempStorage (tpb * items * type_size) + - BlockScanT::TempStorage + TilePrefixCallbackOpT::TempStorage + + The BlockLoad/Store staging is the SMEM bottleneck: + tpb * ipt * accum_size ≤ 48KB + + Additional constraint: WARP_TRANSPOSE load requires + tpb * ipt * sizeof(AccumT) bytes of staging buffer. + + CCCL SM100 scan benchmark winners: + float32/o4: ipt=22, tpb=384 (tile=33792 ≤ 48K) → speedup 1.148 + float64/o4: ipt=23, tpb=416 (tile=76544 > 48K!) → uses NoScaling + int8/o4: ipt=18, tpb=512 (tile=9216 ≤ 48K) + + But wait — CCCL SM100 float64 tile = 416*23*8 = 76544 > 49152! + How does this work? Because SM100 can configure larger SMEM (228KB). + BI-V100 is stuck at 48KB → must reduce items for large types. + """ + max_smem = hw.smem_per_block + + # Start with CCCL SM100 winners, then constrain + cccl_configs = { + 1: (512, 18), # int8: 512*18*1 = 9216 + 2: (512, 13), # int16: 512*13*2 = 13312 + 4: (384, 22), # float32: 384*22*4 = 33792 ✓ + 8: (384, 14), # float64: 384*14*8 = 43008 ✓ (reduced from SM100's 23) + 16: (256, 12), # int128: 256*12*16= 49152 = exactly 48KB + } + + threads, items = cccl_configs.get(accum_size, (384, 16)) + + # Verify SMEM constraint + tile_bytes = threads * items * accum_size + while tile_bytes > max_smem and items > 1: + items -= 1 + tile_bytes = threads * items * accum_size + + # scan delay is IRRELEVANT on BI-V100 + # single_pass_scan_operators.cuh: gridDim.x < 500 → __threadfence_block() + # BI-V100 max grid = ~160 << 500, so ALL delay strategies collapse + delay_effective = "threadfence_block_only" + + return { + "threads": threads, + "items": items, + "tile_size": threads * items, + "tile_bytes": threads * items * accum_size, + "smem_utilization": (threads * items * accum_size) / max_smem, + "smem_limited": True, # scan IS SMEM limited + "delay_strategy": delay_effective, + "load_algorithm": "WARP_TRANSPOSE" if accum_size >= 4 else "DIRECT", + } + + +# ═══════════════════════════════════════════════════════════════ +# Pattern 4: CCCL compound reduce (summary_statistics.cu Welford) +# Source: thrust/examples/summary_statistics.cu +# Maps to: paged_attention_v2 cross-partition merge +# ═══════════════════════════════════════════════════════════════ + +def compound_reduce_merge( + max_a: torch.Tensor, # [H, P_a] partition maxima from partition set A + sum_a: torch.Tensor, # [H, P_a] partition exp-sums + out_a: torch.Tensor, # [H, P_a, d] partition weighted outputs + max_b: torch.Tensor, # [H, P_b] + sum_b: torch.Tensor, # [H, P_b] + out_b: torch.Tensor, # [H, P_b, d] +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """ + Merge two sets of attention partition results. + + Direct translation of summary_statistics.cu binary_op, + adapted for online softmax instead of Welford variance: + + CCCL summary_stats_binary_op (lines 97-125): + n = x.n + y.n + delta = y.mean - x.mean + mean = x.mean + delta * y.n / n + M2 = x.M2 + y.M2 + delta^2 * x.n * y.n / n + + Our attention equivalent: + global_max = max(max_a, max_b) // delta = max_b - max_a + rescale_a = exp(max_a - global_max) // similar to delta normalization + rescale_b = exp(max_b - global_max) + total_sum = sum_a * rescale_a + sum_b * rescale_b + merged_out = (out_a * sum_a * rescale_a + out_b * sum_b * rescale_b) / total_sum + + The Welford parallel merge and log-sum-exp merge are structurally + identical — both need to rescale accumulated statistics when combining + partial results computed with different reference points (mean vs max). + + This function enables incremental/streaming V2: process new KV blocks + without recomputing from scratch. CCCL's ConsumeTiles pattern: + for each tile: ConsumeFullTile → ThreadReduce → update aggregate + becomes: + for each new KV block batch: compute partition → merge with running result + """ + H = max_a.shape[0] + device = max_a.device + + # Concatenate along partition dimension + all_max = torch.cat([max_a, max_b], dim=1) # [H, P_a + P_b] + all_sum = torch.cat([sum_a, sum_b], dim=1) + all_out = torch.cat([out_a, out_b], dim=1) # [H, P_a + P_b, d] + + # Global max for numerical stability + global_max = all_max.max(dim=1, keepdim=True).values # [H, 1] + + # Rescale: exp(partition_max - global_max) * partition_sum + rescale = torch.exp(all_max - global_max) * all_sum # [H, P] + total = rescale.sum(dim=1, keepdim=True) # [H, 1] + + # Weighted merge: bmm(rescale, out) / total + # CCCL norm.cu insight: fuse transform with reduce to minimize traversals + result = torch.bmm(rescale.unsqueeze(1), all_out.float()).squeeze(1) / total # [H, d] + + # Return merged statistics (for further merging if needed) + merged_max = global_max.squeeze(1) # [H] + merged_sum = total.squeeze(1) # [H] + merged_out = result.unsqueeze(1) # [H, 1, d] + + return merged_max, merged_sum, merged_out + + +# ═══════════════════════════════════════════════════════════════ +# Pattern 5: CCCL dispatch_compute_cap policy precomputation +# Source: cc_dispatch.cuh lowest_cc_resolver +# BI-V100: all Qwen3.6 configs precomputed at import time +# ═══════════════════════════════════════════════════════════════ + +# Pre-computed tile configs for all Qwen3.6 data types +# (mirrors CCCL's compile-time policy instantiation) +REDUCE_CONFIGS = { + "float16": reduce_tile_config(2), # KV cache values + "bfloat16": reduce_tile_config(2), + "float32": reduce_tile_config(4), # attention scores + "float64": reduce_tile_config(8), # (rarely used) + "int32": reduce_tile_config(4), # indices +} + +SCAN_CONFIGS = { + "float32": scan_tile_config(4), # softmax denominator + "float64": scan_tile_config(8), + "int32": scan_tile_config(4), +} + +# Qwen3.6 specific: paged attention V2 partition plan +QWEN36_V2_PLAN = grid_even_share( + num_items=100000, # max_model_len + tile_size=1024, # PARTITION_SIZE +) + + +# ═══════════════════════════════════════════════════════════════ +# Pattern 6: CCCL single-tile fast path +# Source: kernel_reduce.cuh line ~270 (DeviceReduceSingleTileKernel) +# dispatch_reduce.cuh Invoke(): if small → InvokeSingleTile +# ═══════════════════════════════════════════════════════════════ + +def should_use_single_tile( + seq_len: int, + partition_size: int = 1024, + reduce_config: Dict = None, +) -> bool: + """ + CCCL dispatch_reduce.cuh decision logic: + if (num_items <= threads * items_per_thread): + InvokeSingleTile() # one CTA, no Phase 2 + else: + InvokePasses() # multi-CTA + reduce + + For paged attention: + - tokens ≤ partition_size → one partition → no Phase 2 merge needed + - This is the common case during early decode (seq_len grows from 1 up) + - Avoids partition overhead for the majority of decode steps + """ + if reduce_config is None: + reduce_config = REDUCE_CONFIGS["float32"] + single_tile_capacity = reduce_config["tile_size"] # e.g. 512 * 24 = 12288 + + # Two conditions (from CCCL): + # 1. Fits in one partition → skip partitioning entirely + # 2. Fits in one CTA's tile → skip GridEvenShare overhead + return seq_len <= partition_size or seq_len <= single_tile_capacity + + +if __name__ == "__main__": + print("=== CCCL Pattern Analysis for BI-V100 ===\n") + + print("Reduce configs (NOT SMEM limited — register pressure only):") + for dtype, cfg in REDUCE_CONFIGS.items(): + print(f" {dtype}: threads={cfg['threads']}, items={cfg['items']}, " + f"vec={cfg['vec_size']}, tile={cfg['tile_size']}, " + f"regs/thread≈{cfg['regs_per_thread']}") + + print("\nScan configs (SMEM limited — strict 48KB constraint):") + for dtype, cfg in SCAN_CONFIGS.items(): + print(f" {dtype}: threads={cfg['threads']}, items={cfg['items']}, " + f"tile_bytes={cfg['tile_bytes']}, " + f"smem_util={cfg['smem_utilization']:.0%}, " + f"delay={cfg['delay_strategy']}") + + print(f"\nQwen3.6 V2 partition plan (100K tokens):") + plan = QWEN36_V2_PLAN + print(f" partitions={plan['grid_size']}, per_cta={plan['items_per_cta']}, " + f"max_blocks={plan['max_blocks']}, single_tile={plan['single_tile']}") + + print(f"\nSingle-tile threshold examples:") + for sl in [100, 500, 1024, 5000, 12288, 50000]: + print(f" seq_len={sl:>6d}: single_tile={should_use_single_tile(sl)}")