Add Kernel Hub build/torch-universal + torch-cpu for get_kernel (szl-maskmod, version 1)

Browse files

Files changed (8) hide show

build/torch-cpu/metadata.json +13 -0
build/torch-cpu/szl_maskmod/__init__.py +7 -0
build/torch-cpu/szl_maskmod/_chain.py +35 -0
build/torch-cpu/szl_maskmod/_ops.py +54 -0
build/torch-universal/metadata.json +13 -0
build/torch-universal/szl_maskmod/__init__.py +7 -0
build/torch-universal/szl_maskmod/_chain.py +35 -0
build/torch-universal/szl_maskmod/_ops.py +54 -0

build/torch-cpu/metadata.json ADDED Viewed

	@@ -0,0 +1,13 @@

+{
+  "name": "szl-maskmod",
+  "version": 1,
+  "summary": "SZL original Flex-silhouette score_mod + block-mask attention with SHA3 receipts (CPU torch)",
+  "license": "Apache-2.0",
+  "universal": true,
+  "python-depends": [],
+  "id": "_szl_maskmod_cpu_010",
+  "backend": {
+    "type": "cpu"
+  },
+  "doctrine": "Lambda=Conjecture 1 (advisory)"
+}

build/torch-cpu/szl_maskmod/__init__.py ADDED Viewed

	@@ -0,0 +1,7 @@

+# SPDX-License-Identifier: Apache-2.0
+# Copyright 2026 SZL Holdings
+from ._chain import ReceiptChain
+from ._ops import maskmod_attn, selfcheck
+__all__ = ["ReceiptChain", "maskmod_attn", "selfcheck"]
+__version__ = "0.1.0"

build/torch-cpu/szl_maskmod/_chain.py ADDED Viewed

	@@ -0,0 +1,35 @@

+# SPDX-License-Identifier: Apache-2.0
+# Copyright 2026 SZL Holdings
+from __future__ import annotations
+import hashlib, json
+from typing import Any, List, Optional, Tuple
+GENESIS = "0" * 64
+def _canon(obj: Any) -> bytes:
+    return json.dumps(obj, sort_keys=True, separators=(",", ":")).encode("utf-8")
+def _sha3(body: bytes) -> str:
+    return hashlib.sha3_256(body).hexdigest()
+class ReceiptChain:
+    def __init__(self) -> None:
+        self._rows: List[dict] = []
+    def emit(self, payload: dict) -> str:
+        prev = self._rows[-1]["digest"] if self._rows else GENESIS
+        body = dict(payload); body["seq"] = len(self._rows); body["prev"] = prev
+        digest = _sha3(_canon(body))
+        self._rows.append({**body, "digest": digest})
+        return digest
+    def verify(self) -> Tuple[bool, int, int]:
+        prev = GENESIS
+        for i, row in enumerate(self._rows):
+            body = {k: v for k, v in row.items() if k != "digest"}
+            if row.get("prev") != prev or _sha3(_canon(body)) != row.get("digest"):
+                return False, i, i
+            prev = row["digest"]
+        return True, len(self._rows), -1
+    def head(self) -> Optional[str]:
+        return self._rows[-1]["digest"] if self._rows else None
+    def __len__(self) -> int:
+        return len(self._rows)

build/torch-cpu/szl_maskmod/_ops.py ADDED Viewed

	@@ -0,0 +1,54 @@

