Update Nix kernel matrix config and safetensors import

build.toml

@@ -1,11 +1,11 @@
 [general]
 name = "higgs-kernels"
-universal = false
+backends = ["cuda"]
 
 [torch]
 src = [
-  "torch-ext/torch_binding.cpp",
-  "torch-ext/torch_binding.h",
+    "torch-ext/torch_binding.cpp",
+    "torch-ext/torch_binding.h",
 ]
 
 [kernel.dequant]
@@ -13,12 +13,12 @@ backend = "cuda"
 depends = ["torch"]
 src = ["csrc/dequant.cu"]
 
-[kernel.quant_f16]
+[kernel.quant_bf16]
 backend = "cuda"
 depends = ["torch"]
-src = ["csrc/quant_f16.cu"]
+src = ["csrc/quant_bf16.cu"]
 
-[kernel.quant_bf16]
+[kernel.quant_f16]
 backend = "cuda"
 depends = ["torch"]
-src = ["csrc/quant_bf16.cu"]
+src = ["csrc/quant_f16.cu"]

flake.lock

@@ -0,0 +1,117 @@
+{
+  "nodes": {
+    "flake-compat": {
+      "locked": {
+        "lastModified": 1765121682,
+        "narHash": "sha256-4VBOP18BFeiPkyhy9o4ssBNQEvfvv1kXkasAYd0+rrA=",
+        "owner": "edolstra",
+        "repo": "flake-compat",
+        "rev": "65f23138d8d09a92e30f1e5c87611b23ef451bf3",
+        "type": "github"
+      },
+      "original": {
+        "owner": "edolstra",
+        "repo": "flake-compat",
+        "type": "github"
+      }
+    },
+    "flake-utils": {
+      "inputs": {
+        "systems": "systems"
+      },
+      "locked": {
+        "lastModified": 1731533236,
+        "narHash": "sha256-l0KFg5HjrsfsO/JpG+r7fRrqm12kzFHyUHqHCVpMMbI=",
+        "owner": "numtide",
+        "repo": "flake-utils",
+        "rev": "11707dc2f618dd54ca8739b309ec4fc024de578b",
+        "type": "github"
+      },
+      "original": {
+        "owner": "numtide",
+        "repo": "flake-utils",
+        "type": "github"
+      }
+    },
+    "kernel-builder": {
+      "inputs": {
+        "flake-compat": "flake-compat",
+        "flake-utils": "flake-utils",
+        "nixpkgs": "nixpkgs",
+        "rust-overlay": "rust-overlay"
+      },
+      "locked": {
+        "lastModified": 1771009495,
+        "narHash": "sha256-/2nRZQbiIvlBoFShfpnoOD27ZgeYBV26bI6E5WsBJws=",
+        "owner": "huggingface",
+        "repo": "kernels",
+        "rev": "b079fd8c66612177cc8edd13292613abb4de994c",
+        "type": "github"
+      },
+      "original": {
+        "owner": "huggingface",
+        "repo": "kernels",
+        "type": "github"
+      }
+    },
+    "nixpkgs": {
+      "locked": {
+        "lastModified": 1766341660,
+        "narHash": "sha256-4yG6vx7Dddk9/zh45Y2KM82OaRD4jO3HA9r98ORzysA=",
+        "owner": "NixOS",
+        "repo": "nixpkgs",
+        "rev": "26861f5606e3e4d1400771b513cc63e5f70151a6",
+        "type": "github"
+      },
+      "original": {
+        "owner": "NixOS",
+        "ref": "nixos-unstable-small",
+        "repo": "nixpkgs",
+        "type": "github"
+      }
+    },
+    "root": {
+      "inputs": {
+        "kernel-builder": "kernel-builder"
+      }
+    },
+    "rust-overlay": {
+      "inputs": {
+        "nixpkgs": [
+          "kernel-builder",
+          "nixpkgs"
+        ]
+      },
+      "locked": {
+        "lastModified": 1769050281,
+        "narHash": "sha256-1H8DN4UZgEUqPUA5ecHOufLZMscJ4IlcGaEftaPtpBY=",
+        "owner": "oxalica",
+        "repo": "rust-overlay",
+        "rev": "6deef0585c52d9e70f96b6121207e1496d4b0c49",
+        "type": "github"
+      },
+      "original": {
+        "owner": "oxalica",
+        "repo": "rust-overlay",
+        "type": "github"
+      }
+    },
+    "systems": {
+      "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",
+  "version": 7
+}

torch-ext/higgs_kernels/__init__.py

@@ -4,14 +4,19 @@ import functools
 import torch
 from ._ops import ops
 
-import safetensors.torch
-
 PKG_PATH = os.path.dirname(os.path.realpath(__file__))
 
 
 @functools.cache
 def load_optimal_grid_2_256(device="cpu", dtype=torch.float16):
-    return safetensors.torch.load_file(
+    try:
+        import safetensors.torch as safetensors_torch
+    except ModuleNotFoundError as exc:
+        raise ModuleNotFoundError(
+            "load_optimal_grid_2_256 requires safetensors"
+        ) from exc
+
+    return safetensors_torch.load_file(
         os.path.join(PKG_PATH, "grids.safetensors"), device=device
     )["2_256"].to(dtype)