[muh_dispatch] 修正 head_dim=128→256 + 删除死代码 + 强制 V1

从 CCCL cc_dispatch.cuh (150行) 读入完整的 compute capability 分派架构:
  dispatch_compute_cap → dispatch_to_cc_list → policy_getter<PolicySelector, CC>
  C++20: policy_constant 做相同 policy 的 CC 去重
  C++17: lowest_cc_resolver 找最低 CC with same policy

从 qwen3_5.py 确认 Qwen3.6-35B-A3B 实际参数:
  head_dim = 256 (NOT 128)
  num_heads = 24, num_kv_heads = 4
  GQA ratio = 6

关键修正:
1. head_dim 128→256
   旧: qwen36_config(head_dim=128) → BLOCK_N=64 → SMEM=64×128×2×2=32KB ✓
   实际: head_dim=256 → BLOCK_N=64 → SMEM=64×256×2×2=64KB > 48KB → CRASH
   修正: BLOCK_N=32 → SMEM=32×256×2×2=32KB ≤ 48KB ✓

2. 删除 _read_reduce_config (依赖 gen_patch, 容器内不可用)
3. 删除 reduce_threads/reduce_items (ixformer 有自己的 reduce, 我们控制不了)
4. 强制 V1 (v1_v2_threshold = max_seq_len + 1)
5. Pre-computed configs at import time (mirrors CCCL compile-time instantiation)

muh_dispatch.py

@@ -2,201 +2,55 @@
 muh_dispatch.py — CCCL-style type-dispatched kernel configuration for BI-V100
 ===============================================================================
 
-This is the key differentiator. Everyone else hardcodes:
+Mirrors CCCL's cc_dispatch.cuh architecture:
-    BLOCK_SIZE = 64
+  cc_dispatch: policy_selector(compute_capability) → policy struct
-    NUM_WARPS = 4
+  muh_dispatch: select_attention_config(hw, dtype, head_dim, ...) → AttentionConfig
-    PARTITION_SIZE = 512
 
-muh_dispatch replaces these with type-dispatched values derived from CCCL's
+Key corrections from CCCL source reading (cc_dispatch.cuh, 150 lines):
-policy_selector architecture. The dispatch key is (dtype, head_dim, seq_len),
+  - CCCL collapses architectures with identical policies (lowest_cc_resolver)
-and the output is a complete kernel configuration tuple.
+  - CCCL dispatches at COMPILE TIME via policy_getter<PolicySelector, CC>
+  - Python equivalent: precompute configs at import time, not per-call
 
-CCCL reference: cub/device/dispatch/tuning/tuning_reduce.cuh
+Source: cccl_upstream/cub/cub/detail/cc_dispatch.cuh
-  Input: (compute_capability, accum_type, op_kind, offset_size, determinism)
+        cccl_upstream/cub/cub/device/dispatch/dispatch_common.cuh
-  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:
+    """Mirrors muh/include/muh/hardware.cuh"""
     warp_size: int = 32
     max_threads_per_block: int = 1024
-    max_shared_memory_per_block: int = 49152  # 48KB
+    max_shared_memory_per_block: int = 49152  # 48KB — confirmed via ixsmi
-    sm_count: int = 16           # CONFIRMED: ixsmi shows 16 SMs per BI-V100 (NOT 50 from spec)
+    sm_count: int = 16           # CONFIRMED: 16 SMs per BI-V100 (NOT 50)
     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.
+    """Complete kernel config — mirrors CCCL's ReducePolicy/ScanPolicy output."""
-    
-    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_block_m: int = 32
+    triton_block_n: int = 32
     triton_num_warps: int = 4
-    
+    triton_num_stages: int = 1
+
     # Paged attention V1/V2 (decode)
     partition_size: int = 512
-    v1_v2_threshold: int = 8192  # seq_len above this → use V2
+    v1_v2_threshold: int = 8192
-    
+
-    # 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_native_v2: bool = False   # native V2 has correctness issues on BI-V100
     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,
