phanerozoic commited on Jul 28

Commit

dea4c9f

verified ·

1 Parent(s): cc96269

Uploaded using `kernel-builder`.

Browse files

Files changed (41) hide show

build/torch211-cxx11-cu126-x86_64-linux/__init__.py +51 -0
build/torch211-cxx11-cu126-x86_64-linux/{_mamba3_cuda_95022f8.abi3.so → _mamba3_cuda_c549070.abi3.so} +2 -2
build/torch211-cxx11-cu126-x86_64-linux/_ops.py +3 -3
build/torch211-cxx11-cu126-x86_64-linux/metadata.json +5 -5
build/torch211-cxx11-cu128-x86_64-linux/__init__.py +51 -0
build/torch211-cxx11-cu128-x86_64-linux/{_mamba3_cuda_95022f8.abi3.so → _mamba3_cuda_c549070.abi3.so} +2 -2
build/torch211-cxx11-cu128-x86_64-linux/_ops.py +3 -3
build/torch211-cxx11-cu128-x86_64-linux/metadata.json +5 -5
build/torch211-cxx11-cu130-x86_64-linux/__init__.py +51 -0
build/torch211-cxx11-cu130-x86_64-linux/{_mamba3_cuda_95022f8.abi3.so → _mamba3_cuda_c549070.abi3.so} +2 -2
build/torch211-cxx11-cu130-x86_64-linux/_ops.py +3 -3
build/torch211-cxx11-cu130-x86_64-linux/metadata.json +5 -5
build/torch212-cxx11-cu126-x86_64-linux/__init__.py +51 -0
build/torch212-cxx11-cu126-x86_64-linux/{_mamba3_cuda_95022f8.abi3.so → _mamba3_cuda_c549070.abi3.so} +2 -2
build/torch212-cxx11-cu126-x86_64-linux/_ops.py +3 -3
build/torch212-cxx11-cu126-x86_64-linux/metadata.json +5 -5
build/torch212-cxx11-cu130-x86_64-linux/__init__.py +51 -0
build/torch212-cxx11-cu130-x86_64-linux/_mamba3_cuda_95022f8.abi3.so +0 -3
build/torch212-cxx11-cu130-x86_64-linux/_mamba3_cuda_c549070.abi3.so +3 -0
build/torch212-cxx11-cu130-x86_64-linux/_ops.py +3 -3
build/torch212-cxx11-cu130-x86_64-linux/metadata.json +5 -5
build/torch212-cxx11-cu132-x86_64-linux/__init__.py +51 -0
build/torch212-cxx11-cu132-x86_64-linux/_mamba3_cuda_95022f8.abi3.so +0 -3
build/torch212-cxx11-cu132-x86_64-linux/_mamba3_cuda_c549070.abi3.so +3 -0
build/torch212-cxx11-cu132-x86_64-linux/_ops.py +3 -3
build/torch212-cxx11-cu132-x86_64-linux/metadata.json +5 -5
build/torch213-cxx11-cu126-x86_64-linux/__init__.py +51 -0
build/torch213-cxx11-cu126-x86_64-linux/_mamba3_cuda_95022f8.abi3.so +0 -3
build/torch213-cxx11-cu126-x86_64-linux/_mamba3_cuda_c549070.abi3.so +3 -0
build/torch213-cxx11-cu126-x86_64-linux/_ops.py +3 -3
build/torch213-cxx11-cu126-x86_64-linux/metadata.json +5 -5
build/torch213-cxx11-cu130-x86_64-linux/__init__.py +51 -0
build/torch213-cxx11-cu130-x86_64-linux/_mamba3_cuda_95022f8.abi3.so +0 -3
build/torch213-cxx11-cu130-x86_64-linux/_mamba3_cuda_c549070.abi3.so +3 -0
build/torch213-cxx11-cu130-x86_64-linux/_ops.py +3 -3
build/torch213-cxx11-cu130-x86_64-linux/metadata.json +5 -5
build/torch213-cxx11-cu132-x86_64-linux/__init__.py +51 -0
build/torch213-cxx11-cu132-x86_64-linux/_mamba3_cuda_95022f8.abi3.so +0 -3
build/torch213-cxx11-cu132-x86_64-linux/_mamba3_cuda_c549070.abi3.so +3 -0
build/torch213-cxx11-cu132-x86_64-linux/_ops.py +3 -3
build/torch213-cxx11-cu132-x86_64-linux/metadata.json +5 -5

build/torch211-cxx11-cu126-x86_64-linux/__init__.py CHANGED Viewed

@@ -13,6 +13,57 @@ from ._ops import ops
 __all__ = ["forward", "DecodeState", "cumulative_angles", "chunk_decay"]
 def cumulative_angles(angles, dt):
     """Per-step rotation rates (B,S,H,Na) and dt (B,H,S) -> cumulative angles, float32."""
     base = torch.tanh(angles.float()) * dt.float().permute(0, 2, 1)[..., None] * math.pi

 __all__ = ["forward", "DecodeState", "cumulative_angles", "chunk_decay"]
+def _register_fakes():
+    """Meta kernels for the three ops, so a compiled graph does not break here.
+    The step op mutates four of its arguments; those carry `Tensor(a!)` in the
+    schema, so functionalization handles them and a fake only has to produce the
+    return value. The library name is read off a live op because the kernel
+    builder and a local JIT build choose different ones.
+    """
+    try:
+        lib = ops.mamba3_mimo_fwd.default._schema.name.split("::")[0]
+    except Exception:
+        return ()
+    def like_f32(t):
+        return torch.empty_like(t, dtype=torch.float32)
+    def register(name, fn):
+        if not hasattr(ops, name):
+            return False
+        torch.library.register_fake(f"{lib}::{name}", fn)
+        return True
+    def _step(q, k, v, z, adt, dt, trap, q_bias, k_bias, angles, mimo_v, mimo_o,
+              mimo_z, D, angle_in, angle_out, S_st, kprev_in, vprev_in,
+              kprev_out, vprev_out):
+        return torch.empty_like(v)
+    def _fwd(q, k, v, z, q_bias, k_bias, mimo_v, mimo_o, mimo_z, D, angles, dt,
+             trap, dA_cs, dA_cs_rev, chunk_size, norm_weight, norm_eps):
+        return torch.empty_like(v)
+    def _bwd(q, k, v, z, q_bias, k_bias, mimo_v, mimo_o, mimo_z, D, angles, dt,
+             trap, dA_cs, dA_cs_rev, chunk_size, norm_weight, norm_eps, dy):
+        H, P = v.shape[-2], v.shape[-1]
+        f32 = torch.float32
+        return [like_f32(q), like_f32(k), like_f32(v), like_f32(v),
+                like_f32(q_bias), like_f32(k_bias), like_f32(mimo_v),
+                like_f32(mimo_o), like_f32(mimo_v),
+                v.new_empty((H,), dtype=f32), v.new_empty((H, P), dtype=f32),
+                like_f32(dA_cs), like_f32(dA_cs_rev), like_f32(dt),
+                like_f32(trap), like_f32(angles)]
+    done = [n for n, f in (("mamba3_mimo_step", _step),
+                           ("mamba3_mimo_fwd", _fwd),
+                           ("mamba3_mimo_bwd", _bwd)) if register(n, f)]
+    return tuple(done)
+FAKE_OPS = _register_fakes()
 def cumulative_angles(angles, dt):
     """Per-step rotation rates (B,S,H,Na) and dt (B,H,S) -> cumulative angles, float32."""
     base = torch.tanh(angles.float()) * dt.float().permute(0, 2, 1)[..., None] * math.pi

