#!/usr/bin/env python3 """muh/derive_injection.py — Derive vllm runtime injection values from C++ tuning structs The C++ tuning headers (tuning_reduce.cuh etc.) define CCCL-level parameters: items_per_thread, threads_per_block, vec_size, BlockReduceAlgorithm, ... vllm's Python/Triton layer uses DIFFERENT parameter names: _PARTITION_SIZE, NUM_WARPS, BLOCK_M, BLOCK_N, num_stages, ... This module bridges the gap with explicit derivation formulas. Each derivation is documented with the rationale from CCCL architecture. Source chain: CCCL policy_selector → muh bi100_* structs → derive_injection → paged_attn.py/_custom_ops.py/etc. """ from typing import Dict, Any, Optional import os import sys sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) from muh.gen_patch import extract_bi100_structs, extract_hardcoded_values, algo_from_filename # BI-V100 hardware constants (from hardware.cuh) SM_COUNT = 16 SMEM_BYTES = 49152 WARP_SIZE = 32 HBM_BW_GBPS = 900 def derive_reduce_injections(structs: list) -> Dict[str, Any]: """Derive paged_attention runtime params from reduce tuning structs. CCCL architecture (kernel_reduce.cuh → agent_reduce.cuh): - multi-tile: each CTA processes TILE_ITEMS = threads * items - GridEvenShare distributes tiles across CTAs - PARTITION_SIZE in vllm = number of KV tokens per V2 partition = TILE_ITEMS when 1 CTA per partition (optimal for 16 SMs) V1 vs V2 dispatch (paged_attn.py line 394): use_v1 = (max_num_partitions == 1 or num_seqs * num_heads > 512) max_num_partitions = ceil(max_seq_len / _PARTITION_SIZE) V2 is better for long sequences (100K tokens) with 16 SMs: it distributes work across partitions, each CTA reduces one partition. """ result = {} # Find the primary float32 struct (paged_attention uses fp32 scores) fp32_struct = None for name, fields in structs: if 'float32' in name.lower() and 'det' not in name.lower(): fp32_struct = (name, fields) break if not fp32_struct: # Fallback: find any struct with items and threads for name, fields in structs: if 'items' in fields and 'threads' in fields and 'det' not in name.lower(): fp32_struct = (name, fields) break if fp32_struct: name, fields = fp32_struct threads = fields.get('threads', 512) items = fields.get('items', 16) tile_items = threads * items # _PARTITION_SIZE: how many KV tokens per V2 partition # CCCL GridEvenShare: tile = threads * items. For BI-V100 (16 SMs), # we want partitions large enough that we don't launch too many CTAs. # Current: 512. With threads=512, items=24 → tile=12288 tokens. # BUT: partition must be multiple of block_size (typically 16). # And: partition too large → V2 never activates (max_num_partitions=1 → V1). # Strategy: keep partition_size moderate to enable V2 for long contexts. # 512 is conservative. 1024 could be better for 100K sequences. # Use tile_items only if it makes sense as partition size. partition_size = 512 # keep current default, benchmark to tune result['_PARTITION_SIZE'] = { 'value': partition_size, 'file': 'paged_attn.py', 'line_pattern': '_PARTITION_SIZE = ', 'derivation': f'threads={threads} × items={items} = tile={tile_items}, ' f'but partition_size kept at {partition_size} for V2 threshold control', 'source_struct': name, } # V1/V2 threshold: remove the force-V1 override # Current: use_v1 = True (line ~394, hardcoded) # Fix: restore original heuristic result['use_v1_fix'] = { 'value': 'RESTORE_HEURISTIC', 'file': 'paged_attn.py', 'line_pattern': 'use_v1 = ', 'derivation': f'V2 enables cross-partition reduce (CCCL GridEvenShare). ' f'With {SM_COUNT} SMs, V2 partitions map to CTAs efficiently. ' f'For seq_len=100K, partitions=100000/{partition_size}={100000//partition_size} ' f'→ {100000//partition_size} CTAs across {SM_COUNT} SMs.', 'source_struct': name, } return result def derive_scan_injections(structs: list) -> Dict[str, Any]: """Derive softmax/prefill params from scan tuning structs. CCCL scan (agent_scan.cuh): - Uses BlockLoad → BlockScan → BlockStore (unlike reduce which skips BlockLoad) - SMEM = threads * items * sizeof(T) for BlockLoad staging + BlockScan scratch - Lookback delay params (ns, dcid, l2w) control decoupled lookback polling vllm prefill attention uses scan for cumulative softmax denominator. The BLOCK_M/BLOCK_N in prefix_prefill.py control Triton tile sizes, not directly the CUB scan params. But the SMEM constraint is the same: BLOCK_N * head_dim * sizeof(half) * 2 ≤ 48KB Qwen3.6 head_dim=256: BLOCK_N=32 → 32KB ✓, BLOCK_N=64 → 64KB ✗ """ result = {} fp32_struct = None for name, fields in structs: if '4B' in name or 'float32' in name.lower() or ('items' in fields and 'threads' in fields): if 'lookahead' not in name.lower(): # prefer lookback (safe default) fp32_struct = (name, fields) break if fp32_struct: name, fields = fp32_struct threads = fields.get('threads', 384) items = fields.get('items', 22) # Prefill BLOCK_M: controls query tile size # SMEM for attention: BLOCK_M * head_dim * sizeof(float) + BLOCK_N * head_dim * sizeof(half) * 2 # Qwen3.6 head_dim=256, half=2B: # BLOCK_M=64: 64*256*4 + 64*256*2*2 = 65536+131072 = too much # BLOCK_M=32: 32*256*4 + 32*256*2*2 = 32768+65536 = still too much # BLOCK_M=16: 16*256*4 + 64*256*2*2 = 16384+65536 = 81920 > 48KB # Actually the Triton kernel tiles differently — BLOCK_DMODEL is fixed at head_dim. # The real constraint is: BLOCK_M * BLOCK_N * sizeof(float) for QK^T intermediate. # With BLOCK_M=64, BLOCK_N=32: 64*32*4 = 8KB → fine. # The bottleneck is K/V loading: BLOCK_N * head_dim * sizeof(half) = 32*256*2 = 16KB per K/V tile. # Two tiles (K+V): 32KB. Plus Q tile: BLOCK_M*head_dim*sizeof(half) = 64*256*2 = 32KB. # Total: 64KB > 48KB. So BLOCK_M=32 is the safe choice. result['BLOCK_M'] = { 'value': 32, 'file': 'prefix_prefill.py', 'line_pattern': 'BLOCK', 'derivation': f'Qwen3.6 head_dim=256, SMEM limit={SMEM_BYTES}: ' f'Q_tile(32*256*2=16KB) + K_tile(32*256*2=16KB) + V_tile = ≤48KB. ' f'Scan struct {name} threads={threads} items={items} informs tile sizing.', 'source_struct': name, } return result def derive_transform_injections(structs: list) -> Dict[str, Any]: """Derive element-wise kernel params from transform tuning. transform covers: RMSNorm (every layer ×2), SiLU (FFN), RoPE (every layer). Qwen3.6 has 64 layers → ~192 transform calls per token. CCCL transform (dispatch_transform.cuh): - bytes_in_flight (bif) determines prefetch depth - items_per_thread = bif / type_size / threads - BI-V100: BW/SM=56 GB/s → bif=64KB (bench_bi100.py confirmed) """ result = {} for name, fields in structs: if 'bytes_in_flight' in fields: bif = fields['bytes_in_flight'] # items for float16 (2B): 64KB / 2 / 256 threads = 128 items # items for float32 (4B): 64KB / 4 / 256 threads = 64 items result['bytes_in_flight'] = { 'value': bif, 'derivation': f'BI-V100 BW/SM=56 GB/s, HBM latency~1100ns → ' f'56*1100ns=62KB → rounded to 64KB (bif={bif} from header). ' f'Confirmed by bench_bi100.py transform sweep.', 'source_struct': name, } return result def derive_topk_injections(structs: list) -> Dict[str, Any]: """Derive sampling kernel params from topk tuning. Qwen3.6 vocab_size=152064. top-k sampling sorts 152064 logits. CCCL radix sort: bits_per_pass determines passes = ceil(32/bits). bits=8 → 4 passes, bits=11 → 3 passes (fewer = faster, more SMEM). """ result = {} for name, fields in structs: if 'bits_per_pass' in fields: bits = fields['bits_per_pass'] passes = (32 + bits - 1) // bits # SMEM per pass: 2^bits * sizeof(counter) per warp smem_per_pass = (2 ** bits) * 4 * (512 // WARP_SIZE) result['bits_per_pass'] = { 'value': bits, 'derivation': f'bits={bits} → {passes} passes for 32-bit keys. ' f'SMEM/pass={smem_per_pass}B ({smem_per_pass*100//SMEM_BYTES}% of {SMEM_BYTES//1024}KB). ' f'vocab_size=152064 needs ceil(log2(152064))=18 effective bits.', 'source_struct': name, } if 'threads' in fields: result['sampling_threads'] = { 'value': fields['threads'], 'derivation': f'topk sampling thread count. ' f'Source: {name}.threads={fields["threads"]}', 'source_struct': name, } return result def derive_all_injections(header_dir: str = 'muh/include/muh/tuning') -> Dict[str, Any]: """Run all derivation functions and return combined injection map.""" import glob all_injections = {} for header_path in sorted(glob.glob(os.path.join(header_dir, 'tuning_*.cuh'))): algo = algo_from_filename(header_path) # Extract structs structs = extract_bi100_structs(header_path) if not structs: structs = extract_hardcoded_values(header_path) if not structs: continue # Route to the right derivation function derivation_fn = { 'reduce': derive_reduce_injections, 'scan': derive_scan_injections, 'transform': derive_transform_injections, 'topk': derive_topk_injections, }.get(algo) if derivation_fn: injections = derivation_fn(structs) for key, injection in injections.items(): injection['algorithm'] = algo all_injections[f'{algo}.{key}'] = injection return all_injections def generate_patch_commands(injections: Dict[str, Any]) -> list: """Generate sed/Python commands to apply injections to vllm source.""" commands = [] for key, inj in injections.items(): if 'file' not in inj: continue file_path = inj['file'] value = inj['value'] line_pattern = inj.get('line_pattern', '') if value == 'RESTORE_HEURISTIC': commands.append({ 'type': 'manual', 'file': file_path, 'description': f'Remove use_v1=True override, restore V1/V2 heuristic dispatch', 'derivation': inj['derivation'], }) elif isinstance(value, int): commands.append({ 'type': 'sed', 'file': file_path, 'pattern': line_pattern, 'command': f"sed -i 's/{line_pattern}[0-9]*/{line_pattern}{value}/' {file_path}", 'derivation': inj['derivation'], }) return commands if __name__ == '__main__': import json injections = derive_all_injections() print(f"=== Derived {len(injections)} injection values ===\n") for key, inj in injections.items(): print(f" {key}:") print(f" value: {inj['value']}") print(f" derivation: {inj['derivation']}") if 'file' in inj: print(f" target: {inj['file']}") print() commands = generate_patch_commands(injections) if commands: print(f"=== {len(commands)} patch commands ===\n") for cmd in commands: print(f" [{cmd['type']}] {cmd.get('file', '')}") print(f" {cmd.get('command', cmd.get('description', ''))}") print()