project_6/muh_dispatch.py

"""
muh_dispatch.py — CCCL-style type-dispatched kernel configuration for BI-V100
===============================================================================

This is the key differentiator. Everyone else hardcodes:
    BLOCK_SIZE = 64
    NUM_WARPS = 4
    PARTITION_SIZE = 512

muh_dispatch replaces these with type-dispatched values derived from CCCL's
policy_selector architecture. The dispatch key is (dtype, head_dim, seq_len),
and the output is a complete kernel configuration tuple.

CCCL reference: cub/device/dispatch/tuning/tuning_reduce.cuh
  Input: (compute_capability, accum_type, op_kind, offset_size, determinism)
  Output: ReducePolicy{multi_tile, single_tile} where each pass has
          (threads, items, vec_size, algorithm, load_modifier)

muh_dispatch translation for paged attention:
  Input: (hardware, dtype, head_dim, seq_len, num_kv_heads)
  Output: AttentionConfig{partition_size, block_size, num_warps, vec_size, 
                          v1_threshold, use_triton}

Deploy: cp muh_dispatch.py /usr/local/corex/.../vllm/muh_dispatch.py
        Then patch paged_attn.py to import and use it.
"""

import os
import sys
import torch
from dataclasses import dataclass
from typing import Optional

# ============================================================
# Hardware descriptor — mirrors muh/include/muh/hardware.cuh
# ============================================================

@dataclass(frozen=True)
class HardwareCapability:
    warp_size: int = 32
    max_threads_per_block: int = 1024
    max_shared_memory_per_block: int = 49152  # 48KB
    sm_count: int = 16           # CONFIRMED: ixsmi shows 16 SMs per BI-V100 (NOT 50 from spec)
    memory_bandwidth_gbps: int = 900
    l2_cache_size_bytes: int = 6 * 1024 * 1024  # 6MB

BI_V100 = HardwareCapability()

# ============================================================
# Type classification — mirrors cub/device/dispatch/tuning/common.cuh
# ============================================================

def classify_dtype(dtype: torch.dtype) -> dict:
    """Classify a torch dtype into CCCL-compatible type descriptors."""
    type_map = {
        torch.float16:  {"size": 2, "type_t": "float16",  "is_float": True},
        torch.bfloat16: {"size": 2, "type_t": "bfloat16", "is_float": True},
        torch.float32:  {"size": 4, "type_t": "float32",  "is_float": True},
        torch.float64:  {"size": 8, "type_t": "float64",  "is_float": True},
        torch.int8:     {"size": 1, "type_t": "int8",     "is_float": False},
        torch.int32:    {"size": 4, "type_t": "int32",     "is_float": False},
        torch.int64:    {"size": 8, "type_t": "int64",     "is_float": False},
    }
    return type_map.get(dtype, {"size": dtype.itemsize, "type_t": "other", "is_float": False})


# ============================================================
# C++ header reader — single source of truth for tuning values
#
# Architecture: "read once, not write twice + assert equal"
# muh_dispatch.py never hand-writes tuning values. It reads them
# from the C++ headers via gen_patch.extract_bi100_structs().
# If headers aren't available (e.g. in a deployed container),
# falls back to compiled-in defaults with a warning.
# ============================================================

_TUNING_CACHE = {}

