9b21a13119
Pipeline: 1. extract.py: Parses all 26 CCCL tuning_*.cuh → 26 YAML schemas in muh/schema/ 2. parse.py: .muh file parser with extends-inheritance + schema validation 3. gen_yaml.py: .muh → computility-run.yaml (verified: matches competition reference) 4. gen_patch.py: .muh → vllm kernel unified diff patches (6 algorithm mappings) 5. baseline.muh: Competition reference config, all tuning values pending BI-V100 benchmarks Schemas extracted: 26 algorithms, 8-19 params each, SM75/80/90/100 reference tunings Priority mapping: reduce→attention, topk→sampling, scan→paged_attention, transform→activations, batch_memcpy→KV_cache, for→RoPE Tested: extract→parse→validate→gen_yaml→gen_patch full pipeline passes
286 lines
9.6 KiB
Python
286 lines
9.6 KiB
Python
#!/usr/bin/env python3
|
|
"""muh/extract.py — Extract CCCL tuning parameter spaces into muh/schema/*.yaml
|
|
|
|
Reads cccl_upstream/cub/cub/device/dispatch/tuning/tuning_*.cuh,
|
|
parses policy struct fields and SM-specific tuning values,
|
|
outputs one YAML file per algorithm under muh/schema/.
|
|
|
|
Usage:
|
|
python3 muh/extract.py [--cccl-root cccl_upstream] [--out-dir muh/schema]
|
|
"""
|
|
|
|
import re
|
|
import os
|
|
import sys
|
|
import glob
|
|
import argparse
|
|
from pathlib import Path
|
|
from collections import OrderedDict
|
|
|
|
# --- Enum value sets (from CCCL headers) ---
|
|
|
|
BLOCK_LOAD_ALGORITHMS = [
|
|
"BLOCK_LOAD_DIRECT",
|
|
"BLOCK_LOAD_VECTORIZE",
|
|
"BLOCK_LOAD_TRANSPOSE",
|
|
"BLOCK_LOAD_WARP_TRANSPOSE",
|
|
"BLOCK_LOAD_WARP_TRANSPOSE_TIMESLICED",
|
|
"BLOCK_LOAD_STRIPED",
|
|
]
|
|
|
|
BLOCK_STORE_ALGORITHMS = [
|
|
"BLOCK_STORE_DIRECT",
|
|
"BLOCK_STORE_WARP_TRANSPOSE",
|
|
"BLOCK_STORE_WARP_TRANSPOSE_TIMESLICED",
|
|
"BLOCK_STORE_STRIPED",
|
|
]
|
|
|
|
BLOCK_REDUCE_ALGORITHMS = [
|
|
"BLOCK_REDUCE_RAKING",
|
|
"BLOCK_REDUCE_RAKING_COMMUTATIVE_ONLY",
|
|
"BLOCK_REDUCE_WARP_REDUCTIONS",
|
|
]
|
|
|
|
BLOCK_SCAN_ALGORITHMS = [
|
|
"BLOCK_SCAN_RAKING",
|
|
"BLOCK_SCAN_RAKING_MEMOIZE",
|
|
"BLOCK_SCAN_WARP_SCANS",
|
|
]
|
|
|
|
CACHE_LOAD_MODIFIERS = [
|
|
"LOAD_DEFAULT",
|
|
"LOAD_CA",
|
|
"LOAD_CG",
|
|
"LOAD_CS",
|
|
"LOAD_CV",
|
|
"LOAD_LDG",
|
|
]
|
|
|
|
LOOKBACK_DELAY_ALGORITHMS = [
|
|
"no_delay",
|
|
"fixed_delay",
|
|
"exponential_backoff",
|
|
"exponential_backoff_jitter",
|
|
"exponential_backoff_jitter_window",
|
|
"exponential_backon_jitter_window",
|
|
"exponential_backon_jitter",
|
|
"exponential_backon",
|
|
]
|
|
|
|
# --- Field type → range/enum mapping ---
|
|
|
|
FIELD_TYPES = {
|
|
"threads_per_block": {"type": "int", "range": [32, 1024], "step": 32},
|
|
"items_per_thread": {"type": "int", "range": [1, 32], "step": 1},
|
|
"vec_size": {"type": "int", "range": [1, 8], "step": 1},
|
|
"bits_per_pass": {"type": "int", "range": [4, 11], "step": 1},
|
|
"radix_bits": {"type": "int", "range": [4, 8], "step": 1},
|
|
"load_algorithm": {"type": "enum", "values": BLOCK_LOAD_ALGORITHMS},
|
|
"store_algorithm": {"type": "enum", "values": BLOCK_STORE_ALGORITHMS},
|
|
"reduce_algorithm": {"type": "enum", "values": BLOCK_REDUCE_ALGORITHMS},
|
|
"scan_algorithm": {"type": "enum", "values": BLOCK_SCAN_ALGORITHMS},
|
|
"load_modifier": {"type": "enum", "values": CACHE_LOAD_MODIFIERS},
|
|
"lookback_delay.kind": {"type": "enum", "values": LOOKBACK_DELAY_ALGORITHMS},
|
|
"lookback_delay.delay": {"type": "int", "range": [0, 2000], "step": 50},
|
|
"lookback_delay.l2_write_latency": {"type": "int", "range": [0, 2000], "step": 50},
|
|
"reduce_and_scan_warps": {"type": "int", "range": [1, 8], "step": 1},
|
|
"lookahead_items_per_thread": {"type": "int", "range": [1, 16], "step": 1},
|
|
}
|
|
|
|
|
|
def extract_policy_fields(content, filename):
|
|
"""Extract policy struct field names from a tuning file."""
|
|
fields = []
|
|
# Match lines like: int threads_per_block; or BlockLoadAlgorithm load_algorithm;
|
|
pattern = re.compile(
|
|
r'^\s+(?:int|BlockLoadAlgorithm|BlockStoreAlgorithm|BlockReduceAlgorithm|'
|
|
r'BlockScanAlgorithm|CacheLoadModifier|LookbackDelayPolicy)\s+'
|
|
r'(\w+)\s*[;=]',
|
|
re.MULTILINE
|
|
)
|
|
for m in pattern.finditer(content):
|
|
field = m.group(1)
|
|
if field not in fields:
|
|
fields.append(field)
|
|
return fields
|
|
|
|
|
|
def extract_sm_tunings(content):
|
|
"""Extract SM-specific tuning values from static constexpr definitions."""
|
|
tunings = {}
|
|
# Match patterns like: sm80_tuning, sm90_tuning, sm100_tuning
|
|
sm_pattern = re.compile(r'struct\s+sm(\d+)_tuning')
|
|
for m in sm_pattern.finditer(content):
|
|
sm = int(m.group(1))
|
|
if sm not in tunings:
|
|
tunings[sm] = []
|
|
|
|
# Extract actual parameter values from constexpr definitions
|
|
# Pattern: static constexpr int threads = 512;
|
|
blocks = re.split(r'(?=struct\s+sm\d+_tuning)', content)
|
|
for block in blocks:
|
|
sm_m = re.match(r'struct\s+sm(\d+)_tuning', block)
|
|
if not sm_m:
|
|
continue
|
|
sm = int(sm_m.group(1))
|
|
vals = {}
|
|
for line in block.split('\n'):
|
|
# int values
|
|
m = re.search(r'static\s+constexpr\s+int\s+(\w+)\s*=\s*(\d+)', line)
|
|
if m:
|
|
vals[m.group(1)] = int(m.group(2))
|
|
# enum values
|
|
m = re.search(r'static\s+constexpr\s+(?:BlockLoadAlgorithm|BlockStoreAlgorithm|CacheLoadModifier)\s+(\w+)\s*=\s*(\w+)', line)
|
|
if m:
|
|
vals[m.group(1)] = m.group(2)
|
|
if vals:
|
|
tunings.setdefault(sm, []).append(vals)
|
|
|
|
return tunings
|
|
|
|
|
|
def extract_inline_tunings(content):
|
|
"""Extract inline tuning values from make_mem_scaled_lookback_scan_policy calls and similar."""
|
|
inline = []
|
|
# Pattern: threads_per_block, items_per_thread in constructor-style calls
|
|
pattern = re.compile(
|
|
r'(?:topk_policy|ReducePassPolicy|ScanLookbackPolicy)\s*\{'
|
|
r'\s*(\d+)\s*,\s*(\d+)',
|
|
re.MULTILINE
|
|
)
|
|
for m in pattern.finditer(content):
|
|
inline.append({
|
|
"threads_per_block": int(m.group(1)),
|
|
"items_per_thread": int(m.group(2)),
|
|
})
|
|
return inline
|
|
|
|
|
|
def algo_name_from_filename(filename):
|
|
"""tuning_topk.cuh → topk"""
|
|
base = os.path.basename(filename)
|
|
return base.replace("tuning_", "").replace(".cuh", "")
|
|
|
|
|
|
def build_schema(algo, fields, sm_tunings, inline_tunings):
|
|
"""Build YAML-serializable schema dict for one algorithm."""
|
|
schema = OrderedDict()
|
|
schema["algorithm"] = algo
|
|
schema["source"] = f"cub/cub/device/dispatch/tuning/tuning_{algo}.cuh"
|
|
|
|
# Parameter space
|
|
params = OrderedDict()
|
|
for field in fields:
|
|
if field in FIELD_TYPES:
|
|
params[field] = dict(FIELD_TYPES[field])
|
|
elif field == "lookback_delay":
|
|
# Expand to sub-fields
|
|
for sub in ["lookback_delay.kind", "lookback_delay.delay", "lookback_delay.l2_write_latency"]:
|
|
params[sub] = dict(FIELD_TYPES[sub])
|
|
else:
|
|
params[field] = {"type": "int", "range": [1, 1024], "step": 1, "note": "unknown_range"}
|
|
schema["parameters"] = dict(params)
|
|
|
|
# Known SM tunings (for reference when tuning BI-V100)
|
|
if sm_tunings:
|
|
ref = OrderedDict()
|
|
for sm, vals_list in sorted(sm_tunings.items()):
|
|
ref[f"sm{sm}"] = vals_list
|
|
schema["reference_tunings"] = dict(ref)
|
|
|
|
# BI-V100 placeholder
|
|
schema["bi_v100"] = {
|
|
"status": "pending_benchmark",
|
|
"note": "Run muh benchmark on Iluvatar BI-V100 to fill these values",
|
|
"threads_per_block": "TBD",
|
|
"items_per_thread": "TBD",
|
|
}
|
|
|
|
return schema
|
|
|
|
|
|
def yaml_dump(data, indent=0):
|
|
"""Simple YAML serializer (no dependency on pyyaml)."""
|
|
lines = []
|
|
prefix = " " * indent
|
|
if isinstance(data, dict):
|
|
for k, v in data.items():
|
|
if isinstance(v, (dict, list)):
|
|
lines.append(f"{prefix}{k}:")
|
|
lines.append(yaml_dump(v, indent + 1))
|
|
else:
|
|
lines.append(f"{prefix}{k}: {v}")
|
|
elif isinstance(data, list):
|
|
for item in data:
|
|
if isinstance(item, dict):
|
|
lines.append(f"{prefix}-")
|
|
lines.append(yaml_dump(item, indent + 1))
|
|
else:
|
|
lines.append(f"{prefix}- {item}")
|
|
else:
|
|
lines.append(f"{prefix}{data}")
|
|
return "\n".join(lines)
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description="Extract CCCL tuning params to muh schema")
|
|
parser.add_argument("--cccl-root", default="cccl_upstream",
|
|
help="Path to CCCL root (default: cccl_upstream)")
|
|
parser.add_argument("--out-dir", default="muh/schema",
|
|
help="Output directory for YAML schemas (default: muh/schema)")
|
|
args = parser.parse_args()
|
|
|
|
tuning_dir = os.path.join(args.cccl_root, "cub", "cub", "device", "dispatch", "tuning")
|
|
out_dir = args.out_dir
|
|
os.makedirs(out_dir, exist_ok=True)
|
|
|
|
tuning_files = sorted(glob.glob(os.path.join(tuning_dir, "tuning_*.cuh")))
|
|
if not tuning_files:
|
|
print(f"ERROR: No tuning_*.cuh files found in {tuning_dir}", file=sys.stderr)
|
|
sys.exit(1)
|
|
|
|
print(f"Found {len(tuning_files)} tuning files in {tuning_dir}")
|
|
|
|
all_algos = []
|
|
for filepath in tuning_files:
|
|
if os.path.basename(filepath) == "common.cuh":
|
|
continue
|
|
|
|
with open(filepath, "r") as f:
|
|
content = f.read()
|
|
|
|
algo = algo_name_from_filename(filepath)
|
|
fields = extract_policy_fields(content, filepath)
|
|
sm_tunings = extract_sm_tunings(content)
|
|
inline_tunings = extract_inline_tunings(content)
|
|
|
|
schema = build_schema(algo, fields, sm_tunings, inline_tunings)
|
|
|
|
out_path = os.path.join(out_dir, f"{algo}.yaml")
|
|
with open(out_path, "w") as f:
|
|
f.write(f"# muh schema for {algo}\n")
|
|
f.write(f"# Auto-extracted from {schema['source']}\n")
|
|
f.write(f"# Generated by muh/extract.py\n\n")
|
|
f.write(yaml_dump(dict(schema)))
|
|
f.write("\n")
|
|
|
|
all_algos.append(algo)
|
|
print(f" {algo}: {len(fields)} params, {len(sm_tunings)} SM tunings → {out_path}")
|
|
|
|
# Write index
|
|
index_path = os.path.join(out_dir, "_index.yaml")
|
|
with open(index_path, "w") as f:
|
|
f.write("# muh schema index — all extracted CCCL tuning algorithms\n\n")
|
|
f.write("algorithms:\n")
|
|
for algo in all_algos:
|
|
f.write(f" - {algo}\n")
|
|
f.write(f"\ntotal: {len(all_algos)}\n")
|
|
f.write(f"source: cccl_upstream/cub/cub/device/dispatch/tuning/\n")
|
|
|
|
print(f"\nDone: {len(all_algos)} schemas → {out_dir}/")
|
|
print(f"Index: {index_path}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|