wafer-cli 0.2.4__py3-none-any.whl → 0.2.5__py3-none-any.whl

wafer/evaluate.py

@@ -14,6 +14,7 @@ logger = logging.getLogger(__name__)
 from wafer_core.utils.kernel_utils.targets.config import (
     BaremetalTarget,
     DigitalOceanTarget,
+    LocalTarget,
     ModalTarget,
     RunPodTarget,
     VMTarget,
@@ -351,6 +352,18 @@ def _build_docker_pip_install_cmd(target: BaremetalTarget | VMTarget) -> str:
     return " && ".join(commands)
 
 
+def _get_wafer_root() -> Path:
+    """Get wafer monorepo root directory.
+
+    Walks up from this file to find the wafer repo root (contains apps/, packages/).
+    """
+    current = Path(__file__).resolve()
+    for parent in [current] + list(current.parents):
+        if (parent / "apps").is_dir() and (parent / "packages").is_dir():
+            return parent
+    raise RuntimeError(f"Could not find wafer root from {__file__}")
+
+
 async def run_evaluate_docker(
     args: EvaluateArgs,
     target: BaremetalTarget | VMTarget,
@@ -488,7 +501,9 @@ async def run_evaluate_docker(
 
         # Build pip install command for torch and other deps, plus wafer-core
         pip_install_cmd = _build_docker_pip_install_cmd(target)
-        install_cmd = f"{pip_install_cmd} && uv pip install --system --break-system-packages wafer-core"
+        install_cmd = (
+            f"{pip_install_cmd} && uv pip install --system --break-system-packages wafer-core"
+        )
 
         # Build evaluate command using installed wafer-core module
         python_cmd_parts = [
@@ -626,6 +641,181 @@ async def run_evaluate_docker(
         )
 
 
+async def run_evaluate_local(
+    args: EvaluateArgs,
+    target: LocalTarget,
+) -> EvaluateResult:
+    """Run evaluation locally on the current machine.
+
+    For LocalTarget - no SSH needed, runs directly.
+
+    Args:
+        args: Evaluate arguments
+        target: Local target config
+
+    Returns:
+        Evaluation result
+    """
+    import os
+    import subprocess
+    import tempfile
+    from datetime import datetime
+
+    # Select GPU
+    gpu_id = args.gpu_id if args.gpu_id is not None else target.gpu_ids[0]
+
+    print(f"Running local evaluation on GPU {gpu_id}...")
+
+    # Create temp directory for eval files
+    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+    with tempfile.TemporaryDirectory(prefix=f"wafer_eval_{timestamp}_") as run_path:
+        run_path = Path(run_path)
+
+        # Write implementation
+        impl_path = run_path / "implementation.py"
+        impl_path.write_text(args.implementation.read_text())
+
+        # Write reference
+        ref_path = run_path / "reference.py"
+        ref_path.write_text(args.reference.read_text())
+
+        # Write custom inputs if provided
+        inputs_path = None
+        if args.inputs:
+            inputs_path = run_path / "custom_inputs.py"
+            inputs_path.write_text(args.inputs.read_text())
+
+        # Write eval script
+        eval_script_path = run_path / "kernelbench_eval.py"
+        eval_script_path.write_text(KERNELBENCH_EVAL_SCRIPT)
+
+        # Write defense module if defensive mode is enabled
+        defense_module_path = None
+        if args.defensive:
+            defense_src = (
+                Path(__file__).parent.parent.parent.parent
+                / "packages"
+                / "wafer-core"
+                / "wafer_core"
+                / "utils"
+                / "kernel_utils"
+                / "defense.py"
+            )
+            if defense_src.exists():
+                defense_module_path = run_path / "defense.py"
+                defense_module_path.write_text(defense_src.read_text())
+            else:
+                print(f"Warning: defense.py not found at {defense_src}")
+
+        # Output file
+        output_path = run_path / "results.json"
+
+        # Build eval command
+        cmd_parts = [
+            "python3",
+            str(eval_script_path),
+            "--impl",
+            str(impl_path),
+            "--reference",
+            str(ref_path),
+            "--output",
+            str(output_path),
+            "--seed",
+            str(args.seed),
+        ]
+
+        if args.benchmark:
+            cmd_parts.append("--benchmark")
+        if args.profile:
+            cmd_parts.append("--profile")
+        if inputs_path:
+            cmd_parts.extend(["--inputs", str(inputs_path)])
+        if args.defensive and defense_module_path:
+            cmd_parts.extend(["--defensive", "--defense-module", str(defense_module_path)])
+
+        # Set environment for GPU selection
+        env = os.environ.copy()
+        if target.vendor == "nvidia":
+            env["CUDA_VISIBLE_DEVICES"] = str(gpu_id)
+        else:  # AMD
+            env["HIP_VISIBLE_DEVICES"] = str(gpu_id)
+            env["ROCM_PATH"] = "/opt/rocm"
+
+        print(f"Running: {' '.join(cmd_parts[:4])} ...")
+
+        # Run evaluation
+        try:
+            result = subprocess.run(
+                cmd_parts,
+                cwd=str(run_path),
+                env=env,
+                capture_output=True,
+                text=True,
+                timeout=args.timeout or 600,
+            )
+        except subprocess.TimeoutExpired:
+            return EvaluateResult(
+                success=False,
+                all_correct=False,
+                correctness_score=0.0,
+                geomean_speedup=0.0,
+                passed_tests=0,
+                total_tests=0,
+                error_message="Evaluation timed out",
+            )
+
+        if result.returncode != 0:
+            error_msg = result.stderr or result.stdout or "Unknown error"
+            # Truncate long errors
+            if len(error_msg) > 1000:
+                error_msg = error_msg[:500] + "\n...\n" + error_msg[-500:]
+            return EvaluateResult(
+                success=False,
+                all_correct=False,
+                correctness_score=0.0,
+                geomean_speedup=0.0,
+                passed_tests=0,
+                total_tests=0,
+                error_message=f"Evaluation failed:\n{error_msg}",
+            )
+
+        # Parse results
+        if not output_path.exists():
+            return EvaluateResult(
+                success=False,
+                all_correct=False,
+                correctness_score=0.0,
+                geomean_speedup=0.0,
+                passed_tests=0,
+                total_tests=0,
+                error_message="No results.json produced",
+            )
+
+        try:
+            results = json.loads(output_path.read_text())
+        except json.JSONDecodeError as e:
+            return EvaluateResult(
+                success=False,
+                all_correct=False,
+                correctness_score=0.0,
+                geomean_speedup=0.0,
+                passed_tests=0,
+                total_tests=0,
+                error_message=f"Failed to parse results: {e}",
+            )
+
+        # Extract results
+        return EvaluateResult(
+            success=True,
+            all_correct=results.get("all_correct", False),
+            correctness_score=results.get("correctness_score", 0.0),
+            geomean_speedup=results.get("geomean_speedup", 0.0),
+            passed_tests=results.get("passed_tests", 0),
+            total_tests=results.get("total_tests", 0),
+            benchmark_results=results.get("benchmark", {}),
+        )
+
+
 async def run_evaluate_ssh(
     args: EvaluateArgs,
     target: BaremetalTarget | VMTarget,
@@ -1025,7 +1215,8 @@ print('Files written')
             return
 
         # Write defense module if defensive mode is enabled
-        if {run_defensive} and "{defense_code_b64}":
+        # NOTE: Check for actual base64 content, not just truthy string (None becomes "None")
+        if {run_defensive} and "{defense_code_b64}" and "{defense_code_b64}" != "None":
             proc = sandbox.exec("python", "-c", f"""
 import base64
 with open('/workspace/defense.py', 'w') as f:
@@ -1072,6 +1263,14 @@ if run_defensive:
         defense = load_fn('defense.py', 'run_all_defenses')
         time_with_defenses = load_fn('defense.py', 'time_execution_with_defenses')
         print('[Defense] Defense module loaded')
+
+        # Wrap kernels for defense API compatibility
+        # Defense API calls kernel(*args), but functional format expects kernel(inputs_tuple)
+        # These wrappers repack the unpacked args back into a tuple
+        def _wrap_for_defense(kernel):
+            return lambda *args: kernel(args)
+        custom_kernel_for_defense = _wrap_for_defense(custom_kernel)
+        ref_kernel_for_defense = _wrap_for_defense(ref_kernel)
     except Exception as e:
         print(f'[Defense] Warning: Could not load defense module: {{e}}')
         defense = None
@@ -1104,30 +1303,30 @@ for tc in test_cases:
         ref_time_ms = 0.0
         if {run_benchmarks}:
             if run_defensive and defense is not None:
-                # Use full defense suite
-                # Run defense checks on implementation kernel
-                all_passed, defense_results, _ = defense(
-                    lambda: custom_kernel(inputs),
-                )
+                # Use full defense suite with wrapped kernels
+                # inputs_list unpacks the tuple so defense can infer dtype/device from tensors
+                inputs_list = list(inputs) if hasattr(inputs, '__iter__') and not isinstance(inputs, torch.Tensor) else [inputs]
+
+                # Run defense checks
+                all_passed, defense_results, _ = defense(custom_kernel_for_defense, *inputs_list)
                 if not all_passed:
                     failed = [name for name, passed, _ in defense_results if not passed]
                     raise ValueError(f"Defense checks failed: {{failed}}")
 
-                # Time with defensive timing
+                # Time with defensive timing (using wrapped kernels)
                 impl_times, _ = time_with_defenses(
-                    lambda: custom_kernel(inputs),
-                    [],
+                    custom_kernel_for_defense,
+                    inputs_list,
                     num_warmup=3,
                     num_trials=10,
                     verbose=False,
-                    run_defenses=False,  # Already ran defenses above
+                    run_defenses=False,
                 )
                 impl_time_ms = sum(impl_times) / len(impl_times)
 
-                # Reference timing (no defense checks needed)
                 ref_times, _ = time_with_defenses(
-                    lambda: ref_kernel(inputs),
-                    [],
+                    ref_kernel_for_defense,
+                    inputs_list,
                     num_warmup=3,
                     num_trials=10,
                     verbose=False,
@@ -1409,7 +1608,8 @@ with open("/tmp/reference.py", "w") as f:
 # Write defense module if available
 run_defensive = {run_defensive}
 defense_b64 = "{defense_b64}"
-if run_defensive and defense_b64:
+# NOTE: Check defense_b64 is not empty and not the string "None" (from None formatting)
+if run_defensive and defense_b64 and defense_b64 != "None":
     defense_code = base64.b64decode(defense_b64).decode()
     with open("/tmp/defense.py", "w") as f:
         f.write(defense_code)
@@ -1429,11 +1629,18 @@ import torch
 
 # Load defense module if available
 defense = None
-if run_defensive and defense_b64:
+if run_defensive and defense_b64 and defense_b64 != "None":
     try:
         defense = load_fn("/tmp/defense.py", "run_all_defenses")
         time_with_defenses = load_fn("/tmp/defense.py", "time_execution_with_defenses")
         print("[Defense] Defense module loaded")
+
+        # Wrap kernels for defense API compatibility
+        # Defense API calls kernel(*args), but functional format expects kernel(inputs_tuple)
+        def _wrap_for_defense(kernel):
+            return lambda *args: kernel(args)
+        custom_kernel_for_defense = _wrap_for_defense(custom_kernel)
+        ref_kernel_for_defense = _wrap_for_defense(ref_kernel)
     except Exception as e:
         print(f"[Defense] Warning: Could not load defense module: {{e}}")
         defense = None
@@ -1466,18 +1673,19 @@ for tc in test_cases:
         ref_time_ms = 0.0
         if {run_benchmarks}:
             if run_defensive and defense is not None:
-                # Use full defense suite
-                all_passed, defense_results, _ = defense(
-                    lambda: custom_kernel(inputs),
-                )
+                # Use full defense suite with wrapped kernels
+                inputs_list = list(inputs) if hasattr(inputs, '__iter__') and not isinstance(inputs, torch.Tensor) else [inputs]
+
+                # Run defense checks
+                all_passed, defense_results, _ = defense(custom_kernel_for_defense, *inputs_list)
                 if not all_passed:
                     failed = [name for name, passed, _ in defense_results if not passed]
                     raise ValueError(f"Defense checks failed: {{failed}}")
 
-                # Time with defensive timing
+                # Time with defensive timing (using wrapped kernels)
                 impl_times, _ = time_with_defenses(
-                    lambda: custom_kernel(inputs),
-                    [],
+                    custom_kernel_for_defense,
+                    inputs_list,
                     num_warmup=3,
                     num_trials=10,
                     verbose=False,
@@ -1485,10 +1693,9 @@ for tc in test_cases:
                 )
                 impl_time_ms = sum(impl_times) / len(impl_times)
 
-                # Reference timing
                 ref_times, _ = time_with_defenses(
-                    lambda: ref_kernel(inputs),
-                    [],
+                    ref_kernel_for_defense,
+                    inputs_list,
                     num_warmup=3,
                     num_trials=10,
                     verbose=False,
@@ -1788,12 +1995,54 @@ async def run_evaluate_runpod(
                         error_message=f"Failed to setup Python environment: {e}",
                     )
 
-                # Install wafer-core in remote venv
-                print("Installing wafer-core...")
-                install_result = await client.exec(
-                    f"{env_state.venv_bin}/uv pip install wafer-core"
-                )
-                if install_result.exit_code != 0:
+                # Upload wafer-core to remote
+                try:
+                    wafer_root = _get_wafer_root()
+                    wafer_core_path = wafer_root / "packages" / "wafer-core"
+                    print(f"Uploading wafer-core from {wafer_core_path}...")
+
+                    wafer_core_remote = f"{REMOTE_WORKSPACE}/wafer-core"
+                    await client.exec(f"mkdir -p {wafer_core_remote}")
+                    wafer_core_workspace = await client.expand_path(wafer_core_remote)
+
+                    upload_result = await client.upload_files(
+                        str(wafer_core_path), wafer_core_workspace, recursive=True
+                    )
+
+                    # Wide event logging for upload result
+                    upload_event = {
+                        "event": "wafer_core_upload",
+                        "target": target.name,
+                        "target_type": "runpod",
+                        "ssh_host": f"{client.user}@{client.host}:{client.port}",
+                        "local_path": str(wafer_core_path),
+                        "remote_path": wafer_core_workspace,
+                        "success": upload_result.success,
+                        "files_copied": upload_result.files_copied,
+                        "duration_seconds": upload_result.duration_seconds,
+                        "error_message": upload_result.error_message,
+                    }
+                    if upload_result.debug_info:
+                        upload_event["debug_info"] = upload_result.debug_info
+                    logger.info(json.dumps(upload_event))
+
+                    # Fail fast if upload failed
+                    if not upload_result.success:
+                        print(f"ERROR: Upload failed: {upload_result.error_message}")
+                        if upload_result.debug_info:
+                            print(f"Debug info: {json.dumps(upload_result.debug_info, indent=2)}")
+                        return EvaluateResult(
+                            success=False,
+                            all_correct=False,
+                            correctness_score=0.0,
+                            geomean_speedup=0.0,
+                            passed_tests=0,
+                            total_tests=0,
+                            error_message=f"Failed to upload wafer-core: {upload_result.error_message}",
+                        )
+
+                    print(f"Uploaded {upload_result.files_copied} files")
+                except Exception as e:
                     return EvaluateResult(
                         success=False,
                         all_correct=False,
@@ -1801,7 +2050,7 @@ async def run_evaluate_runpod(
                         geomean_speedup=0.0,
                         passed_tests=0,
                         total_tests=0,
-                        error_message=f"Failed to install wafer-core: {install_result.stderr}",
+                        error_message=f"Failed to upload wafer-core: {e}",
                     )
 
                 # Select GPU (RunPod pods typically have GPU 0)
@@ -2098,12 +2347,61 @@ async def run_evaluate_digitalocean(
                             error_message=f"Failed to setup Python environment: {e}",
                         )
 
-                    # Install wafer-core in remote venv
-                    print("Installing wafer-core...")
-                    install_result = await client.exec(
-                        f"{env_state.venv_bin}/uv pip install wafer-core"
-                    )
-                    if install_result.exit_code != 0:
+                    # Upload wafer-core to remote
+                    try:
+                        wafer_root = _get_wafer_root()
+                        wafer_core_path = wafer_root / "packages" / "wafer-core"
+                        print(f"Uploading wafer-core from {wafer_core_path}...")
+
+                        wafer_core_remote = f"{REMOTE_WORKSPACE}/wafer-core"
+                        await client.exec(f"mkdir -p {wafer_core_remote}")
+                        wafer_core_workspace = await client.expand_path(wafer_core_remote)
+
+                        # Use SFTP instead of rsync to avoid SSH subprocess timeout issues
+                        # (DigitalOcean may rate-limit new SSH connections)
+                        upload_result = await client.upload_files(
+                            str(wafer_core_path),
+                            wafer_core_workspace,
+                            recursive=True,
+                            use_sftp=True,
+                        )
+
+                        # Wide event logging for upload result
+                        upload_event = {
+                            "event": "wafer_core_upload",
+                            "target": target.name,
+                            "target_type": "digitalocean",
+                            "ssh_host": f"{client.user}@{client.host}:{client.port}",
+                            "local_path": str(wafer_core_path),
+                            "remote_path": wafer_core_workspace,
+                            "success": upload_result.success,
+                            "files_copied": upload_result.files_copied,
+                            "duration_seconds": upload_result.duration_seconds,
+                            "error_message": upload_result.error_message,
+                        }
+                        if upload_result.debug_info:
+                            upload_event["debug_info"] = upload_result.debug_info
+                        logger.info(json.dumps(upload_event))
+
+                        # Fail fast if upload failed
+                        if not upload_result.success:
+                            print(f"ERROR: Upload failed: {upload_result.error_message}")
+                            if upload_result.debug_info:
+                                print(
+                                    f"Debug info: {json.dumps(upload_result.debug_info, indent=2)}"
+                                )
+                            return EvaluateResult(
+                                success=False,
+                                all_correct=False,
+                                correctness_score=0.0,
+                                geomean_speedup=0.0,
+                                passed_tests=0,
+                                total_tests=0,
+                                error_message=f"Failed to upload wafer-core: {upload_result.error_message}",
+                            )
+
+                        print(f"Uploaded {upload_result.files_copied} files")
+                    except Exception as e:
                         return EvaluateResult(
                             success=False,
                             all_correct=False,
@@ -2111,7 +2409,7 @@ async def run_evaluate_digitalocean(
                             geomean_speedup=0.0,
                             passed_tests=0,
                             total_tests=0,
-                            error_message=f"Failed to install wafer-core: {install_result.stderr}",
+                            error_message=f"Failed to upload wafer-core: {e}",
                         )
 
                     # Select GPU (DigitalOcean droplets typically have GPU 0)
@@ -2405,7 +2703,9 @@ async def run_evaluate(args: EvaluateArgs) -> EvaluateResult:
     print(f"Using target: {target_name}")
 
     # Dispatch to appropriate executor
-    if isinstance(target, BaremetalTarget | VMTarget):
+    if isinstance(target, LocalTarget):
+        return await run_evaluate_local(args, target)
+    elif isinstance(target, BaremetalTarget | VMTarget):
         return await run_evaluate_ssh(args, target)
     elif isinstance(target, ModalTarget):
         return await run_evaluate_modal(args, target)
@@ -3531,71 +3831,544 @@ async def run_evaluate_kernelbench_digitalocean(
                 )
 
 
-async def run_evaluate_kernelbench(args: KernelBenchEvaluateArgs) -> EvaluateResult:
-    """Run KernelBench format evaluation on configured target.
-
-    Args:
-        args: KernelBench evaluate arguments
+async def run_evaluate_kernelbench_runpod(
+    args: KernelBenchEvaluateArgs,
+    target: RunPodTarget,
+) -> EvaluateResult:
+    """Run KernelBench format evaluation directly on RunPod AMD GPU.
 
-    Returns:
-        Evaluation result
+    Runs evaluation script directly on host (no Docker) since RunPod pods
+    already have PyTorch/ROCm installed.
     """
-    from .targets import get_default_target, load_target
+    from datetime import datetime
 
-    # Validate input files
-    err = _validate_kernelbench_files(args)
-    if err:
-        return EvaluateResult(
-            success=False,
-            all_correct=False,
-            correctness_score=0.0,
-            geomean_speedup=0.0,
-            passed_tests=0,
-            total_tests=0,
-            error_message=err,
-        )
+    from wafer_core.async_ssh import AsyncSSHClient
+    from wafer_core.targets.runpod import RunPodError, runpod_ssh_context
 
-    # Load target
-    target_name = args.target_name
-    if not target_name:
-        target_name = get_default_target()
-        if not target_name:
-            return EvaluateResult(
-                success=False,
-                all_correct=False,
-                correctness_score=0.0,
-                geomean_speedup=0.0,
-                passed_tests=0,
-                total_tests=0,
-                error_message=(
-                    "No target specified and no default set.\n"
-                    "Set up a target first:\n"
-                    "  wafer config targets init ssh --name my-gpu --host user@host:22\n"
-                    "  wafer config targets init runpod --gpu MI300X\n"
-                    "Then use: --target my-gpu (or set default: wafer config targets default my-gpu)"
-                ),
-            )
+    REMOTE_WORKSPACE_BASE = "/tmp/wafer_eval"
+
+    # Select GPU
+    gpu_id = args.gpu_id if args.gpu_id is not None else target.gpu_ids[0]
+
+    print(f"Provisioning RunPod ({target.gpu_type_id})...")
 
     try:
-        target = load_target(target_name)
-    except FileNotFoundError:
-        return EvaluateResult(
-            success=False,
-            all_correct=False,
-            correctness_score=0.0,
-            geomean_speedup=0.0,
-            passed_tests=0,
-            total_tests=0,
-            error_message=f"Target not found: {target_name}. Run: wafer config targets list",
-        )
+        async with runpod_ssh_context(target) as ssh_info:
+            ssh_target = f"{ssh_info.user}@{ssh_info.host}:{ssh_info.port}"
+            print(f"Connected to RunPod: {ssh_target}")
 
-    print(f"Using target: {target_name}")
+            async with AsyncSSHClient(ssh_target, target.ssh_key) as client:
+                # Create workspace
+                timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+                run_dir = f"kernelbench_eval_{timestamp}"
+                run_path = f"{REMOTE_WORKSPACE_BASE}/{run_dir}"
 
-    # Dispatch to appropriate executor
-    if isinstance(target, DigitalOceanTarget):
-        # DigitalOcean AMD MI300X - uses ROCm Docker with device passthrough
-        return await run_evaluate_kernelbench_digitalocean(args, target)
+                await client.exec(f"mkdir -p {run_path}")
+                print(f"Created run directory: {run_path}")
+
+                # Read and upload files
+                impl_code = args.implementation.read_text()
+                ref_code = args.reference.read_text()
+
+                # Write implementation
+                impl_path = f"{run_path}/implementation.py"
+                write_result = await client.exec(
+                    f"cat > '{impl_path}' << 'IMPL_EOF'\n{impl_code}\nIMPL_EOF"
+                )
+                if write_result.exit_code != 0:
+                    return EvaluateResult(
+                        success=False,
+                        all_correct=False,
+                        correctness_score=0.0,
+                        geomean_speedup=0.0,
+                        passed_tests=0,
+                        total_tests=0,
+                        error_message=f"Failed to write implementation: {write_result.stderr}",
+                    )
+
+                # Write reference
+                ref_path = f"{run_path}/reference.py"
+                write_result = await client.exec(
+                    f"cat > '{ref_path}' << 'REF_EOF'\n{ref_code}\nREF_EOF"
+                )
+                if write_result.exit_code != 0:
+                    return EvaluateResult(
+                        success=False,
+                        all_correct=False,
+                        correctness_score=0.0,
+                        geomean_speedup=0.0,
+                        passed_tests=0,
+                        total_tests=0,
+                        error_message=f"Failed to write reference: {write_result.stderr}",
+                    )
+
+                # Write custom inputs if provided
+                inputs_path = None
+                if args.inputs:
+                    inputs_code = args.inputs.read_text()
+                    inputs_path = f"{run_path}/custom_inputs.py"
+                    write_result = await client.exec(
+                        f"cat > '{inputs_path}' << 'INPUTS_EOF'\n{inputs_code}\nINPUTS_EOF"
+                    )
+                    if write_result.exit_code != 0:
+                        return EvaluateResult(
+                            success=False,
+                            all_correct=False,
+                            correctness_score=0.0,
+                            geomean_speedup=0.0,
+                            passed_tests=0,
+                            total_tests=0,
+                            error_message=f"Failed to write custom inputs: {write_result.stderr}",
+                        )
+
+                # Write eval script
+                eval_script_path = f"{run_path}/kernelbench_eval.py"
+                write_result = await client.exec(
+                    f"cat > '{eval_script_path}' << 'EVAL_EOF'\n{KERNELBENCH_EVAL_SCRIPT}\nEVAL_EOF"
+                )
+                if write_result.exit_code != 0:
+                    return EvaluateResult(
+                        success=False,
+                        all_correct=False,
+                        correctness_score=0.0,
+                        geomean_speedup=0.0,
+                        passed_tests=0,
+                        total_tests=0,
+                        error_message=f"Failed to write eval script: {write_result.stderr}",
+                    )
+
+                # Write defense module if defensive mode is enabled
+                defense_module_path = None
+                if args.defensive:
+                    defense_path = (
+                        Path(__file__).parent.parent.parent.parent
+                        / "packages"
+                        / "wafer-core"
+                        / "wafer_core"
+                        / "utils"
+                        / "kernel_utils"
+                        / "defense.py"
+                    )
+                    if defense_path.exists():
+                        defense_code = defense_path.read_text()
+                        defense_module_path = f"{run_path}/defense.py"
+                        write_result = await client.exec(
+                            f"cat > '{defense_module_path}' << 'DEFENSE_EOF'\n{defense_code}\nDEFENSE_EOF"
+                        )
+                        if write_result.exit_code != 0:
+                            print(f"Warning: Failed to write defense module: {write_result.stderr}")
+                            defense_module_path = None
+                    else:
+                        print(f"Warning: defense.py not found at {defense_path}")
+
+                print("Running KernelBench evaluation (AMD/ROCm)...")
+
+                # Find Python with PyTorch - check common locations on RunPod
+                python_exe = "python3"
+                for candidate in [
+                    "/opt/conda/envs/py_3.10/bin/python3",
+                    "/opt/conda/bin/python3",
+                ]:
+                    check = await client.exec(
+                        f"{candidate} -c 'import torch' 2>/dev/null && echo OK"
+                    )
+                    if "OK" in check.stdout:
+                        python_exe = candidate
+                        print(f"Using Python: {python_exe}")
+                        break
+
+                # Build eval command - run directly on host
+                output_path = f"{run_path}/results.json"
+                python_cmd_parts = [
+                    f"{python_exe} {eval_script_path}",
+                    f"--impl {impl_path}",
+                    f"--reference {ref_path}",
+                    f"--output {output_path}",
+                ]
+
+                if args.benchmark:
+                    python_cmd_parts.append("--benchmark")
+                if args.profile:
+                    python_cmd_parts.append("--profile")
+                if inputs_path:
+                    python_cmd_parts.append(f"--inputs {inputs_path}")
+                if args.defensive and defense_module_path:
+                    python_cmd_parts.append("--defensive")
+                    python_cmd_parts.append(f"--defense-module {defense_module_path}")
+                python_cmd_parts.append(f"--seed {args.seed}")
+
+                eval_cmd = " ".join(python_cmd_parts)
+
+                # Set environment for AMD GPU and run
+                env_vars = f"HIP_VISIBLE_DEVICES={gpu_id} ROCM_PATH=/opt/rocm PYTHONUNBUFFERED=1"
+                full_cmd = f"cd {run_path} && {env_vars} {eval_cmd}"
+
+                # Run and stream output
+                log_lines = []
+                async for line in client.exec_stream(full_cmd):
+                    print(line)
+                    log_lines.append(line)
+
+                # Read results
+                cat_result = await client.exec(f"cat {output_path}")
+
+                if cat_result.exit_code != 0:
+                    log_tail = "\n".join(log_lines[-50:])
+                    return EvaluateResult(
+                        success=False,
+                        all_correct=False,
+                        correctness_score=0.0,
+                        geomean_speedup=0.0,
+                        passed_tests=0,
+                        total_tests=0,
+                        error_message=f"Evaluation failed. Log tail:\n{log_tail}",
+                    )
+
+                # Parse results
+                try:
+                    results_data = json.loads(cat_result.stdout)
+                except json.JSONDecodeError as e:
+                    return EvaluateResult(
+                        success=False,
+                        all_correct=False,
+                        correctness_score=0.0,
+                        geomean_speedup=0.0,
+                        passed_tests=0,
+                        total_tests=0,
+                        error_message=f"Failed to parse results: {e}",
+                    )
+
+                # Convert to EvaluateResult
+                correct = results_data.get("correct", False)
+                speedup = results_data.get("speedup", 0.0) or 0.0
+                error = results_data.get("error")
+
+                if error:
+                    return EvaluateResult(
+                        success=False,
+                        all_correct=False,
+                        correctness_score=0.0,
+                        geomean_speedup=0.0,
+                        passed_tests=0,
+                        total_tests=1,
+                        error_message=error,
+                    )
+
+                return EvaluateResult(
+                    success=True,
+                    all_correct=correct,
+                    correctness_score=1.0 if correct else 0.0,
+                    geomean_speedup=speedup,
+                    passed_tests=1 if correct else 0,
+                    total_tests=1,
+                )
+
+    except RunPodError as e:
+        return EvaluateResult(
+            success=False,
+            all_correct=False,
+            correctness_score=0.0,
+            geomean_speedup=0.0,
+            passed_tests=0,
+            total_tests=0,
+            error_message=f"RunPod error: {e}",
+        )
+
+
+async def run_evaluate_kernelbench_baremetal_amd(
+    args: KernelBenchEvaluateArgs,
+    target: BaremetalTarget,
+) -> EvaluateResult:
+    """Run KernelBench format evaluation directly on AMD baremetal target.
+
+    Runs evaluation script directly on host (no Docker) for AMD GPUs
+    that have PyTorch/ROCm installed.
+    """
+    from datetime import datetime
+
+    from wafer_core.async_ssh import AsyncSSHClient
+
+    REMOTE_WORKSPACE_BASE = "/tmp/wafer_eval"
+
+    # Select GPU
+    gpu_id = args.gpu_id if args.gpu_id is not None else target.gpu_ids[0]
+
+    print(f"Connecting to {target.ssh_target}...")
+
+    async with AsyncSSHClient(target.ssh_target, target.ssh_key) as client:
+        # Create workspace
+        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+        run_dir = f"kernelbench_eval_{timestamp}"
+        run_path = f"{REMOTE_WORKSPACE_BASE}/{run_dir}"
+
+        await client.exec(f"mkdir -p {run_path}")
+        print(f"Created run directory: {run_path}")
+
+        # Read and upload files
+        impl_code = args.implementation.read_text()
+        ref_code = args.reference.read_text()
+
+        # Write implementation
+        impl_path = f"{run_path}/implementation.py"
+        write_result = await client.exec(
+            f"cat > '{impl_path}' << 'IMPL_EOF'\n{impl_code}\nIMPL_EOF"
+        )
+        if write_result.exit_code != 0:
+            return EvaluateResult(
+                success=False,
+                all_correct=False,
+                correctness_score=0.0,
+                geomean_speedup=0.0,
+                passed_tests=0,
+                total_tests=0,
+                error_message=f"Failed to write implementation: {write_result.stderr}",
+            )
+
+        # Write reference
+        ref_path = f"{run_path}/reference.py"
+        write_result = await client.exec(f"cat > '{ref_path}' << 'REF_EOF'\n{ref_code}\nREF_EOF")
+        if write_result.exit_code != 0:
+            return EvaluateResult(
+                success=False,
+                all_correct=False,
+                correctness_score=0.0,
+                geomean_speedup=0.0,
+                passed_tests=0,
+                total_tests=0,
+                error_message=f"Failed to write reference: {write_result.stderr}",
+            )
+
+        # Write custom inputs if provided
+        inputs_path = None
+        if args.inputs:
+            inputs_code = args.inputs.read_text()
+            inputs_path = f"{run_path}/custom_inputs.py"
+            write_result = await client.exec(
+                f"cat > '{inputs_path}' << 'INPUTS_EOF'\n{inputs_code}\nINPUTS_EOF"
+            )
+            if write_result.exit_code != 0:
+                return EvaluateResult(
+                    success=False,
+                    all_correct=False,
+                    correctness_score=0.0,
+                    geomean_speedup=0.0,
+                    passed_tests=0,
+                    total_tests=0,
+                    error_message=f"Failed to write custom inputs: {write_result.stderr}",
+                )
+
+        # Write eval script
+        eval_script_path = f"{run_path}/kernelbench_eval.py"
+        write_result = await client.exec(
+            f"cat > '{eval_script_path}' << 'EVAL_EOF'\n{KERNELBENCH_EVAL_SCRIPT}\nEVAL_EOF"
+        )
+        if write_result.exit_code != 0:
+            return EvaluateResult(
+                success=False,
+                all_correct=False,
+                correctness_score=0.0,
+                geomean_speedup=0.0,
+                passed_tests=0,
+                total_tests=0,
+                error_message=f"Failed to write eval script: {write_result.stderr}",
+            )
+
+        # Write defense module if defensive mode is enabled
+        defense_module_path = None
+        if args.defensive:
+            defense_path = (
+                Path(__file__).parent.parent.parent.parent
+                / "packages"
+                / "wafer-core"
+                / "wafer_core"
+                / "utils"
+                / "kernel_utils"
+                / "defense.py"
+            )
+            if defense_path.exists():
+                defense_code = defense_path.read_text()
+                defense_module_path = f"{run_path}/defense.py"
+                write_result = await client.exec(
+                    f"cat > '{defense_module_path}' << 'DEFENSE_EOF'\n{defense_code}\nDEFENSE_EOF"
+                )
+                if write_result.exit_code != 0:
+                    print(f"Warning: Failed to write defense module: {write_result.stderr}")
+                    defense_module_path = None
+            else:
+                print(f"Warning: defense.py not found at {defense_path}")
+
+        print("Running KernelBench evaluation (AMD/ROCm)...")
+
+        # Find Python with PyTorch - check common locations
+        python_exe = "python3"
+        for candidate in [
+            "/opt/conda/envs/py_3.10/bin/python3",
+            "/opt/conda/bin/python3",
+        ]:
+            check = await client.exec(f"{candidate} -c 'import torch' 2>/dev/null && echo OK")
+            if "OK" in check.stdout:
+                python_exe = candidate
+                print(f"Using Python: {python_exe}")
+                break
+
+        # Build eval command - run directly on host
+        output_path = f"{run_path}/results.json"
+        python_cmd_parts = [
+            f"{python_exe} {eval_script_path}",
+            f"--impl {impl_path}",
+            f"--reference {ref_path}",
+            f"--output {output_path}",
+        ]
+
+        if args.benchmark:
+            python_cmd_parts.append("--benchmark")
+        if args.profile:
+            python_cmd_parts.append("--profile")
+        if inputs_path:
+            python_cmd_parts.append(f"--inputs {inputs_path}")
+        if args.defensive and defense_module_path:
+            python_cmd_parts.append("--defensive")
+            python_cmd_parts.append(f"--defense-module {defense_module_path}")
+        python_cmd_parts.append(f"--seed {args.seed}")
+
+        eval_cmd = " ".join(python_cmd_parts)
+
+        # Set environment for AMD GPU and run
+        env_vars = f"HIP_VISIBLE_DEVICES={gpu_id} ROCM_PATH=/opt/rocm PYTHONUNBUFFERED=1"
+        full_cmd = f"cd {run_path} && {env_vars} {eval_cmd}"
+
+        # Run and stream output
+        log_lines = []
+        async for line in client.exec_stream(full_cmd):
+            print(line)
+            log_lines.append(line)
+
+        # Read results
+        cat_result = await client.exec(f"cat {output_path}")
+
+        if cat_result.exit_code != 0:
+            log_tail = "\n".join(log_lines[-50:])
+            return EvaluateResult(
+                success=False,
+                all_correct=False,
+                correctness_score=0.0,
+                geomean_speedup=0.0,
+                passed_tests=0,
+                total_tests=0,
+                error_message=f"Evaluation failed. Log tail:\n{log_tail}",
+            )
+
+        # Parse results
+        try:
+            results_data = json.loads(cat_result.stdout)
+        except json.JSONDecodeError as e:
+            return EvaluateResult(
+                success=False,
+                all_correct=False,
+                correctness_score=0.0,
+                geomean_speedup=0.0,
+                passed_tests=0,
+                total_tests=0,
+                error_message=f"Failed to parse results: {e}",
+            )
+
+        # Convert to EvaluateResult
+        correct = results_data.get("correct", False)
+        speedup = results_data.get("speedup", 0.0) or 0.0
+        error = results_data.get("error")
+
+        if error:
+            return EvaluateResult(
+                success=False,
+                all_correct=False,
+                correctness_score=0.0,
+                geomean_speedup=0.0,
+                passed_tests=0,
+                total_tests=1,
+                error_message=error,
+            )
+
+        return EvaluateResult(
+            success=True,
+            all_correct=correct,
+            correctness_score=1.0 if correct else 0.0,
+            geomean_speedup=speedup,
+            passed_tests=1 if correct else 0,
+            total_tests=1,
+        )
+
+
+async def run_evaluate_kernelbench(args: KernelBenchEvaluateArgs) -> EvaluateResult:
+    """Run KernelBench format evaluation on configured target.
+
+    Args:
+        args: KernelBench evaluate arguments
+
+    Returns:
+        Evaluation result
+    """
+    from .targets import get_default_target, load_target
+
+    # Validate input files
+    err = _validate_kernelbench_files(args)
+    if err:
+        return EvaluateResult(
+            success=False,
+            all_correct=False,
+            correctness_score=0.0,
+            geomean_speedup=0.0,
+            passed_tests=0,
+            total_tests=0,
+            error_message=err,
+        )
+
+    # Load target
+    target_name = args.target_name
+    if not target_name:
+        target_name = get_default_target()
+        if not target_name:
+            return EvaluateResult(
+                success=False,
+                all_correct=False,
+                correctness_score=0.0,
+                geomean_speedup=0.0,
+                passed_tests=0,
+                total_tests=0,
+                error_message=(
+                    "No target specified and no default set.\n"
+                    "Set up a target first:\n"
+                    "  wafer config targets init ssh --name my-gpu --host user@host:22\n"
+                    "  wafer config targets init runpod --gpu MI300X\n"
+                    "Then use: --target my-gpu (or set default: wafer config targets default my-gpu)"
+                ),
+            )
+
+    try:
+        target = load_target(target_name)
+    except FileNotFoundError:
+        return EvaluateResult(
+            success=False,
+            all_correct=False,
+            correctness_score=0.0,
+            geomean_speedup=0.0,
+            passed_tests=0,
+            total_tests=0,
+            error_message=f"Target not found: {target_name}. Run: wafer config targets list",
+        )
+
+    print(f"Using target: {target_name}")
+
+    # Dispatch to appropriate executor
+    if isinstance(target, DigitalOceanTarget):
+        # DigitalOcean AMD MI300X - uses ROCm Docker with device passthrough
+        return await run_evaluate_kernelbench_digitalocean(args, target)
+    elif isinstance(target, RunPodTarget):
+        # RunPod AMD MI300X - uses ROCm Docker with device passthrough
+        return await run_evaluate_kernelbench_runpod(args, target)
     elif isinstance(target, BaremetalTarget | VMTarget):
+        # Check if this is an AMD target (gfx* compute capability) - run directly
+        if target.compute_capability and target.compute_capability.startswith("gfx"):
+            return await run_evaluate_kernelbench_baremetal_amd(args, target)
         # NVIDIA targets - require docker_image to be set
         if not target.docker_image:
             return EvaluateResult(
@@ -3621,6 +4394,6 @@ async def run_evaluate_kernelbench(args: KernelBenchEvaluateArgs) -> EvaluateRes
             total_tests=0,
             error_message=(
                 f"Target type '{type(target).__name__}' not yet supported for KernelBench format. "
-                "Use a DigitalOcean, Baremetal, or VM target."
+                "Use a DigitalOcean, RunPod, Baremetal, or VM target."
             ),
         )