def _read_reduce_config(accum_size: int) -> dict:
    """Read reduce tuning values from tuning_reduce.cuh.
    
    Returns {"threads": int, "items": int} for the given accum_size.
    Single source of truth: C++ header → Python, no hand-written copy.
    """
    cache_key = f"reduce_{accum_size}"
    if cache_key in _TUNING_CACHE:
        return _TUNING_CACHE[cache_key]
    
    # Try to read from C++ headers
    header_path = os.path.join(
        os.path.dirname(os.path.abspath(__file__)),
        "muh", "include", "muh", "tuning", "tuning_reduce.cuh"
    )
    
    result = None
    if os.path.exists(header_path):
        try:
            sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
            from gen_patch import extract_bi100_structs
            structs = extract_bi100_structs(header_path)
            
            # Select struct by accum_size — names match tuning_reduce.cuh
            target_struct = None
            if accum_size <= 2:
                target_struct = "bi100_plus_accum2_o4"
            elif accum_size <= 4:
                target_struct = "bi100_plus_float32_o4"
            else:
                target_struct = "bi100_plus_float64_o4"
            
            for name, fields in structs:
                if name == target_struct:
                    result = {
                        "threads": fields.get("threads", fields.get("threads_per_block", 256)),
                        "items": fields.get("items", fields.get("items_per_thread", 16)),
                    }
                    break
            
            if result is None:
                # Struct not found — try default
                for name, fields in structs:
                    if "default" in name:
                        result = {
                            "threads": fields.get("threads", 256),
                            "items": fields.get("items", 16),
                        }
                        break
        except Exception as e:
            import warnings
            warnings.warn(
                f"muh_dispatch: failed to read {header_path}: {e}. "
                f"Using compiled-in fallback values.",
                RuntimeWarning, stacklevel=2
            )
    
    # Fallback: compiled-in defaults (last-resort, should not be the normal path)
    if result is None:
        # These values match the C++ headers as of commit 3a2b67c1.
        # If you're seeing this warning in production, the header path is wrong.
        import warnings
        warnings.warn(
            "muh_dispatch: C++ headers not found, using compiled-in fallback. "
            "This means tuning values may be stale.",
            RuntimeWarning, stacklevel=2
        )
        if accum_size <= 4:
            result = {"threads": 512, "items": 16}
        else:
            result = {"threads": 512, "items": 12}
    
    _TUNING_CACHE[cache_key] = result
    return result

# ============================================================
# Attention kernel configuration
# ============================================================

@dataclass
class AttentionConfig:
    """Complete kernel configuration for one paged attention call.
    
    Mirrors CCCL's ReducePolicy / ScanPolicy output structure:
    a single struct containing all parameters the kernel needs.
    """
    # Triton flash attention (prefill)
    triton_block_n: int = 64
    triton_num_warps: int = 4
    
    # Paged attention V1/V2 (decode)
    partition_size: int = 512
    v1_v2_threshold: int = 8192  # seq_len above this → use V2
    
    # Vectorization — derived from dtype
    vec_size: int = 4  # elements per vector load
    
    # Reduce pattern (score reduction per head)
    reduce_threads: int = 512
    reduce_items: int = 16
    
    # Backend selection
    use_native_v1: bool = True
    use_native_v2: bool = False  # V2 native has correctness issues
    use_triton_prefill: bool = True

# ============================================================
# policy_selector — the CCCL-style dispatch function
#
# This is the core: instead of one set of hardcoded constants,
# we dispatch based on (dtype, head_dim, seq_len).
#
# Why this matters for competition:
#   - fp16 attention with head_dim=128: score accum is fp32 (4B)
#     → reduce can use ipt=16, tpb=512 (tile=32KB ≤ 48KB)
#   - fp16 attention with head_dim=256: score tile is 2x larger
#     → reduce must use ipt=8 to fit SMEM
#   - Long sequences (>32K): partition_size=1024 better amortizes
#     the V2 reduce overhead
#   - Short sequences (<1K): V1 always wins, skip V2 entirely
# ============================================================

