LLM Backend

Author: msaroufim

.gitignore

@@ -2,3 +2,4 @@ __pycache__/
 .claude/
 .vscode/
 .ruff_cache/
+generated_kernels/

BackendBench/backends.py

@@ -1,3 +1,8 @@
+import os
+import importlib.util
+from typing import Dict, Callable, List
+
+
 class Backend:
     def __init__(self, name):
         self.name = name
@@ -278,3 +283,229 @@ def __getitem__(self, key):
 
     def __contains__(self, key):
         return key in self.ops
+
+
+class LLMBackend(Backend):
+    def __init__(self) -> None:
+        super().__init__("llm")
+        self.compiled_kernels: Dict[str, Callable] = {}
+
+        # Create generated_kernels directory
+        import datetime
+
+        timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
+        self.kernels_dir = f"generated_kernels/run_{timestamp}"
+        os.makedirs(self.kernels_dir, exist_ok=True)
+
+        # Create README for this run
+        readme_path = os.path.join(self.kernels_dir, "README.md")
+        with open(readme_path, "w") as f:
+            f.write(f"""# Generated Kernels - {timestamp}
+
+This directory contains PyTorch/Triton kernels generated by the LLM Backend.
+
+## Run Info
+- Timestamp: {timestamp}
+- Backend: LLM
+
+## Files
+Each `<op_name>_kernel.py` file contains the complete generated kernel code for that operation, including:
+- All necessary imports
+- Triton kernel implementation (if applicable)
+- Wrapper function that matches PyTorch operation signature
+
+## Usage
+You can inspect these files to debug kernel generation, manually test implementations, or understand what the LLM produced.
+""")
+
+        print(f"Saving generated kernels to: {self.kernels_dir}")
+
+    def compile_kernel_from_string(
+        self, kernel_code: str, op_name: str, attempt: int = 1
+    ) -> Callable:
+        """Compile a kernel from string code and return a callable."""
+        try:
+            is_triton = "triton.jit" in kernel_code or "@triton.jit" in kernel_code
+
+            if is_triton:
+                full_code = self._prepare_triton_code(kernel_code)
+            else:
+                full_code = self._prepare_torch_code(kernel_code)
+
+            kernel_file = os.path.join(self.kernels_dir, f"{op_name}_kernel_attempt_{attempt}.py")
+            with open(kernel_file, "w") as f:
+                f.write(full_code)
+
+            print(f"Saved kernel to: {kernel_file}")
+
+            spec = importlib.util.spec_from_file_location(f"kernel_{op_name}", kernel_file)
+            module = importlib.util.module_from_spec(spec)
+            spec.loader.exec_module(module)
+
+            kernel_func = self._find_kernel_function(module, op_name)
+
+            return kernel_func
+
+        except Exception as e:
+            raise RuntimeError(f"Failed to compile kernel for {op_name}: {str(e)}")
+
+    def _prepare_triton_code(self, kernel_code: str) -> str:
+        """Prepare Triton kernel code with necessary imports."""
+        imports = """
+import torch
+import triton
+import triton.language as tl
+"""
+        if "import torch" not in kernel_code:
+            kernel_code = imports + kernel_code
+        return kernel_code
+
+    def _prepare_torch_code(self, kernel_code: str) -> str:
+        """Prepare regular PyTorch kernel code with necessary imports."""
+        imports = """
+import torch
+import torch.nn.functional as F
+"""
+        if "import torch" not in kernel_code:
+            kernel_code = imports + kernel_code
+        return kernel_code
+
+    def _find_kernel_function(self, module, op_name: str) -> Callable:
+        """Find the main kernel function in the compiled module."""
+        expected_name = f"{op_name}_kernel_impl"
+
+        if hasattr(module, expected_name):
+            return getattr(module, expected_name)
+
+        available_functions = [
+            name
+            for name in dir(module)
+            if callable(getattr(module, name)) and not name.startswith("_")
+        ]
+
+        raise ValueError(
+            f"Expected function '{expected_name}' not found in kernel code for {op_name}. "
+            f"Available functions: {available_functions}. "
+            f"Please ensure the LLM generated code follows the naming convention: {op_name}_kernel_impl"
+        )
+
+    def add_kernel(self, op, kernel_code: str, op_name: str):
+        """Add a kernel implementation for a specific operator."""
+        compiled_kernel = self.compile_kernel_from_string(kernel_code, op_name, attempt=1)
+        self.compiled_kernels[op] = compiled_kernel
+
+    def test_kernel_correctness(
+        self, op, kernel_code: str, test_cases: List, attempt: int = 1
+    ) -> tuple[bool, Dict]:
+        """Test kernel correctness and return detailed feedback."""
+        op_str = str(op)
+        if "aten." in op_str:
+            op_name = op_str.split("aten.")[-1].split(".")[0]
+        else:
+            op_name = op_str.split(".")[-1]
+
+        feedback_info = {
+            "compilation_error": None,
+            "test_errors": [],
+            "summary": None,
+        }
+
+        try:
+            kernel_file = os.path.join(self.kernels_dir, f"{op_name}_kernel_attempt_{attempt}.py")
+
+            if not os.path.exists(kernel_file):
+                is_triton = "triton.jit" in kernel_code or "@triton.jit" in kernel_code
+                if is_triton:
+                    full_code = self._prepare_triton_code(kernel_code)
+                else:
+                    full_code = self._prepare_torch_code(kernel_code)
+
+                with open(kernel_file, "w") as f:
+                    f.write(full_code)
+                print(f"Saved kernel to: {kernel_file}")
+
+            import sys
+            import importlib.util
+
+            spec = importlib.util.spec_from_file_location(
+                f"test_kernel_{op_name}_{attempt}", kernel_file
+            )
+            module = importlib.util.module_from_spec(spec)
+
+            # Add to sys.modules so triton can find it
+            sys.modules[f"test_kernel_{op_name}_{attempt}"] = module
+
+            try:
+                spec.loader.exec_module(module)
+
+                expected_name = f"{op_name}_kernel_impl"
+                if hasattr(module, expected_name):
+                    compiled_kernel = getattr(module, expected_name)
+                else:
+                    available_functions = [
+                        name
+                        for name in dir(module)
+                        if callable(getattr(module, name)) and not name.startswith("_")
+                    ]
+                    raise ValueError(
+                        f"Expected function '{expected_name}' not found. Available: {available_functions}"
+                    )
+
+            finally:
+                if f"test_kernel_{op_name}_{attempt}" in sys.modules:
+                    del sys.modules[f"test_kernel_{op_name}_{attempt}"]
+
+            import torch
+
+            correct_count = 0
+            total_count = 0
+
+            for test in test_cases:
+                try:
+                    args = test.args
+                    kwargs = test.kwargs
+
+                    ref_result = op(*args, **kwargs)
+                    kernel_result = compiled_kernel(*args, **kwargs)
+
+                    torch.testing.assert_close(ref_result, kernel_result, equal_nan=True)
+                    correct_count += 1
+                    print(f"    ✓ Test passed: {ref_result.shape} {ref_result.dtype}")
+
+                except Exception as e:
+                    import traceback
+
+                    print(f"    ✗ Test failed: {str(e)}")
+
+                    feedback_info["test_errors"].append(
+                        {
+                            "test_input": f"args={[arg.shape if hasattr(arg, 'shape') else arg for arg in args]}, kwargs={kwargs}",
+                            "error": str(e),
+                            "error_type": type(e).__name__,
+                            "traceback": traceback.format_exc(),
+                        }
+                    )
+
+                total_count += 1
+
+            is_correct = correct_count == total_count and total_count > 0
+            if not is_correct:
+                feedback_info["summary"] = f"{correct_count}/{total_count} tests passed"
+
+            return is_correct, feedback_info
+
+        except Exception as e:
+            print("    ✗ Compilation failed:")
+            print(f"      Error: {str(e)}")
+
+            feedback_info["compilation_error"] = str(e)
+            feedback_info["summary"] = "Compilation failed"
+            return False, feedback_info
+
+    def __getitem__(self, key):
+        if key in self.compiled_kernels:
+            return self.compiled_kernels[key]
+        raise KeyError(f"No kernel implementation found for {key}")
+
+    def __contains__(self, key):
+        return key in self.compiled_kernels

