[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)
This commit is contained in:
muh-bot
356
muh_dispatch.py
356
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<PolicySelector, CC>
|
||||
- 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)")
|
||||
|
||||
Reference in New Issue
Block a user