Enable Torch 2.8 build

build.toml

@@ -1,61 +1,77 @@
 [general]
 name = "moe"
+universal = false
 
 [torch]
+include = ["."]
+pyext = [
+    "py",
+    "json",
+]
 src = [
-  "core/scalar_type.hpp",
-  "torch-ext/torch_binding.cpp",
-  "torch-ext/torch_binding.h",
+    "core/scalar_type.hpp",
+    "torch-ext/torch_binding.cpp",
+    "torch-ext/torch_binding.h",
 ]
-include = ["."]
-pyext = ["py", "json"]
 
-[kernel.fp8]
-src = [
-  "cuda_compat.h",
-  "dispatch_utils.h",
-  "fp8/amd/hip_float8.h",
-  "fp8/amd/hip_float8_impl.h",
-  "fp8/common.cu",
-  "fp8/common.cuh",
-  "fp8/vectorization.cuh",
+[kernel.moe-marlin]
+backend = "cuda"
+cuda-capabilities = [
+    "8.0",
+    "8.6",
+    "8.7",
+    "8.9",
+    "9.0",
+    "10.0",
+    "10.1",
+    "12.0",
 ]
-include = ["."]
 depends = ["torch"]
-
-
-[kernel.moe]
+include = ["."]
 src = [
-  "cuda_compat.h",
-  "dispatch_utils.h",
-  "moe/moe_align_sum_kernels.cu",
-  "moe/moe_wna16.cu",
-  "moe/moe_wna16_utils.h",
-  "moe/topk_softmax_kernels.cu",
+    "core/exception.hpp",
+    "core/scalar_type.hpp",
+    "marlin-moe/marlin_moe_ops.cu",
+    "marlin-moe/marlin_kernels/marlin_moe_kernel_ku4.cu",
+    "marlin-moe/marlin_kernels/marlin_moe_kernel_ku8b128.cu",
+    "marlin-moe/marlin_kernels/marlin_moe_kernel.h",
+    "marlin-moe/marlin_kernels/marlin_moe_kernel_ku4.h",
+    "marlin-moe/marlin_kernels/marlin_moe_kernel_ku4b8.h",
+    "marlin-moe/marlin_kernels/marlin_moe_kernel_ku4b8.cu",
+    "marlin-moe/marlin_kernels/marlin_moe_kernel_ku8b128.h",
 ]
-depends = ["torch"]
 
-[kernel.moe-marlin]
-cuda-capabilities = ["8.0", "8.6", "8.7", "8.9", "9.0", "10.0", "10.1", "12.0"]
+[kernel.activation]
+backend = "cuda"
+depends = ["torch"]
 src = [
-  "core/exception.hpp",
-  "core/scalar_type.hpp",
-  "marlin-moe/marlin_moe_ops.cu",
-  "marlin-moe/marlin_kernels/marlin_moe_kernel_ku4.cu",
-  "marlin-moe/marlin_kernels/marlin_moe_kernel_ku8b128.cu",
-  "marlin-moe/marlin_kernels/marlin_moe_kernel.h",
-  "marlin-moe/marlin_kernels/marlin_moe_kernel_ku4.h",
-  "marlin-moe/marlin_kernels/marlin_moe_kernel_ku4b8.h",
-  "marlin-moe/marlin_kernels/marlin_moe_kernel_ku4b8.cu",
-  "marlin-moe/marlin_kernels/marlin_moe_kernel_ku8b128.h",
+    "activation/activation_kernels.cu",
+    "activation/cuda_compat.h",
+    "activation/dispatch_utils.h",
 ]
-include = ["."]
-depends = ["torch"]
 
-[kernel.activation]
+[kernel.fp8]
+backend = "cuda"
+depends = ["torch"]
+include = ["."]
 src = [
-  "activation/activation_kernels.cu",
-  "activation/cuda_compat.h",
-  "activation/dispatch_utils.h",
+    "cuda_compat.h",
+    "dispatch_utils.h",
+    "fp8/amd/hip_float8.h",
+    "fp8/amd/hip_float8_impl.h",
+    "fp8/common.cu",
+    "fp8/common.cuh",
+    "fp8/vectorization.cuh",
 ]
+
+[kernel.moe]
+backend = "cuda"
 depends = ["torch"]
+src = [
+    "cuda_compat.h",
+    "dispatch_utils.h",
+    "moe/moe_align_sum_kernels.cu",
+    "moe/moe_wna16.cu",
+    "moe/moe_wna16_utils.h",
+    "moe/topk_softmax_kernels.cu",
+]

flake.lock

@@ -1,6 +1,21 @@
 {
   "nodes": {
     "flake-compat": {
+      "locked": {
+        "lastModified": 1747046372,
+        "narHash": "sha256-CIVLLkVgvHYbgI2UpXvIIBJ12HWgX+fjA8Xf8PUmqCY=",
+        "owner": "edolstra",
+        "repo": "flake-compat",
+        "rev": "9100a0f413b0c601e0533d1d94ffd501ce2e7885",
+        "type": "github"
+      },
+      "original": {
+        "owner": "edolstra",
+        "repo": "flake-compat",
+        "type": "github"
+      }
+    },
+    "flake-compat_2": {
       "locked": {
         "lastModified": 1733328505,
         "narHash": "sha256-NeCCThCEP3eCl2l/+27kNNK7QrwZB1IJCrXfrbv5oqU=",
@@ -33,61 +48,83 @@
         "type": "github"
       }
     },
-    "kernel-builder": {
+    "flake-utils_2": {
       "inputs": {
-        "flake-compat": "flake-compat",
-        "flake-utils": "flake-utils",
-        "nixpkgs": "nixpkgs",
-        "rocm-nix": "rocm-nix"
+        "systems": "systems_2"
       },
       "locked": {
-        "lastModified": 1744976941,
-        "narHash": "sha256-+csrhVaT6Mj2j1FM7P2BDITvf1Xwj2AKdMm0IKZK340=",
-        "owner": "huggingface",
-        "repo": "kernel-builder",
-        "rev": "0a278c2e9aaf6003a4ec6fe35c7158624762de5a",
+        "lastModified": 1731533236,
+        "narHash": "sha256-l0KFg5HjrsfsO/JpG+r7fRrqm12kzFHyUHqHCVpMMbI=",
+        "owner": "numtide",
+        "repo": "flake-utils",
+        "rev": "11707dc2f618dd54ca8739b309ec4fc024de578b",
         "type": "github"
       },
       "original": {
-        "owner": "huggingface",
-        "repo": "kernel-builder",
+        "owner": "numtide",
+        "repo": "flake-utils",
         "type": "github"
       }
     },
-    "nixpkgs": {
+    "hf-nix": {
+      "inputs": {
+        "flake-compat": "flake-compat_2",
+        "flake-utils": "flake-utils_2",
+        "nixpkgs": "nixpkgs"
+      },
       "locked": {
-        "lastModified": 1743559129,
-        "narHash": "sha256-7gpAWsENV3tY2HmeHYQ2MoQxGpys+jQWnkS/BHAMXVk=",
-        "owner": "nixos",
-        "repo": "nixpkgs",
-        "rev": "adae22bea8bcc0aa2fd6e8732044660fb7755f5e",
+        "lastModified": 1753354560,
+        "narHash": "sha256-vmOfRmr0Qm/IbZTWB2sBn+UFrABSTTA/cTg+m27Yt/E=",
+        "owner": "huggingface",
+        "repo": "hf-nix",
+        "rev": "7f2aceda2a2e72cd573bdb25e5c0667fd75f89d3",
         "type": "github"
       },
       "original": {
-        "owner": "nixos",
-        "ref": "nixos-unstable-small",
-        "repo": "nixpkgs",
+        "owner": "huggingface",
+        "repo": "hf-nix",
         "type": "github"
       }
     },
-    "rocm-nix": {
+    "kernel-builder": {
       "inputs": {
+        "flake-compat": "flake-compat",
+        "flake-utils": "flake-utils",
+        "hf-nix": "hf-nix",
         "nixpkgs": [
           "kernel-builder",
+          "hf-nix",
           "nixpkgs"
         ]
       },
       "locked": {
-        "lastModified": 1743085847,
-        "narHash": "sha256-uWG29p+nhZmGRV1LffWwRGjwtPIXeu1F0YTQbXgB+GU=",
+        "lastModified": 1753354632,
+        "narHash": "sha256-31SX3Raiyx0qCuY9JSlx9ZZgxljeUxvW+JdujjxbofQ=",
         "owner": "huggingface",
-        "repo": "rocm-nix",
-        "rev": "245cdc9bfb4bfafa818711c5f5e0b889afe1ba39",
+        "repo": "kernel-builder",
+        "rev": "524b628fd8e58525dbd28455bffb0628092c5265",
         "type": "github"
       },
       "original": {
         "owner": "huggingface",
-        "repo": "rocm-nix",
+        "ref": "torch-2.8",
+        "repo": "kernel-builder",
+        "type": "github"
+      }
+    },
+    "nixpkgs": {
+      "locked": {
+        "lastModified": 1752785354,
+        "narHash": "sha256-Y33ryUz7MPqKrZwlbQcsYCUz2jAJCacRf8jbs0tYUlA=",
+        "owner": "nixos",
+        "repo": "nixpkgs",
+        "rev": "d38025438a6ee456758dc03188ca6873a415463b",
+        "type": "github"
+      },
+      "original": {
+        "owner": "nixos",
+        "repo": "nixpkgs",
+        "rev": "d38025438a6ee456758dc03188ca6873a415463b",
         "type": "github"
       }
     },
@@ -110,6 +147,21 @@
         "repo": "default",
         "type": "github"
       }
+    },
+    "systems_2": {
+      "locked": {
+        "lastModified": 1681028828,
+        "narHash": "sha256-Vy1rq5AaRuLzOxct8nz4T6wlgyUR7zLU309k9mBC768=",
+        "owner": "nix-systems",
+        "repo": "default",
+        "rev": "da67096a3b9bf56a91d16901293e51ba5b49a27e",
+        "type": "github"
+      },
+      "original": {
+        "owner": "nix-systems",
+        "repo": "default",
+        "type": "github"
+      }
     }
   },
   "root": "root",

flake.nix

@@ -2,7 +2,7 @@
   description = "Flake for activation kernels";
 
   inputs = {
-    kernel-builder.url = "github:huggingface/kernel-builder";
+    kernel-builder.url = "github:huggingface/kernel-builder/torch-2.8";
   };
 
   outputs =