Map a CUDA kernel name to an AutoKernel op type.

Determine which kernel to optimize next and print the decision.

Show the full optimization plan with Amdahl's law analysis.

Record an experiment result for a kernel. *status* is one of: kept, revert, failed, crash, timeout

Generate the aggregate report at workspace/aggregate_report.md.

Print the current orchestration status.

Compute fast_p metric. fast_p = (# correct AND speedup >= threshold) / total

Detect whether a kernel file uses the CUDA C++ or Triton backend. Returns 'cuda' or 'triton'.

Try to import detect_gpu from bench.py; fall back to standalone.

Auto-detect current GPU and return its spec.

Extract the KERNEL_TYPE from the source file, if declared.

Apply each kernel replacement individually to find which one causes failure.

Find optimized kernels from the workspace directory. Checks orchestration_state.json first, then scans for *_optimized.py files.

Rough heuristic: is this kernel compute-bound or memory-bound?

Main export pipeline. Parameters ---------- kernel_path : str Path to the AutoKernel kernel file (kernel.py or similar).

Extract the function name passed to compile_cuda(). Looks for patterns like: compile_cuda(CUDA_SRC, "matmul_cuda") compile_cuda(

Extract the Triton kernel code from a starter file, stripping the original module docstring and KERNEL_TYPE declaration (which we replace

Main extraction pipeline.

Extract the Triton kernel code from a Python file. Returns everything from the first import statement onward, skipping the module docstr

Format the verification result into a human-readable report.

Generate the build.toml file for HF Kernels. Parameters ---------- name : str Kernel project name. functions : list

Generate the flake.nix for the HF kernel-builder Nix pipeline.

Generate the Python __init__.py for the HF Kernels module.

Generate appropriate sample input for the model.

Generate the complete kernel file content for extraction.

Build the optimization_plan.json data structure.

Generate a sample input tensor from a shape string like '1,2048'.

Generate a starter kernel.py (ModelNew initially copies Model logic).

Generate deterministic test tensors for all sizes and dtypes.

Generate torch_binding.cpp with TORCH_LIBRARY_EXPAND registration. Uses the HuggingFace Kernels convention: torch/library.h + registration.h

Generate torch_binding.h with forward declarations.

Return a reasonable default shape for a given op_type when parsing fails. Based on the 'large' size from bench.py KERNEL_CONFIGS.

Load existing state, or create one from the optimization plan.

Extract the list of supported (autokernel_supported=True) kernels from the profile report, sorted by rank.

Get current GPU VRAM usage in MB.

Try to determine if the model expects integer token IDs or float tensors.

Check if this op type has a matching kernels/*.py implementation.

List all cached problems with metadata.

Load cached baselines.json if it exists.

Load ModelNew from kernel.py.

Dynamically import a kernel .py file and return the module.

Load active problem metadata.

Unified model loader from CLI args.

Load model according to CLI args. Returns (model, description_string).

Load a model from a Python file by importing it and instantiating the class.

Load a model from an installed Python module (e.g. 'transformers').

Load workspace/orchestration_state.json if it exists.

Load and validate the profile report JSON. Returns None on failure.

Load the reference Model, get_inputs, get_init_inputs from workspace.

Parse results.tsv into a pandas DataFrame, also merging any TSV files found in workspace/results/ (written by orchestrate.py). Returns No

Load the orchestration state, or None if it does not exist.

Create an appropriate input for the model.

Standard matrix multiplication. A @ B.

Parse '1-10' or '1,3,5,7' into a list of problem IDs.

Parse a shape_info string like "M=4096, N=4096, K=4096" into a dict. Handles various formats: - "M=4096, N=4096, K=4096" - "B=1,

Pretty-print the profiling results to the terminal.

Print a concise summary of the experiment session to stdout.

Profile the model and return a list of KernelRecords sorted by GPU time desc. Returns: Tuple of (records, extras) where extras contains p

Read the starter kernel file. Returns None if not found. For backend='triton': reads from kernels/{op_type}.py For backend='cuda': reads fr

Run all correctness stages. Returns dict with results.

KernelBench-compatible correctness checks. Generate n_trials random input sets, run both models, compare within (atol, rtol).

Run 3 times with identical inputs, check bitwise reproducibility.

Run all (or selected) problems in a level.

Run performance benchmarks. Returns dict with metrics.

KernelBench-compatible performance benchmarking via CUDA event timing. Returns speedup = ref_time / kernel_time.

Run torch profiler and save a trace.

Set up and evaluate a single KernelBench problem. Returns a result dict with correctness, speedup, etc.

Test numerical stability: check for NaN/Inf on normal inputs.

Save scores to disk.

Set up workspace for optimizing a KernelBench problem. Creates: workspace/kb_active/reference.py -- original Model + get_inputs

Import kernel.py, run on tiny input, check correctness.

Print GPU specs, PyTorch version, Triton version. Exit on failure.