Update build with torch.compile support

build/torch27-cxx11-rocm63-x86_64-linux/megablocks/{_megablocks_7ca55d1_dirty.abi3.so → _megablocks_9193c11_dirty.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:34fd20a547aa0ef2f6943cf8a996e30ac966ed79e4715d81ce77bd28799b1af6
+oid sha256:9c1bcdc65f532386bf11a2d450077747448f968b9422ed64b0f0fe45955ec9ba
 size 5573528

build/torch27-cxx11-rocm63-x86_64-linux/megablocks/_megablocks_ed040f3_dirty.abi3.so

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

build/torch27-cxx11-rocm63-x86_64-linux/megablocks/_ops.py

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

build/torch27-cxx11-rocm63-x86_64-linux/megablocks/layers.py

@@ -871,6 +871,8 @@ def get_device_mesh(model):
 
 
 class MegaBlocksMoeMLP(torch.nn.Module):
+    
+    can_torch_compile = True
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         moe_top_k = getattr(self.router, "top_k", 4)

build/torch27-cxx11-rocm64-x86_64-linux/megablocks/{_megablocks_7ca55d1_dirty.abi3.so → _megablocks_9193c11_dirty.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:fa7184151ab4959a183ced18df4e1f0f6a0b8261e0009e5db837357e549c1ef7
+oid sha256:8da15cb80ae3e6fdae0341bdefbe64eb8099e4eb20f2c0e5770fa36e44a4acbd
 size 5574080

build/torch27-cxx11-rocm64-x86_64-linux/megablocks/_megablocks_ed040f3_dirty.abi3.so

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

build/torch27-cxx11-rocm64-x86_64-linux/megablocks/_ops.py

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

build/torch27-cxx11-rocm64-x86_64-linux/megablocks/layers.py

@@ -871,6 +871,8 @@ def get_device_mesh(model):
 
 
 class MegaBlocksMoeMLP(torch.nn.Module):
+    
+    can_torch_compile = True
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         moe_top_k = getattr(self.router, "top_k", 4)

build/torch28-cxx11-rocm63-x86_64-linux/megablocks/{_megablocks_7ca55d1_dirty.abi3.so → _megablocks_9193c11_dirty.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:13df995688483e82cab1dfdd1f16cd5ced4a0e8f8aa4143efc09d37e92847c74
+oid sha256:162379f95323d2ee96b87657ed63db30d4af837816f5d77c4daa9e1c7737b8bc
 size 5577848

build/torch28-cxx11-rocm63-x86_64-linux/megablocks/_megablocks_ed040f3_dirty.abi3.so

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

build/torch28-cxx11-rocm63-x86_64-linux/megablocks/_ops.py

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

build/torch28-cxx11-rocm63-x86_64-linux/megablocks/layers.py

@@ -871,6 +871,8 @@ def get_device_mesh(model):
 
 
 class MegaBlocksMoeMLP(torch.nn.Module):
+    
+    can_torch_compile = True
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         moe_top_k = getattr(self.router, "top_k", 4)

build/torch28-cxx11-rocm64-x86_64-linux/megablocks/{_megablocks_7ca55d1_dirty.abi3.so → _megablocks_9193c11_dirty.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:4478dfa4e48d307fdbca4157a971622588ecafeee5c5bff008c0b6dc0bbfc5e3
+oid sha256:f4cac3717fcfc99608da2b0a065203645504372fb8f918541d547a6b7f41d16d
 size 5578400

build/torch28-cxx11-rocm64-x86_64-linux/megablocks/_megablocks_ed040f3_dirty.abi3.so

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

build/torch28-cxx11-rocm64-x86_64-linux/megablocks/_ops.py

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

build/torch28-cxx11-rocm64-x86_64-linux/megablocks/layers.py

@@ -871,6 +871,8 @@ def get_device_mesh(model):
 
 
 class MegaBlocksMoeMLP(torch.nn.Module):
+    
+    can_torch_compile = True
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         moe_top_k = getattr(self.router, "top_k", 4)

flake.lock

@@ -98,11 +98,11 @@
         ]
       },
       "locked": {
-        "lastModified": 1754043127,
-        "narHash": "sha256-U/u8vKscZMwFb5p7hMcSRnSwn5B/7VDDS3hhyrhm//s=",
+        "lastModified": 1754384957,
+        "narHash": "sha256-t55mgR9ySbw+vvxz+exl1o/PbgFOTj1zCKB6Vi9jDW8=",
         "owner": "huggingface",
         "repo": "kernel-builder",
-        "rev": "2a079185c4bd1099f3c6ce53f3168226f0273880",
+        "rev": "a4176c1791b19619d752c2562bb2614509cf2257",
         "type": "github"
       },
       "original": {

result

@@ -1 +1 @@
-/nix/store/m7dzl5drmnsvwv9hma6m58n769x2ib5l-torch-ext-bundle
+/nix/store/8cwl78s3nn8wx38n1qk8nx6014kp0p4g-torch-ext-bundle

torch-ext/megablocks/_megablocks_g6r5q4zwqmcls.abi3.so

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

torch-ext/megablocks/_megablocks_lboajlxwrdgqk.abi3.so

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

torch-ext/megablocks/_megablocks_twnwbxx53zgza.abi3.so

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

torch-ext/megablocks/layers.py

@@ -871,6 +871,8 @@ def get_device_mesh(model):
 
 
 class MegaBlocksMoeMLP(torch.nn.Module):
+    
+    can_torch_compile = True
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         moe_top_k = getattr(self.router, "top_k", 4)