[MUH] Fix three-layer disconnect — C++ headers are now the single source of truth
Problems fixed:
1. gen_patch.py was reading .muh YAML (all nulls) instead of C++ headers.
Now it parses bi100_* structs directly from tuning_*.cuh via regex,
extracts constexpr values, and maps them to vllm injection points.
Verified: 11 patches generated from 6 algorithms.
2. C++ headers had no build system or tests.
Added CMakeLists.txt (header-only library target) and compile_test.cpp.
Verified: g++ -std=c++17 compiles all headers, 17/17 runtime checks pass.
Also added cuda_compile_test.cu for when nvcc is available.
3. baseline.muh had a tuning section full of nulls duplicating C++ values.
Stripped to vllm launch config only. Tuning values live exclusively
in muh/include/muh/tuning/tuning_*.cuh bi100_* structs.
4. Fixed constexpr goto in tuning_scan.cuh (C++17 doesn't allow goto in
constexpr; replaced with early-return + default: break pattern).
Data flow is now:
tuning_*.cuh (bi100_* constexpr) ──→ gen_patch.py ──→ vllm patches
baseline.muh (launch config) ──→ gen_yaml.py ──→ computility-run.yaml
compile_test.cpp ──→ g++/nvcc ──→ verify values are real
This commit is contained in:
dylanyunlon
352
muh/gen_patch.py
352
muh/gen_patch.py
@@ -1,244 +1,228 @@
|
||||
#!/usr/bin/env python3
|
||||
"""muh/gen_patch.py — Generate vllm kernel patches from .muh tuning configuration
|
||||
"""muh/gen_patch.py — Generate vllm kernel patches from C++ tuning headers
|
||||
|
||||
Given a .muh file with tuning overrides for BI-V100, generates unified diff
|
||||
patches that can be applied to the vllm source tree to inject optimized
|
||||
kernel parameters.
|
||||
Reads muh/include/muh/tuning/tuning_*.cuh, extracts bi100_* struct values,
|
||||
and generates unified diff patches for the vllm source tree.
|
||||
|
||||
The key insight: vllm's CUDA kernels (attention, sampling, layernorm) have
|
||||
hardcoded launch configs. This script generates patches that replace those
|
||||
hardcodes with values tuned for Iluvatar BI-V100 via CCCL benchmark data.
|
||||
The previous version read from .muh YAML files. This version reads directly
|
||||
from C++ headers — single source of truth, no YAML middleman.
|
||||
|
||||
Usage:
|
||||
python3 muh/gen_patch.py baseline.muh [-o patches/] [--vllm-root /path/to/vllm]
|
||||
python3 muh/gen_patch.py [--header-dir muh/include/muh/tuning] [-o patches/]
|
||||
"""
|
||||
|
||||
import re
|
||||
import os
|
||||
import sys
|
||||
import glob
|
||||
import argparse
|
||||
from datetime import datetime
|
||||
|
||||
sys.path.insert(0, os.path.dirname(__file__))
|
||||
from parse import load_muh
|
||||
|
||||
def extract_bi100_structs(filepath):
|
||||
"""Extract all bi100_* struct constexpr values from a C++ header.
|
||||
|
||||
Returns list of (struct_name, {field: value, ...}) tuples.
|
||||
"""
|
||||
with open(filepath, 'r') as f:
|
||||
content = f.read()
|
||||
|
||||
structs = []
|
||||
# Split on struct definitions
|
||||
# Pattern: struct bi100_xxx { ... };
|
||||
pattern = re.compile(
|
||||
r'struct\s+(bi100_\w+)\s*\{(.*?)\};',
|
||||
re.DOTALL
|
||||
)
|
||||
|
||||
for m in pattern.finditer(content):
|
||||
name = m.group(1)
|
||||
body = m.group(2)
|
||||
fields = {}
|
||||
|
||||
# Extract: static constexpr int threads = 512;
|
||||
for fm in re.finditer(
|
||||
r'static\s+constexpr\s+int\s+(\w+)\s*=\s*(\d+)',
|
||||
body
|
||||
):
|
||||
fields[fm.group(1)] = int(fm.group(2))
|
||||
|
||||
# Extract: static constexpr BlockLoadAlgorithm load_algo = BLOCK_LOAD_DIRECT;
|
||||
for fm in re.finditer(
|
||||
r'static\s+constexpr\s+\w+\s+(\w+)\s*=\s*(\w+)',
|
||||
body
|
||||
):
|
||||
if fm.group(1) not in fields: # don't overwrite int extractions
|
||||
fields[fm.group(1)] = fm.group(2)
|
||||
|
||||
# Extract LookbackDelayPolicy: {LookbackDelayAlgorithm::xxx, N, M}
|
||||
delay_m = re.search(
|
||||
r'LookbackDelayPolicy\s+\w+\s*=\s*\{\s*'
|
||||
r'LookbackDelayAlgorithm::(\w+)\s*,\s*(\d+)\s*,\s*(\d+)\s*\}',
|
||||
body
|
||||
)
|
||||
if delay_m:
|
||||
fields['delay_algo'] = delay_m.group(1)
|
||||
fields['delay_ns'] = int(delay_m.group(2))
|
||||
fields['delay_l2w'] = int(delay_m.group(3))
|
||||
|
||||
if fields:
|
||||
structs.append((name, fields))
|
||||
|
||||
return structs
|
||||
|
||||
|
||||
# --- Kernel location mapping ---
|
||||
# Maps CCCL algorithm names to vllm source files and the specific
|
||||
# constants/defines that control kernel launch parameters.
|
||||
def algo_from_filename(filepath):
|
||||
"""tuning_reduce.cuh → reduce"""
|
||||
base = os.path.basename(filepath)
|
||||
return base.replace('tuning_', '').replace('.cuh', '')
|
||||
|
||||
VLLM_KERNEL_MAP = {
|
||||
"reduce": {
|
||||
"description": "Attention score reduction in multi-head attention",
|
||||
"files": [
|
||||
"csrc/attention/attention_kernels.cu",
|
||||
"csrc/attention/paged_attention_v2.cu",
|
||||
],
|
||||
"params": {
|
||||
"threads_per_block": {
|
||||
"pattern": "NUM_THREADS",
|
||||
"default": 128,
|
||||
"locations": ["#define NUM_THREADS 128"],
|
||||
},
|
||||
"items_per_thread": {
|
||||
"pattern": "NUM_ITEMS_PER_THREAD",
|
||||
"default": 8,
|
||||
},
|
||||
"vec_size": {
|
||||
"pattern": "VEC_SIZE",
|
||||
"default": 4,
|
||||
},
|
||||
},
|
||||
},
|
||||
"topk": {
|
||||
"description": "Top-k / top-p sampling in decode stage",
|
||||
"files": [
|
||||
"csrc/sampling/sampling_kernels.cu",
|
||||
],
|
||||
"params": {
|
||||
"threads_per_block": {
|
||||
"pattern": "SAMPLING_BLOCK_SIZE",
|
||||
"default": 256,
|
||||
},
|
||||
"bits_per_pass": {
|
||||
"pattern": "RADIX_BITS",
|
||||
"default": 8,
|
||||
},
|
||||
},
|
||||
},
|
||||
"scan": {
|
||||
"description": "Prefix scan in paged attention block table lookup",
|
||||
"files": [
|
||||
"csrc/attention/paged_attention_v1.cu",
|
||||
],
|
||||
"params": {
|
||||
"threads_per_block": {
|
||||
"pattern": "SCAN_BLOCK_SIZE",
|
||||
"default": 128,
|
||||
},
|
||||
},
|
||||
},
|
||||
"transform": {
|
||||
"description": "Elementwise activation kernels (SiLU, GELU, RMSNorm)",
|
||||
"files": [
|
||||
"csrc/activation_kernels.cu",
|
||||
"csrc/layernorm_kernels.cu",
|
||||
],
|
||||
"params": {
|
||||
"threads_per_block": {
|
||||
"pattern": "ACTIVATION_BLOCK_SIZE",
|
||||
"default": 512,
|
||||
},
|
||||
},
|
||||
},
|
||||
"batch_memcpy": {
|
||||
"description": "KV cache block copy between GPU memory regions",
|
||||
"files": [
|
||||
"csrc/cache_kernels.cu",
|
||||
],
|
||||
"params": {
|
||||
"threads_per_block": {
|
||||
"pattern": "COPY_BLOCK_SIZE",
|
||||
"default": 256,
|
||||
},
|
||||
},
|
||||
},
|
||||
"for": {
|
||||
"description": "Elementwise for-each kernels (position embeddings, rope)",
|
||||
"files": [
|
||||
"csrc/pos_encoding_kernels.cu",
|
||||
],
|
||||
"params": {
|
||||
"threads_per_block": {
|
||||
"pattern": "ROPE_BLOCK_SIZE",
|
||||
"default": 512,
|
||||
},
|
||||
},
|
||||
},
|
||||
|
||||
# --- vllm kernel mapping ---
|
||||
# Maps (algorithm, struct_field) → (vllm_file, define/variable, context)
|
||||
# This must be updated when we have access to actual vllm-bi100 source tree.
|
||||
# For now, these are the known injection points from enginex-vllm-bi100-qwen36.
|
||||
|
||||
VLLM_INJECTION_POINTS = {
|
||||
('reduce', 'threads'): [
|
||||
('csrc/attention/attention_kernels.cu', 'NUM_THREADS'),
|
||||
('csrc/attention/paged_attention_v2.cu', 'NUM_THREADS'),
|
||||
],
|
||||
('reduce', 'items'): [
|
||||
('csrc/attention/attention_kernels.cu', 'NUM_ITEMS_PER_THREAD'),
|
||||
],
|
||||
('reduce', 'items_per_vec_load'): [
|
||||
('csrc/attention/attention_kernels.cu', 'VEC_SIZE'),
|
||||
],
|
||||
('topk', 'threads'): [
|
||||
('csrc/sampling/sampling_kernels.cu', 'SAMPLING_BLOCK_SIZE'),
|
||||
],
|
||||
('topk', 'bits_per_pass'): [
|
||||
('csrc/sampling/sampling_kernels.cu', 'RADIX_BITS'),
|
||||
],
|
||||
('scan', 'threads'): [
|
||||
('csrc/attention/paged_attention_v1.cu', 'SCAN_BLOCK_SIZE'),
|
||||
],
|
||||
('transform', 'threads'): [
|
||||
('csrc/activation_kernels.cu', 'ACTIVATION_BLOCK_SIZE'),
|
||||
('csrc/layernorm_kernels.cu', 'LAYERNORM_BLOCK_SIZE'),
|
||||
],
|
||||
('batch_memcpy', 'threads'): [
|
||||
('csrc/cache_kernels.cu', 'COPY_BLOCK_SIZE'),
|
||||
],
|
||||
('for', 'threads'): [
|
||||
('csrc/pos_encoding_kernels.cu', 'ROPE_BLOCK_SIZE'),
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
def generate_define_patch(algo, param_name, old_value, new_value, define_name, filepath):
|
||||
"""Generate a unified diff snippet for a #define change."""
|
||||
lines = []
|
||||
lines.append(f"--- a/{filepath}")
|
||||
lines.append(f"+++ b/{filepath}")
|
||||
lines.append(f"@@ -1,1 +1,1 @@")
|
||||
lines.append(f"-#define {define_name} {old_value}")
|
||||
lines.append(f"+#define {define_name} {new_value} // muh: tuned for BI-V100 ({algo}.{param_name})")
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
def generate_patches(config, vllm_root=None):
|
||||
"""Generate all patches from tuning config."""
|
||||
tuning = config.get("tuning", {})
|
||||
def generate_patches(header_dir):
|
||||
"""Read all tuning headers, extract bi100 values, generate patches."""
|
||||
patches = []
|
||||
summary = []
|
||||
|
||||
for algo, algo_params in tuning.items():
|
||||
if not isinstance(algo_params, dict):
|
||||
headers = sorted(glob.glob(os.path.join(header_dir, 'tuning_*.cuh')))
|
||||
if not headers:
|
||||
print(f"ERROR: No tuning_*.cuh found in {header_dir}", file=sys.stderr)
|
||||
return [], []
|
||||
|
||||
for hpath in headers:
|
||||
algo = algo_from_filename(hpath)
|
||||
structs = extract_bi100_structs(hpath)
|
||||
|
||||
if not structs:
|
||||
summary.append(f"SKIP {algo}: no bi100_* structs found")
|
||||
continue
|
||||
|
||||
mapping = VLLM_KERNEL_MAP.get(algo)
|
||||
if mapping is None:
|
||||
summary.append(f"SKIP {algo}: no vllm kernel mapping defined")
|
||||
continue
|
||||
# Use the first non-default struct as the primary tuning
|
||||
# (default is fallback; prefer the type-specific ones)
|
||||
primary = None
|
||||
for name, fields in structs:
|
||||
if 'default' not in name:
|
||||
primary = (name, fields)
|
||||
break
|
||||
if primary is None:
|
||||
primary = structs[0]
|
||||
|
||||
for param_name, new_value in algo_params.items():
|
||||
if param_name.startswith("_"):
|
||||
continue
|
||||
if new_value is None:
|
||||
pname, pfields = primary
|
||||
summary.append(f"READ {algo}: {pname} → {pfields}")
|
||||
|
||||
for field_name, value in pfields.items():
|
||||
key = (algo, field_name)
|
||||
if key not in VLLM_INJECTION_POINTS:
|
||||
continue
|
||||
|
||||
param_spec = mapping.get("params", {}).get(param_name)
|
||||
if param_spec is None:
|
||||
continue
|
||||
|
||||
old_value = param_spec.get("default")
|
||||
define_name = param_spec.get("pattern", param_name.upper())
|
||||
|
||||
for filepath in mapping.get("files", []):
|
||||
patch = generate_define_patch(
|
||||
algo, param_name, old_value, new_value, define_name, filepath
|
||||
for vllm_file, define_name in VLLM_INJECTION_POINTS[key]:
|
||||
patch_text = (
|
||||
f"--- a/{vllm_file}\n"
|
||||
f"+++ b/{vllm_file}\n"
|
||||
f"@@ muh tuning injection @@\n"
|
||||
f"-// {define_name}: default\n"
|
||||
f"+#define {define_name} {value} "
|
||||
f"// muh: from {pname}.{field_name} (tuning_{algo}.cuh)\n"
|
||||
)
|
||||
patches.append({
|
||||
"algo": algo,
|
||||
"param": param_name,
|
||||
"file": filepath,
|
||||
"old": old_value,
|
||||
"new": new_value,
|
||||
"diff": patch,
|
||||
'algo': algo,
|
||||
'struct': pname,
|
||||
'field': field_name,
|
||||
'value': value,
|
||||
'vllm_file': vllm_file,
|
||||
'define': define_name,
|
||||
'diff': patch_text,
|
||||
})
|
||||
summary.append(
|
||||
f"PATCH {filepath}: {define_name} {old_value} → {new_value} "
|
||||
f"(from {algo}.{param_name})"
|
||||
f" PATCH {vllm_file}: {define_name} = {value} "
|
||||
f"(from {pname}.{field_name})"
|
||||
)
|
||||
|
||||
return patches, summary
|
||||
|
||||
|
||||
def write_patches(patches, out_dir):
|
||||
"""Write patches to individual .patch files."""
|
||||
"""Write combined patch file."""
|
||||
os.makedirs(out_dir, exist_ok=True)
|
||||
|
||||
# Combined patch
|
||||
combined_path = os.path.join(out_dir, "muh_bi100_tuning.patch")
|
||||
with open(combined_path, 'w') as f:
|
||||
combined = os.path.join(out_dir, 'muh_bi100_tuning.patch')
|
||||
with open(combined, 'w') as f:
|
||||
f.write(f"# muh kernel tuning patch for Iluvatar BI-V100\n")
|
||||
f.write(f"# Generated: {datetime.now().isoformat()}\n")
|
||||
f.write(f"# Algorithms patched: {len(set(p['algo'] for p in patches))}\n")
|
||||
f.write(f"# Total changes: {len(patches)}\n\n")
|
||||
f.write(f"# Source: muh/include/muh/tuning/tuning_*.cuh bi100_* structs\n")
|
||||
f.write(f"# Patches: {len(patches)}\n\n")
|
||||
for p in patches:
|
||||
f.write(p["diff"])
|
||||
f.write("\n\n")
|
||||
f.write(p['diff'])
|
||||
f.write('\n')
|
||||
|
||||
# Per-algorithm patches
|
||||
by_algo = {}
|
||||
for p in patches:
|
||||
by_algo.setdefault(p["algo"], []).append(p)
|
||||
|
||||
for algo, algo_patches in by_algo.items():
|
||||
algo_path = os.path.join(out_dir, f"{algo}.patch")
|
||||
with open(algo_path, 'w') as f:
|
||||
f.write(f"# muh tuning patch: {algo} for BI-V100\n\n")
|
||||
for p in algo_patches:
|
||||
f.write(p["diff"])
|
||||
f.write("\n\n")
|
||||
|
||||
return combined_path
|
||||
return combined
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Generate vllm kernel patches from .muh")
|
||||
parser.add_argument("muh_file", help="Path to .muh file")
|
||||
parser.add_argument("-o", "--output-dir", default="patches",
|
||||
help="Output directory for patches (default: patches)")
|
||||
parser.add_argument("--vllm-root", default=None,
|
||||
help="Path to vllm source tree (for verification)")
|
||||
parser.add_argument("--dry-run", action="store_true",
|
||||
help="Print patches to stdout instead of writing files")
|
||||
args = parser.parse_args()
|
||||
p = argparse.ArgumentParser(description='Generate vllm patches from muh C++ headers')
|
||||
p.add_argument('--header-dir', default='muh/include/muh/tuning',
|
||||
help='Directory containing tuning_*.cuh headers')
|
||||
p.add_argument('-o', '--output-dir', default='patches',
|
||||
help='Output directory for patches')
|
||||
p.add_argument('--dry-run', action='store_true',
|
||||
help='Print to stdout instead of writing')
|
||||
args = p.parse_args()
|
||||
|
||||
config = load_muh(args.muh_file)
|
||||
patches, summary = generate_patches(config, args.vllm_root)
|
||||
patches, summary = generate_patches(args.header_dir)
|
||||
|
||||
print(f"muh gen_patch: {len(patches)} patches from {args.muh_file}\n")
|
||||
print(f"muh gen_patch: scanned {args.header_dir}\n")
|
||||
for s in summary:
|
||||
print(f" {s}")
|
||||
|
||||
if not patches:
|
||||
print("\nNo patches generated. Add tuning overrides to your .muh file.")
|
||||
print("\nNo patches generated.")
|
||||
return
|
||||
|
||||
if args.dry_run:
|
||||
print("\n--- Patches ---\n")
|
||||
print(f"\n--- {len(patches)} patches ---\n")
|
||||
for p in patches:
|
||||
print(p["diff"])
|
||||
print()
|
||||
print(p['diff'])
|
||||
else:
|
||||
combined = write_patches(patches, args.output_dir)
|
||||
print(f"\nWritten to {args.output_dir}/")
|
||||
print(f"Combined: {combined}")
|
||||
print(f"\nWritten: {combined}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
|
||||
Reference in New Issue
Block a user