build/torch211-cxx11-cu126-x86_64-linux/{_mamba3_cuda_95022f8.abi3.so → _mamba3_cuda_c549070.abi3.so} RENAMED Viewed

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:17e6f3474429a73c6a6c9074b6f2a2f973277110e6dbca3b1339c9a1dc5f4c8a
-size 6356800

 version https://git-lfs.github.com/spec/v1
+oid sha256:141f1ccdf30fd91844b71a08eaa9919be3c2fe6258a60523a3abd23afe944a26
+size 6357160

build/torch211-cxx11-cu126-x86_64-linux/_ops.py CHANGED Viewed

@@ -1,9 +1,9 @@
 import torch
-from . import _mamba3_cuda_95022f8
-ops = torch.ops._mamba3_cuda_95022f8
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_mamba3_cuda_95022f8::{op_name}"

 import torch
+from . import _mamba3_cuda_c549070
+ops = torch.ops._mamba3_cuda_c549070
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
+    return f"_mamba3_cuda_c549070::{op_name}"

build/torch211-cxx11-cu126-x86_64-linux/metadata.json CHANGED Viewed

@@ -1,6 +1,6 @@
 {
   "name": "mamba3",
-  "id": "_mamba3_cuda_95022f8",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],
@@ -16,9 +16,9 @@
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "__init__.py": "zWrrGqu25MORtMITQd22NHExxrHe4vYeSx4kAjDl/94=",
-      "_mamba3_cuda_95022f8.abi3.so": "F+bzR0QppzxqbJB0tvKi+XMncRDm28o7EznJodxfTIo=",
-      "_ops.py": "5TOqHeKVQS1MnkoVuOdCv8fNnWQNe+S86igdxnRw8G8=",
       "mamba3/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
     }
   },
@@ -29,7 +29,7 @@
       "dirty": false
     },
     "kernel": {
-      "sha": "95022f83b9b4c5d080443e6e91a8549afd7ae80b",
       "dirty": false
     }
   }

 {
   "name": "mamba3",
+  "id": "_mamba3_cuda_c549070",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],
   "digest": {
     "algorithm": "sha256",
     "files": {
+      "__init__.py": "zKP4l3anEqOcUuw5TEp4mkMhr0OFWQe9cUTq0Wu38aQ=",
+      "_mamba3_cuda_c549070.abi3.so": "FB8czfMP2RhEtxoI6qmRm+PC/mJYpgUjo6vSOv6USiY=",
+      "_ops.py": "Hw1EtZGXlC4SXcGIUpbQwVl8+sKMlb72u1JIHgKqHT8=",
       "mamba3/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
     }
   },
       "dirty": false
     },
     "kernel": {
+      "sha": "c54907036afbc32fae8bddb75ae7aa2bbaf3dde3",
       "dirty": false
     }
   }

build/torch211-cxx11-cu128-x86_64-linux/__init__.py CHANGED Viewed

@@ -13,6 +13,57 @@ from ._ops import ops
 __all__ = ["forward", "DecodeState", "cumulative_angles", "chunk_decay"]
 def cumulative_angles(angles, dt):
     """Per-step rotation rates (B,S,H,Na) and dt (B,H,S) -> cumulative angles, float32."""
     base = torch.tanh(angles.float()) * dt.float().permute(0, 2, 1)[..., None] * math.pi

 __all__ = ["forward", "DecodeState", "cumulative_angles", "chunk_decay"]
+def _register_fakes():
+    """Meta kernels for the three ops, so a compiled graph does not break here.
+    The step op mutates four of its arguments; those carry `Tensor(a!)` in the
+    schema, so functionalization handles them and a fake only has to produce the
+    return value. The library name is read off a live op because the kernel
+    builder and a local JIT build choose different ones.
+    """
+    try:
+        lib = ops.mamba3_mimo_fwd.default._schema.name.split("::")[0]
+    except Exception:
+        return ()
+    def like_f32(t):
+        return torch.empty_like(t, dtype=torch.float32)
+    def register(name, fn):
+        if not hasattr(ops, name):
+            return False
+        torch.library.register_fake(f"{lib}::{name}", fn)
+        return True
+    def _step(q, k, v, z, adt, dt, trap, q_bias, k_bias, angles, mimo_v, mimo_o,
+              mimo_z, D, angle_in, angle_out, S_st, kprev_in, vprev_in,
+              kprev_out, vprev_out):
+        return torch.empty_like(v)
+    def _fwd(q, k, v, z, q_bias, k_bias, mimo_v, mimo_o, mimo_z, D, angles, dt,
+             trap, dA_cs, dA_cs_rev, chunk_size, norm_weight, norm_eps):
+        return torch.empty_like(v)
+    def _bwd(q, k, v, z, q_bias, k_bias, mimo_v, mimo_o, mimo_z, D, angles, dt,
+             trap, dA_cs, dA_cs_rev, chunk_size, norm_weight, norm_eps, dy):
+        H, P = v.shape[-2], v.shape[-1]
+        f32 = torch.float32
+        return [like_f32(q), like_f32(k), like_f32(v), like_f32(v),
+                like_f32(q_bias), like_f32(k_bias), like_f32(mimo_v),
+                like_f32(mimo_o), like_f32(mimo_v),
+                v.new_empty((H,), dtype=f32), v.new_empty((H, P), dtype=f32),
+                like_f32(dA_cs), like_f32(dA_cs_rev), like_f32(dt),
+                like_f32(trap), like_f32(angles)]
+    done = [n for n, f in (("mamba3_mimo_step", _step),
+                           ("mamba3_mimo_fwd", _fwd),
+                           ("mamba3_mimo_bwd", _bwd)) if register(n, f)]
+    return tuple(done)
+FAKE_OPS = _register_fakes()
 def cumulative_angles(angles, dt):
     """Per-step rotation rates (B,S,H,Na) and dt (B,H,S) -> cumulative angles, float32."""
     base = torch.tanh(angles.float()) * dt.float().permute(0, 2, 1)[..., None] * math.pi

build/torch211-cxx11-cu128-x86_64-linux/{_mamba3_cuda_95022f8.abi3.so → _mamba3_cuda_c549070.abi3.so} RENAMED Viewed

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:589cc0b3472a805cc1b2efd9fcdb9afd0f68ea36279e8f669538932b01e3d950
-size 9965472

 version https://git-lfs.github.com/spec/v1
+oid sha256:88e544b58ad03bcd92a26f3f2f0be770bd65fb274a79a40f427f55546f8d5990
+size 9969928

build/torch211-cxx11-cu128-x86_64-linux/_ops.py CHANGED Viewed

@@ -1,9 +1,9 @@
 import torch
-from . import _mamba3_cuda_95022f8
-ops = torch.ops._mamba3_cuda_95022f8
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_mamba3_cuda_95022f8::{op_name}"

 import torch
+from . import _mamba3_cuda_c549070
+ops = torch.ops._mamba3_cuda_c549070
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
+    return f"_mamba3_cuda_c549070::{op_name}"