+# SPDX-License-Identifier: Apache-2.0
+# Copyright 2026 SZL Holdings
+"""Original score_mod + block-mask attention. Not copied from flex_attention.py."""
+from __future__ import annotations
+from typing import Callable, Optional
+import torch
+from ._chain import ReceiptChain
+ScoreMod = Callable[[torch.Tensor, torch.Tensor, torch.Tensor], torch.Tensor]
+def _dense_attn(q, k, v, *, causal, score_mod, block_mask, scale):
+    b, h, tq, d = q.shape
+    tkv = k.shape[2]
+    sm = (d ** -0.5) if scale is None else scale
+    scores = torch.matmul(q, k.transpose(-2, -1)) * sm
+    q_idx = torch.arange(tq, device=q.device)[:, None]
+    k_idx = torch.arange(tkv, device=q.device)[None, :]
+    if causal:
+        scores = scores.masked_fill(q_idx < k_idx, float("-inf"))
+    if block_mask is not None:
+        scores = scores.masked_fill(~block_mask, float("-inf"))
+    if score_mod is not None:
+        scores = score_mod(scores, q_idx, k_idx)
+    p = torch.softmax(scores, dim=-1)
+    p = torch.nan_to_num(p, nan=0.0)
+    return torch.matmul(p, v)
+def maskmod_attn(
+    q, k, v, *,
+    score_mod: Optional[ScoreMod] = None,
+    block_mask: Optional[torch.Tensor] = None,
+    causal: bool = False,
+    chain: Optional[ReceiptChain] = None,
+    scale: Optional[float] = None,
+):
+    y = _dense_attn(q, k, v, causal=causal, score_mod=score_mod, block_mask=block_mask, scale=scale)
+    if chain is not None:
+        mid = "none" if score_mod is None else getattr(score_mod, "__name__", type(score_mod).__name__)
+        bdigest = "none" if block_mask is None else f"sum={float(block_mask.to(torch.float32).sum().item()):.6g}"
+        chain.emit({"op": "maskmod_attn", "score_mod": mid, "block_mask": bdigest,
+                    "causal": causal, "q_shape": list(q.shape), "lambda": "Conjecture 1"})
+    return y
+def selfcheck() -> dict:
+    torch.manual_seed(20260828)
+    q = k = v = torch.randn(1, 2, 8, 16)
+    chain = ReceiptChain()
+    y = maskmod_attn(q, k, v, causal=True, chain=chain)
+    ref = torch.nn.functional.scaled_dot_product_attention(q, k, v, is_causal=True, dropout_p=0.0)
+    err = float((y - ref).abs().max().item())
+    ok_c, depth, brk = chain.verify()
+    ok = bool(err < 1e-5 and ok_c)
+    return {"ok": ok, "max_abs_vs_sdpa_causal": err, "chain_ok": ok_c, "chain_depth": depth,
+            "lambda": "Conjecture 1", "note": "correctness vs SDPA causal; no speedup claimed"}

build/torch-universal/metadata.json ADDED Viewed

	@@ -0,0 +1,13 @@

+{
+  "name": "szl-maskmod",
+  "version": 1,
+  "summary": "SZL original Flex-silhouette score_mod + block-mask attention with SHA3 receipts (CPU torch)",
+  "license": "Apache-2.0",
+  "universal": true,
+  "python-depends": [],
+  "id": "_szl_maskmod_universal_010",
+  "backend": {
+    "type": "cpu"
+  },
+  "doctrine": "Lambda=Conjecture 1 (advisory)"
+}

build/torch-universal/szl_maskmod/__init__.py ADDED Viewed

	@@ -0,0 +1,7 @@

+# SPDX-License-Identifier: Apache-2.0
+# Copyright 2026 SZL Holdings
+from ._chain import ReceiptChain
+from ._ops import maskmod_attn, selfcheck
+__all__ = ["ReceiptChain", "maskmod_attn", "selfcheck"]
+__version__ = "0.1.0"

build/torch-universal/szl_maskmod/_chain.py ADDED Viewed

	@@ -0,0 +1,35 @@