BackendBench/eval.py

@@ -1,8 +1,10 @@
 import logging
 
 import torch
+
 from triton.testing import do_bench
 
+
 logger = logging.getLogger(__name__)
 
 
@@ -66,7 +68,6 @@ def eval_performance(op, impl, tests):
         test_times = [cpu_bench(lambda: impl(*test.args, **test.kwargs)) for test in tests]
 
     speedups = torch.tensor(test_times) / torch.tensor(base_times)
-    # geometric mean of speedups
     return speedups.log().mean().exp()
 
 

BackendBench/kernel_templates.py

@@ -0,0 +1,118 @@
+"""
+Kernel code templates and prompt engineering for LLM-based kernel generation.
+"""
+
+from typing import Dict
+from .prompts import (
+    TRITON_KERNEL_PROMPT,
+    PYTORCH_KERNEL_PROMPT,
+    TRITON_OPTIMIZATIONS,
+    TRITON_EXAMPLE_TEMPLATES,
+)
+
+
+class KernelTemplate:
+    """Base class for kernel templates."""
+
+    def __init__(self, name: str, framework: str):
+        self.name = name
+        self.framework = framework
+
+    def create_prompt(self, op_name: str, op_signature: str, op_description: str) -> str:
+        """Create a prompt for kernel generation."""
+        raise NotImplementedError
+
+
+class TritonKernelTemplate(KernelTemplate):
+    """Template for Triton kernel generation."""
+
+    def __init__(self):
+        super().__init__("triton", "triton")
+
+    def create_prompt(self, op_name: str, op_signature: str, op_description: str) -> str:
+        """Create a specialized prompt for Triton kernel generation."""
+
+        # Get operation-specific optimizations
+        optimizations = self._get_optimizations(op_name)
+
+        # Get example template
+        example = self._get_example_template(op_name)
+
+        return TRITON_KERNEL_PROMPT.format(
+            op_name=op_name,
+            op_signature=op_signature,
+            op_description=op_description,
+            optimizations=optimizations,
+            example=example,
+        )
+
+    def _get_optimizations(self, op_name: str) -> str:
+        """Get operation-specific optimization guidelines."""
+        return TRITON_OPTIMIZATIONS.get(op_name, TRITON_OPTIMIZATIONS["default"])
+
+    def _get_example_template(self, op_name: str) -> str:
+        """Get operation-specific code template."""
+        return TRITON_EXAMPLE_TEMPLATES["default"]
+
+
+class PyTorchKernelTemplate(KernelTemplate):
+    """Template for pure PyTorch kernel generation."""
+
+    def __init__(self):
+        super().__init__("pytorch", "pytorch")
+
+    def create_prompt(self, op_name: str, op_signature: str, op_description: str) -> str:
+        """Create a prompt for PyTorch kernel generation."""
+
+        return PYTORCH_KERNEL_PROMPT.format(
+            op_name=op_name, op_signature=op_signature, op_description=op_description
+        )
+
+
+class KernelTemplateManager:
+    """Manages kernel templates for different frameworks."""
+
+    def __init__(self):
+        self.templates: Dict[str, KernelTemplate] = {
+            "triton": TritonKernelTemplate(),
+            "pytorch": PyTorchKernelTemplate(),
+            # TODO: Add cuda, cutile, whatever we want
+        }
+
+    def get_template(self, framework: str) -> KernelTemplate:
+        """Get template for specified framework."""
+        if framework not in self.templates:
+            raise ValueError(f"Unknown framework: {framework}")
+        return self.templates[framework]
+
+    def create_prompt(
+        self, op_name: str, op_signature: str, op_description: str, framework: str = "triton"
+    ) -> str:
+        """Create a prompt using the specified template."""
+        template = self.get_template(framework)
+        return template.create_prompt(op_name, op_signature, op_description)
+
+    def create_refinement_prompt(
+        self,
+        op_name: str,
+        op_signature: str,
+        op_description: str,
+        framework: str = "triton",
+        feedback: str = "",
+    ) -> str:
+        """Create a refinement prompt with feedback from previous attempts."""
+        base_prompt = self.create_prompt(op_name, op_signature, op_description, framework)
+
+        if feedback and feedback.strip():
+            refinement_prompt = f"""{feedback}
+
+{base_prompt}
+
+Fix the above errors and generate corrected code."""
+        else:
+            # Fallback if no feedback
+            refinement_prompt = f"""{base_prompt}
+
+The previous attempt failed. Please generate a corrected version."""
+
+        return refinement_prompt