build/torch211-cxx11-cu128-x86_64-linux/metadata.json CHANGED Viewed

@@ -1,6 +1,6 @@
 {
   "name": "mamba3",
-  "id": "_mamba3_cuda_95022f8",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],
@@ -18,9 +18,9 @@
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "__init__.py": "zWrrGqu25MORtMITQd22NHExxrHe4vYeSx4kAjDl/94=",
-      "_mamba3_cuda_95022f8.abi3.so": "WJzAs0cqgFzBsu/Z/Nua/Q9o6jYnno9mlTiTKwHj2VA=",
-      "_ops.py": "5TOqHeKVQS1MnkoVuOdCv8fNnWQNe+S86igdxnRw8G8=",
       "mamba3/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
     }
   },
@@ -31,7 +31,7 @@
       "dirty": false
     },
     "kernel": {
-      "sha": "95022f83b9b4c5d080443e6e91a8549afd7ae80b",
       "dirty": false
     }
   }

 {
   "name": "mamba3",
+  "id": "_mamba3_cuda_c549070",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],
   "digest": {
     "algorithm": "sha256",
     "files": {
+      "__init__.py": "zKP4l3anEqOcUuw5TEp4mkMhr0OFWQe9cUTq0Wu38aQ=",
+      "_mamba3_cuda_c549070.abi3.so": "iOVEtYrQO82Som8/LwvncL1l+ydKeaQPQn9VVG+NWZA=",
+      "_ops.py": "Hw1EtZGXlC4SXcGIUpbQwVl8+sKMlb72u1JIHgKqHT8=",
       "mamba3/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
     }
   },
       "dirty": false
     },
     "kernel": {
+      "sha": "c54907036afbc32fae8bddb75ae7aa2bbaf3dde3",
       "dirty": false
     }
   }

build/torch211-cxx11-cu130-x86_64-linux/__init__.py CHANGED Viewed

@@ -13,6 +13,57 @@ from ._ops import ops
 __all__ = ["forward", "DecodeState", "cumulative_angles", "chunk_decay"]
 def cumulative_angles(angles, dt):
     """Per-step rotation rates (B,S,H,Na) and dt (B,H,S) -> cumulative angles, float32."""
     base = torch.tanh(angles.float()) * dt.float().permute(0, 2, 1)[..., None] * math.pi

 __all__ = ["forward", "DecodeState", "cumulative_angles", "chunk_decay"]
+def _register_fakes():
+    """Meta kernels for the three ops, so a compiled graph does not break here.
+    The step op mutates four of its arguments; those carry `Tensor(a!)` in the
+    schema, so functionalization handles them and a fake only has to produce the
+    return value. The library name is read off a live op because the kernel
+    builder and a local JIT build choose different ones.
+    """
+    try:
+        lib = ops.mamba3_mimo_fwd.default._schema.name.split("::")[0]
+    except Exception:
+        return ()
+    def like_f32(t):
+        return torch.empty_like(t, dtype=torch.float32)
+    def register(name, fn):
+        if not hasattr(ops, name):
+            return False
+        torch.library.register_fake(f"{lib}::{name}", fn)
+        return True
+    def _step(q, k, v, z, adt, dt, trap, q_bias, k_bias, angles, mimo_v, mimo_o,
+              mimo_z, D, angle_in, angle_out, S_st, kprev_in, vprev_in,
+              kprev_out, vprev_out):
+        return torch.empty_like(v)
+    def _fwd(q, k, v, z, q_bias, k_bias, mimo_v, mimo_o, mimo_z, D, angles, dt,
+             trap, dA_cs, dA_cs_rev, chunk_size, norm_weight, norm_eps):
+        return torch.empty_like(v)
+    def _bwd(q, k, v, z, q_bias, k_bias, mimo_v, mimo_o, mimo_z, D, angles, dt,
+             trap, dA_cs, dA_cs_rev, chunk_size, norm_weight, norm_eps, dy):
+        H, P = v.shape[-2], v.shape[-1]
+        f32 = torch.float32
+        return [like_f32(q), like_f32(k), like_f32(v), like_f32(v),
+                like_f32(q_bias), like_f32(k_bias), like_f32(mimo_v),
+                like_f32(mimo_o), like_f32(mimo_v),
+                v.new_empty((H,), dtype=f32), v.new_empty((H, P), dtype=f32),
+                like_f32(dA_cs), like_f32(dA_cs_rev), like_f32(dt),
+                like_f32(trap), like_f32(angles)]
+    done = [n for n, f in (("mamba3_mimo_step", _step),
+                           ("mamba3_mimo_fwd", _fwd),
+                           ("mamba3_mimo_bwd", _bwd)) if register(n, f)]
+    return tuple(done)
+FAKE_OPS = _register_fakes()
 def cumulative_angles(angles, dt):
     """Per-step rotation rates (B,S,H,Na) and dt (B,H,S) -> cumulative angles, float32."""
     base = torch.tanh(angles.float()) * dt.float().permute(0, 2, 1)[..., None] * math.pi

build/torch211-cxx11-cu130-x86_64-linux/{_mamba3_cuda_95022f8.abi3.so → _mamba3_cuda_c549070.abi3.so} RENAMED Viewed

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:21cd09f0d1ed8df1c84e3a60ffedfdb41713ada4035db42dd4e4853c3758ea5e
-size 9776800

 version https://git-lfs.github.com/spec/v1
+oid sha256:54b23ced497150c6046cc5ac4cbb539aa994c0914e903edbc693840b92a2fed6
+size 9781256

build/torch211-cxx11-cu130-x86_64-linux/_ops.py CHANGED Viewed

@@ -1,9 +1,9 @@
 import torch
-from . import _mamba3_cuda_95022f8
-ops = torch.ops._mamba3_cuda_95022f8
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_mamba3_cuda_95022f8::{op_name}"

 import torch
+from . import _mamba3_cuda_c549070
+ops = torch.ops._mamba3_cuda_c549070
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
+    return f"_mamba3_cuda_c549070::{op_name}"

build/torch211-cxx11-cu130-x86_64-linux/metadata.json CHANGED Viewed

@@ -1,6 +1,6 @@
 {
   "name": "mamba3",
-  "id": "_mamba3_cuda_95022f8",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],
@@ -18,9 +18,9 @@
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "__init__.py": "zWrrGqu25MORtMITQd22NHExxrHe4vYeSx4kAjDl/94=",
-      "_mamba3_cuda_95022f8.abi3.so": "Ic0J8NHtjfHITjpg/+39tBcTraQDXbQt1OSFPDdY6l4=",
-      "_ops.py": "5TOqHeKVQS1MnkoVuOdCv8fNnWQNe+S86igdxnRw8G8=",
       "mamba3/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
     }
   },
