Update build

build.toml

@@ -1,5 +1,5 @@
 [general]
-name = "megablocks"
+name = "megablocks_rocm"
 universal = false
 
 [torch]

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

flake.lock

@@ -73,11 +73,11 @@
         "nixpkgs": "nixpkgs"
       },
       "locked": {
-        "lastModified": 1753867039,
-        "narHash": "sha256-lROFmwSsxtjNsf7U/mQ/jW4GQbnwUQstu0MJrhPdGhg=",
+        "lastModified": 1754038838,
+        "narHash": "sha256-oHigCT4z0ayyLyEuxdZooSXRAZP8lfOkZHzY1lx1U50=",
         "owner": "huggingface",
         "repo": "hf-nix",
-        "rev": "dab75553a43839d0dad876c4e59bbcd8b41acd72",
+        "rev": "336f781fa284e193baa3d4c3ce3f95fb34e9ffad",
         "type": "github"
       },
       "original": {
@@ -98,11 +98,11 @@
         ]
       },
       "locked": {
-        "lastModified": 1753867163,
-        "narHash": "sha256-x3D5QthCR+buQ7oX9+HRhYt4sqxHzYqLYdK+HHAf0Qc=",
+        "lastModified": 1754043127,
+        "narHash": "sha256-U/u8vKscZMwFb5p7hMcSRnSwn5B/7VDDS3hhyrhm//s=",
         "owner": "huggingface",
         "repo": "kernel-builder",
-        "rev": "768759f1423ad96ec1232fb16f936718a6742cc9",
+        "rev": "2a079185c4bd1099f3c6ce53f3168226f0273880",
         "type": "github"
       },
       "original": {

flake.nix

@@ -1,62 +1,34 @@
 {
+  description = "Flake for megablocks_moe kernel";
+
   inputs = {
-    hf-nix.url = "github:huggingface/hf-nix/mktestdocs-0.2.5";
-    nixpkgs.follows = "hf-nix/nixpkgs";
-    flake-utils.url = "github:numtide/flake-utils";
+    kernel-builder.url = "github:huggingface/kernel-builder/torch-2.8";
   };
+
   outputs =
     {
       self,
-      nixpkgs,
-      flake-utils,
-      hf-nix,
+      kernel-builder,
     }:
-    flake-utils.lib.eachDefaultSystem (
-      system:
-      let
-        pkgs = import nixpkgs {
-          inherit system;
-          config = hf-nix.lib.config system;
-          overlays = [
-            hf-nix.overlays.default
-          ];
-        };
-      in
-      {
-        formatter = pkgs.nixfmt-tree;
-        devShells = with pkgs; rec {
-          default = mkShell {
-            nativeBuildInputs = [
-              # For hf-doc-builder.
-              nodejs
-            ];
-            buildInputs =
-              [
-                black
-                mypy
-                pyright
-                ruff
-              ]
-              ++ (with python3.pkgs; [
-                docutils
-                huggingface-hub
-                mktestdocs
-                pytest
-                pytest-benchmark
-                pyyaml
-                torch
-                types-pyyaml
-                venvShellHook
-              ]);
+    kernel-builder.lib.genFlakeOutputs {
+      path = ./.;
+      rev = self.shortRev or self.dirtyShortRev or self.lastModifiedDate;
 
-            venvDir = "./.venv";
+      pythonCheckInputs = pkgs: with pkgs; [ 
+        tqdm
+        py-cpuinfo
+        importlib-metadata
+        torchmetrics
+      ];
 
-            postVenvCreation = ''
-              unset SOURCE_DATE_EPOCH
-              ( python -m pip install --no-build-isolation --no-dependencies -e . )
-            '';
-          };
-        };
-      }
-    );
+      torchVersions = defaultVersions: defaultVersions ++ [
+        {
+          torchVersion = "2.7";
+          rocmVersion = "6.4.2";
+          cxx11Abi = true;
+          systems = [ "x86_64-linux" ];
+          bundleBuild = true;
+        }
+      ];
+    };
 }

result

@@ -0,0 +1 @@
+/nix/store/ssm0vfvdzhb48ifk4zmky6hpzfccwggn-torch-ext-bundle