From 3426d8185aa6e59932b14811aadcb4e011b2c775 Mon Sep 17 00:00:00 2001 From: muh-bot Date: Wed, 5 Aug 2026 07:12:26 +0000 Subject: [PATCH] =?UTF-8?q?[muh=5Fdispatch]=20=E4=BF=AE=E6=AD=A3=20head=5F?= =?UTF-8?q?dim=3D128=E2=86=92256=20+=20=E5=88=A0=E9=99=A4=E6=AD=BB?= =?UTF-8?q?=E4=BB=A3=E7=A0=81=20+=20=E5=BC=BA=E5=88=B6=20V1?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 从 CCCL cc_dispatch.cuh (150行) 读入完整的 compute capability 分派架构: dispatch_compute_cap → dispatch_to_cc_list → policy_getter 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 | 356 +++++++++++------------------------------------- 1 file changed, 82 insertions(+), 274 deletions(-) diff --git a/muh_dispatch.py b/muh_dispatch.py index e59c3745..b86994bb 100644 --- a/muh_dispatch.py +++ b/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: - BLOCK_SIZE = 64 - NUM_WARPS = 4 - PARTITION_SIZE = 512 +Mirrors CCCL's cc_dispatch.cuh architecture: + cc_dispatch: policy_selector(compute_capability) → policy struct + muh_dispatch: select_attention_config(hw, dtype, head_dim, ...) → AttentionConfig -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. +Key corrections from CCCL source reading (cc_dispatch.cuh, 150 lines): + - CCCL collapses architectures with identical policies (lowest_cc_resolver) + - CCCL dispatches at COMPILE TIME via policy_getter + - Python equivalent: precompute configs at import time, not per-call -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. +Source: cccl_upstream/cub/cub/detail/cc_dispatch.cuh + cccl_upstream/cub/cub/device/dispatch/dispatch_common.cuh """ -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 - sm_count: int = 16 # CONFIRMED: ixsmi shows 16 SMs per BI-V100 (NOT 50 from spec) + max_shared_memory_per_block: int = 49152 # 48KB — confirmed via ixsmi + 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. - - Mirrors CCCL's ReducePolicy / ScanPolicy output structure: - a single struct containing all parameters the kernel needs. - """ + """Complete kernel config — mirrors CCCL's ReducePolicy/ScanPolicy output.""" # 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 - - # 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 - + v1_v2_threshold: int = 8192 + # 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): - - 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) + + CCCL dispatch axes: (compute_capability, type_t, op_kind_t, offset_size) + Our dispatch axes: (hardware, dtype, head_dim, max_seq_len, num_kv_heads) """ - info = classify_dtype(dtype) - elem_size = info["size"] - + elem_size = dtype.itemsize if hasattr(dtype, 'itemsize') else torch.tensor([], dtype=dtype).element_size() + smem = hw.max_shared_memory_per_block + # --- 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: + # SMEM = BLOCK_N × head_dim × elem_size × 2 (K + V staging) + # Must fit in 48KB with margin for softmax accumulators + # Qwen3.6: head_dim=256, bf16 → elem_size=2 + # BLOCK_N=64: 64×256×2×2 = 64KB > 48KB → CRASH + # 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 - - # 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 - + + # BLOCK_M: same as BLOCK_N for square tiles (simplifies causal mask) + # BI-V100: 4 warps, not 8 (BLOCK=32 → 32 rows, 8 warps = 256 threads + # means only 32/256=0.125 rows/thread — wasteful) + triton_block_m = triton_block_n + 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 - 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 - + # V1 threshold: for seq_len > threshold, V2 would be better IF V2 were native C++ + # Currently V2 is PyTorch → always slower than V1 ixformer + # So threshold is effectively infinite (always V1) + v1_threshold = max_seq_len + 1 # force V1 + 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, ) -# ============================================================ -# Convenience: get config for Qwen3.6 on BI-V100 -# ============================================================ +# Pre-computed configs (mirrors CCCL's compile-time policy instantiation) +# CCCL does this via template instantiation; we do it at import time. -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, - ) +QWEN36_BF16 = select_attention_config( + hw=BI_V100, + dtype=torch.bfloat16, + head_dim=256, # CONFIRMED from qwen3_5.py: text_cfg.head_dim = 256 + max_seq_len=100000, # from computility-run.yaml: --max-model-len 100000 + num_kv_heads=4, # CONFIRMED: num_key_value_heads = 4 +) +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") - - 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() - + print(f"Qwen3.6 bf16 (head_dim=256):") + print(f" triton: BLOCK_M={QWEN36_BF16.triton_block_m} BLOCK_N={QWEN36_BF16.triton_block_n}" + f" warps={QWEN36_BF16.triton_num_warps} stages={QWEN36_BF16.triton_num_stages}") + print(f" decode: partition={QWEN36_BF16.partition_size} v1_thresh={QWEN36_BF16.v1_v2_threshold}") + print(f" SMEM: {QWEN36_BF16.triton_block_n}×256×2×2 = {QWEN36_BF16.triton_block_n*256*2*2} bytes" + f" ({QWEN36_BF16.triton_block_n*256*2*2/1024:.0f}KB ≤ 48KB)") + print(f" V1 forced: {QWEN36_BF16.use_native_v1} (V2 native has correctness issues)")