Add Builds

CMakeLists.txt

@@ -1,5 +1,5 @@
 cmake_minimum_required(VERSION 3.26)
-project(gptoss_kernels LANGUAGES CXX)
+project(gptoss_kernels LANGUAGES CXX C OBJC OBJCXX)
 
 set(CMAKE_OSX_DEPLOYMENT_TARGET "15.0" CACHE STRING "Minimum macOS deployment version")
 
@@ -89,8 +89,8 @@ endif()
 include(${CMAKE_CURRENT_LIST_DIR}/cmake/compile-metal.cmake)
 
 define_gpu_extension_target(
-  _gptoss_kernels_931bc1b_dirty
-  DESTINATION _gptoss_kernels_931bc1b_dirty
+  _gptoss_kernels_3f5b0c9
+  DESTINATION _gptoss_kernels_3f5b0c9
   LANGUAGE ${GPU_LANG}
   SOURCES ${SRC}
   COMPILE_FLAGS ${GPU_FLAGS}
@@ -100,5 +100,5 @@ define_gpu_extension_target(
 
 # Compile Metal shaders if any were found
 if(ALL_METAL_SOURCES)
-  compile_metal_shaders(_gptoss_kernels_931bc1b_dirty "${ALL_METAL_SOURCES}")
+  compile_metal_shaders(_gptoss_kernels_3f5b0c9 "${ALL_METAL_SOURCES}")
 endif()

build.toml

@@ -14,24 +14,42 @@ depends = ["torch"]
 backend = "metal"
 
 src = [
-  "gptoss_kernels/source/accumulate.metal",
-  "gptoss_kernels/source/expert_routing_metadata.metal",
-  "gptoss_kernels/source/metal.m",
-  "gptoss_kernels/source/scatter.metal",
-  "gptoss_kernels/source/topk.metal",
-  "gptoss_kernels/source/embeddings.metal",
-  "gptoss_kernels/source/metal-kernels.c",
-  "gptoss_kernels/source/random.metal",
-  "gptoss_kernels/source/sdpa.metal",
-  "gptoss_kernels/source/matmul.metal",
-  "gptoss_kernels/source/rmsnorm.metal",
-  "gptoss_kernels/source/sample.metal",
-  "gptoss_kernels/source/moematmul.metal",
-  "gptoss_kernels/source/convert.metal",
-  "gptoss_kernels/source/rope.metal",
-  "gptoss_kernels/source/gather_and_accumulate.metal",
-  "gptoss_kernels/source/tensor_wrappers.cpp",
-  "gptoss_kernels/source/log.c",
+  "gptoss_kernels/include/gpt-oss.h",
+"gptoss_kernels/include/gpt-oss/types.h",
+"gptoss_kernels/include/gpt-oss/macros.h",
+"gptoss_kernels/include/gpt-oss/functions.h",
+"gptoss_kernels/source/accumulate.metal",
+"gptoss_kernels/source/log.c",
+"gptoss_kernels/source/expert_routing_metadata.metal",
+"gptoss_kernels/source/metal.m",
+"gptoss_kernels/source/scatter.metal",
+"gptoss_kernels/source/topk.metal",
+"gptoss_kernels/source/embeddings.metal",
+"gptoss_kernels/source/metal-kernels.c",
+"gptoss_kernels/source/tensor_wrappers.cpp",
+"gptoss_kernels/source/random.metal",
+"gptoss_kernels/source/sdpa.metal",
+"gptoss_kernels/source/matmul.metal",
+"gptoss_kernels/source/rmsnorm.metal",
+"gptoss_kernels/source/sample.metal",
+"gptoss_kernels/source/moematmul.metal",
+"gptoss_kernels/source/convert.metal",
+"gptoss_kernels/source/rope.metal",
+"gptoss_kernels/source/gather_and_accumulate.metal",
+"gptoss_kernels/source/include/internal/uuid.h",
+"gptoss_kernels/source/include/internal/metal.hpp",
+"gptoss_kernels/source/include/internal/datatype.h",
+"gptoss_kernels/source/include/internal/rng.h",
+"gptoss_kernels/source/include/internal/rng.hpp",
+"gptoss_kernels/source/include/internal/log.h",
+"gptoss_kernels/source/include/internal/macros.h",
+"gptoss_kernels/source/include/internal/storage.h",
+"gptoss_kernels/source/include/internal/model.h",
+"gptoss_kernels/source/include/internal/math.h",
+"gptoss_kernels/source/include/internal/metal.h",
+"gptoss_kernels/source/include/internal/kernel-args.h",
+"gptoss_kernels/source/include/internal/datatype.hpp",
+"gptoss_kernels/source/include/internal/metal-kernels.h",
 ]
 
 include = ["gptoss_kernels/source/include", "gptoss_kernels/include", "."]

build/torch28-metal-aarch64-darwin/gptoss_kernels/__init__.py

@@ -0,0 +1,8 @@
+from ._ops import ops
+import torch
+
+def f32_bf16w_matmul(input: torch.Tensor, weight_bf16: torch.Tensor, bias_bf16: torch.Tensor, output: torch.Tensor, num_tokens: int, num_cols: int, num_rows: int, threadgroup_size: int) -> None:
+    ops.f32_bf16w_matmul_torch(input, weight_bf16, bias_bf16, output, num_tokens, num_cols, num_rows, threadgroup_size)
+    return output
+
+__all__ = ["f32_bf16w_matmul"]

build/torch28-metal-aarch64-darwin/gptoss_kernels/_gptoss_kernels_b5918a0_dirty.abi3.so

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

{torch-ext → build/torch28-metal-aarch64-darwin}/gptoss_kernels/_ops.py

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

build/torch28-metal-aarch64-darwin/gptoss_kernels/test.py

@@ -0,0 +1,6 @@
+import _gptoss_kernels_931bc1b_dirty 
+import torch
+
+print(dir(_gptoss_kernels_931bc1b_dirty))
+
+from gptoss_kernels import _gptoss_kernels_931bc1b_dirty

build/torch29-metal-aarch64-darwin/gptoss_kernels/__init__.py

@@ -0,0 +1,8 @@
+from ._ops import ops
+import torch
+
+def f32_bf16w_matmul(input: torch.Tensor, weight_bf16: torch.Tensor, bias_bf16: torch.Tensor, output: torch.Tensor, num_tokens: int, num_cols: int, num_rows: int, threadgroup_size: int) -> None:
+    ops.f32_bf16w_matmul_torch(input, weight_bf16, bias_bf16, output, num_tokens, num_cols, num_rows, threadgroup_size)
+    return output
+
+__all__ = ["f32_bf16w_matmul"]

torch-ext/gptoss_kernels/_gptoss_kernels_931bc1b_dirty.abi3.so → build/torch29-metal-aarch64-darwin/gptoss_kernels/_gptoss_kernels_b5918a0_dirty.abi3.so

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:31cddc1925c6c7901a5619ff55420ae6249d2c48de202a23a7c4534e4ccdcd4c
+oid sha256:257007386e3cbb53fa1bab940371899d8511fb1e429cfc476a9aea87ace35aab
 size 126536

build/torch29-metal-aarch64-darwin/gptoss_kernels/_ops.py

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

build/torch29-metal-aarch64-darwin/gptoss_kernels/test.py

@@ -0,0 +1,6 @@
+import _gptoss_kernels_931bc1b_dirty 
+import torch
+
+print(dir(_gptoss_kernels_931bc1b_dirty))
+
+from gptoss_kernels import _gptoss_kernels_931bc1b_dirty

cmake/compile-metal.cmake

@@ -4,12 +4,22 @@ function(compile_metal_shaders TARGET_NAME METAL_SOURCES)
     find_program(METAL_COMPILER xcrun REQUIRED)
     
     # Set Metal compiler flags
-    set(METAL_FLAGS "-std=metal3.0" "-O2")
+    set(METAL_FLAGS "-std=metal3.2" "-O2")
     
     # Output directory for compiled metallib
     set(METALLIB_OUTPUT_DIR "${CMAKE_BINARY_DIR}/metallib")
     file(MAKE_DIRECTORY ${METALLIB_OUTPUT_DIR})
     
+    set(METAL_INCLUDE_DIRS
+        "${CMAKE_SOURCE_DIR}/gptoss_kernels/source/include"
+        "${CMAKE_SOURCE_DIR}/gptoss_kernels/include"
+        "${CMAKE_SOURCE_DIR}/."
+    )
+
+    foreach(INC ${METAL_INCLUDE_DIRS})
+        list(APPEND METAL_FLAGS "-I${INC}")
+    endforeach()
+
     # Separate .metal files from .h files and compile .metal files to .air
     set(AIR_FILES)
     set(METAL_FILES)

flake.lock

@@ -98,11 +98,11 @@
         ]
       },
       "locked": {
-        "lastModified": 1761991868,
+        "lastModified": 1761998455,
         "narHash": "sha256-+csvkWC9jC4mwq1LNfK4O6m3Qg4dCCXjP5JGdPa3TEo=",
         "owner": "huggingface",
         "repo": "kernel-builder",
-        "rev": "79cbfcdfde82c8847551f67f4b951a410794a5c6",
+        "rev": "25ea190b1dca356f8bec2ec7cb92507b11ca62a1",
         "type": "github"
       },
       "original": {

flake.nix

@@ -2,7 +2,7 @@
   description = "Flake for Torch kernel extension";
 
   inputs = {
-    kernel-builder.url = "github:huggingface/kernel-builder?ref=metal_kernels";
+    kernel-builder.url = "github:huggingface/kernel-builder/metal_kernels";
   };
 
   outputs = { self, kernel-builder, }:

setup.py

@@ -108,7 +108,7 @@ setup(
     name="gptoss_kernels",
     # The version is just a stub, it's not used by the final build artefact.
     version="0.1.0",
-    ext_modules=[CMakeExtension("gptoss_kernels._gptoss_kernels_931bc1b_dirty")],
+    ext_modules=[CMakeExtension("gptoss_kernels._gptoss_kernels_3f5b0c9")],
     cmdclass={"build_ext": CMakeBuild},
     packages=find_packages(where="torch-ext", include=["gptoss_kernels*"]),
     package_dir={"": "torch-ext"},