wafer-core 0.1.18__tar.gz → 0.1.39__tar.gz

{wafer_core-0.1.18 → wafer_core-0.1.39}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: wafer-core
-Version: 0.1.18
+Version: 0.1.39
 Summary: Core utilities and environments for Wafer GPU kernel optimization
 Requires-Python: >=3.10
 Requires-Dist: aiohttp>=3.9.0
@@ -15,12 +15,14 @@ Requires-Dist: dash-svg>=0.0.11
 Requires-Dist: dash>=3.0.0
 Requires-Dist: dnspython>=2.8.0
 Requires-Dist: httpx>=0.25.0
+Requires-Dist: ijson>=3.2.0
 Requires-Dist: kaleido==0.2.1
 Requires-Dist: markdownify>=0.11.0
 Requires-Dist: matplotlib>=3.0.0
 Requires-Dist: modal>=0.64.0
 Requires-Dist: numpy>=1.17.5
 Requires-Dist: openai>=1.0.0
+Requires-Dist: orjson>=3.9.0
 Requires-Dist: pandas~=3.0.0
 Requires-Dist: paramiko>=3.0.0
 Requires-Dist: paramiko>=3.4.0

{wafer_core-0.1.18 → wafer_core-0.1.39}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "wafer-core"
-version = "0.1.18"
+version = "0.1.39"
 description = "Core utilities and environments for Wafer GPU kernel optimization"
 requires-python = ">=3.10"
 dependencies = [
@@ -28,6 +28,8 @@ dependencies = [
     "dash-bootstrap-components>=1.0.0", # UI components
     "dash-svg>=0.0.11", # SVG components for Dash
     "pandas~=3.0.0", # Data processing - allows 3.0.x patches
+    "orjson>=3.9.0", # Fast JSON parsing (10-30x faster than stdlib)
+    "ijson>=3.2.0", # Streaming JSON parser for progressive loading
     "numpy>=1.17.5", # Numerical computing
     "matplotlib>=3.0.0", # Additional plotting
     "tabulate>=0.8.0", # Table formatting
@@ -98,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.18 → wafer_core-0.1.39}/pyproject.toml.bak

@@ -28,6 +28,8 @@ dependencies = [
     "dash-bootstrap-components>=1.0.0", # UI components
     "dash-svg>=0.0.11", # SVG components for Dash
     "pandas~=3.0.0", # Data processing - allows 3.0.x patches
+    "orjson>=3.9.0", # Fast JSON parsing (10-30x faster than stdlib)
+    "ijson>=3.2.0", # Streaming JSON parser for progressive loading
     "numpy>=1.17.5", # Numerical computing
     "matplotlib>=3.0.0", # Additional plotting
     "tabulate>=0.8.0", # Table formatting
@@ -98,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.39/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.39/tests/trace_compare/__init__.py

@@ -0,0 +1 @@
+"""Tests for trace comparison functionality."""

wafer_core-0.1.39/tests/trace_compare/benchmark_trace_compare.py

@@ -0,0 +1,183 @@
+"""Benchmark script for trace comparison performance.
+
+Measures performance for:
+- JSON parsing (orjson)
+- Phase lookup (binary search)
+- Full trace loading
+
+Run with:
+    python -m tests.trace_compare.benchmark_trace_compare
+"""
+
+import json
+import time
+from pathlib import Path
+
+import orjson
+
+from wafer_core.lib.trace_compare.loader import load_trace, _process_events_single_pass
+
+
+TRACE_EXAMPLES_DIR = Path("/root/wafer/experiments/ian/vllm-trace-compare/examples")
+TRACE_FILES = {
+    "small": TRACE_EXAMPLES_DIR / "amd_llama.json",
+    "medium": TRACE_EXAMPLES_DIR / "nvidia_llama.json",
+    "large": TRACE_EXAMPLES_DIR / "amd_oss.json",
+}
+
+
+def benchmark_json_parsing(file_path: Path, num_runs: int = 3) -> dict[str, float]:
+    """Benchmark JSON parsing: orjson vs stdlib json."""
+    print(f"\n📊 Benchmarking JSON parsing: {file_path.name}")
+    print(f"   File size: {file_path.stat().st_size / (1024**2):.1f} MB")
+    
+    with open(file_path, "rb") as f:
+        content_bytes = f.read()
+    content_str = content_bytes.decode("utf-8")
+    
+    orjson_times = []
+    for _ in range(num_runs):
+        start = time.perf_counter()
+        _ = orjson.loads(content_bytes)
+        orjson_times.append(time.perf_counter() - start)
+    
+    json_times = []
+    for _ in range(num_runs):
+        start = time.perf_counter()
+        _ = json.loads(content_str)
+        json_times.append(time.perf_counter() - start)
+    
+    orjson_avg = sum(orjson_times) / len(orjson_times)
+    json_avg = sum(json_times) / len(json_times)
+    speedup = json_avg / orjson_avg if orjson_avg > 0 else 0
+    
+    print(f"   orjson:     {orjson_avg:.2f}s (avg of {num_runs} runs)")
+    print(f"   stdlib json: {json_avg:.2f}s (avg of {num_runs} runs)")
+    print(f"   Speedup:    {speedup:.1f}x")
+    
+    return {
+        "orjson_time": orjson_avg,
+        "json_time": json_avg,
+        "speedup": speedup,
+    }
+
+
+def benchmark_phase_lookup(file_path: Path, num_samples: int = 10000) -> dict[str, float]:
+    """Benchmark phase lookup using binary search."""
+    print(f"\n📊 Benchmarking phase lookup: {file_path.name}")
+    
+    with open(file_path, "rb") as f:
+        trace = orjson.loads(f.read())
+    
+    events = trace.get("traceEvents", [])
+    pass_result = _process_events_single_pass(events, include_stacks=False)
+    sorted_phases = sorted(pass_result.phases, key=lambda p: p["ts_start"])
+    
+    import bisect
+    phase_starts = [p["ts_start"] for p in sorted_phases]
+    phase_types = [p["type"] for p in sorted_phases]
+    phase_ends = [p["ts_end"] for p in sorted_phases]
+    
+    def _get_phase_binary(ts: int) -> str:
+        if not phase_starts:
+            return "decode"
+        idx = bisect.bisect_right(phase_starts, ts) - 1
+        if idx >= 0 and phase_starts[idx] <= ts <= phase_ends[idx]:
+            return phase_types[idx]
+        return "decode"
+    
+    kernel_timestamps = [
+        ev["ts"] for ev in pass_result.kernel_events[:num_samples]
+    ]
+    
+    start = time.perf_counter()
+    for ts in kernel_timestamps:
+        _ = _get_phase_binary(ts)
+    binary_time = time.perf_counter() - start
+    
+    print(f"   Samples:    {len(kernel_timestamps)} kernel timestamps")
+    print(f"   Binary:     {binary_time:.3f}s")
+    
+    return {
+        "binary_time": binary_time,
+        "samples": len(kernel_timestamps),
+    }
+
+
+def benchmark_load_trace(file_path: Path, num_runs: int = 2) -> dict[str, float]:
+    """Benchmark full trace loading."""
+    print(f"\n📊 Benchmarking full trace loading: {file_path.name}")
+    
+    opt_times = []
+    for _ in range(num_runs):
+        start = time.perf_counter()
+        _ = load_trace(file_path, include_stacks=True)
+        opt_times.append(time.perf_counter() - start)
+    
+    opt_avg = sum(opt_times) / len(opt_times)
+    
+    print(f"   Load time:  {opt_avg:.2f}s (avg of {num_runs} runs)")
+    
+    return {
+        "load_time": opt_avg,
+    }
+
+
+def main() -> None:
+    """Run all benchmarks."""
+    print("=" * 70)
+    print("Trace Compare Performance Benchmarks")
+    print("=" * 70)
+    
+    available_files = {name: path for name, path in TRACE_FILES.items() if path.exists()}
+    
+    if not available_files:
+        print("\n❌ No trace files found in:", TRACE_EXAMPLES_DIR)
+        return
+    
+    print(f"\n✅ Found {len(available_files)} trace file(s)")
+    
+    test_file = min(available_files.values(), key=lambda p: p.stat().st_size)
+    print(f"\n📁 Using test file: {test_file.name} ({test_file.stat().st_size / (1024**2):.1f} MB)")
+    
+    results = {}
+    
+    try:
+        results["json_parsing"] = benchmark_json_parsing(test_file, num_runs=3)
+    except Exception as e:
+        print(f"   ❌ Error: {e}")
+        results["json_parsing"] = None
+    
+    try:
+        results["phase_lookup"] = benchmark_phase_lookup(test_file, num_samples=10000)
+    except Exception as e:
+        print(f"   ❌ Error: {e}")
+        results["phase_lookup"] = None
+    
+    try:
+        results["load_trace"] = benchmark_load_trace(test_file, num_runs=1)
+    except Exception as e:
+        print(f"   ❌ Error: {e}")
+        results["load_trace"] = None
+    
+    print("\n" + "=" * 70)
+    print("Summary")
+    print("=" * 70)
+    
+    json_parsing = results.get("json_parsing")
+    if json_parsing is not None:
+        print(f"JSON Parsing Speedup: {json_parsing['speedup']:.1f}x")
+    
+    phase_lookup = results.get("phase_lookup")
+    if phase_lookup is not None:
+        print(f"Phase Lookup: {phase_lookup['binary_time']:.3f}s ({phase_lookup['samples']} samples)")
+    
+    load_trace_result = results.get("load_trace")
+    if load_trace_result is not None:
+        print(f"Full Load Time: {load_trace_result['load_time']:.2f}s")
+    
+    print("\n✅ Benchmarks complete!")
+
+
+if __name__ == "__main__":
+    main()