project_6/muh/gen_patch.py

#!/usr/bin/env python3
"""muh/gen_patch.py — Generate vllm kernel patches from C++ tuning headers

Reads muh/include/muh/tuning/tuning_*.cuh, extracts bi100_* struct values,
and generates unified diff patches for the vllm source tree.

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 [--header-dir muh/include/muh/tuning] [-o patches/]
"""

import re
import os
import sys
import glob
import argparse
from datetime import datetime


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


def algo_from_filename(filepath):
    """tuning_reduce.cuh → reduce"""
    base = os.path.basename(filepath)
    return base.replace('tuning_', '').replace('.cuh', '')


# --- 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 = {
    # ═══════════════════════════════════════════════════════════════════
    # enginex ships Python + precompiled .so + Triton — NO .cu source.
    # All injection is via Python runtime values and Triton JIT configs.
    #
    # DEAD (csrc/*.cu) paths preserved as comments for when/if EngineX
    # exposes CUDA source in future releases.
    # ═══════════════════════════════════════════════════════════════════

    # ─── 1. PAGED ATTENTION DECODE (Output TPS × 16.796 = 83%) ─────
    # paged_attn.py: controls V1/V2 dispatch and partition granularity.
    # CCCL parallel: compound reduce (summary_statistics.cu Welford pattern).
    # The ixformer .so has NUM_THREADS baked in — we control PARTITION_SIZE
    # and V1/V2 threshold from Python, which determines how many CTAs launch.
    # _PARTITION_SIZE = number of KV tokens per partition in V2.
    # NOT the same as items_per_thread. Currently hardcoded 512 in paged_attn.py.
    # Tuning: larger partition → fewer inter-partition reduce passes (good for 16 SMs).
    # Smaller partition → more parallelism across CTAs (good for many SMs).
    # BI-V100 with 16 SMs: partition=512 is a reasonable balance.
    # To change, must also update max_num_partitions calculation.
    ('reduce', 'partition_size'): [
        ('paged_attn.py', '_PARTITION_SIZE'),
    ],

    # ─── 2. TRITON PREFILL ATTENTION ───────────────────────────────
    # prefix_prefill.py: Triton JIT kernel for context (prefill) attention.
    # CCCL parallel: scan + reduce + transform (softmax + QKV matmul).
    # SMEM constraint: BLOCK_N × head_dim × elem_size × 2 ≤ 48KB.
    # Qwen3.6 head_dim=256, bf16: BLOCK_N=32 → 32KB ✓, BLOCK_N=64 → 64KB ✗
    ('prefill', 'BLOCK_M'): [
        ('prefix_prefill.py', 'BLOCK'),
    ],
    ('prefill', 'NUM_WARPS'): [
        ('prefix_prefill.py', 'NUM_WARPS'),
    ],

    # ─── 3. TRITON FLASH ATTENTION (autotune) ─────────────────────
    # triton_flash_attention.py: @triton.autotune with 20+ Config entries.
    # We added BI-V100 specific configs (BLOCK_M=32/64, num_stages=2,
    # num_warps=2/4) based on CCCL transform benchmark bytes_in_flight=64KB.
    # Autotune picks the fastest at runtime — our configs compete fairly.
    ('flash_attn', 'BLOCK_M'): [
        ('vllm/attention/ops/triton_flash_attention.py', 'BLOCK_M'),
    ],
    ('flash_attn', 'BLOCK_N'): [
        ('vllm/attention/ops/triton_flash_attention.py', 'BLOCK_N'),
    ],

    # ─── 4. MoE ROUTING (Qwen3.6 is MoE: 256 experts, top-8) ────
    # fused_moe.py: GEMM tiling for expert-parallel matmul.
    # CCCL parallel: batch_memcpy (expert weight scatter) + transform (gate).
    ('moe', 'BLOCK_SIZE_M'): [
        ('vllm/model_executor/layers/fused_moe/fused_moe.py', 'BLOCK_SIZE_M'),
    ],

    # ─── 5. RUNTIME HARDWARE OVERRIDES ────────────────────────────
    # _custom_ops.py: BI-V100 SMEM was hardcoded 32KB → fixed to 48KB.
    # This unblocks all Triton kernels that tile by SMEM availability.
    ('runtime', 'SMEM'): [
        ('vllm/_custom_ops.py', 'get_max_shared_memory'),
    ],

    # ─── 6. LAUNCH CONFIGURATION (computility-run.yaml) ──────────
    # Server-level tuning: max-model-len, gpu-memory-utilization, tp,
    # max-num-seqs, batched-tokens, chunked-prefill, prefix-caching.
    # CCCL parallel: these control the problem size fed to all kernels.
    ('scheduler', 'num_steps'): [
        ('computility-run.yaml', 'num-scheduler-steps'),
    ],
    ('scheduler', 'max_num_seqs'): [
        ('computility-run.yaml', '--max-num-seqs'),
    ],
    ('scheduler', 'max_batched_tokens'): [
        ('computility-run.yaml', '--max-num-batched-tokens'),
    ],
    ('scheduler', 'gpu_mem_util'): [
        ('computility-run.yaml', '--gpu-memory-utilization'),
    ],

    # ─── 7. PAGED ATTENTION V2 ENABLE (currently force-disabled) ──
    # paged_attn.py line ~99: use_v1 = True disables V2 for all seq_lens.
    # V2 partitions long sequences across CTAs (CCCL GridEvenShare pattern).
    # For seq_len > 8K, V2 should be faster — but needs native C++ impl,
    # not the PyTorch fallback currently in paged_attention_v2_pytorch.py.
    ('reduce', 'v1_v2_threshold'): [
        ('paged_attn.py', 'use_v1'),
    ],

    # ═══════════════════════════════════════════════════════════════════
    # DEAD csrc/*.cu injection points (no .cu source in enginex):
    # ('reduce', 'threads'): [('csrc/attention/attention_kernels.cu', 'NUM_THREADS')],
    # ('topk', 'threads'): [('csrc/sampling/sampling_kernels.cu', 'SAMPLING_BLOCK_SIZE')],
    # ('scan', 'threads'): [('csrc/attention/paged_attention_v1.cu', 'SCAN_BLOCK_SIZE')],
    # ('transform', 'threads'): [('csrc/activation_kernels.cu', 'ACTIVATION_BLOCK_SIZE')],
    # ('batch_memcpy', 'threads'): [('csrc/cache_kernels.cu', 'COPY_BLOCK_SIZE')],
    # ('for', 'threads'): [('csrc/pos_encoding_kernels.cu', 'ROPE_BLOCK_SIZE')],
    # ═══════════════════════════════════════════════════════════════════
}


