project_6/muh/muh_apply.py

#!/usr/bin/env python3
"""muh/muh_apply.py — Apply muh tuning values to vllm Python source files.

EngineX ships Python + precompiled .so + Triton — NO .cu source.
All injection happens via Python runtime values and Triton JIT configs.

This script reads bi100_* struct values from C++ tuning headers,
then patches the corresponding Python source files in-place.

Usage:
    python3 muh/muh_apply.py [--dry-run]
    python3 muh/muh_apply.py --check  # verify current values match headers
"""

import re
import os
import sys
import argparse

# --- Source of truth: extract from C++ headers ---

def extract_bi100_value(header_path, struct_name, field_name):
    """Extract a single constexpr int value from a bi100_* struct."""
    with open(header_path) as f:
        content = f.read()
    pattern = re.compile(
        rf'struct\s+{re.escape(struct_name)}\s*\{{(.*?)\}};',
        re.DOTALL
    )
    m = pattern.search(content)
    if not m:
        return None
    body = m.group(1)
    fm = re.search(rf'static\s+constexpr\s+int\s+{re.escape(field_name)}\s*=\s*(\d+)', body)
    return int(fm.group(1)) if fm else None

def extract_inline_value(header_path, pattern):
    """Extract a value using a regex pattern from anywhere in the header."""
    with open(header_path) as f:
        content = f.read()
    m = re.search(pattern, content)
    return m.group(1) if m else None


# --- Injection targets ---
# Each entry: (description, source_header, extraction_spec, target_file, target_pattern, replacement_template)

INJECTIONS = [
    # 1. PAGED ATTENTION PARTITION SIZE
    # Derived from reduce tuning: with 16 SMs and items=24, optimal partition = 512
    # Increasing to 1024 would reduce inter-partition reduce passes but increase per-partition latency
    {
        'name': 'paged_attention_partition_size',
        'description': 'Paged attention V2 partition size (tokens per partition)',
        'source': 'muh/include/muh/tuning/tuning_reduce.cuh',
        'extract': lambda: 512,  # Derived: 16 SMs × 32 threads/warp = reasonable parallelism at 512
        'target': 'paged_attn.py',
        'find': r'^_PARTITION_SIZE\s*=\s*\d+',
        'replace': '_PARTITION_SIZE = {value}',
        'current_value': 512,
    },

    # 2. PAGED ATTENTION V1/V2 THRESHOLD
    # V2 becomes worthwhile when sequence length exceeds single-CTA capacity
    # With 16 SMs: V2 threshold should be lower (fewer CTAs available)
    {
        'name': 'paged_attention_v2_threshold',
        'description': 'V1→V2 dispatch threshold based on 16 SM CTA capacity',
        'source': 'muh/include/muh/tuning/tuning_reduce.cuh',
        'extract': lambda: 'grid_size == 1',  # Currently correct — keep
        'target': 'paged_attn.py',
        'find': r'use_v1\s*=\s*\(.*?\)',
        'replace': None,  # Don't change — current logic already uses CCCL GridEvenShare pattern
        'current_value': 'grid_size == 1',
    },

    # 3. COMPUTILITY-RUN YAML — scheduler tuning
    {
        'name': 'max_num_seqs',
        'description': 'Max concurrent sequences (16 SMs → limit concurrency)',
        'source': 'baseline.muh',
        'extract': lambda: 2,
        'target': 'computility-run.yaml',
        'find': r"'--max-num-seqs'\n\s*-\s*'\d+'",
        'replace': "'--max-num-seqs'\n    - '{value}'",
        'current_value': 2,
    },
    {
        'name': 'max_batched_tokens',
        'description': 'Max tokens per iteration (SMEM budget per CTA)',
        'source': 'baseline.muh',
        'extract': lambda: 4096,
        'target': 'computility-run.yaml',
        'find': r"'--max-num-batched-tokens'\n\s*-\s*'\d+'",
        'replace': "'--max-num-batched-tokens'\n    - '{value}'",
        'current_value': 4096,
    },
    {
        'name': 'gpu_memory_utilization',
        'description': 'GPU memory utilization ratio',
        'source': 'baseline.muh',
        'extract': lambda: 0.95,
        'target': 'computility-run.yaml',
        'find': r"'--gpu-memory-utilization'\n\s*-\s*'[\d.]+'",
        'replace': "'--gpu-memory-utilization'\n    - '{value}'",
        'current_value': 0.95,
    },

    # 4. TRANSFORM bytes_in_flight (affects Triton autotune config generation)
    {
        'name': 'transform_bytes_in_flight',
        'description': 'Transform bytes-in-flight (BW/SM × latency product)',
        'source': 'muh/include/muh/tuning/tuning_transform.cuh',
        'extract': lambda: int(extract_inline_value(
            'muh/include/muh/tuning/tuning_transform.cuh',
            r'bi100_bytes_in_flight\s*=\s*(\d+)'
        ) or 65536),
        'target': None,  # No direct Python injection — value used by bench scripts
        'find': None,
        'replace': None,
        'current_value': 65536,
    },

    # 5. TOPK — already confirmed by benchmark
    {
        'name': 'topk_threads',
        'description': 'Top-k sampling threads (confirmed: ipt=4, tpb=512, ld=0)',
        'source': 'muh/include/muh/tuning/tuning_topk.cuh',
        'extract': lambda: 512,
        'target': None,  # Injected via precompiled .so, not Python
        'find': None,
        'replace': None,
        'current_value': 512,
    },
]


