migrate to new huggingface/kernels builder schema

build.toml

@@ -1,10 +1,16 @@
+# Hugging Face kernel manifest (new huggingface/kernels schema).
+# Pure-Python kernel (CUTLASS Python DSL, JIT-compiled on the GPU at runtime):
+# no compiled sources -> declared with [torch-noarch]. CUDA-only (CUTLASS is NVIDIA).
 [general]
-name = "my_softmax_function"
+name = "my-softmax-function"
+version = 1
+edition = 5
+license = "Apache-2.0"
 backends = ["cuda"]
 
 [general.cuda]
+# nvidia-cutlass-dsl 4.3.0 requires CUDA >= 12.8.
 minver = "12.8"
 python-depends = ["nvidia-cutlass-dsl"]
 
-[torch]
-src = []
+[torch-noarch]

build/torch-cuda/_kernel.py

@@ -10,12 +10,11 @@ import math
 import torch
 from cutlass.cute.runtime import from_dlpack
 
+from ._ops import add_op_namespace_prefix
 from .attention_v3 import solve
 
-_OP_NAME = "my_softmax_function::softmax_attention"
 
-
-@torch.library.custom_op(_OP_NAME, mutates_args=())
+@torch.library.custom_op(add_op_namespace_prefix("softmax_attention"), mutates_args=())
 def softmax_attention(
     Q: torch.Tensor,        # (M, d)
     K: torch.Tensor,        # (N, d)
@@ -65,4 +64,5 @@ def attention(Q, K, V, scale=None):
     """Convenience entry point with a default 1/sqrt(d) scale."""
     if scale is None:
         scale = 1.0 / math.sqrt(Q.shape[-1])
-    return torch.ops.my_softmax_function.softmax_attention(Q, K, V, scale)
+    # Call the registered op directly (its namespace is build-unique).
+    return softmax_attention(Q, K, V, scale)

build/torch-cuda/_ops.py

@@ -1,8 +1,38 @@
 import torch
-ops = torch.ops._my_softmax_function_2af76fc
 
-def add_op_namespace_prefix(op_name: str):
+def get_backend() -> str:
+    """Detect the backend by inspecting torch."""
+    import torch
+
+    if hasattr(torch, "neuron"):
+        # Needs to be sorted before specific Torch builds, since Neuron
+        # extension can be loaded into e.g. CUDA Torch builds.
+        return "neuron"
+    elif torch.version.cuda is not None:
+        return "cuda"
+    elif torch.version.hip is not None:
+        return "rocm"
+    elif torch.backends.mps.is_available():
+        return "metal"
+    elif hasattr(torch.version, "xpu") and torch.version.xpu is not None:
+        return "xpu"
+    else:
+        return "cpu"
+
+
+def _find_ops_name() -> str:
+    kernel_name = "my_softmax_function"
+    unique_id = "490d9fe_dirty"
+    backend = get_backend()
+    return f"_{kernel_name}_{backend}_{unique_id}"
+
+
+_OPS_NAME = _find_ops_name()
+
+ops = getattr(torch.ops, _OPS_NAME)
+
+def add_op_namespace_prefix(op_name: str) -> str:
     """
     Prefix op by namespace.
     """
-    return f"_my_softmax_function_2af76fc::{op_name}"
+    return f"{_OPS_NAME}::{op_name}"

build/torch-cuda/metadata.json

@@ -1,5 +1,32 @@
 {
+  "name": "my-softmax-function",
+  "id": "_my_softmax_function_cuda_490d9fe_dirty",
+  "version": 1,
+  "license": "Apache-2.0",
   "python-depends": [
     "nvidia-cutlass-dsl"
-  ]
+  ],
+  "backend": {
+    "type": "cuda"
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "ZpvyIhcU8bz4Xcc1WhVskt3WR3Yx9CCzS8+Jl1DOXGU=",
+      "_kernel.py": "xDfkK2UVoOoZY8Jn4VtzTxrzdo6MfXVm72qu2S+/jp8=",
+      "_ops.py": "ZkMq9OCa4I7SzsNvw++4/jnPxnmILtflnvy5frJ1Pgc=",
+      "attention_v3.py": "xWyf2XnBmPKuOe5D6h8/3S3QwQXpbYFZTYCRh49Fkso="
+    }
+  },
+  "provenance": {
+    "kernel-builder": {
+      "version": "0.17.0-dev0",
+      "sha": "d0610aa58db33b142c86b59598a2a1c730f52996",
+      "dirty": false
+    },
+    "kernel": {
+      "sha": "490d9fe259aa1f38cdfd8c9646791796518182ee",
+      "dirty": true
+    }
+  }
 }

build/torch-cuda/my_softmax_function/__init__.py

@@ -1,26 +0,0 @@
-import ctypes
-import sys
-
-import importlib
-from pathlib import Path
-from types import ModuleType
-
-def _import_from_path(file_path: Path) -> ModuleType:
-    # We cannot use the module name as-is, after adding it to `sys.modules`,
-    # it would also be used for other imports. So, we make a module name that
-    # depends on the path for it to be unique using the hex-encoded hash of
-    # the path.
-    path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value)
-    module_name = path_hash
-    spec = importlib.util.spec_from_file_location(module_name, file_path)
-    if spec is None:
-        raise ImportError(f"Cannot load spec for {module_name} from {file_path}")
-    module = importlib.util.module_from_spec(spec)
-    if module is None:
-        raise ImportError(f"Cannot load module {module_name} from spec")
-    sys.modules[module_name] = module
-    spec.loader.exec_module(module)  # type: ignore
-    return module
-
-
-globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py")))