@@ -31,7 +31,7 @@
       "dirty": false
     },
     "kernel": {
-      "sha": "95022f83b9b4c5d080443e6e91a8549afd7ae80b",
       "dirty": false
     }
   }

 {
   "name": "mamba3",
+  "id": "_mamba3_cuda_c549070",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],
   "digest": {
     "algorithm": "sha256",
     "files": {
+      "__init__.py": "zKP4l3anEqOcUuw5TEp4mkMhr0OFWQe9cUTq0Wu38aQ=",
+      "_mamba3_cuda_c549070.abi3.so": "VLI87UlxUMYEbMWsTLtTmqmUwJFOkD7bxpOEC5Ki/tY=",
+      "_ops.py": "Hw1EtZGXlC4SXcGIUpbQwVl8+sKMlb72u1JIHgKqHT8=",
       "mamba3/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
     }
   },
       "dirty": false
     },
     "kernel": {
+      "sha": "c54907036afbc32fae8bddb75ae7aa2bbaf3dde3",
       "dirty": false
     }
   }

build/torch212-cxx11-cu126-x86_64-linux/__init__.py CHANGED Viewed

@@ -13,6 +13,57 @@ from ._ops import ops
 __all__ = ["forward", "DecodeState", "cumulative_angles", "chunk_decay"]
 def cumulative_angles(angles, dt):
     """Per-step rotation rates (B,S,H,Na) and dt (B,H,S) -> cumulative angles, float32."""
     base = torch.tanh(angles.float()) * dt.float().permute(0, 2, 1)[..., None] * math.pi

 __all__ = ["forward", "DecodeState", "cumulative_angles", "chunk_decay"]
+def _register_fakes():
+    """Meta kernels for the three ops, so a compiled graph does not break here.
+    The step op mutates four of its arguments; those carry `Tensor(a!)` in the
+    schema, so functionalization handles them and a fake only has to produce the
+    return value. The library name is read off a live op because the kernel
+    builder and a local JIT build choose different ones.
+    """
+    try:
+        lib = ops.mamba3_mimo_fwd.default._schema.name.split("::")[0]
+    except Exception:
+        return ()
+    def like_f32(t):
+        return torch.empty_like(t, dtype=torch.float32)
+    def register(name, fn):
+        if not hasattr(ops, name):
+            return False
+        torch.library.register_fake(f"{lib}::{name}", fn)
+        return True
+    def _step(q, k, v, z, adt, dt, trap, q_bias, k_bias, angles, mimo_v, mimo_o,
+              mimo_z, D, angle_in, angle_out, S_st, kprev_in, vprev_in,
+              kprev_out, vprev_out):
+        return torch.empty_like(v)
+    def _fwd(q, k, v, z, q_bias, k_bias, mimo_v, mimo_o, mimo_z, D, angles, dt,
+             trap, dA_cs, dA_cs_rev, chunk_size, norm_weight, norm_eps):
+        return torch.empty_like(v)
+    def _bwd(q, k, v, z, q_bias, k_bias, mimo_v, mimo_o, mimo_z, D, angles, dt,
+             trap, dA_cs, dA_cs_rev, chunk_size, norm_weight, norm_eps, dy):
+        H, P = v.shape[-2], v.shape[-1]
+        f32 = torch.float32
+        return [like_f32(q), like_f32(k), like_f32(v), like_f32(v),
+                like_f32(q_bias), like_f32(k_bias), like_f32(mimo_v),
+                like_f32(mimo_o), like_f32(mimo_v),
+                v.new_empty((H,), dtype=f32), v.new_empty((H, P), dtype=f32),
+                like_f32(dA_cs), like_f32(dA_cs_rev), like_f32(dt),
+                like_f32(trap), like_f32(angles)]
+    done = [n for n, f in (("mamba3_mimo_step", _step),
+                           ("mamba3_mimo_fwd", _fwd),
+                           ("mamba3_mimo_bwd", _bwd)) if register(n, f)]
+    return tuple(done)
+FAKE_OPS = _register_fakes()
 def cumulative_angles(angles, dt):
     """Per-step rotation rates (B,S,H,Na) and dt (B,H,S) -> cumulative angles, float32."""
     base = torch.tanh(angles.float()) * dt.float().permute(0, 2, 1)[..., None] * math.pi

build/torch212-cxx11-cu126-x86_64-linux/{_mamba3_cuda_95022f8.abi3.so → _mamba3_cuda_c549070.abi3.so} RENAMED Viewed

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:afa8423f1bf4f24eb0633e541c5977adda483a3f846dd791962f237c2051fc61
-size 6356696

 version https://git-lfs.github.com/spec/v1
+oid sha256:6fa7abbfce2cdaa2019648bfd060b9c89c921cc6d063acce57ba3c36bddba519
+size 6357056

build/torch212-cxx11-cu126-x86_64-linux/_ops.py CHANGED Viewed

@@ -1,9 +1,9 @@
 import torch
-from . import _mamba3_cuda_95022f8
-ops = torch.ops._mamba3_cuda_95022f8
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_mamba3_cuda_95022f8::{op_name}"

 import torch
+from . import _mamba3_cuda_c549070
+ops = torch.ops._mamba3_cuda_c549070
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
+    return f"_mamba3_cuda_c549070::{op_name}"

build/torch212-cxx11-cu126-x86_64-linux/metadata.json CHANGED Viewed

@@ -1,6 +1,6 @@
 {
   "name": "mamba3",
-  "id": "_mamba3_cuda_95022f8",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],
@@ -16,9 +16,9 @@
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "__init__.py": "zWrrGqu25MORtMITQd22NHExxrHe4vYeSx4kAjDl/94=",
-      "_mamba3_cuda_95022f8.abi3.so": "r6hCPxv08k6wYz5UHFl3rdpIOj+EbdeRli8jfCBR/GE=",
-      "_ops.py": "5TOqHeKVQS1MnkoVuOdCv8fNnWQNe+S86igdxnRw8G8=",
       "mamba3/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
     }
   },
@@ -29,7 +29,7 @@
       "dirty": false
     },
     "kernel": {
-      "sha": "95022f83b9b4c5d080443e6e91a8549afd7ae80b",
       "dirty": false
     }
   }

 {
   "name": "mamba3",
+  "id": "_mamba3_cuda_c549070",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],
   "digest": {
     "algorithm": "sha256",
     "files": {
+      "__init__.py": "zKP4l3anEqOcUuw5TEp4mkMhr0OFWQe9cUTq0Wu38aQ=",
+      "_mamba3_cuda_c549070.abi3.so": "b6erv84s2qIBlki/0GC5yJySHMbQY6zOV7o8Nr3bpRk=",
+      "_ops.py": "Hw1EtZGXlC4SXcGIUpbQwVl8+sKMlb72u1JIHgKqHT8=",
       "mamba3/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
     }
   },
       "dirty": false
     },
     "kernel": {
+      "sha": "c54907036afbc32fae8bddb75ae7aa2bbaf3dde3",
       "dirty": false
     }
   }

build/torch212-cxx11-cu130-x86_64-linux/__init__.py CHANGED Viewed

@@ -13,6 +13,57 @@ from ._ops import ops
 __all__ = ["forward", "DecodeState", "cumulative_angles", "chunk_decay"]
 def cumulative_angles(angles, dt):
     """Per-step rotation rates (B,S,H,Na) and dt (B,H,S) -> cumulative angles, float32."""
     base = torch.tanh(angles.float()) * dt.float().permute(0, 2, 1)[..., None] * math.pi

 __all__ = ["forward", "DecodeState", "cumulative_angles", "chunk_decay"]
