megakernel-optimization

AutoMegaKernel (AMK), megakernel schedule optimization

AMK compiles a HuggingFace Llama-family model into ONE persistent CUDA megakernel and tunes it with an AutoKernel-style loop: read the edit surface -> propose ONE knob change -> eval -> keep/revert -> record -> repeat. This skill drives Loop 2 (schedule + kernel_knobs search). You never write kernel code; you only edit a structured ScheduleConfig (plus its reserved kernel_knobs sub-object). The frozen VM lowers your config deterministically and the CPU ReferenceVM judges correctness vs eager PyTorch.

HARD HONESTY RULES (state and obey these every time)

• Correctness FIRST. A latency is NEVER reported without a correctness PASS vs the CPU ReferenceVM. Keep a candidate only if it is correct AND >= 1% faster than the incumbent.
• validate-before-launch. An unsafe ScheduleConfig is a clean REJECTED (a deadlock/race-free proof rejects it before launch), never a hung GPU.
• The edit surface is ScheduleConfig + kernel_knobs ONLY, never raw kernel code, never vm/, never the frozen ABI.
• Measured-gpu latency is drift-robust; physically-impossible sub-roofline latencies are withheld as artifacts.
• All speedups are vs AMK's OWN baseline (default schedule), NOT a claim of beating cuBLAS/vLLM. AMK is currently within ~13% of cuBLAS at batch-1, behind it.

The edit surface (read it before proposing)

Read the surface programmatically, never guess knob names. Prefer the canonical MCP tool; fall back to the CLI if MCP is unavailable.

• MCP: amk_propose(model, gpu="rtx5090") -> { schedule_config, schedule_id, search_space, ... }. search_space includes the kernel_knobs.* sub-surface.
• CLI: amk propose <model> --gpu <arch> (or uv run python amk_cli.py propose <model> --gpu <arch>) prints the same surface as JSON on stdout.

The ScheduleConfig knobs (edit ONE per trial): tiling.gemv.N_tile, tiling.attention.kv_block, fusion_grouping, sm_assignment, pipelining_depth, page_allocation, threads_per_block, smem_bytes_per_block. The reserved kernel_knobs object holds GEMV build knobs: cols_per_warp, cpasync, cpa_stages, cpa_cols (these move MEASURED latency under device=cuda; the predicted/CPU path does not model them). A config WITHOUT kernel_knobs is byte-identical to the production incumbent.