@@ -206,130 +60,84 @@ def select_attention_config(
     num_kv_heads: int,
 ) -> AttentionConfig:
     """CCCL-style policy selector for paged attention.
-    
+
-    Dispatch axes (matching CCCL's type_t × op_kind_t × offset_size):
+    CCCL dispatch axes: (compute_capability, type_t, op_kind_t, offset_size)
-      - dtype → determines accum_size, SMEM per element
+    Our dispatch axes: (hardware, dtype, head_dim, max_seq_len, num_kv_heads)
-      - 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 = dtype.itemsize if hasattr(dtype, 'itemsize') else torch.tensor([], dtype=dtype).element_size()
-    elem_size = info["size"]
+    smem = hw.max_shared_memory_per_block
-    
+
     # --- Triton prefill config ---
-    # SMEM for flash attention = BLOCK_N × head_dim × elem_size × 2 (K+V)
+    # SMEM = BLOCK_N × head_dim × elem_size × 2 (K + V staging)
-    # Must fit in 48KB
+    # Must fit in 48KB with margin for softmax accumulators
-    triton_block_n = 128
+    # Qwen3.6: head_dim=256, bf16 → elem_size=2
-    triton_smem = triton_block_n * head_dim * elem_size * 2
+    #   BLOCK_N=64: 64×256×2×2 = 64KB > 48KB → CRASH
-    while triton_smem > hw.max_shared_memory_per_block and triton_block_n > 16:
+    #   BLOCK_N=32: 32×256×2×2 = 32KB ≤ 48KB ✓
+    #   BLOCK_N=64 only safe for head_dim≤128: 64×128×2×2 = 32KB
+
+    triton_block_n = 64
+    while triton_block_n * head_dim * elem_size * 2 > smem and triton_block_n > 16:
         triton_block_n //= 2
-        triton_smem = triton_block_n * head_dim * elem_size * 2
+
-    
+    # BLOCK_M: same as BLOCK_N for square tiles (simplifies causal mask)
-    # NUM_WARPS: bandwidth-limited GPU → fewer warps, more blocks
+    # BI-V100: 4 warps, not 8 (BLOCK=32 → 32 rows, 8 warps = 256 threads
-    # CCCL analogy: transform policy uses 128 threads (4 warps) for SM100
+    # means only 32/256=0.125 rows/thread — wasteful)
-    # because bulk operations are BW-limited
+    triton_block_m = triton_block_n
-    triton_num_warps = 4 if hw.memory_bandwidth_gbps < 1500 else 8
+    triton_num_warps = 4
-    
+
+    # fp32 halves the block (element size doubles → SMEM doubles)
+    if dtype == torch.float32:
+        triton_block_m //= 2
+        triton_block_n //= 2
+
+    # num_stages=1 on BI-V100: no async copy hardware (needs SM80+ cp.async)
+    triton_num_stages = 1
+
     # --- Paged attention decode config ---
-    # Score accumulator is always fp32 (4 bytes) regardless of KV dtype
+    # V1 threshold: for seq_len > threshold, V2 would be better IF V2 were native C++
-    accum_size = 4
+    # Currently V2 is PyTorch → always slower than V1 ixformer
-    
+    # So threshold is effectively infinite (always V1)
-    # Partition size for V2:
+    v1_threshold = max_seq_len + 1  # force V1
-    # 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_m=triton_block_m,
         triton_block_n=triton_block_n,
         triton_num_warps=triton_num_warps,
-        partition_size=partition_size,
+        triton_num_stages=triton_num_stages,
+        partition_size=512,
         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_native_v2=False,
         use_triton_prefill=True,
     )
 
 
-# ============================================================
+# Pre-computed configs (mirrors CCCL's compile-time policy instantiation)
-# Convenience: get config for Qwen3.6 on BI-V100
+# CCCL does this via template instantiation; we do it at import time.
-# ============================================================
 
-def qwen36_config() -> AttentionConfig:
+QWEN36_BF16 = select_attention_config(
-    """Pre-computed config for Qwen3.6-35B-A3B on BI-V100.
+    hw=BI_V100,
-    
+    dtype=torch.bfloat16,
-    Qwen3.6 uses:
+    head_dim=256,        # CONFIRMED from qwen3_5.py: text_cfg.head_dim = 256
-      - head_dim = 128
+    max_seq_len=100000,  # from computility-run.yaml: --max-model-len 100000
-      - num_heads = 64, num_kv_heads = 8 (GQA 8:1)
+    num_kv_heads=4,      # CONFIRMED: num_key_value_heads = 4
-      - 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,
-    )
 
+QWEN36_FP16 = select_attention_config(
+    hw=BI_V100,
+    dtype=torch.float16,
+    head_dim=256,
+    max_seq_len=100000,
+    num_kv_heads=4,
+)
 
-# ============================================================
-# Self-test
-# ============================================================
 
 if __name__ == "__main__":
     print("=== muh_dispatch: CCCL-style type-dispatched kernel config ===\n")
-    
+    print(f"Qwen3.6 bf16 (head_dim=256):")
-    configs = [
+    print(f"  triton: BLOCK_M={QWEN36_BF16.triton_block_m} BLOCK_N={QWEN36_BF16.triton_block_n}"
-        ("Qwen3.6 bf16 h128 100K", torch.bfloat16, 128, 100000, 8),
+          f" warps={QWEN36_BF16.triton_num_warps} stages={QWEN36_BF16.triton_num_stages}")
-        ("Qwen3.6 fp16 h128 100K", torch.float16, 128, 100000, 8),
+    print(f"  decode: partition={QWEN36_BF16.partition_size} v1_thresh={QWEN36_BF16.v1_v2_threshold}")
-        ("Qwen3.6 bf16 h256 100K", torch.bfloat16, 256, 100000, 8),
+    print(f"  SMEM: {QWEN36_BF16.triton_block_n}×256×2×2 = {QWEN36_BF16.triton_block_n*256*2*2} bytes"
-        ("Short context bf16 h128 2K", torch.bfloat16, 128, 2048, 8),
+          f" ({QWEN36_BF16.triton_block_n*256*2*2/1024:.0f}KB ≤ 48KB)")
-        ("fp32 fallback h128 32K", torch.float32, 128, 32768, 8),
+    print(f"  V1 forced: {QWEN36_BF16.use_native_v1} (V2 native has correctness issues)")
-    ]
-    
-    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()
-