wafer-core 0.1.31__py3-none-any.whl → 0.1.33__py3-none-any.whl

wafer_core/utils/kernel_utils/test_reward_hacks.py

@@ -0,0 +1,140 @@
+"""
+Reward hack attack implementations for testing defense mechanisms.
+
+Each attack is a callable that produces correct-looking output while cheating.
+Sources: Makora taxonomy (Jan 2026), CUDA-L2, Kevin paper.
+"""
+
+import threading
+
+import torch
+
+# -- Legitimate --
+
+
+def legit_matmul(A: torch.Tensor, B: torch.Tensor) -> torch.Tensor:
+    return torch.matmul(A, B)
+
+
+def legit_relu(x: torch.Tensor) -> torch.Tensor:
+    return torch.relu(x)
+
+
+# -- Timing attacks (CUDA-L2) --
+
+
+def stream_injection(A: torch.Tensor, B: torch.Tensor) -> torch.Tensor:
+    stream = torch.cuda.Stream()
+    with torch.cuda.stream(stream):
+        return torch.matmul(A, B)
+
+
+def thread_injection(A: torch.Tensor, B: torch.Tensor) -> torch.Tensor:
+    out = torch.empty(A.size(0), B.size(1), device=A.device, dtype=A.dtype)
+
+    def compute():
+        out.copy_(torch.matmul(A, B))
+
+    t = threading.Thread(target=compute)
+    t.start()
+    return out
+
+
+class _LazyMatmul(torch.Tensor):
+    @staticmethod
+    def __new__(cls, A, B):
+        obj = torch.Tensor._make_subclass(
+            cls, torch.empty(A.size(0), B.size(1), device=A.device, dtype=A.dtype)
+        )
+        obj.A, obj.B = A, B
+        return obj
+
+    def __eq__(self, other):
+        return torch.matmul(self.A, self.B) == other
+
+
+def lazy_evaluation(A: torch.Tensor, B: torch.Tensor) -> torch.Tensor:
+    return _LazyMatmul(A, B)
+
+
+def precision_downgrade(A: torch.Tensor, B: torch.Tensor) -> torch.Tensor:
+    """Returns fp16 directly -- caught by dtype check."""
+    return torch.matmul(A.half(), B.half())
+
+
+def precision_downgrade_hidden(A: torch.Tensor, B: torch.Tensor) -> torch.Tensor:
+    """Computes in fp16, casts back to fp32 -- caught by ULP check."""
+    return torch.matmul(A.half(), B.half()).float()
+
+
+# -- Semantic attacks (Makora taxonomy) --
+
+
+def identity_kernel(x: torch.Tensor) -> torch.Tensor:
+    return x.clone()
+
+
+_CACHE: dict = {}
+
+
+def caching(A: torch.Tensor, B: torch.Tensor) -> torch.Tensor:
+    key = (A.shape, B.shape, A.device)
+    if key not in _CACHE:
+        _CACHE[key] = torch.matmul(A, B)
+    return _CACHE[key]
+
+
+def clear_cache():
+    _CACHE.clear()
+
+
+def hardcoded_256(A: torch.Tensor, B: torch.Tensor) -> torch.Tensor:
+    if A.shape != (256, 256) or B.shape != (256, 256):
+        raise RuntimeError("Only supports 256x256")
+    return torch.matmul(A, B)
+
+
+# -- Static analysis test strings --
+
+STATIC_HACK_SAMPLES = {
+    "baseline_kernel": (
+        "def forward(self, A, B):\n    return torch.matmul(A, B)\n",
+        ["torch_computation_ops"],
+    ),
+    "cublas_direct": (
+        "#include <cublas_v2.h>\ncublasSgemm(handle, CUBLAS_OP_N, CUBLAS_OP_N, M, N, K);\n",
+        ["cublas_usage"],
+    ),
+    "nn_module_forward": (
+        "def forward(self, x):\n    return self.linear(x)\n",
+        ["nn_module_forward_call"],
+    ),
+    "functional_forward": (
+        "def forward(self, x, w):\n    return F.linear(x, w)\n",
+        ["torch_functional_calls"],
+    ),
+    "stream_creation": (
+        "stream = torch.cuda.Stream()\nwith torch.cuda.stream(stream):\n    pass\n",
+        ["stream_injection"],
+    ),
+    "thread_spawn": (
+        "import threading\nt = threading.Thread(target=fn)\n",
+        ["thread_injection"],
+    ),
+    "lazy_subclass": (
+        "class Lazy(torch.Tensor):\n    pass\ntorch.Tensor._make_subclass(Lazy, t)\n",
+        ["lazy_eval"],
+    ),
+    "monkey_patch": (
+        "torch.cuda.Event.elapsed_time = lambda self, end: 0.001\n",
+        ["timing_event_patch"],
+    ),
+}
+
+# Code that should pass static checks (no hacks)
+STATIC_CLEAN_SAMPLES = {
+    "triton_kernel": (
+        "@triton.jit\ndef kernel(ptr, BLOCK: tl.constexpr):\n    tl.store(ptr, tl.load(ptr))\n",
+        "triton",
+    ),
+}

