wafer-core 0.1.22__py3-none-any.whl → 0.1.23__py3-none-any.whl

wafer_core/environments/coding.py

@@ -34,6 +34,7 @@ from wafer_core.tools import (
     GLOB_TOOL,
     GREP_TOOL,
     READ_TOOL,
+    SEARCH_DOCS_TOOL,
     SKILL_TOOL,
     WRITE_TOOL,
     ApprovalCallback,
@@ -42,6 +43,7 @@ from wafer_core.tools import (
     exec_glob,
     exec_grep,
     exec_read,
+    exec_search_docs,
     exec_skill,
     exec_write,
 )
@@ -63,6 +65,7 @@ ALL_TOOLS = {
     "glob": GLOB_TOOL,
     "grep": GREP_TOOL,
     "bash": BASH_TOOL,
+    "search_docs": SEARCH_DOCS_TOOL,
     "skill": SKILL_TOOL,
     # TODO(wafer-tool): "wafer": WAFER_TOOL,
 }
@@ -211,6 +214,7 @@ class CodingEnvironment:
                 self.bash_approval_callback,
                 self._sandbox_policy,
             ),
+            "search_docs": lambda tc: exec_search_docs(tc),
             "skill": lambda tc: exec_skill(tc),
             # TODO(wafer-tool): "wafer": lambda tc: exec_wafer(
             #     tc, self.working_dir, self.enabled_tools, self.allow_spawn, cancel_scope

wafer_core/targets/runpod.py

@@ -14,7 +14,6 @@ from __future__ import annotations
 
 import json
 import logging
-import os
 import time
 from contextlib import asynccontextmanager
 from dataclasses import dataclass
@@ -250,10 +249,17 @@ async def provision_pod(target: RunPodTarget) -> tuple[str, str, int, str]:
         "ports": "22/tcp",
         "startSsh": True,
         "startJupyter": False,
-        "imageName": target.image,
         "env": [],
     }
 
+    if target.template_id:
+        # Template defines image, dockerArgs (sshd setup), and ports.
+        # Required for non-RunPod images (e.g. rocm/pytorch) that don't
+        # have RunPod's built-in SSH handler.
+        pod_input["templateId"] = target.template_id
+    else:
+        pod_input["imageName"] = target.image
+
     variables = {"input": pod_input}
 
     logger.info(f"Provisioning RunPod pod: {pod_name}")
