#!/usr/bin/env python3 """muh/dispatch.py — Runtime policy dispatch for vllm kernel configuration This is the core of the muh competitive moat. CCCL's policy_selector is a compile-time C++ template that maps: (type_t, op_kind_t, accum_size, offset_size, compute_capability) → (threads_per_block, items_per_thread, vec_size, load_algorithm, ...) vllm doesn't use CUB directly — it uses PyTorch/Triton/custom CUDA kernels. But those kernels have the SAME tuning dimensions: - BLOCK_SIZE (= threads_per_block) - NUM_WARPS (= threads_per_block / 32) - PARTITION_SIZE (= threads_per_block * items_per_thread) - TILE_SIZE for shared memory This module provides a Python-side policy_selector that: 1. Reads bi100_* values from C++ headers (via gen_patch.extract_bi100_structs) 2. Maps CCCL algorithm→vllm kernel paths (the INJECTION_POINTS) 3. Applies SMEM constraints for BI-V100 (48KB limit) 4. Outputs the concrete values to inject into vllm source The moat is NOT the parameter values (anyone can benchmark those). The moat is: a) Knowing WHICH 7 dimensions to search (from CCCL's policy structs) b) Knowing the CONSTRAINTS (SMEM ≤ 48KB, occupancy, L2 coherence delay) c) Knowing WHERE in vllm each algorithm appears (the injection mapping) d) Having the infrastructure to iterate: benchmark → update header → gen_patch → rebuild """ import os import sys import json # Add parent dir for imports sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) from gen_patch import extract_bi100_structs, algo_from_filename # ────────────────────────────────────────────────────────────── # BI-V100 hardware constraints (from hardware.cuh) # These are the hard limits that make our tuning values different # from every other GPU — and why copy-pasting SM100 values crashes. # ────────────────────────────────────────────────────────────── BI_V100 = { "warp_size": 32, "max_threads_per_block": 1024, "max_shared_memory_per_block": 49152, # 48 KiB "max_registers_per_thread": 255, "l2_cache_size_bytes": 6 * 1024 * 1024, # 6 MiB "memory_bandwidth_gbps": 900, "sm_count": 16 # CONFIRMED 2026-08-01, # Derived "bandwidth_per_sm_gbps": 900 / 16 # 56.25 GB/s per SM, # 18 GB/s ≈ A100 level } SM100 = { "max_shared_memory_per_block": 49152, # same default, but can configure higher "l2_cache_size_bytes": 50 * 1024 * 1024, # 50 MiB "memory_bandwidth_gbps": 8000, "sm_count": 148, "bandwidth_per_sm_gbps": 8000 / 148, # 54 GB/s } # ────────────────────────────────────────────────────────────── # SMEM constraint checker # This is the single most important function in muh. # Every bi100_* struct MUST pass this check or the kernel will crash. # ────────────────────────────────────────────────────────────── def check_smem(threads: int, items: int, elem_bytes: int, smem_limit: int = BI_V100["max_shared_memory_per_block"]) -> dict: """Check if a tile fits in shared memory. Returns dict with: tile_bytes: actual shared memory usage fits: True if tile_bytes <= smem_limit utilization: tile_bytes / smem_limit (higher = more efficient but riskier) max_items: maximum items_per_thread that fits """ tile_bytes = threads * items * elem_bytes max_items = smem_limit // (threads * elem_bytes) if threads * elem_bytes > 0 else 0 return { "tile_bytes": tile_bytes, "fits": tile_bytes <= smem_limit, "utilization": tile_bytes / smem_limit if smem_limit > 0 else 0, "max_items": max_items, "overflow_bytes": max(0, tile_bytes - smem_limit), } def scale_mem_bound(nominal_4B_threads: int, nominal_4B_items: int, type_size: int) -> tuple: """Scale items and threads for a given type size, matching CCCL exactly. Mirrors cub::detail::scale_mem_bound() from util_arch.cuh lines 153-161. Returns (items_per_thread, threads_per_block) — items-first, matching CCCL's scaling_result struct field order. Three differences from the old muh version (all were bugs): 1. Return order: (items, threads) not (threads, items) 2. Items clamp upper bound: nominal * 2, not nominal * 1 (CCCL allows small types like char to double items_per_thread) 3. Threads SMEM cap: min(nominal, round_up(max_smem/(type*items), 32)) (prevents launching more threads than SMEM can feed) Verified against all 18 CCCL test cases in catch2_test_util_arch.cu. """ MAX_SMEM = 48 * 1024 # 49152 bytes, hardcoded in CCCL as max_smem_per_block # Step 1: scale items inversely with type size items = nominal_4B_items * 4 // type_size items = max(1, min(items, nominal_4B_items * 2)) # clamp: [1, 2*nominal] # Step 2: cap threads by SMEM constraint # round_up(x, 32) aligns to warp boundary smem_per_item = type_size * items if smem_per_item > 0: max_threads_by_smem = ((MAX_SMEM // smem_per_item + 31) // 32) * 32 else: max_threads_by_smem = nominal_4B_threads threads = min(nominal_4B_threads, max_threads_by_smem) return (items, threads) # items-first, matching CCCL scaling_result def scale_delay_for_l2(sm100_delay_ns: int, sm100_l2w: int) -> tuple: """Scale lookback delay parameters for BI-V100's smaller L2. SM100 L2 = 50MB, BI-V100 L2 = 6MB (8.3x smaller). Smaller L2 → faster coherence → shorter delays needed. Heuristic: ns *= 0.5, l2w *= 0.6 (to be refined by benchmark). """ bi100_ns = int(sm100_delay_ns * 0.5) bi100_l2w = int(sm100_l2w * 0.6) return (bi100_ns, bi100_l2w) # ────────────────────────────────────────────────────────────── # vllm kernel → CCCL algorithm mapping # # This is the strategic knowledge that makes CCCL useful for vllm. # Each entry maps a vllm kernel file to: # - The CCCL algorithm it implements (reduce, scan, sort, etc.) # - The data types it operates on (determines which bi100_* struct to use) # - The tuning dimensions that appear in the kernel code # # Built from reading: # - paged_attn.py (PagedAttention V1/V2 dispatch) # - prefix_prefill.py (Triton/PyTorch context attention) # - vllm/model_executor/layers/sampler.py (top-k/top-p) # - paged_attention_kernel_architecture.md (CCCL pattern mapping) # ────────────────────────────────────────────────────────────── VLLM_KERNEL_MAP = { # === DECODE HOT PATH (Output TPS × 16.796 = 83%) === "paged_attention_v1": { "cccl_algorithms": ["reduce"], "description": "Single-pass decode attention for seq_len ≤ 8192", "data_types": { "query": "float16", # Q: [num_seqs, num_heads, head_dim] "key_cache": "float16", # K: [num_blocks, num_kv_heads, head_dim//x, block_size, x] "score": "float32", # QK^T intermediate: always fp32 for precision "output": "float16", # weighted V sum }, "tuning_dimensions": { "NUM_THREADS": {"cccl_field": "threads_per_block", "range": [128, 256, 512]}, "NUM_WARPS": {"derived_from": "NUM_THREADS / 32"}, "_PARTITION_SIZE": {"value": 512, "note": "hardcoded in paged_attn.py, affects V2 threshold"}, }, "cccl_pattern": "compound reduce: summary_statistics.cu binary op pattern", "smem_formula": "NUM_THREADS * head_dim * sizeof(float) + head_dim * block_size * sizeof(half) * 2", }, "paged_attention_v2": { "cccl_algorithms": ["reduce", "scan"], "description": "Two-pass partitioned attention for seq_len > 8192", "data_types": { "score": "float32", "exp_sum": "float32", "max_logits": "float32", }, "tuning_dimensions": { "NUM_THREADS": {"cccl_field": "threads_per_block"}, "PARTITION_SIZE": {"cccl_field": "threads_per_block * items_per_thread"}, }, "cccl_pattern": "reduce pass 1 (per-partition) + reduce pass 2 (cross-partition merge)", }, "context_attention_fwd": { "cccl_algorithms": ["scan", "reduce", "transform"], "description": "Prefill attention (Triton kernel, bypassed on BI-V100)", "status": "BYPASSED — Triton hangs BI-V100, using _forward_prefix_pytorch", "tuning_dimensions": { "BLOCK_M": {"value": 64, "note": "query tile"}, "BLOCK_N": {"value": 64, "note": "KV tile"}, "BLOCK_DMODEL": {"value": 256, "note": "head_dim, must match model"}, }, "note": "PyTorch fallback has no tunable block sizes — optimization comes from algorithmic changes (K-tiling)", }, "sampling_topk": { "cccl_algorithms": ["topk", "radix_sort"], "description": "Top-k token selection from logits", "data_types": { "logits": "float32", # [batch, vocab_size=152064] "indices": "int32", }, "tuning_dimensions": { "BLOCK_SIZE": {"cccl_field": "threads_per_block"}, "RADIX_BITS": {"cccl_field": "bits_per_pass"}, }, }, "activation_kernels": { "cccl_algorithms": ["transform"], "description": "SiLU, GELU, element-wise activations", "data_types": {"input": "float16", "output": "float16"}, "tuning_dimensions": { "BLOCK_SIZE": {"cccl_field": "threads_per_block"}, "VEC_SIZE": {"cccl_field": "vec_size"}, }, }, "layernorm_kernels": { "cccl_algorithms": ["reduce", "transform"], "description": "RMSNorm / LayerNorm: reduce for variance, transform for normalize", "data_types": {"input": "float16", "accum": "float32"}, "tuning_dimensions": { "BLOCK_SIZE": {"cccl_field": "threads_per_block"}, }, }, "rotary_embedding": { "cccl_algorithms": ["for_each", "transform"], "description": "RoPE position encoding", "data_types": {"input": "float16"}, "tuning_dimensions": { "BLOCK_SIZE": {"cccl_field": "threads_per_block"}, }, }, # === CACHE PATH (Cache TPS × 0.56 = 3%) === "cache_kernels": { "cccl_algorithms": ["batch_memcpy"], "description": "KV cache block copy/swap operations", "data_types": {"kv_cache": "float16"}, "tuning_dimensions": { "BLOCK_SIZE": {"cccl_field": "threads_per_block"}, }, }, } # ────────────────────────────────────────────────────────────── # Policy dispatch: given a vllm kernel, return optimal BI-V100 config # ────────────────────────────────────────────────────────────── def dispatch_policy(kernel_name: str, tuning_headers_dir: str = "muh/include/muh/tuning") -> dict: """Given a vllm kernel name, return the optimal BI-V100 tuning parameters. This is the Python equivalent of CCCL's policy_selector::operator()(). It reads the C++ headers, applies SMEM constraints, and returns the concrete values to inject into the vllm kernel. """ if kernel_name not in VLLM_KERNEL_MAP: return {"error": f"Unknown kernel: {kernel_name}"} kernel_info = VLLM_KERNEL_MAP[kernel_name] cccl_algos = kernel_info["cccl_algorithms"] result = { "kernel": kernel_name, "description": kernel_info.get("description", ""), "policies": {}, "smem_checks": [], } for algo in cccl_algos: header_path = os.path.join(tuning_headers_dir, f"tuning_{algo}.cuh") if algo == "for_each": header_path = os.path.join(tuning_headers_dir, "tuning_for.cuh") if not os.path.exists(header_path): result["policies"][algo] = {"status": "NO_HEADER", "fallback": "CCCL_DEFAULT"} continue structs = extract_bi100_structs(header_path) if not structs: result["policies"][algo] = {"status": "NO_BI100_STRUCTS"} continue # Select the most relevant struct for this kernel's data types algo_policies = {} for name, fields in structs: # Check SMEM constraint threads = fields.get("threads", fields.get("threads_per_block", 256)) items = fields.get("items", fields.get("items_per_thread", 16)) # Determine element size from kernel data types elem_bytes = 4 # default to float32 if "float16" in str(kernel_info.get("data_types", {}).values()): elem_bytes = 2 if "score" in kernel_info.get("data_types", {}): elem_bytes = 4 # scores are always fp32 smem = check_smem(threads, items, elem_bytes) algo_policies[name] = {**fields, "_smem_check": smem} if not smem["fits"]: result["smem_checks"].append({ "struct": name, "OVERFLOW": True, "tile_bytes": smem["tile_bytes"], "limit": BI_V100["max_shared_memory_per_block"], "max_safe_items": smem["max_items"], }) result["policies"][algo] = algo_policies return result def dispatch_all(tuning_headers_dir: str = "muh/include/muh/tuning") -> dict: """Dispatch policies for ALL vllm kernels. Used by gen_patch.""" results = {} for kernel_name in VLLM_KERNEL_MAP: results[kernel_name] = dispatch_policy(kernel_name, tuning_headers_dir) return results # ────────────────────────────────────────────────────────────── # CLI: dump all dispatch results for inspection # ────────────────────────────────────────────────────────────── if __name__ == "__main__": import argparse p = argparse.ArgumentParser(description="muh policy dispatch for vllm kernels") p.add_argument("--headers", default="muh/include/muh/tuning") p.add_argument("--kernel", default=None, help="Specific kernel to dispatch") p.add_argument("--json", action="store_true", help="JSON output") args = p.parse_args() if args.kernel: result = dispatch_policy(args.kernel, args.headers) else: result = dispatch_all(args.headers) if args.json: print(json.dumps(result, indent=2, default=str)) else: for kernel_name, policy in (result.items() if isinstance(result, dict) and "kernel" not in result else [(result.get("kernel","?"), result)]): if isinstance(policy, dict) and "kernel" in policy: kernel_name = policy["kernel"] print(f"\n{'='*60}") print(f"Kernel: {kernel_name}") if isinstance(policy, dict): print(f" Description: {policy.get('description','')}") for algo, algo_policy in policy.get("policies", {}).items(): print(f" [{algo}]:") if isinstance(algo_policy, dict) and "status" in algo_policy: print(f" {algo_policy}") elif isinstance(algo_policy, dict): for struct_name, fields in algo_policy.items(): smem = fields.pop("_smem_check", {}) print(f" {struct_name}: {fields}") if smem: status = "✓" if smem.get("fits") else "✗ OVERFLOW" print(f" SMEM: {smem.get('tile_bytes',0)} bytes ({status})") for check in policy.get("smem_checks", []): print(f" ⚠ SMEM OVERFLOW: {check}")