flake.lock

@@ -2,11 +2,11 @@
   "nodes": {
     "flake-compat": {
       "locked": {
-        "lastModified": 1765121682,
-        "narHash": "sha256-4VBOP18BFeiPkyhy9o4ssBNQEvfvv1kXkasAYd0+rrA=",
+        "lastModified": 1767039857,
+        "narHash": "sha256-vNpUSpF5Nuw8xvDLj2KCwwksIbjua2LZCqhV1LNRDns=",
         "owner": "edolstra",
         "repo": "flake-compat",
-        "rev": "65f23138d8d09a92e30f1e5c87611b23ef451bf3",
+        "rev": "5edf11c44bc78a0d334f6334cdaf7d60d732daab",
         "type": "github"
       },
       "original": {
@@ -37,35 +37,36 @@
       "inputs": {
         "flake-compat": "flake-compat",
         "flake-utils": "flake-utils",
-        "nixpkgs": "nixpkgs"
+        "nixpkgs": "nixpkgs",
+        "rust-overlay": "rust-overlay"
       },
       "locked": {
-        "lastModified": 1775482375,
-        "narHash": "sha256-RUbxfJGs96jwnwSci3+8h08GB2/katOI67yZTCaV+aE=",
+        "lastModified": 1785427612,
+        "narHash": "sha256-hUFEwQKDjVHw9v6xgC6Lx402TYfUA8zXEi8hDE/F4y0=",
         "owner": "huggingface",
-        "repo": "kernel-builder",
-        "rev": "dffbce5a048648febb96bbed79f3811ab6d7a577",
+        "repo": "kernels",
+        "rev": "d0610aa58db33b142c86b59598a2a1c730f52996",
         "type": "github"
       },
       "original": {
         "owner": "huggingface",
-        "repo": "kernel-builder",
+        "repo": "kernels",
         "type": "github"
       }
     },
     "nixpkgs": {
       "locked": {
-        "lastModified": 1766341660,
-        "narHash": "sha256-4yG6vx7Dddk9/zh45Y2KM82OaRD4jO3HA9r98ORzysA=",
+        "lastModified": 1783284758,
+        "narHash": "sha256-tiQ8/qi8I45OOaBBYlVbXoAVkeQzvvTQOv5I45rMw5o=",
         "owner": "NixOS",
         "repo": "nixpkgs",
-        "rev": "26861f5606e3e4d1400771b513cc63e5f70151a6",
+        "rev": "ec1a11210589d294f0ac99d3290a27e6c73dfa1d",
         "type": "github"
       },
       "original": {
         "owner": "NixOS",
-        "ref": "nixos-unstable-small",
         "repo": "nixpkgs",
+        "rev": "ec1a11210589d294f0ac99d3290a27e6c73dfa1d",
         "type": "github"
       }
     },
@@ -74,6 +75,27 @@
         "kernel-builder": "kernel-builder"
       }
     },
+    "rust-overlay": {
+      "inputs": {
+        "nixpkgs": [
+          "kernel-builder",
+          "nixpkgs"
+        ]
+      },
+      "locked": {
+        "lastModified": 1783320166,
+        "narHash": "sha256-l7C/OsjcnWDOk2K3ssj+SBduwL67LashjBqis9+t468=",
+        "owner": "oxalica",
+        "repo": "rust-overlay",
+        "rev": "20ee15370c9256669d66968b89ee20a4b0a4e673",
+        "type": "github"
+      },
+      "original": {
+        "owner": "oxalica",
+        "repo": "rust-overlay",
+        "type": "github"
+      }
+    },
     "systems": {
       "locked": {
         "lastModified": 1681028828,

flake.nix

@@ -2,7 +2,9 @@
   description = "my_softmax_function — fused softmax-attention kernel";
 
   inputs = {
-    kernel-builder.url = "github:huggingface/kernel-builder";
+    # New merged repo (kernel-builder folded into huggingface/kernels).
+    # Its metadata schema matches the current `kernels` loader library.
+    kernel-builder.url = "github:huggingface/kernels";
   };
 
   outputs =
@@ -10,7 +12,8 @@
       self,
       kernel-builder,
     }:
-    kernel-builder.lib.genFlakeOutputs {
+    # genFlakeOutputs was renamed to genKernelFlakeOutputs in the new repo.
+    kernel-builder.lib.genKernelFlakeOutputs {
       inherit self;
       path = ./.;
     };

torch-ext/my_softmax_function/_kernel.py

@@ -10,12 +10,11 @@ import math
 import torch
 from cutlass.cute.runtime import from_dlpack
 
+from ._ops import add_op_namespace_prefix
 from .attention_v3 import solve
 
-_OP_NAME = "my_softmax_function::softmax_attention"
 
-
-@torch.library.custom_op(_OP_NAME, mutates_args=())
+@torch.library.custom_op(add_op_namespace_prefix("softmax_attention"), mutates_args=())
 def softmax_attention(
     Q: torch.Tensor,        # (M, d)
     K: torch.Tensor,        # (N, d)
@@ -65,4 +64,5 @@ def attention(Q, K, V, scale=None):
     """Convenience entry point with a default 1/sqrt(d) scale."""
     if scale is None:
         scale = 1.0 / math.sqrt(Q.shape[-1])
-    return torch.ops.my_softmax_function.softmax_attention(Q, K, V, scale)
+    # Call the registered op directly (its namespace is build-unique).
+    return softmax_attention(Q, K, V, scale)