@@ -334,7 +340,8 @@ async def _wait_for_ssh(pod_id: str, timeout_seconds: int) -> tuple[str, int, st
         # Check for SSH port
         runtime = pod.get("runtime")
         if runtime and status == "running":
-            for port in runtime.get("ports", []):
+            # ports can be null in JSON response, so use 'or []' instead of default
+            for port in runtime.get("ports") or []:
                 if (
                     port.get("privatePort") == 22
                     and port.get("isIpPublic")
@@ -378,6 +385,55 @@ async def terminate_pod(pod_id: str) -> bool:
         return False
 
 
+# =============================================================================
+# Template Management (not yet implemented)
+# =============================================================================
+#
+# The saveTemplate mutation allows creating reusable pod templates with custom
+# configurations. Templates can specify docker images, environment setup,
+# container disk size, and other pod settings.
+#
+# Example mutation:
+#
+#     mutation saveTemplate($input: SaveTemplateInput) {
+#         saveTemplate(input: $input) {
+#             id
+#             name
+#             imageName
+#             containerDiskInGb
+#             ports
+#             dockerArgs
+#             startSsh
+#             startJupyter
+#         }
+#     }
+#
+# Example variables:
+#
+#     {
+#         "input": {
+#             "containerDiskInGb": 50,
+#             "dockerArgs": "bash -c \"apt-get update && apt-get install -y openssh-server && ...\"",
+#             "env": [],
+#             "isPublic": false,
+#             "isServerless": false,
+#             "name": "template-name",
+#             "ports": "22/tcp",
+#             "portsConfig": [{"name": "SSH", "port": "22"}],
+#             "readme": "",
+#             "volumeInGb": 0,
+#             "volumeMountPath": "",
+#             "config": {},
+#             "category": "AMD",
+#             "imageName": "rocm/pytorch:rocm7.0.2_ubuntu24.04_py3.12_pytorch_release_2.7.1"
+#         }
+#     }
+#
+# Note: Template creation is not currently implemented in this module.
+# If needed, implement a save_template() function following the pattern of
+# provision_pod() and terminate_pod() above.
+
+
 # =============================================================================
 # Context Manager
 # =============================================================================
@@ -482,20 +538,103 @@ async def cleanup_target(target_name: str) -> bool:
     return success
 
 
+async def sync_pods_from_api() -> list[PodState]:
+    """Query RunPod API for all running pods and update local state.
+
+    This discovers pods that exist on the account but aren't in our state file
+    (e.g., created manually or by another machine). Updates the state file with
+    any wafer-created pods found.
+
+    Returns list of all running pods with SSH info.
+    """
+    query = """
+    query {
+        myself {
+            pods {
+                id
+                name
+                desiredStatus
+                runtime {
+                    ports {
+                        ip
+                        isIpPublic
+                        privatePort
+                        publicPort
+                    }
+                }
+            }
+        }
+    }
+    """
+
+    try:
+        data = await _graphql_request_async(query)
+    except Exception as e:
+        logger.warning(f"Failed to query pods from API: {e}")
+        return []
+
+    pods = data.get("myself", {}).get("pods", [])
+    running_pods = []
+
+    for pod in pods:
+        status = pod.get("desiredStatus", "").lower()
+        if status != "running":
+            continue
+
+        pod_id = pod["id"]
+        pod_name = pod.get("name", "")
+
+        # Extract SSH info
+        runtime = pod.get("runtime")
+        if not runtime:
+            continue
+
+        public_ip = None
+        ssh_port = None
+        for port in runtime.get("ports") or []:
+            if port.get("privatePort") == 22 and port.get("isIpPublic"):
+                public_ip = port.get("ip")
+                ssh_port = port.get("publicPort")
+                break
+
+        if not public_ip or not ssh_port:
+            continue
+
+        # Extract target name from pod name (wafer-{target_name}-{timestamp})
+        target_name = None
+        if pod_name.startswith("wafer-"):
+            parts = pod_name.split("-")
+            if len(parts) >= 3:
+                # Handle target names with hyphens: wafer-runpod-mi300x-1234567
+                target_name = "-".join(parts[1:-1])
+
+        pod_state = PodState(
+            pod_id=pod_id,
+            target_name=target_name or pod_name,
+            public_ip=public_ip,
+            ssh_port=ssh_port,
+            ssh_username="root",
+            created_at=datetime.now(timezone.utc).isoformat(),
+        )
+        running_pods.append(pod_state)
+
+        # Update state file if this is a wafer-created pod
+        if target_name:
+            existing = get_pod_state(target_name)
+            if not existing or existing.pod_id != pod_id:
+                logger.info(f"Syncing pod {pod_id} to state for target {target_name}")
+                _add_pod_to_state(target_name, pod_id, public_ip, ssh_port, "root")
+
+    return running_pods
+
+
 async def list_running_pods() -> list[PodState]:
-    """List all pods in state file that are still running."""
-    state = _load_state()
-    running = []
+    """List all running pods by querying the RunPod API.
 
-    for name, pod_state in state.items():
-        if await check_pod_running(pod_state.pod_id):
-            running.append(pod_state)
-        else:
-            # Clean up stale entry
-            logger.info(f"Removing stale state for {name} (pod {pod_state.pod_id})")
-            _remove_pod_from_state(name)
-
-    return running
+    Syncs state file with API to discover pods not in local state.
+    Returns list of running pods with SSH info.
+    """
+    return await sync_pods_from_api()
 
 
 async def cleanup_all_pods() -> int:

wafer_core/tools/__init__.py

@@ -72,6 +72,10 @@ from wafer_core.tools.write_kernel_tool import (
     KernelSubmission,
     exec_write_kernel,
 )
+from wafer_core.tools.search_docs_tool import (
+    SEARCH_DOCS_TOOL,
+    exec_search_docs,
+)
 
 __all__ = [
     # File tools
@@ -133,4 +137,7 @@ __all__ = [
     "exec_tracelens_report",
     "exec_tracelens_compare",
     "exec_tracelens_collective",
+    # Search docs tool
+    "SEARCH_DOCS_TOOL",
+    "exec_search_docs",
 ]

wafer_core/tools/file_tools/grep_tool.py

@@ -1,4 +1,4 @@
-"""Grep tool using ripgrep for fast content search."""
+"""Grep tool using ripgrep (with fallback to standard grep)."""
 
 from pathlib import Path
 
@@ -15,7 +15,7 @@ GREP_TOOL = Tool(
     function=ToolFunction(
         name="grep",
         description=(
-            "Search for a pattern in files using ripgrep. "
+            "Search for a pattern in files. "
             "Returns matching lines with file paths and line numbers. "
             "Supports regex patterns by default."
         ),
@@ -54,7 +54,7 @@ GREP_TOOL = Tool(
 
 
 async def exec_grep(tool_call: ToolCall, working_dir: Path) -> ToolResult:
-    """Execute grep using ripgrep."""
+    """Execute grep using ripgrep (preferred) or standard grep (fallback)."""
     import shutil
     import subprocess
 
@@ -74,35 +74,55 @@ async def exec_grep(tool_call: ToolCall, working_dir: Path) -> ToolResult:
             error="'pattern' is required",
         )
 
-    # Find ripgrep
+    # Try ripgrep first, fall back to standard grep
     rg_path = shutil.which("rg")
-    if not rg_path:
+    grep_path = shutil.which("grep")
+    
+    if rg_path:
+        # Use ripgrep (faster, better defaults)
+        cmd = [rg_path, "--line-number", "--no-heading", "--color=never"]
+        
+        if case_insensitive:
+            cmd.append("--ignore-case")
+        
+        if context_lines:
+            cmd.extend(["--context", str(context_lines)])
+        
+        if glob_pattern:
+            cmd.extend(["--glob", glob_pattern])
+        
+        # Limit results
+        cmd.extend(["--max-count", str(max_results)])
+        
+        cmd.append(pattern)
+        cmd.append(search_path)
+        use_ripgrep = True
+    elif grep_path:
+        # Fallback to standard grep
+        cmd = [grep_path, "-r", "-n", "--color=never"]
+        
+        if case_insensitive:
+            cmd.append("-i")
+        
+        if context_lines:
+            cmd.extend(["-C", str(context_lines)])
+        
+        if glob_pattern:
+            # Standard grep uses --include for glob patterns
+            cmd.extend(["--include", glob_pattern])
+        
+        cmd.append(pattern)
+        cmd.append(search_path)
+        use_ripgrep = False
+    else:
         return ToolResult(
             tool_call_id=tool_call.id,
             is_error=True,
             content="",
-            error="ripgrep (rg) not found. Please install it: brew install ripgrep",
+            error="Neither ripgrep (rg) nor grep found. Please install one.",
         )
 
-    # Build command
-    cmd = [rg_path, "--line-number", "--no-heading", "--color=never"]
-
-    if case_insensitive:
-        cmd.append("--ignore-case")
-
-    if context_lines:
-        cmd.extend(["--context", str(context_lines)])
-
-    if glob_pattern:
-        cmd.extend(["--glob", glob_pattern])
-
-    # Limit results
-    cmd.extend(["--max-count", str(max_results)])
-
-    cmd.append(pattern)
-    cmd.append(search_path)
-
-    # Run ripgrep
+    # Run the search
     try:
         result = subprocess.run(
             cmd,
@@ -126,13 +146,14 @@ async def exec_grep(tool_call: ToolCall, working_dir: Path) -> ToolResult:
             error=f"Search failed: {e}",
         )
 
-    # ripgrep returns exit code 1 for no matches (not an error)
+    # Both ripgrep and grep return exit code 1 for no matches (not an error)
     if result.returncode not in (0, 1):
+        tool_name = "ripgrep" if use_ripgrep else "grep"
         return ToolResult(
             tool_call_id=tool_call.id,
             is_error=True,
             content="",
-            error=result.stderr or f"ripgrep exited with code {result.returncode}",
+            error=result.stderr or f"{tool_name} exited with code {result.returncode}",
         )
 
     output = result.stdout.strip()
@@ -143,8 +164,14 @@ async def exec_grep(tool_call: ToolCall, working_dir: Path) -> ToolResult:
             content=f"No matches found for pattern: {pattern}",
         )
 
-    # Count matches
-    match_count = len(output.split("\n"))
+    # Count matches and limit output for standard grep
+    lines = output.split("\n")
+    if not use_ripgrep and len(lines) > max_results:
+        lines = lines[:max_results]
+        output = "\n".join(lines)
+        output += f"\n... (truncated to {max_results} results)"
+    
+    match_count = min(len(lines), max_results)
     header = f"Found {match_count} matches:\n\n"
 
     return ToolResult(

wafer_core/tools/search_docs_tool.py

@@ -0,0 +1,196 @@
+"""Search documentation tool for GPU programming corpora.
+
+Provides semantic and keyword search over documentation for CuTeDSL, CUDA, etc.
+
+Corpora are downloaded via `wafer corpus download <name>` and stored in ~/.cache/wafer/corpora/.
+"""
+
+import re
+from pathlib import Path
+
+from wafer_core.rollouts.dtypes import Tool, ToolCall, ToolFunction, ToolFunctionParameter, ToolResult
+
+# Cache directory where wafer corpus download stores files
+CACHE_DIR = Path.home() / ".cache" / "wafer" / "corpora"
+
+# Available corpora (names match wafer corpus download)
+AVAILABLE_CORPORA = ["cutlass", "cutedsl", "cuda", "hip", "amd"]
+
+SEARCH_DOCS_TOOL = Tool(
+    type="function",
+    function=ToolFunction(
+        name="search_docs",
+        description="""Search GPU programming documentation for relevant information.
+
+Use this tool to find documentation about:
+- CUTLASS C++ (cute:: namespace, gemm tutorials, tensor cores, TMA, Blackwell)
+- CuTeDSL Python API (@cute.kernel, @cute.jit, cute.arch functions)
+- CUDA programming concepts
+- GPU kernel optimization techniques
+- Code examples and patterns
+
+Available corpora:
+- 'cutlass' - NVIDIA CUTLASS C++ docs + GitHub examples (gemm, hopper, blackwell)
+- 'cutedsl' - CuTeDSL Python documentation
+- 'cuda' - General CUDA programming docs
+- 'hip' - AMD HIP programming docs
+- 'amd' - AMD GPU kernel development (rocWMMA, CK, etc.)
+
+Note: Corpora must be downloaded first with `wafer corpus download <name>`.
+Returns relevant documentation snippets with file paths.""",
+        parameters=ToolFunctionParameter(
+            type="object",
+            properties={
+                "query": {
+                    "type": "string",
+                    "description": "Search query - describe what you're looking for",
+                },
+                "corpus": {
+                    "type": "string",
+                    "description": "Which docs to search: 'cutlass', 'cutedsl', 'cuda', 'hip', 'amd' (default: cutlass)",
+                },
+                "max_results": {
+                    "type": "integer",
+                    "description": "Maximum number of results to return (default: 5)",
+                },
+            },
+        ),
+        required=["query"],
+    )
+)
+
+
+def _get_corpus_path(corpus_name: str) -> Path | None:
+    """Get the path to a corpus in the cache directory.
+    
+    Corpora are stored at ~/.cache/wafer/corpora/<corpus_name>/
+    """
+    if corpus_name not in AVAILABLE_CORPORA:
+        return None
+    
+    corpus_path = CACHE_DIR / corpus_name
+    if corpus_path.exists():
+        return corpus_path
+    
+    return None
+
+
+def _search_files(corpus_path: Path, query: str, max_results: int = 5) -> list[dict]:
+    """Simple keyword search through documentation files."""
+    results = []
+    query_terms = query.lower().split()
+    
+    # Search .md, .py, .cu, .hpp, and .h files (for CUTLASS examples)
+    for pattern in ["**/*.md", "**/*.py", "**/*.cu", "**/*.hpp", "**/*.h", "**/*.cuh"]:
+        for file_path in corpus_path.glob(pattern):
+            if file_path.is_file():
+                try:
+                    content = file_path.read_text(encoding="utf-8", errors="ignore")
+                    content_lower = content.lower()
+                    
+                    # Score based on term matches
+                    score = sum(content_lower.count(term) for term in query_terms)
+                    
+                    if score > 0:
+                        # Extract relevant snippets
+                        snippets = _extract_snippets(content, query_terms)
+                        results.append({
+                            "file": str(file_path),  # Return absolute path so read tool can access it
+                            "score": score,
+                            "snippets": snippets[:3],  # Top 3 snippets
+                        })
+                except Exception:
+                    continue
+    
+    # Sort by score and return top results
+    results.sort(key=lambda x: x["score"], reverse=True)
+    return results[:max_results]
+
+
+def _extract_snippets(content: str, terms: list[str], context_lines: int = 5) -> list[str]:
+    """Extract snippets containing search terms."""
+    snippets = []
+    lines = content.split("\n")
+    
+    for i, line in enumerate(lines):
+        line_lower = line.lower()
+        if any(term in line_lower for term in terms):
+            # Get context around the match
+            start = max(0, i - context_lines)
+            end = min(len(lines), i + context_lines + 1)
+            snippet = "\n".join(lines[start:end])
+            
+            # Skip very short snippets
+            if len(snippet.strip()) > 50:
+                snippets.append(snippet)
+    
+    return snippets
+
+
+async def exec_search_docs(
+    tool_call: ToolCall,
+    corpus_override: str | None = None,
+) -> ToolResult:
+    """Execute search_docs tool.
+    
+    Args:
+        tool_call: The tool call with query and optional corpus
+        corpus_override: Override corpus path (for testing)
+    """
+    query = tool_call.args.get("query", "")
+    corpus_name = tool_call.args.get("corpus", "cutlass")
+    max_results = tool_call.args.get("max_results", 5)
+    
+    if not query:
+        return ToolResult(
+            tool_call_id=tool_call.id,
+            content="",
+            error="query parameter is required",
+        )
+    
+    # Find corpus path
+    if corpus_override:
+        corpus_path = Path(corpus_override)
+    else:
+        corpus_path = _get_corpus_path(corpus_name)
+        if corpus_path is None:
+            return ToolResult(
+                tool_call_id=tool_call.id,
+                content="",
+                error=f"Unknown corpus: {corpus_name}. Available: {AVAILABLE_CORPORA}",
+            )
+    
+    if not corpus_path.exists():
+        return ToolResult(
+            tool_call_id=tool_call.id,
+            content="",
+            error=f"Corpus '{corpus_name}' not downloaded. Run: wafer corpus download {corpus_name}",
+        )
+    
+    # Search
+    results = _search_files(corpus_path, query, max_results)
+    
+    if not results:
+        return ToolResult(
+            tool_call_id=tool_call.id,
+            content=f"No results found for query: {query}",
+            error=None,
+        )
+    
+    # Format output
+    output_parts = [f"Found {len(results)} results for: {query}\n"]
+    
+    for i, result in enumerate(results, 1):
+        output_parts.append(f"\n{'='*60}")
+        output_parts.append(f"[{i}] {result['file']} (score: {result['score']})")
+        output_parts.append("=" * 60)
+        
+        for snippet in result["snippets"]:
+            output_parts.append(snippet)
+            output_parts.append("-" * 40)
+    
+    return ToolResult(
+        tool_call_id=tool_call.id,
+        content="\n".join(output_parts),
+        error=None,
+    )

wafer_core/utils/kernel_utils/static_checker.py

@@ -155,6 +155,161 @@ def check_torch_computation_ops(code: str) -> tuple[bool, str]:
     return (False, "")
 
 
+# =============================================================================
+# NN.MODULE FORWARD CALL CHECKS (Reward Hacking in forward())
+# =============================================================================
+
+# These patterns detect calling PyTorch nn.Module forward methods inside forward()
+# e.g., self.conv(x), self.linear(x), self.bn(x) - these invoke cuBLAS/cuDNN
+# 
+# This is different from:
+# - nn.Linear(...) in __init__ = OK (just creates parameter container)
+# - self.linear.weight in forward() = OK (accessing weights for custom kernel)
+# - self.linear(x) in forward() = BAD (invokes PyTorch's matmul via cuBLAS)
+
+NN_MODULE_FORWARD_PATTERNS = [
+    # Common layer types being called as functions
+    r"self\.(conv\d*d?|linear|bn|batch_norm|layer_norm|group_norm|instance_norm)\s*\(",
+    # More generic pattern: self.<name>(x) or self.<name>(input)
+    # But we need to be careful not to match custom module calls
+]
+
+# =============================================================================
+# TORCH.NN.FUNCTIONAL CHECKS (Reward Hacking)
+# =============================================================================
+
+# Patterns for torch.nn.functional / F.* calls that bypass custom kernel requirement
+# These call into cuBLAS/cuDNN under the hood
+TORCH_FUNCTIONAL_PATTERNS = [
+    # F.linear, F.conv*, F.batch_norm etc. (common alias)
+    r"\bF\.(linear|conv[123]d|conv_transpose[123]d|batch_norm|layer_norm|group_norm|instance_norm)\s*\(",
+    # Full path torch.nn.functional.*
+    r"\btorch\.nn\.functional\.(linear|conv[123]d|conv_transpose[123]d|batch_norm|layer_norm|group_norm|instance_norm)\s*\(",
+]
+
+
+def check_torch_functional_calls(code: str) -> tuple[bool, str]:
+    """Check for torch.nn.functional / F.* calls in forward() method (reward hacking).
+    
+    Detects patterns like F.linear(x, weight), F.conv2d(x, weight) which invoke
+    PyTorch's built-in operations (backed by cuBLAS/cuDNN) instead of custom kernels.
+    
+    Note: We only check inside forward() to avoid false positives from imports or
+    training-mode fallbacks that aren't used during inference.
+    """
+    # Only check inside forward() method
+    forward_code = _extract_forward_method(code)
+    if not forward_code:
+        return (False, "")
+    
+    forward_code = _strip_comments(forward_code)
+    
+    for pattern in TORCH_FUNCTIONAL_PATTERNS:
+        match = re.search(pattern, forward_code)
+        if match:
+            return (True, f"Uses torch.nn.functional in forward(): {match.group(0)} (reward hacking - must use custom kernel)")
+    
+    return (False, "")
+
+
+def _extract_forward_method(code: str) -> str:
+    """Extract only the forward() method content from code."""
+    lines = code.split('\n')
+    result = []
+    in_forward = False
+    base_indent = 0
+    
+    for i, line in enumerate(lines):
+        # Look for forward method definition
+        if re.search(r'\bdef\s+forward\s*\(\s*self', line):
+            in_forward = True
+            # Get the indentation level of the def line
+            base_indent = len(line) - len(line.lstrip())
+            result.append(line)
+            continue
+        
+        if in_forward:
+            # Check if we've exited the forward method (new method/class at same or lower indent)
+            stripped = line.strip()
+            if stripped and not stripped.startswith('#') and not stripped.startswith('"""') and not stripped.startswith("'''"):
+                current_indent = len(line) - len(line.lstrip())
+                # If we hit a new def/class at the same or lower indentation, we're done
+                if current_indent <= base_indent and (stripped.startswith('def ') or stripped.startswith('class ')):
+                    break
+            result.append(line)
+    
+    return '\n'.join(result)
+
+
+def check_nn_module_forward_call(code: str) -> tuple[bool, str]:
+    """Check for nn.Module forward calls inside forward() method (reward hacking).
+    
+    Detects patterns like self.conv(x), self.linear(x) which invoke PyTorch's
+    built-in layers (backed by cuBLAS/cuDNN) instead of custom kernels.
+    """
+    # Only check inside forward() method
+    forward_code = _extract_forward_method(code)
+    if not forward_code:
+        return (False, "")
+    
+    forward_code = _strip_comments(forward_code)
+    
+    for pattern in NN_MODULE_FORWARD_PATTERNS:
+        match = re.search(pattern, forward_code)
+        if match:
+            return (True, f"Calls PyTorch nn.Module in forward(): {match.group(0)} (reward hacking - must use custom kernel)")
+    
+    return (False, "")
+
+
+# =============================================================================
+# CUBLAS/CUDNN DIRECT USAGE CHECKS (Reward Hacking)
+# =============================================================================
+
+# Direct cuBLAS calls bypass custom kernel requirement
+CUBLAS_PATTERNS = [
+    r"\bcublas[A-Z]\w+\s*\(",  # cublasSgemm, cublasGemmEx, etc.
+    r"\bcublasCreate\b",
+    r"\bcublasDestroy\b",
+    r"\bcublasSetStream\b",
+    r"\bcublasSetMathMode\b",
+    r"#include\s*[<\"]cublas",  # #include <cublas_v2.h>
+    r"CUBLAS_TENSOR_OP_MATH",
+]
+
+# Direct cuDNN calls bypass custom kernel requirement
+CUDNN_PATTERNS = [
+    r"\bcudnn[A-Z]\w+\s*\(",  # cudnnConvolutionForward, etc.
+    r"\bcudnnCreate\b",
+    r"\bcudnnDestroy\b",
+    r"#include\s*[<\"]cudnn",  # #include <cudnn.h>
+]
+
+
+def check_cublas_usage(code: str) -> tuple[bool, str]:
+    """Check for direct cuBLAS API usage (reward hacking)."""
+    code = _strip_comments(code)
+    
+    for pattern in CUBLAS_PATTERNS:
+        match = re.search(pattern, code)
+        if match:
+            return (True, f"Uses cuBLAS directly: {match.group(0)} (reward hacking - must write custom kernel)")
+    
+    return (False, "")
+
+
+def check_cudnn_usage(code: str) -> tuple[bool, str]:
+    """Check for direct cuDNN API usage (reward hacking)."""
+    code = _strip_comments(code)
+    
+    for pattern in CUDNN_PATTERNS:
+        match = re.search(pattern, code)
+        if match:
+            return (True, f"Uses cuDNN directly: {match.group(0)} (reward hacking - must write custom kernel)")
+    
+    return (False, "")
+
+
 # =============================================================================
 # TIMING MANIPULATION CHECKS
 # =============================================================================
@@ -310,8 +465,16 @@ def check_tk_impl(code: str) -> tuple[bool, str]:
 def check_cute_impl(code: str) -> tuple[bool, str]:
     """Check for valid CUTLASS/CuTe kernel implementation."""
     code = _strip_comments(code)
-    if not any(p in code for p in ["cute::", "cutlass::", "from cutlass"]):
-        return (True, "Missing cute:: or cutlass:: namespace")
+    # Accept explicit namespace qualifiers OR using namespace declarations
+    valid_patterns = [
+        "cute::",
+        "cutlass::",
+        "from cutlass",
+        "using namespace cute",
+        "using namespace cutlass",
+    ]
+    if not any(p in code for p in valid_patterns):
+        return (True, "Missing cute:: or cutlass:: namespace (or 'using namespace')")
     return (False, "")
 
 
@@ -334,6 +497,11 @@ CHECK_FUNCTIONS: dict[str, Callable[[str], tuple[bool, str]]] = {
     "timing_event_patch": check_timing_event_patch,
     # Torch ops
     "torch_computation_ops": check_torch_computation_ops,
+    # Reward hacking checks
+    "cublas_usage": check_cublas_usage,
+    "cudnn_usage": check_cudnn_usage,
+    "nn_module_forward_call": check_nn_module_forward_call,
+    "torch_functional_calls": check_torch_functional_calls,
     # Timing manipulation
     "stream_injection": check_stream_injection,
     "thread_injection": check_thread_injection,
@@ -363,12 +531,16 @@ STRICT_CHECKS = [
     "timing_event_patch",
     "thread_injection",
     "lazy_eval",
+    "cublas_usage",  # Direct cuBLAS is reward hacking
+    "cudnn_usage",   # Direct cuDNN is reward hacking
+    "nn_module_forward_call",  # Calling self.conv(x), self.linear(x) in forward() is reward hacking
+    "torch_functional_calls",  # Calling F.linear(), F.conv2d() in forward() is reward hacking
+    "torch_computation_ops",  # torch.mm, torch.matmul, torch.conv* etc. are reward hacking
 ]
 
 # Checks that emit warnings but don't fail
 WARNING_CHECKS = [
     "pytorch_wrap",
-    "torch_computation_ops",
     "stream_injection",
 ]
 

wafer_core/utils/kernel_utils/targets/config.py

@@ -290,7 +290,7 @@ class RunPodTarget:
             ssh_key="~/.ssh/id_ed25519",
             gpu_type_id="AMD Instinct MI300X OAM",
             gpu_count=1,
-            image="runpod/pytorch:2.4.0-py3.10-rocm6.1.0-ubuntu22.04",
+            image="rocm/pytorch:rocm7.0.2_ubuntu24.04_py3.12_pytorch_release_2.7.1",
             keep_alive=True,  # Don't terminate after eval
         )
 
@@ -304,7 +304,21 @@ class RunPodTarget:
     gpu_type_id: str = AMD_MI300X_GPU_ID  # RunPod GPU type identifier
     gpu_count: int = 1
     container_disk_gb: int = 50
-    image: str = "runpod/pytorch:2.4.0-py3.10-rocm6.1.0-ubuntu22.04"
+    # TODO: Consider creating a custom Docker image with HipKittens pre-installed
+    # to avoid needing `wafer config targets install <target> hipkittens`.
+    # HipKittens repo: https://github.com/HazyResearch/hipkittens
+    # CK (Composable Kernel) is already included in ROCm 7.0.
+    #
+    # WARNING: PyTorch's hipify can corrupt /opt/rocm/include/thrust/ headers.
+    # If you see "cuda/__cccl_config not found" errors, run:
+    #   apt-get install --reinstall -y rocthrust
+    # See docker/rocm7-runpod/README.md for details.
+    image: str = "rocm/pytorch:rocm7.0.2_ubuntu24.04_py3.12_pytorch_release_2.7.1"
+
+    # RunPod template ID — required for non-RunPod images that need custom
+    # dockerArgs (e.g. to install and start sshd). When set, takes priority
+    # over `image` in the deploy mutation.
+    template_id: str | None = None
 
     # Timeouts
     provision_timeout: int = 900  # 15 min for SSH to be ready
@@ -426,7 +440,7 @@ class DigitalOceanTarget:
     # DigitalOcean instance configuration
     region: str = "atl1"  # Atlanta (AMD GPUs available here)
     size_slug: str = "gpu-mi300x1-192gb-devcloud"  # Single MI300X GPU
-    image: str = "gpu-amd-base"  # AMD AI/ML Ready image with ROCm
+    image: str = "amd-pytorchrocm7"  # PyTorch (ROCm7) marketplace image
 
     # Timeouts
     provision_timeout: int = 600  # 10 min for droplet to be ready

wafer_core/utils/modal_execution/modal_app.py

@@ -20,35 +20,17 @@ import modal
 
 # Build Modal image with all dependencies
 # This image is cached and reused across function invocations
-def build_modal_image(
-    gpu_type: str = "B200",
-    compute_capability: str = "10.0",
-) -> modal.Image:
+def build_modal_image() -> modal.Image:
     """Build Modal image with PyTorch, CUTLASS, and evaluation dependencies.
 
     Uses explicit local code inclusion to avoid pulling in SSH deployment code.
 
-    Phase 2 solution from MODAL_HANDOFF.md:
-    - Use add_local_dir with ignore parameter to exclude deployment files
-    - Only include files needed for kernel evaluation
-
-    Args:
-        gpu_type: GPU type (determines PyTorch index URL)
-        compute_capability: CUDA compute capability
-
     Returns:
         Modal Image ready for kernel evaluation
     """
-    # Determine PyTorch index based on GPU type
-    # Match logic from configs/base_config.py
-    if gpu_type in ["B200", "GB200"] or compute_capability.startswith("10."):
-        # Blackwell requires PyTorch 2.8+ with CUDA 12.8
-        torch_index = "https://download.pytorch.org/whl/nightly/cu128"
-        torch_version = "torch>=2.8.0"
-    else:
-        # Older GPUs (H100, A100) use stable PyTorch
-        torch_index = "https://download.pytorch.org/whl/cu124"
-        torch_version = "torch>=2.4.0"
+    # Use CUDA 13.0 for all GPUs (H100, A100, B200, GB200)
+    torch_index = "https://download.pytorch.org/whl/cu130"
+    torch_version = "torch>=2.6.0"
 
     # Build image with dependencies
     image = (
@@ -74,6 +56,15 @@ def build_modal_image(
             "scipy",
             "pytest",
         )
+        # Install CUTLASS headers for C++ kernel compilation (v4.3.5)
+        .run_commands(
+            "git clone --depth 1 --branch v4.3.5 https://github.com/NVIDIA/cutlass.git /usr/local/cutlass",
+            # Verify CUTLASS was installed correctly
+            "ls -la /usr/local/cutlass/include/cutlass/util/ | head -20",
+            "test -f /usr/local/cutlass/include/cutlass/util/packed_stride.hpp && echo 'CUTLASS headers OK' || echo 'CUTLASS headers MISSING'",
+        )
+        # Set CUTLASS_PATH environment variable
+        .env({"CUTLASS_PATH": "/usr/local/cutlass/include"})
         # Create empty __init__.py files for proper Python package structure
         # MUST run before add_local_* commands (Modal restriction)
         .run_commands(
@@ -111,20 +102,16 @@ def build_modal_image(
 # Create app (can be customized per target)
 def create_modal_app(
     app_name: str = "test-kernel-eval",  # Match test script default
-    gpu_type: str = "B200",
-    compute_capability: str = "10.0",
 ) -> modal.App:
     """Create Modal app for kernel evaluation.
 
     Args:
         app_name: Modal app name
-        gpu_type: GPU type for image building
-        compute_capability: CUDA compute capability
 
     Returns:
         Modal App instance
     """
-    image = build_modal_image(gpu_type=gpu_type, compute_capability=compute_capability)
+    image = build_modal_image()
     return modal.App(name=app_name, image=image)
 
 

{wafer_core-0.1.22.dist-info → wafer_core-0.1.23.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: wafer-core
-Version: 0.1.22
+Version: 0.1.23
 Summary: Core utilities and environments for Wafer GPU kernel optimization
 Requires-Python: >=3.10
 Requires-Dist: aiohttp>=3.9.0

{wafer_core-0.1.22.dist-info → wafer_core-0.1.23.dist-info}/RECORD

@@ -12,7 +12,7 @@ wafer_core/config/__init__.py,sha256=hKywfjA4YXd4lBeBFEcBoMwFoflPHJTiBnkTq7_JYOQ
 wafer_core/config/loader.py,sha256=k7JnILmO13TWUzIv9Lm8fvmj3UfYHZDgaFurjQ-GXpY,6623
 wafer_core/config/schema.py,sha256=2WhFlnG0VYYX4T-70BLeJK8Janvi4KEa8KKGZA7331w,3898
 wafer_core/environments/__init__.py,sha256=SIsResVtm22tr_d-oHPeeSxrkhFdmPOFico3DqtRqK8,238
-wafer_core/environments/coding.py,sha256=T-_JFU-n5OxPR8xAWp8qar4Y5xyC-TWTIBjRy4PDel8,8418
+wafer_core/environments/coding.py,sha256=N-ELZwJu5vKLCVtwO25c6JSty6fmqf85VR2d3WJ4RXw,8559
 wafer_core/environments/gpumode.py,sha256=8Da08nltvN_YloNyYI6-omN2D4n5C7aptKDCtUgT2bQ,17191
 wafer_core/lib/__init__.py,sha256=4-4p3mhwlquejWGglYXU8_nHdA0LoPaa_jGzcm13USA,1325
 wafer_core/lib/kernel_scope/__init__.py,sha256=WW2vu8jUlqOu-MCpgO40lIYacCA9N2u-uuECIs_JO2w,2817
@@ -585,11 +585,12 @@ wafer_core/sessions/dtypes.py,sha256=K6nOjvL6sjCGY7GTtdEygf1IZY_18R9YkHGqFyMd8wY
 wafer_core/sessions/hooks.py,sha256=A-txm6ufnRGQCdtP3vwh7oEOdlLN9Tv0XsjORMihuAI,4295
 wafer_core/targets/__init__.py,sha256=sHndC7AAOaHXlrmDXFLB53a5Y8DBjuyqS6nwsO2nj-Y,1728
 wafer_core/targets/digitalocean.py,sha256=cvoYpYjtSyy5t2lQAPi7ERruuuibronah_ivOiduAHQ,16550
-wafer_core/targets/runpod.py,sha256=bYTLVRaASrewJLIcZRtPMfMU2en1McE6W5w-Edo_OPQ,15785
-wafer_core/tools/__init__.py,sha256=wBQD45GdSfkxcT6NHzIv0IMeXCc0enwwkpm3T_9j1X8,3341
+wafer_core/targets/runpod.py,sha256=LrVmNvA6qjzL5nbGSWvtw7CHrK6bDu7_o3vKIek00Tc,20286
+wafer_core/tools/__init__.py,sha256=deGQQlcdSD6zQx8JHizfSXgF5-EntdBOF_ngtob1-VU,3506
 wafer_core/tools/bash_tool.py,sha256=daoKOVGSgL0x9X_3l8Apd6-wFH4VMXMGJwVemw2FIfc,16828
 wafer_core/tools/glob_tool.py,sha256=9X5PdOjQJj7kiVNqqCZC0-1LmnE6wHx3Zc9zfMjtXdc,3533
 wafer_core/tools/grep_tool.py,sha256=cStyDz-J47oDLLZCL83yOvYo8Ijv4qu3D372JKT_ptM,4580
+wafer_core/tools/search_docs_tool.py,sha256=WY4hY83sseX8Fpxvw6DZxiG-F95F2t3-4PyfMD1Lpkg,6809
 wafer_core/tools/skill_tool.py,sha256=JXsT5hBTUH5U4tmzHEywU7eHHt5xCEF79tL2tsuk4-c,2067
 wafer_core/tools/wafer_tool.py,sha256=-dgPTHbWXq3I3wFj0mP7-lj5iZqGRoFvFf9IEEo3plQ,6345
 wafer_core/tools/write_kernel_tool.py,sha256=dJjhr-WBhVNe06hcJQVmBZTbS8mid64KF1MwlE2s2R4,21547
@@ -615,7 +616,7 @@ wafer_core/tools/capture_tool/metrics.py,sha256=BFZNmdE-kh3LneYdWXTNZmlLuo-DCrP5
 wafer_core/tools/file_tools/__init__.py,sha256=2H7Rq5bijNQHGO4W6jjQAShkrcmdcHC0EQ8mBpgrApI,632
 wafer_core/tools/file_tools/edit_tool.py,sha256=Efx83pM1Ljb07cJmAGVhPX4YiPJICK70sZM6uCjRWB0,4109
 wafer_core/tools/file_tools/glob_tool.py,sha256=Av4LfC21fHXbnSsgh_9zDxlY9Qhb48aApaGos4j3B4g,3437
-wafer_core/tools/file_tools/grep_tool.py,sha256=FRIYeBfCcywHqkiT8OqVL8xOtMiZDOV6EhUU2LN5fqM,4444
+wafer_core/tools/file_tools/grep_tool.py,sha256=42eFj2pxBBrs5eg_GhyYJ-j2fNWkmGPvrEqXFmi5E10,5539
 wafer_core/tools/file_tools/read_tool.py,sha256=K0Hd8zwyL4Yva5YO9spXDfTRfXvfjqh9ztVrA8s1bJE,3961
 wafer_core/tools/file_tools/utils.py,sha256=HgaqYan2Pky4hTLX2L9d2Gj9oS325H7rFbJj-jryNtc,2576
 wafer_core/tools/file_tools/write_tool.py,sha256=X4N8y8wB-k9d5PcMRmZMRKIXlG9jHJiRdlEFFRLdZzs,2083
@@ -661,18 +662,18 @@ wafer_core/utils/kernel_utils/evaluate.py,sha256=1kxFNMl9VCXfKfk_BIiuA_zFfvDB1sl
 wafer_core/utils/kernel_utils/gpu_validation.py,sha256=LRiDjW_xAK4fXf1Vw1aYHG54B1W0J6b5L0K6PXzM2tI,3759
 wafer_core/utils/kernel_utils/reference_cache.py,sha256=4IQ2gND1StHULRO7geyAElEStbjQxwOeP6X09E5wCB0,11283
 wafer_core/utils/kernel_utils/results.py,sha256=QJGeah_41LSzxyYwGl9VxHPxTVAN2bLtk5bWdWLIpL4,6705
-wafer_core/utils/kernel_utils/static_checker.py,sha256=GWC7RZdwL4WqMLK0nO-wnZd-VmlMahCVKP6l9zVAW30,12909
+wafer_core/utils/kernel_utils/static_checker.py,sha256=XIQkzAOkGH5xtrOuZM4tNUqVJ0QRkYeJ7_8DosDOtkw,19886
 wafer_core/utils/kernel_utils/task.py,sha256=XcmKxKUWh5It6nX3zGqj77tWgA32uPfQMqNOqyD5T48,2682
 wafer_core/utils/kernel_utils/utils.py,sha256=uDZoJDxh07hJeLNlPdKN2vgB15pqIr1LbXf0YIBHU4E,43056
 wafer_core/utils/kernel_utils/targets/__init__.py,sha256=4NwRLsuJ__S4xKAfda4Ag82C5MQ3Qio-4xA5S-mQGlU,2067
-wafer_core/utils/kernel_utils/targets/config.py,sha256=3TvT2Hp3TV-cjsSiZ8NOmAz8epDoVii76p6DDAI2V64,19134
+wafer_core/utils/kernel_utils/targets/config.py,sha256=sNXyYTZ9rL9OET4xqbHZ0d4b8ChzST1yUKvNOv8JSQs,19933
 wafer_core/utils/kernel_utils/targets/execution.py,sha256=bZuNXCo0sIdD6hFhetLPrtDC-zMSiIsAx_aml49VVL0,15033
 wafer_core/utils/kernel_utils/targets/selection.py,sha256=5I_RG_7cfhq7uaeR28meC2EeNNKssFsK-Tc3QFG6Ze0,3590
 wafer_core/utils/modal_execution/__init__.py,sha256=jkVqYOLzCT5K73N9Od0UIUsx-99A0m6bpDrxfyXxQZ8,945
-wafer_core/utils/modal_execution/modal_app.py,sha256=ibBllC59R9bP9w4QweFocVGeSX5kDmUuDl39PPYgCpE,11796
+wafer_core/utils/modal_execution/modal_app.py,sha256=VfS2cX8gHtnlPXemmMcEwDPeQdhiv2tly3CifOyh9f4,11455
 wafer_core/utils/modal_execution/modal_config.py,sha256=7cGX9TGqilQ3qxI3OFGXV5orjtyRU-PEDOJ4vP2oxno,4421
 wafer_core/utils/modal_execution/modal_execution.py,sha256=gChjnV6jqA3A7IRP3DfvV5cSfm_MN0X4f7JZufXgdZE,24594
 wafer_core/utils/modal_execution/test_modal.py,sha256=_jqou_hrLs1Daf1590Pnb0a_lXMMa2rczAPpW9HpoNQ,8153
-wafer_core-0.1.22.dist-info/METADATA,sha256=wV6MLEufRIKPRacW_ErMoqgymAFNI2XgP4wobQoKjnM,1420
-wafer_core-0.1.22.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-wafer_core-0.1.22.dist-info/RECORD,,
+wafer_core-0.1.23.dist-info/METADATA,sha256=HnIqBmqEQ6t_dc54Rnyg_Wyy-HKuAr3XTmsEoJkjJLo,1420
+wafer_core-0.1.23.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+wafer_core-0.1.23.dist-info/RECORD,,