def select_attention_config(
    hw: HardwareCapability,
    dtype: torch.dtype,
    head_dim: int,
    max_seq_len: int,
    num_kv_heads: int,
) -> AttentionConfig:
    """CCCL-style policy selector for paged attention.
    
    Dispatch axes (matching CCCL's type_t × op_kind_t × offset_size):
      - dtype → determines accum_size, SMEM per element
      - head_dim → determines tile width
      - max_seq_len → determines V1/V2 threshold and partition_size
      - num_kv_heads → determines GQA ratio (affects memory pattern)
    """
    info = classify_dtype(dtype)
    elem_size = info["size"]
    
    # --- Triton prefill config ---
    # SMEM for flash attention = BLOCK_N × head_dim × elem_size × 2 (K+V)
    # Must fit in 48KB
    triton_block_n = 128
    triton_smem = triton_block_n * head_dim * elem_size * 2
    while triton_smem > hw.max_shared_memory_per_block and triton_block_n > 16:
        triton_block_n //= 2
        triton_smem = triton_block_n * head_dim * elem_size * 2
    
    # NUM_WARPS: bandwidth-limited GPU → fewer warps, more blocks
    # CCCL analogy: transform policy uses 128 threads (4 warps) for SM100
    # because bulk operations are BW-limited
    triton_num_warps = 4 if hw.memory_bandwidth_gbps < 1500 else 8
    
    # --- Paged attention decode config ---
    # Score accumulator is always fp32 (4 bytes) regardless of KV dtype
    accum_size = 4
    
    # Partition size for V2:
    # Larger partition = fewer partitions = less reduce overhead
    # But each partition must fit: partition_size × head_dim × accum_size in SMEM
    # CCCL parallel: reduce tile_size = threads × items × accum_size ≤ SMEM
    partition_smem = lambda ps: ps * head_dim * accum_size
    partition_size = 1024
    while partition_smem(partition_size) > hw.max_shared_memory_per_block:
        partition_size //= 2
    if partition_size < 256:
        partition_size = 256  # minimum for occupancy
    
    # V1/V2 threshold:
    # V1 is one block per (seq, head) — good for short seq
    # V2 splits into partitions — good for long seq
    # Crossover depends on SM count (more SMs → V2 wins earlier)
    # CCCL parallel: single_tile vs multi_tile in ReducePolicy
    if max_seq_len <= 2048:
        v1_threshold = max_seq_len + 1  # always V1
    elif hw.sm_count >= 80:
        v1_threshold = 4096  # high SM count → V2 wins earlier
    else:
        v1_threshold = 8192  # 50 SMs → V2 wins later
    
    # Vec size for score loads:
    # CCCL analogy: reduce uses vec_size=2 for fp32, vec_size=1 for fp64
    # because 128-bit loads = 4×fp32 = 2×fp64
    vec_size = min(16 // accum_size, 4)  # 128-bit / accum_size
    
    # Reduce config (for V2's final reduction across partitions):
    # Read from C++ headers — single source of truth, no hand-written copy.
    reduce_cfg = _read_reduce_config(accum_size)
    reduce_threads = reduce_cfg["threads"]
    reduce_items = reduce_cfg["items"]
    
    # Sanity check: reduce tile fits SMEM
    reduce_tile = reduce_threads * reduce_items * accum_size
    while reduce_tile > hw.max_shared_memory_per_block:
        reduce_items -= 1
        reduce_tile = reduce_threads * reduce_items * accum_size
    
    return AttentionConfig(
        triton_block_n=triton_block_n,
        triton_num_warps=triton_num_warps,
        partition_size=partition_size,
        v1_v2_threshold=v1_threshold,
        vec_size=vec_size,
        reduce_threads=reduce_threads,
        reduce_items=reduce_items,
        use_native_v1=True,
        use_native_v2=False,  # still correctness issues
        use_triton_prefill=True,
    )


# ============================================================
# Convenience: get config for Qwen3.6 on BI-V100
# ============================================================

def qwen36_config() -> AttentionConfig:
    """Pre-computed config for Qwen3.6-35B-A3B on BI-V100.
    
    Qwen3.6 uses:
      - head_dim = 128
      - num_heads = 64, num_kv_heads = 8 (GQA 8:1)
      - dtype = bfloat16 / float16
      - max_model_len = 100000
    """
    return select_attention_config(
        hw=BI_V100,
        dtype=torch.bfloat16,
        head_dim=128,
        max_seq_len=100000,
        num_kv_heads=8,
    )


# ============================================================
# Self-test
# ============================================================

if __name__ == "__main__":
    print("=== muh_dispatch: CCCL-style type-dispatched kernel config ===\n")
    
    configs = [
        ("Qwen3.6 bf16 h128 100K", torch.bfloat16, 128, 100000, 8),
        ("Qwen3.6 fp16 h128 100K", torch.float16, 128, 100000, 8),
        ("Qwen3.6 bf16 h256 100K", torch.bfloat16, 256, 100000, 8),
        ("Short context bf16 h128 2K", torch.bfloat16, 128, 2048, 8),
        ("fp32 fallback h128 32K", torch.float32, 128, 32768, 8),
    ]
    
    for name, dtype, hdim, seqlen, kvh in configs:
        cfg = select_attention_config(BI_V100, dtype, hdim, seqlen, kvh)
        print(f"  {name}:")
        print(f"    triton: block_n={cfg.triton_block_n} warps={cfg.triton_num_warps}")
        print(f"    decode: partition={cfg.partition_size} v1_thresh={cfg.v1_v2_threshold}")
        print(f"    reduce: threads={cfg.reduce_threads} items={cfg.reduce_items} vec={cfg.vec_size}")
        print()