# --- Complete tuning algorithm registry ---
# All 26 algorithms with muh tuning headers.
# 'injection': algorithms with known vllm kernel injection points
# 'library': algorithms used via CCCL library calls (no direct vllm injection)
# 'struct_mode': 'named' = has bi100_* structs, 'inline' = computes in policy_selector

TUNING_REGISTRY = {
    # === 6 algorithms with vllm injection points (struct_mode='named') ===
    'reduce':       {'mode': 'injection', 'struct_mode': 'named',  'vllm_files': ['csrc/attention/attention_kernels.cu', 'csrc/attention/paged_attention_v2.cu']},
    'scan':         {'mode': 'injection', 'struct_mode': 'named',  'vllm_files': ['csrc/attention/paged_attention_v1.cu']},
    'topk':         {'mode': 'injection', 'struct_mode': 'named',  'vllm_files': ['csrc/sampling/sampling_kernels.cu']},
    'transform':    {'mode': 'injection', 'struct_mode': 'named',  'vllm_files': ['csrc/activation_kernels.cu', 'csrc/layernorm_kernels.cu']},
    'batch_memcpy': {'mode': 'injection', 'struct_mode': 'named',  'vllm_files': ['csrc/cache_kernels.cu']},
    'for':          {'mode': 'injection', 'struct_mode': 'named',  'vllm_files': ['csrc/pos_encoding_kernels.cu']},

    # === 20 algorithms without direct vllm injection (struct_mode='inline') ===
    # These are used via CCCL device-level APIs, not via #define injection.
    # Their tuning values affect performance when vllm calls CUB functions.
    'adjacent_difference':       {'mode': 'library', 'struct_mode': 'inline'},
    'batched_topk':              {'mode': 'library', 'struct_mode': 'inline'},
    'find':                      {'mode': 'library', 'struct_mode': 'inline'},
    'find_bound_sorted_values':  {'mode': 'library', 'struct_mode': 'inline'},
    'histogram':                 {'mode': 'library', 'struct_mode': 'inline'},
    'merge':                     {'mode': 'library', 'struct_mode': 'inline'},
    'merge_sort':                {'mode': 'library', 'struct_mode': 'inline'},
    'radix_sort':                {'mode': 'library', 'struct_mode': 'inline'},
    'reduce_by_key':             {'mode': 'library', 'struct_mode': 'inline'},
    'rle_encode':                {'mode': 'library', 'struct_mode': 'inline'},
    'rle_non_trivial_runs':      {'mode': 'library', 'struct_mode': 'inline'},
    'scan_by_key':               {'mode': 'library', 'struct_mode': 'inline'},
    'segmented_radix_sort':      {'mode': 'library', 'struct_mode': 'inline'},
    'segmented_reduce':          {'mode': 'library', 'struct_mode': 'inline'},
    'segmented_scan':            {'mode': 'library', 'struct_mode': 'inline'},
    'segmented_sort':            {'mode': 'library', 'struct_mode': 'inline'},
    'select_if':                 {'mode': 'library', 'struct_mode': 'inline'},
    'three_way_partition':       {'mode': 'library', 'struct_mode': 'inline'},
    'transform_tile':            {'mode': 'library', 'struct_mode': 'inline'},
    'unique_by_key':             {'mode': 'library', 'struct_mode': 'inline'},
}



