Compile CUDA C++ source into a callable PyTorch extension module. Parameters ---------- cuda_src : str The CUDA C++ source c

Return byte-width for a dtype.

Check for NaN or Inf.

Entry point called by bench.py. Must match reference.matmul_ref signature.

Persist the orchestration state to disk.

Compare two tensors and return statistics.

Extract a single tensor from model output, which might be a tuple, dict, or ModelOutput-like object.

Convert a shape dict to a display string like 'M=4096, N=4096, K=4096'.

Create workspace directories if they do not exist.

Determine baseline throughput from the first row in results or from baselines.json.

Identify operations, shapes, parameter usage, and estimate difficulty.

Amdahl's law: S = 1 / ((1 - p) + p / s) where p = fraction of total GPU time in optimized kernels, s = speedup of those kernels.

Load workspace/optimization_plan.json. Returns None if missing.

Benchmark a function and return median time in milliseconds. Uses triton.testing.do_bench if available, otherwise manual implementation.

Short string tag for a dtype, e.g. 'fp16', 'bf16'.

Return the index of the next pending kernel after current_idx, or None.

Derive the per-kernel results TSV path from a kernel file name.

Dynamically load a Python module from a file path.

Load a single TSV file into a DataFrame. Returns None if missing/empty.

Parse a dtype string into a torch.dtype.

Print the move-on decision.

Print greppable summary for agent log parsing.

Trimmed median: median of the middle 80% of measurements.

Run a single forward pass with the correct calling convention.

Time a model using CUDA events.

Move the orchestrator to a new kernel index.

Benchmark model inference. Returns (output, median_latency_ms). Uses CUDA events for precise GPU timing.

Classify an experiment row (dict or pandas Series) into one of: 'kept' -- correctness PASS and tagged as kept / speedup > 1 'failed

Compare reference and optimized outputs. Returns comparison metrics.

Compute fast_p at all standard thresholds.

Extract the CUDA_SRC string from a Python kernel file. Handles: - CUDA_SRC = r\"\"\"...\"\"\" - CUDA_SRC = \"\"\"...\"\"\"

Find all torch::Tensor-returning function declarations in the CUDA source. Returns a list of dicts with keys: - 'return_type': e.g. 't

Generate a markdown report summarizing the session.

Retrieve a problem from cache. Returns None if not cached.

Load a single problem from a standalone .py file.

Load problems from HuggingFace ScalingIntelligence/KernelBench.

Load problems from a local KernelBench git clone.

Load accumulated scores from disk.

Generate the scatter plot and save to progress.png.

Print a leaderboard-style report.

Save verification results as JSON for programmatic consumption.

Scale all shape dimensions by a factor, rounding to nearest integer. Ensures all values are at least 1.

Append a single result row to the per-kernel TSV.

Replace matching modules in the model. Returns number of modules replaced.

Benchmark *fn* using CUDA events. Returns median latency in microseconds.

Build a torch.library ops.def() schema string for a function.

Compare two tensors. Returns match info.

Compare outputs: tensors, tuples/lists of tensors, or scalars.

Return a default kernel entry for the orchestration state.

Scan kernels/ directory for supported kernel types.

Export a CUDA C++ kernel to HF Kernels format.

Export a Triton kernel to HF Kernels format. Triton kernels are already Python, so the export is simpler: package the Triton code as a P

Extract Model class source and produce a renamed ModelNew copy.

Extract a class definition (including body) from source.

Extract a forward declaration for func_name from the CUDA source. load_inline auto-generates pybind11 bindings in main.cpp when ``functi

Standalone GPU detection when bench.py is not importable.

Generate a pybind11-compatible C++ wrapper that forwards torch::Tensor arguments to the CUDA kernel launcher. arg_specs: list of (name,

Generate actionable suggestions based on experiment history.

Generate -gencode flags for the current GPU architecture.

Get current GPU name.

Create a deterministic hash of the source + flags for cache keying.

What-if analysis: if we achieve *assumed_speedup* on the top-N kernels by pct_total.

Build a fresh orchestration state from an optimization plan.

Heuristic: does the model expect input_ids (integer tokens)?

Load a model class from a Python file and instantiate it.

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

FLOPs for a single matmul C[M,N] = A[M,K] @ B[K,N].

Convert a C++ parameter type to a torch.library schema type.

Parse a C++ parameter list into (type, name) pairs. E.g. "torch::Tensor A, torch::Tensor B" -> [("torch::Tensor", "A"), ...]

Move model to device, generate input, validate forward pass. Handles OOM by trying smaller batch sizes.

Replace all nn.LayerNorm modules with optimized wrapper.

Replace all nn.Linear modules with optimized matmul wrapper.

Replace RMSNorm modules. Since there is no standard nn.RMSNorm, we look for common class names and attributes.

Append results to workspace/kb_active/results.json.

Evaluate move-on criteria for the current kernel. Returns (should_move, reason).

Attempt a forward pass. Returns True on success.

Benchmark torch.matmul at all sizes and dtypes. Returns results dict.

Build the profile_report.json structure.