update Builds

CMakeLists.txt

@@ -39,13 +39,19 @@ set(TORCH_gptoss_kernels_SRC
 
 list(APPEND SRC "${TORCH_gptoss_kernels_SRC}")
 set(gptoss_kernels_SRC
-  "gptoss_kernels/source/accumulate.metal"
+  "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/metal.mm"
 "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"
@@ -55,8 +61,20 @@ set(gptoss_kernels_SRC
 "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/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"
 )
 
 # Separate Metal shader files from other sources
@@ -85,12 +103,18 @@ list(APPEND SRC "${gptoss_kernels_CPP_SRC}")
 if(gptoss_kernels_METAL_SRC)
   list(APPEND ALL_METAL_SOURCES "${gptoss_kernels_METAL_SRC}")
 endif()
+
+# Keep the includes directory for the Metal sources
+if(gptoss_kernels_METAL_SRC)
+  list(APPEND METAL_INCLUDE_DIRS ${CMAKE_SOURCE_DIR}/gptoss_kernels/source/include;${CMAKE_SOURCE_DIR}/gptoss_kernels/include;${CMAKE_SOURCE_DIR}/.)
+endif()
+
 # Include Metal shader compilation utilities
 include(${CMAKE_CURRENT_LIST_DIR}/cmake/compile-metal.cmake)
 
 define_gpu_extension_target(
-  _gptoss_kernels_3f5b0c9
-  DESTINATION _gptoss_kernels_3f5b0c9
+  _gptoss_kernels_3a886f8_dirty
+  DESTINATION _gptoss_kernels_3a886f8_dirty
   LANGUAGE ${GPU_LANG}
   SOURCES ${SRC}
   COMPILE_FLAGS ${GPU_FLAGS}
@@ -100,5 +124,5 @@ define_gpu_extension_target(
 
 # Compile Metal shaders if any were found
 if(ALL_METAL_SOURCES)
-  compile_metal_shaders(_gptoss_kernels_3f5b0c9 "${ALL_METAL_SOURCES}")
+  compile_metal_shaders(_gptoss_kernels_3a886f8_dirty "${ALL_METAL_SOURCES}" "${METAL_INCLUDE_DIRS}")
 endif()

build/torch28-metal-aarch64-darwin/gptoss_kernels/{_gptoss_kernels_6df4712_dirty.abi3.so → _gptoss_kernels_5341d17_dirty.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:4bdba37d899fba53e0f7134eb2461f69eaa80728bd00305064f8eb6594665d32
+oid sha256:fa19b7a232893fc5ac4ef189ae0973e3e672efac424580f68fd2873cb2a7fbc8
 size 291032

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

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

build/torch29-metal-aarch64-darwin/gptoss_kernels/{_gptoss_kernels_6df4712_dirty.abi3.so → _gptoss_kernels_5341d17_dirty.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:7dc7a799252e70b808880ffd8c685887b55c9d6a1354424152645aff6dbda64b
+oid sha256:679729e810bc2a360f49eed34299e2d63f0eb24489d0f6a032ee12175e7831a3
 size 292040

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

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

cmake/compile-metal.cmake

@@ -1,5 +1,5 @@
 # Metal shader compilation function
-function(compile_metal_shaders TARGET_NAME METAL_SOURCES)
+function(compile_metal_shaders TARGET_NAME METAL_SOURCES EXTRA_INCLUDE_DIRS)
     # Find the Metal compiler
     find_program(METAL_COMPILER xcrun REQUIRED)
     
@@ -9,14 +9,8 @@ function(compile_metal_shaders TARGET_NAME METAL_SOURCES)
     # 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})
+    foreach(INC ${EXTRA_INCLUDE_DIRS})
         list(APPEND METAL_FLAGS "-I${INC}")
     endforeach()
 

flake.lock

@@ -98,11 +98,11 @@
         ]
       },
       "locked": {
-        "lastModified": 1761998455,
-        "narHash": "sha256-+csvkWC9jC4mwq1LNfK4O6m3Qg4dCCXjP5JGdPa3TEo=",
+        "lastModified": 1762017065,
+        "narHash": "sha256-1xkz9M9IozOoSE6+SE3l8Jv/5U57UaQg6czjW7AXj9Y=",
         "owner": "huggingface",
         "repo": "kernel-builder",
-        "rev": "25ea190b1dca356f8bec2ec7cb92507b11ca62a1",
+        "rev": "64417e893b6aae48edf47a788a5560116708de08",
         "type": "github"
       },
       "original": {

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_3f5b0c9")],
+    ext_modules=[CMakeExtension("gptoss_kernels._gptoss_kernels_3a886f8_dirty")],
     cmdclass={"build_ext": CMakeBuild},
     packages=find_packages(where="torch-ext", include=["gptoss_kernels*"]),
     package_dir={"": "torch-ext"},

torch-ext/gptoss_kernels/_ops.py

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