+def _register_fakes():
+    """Meta kernels for the three ops, so a compiled graph does not break here.
+    The step op mutates four of its arguments; those carry `Tensor(a!)` in the
+    schema, so functionalization handles them and a fake only has to produce the
+    return value. The library name is read off a live op because the kernel
+    builder and a local JIT build choose different ones.
+    """
+    try:
+        lib = ops.mamba3_mimo_fwd.default._schema.name.split("::")[0]
+    except Exception:
+        return ()
+    def like_f32(t):
+        return torch.empty_like(t, dtype=torch.float32)
+    def register(name, fn):
+        if not hasattr(ops, name):
+            return False
+        torch.library.register_fake(f"{lib}::{name}", fn)
+        return True
+    def _step(q, k, v, z, adt, dt, trap, q_bias, k_bias, angles, mimo_v, mimo_o,
+              mimo_z, D, angle_in, angle_out, S_st, kprev_in, vprev_in,
+              kprev_out, vprev_out):
+        return torch.empty_like(v)
+    def _fwd(q, k, v, z, q_bias, k_bias, mimo_v, mimo_o, mimo_z, D, angles, dt,
+             trap, dA_cs, dA_cs_rev, chunk_size, norm_weight, norm_eps):
+        return torch.empty_like(v)
+    def _bwd(q, k, v, z, q_bias, k_bias, mimo_v, mimo_o, mimo_z, D, angles, dt,
+             trap, dA_cs, dA_cs_rev, chunk_size, norm_weight, norm_eps, dy):
+        H, P = v.shape[-2], v.shape[-1]
+        f32 = torch.float32
+        return [like_f32(q), like_f32(k), like_f32(v), like_f32(v),
+                like_f32(q_bias), like_f32(k_bias), like_f32(mimo_v),
+                like_f32(mimo_o), like_f32(mimo_v),
+                v.new_empty((H,), dtype=f32), v.new_empty((H, P), dtype=f32),
+                like_f32(dA_cs), like_f32(dA_cs_rev), like_f32(dt),
+                like_f32(trap), like_f32(angles)]
+    done = [n for n, f in (("mamba3_mimo_step", _step),
+                           ("mamba3_mimo_fwd", _fwd),
+                           ("mamba3_mimo_bwd", _bwd)) if register(n, f)]
+    return tuple(done)
+FAKE_OPS = _register_fakes()
 def cumulative_angles(angles, dt):
     """Per-step rotation rates (B,S,H,Na) and dt (B,H,S) -> cumulative angles, float32."""
     base = torch.tanh(angles.float()) * dt.float().permute(0, 2, 1)[..., None] * math.pi

build/torch212-cxx11-cu130-x86_64-linux/_mamba3_cuda_95022f8.abi3.so DELETED Viewed

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:0e7fc866c4a13501771152158b29f29847fcf49591b8a97380d33ea0b514887e
-size 9781160

build/torch212-cxx11-cu130-x86_64-linux/_mamba3_cuda_c549070.abi3.so ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:f682d95fc61b6b2b35da76007e0d79be29b21e53ac88c81202b607c5ba03cf34
+size 9781520

build/torch212-cxx11-cu130-x86_64-linux/_ops.py CHANGED Viewed

@@ -1,9 +1,9 @@
 import torch
-from . import _mamba3_cuda_95022f8
-ops = torch.ops._mamba3_cuda_95022f8
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_mamba3_cuda_95022f8::{op_name}"

 import torch
+from . import _mamba3_cuda_c549070
+ops = torch.ops._mamba3_cuda_c549070
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
+    return f"_mamba3_cuda_c549070::{op_name}"

build/torch212-cxx11-cu130-x86_64-linux/metadata.json CHANGED Viewed

@@ -1,6 +1,6 @@
 {
   "name": "mamba3",
-  "id": "_mamba3_cuda_95022f8",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],
@@ -18,9 +18,9 @@
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "__init__.py": "zWrrGqu25MORtMITQd22NHExxrHe4vYeSx4kAjDl/94=",
-      "_mamba3_cuda_95022f8.abi3.so": "Dn/IZsShNQF3EVIViynymEf89JWRuKlzgNM+oLUUiH4=",
-      "_ops.py": "5TOqHeKVQS1MnkoVuOdCv8fNnWQNe+S86igdxnRw8G8=",
       "mamba3/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
     }
   },
@@ -31,7 +31,7 @@
       "dirty": false
     },
     "kernel": {
-      "sha": "95022f83b9b4c5d080443e6e91a8549afd7ae80b",
       "dirty": false
     }
   }

 {
   "name": "mamba3",
+  "id": "_mamba3_cuda_c549070",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],
   "digest": {
     "algorithm": "sha256",
     "files": {
+      "__init__.py": "zKP4l3anEqOcUuw5TEp4mkMhr0OFWQe9cUTq0Wu38aQ=",
+      "_mamba3_cuda_c549070.abi3.so": "9oLZX8Ybays12nYAfg15vimyHlOsiMgSArYHxboDzzQ=",
+      "_ops.py": "Hw1EtZGXlC4SXcGIUpbQwVl8+sKMlb72u1JIHgKqHT8=",
       "mamba3/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
     }
   },
       "dirty": false
     },
     "kernel": {
+      "sha": "c54907036afbc32fae8bddb75ae7aa2bbaf3dde3",
       "dirty": false
     }
   }

build/torch212-cxx11-cu132-x86_64-linux/__init__.py CHANGED Viewed

@@ -13,6 +13,57 @@ from ._ops import ops
 __all__ = ["forward", "DecodeState", "cumulative_angles", "chunk_decay"]
 def cumulative_angles(angles, dt):
     """Per-step rotation rates (B,S,H,Na) and dt (B,H,S) -> cumulative angles, float32."""
     base = torch.tanh(angles.float()) * dt.float().permute(0, 2, 1)[..., None] * math.pi

 __all__ = ["forward", "DecodeState", "cumulative_angles", "chunk_decay"]
+def _register_fakes():
+    """Meta kernels for the three ops, so a compiled graph does not break here.
+    The step op mutates four of its arguments; those carry `Tensor(a!)` in the
+    schema, so functionalization handles them and a fake only has to produce the
+    return value. The library name is read off a live op because the kernel
+    builder and a local JIT build choose different ones.
+    """
+    try:
+        lib = ops.mamba3_mimo_fwd.default._schema.name.split("::")[0]
+    except Exception:
+        return ()
+    def like_f32(t):
+        return torch.empty_like(t, dtype=torch.float32)
+    def register(name, fn):
+        if not hasattr(ops, name):
+            return False
+        torch.library.register_fake(f"{lib}::{name}", fn)
+        return True
+    def _step(q, k, v, z, adt, dt, trap, q_bias, k_bias, angles, mimo_v, mimo_o,
+              mimo_z, D, angle_in, angle_out, S_st, kprev_in, vprev_in,
+              kprev_out, vprev_out):
+        return torch.empty_like(v)
+    def _fwd(q, k, v, z, q_bias, k_bias, mimo_v, mimo_o, mimo_z, D, angles, dt,
+             trap, dA_cs, dA_cs_rev, chunk_size, norm_weight, norm_eps):
+        return torch.empty_like(v)
+    def _bwd(q, k, v, z, q_bias, k_bias, mimo_v, mimo_o, mimo_z, D, angles, dt,
+             trap, dA_cs, dA_cs_rev, chunk_size, norm_weight, norm_eps, dy):
+        H, P = v.shape[-2], v.shape[-1]
+        f32 = torch.float32
+        return [like_f32(q), like_f32(k), like_f32(v), like_f32(v),
+                like_f32(q_bias), like_f32(k_bias), like_f32(mimo_v),
+                like_f32(mimo_o), like_f32(mimo_v),
+                v.new_empty((H,), dtype=f32), v.new_empty((H, P), dtype=f32),
+                like_f32(dA_cs), like_f32(dA_cs_rev), like_f32(dt),
+                like_f32(trap), like_f32(angles)]
+    done = [n for n, f in (("mamba3_mimo_step", _step),
+                           ("mamba3_mimo_fwd", _fwd),
+                           ("mamba3_mimo_bwd", _bwd)) if register(n, f)]
+    return tuple(done)
+FAKE_OPS = _register_fakes()
 def cumulative_angles(angles, dt):
     """Per-step rotation rates (B,S,H,Na) and dt (B,H,S) -> cumulative angles, float32."""
     base = torch.tanh(angles.float()) * dt.float().permute(0, 2, 1)[..., None] * math.pi