+# SPDX-License-Identifier: Apache-2.0
+# Copyright 2026 SZL Holdings
+from __future__ import annotations
+import hashlib, json
+from typing import Any, List, Optional, Tuple
+GENESIS = "0" * 64
+def _canon(obj: Any) -> bytes:
+    return json.dumps(obj, sort_keys=True, separators=(",", ":")).encode("utf-8")
+def _sha3(body: bytes) -> str:
+    return hashlib.sha3_256(body).hexdigest()
+class ReceiptChain:
+    def __init__(self) -> None:
+        self._rows: List[dict] = []
+    def emit(self, payload: dict) -> str:
+        prev = self._rows[-1]["digest"] if self._rows else GENESIS
+        body = dict(payload); body["seq"] = len(self._rows); body["prev"] = prev
+        digest = _sha3(_canon(body))
+        self._rows.append({**body, "digest": digest})
+        return digest
+    def verify(self) -> Tuple[bool, int, int]:
+        prev = GENESIS
+        for i, row in enumerate(self._rows):
+            body = {k: v for k, v in row.items() if k != "digest"}
+            if row.get("prev") != prev or _sha3(_canon(body)) != row.get("digest"):
+                return False, i, i
+            prev = row["digest"]
+        return True, len(self._rows), -1
+    def head(self) -> Optional[str]:
+        return self._rows[-1]["digest"] if self._rows else None
+    def __len__(self) -> int:
+        return len(self._rows)

build/torch-universal/szl_maskmod/_ops.py ADDED Viewed

	@@ -0,0 +1,54 @@

+# SPDX-License-Identifier: Apache-2.0
+# Copyright 2026 SZL Holdings
+"""Original score_mod + block-mask attention. Not copied from flex_attention.py."""
+from __future__ import annotations
+from typing import Callable, Optional
+import torch
+from ._chain import ReceiptChain
+ScoreMod = Callable[[torch.Tensor, torch.Tensor, torch.Tensor], torch.Tensor]
+def _dense_attn(q, k, v, *, causal, score_mod, block_mask, scale):
+    b, h, tq, d = q.shape
+    tkv = k.shape[2]
+    sm = (d ** -0.5) if scale is None else scale
+    scores = torch.matmul(q, k.transpose(-2, -1)) * sm
+    q_idx = torch.arange(tq, device=q.device)[:, None]
+    k_idx = torch.arange(tkv, device=q.device)[None, :]
+    if causal:
+        scores = scores.masked_fill(q_idx < k_idx, float("-inf"))
+    if block_mask is not None:
+        scores = scores.masked_fill(~block_mask, float("-inf"))
+    if score_mod is not None:
+        scores = score_mod(scores, q_idx, k_idx)
+    p = torch.softmax(scores, dim=-1)
+    p = torch.nan_to_num(p, nan=0.0)
+    return torch.matmul(p, v)
+def maskmod_attn(
+    q, k, v, *,
+    score_mod: Optional[ScoreMod] = None,
+    block_mask: Optional[torch.Tensor] = None,
+    causal: bool = False,
+    chain: Optional[ReceiptChain] = None,
+    scale: Optional[float] = None,
+):
+    y = _dense_attn(q, k, v, causal=causal, score_mod=score_mod, block_mask=block_mask, scale=scale)
+    if chain is not None:
+        mid = "none" if score_mod is None else getattr(score_mod, "__name__", type(score_mod).__name__)
+        bdigest = "none" if block_mask is None else f"sum={float(block_mask.to(torch.float32).sum().item()):.6g}"
+        chain.emit({"op": "maskmod_attn", "score_mod": mid, "block_mask": bdigest,
+                    "causal": causal, "q_shape": list(q.shape), "lambda": "Conjecture 1"})
+    return y
+def selfcheck() -> dict:
+    torch.manual_seed(20260828)
+    q = k = v = torch.randn(1, 2, 8, 16)
+    chain = ReceiptChain()
+    y = maskmod_attn(q, k, v, causal=True, chain=chain)
+    ref = torch.nn.functional.scaled_dot_product_attention(q, k, v, is_causal=True, dropout_p=0.0)
+    err = float((y - ref).abs().max().item())
+    ok_c, depth, brk = chain.verify()
+    ok = bool(err < 1e-5 and ok_c)
+    return {"ok": ok, "max_abs_vs_sdpa_causal": err, "chain_ok": ok_c, "chain_depth": depth,
+            "lambda": "Conjecture 1", "note": "correctness vs SDPA causal; no speedup claimed"}