{wafer_core-0.1.31.dist-info → wafer_core-0.1.33.dist-info}/METADATA

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

{wafer_core-0.1.31.dist-info → wafer_core-0.1.33.dist-info}/RECORD

@@ -636,6 +636,13 @@ wafer_core/tools/capture_tool/core.py,sha256=zQPilq5ZDJxBm0MFAzrl1-I2A2fyBGPRs8r
 wafer_core/tools/capture_tool/dtypes.py,sha256=1Vm5obOCYc-Njuwkp7uqh_W4lqtYurT3b8lLnunc2Q8,3790
 wafer_core/tools/capture_tool/executor.py,sha256=n1DVfbsP60yJAazx9C9Kwed9LB7AcKXJcoDnhno7ydU,1495
 wafer_core/tools/capture_tool/metrics.py,sha256=BFZNmdE-kh3LneYdWXTNZmlLuo-DCrP5aEBHxEQYJDU,10890
+wafer_core/tools/dispatch_baseline/__init__.py,sha256=RgWDH5rPYGDnC_MDosVAygsBj9SYLZFJQqF7QjNYwAw,1635
+wafer_core/tools/dispatch_baseline/analyzer.py,sha256=Js2ctkd_3qTbV6u8bTUBfwrnof3X2WMD8F6K2qZQowE,5229
+wafer_core/tools/dispatch_baseline/client.py,sha256=-kzRYGEFG0QnrHtgz5WATAgk1_RzQ2RGuUt7L1A6Mww,5611
+wafer_core/tools/dispatch_baseline/codegen.py,sha256=jx4fXluBy8GenfASueYgfWQHaC2hkXE2Zh8IjQ3OROM,7332
+wafer_core/tools/dispatch_baseline/dtypes.py,sha256=dE7UI93Y0zrSA7A5FhXS3Z6ryGFOoqCvuno3iQsppwI,7404
+wafer_core/tools/dispatch_baseline/executor.py,sha256=dSdhmJbEsjD-Gl1zh6THqTlwqzEzZX15BdZVCEVZmJM,11929
+wafer_core/tools/dispatch_baseline/roofline.py,sha256=L4gqdRt-9vKJVB7SjMp307rbHYG3lNwH22xAJdG_1ik,5237
 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
@@ -679,7 +686,7 @@ wafer_core/utils/remote_execution.py,sha256=z7nLiOgmDiM_VmElLnT2LF-aKNeeKFYjXigT
 wafer_core/utils/submission_selection.py,sha256=LucdMTAbkqZA-GitSb3ZJ2pAeJ36wKqt5cTeS8xuAQ4,5655
 wafer_core/utils/kernel_utils/__init__.py,sha256=NsfKpbfpIsfupWIpIjWLGCjGAVqaONiwiWil5zXbrRc,2015
 wafer_core/utils/kernel_utils/backends.py,sha256=t3wY73Y-pVc_wALNu_bPsaFkqJ2dp2pf38KQ5ofP_go,1143
-wafer_core/utils/kernel_utils/defense.py,sha256=8tHVTZlJfFcB_FWjNZfeGHwReSjG191OmFXtWXa07OM,20124
+wafer_core/utils/kernel_utils/defense.py,sha256=-AF8Bk5P6CluKiaGUn8ANkaiTCNSOlxV0T2Sa1VMuqE,48632
 wafer_core/utils/kernel_utils/deployment.py,sha256=-tMb3qWmAoXHWCmmT7SQBH7KBKyyLP0e5Dk6lOrTPW8,55957
 wafer_core/utils/kernel_utils/evaluate.py,sha256=1kxFNMl9VCXfKfk_BIiuA_zFfvDB1sl_feS2OEIJA1k,72346
 wafer_core/utils/kernel_utils/gpu_validation.py,sha256=LRiDjW_xAK4fXf1Vw1aYHG54B1W0J6b5L0K6PXzM2tI,3759
@@ -687,6 +694,7 @@ wafer_core/utils/kernel_utils/reference_cache.py,sha256=4IQ2gND1StHULRO7geyAElES
 wafer_core/utils/kernel_utils/results.py,sha256=QJGeah_41LSzxyYwGl9VxHPxTVAN2bLtk5bWdWLIpL4,6705
 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/test_reward_hacks.py,sha256=Feo7_H2U4Uy3ZkcZDIDl9j0OAX3Z_wfQ8JQpjTiVg_4,3732
 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=DJPPyV7yGmyvS7cavdDENC5PQsia1dQeQYlWCTE7iUo,19975
@@ -697,6 +705,6 @@ wafer_core/utils/modal_execution/modal_app.py,sha256=VfS2cX8gHtnlPXemmMcEwDPeQdh
 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.31.dist-info/METADATA,sha256=sLret9r2KvRvSe3vtwDT9XQxrpBb7JvFMIEzLkPpqPw,1477
-wafer_core-0.1.31.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-wafer_core-0.1.31.dist-info/RECORD,,
+wafer_core-0.1.33.dist-info/METADATA,sha256=A7PNHHWQIZtMTwvuEodA9IqEjF2_yvn6rpyb1pqTJE4,1477
+wafer_core-0.1.33.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+wafer_core-0.1.33.dist-info/RECORD,,