build/torch212-cxx11-cu132-x86_64-linux/_mamba3_cuda_95022f8.abi3.so DELETED Viewed

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:72d692a106fe0e6827903413f69291a2b7dca0971e936b09db1f57cd16a6c510
-size 9899944

build/torch212-cxx11-cu132-x86_64-linux/_mamba3_cuda_c549070.abi3.so ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:d767da8589280130d46c1bde9e5ad6fb81d943b4e6cf6c40bace41c288d416ef
+size 9904400

build/torch212-cxx11-cu132-x86_64-linux/_ops.py CHANGED Viewed

@@ -1,9 +1,9 @@
 import torch
-from . import _mamba3_cuda_95022f8
-ops = torch.ops._mamba3_cuda_95022f8
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_mamba3_cuda_95022f8::{op_name}"

 import torch
+from . import _mamba3_cuda_c549070
+ops = torch.ops._mamba3_cuda_c549070
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
+    return f"_mamba3_cuda_c549070::{op_name}"

build/torch212-cxx11-cu132-x86_64-linux/metadata.json CHANGED Viewed

@@ -1,6 +1,6 @@
 {
   "name": "mamba3",
-  "id": "_mamba3_cuda_95022f8",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],
@@ -18,9 +18,9 @@
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "__init__.py": "zWrrGqu25MORtMITQd22NHExxrHe4vYeSx4kAjDl/94=",
-      "_mamba3_cuda_95022f8.abi3.so": "ctaSoQb+DmgnkDQT9pKRorfcoJcek2sJ2x9XzRamxRA=",
-      "_ops.py": "5TOqHeKVQS1MnkoVuOdCv8fNnWQNe+S86igdxnRw8G8=",
       "mamba3/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
     }
   },
@@ -31,7 +31,7 @@
       "dirty": false
     },
     "kernel": {
-      "sha": "95022f83b9b4c5d080443e6e91a8549afd7ae80b",
       "dirty": false
     }
   }

 {
   "name": "mamba3",
+  "id": "_mamba3_cuda_c549070",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],
   "digest": {
     "algorithm": "sha256",
     "files": {
+      "__init__.py": "zKP4l3anEqOcUuw5TEp4mkMhr0OFWQe9cUTq0Wu38aQ=",
+      "_mamba3_cuda_c549070.abi3.so": "12fahYkoATDUbBvenlrW+4HZQ7Tmz2xAus5BwojUFu8=",
+      "_ops.py": "Hw1EtZGXlC4SXcGIUpbQwVl8+sKMlb72u1JIHgKqHT8=",
       "mamba3/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
     }
   },
       "dirty": false
     },
     "kernel": {
+      "sha": "c54907036afbc32fae8bddb75ae7aa2bbaf3dde3",
       "dirty": false
     }
   }

build/torch213-cxx11-cu126-x86_64-linux/__init__.py CHANGED Viewed

@@ -13,6 +13,57 @@ from ._ops import ops
 __all__ = ["forward", "DecodeState", "cumulative_angles", "chunk_decay"]
 def cumulative_angles(angles, dt):
     """Per-step rotation rates (B,S,H,Na) and dt (B,H,S) -> cumulative angles, float32."""
     base = torch.tanh(angles.float()) * dt.float().permute(0, 2, 1)[..., None] * math.pi

 __all__ = ["forward", "DecodeState", "cumulative_angles", "chunk_decay"]
+def _register_fakes():
+    """Meta kernels for the three ops, so a compiled graph does not break here.
+    The step op mutates four of its arguments; those carry `Tensor(a!)` in the
+    schema, so functionalization handles them and a fake only has to produce the
+    return value. The library name is read off a live op because the kernel
+    builder and a local JIT build choose different ones.
+    """
+    try:
+        lib = ops.mamba3_mimo_fwd.default._schema.name.split("::")[0]
+    except Exception:
+        return ()
+    def like_f32(t):
+        return torch.empty_like(t, dtype=torch.float32)
+    def register(name, fn):
+        if not hasattr(ops, name):
+            return False
+        torch.library.register_fake(f"{lib}::{name}", fn)
+        return True
+    def _step(q, k, v, z, adt, dt, trap, q_bias, k_bias, angles, mimo_v, mimo_o,
+              mimo_z, D, angle_in, angle_out, S_st, kprev_in, vprev_in,
+              kprev_out, vprev_out):
+        return torch.empty_like(v)
+    def _fwd(q, k, v, z, q_bias, k_bias, mimo_v, mimo_o, mimo_z, D, angles, dt,
+             trap, dA_cs, dA_cs_rev, chunk_size, norm_weight, norm_eps):
+        return torch.empty_like(v)
+    def _bwd(q, k, v, z, q_bias, k_bias, mimo_v, mimo_o, mimo_z, D, angles, dt,
+             trap, dA_cs, dA_cs_rev, chunk_size, norm_weight, norm_eps, dy):
+        H, P = v.shape[-2], v.shape[-1]
+        f32 = torch.float32
+        return [like_f32(q), like_f32(k), like_f32(v), like_f32(v),
+                like_f32(q_bias), like_f32(k_bias), like_f32(mimo_v),
+                like_f32(mimo_o), like_f32(mimo_v),
+                v.new_empty((H,), dtype=f32), v.new_empty((H, P), dtype=f32),
+                like_f32(dA_cs), like_f32(dA_cs_rev), like_f32(dt),
+                like_f32(trap), like_f32(angles)]
+    done = [n for n, f in (("mamba3_mimo_step", _step),
+                           ("mamba3_mimo_fwd", _fwd),
+                           ("mamba3_mimo_bwd", _bwd)) if register(n, f)]
+    return tuple(done)
+FAKE_OPS = _register_fakes()
 def cumulative_angles(angles, dt):
     """Per-step rotation rates (B,S,H,Na) and dt (B,H,S) -> cumulative angles, float32."""
     base = torch.tanh(angles.float()) * dt.float().permute(0, 2, 1)[..., None] * math.pi

build/torch213-cxx11-cu126-x86_64-linux/_mamba3_cuda_95022f8.abi3.so DELETED Viewed

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:b2a46b3eb53a3dea7570012ebd536aecf803763a423038d38e8d466b57507c8e
-size 6356984

