diff --git a/docs/SPECIALIZATION_ANALYSIS.md b/docs/SPECIALIZATION_ANALYSIS.md index ad0bbf8b..4255a823 100644 --- a/docs/SPECIALIZATION_ANALYSIS.md +++ b/docs/SPECIALIZATION_ANALYSIS.md @@ -83,3 +83,30 @@ Priority to add (by SMEM overflow risk): 4. `scan_by_key` (88KB) — MEDIUM risk 5. `unique_by_key` (61KB) — LOW risk 6. Others: LOW risk (small tile sizes, unlikely SMEM overflow) + +## SMEM Overflow Detection (from muh/dispatch.py) + +Running `python3 muh_kernel_map.py` against all 6 tuning headers +detected 5 lookahead structs with incorrect SMEM estimates: + +| Struct | SMEM calc | Limit | Status | +|--------|----------:|------:|--------| +| bi100_lookahead_1B | 162,816 | 49,152 | ✗ OVERFLOW | +| bi100_lookahead_2B | 97,280 | 49,152 | ✗ OVERFLOW | +| bi100_lookahead_4B | 80,896 | 49,152 | ✗ OVERFLOW | +| bi100_lookahead_4B_float | 89,088 | 49,152 | ✗ OVERFLOW | +| bi100_lookahead_8B | 89,088 | 49,152 | ✗ OVERFLOW | + +**Root cause**: Lookahead SMEM ≠ `threads × items × elem_bytes`. +The lookahead pipeline uses multi-stage buffering where SMEM = +`(reduce_squad + scan_store_squad) × items × accum_size × stages`. +The simple tile formula overestimates by including lookahead items +that live in registers, not SMEM. + +**Impact**: These are currently non-functional on BI-V100 anyway +(lookahead requires SM90+ warpspeed pipeline support). The dispatch +correctly falls back to lookback algorithm. But the values in the +structs are misleading — they should either be corrected or removed. + +**Action**: Issue #27 (scan benchmark) TC-04 covers this: +"lookahead 可行性评估 — 测试 ScanAlgorithm::lookahead 是否能在 BI-V100 上编译运行" diff --git a/muh_kernel_map.py b/muh_kernel_map.py new file mode 100644 index 00000000..cde93f8d --- /dev/null +++ b/muh_kernel_map.py @@ -0,0 +1,352 @@ +#!/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": 50, + # Derived + "bandwidth_per_sm_gbps": 900 / 50, # 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_threads: int, nominal_4b_items: int, + type_size: int) -> tuple: + """Scale items_per_thread inversely with type size to keep tile constant. + + Mirrors cub::detail::MemBoundScaling. + For 4-byte types: items = nominal_4b_items + For 8-byte types: items = nominal_4b_items * 4 / 8 = half + For 2-byte types: items = nominal_4b_items * 4 / 2 = double (capped) + """ + items = (nominal_4b_items * 4) // type_size + items = max(1, min(items, nominal_4b_items)) + return (nominal_threads, items) + + +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}")