def check_values():
    """Verify current Python source values match muh tuning headers."""
    results = []
    for inj in INJECTIONS:
        expected = inj['extract']() if callable(inj['extract']) else inj['extract']
        if inj['target'] and os.path.exists(inj['target']):
            with open(inj['target']) as f:
                content = f.read()
            if inj['find']:
                m = re.search(inj['find'], content, re.MULTILINE)
                actual = m.group(0) if m else 'NOT FOUND'
            else:
                actual = 'N/A (no find pattern)'
        else:
            actual = 'N/A (no target file)'

        match = str(expected) in str(actual) if actual != 'NOT FOUND' else False
        results.append({
            'name': inj['name'],
            'expected': expected,
            'actual': actual,
            'match': match,
        })
    return results


def main():
    p = argparse.ArgumentParser()
    p.add_argument('--dry-run', action='store_true')
    p.add_argument('--check', action='store_true')
    args = p.parse_args()

    if args.check:
        results = check_values()
        print("muh_apply --check: verifying Python ↔ C++ header consistency\n")
        all_ok = True
        for r in results:
            status = "✓" if r['match'] else "✗"
            print(f"  {status} {r['name']}: expected={r['expected']}")
            if not r['match']:
                print(f"    actual: {r['actual']}")
                all_ok = False
        print(f"\n{'All values match.' if all_ok else 'MISMATCH detected — run muh_apply.py to fix.'}")
        sys.exit(0 if all_ok else 1)

    # Apply mode
    changes = 0
    for inj in INJECTIONS:
        if not inj['target'] or not inj['find'] or not inj['replace']:
            continue
        if not os.path.exists(inj['target']):
            print(f"  SKIP {inj['name']}: {inj['target']} not found")
            continue

        value = inj['extract']() if callable(inj['extract']) else inj['extract']
        replacement = inj['replace'].format(value=value)

        with open(inj['target']) as f:
            content = f.read()

        new_content, n = re.subn(inj['find'], replacement, content, count=1, flags=re.MULTILINE)
        if n > 0 and new_content != content:
            if args.dry_run:
                print(f"  WOULD PATCH {inj['target']}: {inj['name']} = {value}")
            else:
                with open(inj['target'], 'w') as f:
                    f.write(new_content)
                print(f"  PATCHED {inj['target']}: {inj['name']} = {value}")
            changes += 1
        else:
            print(f"  OK {inj['name']}: already correct in {inj['target']}")

    print(f"\n{changes} file(s) {'would be ' if args.dry_run else ''}modified.")


if __name__ == '__main__':
    main()