def extract_hardcoded_values(filepath):
    """Fallback: extract key values from policy_selector return statements.
    
    For algorithms where bi100_* structs don't exist (values computed inline).
    Handles multiple patterns:
      - topk: return {threads, items, load_algo, scan_algo, bits}
      - transform: constexpr int bi100_bytes_in_flight = N;
      - batch_memcpy: return {threads, items, ...}
      - generic: first integer in return {} is threads_per_block
    """
    with open(filepath, 'r') as f:
        content = f.read()
    
    algo = algo_from_filename(filepath)
    
    # --- topk special case: extract bits_per_pass from calc_bits_per_pass ---
    if algo == 'topk':
        # Extract the return statement: return {threads, items, ..., bits};
        iluvatar_match = re.search(
            r'hw\.at_least\(.*iluvatar.*?\)\s*\{(.*?)return\s*\{([^}]+)\}',
            content, re.DOTALL
        )
        if iluvatar_match:
            return_args = iluvatar_match.group(2).strip()
            # Pattern: {512, items, BLOCK_LOAD_VECTORIZE, BLOCK_SCAN_WARP_SCANS, calc_bits_per_pass(key_size)}
            parts = [p.strip() for p in return_args.split(',')]
            fields = {}
            if len(parts) >= 1 and parts[0].isdigit():
                fields['threads'] = int(parts[0])
            # calc_bits_per_pass for float32 (key_size=4) = 11
            bits_match = re.search(r'calc_bits_per_pass', return_args)
            if bits_match:
                fields['bits_per_pass'] = 11  # key_size=4 for float32 logits
            return [('__inline_topk__', fields)]
    
    # --- transform special case: extract bytes_in_flight + thread config ---
    if algo == 'transform':
        fields = {}
        bif_match = re.search(r'bi100_bytes_in_flight\s*=\s*(\d+)', content)
        if bif_match:
            fields['bytes_in_flight'] = int(bif_match.group(1))
        # Look for thread count in vectorized policy or return statement
        vec_threads = re.search(
            r'VectorizedPolicy\s*\{?\s*(\d+)\s*,\s*(\d+)',
            content
        )
        if vec_threads:
            fields['threads'] = int(vec_threads.group(1))
            fields['items'] = int(vec_threads.group(2))
        elif not fields:
            # Fallback: find any constexpr threads
            t_match = re.search(r'threads_per_block\s*=?\s*(\d+)', content)
            if t_match:
                fields['threads'] = int(t_match.group(1))
        if fields:
            return [('__inline_transform__', fields)]
    
    # --- batch_memcpy special case ---
    if algo == 'batch_memcpy':
        fields = {}
        t_match = re.search(r'threads_per_block\s*[=:]\s*(\d+)', content)
        if t_match:
            fields['threads'] = int(t_match.group(1))
        if not fields:
            t_match = re.search(r'return\s*\{?\s*(\d+)', content)
            if t_match:
                fields['threads'] = int(t_match.group(1))
        if fields:
            return [('__inline_batch_memcpy__', fields)]
    
    # --- Generic fallback: find iluvatar branch return value ---
    iluvatar_match = re.search(
        r'hw\.at_least\(.*iluvatar.*?\)\s*\{(.*?)(?=\n\s{2,4}\})',
        content, re.DOTALL
    )
    if not iluvatar_match:
        return []
    
    branch = iluvatar_match.group(1)
    
    # Find return {N, ...} — first integer is typically threads_per_block
    return_match = re.search(r'return\s*\{(\d+)', branch)
    if not return_match:
        return []
    
    threads = int(return_match.group(1))
    return [('__inline__', {'threads': threads})]


