wafer-core 0.1.30__tar.gz → 0.1.32__tar.gz

{wafer_core-0.1.30 → wafer_core-0.1.32}/PKG-INFO

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

{wafer_core-0.1.30 → wafer_core-0.1.32}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "wafer-core"
-version = "0.1.30"
+version = "0.1.32"
 description = "Core utilities and environments for Wafer GPU kernel optimization"
 requires-python = ">=3.10"
 dependencies = [
@@ -100,13 +100,16 @@ ignore = [
 ]
 
 [tool.ruff.lint.per-file-ignores]
-"tests/**/*.py" = ["ANN201"]  # Don't require return type annotations in tests
+"tests/**/*.py" = ["ANN201", "ANN202"]  # Don't require return type annotations in tests
 "wafer_core/utils/modal_execution/*.py" = ["PLR0913"]  # Modal RPC functions have many parameters
 "wafer_core/utils/kernel_utils/targets/execution.py" = ["PLR0915"]  # run_operation_on_target has Modal+SSH paths (71 statements)
 "wafer_core/environments/gpumode.py" = ["ANN401"]  # Environment.exec_tool uses Any for rollouts compatibility
 "wafer_core/rollouts/providers/anthropic.py" = ["PLR0915", "PLR1702", "TRY002", "TRY004"]  # Complex provider with streaming, retry logic, and try/finally cleanup
 "wafer_core/sandbox/*.py" = ["ASYNC109", "ASYNC240"]  # Sandbox uses subprocess timeout (not trio.fail_after) intentionally
 "wafer_core/tools/bash_tool.py" = ["ASYNC109", "ASYNC221"]  # Bash tool uses subprocess with timeout
+"wafer_core/utils/kernel_utils/defense.py" = ["ANN002", "ANN003", "ANN202", "PLR0913"]  # Defense functions use *args/**kwargs passthrough and have many config params
+"wafer_core/utils/kernel_utils/test_reward_hacks.py" = ["ANN"]  # Test fixture file, annotations add noise
+"tests/test_reward_hack_defenses.py" = ["ANN001"]  # Parametrized test args don't need annotations
 
 [tool.ruff.lint.pylint]
 max-args = 7             # Max function arguments (Tiger Style: few parameters)

{wafer_core-0.1.30 → wafer_core-0.1.32}/pyproject.toml.bak

@@ -100,13 +100,16 @@ ignore = [
 ]
 
 [tool.ruff.lint.per-file-ignores]
-"tests/**/*.py" = ["ANN201"]  # Don't require return type annotations in tests
+"tests/**/*.py" = ["ANN201", "ANN202"]  # Don't require return type annotations in tests
 "wafer_core/utils/modal_execution/*.py" = ["PLR0913"]  # Modal RPC functions have many parameters
 "wafer_core/utils/kernel_utils/targets/execution.py" = ["PLR0915"]  # run_operation_on_target has Modal+SSH paths (71 statements)
 "wafer_core/environments/gpumode.py" = ["ANN401"]  # Environment.exec_tool uses Any for rollouts compatibility
 "wafer_core/rollouts/providers/anthropic.py" = ["PLR0915", "PLR1702", "TRY002", "TRY004"]  # Complex provider with streaming, retry logic, and try/finally cleanup
 "wafer_core/sandbox/*.py" = ["ASYNC109", "ASYNC240"]  # Sandbox uses subprocess timeout (not trio.fail_after) intentionally
 "wafer_core/tools/bash_tool.py" = ["ASYNC109", "ASYNC221"]  # Bash tool uses subprocess with timeout
+"wafer_core/utils/kernel_utils/defense.py" = ["ANN002", "ANN003", "ANN202", "PLR0913"]  # Defense functions use *args/**kwargs passthrough and have many config params
+"wafer_core/utils/kernel_utils/test_reward_hacks.py" = ["ANN"]  # Test fixture file, annotations add noise
+"tests/test_reward_hack_defenses.py" = ["ANN001"]  # Parametrized test args don't need annotations
 
 [tool.ruff.lint.pylint]
 max-args = 7             # Max function arguments (Tiger Style: few parameters)

wafer_core-0.1.32/tests/test_reward_hack_defenses.py

@@ -0,0 +1,185 @@
+"""
+Integration tests for reward hack defenses.
+
+Tests the two cut points:
+1. validate_kernel_static() - catches hacks in source code
+2. run_all_defenses_extended() - catches hacks at runtime
+
+Each test runs a known attack through the full defense pipeline.
+Requires torch + CUDA — skipped entirely in CI environments without GPU.
+"""
+
+import pytest
+
+torch = pytest.importorskip("torch", reason="torch required for reward hack defense tests")
+
+from wafer_core.utils.kernel_utils.defense import (  # noqa: E402
+    run_all_defenses,
+    run_all_defenses_extended,
+)
+from wafer_core.utils.kernel_utils.static_checker import validate_kernel_static  # noqa: E402
+from wafer_core.utils.kernel_utils.test_reward_hacks import (  # noqa: E402
+    STATIC_HACK_SAMPLES,
+    caching,
+    clear_cache,
+    hardcoded_256,
+    identity_kernel,
+    lazy_evaluation,
+    legit_matmul,
+    precision_downgrade,
+    thread_injection,
+)
+
+requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
+
+
+# =============================================================================
+# Cut point 1: validate_kernel_static catches code-level hacks
+# =============================================================================
+
+
+@pytest.mark.parametrize(
+    "name,code,expected_checks",
+    [(name, code, checks) for name, (code, checks) in STATIC_HACK_SAMPLES.items()],
+    ids=list(STATIC_HACK_SAMPLES.keys()),
+)
+def test_static_checker_catches_hack(name, code, expected_checks):
+    """Each hack code sample should fail validate_kernel_static."""
+    valid, errors, warnings = validate_kernel_static(code)
+    # At least one of the expected checks should have fired
+    assert not valid or len(warnings) > 0, (
+        f"Static checker missed {name}: expected checks {expected_checks} to fire, "
+        f"got errors={errors}, warnings={warnings}"
+    )
+
+
+def test_static_checker_passes_clean_code():
+    """Non-hack code should pass static checks."""
+    clean = "@triton.jit\ndef kernel(ptr, BLOCK: tl.constexpr):\n    tl.store(ptr, tl.load(ptr))\n"
+    valid, errors, warnings = validate_kernel_static(clean, backend="triton")
+    assert valid, f"Clean code should pass: errors={errors}"
+
+
+# =============================================================================
+# Cut point 2: run_all_defenses / run_all_defenses_extended
+# =============================================================================
+
+
+@requires_cuda
+@pytest.mark.parametrize(
+    "name,attack_fn",
+    [
+        ("thread_injection", thread_injection),
+        ("lazy_evaluation", lazy_evaluation),
+        ("precision_downgrade", precision_downgrade),
+    ],
+)
+def test_runtime_defense_catches_timing_attack(name, attack_fn):
+    """Timing attacks are caught by run_all_defenses."""
+    A = torch.randn(256, 256, device="cuda")
+    B = torch.randn(256, 256, device="cuda")
+    torch.cuda.synchronize()
+
+    all_passed, results, _ = run_all_defenses(legit_matmul, A, B)
+    # Sanity: legit passes
+    assert all_passed
+
+    all_passed, results, _ = run_all_defenses(attack_fn, A, B)
+    assert not all_passed, f"{name} should be caught"
+
+
+@requires_cuda
+def test_monkey_patch_detected():
+    """Monkey-patching of timing functions is caught."""
+    original = torch.cuda.Event.elapsed_time
+    try:
+        torch.cuda.Event.elapsed_time = lambda self, end: 0.001
+        A = torch.randn(256, 256, device="cuda")
+        B = torch.randn(256, 256, device="cuda")
+        all_passed, results, _ = run_all_defenses(legit_matmul, A, B)
+        assert not all_passed
+    finally:
+        torch.cuda.Event.elapsed_time = original
+
+
+@requires_cuda
+def test_caching_attack_caught():
+    """Caching/memoization is caught by multi-input validation."""
+    clear_cache()
+
+    def input_gen():
+        return (
+            torch.randn(256, 256, device="cuda"),
+            torch.randn(256, 256, device="cuda"),
+        )
+
+    A, B = input_gen()
+    all_passed, results, _ = run_all_defenses_extended(
+        caching,
+        A,
+        B,
+        reference_fn=torch.matmul,
+        input_generator=input_gen,
+    )
+    assert not all_passed, "Caching attack should be caught"
+
+
+@requires_cuda
+def test_hardcoded_shapes_caught():
+    """Hardcoded optimization is caught by multi-shape validation."""
+    A = torch.randn(256, 256, device="cuda")
+    B = torch.randn(256, 256, device="cuda")
+
+    all_passed, results, _ = run_all_defenses_extended(
+        hardcoded_256,
+        A,
+        B,
+        reference_fn=torch.matmul,
+        test_shapes=[(128, 128), (256, 256), (512, 512)],
+    )
+    assert not all_passed, "Hardcoded shapes attack should be caught"
+
+
+@requires_cuda
+def test_identity_kernel_caught():
+    """Identity kernel is caught when reference produces different output."""
+    x = torch.randn(256, 256, device="cuda")
+
+    all_passed, results, _ = run_all_defenses_extended(
+        identity_kernel,
+        x,
+        reference_fn=torch.relu,
+    )
+    assert not all_passed, "Identity kernel should be caught"
+
+
+@requires_cuda
+def test_extended_pipeline_legit_with_all_defenses():
+    """Legit kernel passes full extended pipeline with all defense flags enabled.
+
+    This is the actual code path used by `wafer evaluate kernelbench`.
+    Exercises: base timing defenses, caching, shape, identity, GEMM exact
+    correctness, ULP precision, and memory guards — all through the single
+    run_all_defenses_extended() entry point.
+    """
+
+    def input_gen():
+        return (
+            torch.randn(256, 256, device="cuda"),
+            torch.randn(256, 256, device="cuda"),
+        )
+
+    A, B = input_gen()
+    all_passed, results, _ = run_all_defenses_extended(
+        legit_matmul,
+        A,
+        B,
+        reference_fn=torch.matmul,
+        input_generator=input_gen,
+        test_shapes=[(128, 128), (256, 256)],
+        is_gemm=True,
+        gemm_dims=(128, 128, 128),
+        check_precision_ulp=True,
+    )
+    failed = [(name, msg) for name, passed, msg in results if not passed]
+    assert all_passed, f"Legit kernel failed full extended pipeline: {failed}"

{wafer_core-0.1.30 → wafer_core-0.1.32}/tests/trace_compare/test_trace_compare_correctness.py

@@ -34,6 +34,18 @@ NVIDIA_LLAMA_TRACE = TRACE_EXAMPLES_DIR / "nvidia_llama.json"
 GOLDEN_FILE = Path(__file__).parent / "expected_output_amd_llama.json"
 
 
+def _path_exists_safe(path: Path) -> bool:
+    """Check if path exists, returning False on permission errors.
+    
+    CI runners may not have permission to stat /root/... paths,
+    so we treat PermissionError the same as file not existing.
+    """
+    try:
+        return path.exists()
+    except PermissionError:
+        return False
+
+
 def load_golden_file() -> dict[str, Any]:
     """Load expected output from golden file."""
     if not GOLDEN_FILE.exists():
@@ -50,7 +62,7 @@ def load_golden_file() -> dict[str, Any]:
 @pytest.fixture(scope="session")
 def amd_trace_result():
     """Load AMD trace once per test session (~17s)."""
-    if not AMD_LLAMA_TRACE.exists():
+    if not _path_exists_safe(AMD_LLAMA_TRACE):
         pytest.skip(f"Trace file not found: {AMD_LLAMA_TRACE}")
     return load_trace(AMD_LLAMA_TRACE, include_stacks=True)
 
@@ -58,7 +70,7 @@ def amd_trace_result():
 @pytest.fixture(scope="session")
 def amd_trace_streaming_result():
     """Load AMD trace with streaming loader once per test session."""
-    if not AMD_LLAMA_TRACE.exists():
+    if not _path_exists_safe(AMD_LLAMA_TRACE):
         pytest.skip(f"Trace file not found: {AMD_LLAMA_TRACE}")
     return load_trace_streaming(AMD_LLAMA_TRACE, include_stacks=True)
 
@@ -66,7 +78,7 @@ def amd_trace_streaming_result():
 @pytest.fixture(scope="session")
 def analyze_result():
     """Run analyze_traces once per test session (~35s)."""
-    if not AMD_LLAMA_TRACE.exists() or not NVIDIA_LLAMA_TRACE.exists():
+    if not _path_exists_safe(AMD_LLAMA_TRACE) or not _path_exists_safe(NVIDIA_LLAMA_TRACE):
         pytest.skip("Trace files not found")
     return analyze_traces(
         AMD_LLAMA_TRACE,
@@ -165,7 +177,7 @@ def test_streaming_loader_matches_regular_loader(amd_trace_result, amd_trace_str
 
 def test_metadata_extraction_matches_full_load(amd_trace_result) -> None:
     """Test that fast metadata extraction matches full load metadata."""
-    if not AMD_LLAMA_TRACE.exists():
+    if not _path_exists_safe(AMD_LLAMA_TRACE):
         pytest.skip(f"Trace file not found: {AMD_LLAMA_TRACE}")
     
     # Extract metadata fast (~2ms)
@@ -186,7 +198,7 @@ def test_metadata_extraction_matches_full_load(amd_trace_result) -> None:
 
 def test_streaming_metadata_callback_fires() -> None:
     """Test that streaming loader fires metadata callback (fast, no full load)."""
-    if not AMD_LLAMA_TRACE.exists():
+    if not _path_exists_safe(AMD_LLAMA_TRACE):
         pytest.skip(f"Trace file not found: {AMD_LLAMA_TRACE}")
     
     # Only test the metadata callback - don't do full load
@@ -329,7 +341,7 @@ def test_architecture_detection_unknown() -> None:
 
 def test_analyze_trace_pair_metadata_callback() -> None:
     """Test that metadata extraction works for both traces (fast, no full analysis)."""
-    if not AMD_LLAMA_TRACE.exists() or not NVIDIA_LLAMA_TRACE.exists():
+    if not _path_exists_safe(AMD_LLAMA_TRACE) or not _path_exists_safe(NVIDIA_LLAMA_TRACE):
         pytest.skip("Trace files not found")
     
     # Just test metadata extraction - much faster than full analysis
@@ -415,7 +427,7 @@ def test_loaded_trace_contains_all_data() -> None:
     """Test that LoadedTrace contains all required fields."""
     from wafer_core.lib.trace_compare.loader import load_trace_full, LoadedTrace
     
-    if not AMD_LLAMA_TRACE.exists():
+    if not _path_exists_safe(AMD_LLAMA_TRACE):
         pytest.skip(f"Trace file not found: {AMD_LLAMA_TRACE}")
     
     loaded = load_trace_full(AMD_LLAMA_TRACE, include_stacks=True)
@@ -441,7 +453,7 @@ def test_analyze_from_loaded_matches_path_api() -> None:
     from wafer_core.lib.trace_compare.loader import load_trace_full
     from wafer_core.lib.trace_compare.analyzer import analyze_traces_from_loaded
     
-    if not AMD_LLAMA_TRACE.exists() or not NVIDIA_LLAMA_TRACE.exists():
+    if not _path_exists_safe(AMD_LLAMA_TRACE) or not _path_exists_safe(NVIDIA_LLAMA_TRACE):
         pytest.skip("Trace files not found")
     
     # Load traces once
@@ -472,7 +484,7 @@ def test_analyze_trace_pair_timing() -> None:
     """Test that analyze_trace_pair completes in reasonable time (<25s)."""
     import time
     
-    if not AMD_LLAMA_TRACE.exists() or not NVIDIA_LLAMA_TRACE.exists():
+    if not _path_exists_safe(AMD_LLAMA_TRACE) or not _path_exists_safe(NVIDIA_LLAMA_TRACE):
         pytest.skip("Trace files not found")
     
     start = time.perf_counter()

wafer_core-0.1.32/wafer_core/tools/dispatch_baseline/__init__.py

@@ -0,0 +1,73 @@
+"""Kernel trace tool for discovering SOTA kernels.
+
+Given a PyTorch operation, traces what kernel PyTorch dispatches to on target hardware.
+
+Usage:
+    from wafer_core.tools.dispatch_baseline import quick_trace
+
+    result = quick_trace(
+        "torch.matmul(A, B)",
+        {"A": (4096, 4096), "B": (4096, 4096)},
+        hardware="H100",
+    )
+    print(result.summary())
+"""
+
+from wafer_core.tools.dispatch_baseline.client import (
+    lookup_baseline,
+    store_baseline,
+)
+from wafer_core.tools.dispatch_baseline.codegen import (
+    generate_trace_script,
+    parse_op_string,
+    update_dtypes,
+    update_shapes,
+)
+from wafer_core.tools.dispatch_baseline.dtypes import (
+    HardwareSpec,
+    KernelInfo,
+    KernelTraceConfig,
+    KernelTraceResult,
+    OpSpec,
+    RooflineAnalysis,
+    TensorSpec,
+)
+from wafer_core.tools.dispatch_baseline.executor import (
+    TraceExecutionResult,
+    quick_trace,
+    trace_kernel_local,
+    trace_kernel_remote,
+)
+from wafer_core.tools.dispatch_baseline.roofline import (
+    HARDWARE_SPECS,
+    compute_roofline,
+    get_hardware_spec,
+)
+
+__all__ = [
+    # Data types
+    "HardwareSpec",
+    "KernelInfo",
+    "KernelTraceConfig",
+    "KernelTraceResult",
+    "OpSpec",
+    "RooflineAnalysis",
+    "TensorSpec",
+    "TraceExecutionResult",
+    # Codegen
+    "generate_trace_script",
+    "parse_op_string",
+    "update_dtypes",
+    "update_shapes",
+    # Execution
+    "quick_trace",
+    "trace_kernel_local",
+    "trace_kernel_remote",
+    # Database
+    "lookup_baseline",
+    "store_baseline",
+    # Roofline
+    "HARDWARE_SPECS",
+    "compute_roofline",
+    "get_hardware_spec",
+]

wafer_core-0.1.32/wafer_core/tools/dispatch_baseline/analyzer.py

@@ -0,0 +1,174 @@
+"""Analyzer for kernel trace output.
+
+Parses profiler output to extract kernel information and identify the primary kernel.
+"""
+
+import json
+import re
+from dataclasses import dataclass
+from typing import Any
+
+from wafer_core.tools.dispatch_baseline.dtypes import KernelInfo, KernelTraceResult, OpSpec
+
+
+@dataclass(frozen=True)
+class ParsedTraceResult:
+    """Parsed trace output with environment info for caching."""
+
+    result: KernelTraceResult
+    pytorch_version: str
+    runtime_version: str
+    gpu_arch: str
+
+
+def parse_trace_output(output: str, op_spec: OpSpec, hardware: str) -> ParsedTraceResult:
+    """Parse kernel trace script output and extract kernel information.
+
+    Args:
+        output: Raw output from the trace script (stdout)
+        op_spec: The operation that was traced
+        hardware: Target hardware name
+
+    Returns:
+        ParsedTraceResult with extracted kernel information and environment info
+    """
+    # Look for our JSON marker in the output
+    json_match = re.search(r"KERNEL_TRACE_RESULT_JSON:(.+)$", output, re.MULTILINE)
+
+    if not json_match:
+        return ParsedTraceResult(
+            result=KernelTraceResult(
+                op_spec=op_spec,
+                hardware=hardware,
+                kernels=[],
+                primary_kernel=None,
+                raw_output=output,
+                error="Could not find KERNEL_TRACE_RESULT_JSON marker in output",
+            ),
+            pytorch_version="unknown",
+            runtime_version="unknown",
+            gpu_arch="unknown",
+        )
+
+    try:
+        result_json = json.loads(json_match.group(1))
+    except json.JSONDecodeError as e:
+        return ParsedTraceResult(
+            result=KernelTraceResult(
+                op_spec=op_spec,
+                hardware=hardware,
+                kernels=[],
+                primary_kernel=None,
+                raw_output=output,
+                error=f"Failed to parse JSON: {e}",
+            ),
+            pytorch_version="unknown",
+            runtime_version="unknown",
+            gpu_arch="unknown",
+        )
+
+    # Extract environment info
+    env_info = result_json.get("environment", {})
+    pytorch_version = env_info.get("pytorch_version", "unknown")
+    runtime_version = env_info.get("runtime_version", "unknown")
+    gpu_arch = env_info.get("gpu_arch", "unknown")
+
+    # Extract kernels from JSON
+    kernels = []
+    for k in result_json.get("kernels", []):
+        kernel_info = _parse_kernel_dict(k)
+        if kernel_info:
+            kernels.append(kernel_info)
+
+    # Identify primary kernel (longest duration)
+    primary_kernel = kernels[0] if kernels else None
+
+    return ParsedTraceResult(
+        result=KernelTraceResult(
+            op_spec=op_spec,
+            hardware=hardware,
+            kernels=kernels,
+            primary_kernel=primary_kernel,
+            raw_output=output,
+        ),
+        pytorch_version=pytorch_version,
+        runtime_version=runtime_version,
+        gpu_arch=gpu_arch,
+    )
+
+
+def _parse_kernel_dict(kernel_dict: dict[str, Any]) -> KernelInfo | None:
+    """Parse a kernel dictionary into KernelInfo."""
+    name = kernel_dict.get("name", "")
+    if not name:
+        return None
+
+    return KernelInfo(
+        name=name,
+        duration_us=kernel_dict.get("duration_us", 0.0),
+    )
+
+
+def extract_kernels_from_nsys(nsys_output: str) -> list[dict[str, Any]]:
+    """Extract kernel information from nsys output.
+
+    This is a fallback for when torch.profiler doesn't capture all kernels.
+
+    Args:
+        nsys_output: Output from `nsys stats` or similar
+
+    Returns:
+        List of kernel dictionaries
+    """
+    kernels = []
+
+    # Look for GPU kernel summary lines
+    # Format varies, but typically: name, duration, count, etc.
+    lines = nsys_output.split("\n")
+
+    for line in lines:
+        # Skip headers and empty lines
+        if not line.strip() or "---" in line or "Name" in line:
+            continue
+
+        # Try to parse as nsys kernel line
+        # This is a simplified parser - real nsys output varies
+        parts = line.split()
+        if len(parts) >= 2:
+            # Heuristic: last column is often the kernel name
+            name = parts[-1]
+            if any(c.isalpha() for c in name) and "_" in name:
+                kernels.append({"name": name, "duration_us": 0.0})
+
+    return kernels
+
+
+def merge_kernel_infos(
+    profiler_kernels: list[KernelInfo], nsys_kernels: list[dict[str, Any]]
+) -> list[KernelInfo]:
+    """Merge kernel info from multiple sources.
+
+    Prefers profiler data when available, but adds any missing kernels from nsys.
+
+    Args:
+        profiler_kernels: Kernels from torch.profiler
+        nsys_kernels: Kernels from nsys
+
+    Returns:
+        Merged list of KernelInfo objects
+    """
+    seen_names = {k.name for k in profiler_kernels}
+    result = list(profiler_kernels)
+
+    for nsys_kernel in nsys_kernels:
+        name = nsys_kernel.get("name", "")
+        if name and name not in seen_names:
+            result.append(
+                KernelInfo(
+                    name=name,
+                    duration_us=nsys_kernel.get("duration_us", 0.0),
+                )
+            )
+            seen_names.add(name)
+
+    return result