build/torch213-cxx11-cu126-x86_64-linux/_mamba3_cuda_c549070.abi3.so ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:54292496960e8d5e438bf4f60a80e7d4349f4f0d75457be74cb633e2b175ba17
+size 6357344

build/torch213-cxx11-cu126-x86_64-linux/_ops.py CHANGED Viewed

@@ -1,9 +1,9 @@
 import torch
-from . import _mamba3_cuda_95022f8
-ops = torch.ops._mamba3_cuda_95022f8
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_mamba3_cuda_95022f8::{op_name}"

 import torch
+from . import _mamba3_cuda_c549070
+ops = torch.ops._mamba3_cuda_c549070
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
+    return f"_mamba3_cuda_c549070::{op_name}"

build/torch213-cxx11-cu126-x86_64-linux/metadata.json CHANGED Viewed

@@ -1,6 +1,6 @@
 {
   "name": "mamba3",
-  "id": "_mamba3_cuda_95022f8",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],
@@ -16,9 +16,9 @@
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "__init__.py": "zWrrGqu25MORtMITQd22NHExxrHe4vYeSx4kAjDl/94=",
-      "_mamba3_cuda_95022f8.abi3.so": "sqRrPrU6Pep1cAEuvVNq7PgDdjpCMDjTjo1Ga1dQfI4=",
-      "_ops.py": "5TOqHeKVQS1MnkoVuOdCv8fNnWQNe+S86igdxnRw8G8=",
       "mamba3/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
     }
   },
@@ -29,7 +29,7 @@
       "dirty": false
     },
     "kernel": {
-      "sha": "95022f83b9b4c5d080443e6e91a8549afd7ae80b",
       "dirty": false
     }
   }

 {
   "name": "mamba3",
+  "id": "_mamba3_cuda_c549070",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],
   "digest": {
     "algorithm": "sha256",
     "files": {
+      "__init__.py": "zKP4l3anEqOcUuw5TEp4mkMhr0OFWQe9cUTq0Wu38aQ=",
+      "_mamba3_cuda_c549070.abi3.so": "VCkklpYOjV5Di/T2CoDn1DSfTw11RXvnTLYz4rF1uhc=",
+      "_ops.py": "Hw1EtZGXlC4SXcGIUpbQwVl8+sKMlb72u1JIHgKqHT8=",
       "mamba3/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
     }
   },
       "dirty": false
     },
     "kernel": {
+      "sha": "c54907036afbc32fae8bddb75ae7aa2bbaf3dde3",
       "dirty": false
     }
   }

build/torch213-cxx11-cu130-x86_64-linux/__init__.py CHANGED Viewed

@@ -13,6 +13,57 @@ from ._ops import ops
 __all__ = ["forward", "DecodeState", "cumulative_angles", "chunk_decay"]
 def cumulative_angles(angles, dt):
     """Per-step rotation rates (B,S,H,Na) and dt (B,H,S) -> cumulative angles, float32."""
     base = torch.tanh(angles.float()) * dt.float().permute(0, 2, 1)[..., None] * math.pi

 __all__ = ["forward", "DecodeState", "cumulative_angles", "chunk_decay"]
+def _register_fakes():
+    """Meta kernels for the three ops, so a compiled graph does not break here.
+    The step op mutates four of its arguments; those carry `Tensor(a!)` in the
+    schema, so functionalization handles them and a fake only has to produce the
+    return value. The library name is read off a live op because the kernel
+    builder and a local JIT build choose different ones.
+    """
+    try:
+        lib = ops.mamba3_mimo_fwd.default._schema.name.split("::")[0]
+    except Exception:
+        return ()
+    def like_f32(t):
+        return torch.empty_like(t, dtype=torch.float32)
+    def register(name, fn):
+        if not hasattr(ops, name):
+            return False
+        torch.library.register_fake(f"{lib}::{name}", fn)
+        return True
+    def _step(q, k, v, z, adt, dt, trap, q_bias, k_bias, angles, mimo_v, mimo_o,
+              mimo_z, D, angle_in, angle_out, S_st, kprev_in, vprev_in,
+              kprev_out, vprev_out):
+        return torch.empty_like(v)
+    def _fwd(q, k, v, z, q_bias, k_bias, mimo_v, mimo_o, mimo_z, D, angles, dt,
+             trap, dA_cs, dA_cs_rev, chunk_size, norm_weight, norm_eps):
+        return torch.empty_like(v)
+    def _bwd(q, k, v, z, q_bias, k_bias, mimo_v, mimo_o, mimo_z, D, angles, dt,
+             trap, dA_cs, dA_cs_rev, chunk_size, norm_weight, norm_eps, dy):
+        H, P = v.shape[-2], v.shape[-1]
+        f32 = torch.float32
+        return [like_f32(q), like_f32(k), like_f32(v), like_f32(v),
+                like_f32(q_bias), like_f32(k_bias), like_f32(mimo_v),
+                like_f32(mimo_o), like_f32(mimo_v),
+                v.new_empty((H,), dtype=f32), v.new_empty((H, P), dtype=f32),
+                like_f32(dA_cs), like_f32(dA_cs_rev), like_f32(dt),
+                like_f32(trap), like_f32(angles)]
+    done = [n for n, f in (("mamba3_mimo_step", _step),
+                           ("mamba3_mimo_fwd", _fwd),
+                           ("mamba3_mimo_bwd", _bwd)) if register(n, f)]
+    return tuple(done)
+FAKE_OPS = _register_fakes()
 def cumulative_angles(angles, dt):
     """Per-step rotation rates (B,S,H,Na) and dt (B,H,S) -> cumulative angles, float32."""
     base = torch.tanh(angles.float()) * dt.float().permute(0, 2, 1)[..., None] * math.pi

build/torch213-cxx11-cu130-x86_64-linux/_mamba3_cuda_95022f8.abi3.so DELETED Viewed

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:8397f55e5c3b107d74dbfc7dcad82645104fb180b04efc7e17c5986006b5bc53
-size 9781448

build/torch213-cxx11-cu130-x86_64-linux/_mamba3_cuda_c549070.abi3.so ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:a081e8b4d7b21bfaa50c82b0163e155916f42d38fed8aa1a035fd29b60e00066
+size 9781808

build/torch213-cxx11-cu130-x86_64-linux/_ops.py CHANGED Viewed

@@ -1,9 +1,9 @@
 import torch
-from . import _mamba3_cuda_95022f8
-ops = torch.ops._mamba3_cuda_95022f8
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_mamba3_cuda_95022f8::{op_name}"

 import torch
+from . import _mamba3_cuda_c549070
+ops = torch.ops._mamba3_cuda_c549070
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
+    return f"_mamba3_cuda_c549070::{op_name}"

build/torch213-cxx11-cu130-x86_64-linux/metadata.json CHANGED Viewed

@@ -1,6 +1,6 @@
 {
   "name": "mamba3",
-  "id": "_mamba3_cuda_95022f8",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],
@@ -18,9 +18,9 @@
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "__init__.py": "zWrrGqu25MORtMITQd22NHExxrHe4vYeSx4kAjDl/94=",
-      "_mamba3_cuda_95022f8.abi3.so": "g5f1Xlw7EH102/x9ytgmRRBPsYCwTvx+F8WYYAa1vFM=",
-      "_ops.py": "5TOqHeKVQS1MnkoVuOdCv8fNnWQNe+S86igdxnRw8G8=",
       "mamba3/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
     }
   },
