Method generate_starter

kernelbench/bridge.py:206–303

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

(self, backend: str = "cuda")

    # ----- Starter generation -----

    def generate_starter(self, backend: str = "cuda") -> str:
        """Generate a starter kernel.py (ModelNew initially copies Model logic)."""
        analysis = self.analyze()
        ops_str = ", ".join(analysis["operations"]) or "unknown"

        header = f'''"""
KernelBench Problem {self.uid}: {self.name}
Level: {self.level} | Problem ID: {self.problem_id}
Operations: {ops_str}
Difficulty: {analysis["estimated_difficulty"]}

Source: ScalingIntelligence/KernelBench
Optimized with AutoKernel (https://github.com/RightNow-AI/autokernel)

The agent optimizes ModelNew to outperform the PyTorch reference (Model).
Edit ModelNew.forward() -- use CUDA C++ via compile_cuda() or Triton @jit.
Run `uv run kernelbench/bench_kb.py` to evaluate correctness + speedup.
"""

KERNELBENCH_PROBLEM = {{
    "level": {self.level},
    "problem_id": {self.problem_id},
    "name": {self.name!r},
}}

import torch
import torch.nn as nn
import torch.nn.functional as F
'''

        # Deduplicate imports already in header
        skip_imports = {
            "import torch",
            "import torch.nn as nn",
            "import torch.nn.functional as F",
            "from torch import nn",
        }
        filtered_lines = []
        for line in self.source_code.split("\n"):
            if line.strip() in skip_imports:
                continue
            # Also skip `from torch.nn import functional as F` and similar
            if re.match(r"^\s*import\s+torch\.nn\.functional\s+as\s+F\s*$", line.strip()):
                continue
            filtered_lines.append(line)
        remaining_source = "\n".join(filtered_lines).strip()

        # Build ModelNew by copying Model class
        model_new_source = self._extract_and_rename_model()

        compile_hint = ""
        if backend == "cuda":
            compile_hint = """
# Optional: use AutoKernel's CUDA compilation utility for custom CUDA C++ kernels
# from kernels.cuda._compile import compile_cuda
#
# CUDA_SRC = r\"""
# #include <torch/extension.h>