Add LLM Relay Backend for local plugboard server integration for model evals (#64)

* Add LLMRelayBackend class in backends/llm_relay.py
* Update backends/init.py to include LLMRelayBackend
* Add llm-relay backend support to main.py with setup_llm_relay_backend function
* Add --llm-relay-model CLI option for model configuration
* Update LLMKernelGenerator to support configurable model parameter

Author: shaahins

BackendBench/backends/__init__.py

@@ -12,12 +12,14 @@
 from .flag_gems import FlagGemsBackend
 from .kernel_agent import KernelAgentBackend
 from .llm import LLMBackend
+from .llm_relay import LLMRelayBackend
 
 __all__ = [
     "Backend",
     "DirectoryBackend",
     "AtenBackend",
     "FlagGemsBackend",
     "LLMBackend",
+    "LLMRelayBackend",
     "KernelAgentBackend",
 ]

BackendBench/backends/llm_relay.py

@@ -0,0 +1,262 @@
+import datetime
+import importlib.util
+import logging
+import os
+import sys
+import torch
+import traceback
+from typing import Callable, Dict, List
+
+from .base import Backend
+
+logger = logging.getLogger(__name__)
+
+
+class LLMRelayBackend(Backend):
+    """
+    Backend that uses LLMKernelGenerator to communicate with local plugboard server.
+    This backend will eventually replace the LLMBackend that uses direct Anthropic API calls.
+    """
+
+    def __init__(self, model: str = "gcp-claude-4-sonnet") -> None:
+        super().__init__("llm-relay")
+        self.compiled_kernels: Dict[str, Callable] = {}
+        self.model = model
+
+        # Create generated_kernels directory
+        timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
+        self.kernels_dir = f"generated_kernels/llm_relay_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 - LLM Relay - {timestamp}
+
+This directory contains PyTorch/Triton kernels generated by the LLM Relay Backend.
+
+## Run Info
+- Timestamp: {timestamp}
+- Backend: LLM Relay
+- Model: {model}
+- Server: Local plugboard server (localhost:11434)
+
+## 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
+
+## Server Setup
+This backend requires the plugboard server to be running:
+```
+buck run @//mode/inplace run_plugboard_server -- --model gcp-claude-4-sonnet --pipeline usecase-dev-ai-user
+```
+
+## Usage
+You can inspect these files to debug kernel generation, manually test implementations, or understand what the LLM produced.
+"""
+            )
+
+        logger.info(f"Saving LLM Relay 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)
+
+            logger.debug(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)
+                logger.debug(f"Saved kernel to: {kernel_file}")
+
+            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}"]
+
+                # Clear CUDA cache and synchronize to prevent memory buildup
+                if torch.cuda.is_available():
+                    torch.cuda.empty_cache()
+                    torch.cuda.synchronize()
+
+            correct_count = 0
+            total_count = 0
+
+            for test in test_cases:
+                try:
+                    args = test.args
+                    kwargs = test.kwargs
+
+                    ref_result = op(*args, **kwargs)
+
+                    # Clear CUDA cache after running each kernel to prevent grabbing previous solutions
+                    if torch.cuda.is_available():
+                        torch.cuda.empty_cache()
+
+                    kernel_result = compiled_kernel(*args, **kwargs)
+
+                    torch.testing.assert_close(ref_result, kernel_result, equal_nan=True)
+                    correct_count += 1
+                    logger.debug(f"    ✓ Test passed: {ref_result.shape} {ref_result.dtype}")
+
+                except Exception as e:
+                    logger.debug(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(),
+                        }
+                    )
+
+                finally:
+                    # Clean up memory by deleting args and kwargs if they exist
+                    if "args" in locals():
+                        del args
+                    if "kwargs" in locals():
+                        del kwargs
+
+                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:
+            logger.error("    ✗ Compilation failed:")
+            logger.error(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