@@ -31,7 +31,7 @@
       "dirty": false
     },
     "kernel": {
-      "sha": "95022f83b9b4c5d080443e6e91a8549afd7ae80b",
       "dirty": false
     }
   }

 {
   "name": "mamba3",
+  "id": "_mamba3_cuda_c549070",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],
   "digest": {
     "algorithm": "sha256",
     "files": {
+      "__init__.py": "zKP4l3anEqOcUuw5TEp4mkMhr0OFWQe9cUTq0Wu38aQ=",
+      "_mamba3_cuda_c549070.abi3.so": "oIHotNeyG/qlDIKwFj4VWRb0LTj+2KoaA1/Sm2DgAGY=",
+      "_ops.py": "Hw1EtZGXlC4SXcGIUpbQwVl8+sKMlb72u1JIHgKqHT8=",
       "mamba3/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
     }
   },
       "dirty": false
     },
     "kernel": {
+      "sha": "c54907036afbc32fae8bddb75ae7aa2bbaf3dde3",
       "dirty": false
     }
   }

build/torch213-cxx11-cu132-x86_64-linux/__init__.py CHANGED Viewed

@@ -13,6 +13,57 @@ from ._ops import ops
 __all__ = ["forward", "DecodeState", "cumulative_angles", "chunk_decay"]
 def cumulative_angles(angles, dt):
     """Per-step rotation rates (B,S,H,Na) and dt (B,H,S) -> cumulative angles, float32."""
     base = torch.tanh(angles.float()) * dt.float().permute(0, 2, 1)[..., None] * math.pi

 __all__ = ["forward", "DecodeState", "cumulative_angles", "chunk_decay"]
+def _register_fakes():
+    """Meta kernels for the three ops, so a compiled graph does not break here.
+    The step op mutates four of its arguments; those carry `Tensor(a!)` in the
+    schema, so functionalization handles them and a fake only has to produce the
+    return value. The library name is read off a live op because the kernel
+    builder and a local JIT build choose different ones.
+    """
+    try:
+        lib = ops.mamba3_mimo_fwd.default._schema.name.split("::")[0]
+    except Exception:
+        return ()
+    def like_f32(t):
+        return torch.empty_like(t, dtype=torch.float32)
+    def register(name, fn):
+        if not hasattr(ops, name):
+            return False
+        torch.library.register_fake(f"{lib}::{name}", fn)
+        return True
+    def _step(q, k, v, z, adt, dt, trap, q_bias, k_bias, angles, mimo_v, mimo_o,
+              mimo_z, D, angle_in, angle_out, S_st, kprev_in, vprev_in,
+              kprev_out, vprev_out):
+        return torch.empty_like(v)
+    def _fwd(q, k, v, z, q_bias, k_bias, mimo_v, mimo_o, mimo_z, D, angles, dt,
+             trap, dA_cs, dA_cs_rev, chunk_size, norm_weight, norm_eps):
+        return torch.empty_like(v)
+    def _bwd(q, k, v, z, q_bias, k_bias, mimo_v, mimo_o, mimo_z, D, angles, dt,
+             trap, dA_cs, dA_cs_rev, chunk_size, norm_weight, norm_eps, dy):
+        H, P = v.shape[-2], v.shape[-1]
+        f32 = torch.float32
+        return [like_f32(q), like_f32(k), like_f32(v), like_f32(v),
+                like_f32(q_bias), like_f32(k_bias), like_f32(mimo_v),
+                like_f32(mimo_o), like_f32(mimo_v),
+                v.new_empty((H,), dtype=f32), v.new_empty((H, P), dtype=f32),
+                like_f32(dA_cs), like_f32(dA_cs_rev), like_f32(dt),
+                like_f32(trap), like_f32(angles)]
+    done = [n for n, f in (("mamba3_mimo_step", _step),
+                           ("mamba3_mimo_fwd", _fwd),
+                           ("mamba3_mimo_bwd", _bwd)) if register(n, f)]
+    return tuple(done)
+FAKE_OPS = _register_fakes()
 def cumulative_angles(angles, dt):
     """Per-step rotation rates (B,S,H,Na) and dt (B,H,S) -> cumulative angles, float32."""
     base = torch.tanh(angles.float()) * dt.float().permute(0, 2, 1)[..., None] * math.pi

build/torch213-cxx11-cu132-x86_64-linux/_mamba3_cuda_95022f8.abi3.so DELETED Viewed

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:30a1344dae5df07c7eb0a2498cdfe1fa192b135ded4bb438b33f633d32b38982
-size 9900232

build/torch213-cxx11-cu132-x86_64-linux/_mamba3_cuda_c549070.abi3.so ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:74b0abe576e73608430da49662a126fd8cca0453a390369a62f3457050517690
+size 9904688

build/torch213-cxx11-cu132-x86_64-linux/_ops.py CHANGED Viewed

@@ -1,9 +1,9 @@
 import torch
-from . import _mamba3_cuda_95022f8
-ops = torch.ops._mamba3_cuda_95022f8
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_mamba3_cuda_95022f8::{op_name}"

 import torch
+from . import _mamba3_cuda_c549070
+ops = torch.ops._mamba3_cuda_c549070
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
+    return f"_mamba3_cuda_c549070::{op_name}"

build/torch213-cxx11-cu132-x86_64-linux/metadata.json CHANGED Viewed

@@ -1,6 +1,6 @@
 {
   "name": "mamba3",
-  "id": "_mamba3_cuda_95022f8",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],
@@ -18,9 +18,9 @@
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "__init__.py": "zWrrGqu25MORtMITQd22NHExxrHe4vYeSx4kAjDl/94=",
-      "_mamba3_cuda_95022f8.abi3.so": "MKE0Ta5d8Hx+sKJJjN/h+hkrE13tS7Q4sz9jPTKziYI=",
-      "_ops.py": "5TOqHeKVQS1MnkoVuOdCv8fNnWQNe+S86igdxnRw8G8=",
       "mamba3/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
     }
   },
@@ -31,7 +31,7 @@
       "dirty": false
     },
     "kernel": {
-      "sha": "95022f83b9b4c5d080443e6e91a8549afd7ae80b",
       "dirty": false
     }
   }

 {
   "name": "mamba3",
+  "id": "_mamba3_cuda_c549070",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],
   "digest": {
     "algorithm": "sha256",
     "files": {
+      "__init__.py": "zKP4l3anEqOcUuw5TEp4mkMhr0OFWQe9cUTq0Wu38aQ=",
+      "_mamba3_cuda_c549070.abi3.so": "dLCr5XbnNghDDaSWYqEm/YzKBFOjkDaaYvNFcFBRdpA=",
+      "_ops.py": "Hw1EtZGXlC4SXcGIUpbQwVl8+sKMlb72u1JIHgKqHT8=",
       "mamba3/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
     }
   },
       "dirty": false
     },
     "kernel": {
+      "sha": "c54907036afbc32fae8bddb75ae7aa2bbaf3dde3",
       "dirty": false
     }
   }