def generate_patches(header_dir):
    """Read all tuning headers, extract bi100 values, generate patches."""
    patches = []
    summary = []

    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:
            # Fallback: try extracting inline values from policy_selector
            structs = extract_hardcoded_values(hpath)
            if not structs:
                summary.append(f"SKIP {algo}: no bi100_* structs and no inline values found")
                continue

        # Select the struct that matches each vllm kernel's data type.
        #
        # CCCL's policy_selector dispatches by (accum_size, type_t, offset_size).
        # gen_patch must do the same: when injecting into paged_attention
        # (float32 scores), use bi100_plus_float32_o4, not bi100_plus_accum1_o4.
        #
        # The VLLM_KERNEL_MAP in muh_kernel_map.py defines each kernel's
        # data_types. This mapping encodes the primary data type per algorithm:
        ALGO_PRIMARY_TYPE = {
            'reduce':       ('float32', 4),   # paged_attention scores
            'scan':         ('float32', 4),   # softmax denominator
            'topk':         ('float32', 4),   # logits
            'transform':    ('float16', 2),   # activations (SiLU, RMSNorm input)
            'batch_memcpy': ('float16', 2),   # KV cache blocks
            'for':          ('float16', 2),   # RoPE
        }

        target_type, target_size = ALGO_PRIMARY_TYPE.get(algo, ('float32', 4))

        # Score each struct by match quality
        def struct_score(name, fields):
            score = 0
            name_lower = name.lower()
            # Exact type name match (best)
            if target_type.replace('float', 'f') in name_lower or target_type in name_lower:
                score += 100
            # Accum/type size match in name (e.g. "_4B_", "_accum4_", "float32")
            size_tags = [f'_{target_size}B', f'_accum{target_size}', f'float{target_size*8}']
            for tag in size_tags:
                if tag.lower() in name_lower:
                    score += 50
            # Offset size 4 preferred (most common in vllm)
            if '_o4' in name_lower:
                score += 10
            # Penalize 'default' and 'det' (deterministic) structs
            if 'default' in name_lower:
                score -= 200
            if 'det' in name_lower:
                score -= 50
            # Penalize 1-byte type structs for float32 targets
            if target_size >= 4 and ('_1B' in name or 'accum1' in name_lower):
                score -= 100
            return score

        scored = [(struct_score(n, f), n, f) for n, f in structs]
        scored.sort(key=lambda x: -x[0])
        _, pname, pfields = scored[0]

        summary.append(f"READ {algo}: {pname} → {pfields} (target: {target_type})")

        for field_name, value in pfields.items():
            key = (algo, field_name)
            if key not in VLLM_INJECTION_POINTS:
                continue

            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,
                    'struct': pname,
                    'field': field_name,
                    'value': value,
                    'vllm_file': vllm_file,
                    'define': define_name,
                    'diff': patch_text,
                })
                summary.append(
                    f"  PATCH {vllm_file}: {define_name} = {value} "
                    f"(from {pname}.{field_name})"
                )

    return patches, summary


def write_patches(patches, out_dir):
    """Write combined patch file."""
    os.makedirs(out_dir, exist_ok=True)

    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"# 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')

    return combined


def main():
    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()

    patches, summary = generate_patches(args.header_dir)

    print(f"muh gen_patch: scanned {args.header_dir}\n")
    for s in summary:
        print(f"  {s}")

    if not patches:
        print("\nNo patches generated.")
        return

    if args.dry_run:
        print(f"\n--- {len(patches)} patches ---\n")
        for p in patches:
            print(p['diff'])
    else:
        combined = write_patches(patches, args.output_dir)
        print(f"\nWritten: {combined}")


if __name__ == '__main__':
    main()