Clean

.gitignore

@@ -0,0 +1,2 @@
+.venv/
+**/_pycache__/

CMakeLists.txt

@@ -1,128 +0,0 @@
-cmake_minimum_required(VERSION 3.26)
-project(gptoss_kernels LANGUAGES CXX C OBJC OBJCXX)
-
-set(CMAKE_OSX_DEPLOYMENT_TARGET "15.0" CACHE STRING "Minimum macOS deployment version")
-
-install(CODE "set(CMAKE_INSTALL_LOCAL_ONLY TRUE)" ALL_COMPONENTS)
-
-include(FetchContent)
-file(MAKE_DIRECTORY ${FETCHCONTENT_BASE_DIR}) # Ensure the directory exists
-message(STATUS "FetchContent base directory: ${FETCHCONTENT_BASE_DIR}")
-
-include(${CMAKE_CURRENT_LIST_DIR}/cmake/utils.cmake)
-
-if(DEFINED Python3_EXECUTABLE)
-  # Allow passing through the interpreter (e.g. from setup.py).
-  find_package(Python3 COMPONENTS Development Development.SABIModule Interpreter)
-  if (NOT Python3_FOUND)
-    message(FATAL_ERROR "Unable to find python matching: ${EXECUTABLE}.")
-  endif()
-else()
-  find_package(Python3 REQUIRED COMPONENTS Development Development.SABIModule Interpreter)
-endif()
-
-append_cmake_prefix_path("torch" "torch.utils.cmake_prefix_path")
-
-find_package(Torch REQUIRED)
-
-add_compile_definitions(METAL_KERNEL)
-
-# Initialize list for Metal shader sources
-set(ALL_METAL_SOURCES)
-#get_torch_gpu_compiler_flags(TORCH_GPU_FLAGS ${GPU_LANG})
-#list(APPEND GPU_FLAGS ${TORCH_GPU_FLAGS})
-
-set(TORCH_gptoss_kernels_SRC
-  torch-ext/torch_binding.cpp torch-ext/torch_binding.h
-)
-
-
-list(APPEND SRC "${TORCH_gptoss_kernels_SRC}")
-set(gptoss_kernels_SRC
-  "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.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"
-"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"
-)
-
-# Separate Metal shader files from other sources
-set(gptoss_kernels_METAL_SRC)
-set(gptoss_kernels_CPP_SRC)
-
-foreach(src_file IN LISTS gptoss_kernels_SRC)
-  if(src_file MATCHES "\\.(metal|h)$")
-    list(APPEND gptoss_kernels_METAL_SRC ${src_file})
-  else()
-    list(APPEND gptoss_kernels_CPP_SRC ${src_file})
-  endif()
-endforeach()
-
-# TODO: check if CLion support this:
-# https://youtrack.jetbrains.com/issue/CPP-16510/CLion-does-not-handle-per-file-include-directories
-set_source_files_properties(
-  ${gptoss_kernels_CPP_SRC}
-  PROPERTIES INCLUDE_DIRECTORIES "${CMAKE_SOURCE_DIR}/gptoss_kernels/source/include;${CMAKE_SOURCE_DIR}/gptoss_kernels/include;${CMAKE_SOURCE_DIR}/.")
-
-
-# Add C++ sources to main source list
-list(APPEND SRC "${gptoss_kernels_CPP_SRC}")
-
-# Keep track of Metal sources for later compilation
-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_3a886f8_dirty
-  DESTINATION _gptoss_kernels_3a886f8_dirty
-  LANGUAGE ${GPU_LANG}
-  SOURCES ${SRC}
-  COMPILE_FLAGS ${GPU_FLAGS}
-  ARCHITECTURES ${GPU_ARCHES}
-  USE_SABI 3
-  WITH_SOABI)
-
-# Compile Metal shaders if any were found
-if(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_9964bae_dirty.abi3.so → _gptoss_kernels_9ffd725_dirty.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:b52d3924ac74e614664fd9ec72e9673807ed170e57277b81c1922c0b54a88a6a
+oid sha256:8e9880a41249b5e1859134140111e0515925129b0c70d4587a70324b489bb9e3
 size 391752

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

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

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

@@ -1,15 +0,0 @@
-import _gptoss_kernels_931bc1b_dirty 
-import torch
-
-print(dir(_gptoss_kernels_931bc1b_dirty))
-
-from gptoss_kernels import _gptoss_kernels_931bc1b_dirty
-
-print(dir(f32_bf16w_matmul))
-
-input = torch.randn(10, 10)
-weight_bf16 = torch.randn(10, 10)
-bias_bf16 = torch.randn(10)
-output = torch.randn(10, 10)
-f32_bf16w_matmul(input, weight_bf16, bias_bf16, output, 10, 10, 10, 10)
-print(output)

build/torch29-metal-aarch64-darwin/gptoss_kernels/{_gptoss_kernels_9964bae_dirty.abi3.so → _gptoss_kernels_9ffd725_dirty.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:dc170dbf45587f9a1091e9b6c92ab02ebe4dc3cdd13be8e56a9a8d3a353d8c86
+oid sha256:fd8bc3d9d6d40953740d4566aeb3543b1025a647f8df05a9d296299bafbd8c31
 size 392840

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

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

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

@@ -1,15 +0,0 @@
-import _gptoss_kernels_931bc1b_dirty 
-import torch
-
-print(dir(_gptoss_kernels_931bc1b_dirty))
-
-from gptoss_kernels import _gptoss_kernels_931bc1b_dirty
-
-print(dir(f32_bf16w_matmul))
-
-input = torch.randn(10, 10)
-weight_bf16 = torch.randn(10, 10)
-bias_bf16 = torch.randn(10)
-output = torch.randn(10, 10)
-f32_bf16w_matmul(input, weight_bf16, bias_bf16, output, 10, 10, 10, 10)
-print(output)

cmake/compile-metal.cmake

@@ -1,90 +0,0 @@
-# Metal shader compilation function
-function(compile_metal_shaders TARGET_NAME METAL_SOURCES EXTRA_INCLUDE_DIRS)
-    # Find the Metal compiler
-    find_program(METAL_COMPILER xcrun REQUIRED)
-    
-    # Set Metal compiler flags
-    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})
-
-    foreach(INC ${EXTRA_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)
-    set(HEADER_FILES)
-    
-    foreach(SOURCE_FILE ${METAL_SOURCES})
-        if(SOURCE_FILE MATCHES "\\.metal$")
-            list(APPEND METAL_FILES ${SOURCE_FILE})
-        elseif(SOURCE_FILE MATCHES "\\.h$")
-            list(APPEND HEADER_FILES ${SOURCE_FILE})
-        endif()
-    endforeach()
-    
-    foreach(METAL_FILE ${METAL_FILES})
-        get_filename_component(METAL_NAME ${METAL_FILE} NAME_WE)
-        set(AIR_FILE "${CMAKE_BINARY_DIR}/${METAL_NAME}.air")
-        
-        # Include header files as dependencies
-        set(ALL_DEPENDENCIES ${CMAKE_CURRENT_SOURCE_DIR}/${METAL_FILE})
-        foreach(HEADER_FILE ${HEADER_FILES})
-            list(APPEND ALL_DEPENDENCIES ${CMAKE_CURRENT_SOURCE_DIR}/${HEADER_FILE})
-        endforeach()
-        
-        add_custom_command(
-            OUTPUT ${AIR_FILE}
-            COMMAND ${METAL_COMPILER} -sdk macosx metal ${METAL_FLAGS}
-                    -c ${CMAKE_CURRENT_SOURCE_DIR}/${METAL_FILE}
-                    -o ${AIR_FILE}
-            DEPENDS ${ALL_DEPENDENCIES}
-            COMMENT "Compiling Metal shader ${METAL_FILE} to ${AIR_FILE}"
-            VERBATIM
-        )
-        
-        list(APPEND AIR_FILES ${AIR_FILE})
-    endforeach()
-    
-    # Link all .air files into a single .metallib
-    set(METALLIB_FILE "${METALLIB_OUTPUT_DIR}/${TARGET_NAME}.metallib")
-    add_custom_command(
-        OUTPUT ${METALLIB_FILE}
-        COMMAND ${METAL_COMPILER} -sdk macosx metallib ${AIR_FILES}
-                -o ${METALLIB_FILE}
-        DEPENDS ${AIR_FILES}
-        COMMENT "Linking Metal library ${METALLIB_FILE}"
-        VERBATIM
-    )
-    
-    # Generate C++ header with embedded metallib data
-    set(METALLIB_HEADER "${CMAKE_BINARY_DIR}/${TARGET_NAME}_metallib.h")
-    set(METALLIB_TO_HEADER_SCRIPT "${CMAKE_CURRENT_SOURCE_DIR}/cmake/metallib_to_header.py")
-    
-    add_custom_command(
-        OUTPUT ${METALLIB_HEADER}
-        COMMAND ${Python_EXECUTABLE} ${METALLIB_TO_HEADER_SCRIPT} ${METALLIB_FILE} ${METALLIB_HEADER} ${TARGET_NAME}
-        DEPENDS ${METALLIB_FILE} ${METALLIB_TO_HEADER_SCRIPT}
-        COMMENT "Generating embedded Metal library header ${METALLIB_HEADER}"
-        VERBATIM
-    )
-    
-    # Create a custom target for the metallib
-    add_custom_target(${TARGET_NAME}_metallib ALL DEPENDS ${METALLIB_FILE} ${METALLIB_HEADER})
-    
-    # Add dependency to main target
-    add_dependencies(${TARGET_NAME} ${TARGET_NAME}_metallib)
-    
-    # Add the generated header to include directories
-    target_include_directories(${TARGET_NAME} PRIVATE ${CMAKE_BINARY_DIR})
-    
-    # Pass the metallib header and namespace as compile definitions
-    target_compile_definitions(${TARGET_NAME} PRIVATE 
-        EMBEDDED_METALLIB_HEADER="${TARGET_NAME}_metallib.h"
-        EMBEDDED_METALLIB_NAMESPACE=${TARGET_NAME}_metal
-    )
-endfunction()

cmake/metallib_to_header.py

@@ -1,73 +0,0 @@
-#!/usr/bin/env python3
-import sys
-import os
-
-def convert_metallib_to_header(metallib_path: str, header_path: str, target_name: str) -> None:
-    """Convert a metallib binary file to a C++ header with embedded data."""
-    
-    # Read the metallib binary data
-    with open(metallib_path, 'rb') as f:
-        data: bytes = f.read()
-    
-    # Generate the header content
-    header_content: str = """// Auto-generated file containing embedded Metal library
-#pragma once
-#include <cstddef>
-#include <Metal/Metal.h>
-
-namespace """ + target_name + """_metal {
-    static const unsigned char metallib_data[] = {
-"""
-    
-    # Convert binary data to C array format
-    bytes_per_line: int = 16
-    for i in range(0, len(data), bytes_per_line):
-        chunk: bytes = data[i:i + bytes_per_line]
-        hex_values: str = ', '.join('0x{:02x}'.format(b) for b in chunk)
-        header_content += "        " + hex_values + ","
-        if i + bytes_per_line < len(data):
-            header_content += "\n"
-    
-    header_content += """
-    };
-    static const size_t metallib_data_len = """ + str(len(data)) + """;
-    
-    // Convenience function to create Metal library from embedded data
-    inline id<MTLLibrary> createLibrary(id<MTLDevice> device, NSError** error = nullptr) {
-        dispatch_data_t libraryData = dispatch_data_create(
-            metallib_data,
-            metallib_data_len,
-            dispatch_get_main_queue(),
-            ^{ /* No cleanup needed for static data */ });
-        
-        NSError* localError = nil;
-        id<MTLLibrary> library = [device newLibraryWithData:libraryData error:&localError];
-        
-        if (error) {
-            *error = localError;
-        }
-        
-        return library;
-    }
-} // namespace """ + target_name + """_metal
-"""
-    
-    # Write the header file
-    dir_path: str = os.path.dirname(header_path)
-    if dir_path:
-        os.makedirs(dir_path, exist_ok=True)
-    with open(header_path, 'w') as f:
-        f.write(header_content)
-    
-    print("Generated {} ({} bytes)".format(header_path, len(data)))
-
-if __name__ == "__main__":
-    if len(sys.argv) != 4:
-        print("Usage: metallib_to_header.py <metallib_path> <header_path> <target_name>")
-        sys.exit(1)
-    
-    metallib_path: str = sys.argv[1]
-    header_path: str = sys.argv[2]
-    target_name: str = sys.argv[3]
-    
-    convert_metallib_to_header(metallib_path, header_path, target_name)

cmake/utils.cmake

@@ -1,557 +0,0 @@
-# Vendored from vLLM:
-#
-# https://github.com/vllm-project/vllm/blob/main/cmake/utils.cmake
-#
-# Attempt to find the python package that uses the same python executable as
-# `EXECUTABLE` and is one of the `SUPPORTED_VERSIONS`.
-#
-macro (find_python_from_executable EXECUTABLE SUPPORTED_VERSIONS)
-  file(REAL_PATH ${EXECUTABLE} EXECUTABLE)
-  set(Python3_EXECUTABLE ${EXECUTABLE})
-  find_package(Python3 COMPONENTS Interpreter Development.Module Development.SABIModule)
-  if (NOT Python3_FOUND)
-    message(FATAL_ERROR "Unable to find python matching: ${EXECUTABLE}.")
-  endif()
-  set(_VER "${Python3_VERSION_MAJOR}.${Python3_VERSION_MINOR}")
-  set(_SUPPORTED_VERSIONS_LIST ${SUPPORTED_VERSIONS} ${ARGN})
-  if (NOT _VER IN_LIST _SUPPORTED_VERSIONS_LIST)
-    message(FATAL_ERROR
-      "Python version (${_VER}) is not one of the supported versions: "
-      "${_SUPPORTED_VERSIONS_LIST}.")
-  endif()
-  message(STATUS "Found python matching: ${EXECUTABLE}.")
-endmacro()
-
-#
-# Run `EXPR` in python.  The standard output of python is stored in `OUT` and
-# has trailing whitespace stripped.  If an error is encountered when running
-# python, a fatal message `ERR_MSG` is issued.
-#
-function (run_python OUT EXPR ERR_MSG)
-  execute_process(
-    COMMAND
-    "${Python3_EXECUTABLE}" "-c" "${EXPR}"
-    OUTPUT_VARIABLE PYTHON_OUT
-    RESULT_VARIABLE PYTHON_ERROR_CODE
-    ERROR_VARIABLE PYTHON_STDERR
-    OUTPUT_STRIP_TRAILING_WHITESPACE)
-
-  if(NOT PYTHON_ERROR_CODE EQUAL 0)
-    message(FATAL_ERROR "${ERR_MSG}: ${PYTHON_STDERR}")
-  endif()
-  set(${OUT} ${PYTHON_OUT} PARENT_SCOPE)
-endfunction()
-
-# Run `EXPR` in python after importing `PKG`. Use the result of this to extend
-# `CMAKE_PREFIX_PATH` so the torch cmake configuration can be imported.
-macro (append_cmake_prefix_path PKG EXPR)
-  run_python(_PREFIX_PATH
-    "import ${PKG}; print(${EXPR})" "Failed to locate ${PKG} path")
-  list(APPEND CMAKE_PREFIX_PATH ${_PREFIX_PATH})
-endmacro()
-
-#
-# Add a target named `hipify${NAME}` that runs the hipify preprocessor on a set
-# of CUDA source files. The names of the corresponding "hipified" sources are
-# stored in `OUT_SRCS`.
-#
-function (hipify_sources_target OUT_SRCS NAME ORIG_SRCS)
-  #
-  # Split into C++ and non-C++ (i.e. CUDA) sources.
-  #
-  set(NODUP_SRCS ${ORIG_SRCS})
-  list(REMOVE_DUPLICATES NODUP_SRCS)
-  set(SRCS ${NODUP_SRCS})
-  set(CXX_SRCS ${NODUP_SRCS})
-  list(FILTER SRCS INCLUDE REGEX "\.cu$")
-  list(FILTER CXX_SRCS EXCLUDE REGEX "\.cu$")
-
-  #
-  # Generate ROCm/HIP source file names from CUDA file names.
-  # Since HIP files are generated code, they will appear in the build area
-  # `CMAKE_CURRENT_BINARY_DIR` directory rather than the original csrc dir.
-  #
-  set(HIP_SRCS)
-  foreach (SRC ${SRCS})
-    get_source_file_property(include_dirs "${SRC}" INCLUDE_DIRECTORIES)
-    get_source_file_property(compile_options "${SRC}" COMPILE_OPTIONS)
-    string(REGEX REPLACE "\.cu$" "\.hip" SRC ${SRC})
-    string(REGEX REPLACE "cuda" "hip" SRC ${SRC})
-
-    if(include_dirs)
-      # Copy over include directories from the original CUDA file.
-      set_source_files_properties(
-        ${SRC}
-        PROPERTIES INCLUDE_DIRECTORIES "${include_dirs}")
-    endif()
-
-    if(compile_options)
-      set_source_files_properties(
-        ${SRC}
-        PROPERTIES COMPILE_OPTIONS "${compile_options}")
-    endif()
-
-    list(APPEND HIP_SRCS "${CMAKE_CURRENT_BINARY_DIR}/${SRC}")
-  endforeach()
-
-  add_custom_target(
-    hipify${NAME}
-    COMMAND "${Python3_EXECUTABLE}" ${CMAKE_SOURCE_DIR}/cmake/hipify.py -p ${CMAKE_SOURCE_DIR} -o ${CMAKE_CURRENT_BINARY_DIR} ${SRCS}
-    DEPENDS ${CMAKE_SOURCE_DIR}/cmake/hipify.py ${SRCS}
-    BYPRODUCTS ${HIP_SRCS}
-    COMMENT "Running hipify on ${NAME} extension source files.")
-
-  # Swap out original extension sources with hipified sources.
-  list(APPEND HIP_SRCS ${CXX_SRCS})
-  set(${OUT_SRCS} ${HIP_SRCS} PARENT_SCOPE)
-endfunction()
-
-#
-# Get additional GPU compiler flags from torch.
-#
-function (get_torch_gpu_compiler_flags OUT_GPU_FLAGS GPU_LANG)
-  if (${GPU_LANG} STREQUAL "CUDA")
-    #
-    # Get common NVCC flags from torch.
-    #
-    run_python(GPU_FLAGS
-      "from torch.utils.cpp_extension import COMMON_NVCC_FLAGS; print(';'.join(COMMON_NVCC_FLAGS))"
-      "Failed to determine torch nvcc compiler flags")
-
-    if (CUDA_VERSION VERSION_GREATER_EQUAL 11.8)
-      list(APPEND GPU_FLAGS "-DENABLE_FP8")
-      list(REMOVE_ITEM GPU_FLAGS
-        "-D__CUDA_NO_HALF_OPERATORS__"
-        "-D__CUDA_NO_HALF_CONVERSIONS__"
-        "-D__CUDA_NO_BFLOAT16_CONVERSIONS__"
-        "-D__CUDA_NO_HALF2_OPERATORS__")
-    endif()
-
-  elseif(${GPU_LANG} STREQUAL "HIP")
-    #
-    # Get common HIP/HIPCC flags from torch.
-    #
-    run_python(GPU_FLAGS
-      "import torch.utils.cpp_extension as t; print(';'.join(t.COMMON_HIP_FLAGS + t.COMMON_HIPCC_FLAGS))"
-      "Failed to determine torch nvcc compiler flags")
-
-    list(APPEND GPU_FLAGS
-      "-DUSE_ROCM"
-      "-DENABLE_FP8"
-      "-U__HIP_NO_HALF_CONVERSIONS__"
-      "-U__HIP_NO_HALF_OPERATORS__"
-      "-fno-gpu-rdc")
-
-  endif()
-  set(${OUT_GPU_FLAGS} ${GPU_FLAGS} PARENT_SCOPE)
-endfunction()
-
-# Macro for converting a `gencode` version number to a cmake version number.
-macro(string_to_ver OUT_VER IN_STR)
-  string(REGEX REPLACE "\([0-9]+\)\([0-9]\)" "\\1.\\2" ${OUT_VER} ${IN_STR})
-endmacro()
-
-#
-# Clear all `-gencode` flags from `CMAKE_CUDA_FLAGS` and store them in
-# `CUDA_ARCH_FLAGS`.
-#
-# Example:
-#   CMAKE_CUDA_FLAGS="-Wall -gencode arch=compute_70,code=sm_70 -gencode arch=compute_75,code=sm_75"
-#   clear_cuda_arches(CUDA_ARCH_FLAGS)
-#   CUDA_ARCH_FLAGS="-gencode arch=compute_70,code=sm_70;-gencode arch=compute_75,code=sm_75"
-#   CMAKE_CUDA_FLAGS="-Wall"
-#
-macro(clear_cuda_arches CUDA_ARCH_FLAGS)
-    # Extract all `-gencode` flags from `CMAKE_CUDA_FLAGS`
-    string(REGEX MATCHALL "-gencode arch=[^ ]+" CUDA_ARCH_FLAGS
-      ${CMAKE_CUDA_FLAGS})
-
-    # Remove all `-gencode` flags from `CMAKE_CUDA_FLAGS` since they will be modified
-    # and passed back via the `CUDA_ARCHITECTURES` property.
-    string(REGEX REPLACE "-gencode arch=[^ ]+ *" "" CMAKE_CUDA_FLAGS
-      ${CMAKE_CUDA_FLAGS})
-endmacro()
-
-#
-# Extract unique CUDA architectures from a list of compute capabilities codes in 
-# the form `<major><minor>[<letter>]`, convert them to the form sort 
-# `<major>.<minor>`, dedupes them and then sorts them in ascending order and 
-# stores them in `OUT_ARCHES`.
-#
-# Example:
-#   CUDA_ARCH_FLAGS="-gencode arch=compute_75,code=sm_75;...;-gencode arch=compute_90a,code=sm_90a" 
-#   extract_unique_cuda_archs_ascending(OUT_ARCHES CUDA_ARCH_FLAGS)
-#   OUT_ARCHES="7.5;...;9.0"
-function(extract_unique_cuda_archs_ascending OUT_ARCHES CUDA_ARCH_FLAGS)
-  set(_CUDA_ARCHES)
-  foreach(_ARCH ${CUDA_ARCH_FLAGS})
-    string(REGEX MATCH "arch=compute_\([0-9]+a?\)" _COMPUTE ${_ARCH})
-    if (_COMPUTE)
-      set(_COMPUTE ${CMAKE_MATCH_1})
-    endif()
-
-    string_to_ver(_COMPUTE_VER ${_COMPUTE})
-    list(APPEND _CUDA_ARCHES ${_COMPUTE_VER})
-  endforeach()
-
-  list(REMOVE_DUPLICATES _CUDA_ARCHES)
-  list(SORT _CUDA_ARCHES COMPARE NATURAL ORDER ASCENDING)
-  set(${OUT_ARCHES} ${_CUDA_ARCHES} PARENT_SCOPE)
-endfunction()
-
-#
-# For a specific file set the `-gencode` flag in compile options conditionally 
-# for the CUDA language. 
-#
-# Example:
-#   set_gencode_flag_for_srcs(
-#     SRCS "foo.cu"
-#     ARCH "compute_75"
-#     CODE "sm_75")
-#   adds: "-gencode arch=compute_75,code=sm_75" to the compile options for 
-#    `foo.cu` (only for the CUDA language).
-#
-macro(set_gencode_flag_for_srcs)
-  set(options)
-  set(oneValueArgs ARCH CODE)
-  set(multiValueArgs SRCS)
-  cmake_parse_arguments(arg "${options}" "${oneValueArgs}"
-                        "${multiValueArgs}" ${ARGN} )
-  set(_FLAG -gencode arch=${arg_ARCH},code=${arg_CODE})
-  set_property(
-    SOURCE ${arg_SRCS}
-    APPEND PROPERTY
-    COMPILE_OPTIONS "$<$<COMPILE_LANGUAGE:CUDA>:${_FLAG}>"
-  )
-
-  message(DEBUG "Setting gencode flag for ${arg_SRCS}: ${_FLAG}")
-endmacro(set_gencode_flag_for_srcs)
-
-#
-# For a list of source files set the `-gencode` flags in the files specific 
-#  compile options (specifically for the CUDA language).
-#
-# arguments are:
-#  SRCS: list of source files
-#  CUDA_ARCHS: list of CUDA architectures in the form `<major>.<minor>[letter]`
-#  BUILD_PTX_FOR_ARCH: if set to true, then the PTX code will be built
-#    for architecture `BUILD_PTX_FOR_ARCH` if there is a CUDA_ARCH in CUDA_ARCHS 
-#    that is larger than BUILD_PTX_FOR_ARCH.
-#
-macro(set_gencode_flags_for_srcs)
-  set(options)
-  set(oneValueArgs BUILD_PTX_FOR_ARCH)
-  set(multiValueArgs SRCS CUDA_ARCHS)
-  cmake_parse_arguments(arg "${options}" "${oneValueArgs}"
-                        "${multiValueArgs}" ${ARGN} )
-
-  foreach(_ARCH ${arg_CUDA_ARCHS})
-    # handle +PTX suffix: generate both sm and ptx codes if requested
-    string(FIND "${_ARCH}" "+PTX" _HAS_PTX)
-    if(NOT _HAS_PTX EQUAL -1)
-      string(REPLACE "+PTX" "" _BASE_ARCH "${_ARCH}")
-      string(REPLACE "." "" _STRIPPED_ARCH "${_BASE_ARCH}")
-      set_gencode_flag_for_srcs(
-        SRCS ${arg_SRCS}
-        ARCH "compute_${_STRIPPED_ARCH}"
-        CODE "sm_${_STRIPPED_ARCH}")
-      set_gencode_flag_for_srcs(
-        SRCS ${arg_SRCS}
-        ARCH "compute_${_STRIPPED_ARCH}"
-        CODE "compute_${_STRIPPED_ARCH}")
-    else()
-      string(REPLACE "." "" _STRIPPED_ARCH "${_ARCH}")
-      set_gencode_flag_for_srcs(
-        SRCS ${arg_SRCS}
-        ARCH "compute_${_STRIPPED_ARCH}"
-        CODE "sm_${_STRIPPED_ARCH}")
-    endif()
-  endforeach()
-
-  if (${arg_BUILD_PTX_FOR_ARCH})
-    list(SORT arg_CUDA_ARCHS COMPARE NATURAL ORDER ASCENDING)
-    list(GET arg_CUDA_ARCHS -1 _HIGHEST_ARCH)
-    if (_HIGHEST_ARCH VERSION_GREATER_EQUAL ${arg_BUILD_PTX_FOR_ARCH})
-      string(REPLACE "." "" _PTX_ARCH "${arg_BUILD_PTX_FOR_ARCH}")
-      set_gencode_flag_for_srcs(
-        SRCS ${arg_SRCS}
-        ARCH "compute_${_PTX_ARCH}"
-        CODE "compute_${_PTX_ARCH}")
-    endif()
-  endif()
-endmacro()
-
-#
-# For the given `SRC_CUDA_ARCHS` list of gencode versions in the form
-#  `<major>.<minor>[letter]` compute the "loose intersection" with the
-#  `TGT_CUDA_ARCHS` list of gencodes. We also support the `+PTX` suffix in
-#  `SRC_CUDA_ARCHS` which indicates that the PTX code should be built when there
-#  is a CUDA_ARCH in `TGT_CUDA_ARCHS` that is equal to or larger than the
-#  architecture in `SRC_CUDA_ARCHS`.
-# The loose intersection is defined as:
-#   { max{ x \in tgt | x <= y } | y \in src, { x \in tgt | x <= y } != {} }
-#  where `<=` is the version comparison operator.
-# In other words, for each version in `TGT_CUDA_ARCHS` find the highest version
-#  in `SRC_CUDA_ARCHS` that is less or equal to the version in `TGT_CUDA_ARCHS`.
-# We have special handling for x.0a, if x.0a is in `SRC_CUDA_ARCHS` and x.0 is
-#  in `TGT_CUDA_ARCHS` then we should remove x.0a from `SRC_CUDA_ARCHS` and add
-#  x.0a to the result (and remove x.0 from TGT_CUDA_ARCHS).
-# The result is stored in `OUT_CUDA_ARCHS`.
-#
-# Example:
-#   SRC_CUDA_ARCHS="7.5;8.0;8.6;9.0;9.0a"
-#   TGT_CUDA_ARCHS="8.0;8.9;9.0"
-#   cuda_archs_loose_intersection(OUT_CUDA_ARCHS SRC_CUDA_ARCHS TGT_CUDA_ARCHS)
-#   OUT_CUDA_ARCHS="8.0;8.6;9.0;9.0a"
-#
-# Example With PTX:
-#   SRC_CUDA_ARCHS="8.0+PTX"
-#   TGT_CUDA_ARCHS="9.0"
-#   cuda_archs_loose_intersection(OUT_CUDA_ARCHS SRC_CUDA_ARCHS TGT_CUDA_ARCHS)
-#   OUT_CUDA_ARCHS="8.0+PTX"
-#
-function(cuda_archs_loose_intersection OUT_CUDA_ARCHS SRC_CUDA_ARCHS TGT_CUDA_ARCHS)
-  set(_SRC_CUDA_ARCHS "${SRC_CUDA_ARCHS}")
-  set(_TGT_CUDA_ARCHS ${TGT_CUDA_ARCHS})
-
-  # handle +PTX suffix: separate base arch for matching, record PTX requests
-  set(_PTX_ARCHS)
-  foreach(_arch ${_SRC_CUDA_ARCHS})
-    if(_arch MATCHES "\\+PTX$")
-      string(REPLACE "+PTX" "" _base "${_arch}")
-      list(APPEND _PTX_ARCHS "${_base}")
-      list(REMOVE_ITEM _SRC_CUDA_ARCHS "${_arch}")
-      list(APPEND _SRC_CUDA_ARCHS "${_base}")
-    endif()
-  endforeach()
-  list(REMOVE_DUPLICATES _PTX_ARCHS)
-  list(REMOVE_DUPLICATES _SRC_CUDA_ARCHS)
-
-  # if x.0a is in SRC_CUDA_ARCHS and x.0 is in CUDA_ARCHS then we should
-  # remove x.0a from SRC_CUDA_ARCHS and add x.0a to _CUDA_ARCHS
-  set(_CUDA_ARCHS)
-  foreach(_arch ${_SRC_CUDA_ARCHS})
-    if(_arch MATCHES "\\a$")
-      list(REMOVE_ITEM _SRC_CUDA_ARCHS "${_arch}")
-      string(REPLACE "a" "" _base "${_arch}")
-      if ("${_base}" IN_LIST TGT_CUDA_ARCHS)
-        list(REMOVE_ITEM _TGT_CUDA_ARCHS "${_base}")
-        list(APPEND _CUDA_ARCHS "${_arch}")
-      endif()
-    endif()
-  endforeach()
-
-  list(SORT _SRC_CUDA_ARCHS COMPARE NATURAL ORDER ASCENDING)
-
-  # for each ARCH in TGT_CUDA_ARCHS find the highest arch in SRC_CUDA_ARCHS that
-  # is less or equal to ARCH (but has the same major version since SASS binary
-  # compatibility is only forward compatible within the same major version).
-  foreach(_ARCH ${_TGT_CUDA_ARCHS})
-    set(_TMP_ARCH)
-    # Extract the major version of the target arch
-    string(REGEX REPLACE "^([0-9]+)\\..*$" "\\1" TGT_ARCH_MAJOR "${_ARCH}")
-    foreach(_SRC_ARCH ${_SRC_CUDA_ARCHS})
-      # Extract the major version of the source arch
-      string(REGEX REPLACE "^([0-9]+)\\..*$" "\\1" SRC_ARCH_MAJOR "${_SRC_ARCH}")
-      # Check version-less-or-equal, and allow PTX arches to match across majors
-      if (_SRC_ARCH VERSION_LESS_EQUAL _ARCH)
-        if (_SRC_ARCH IN_LIST _PTX_ARCHS OR SRC_ARCH_MAJOR STREQUAL TGT_ARCH_MAJOR)
-          set(_TMP_ARCH "${_SRC_ARCH}")
-        endif()
-      else()
-        # If we hit a version greater than the target, we can break
-        break()
-      endif()
-    endforeach()
-
-    # If we found a matching _TMP_ARCH, append it to _CUDA_ARCHS
-    if (_TMP_ARCH)
-      list(APPEND _CUDA_ARCHS "${_TMP_ARCH}")
-    endif()
-  endforeach()
-
-  list(REMOVE_DUPLICATES _CUDA_ARCHS)
-
-  # reapply +PTX suffix to architectures that requested PTX
-  set(_FINAL_ARCHS)
-  foreach(_arch ${_CUDA_ARCHS})
-    if(_arch IN_LIST _PTX_ARCHS)
-      list(APPEND _FINAL_ARCHS "${_arch}+PTX")
-    else()
-      list(APPEND _FINAL_ARCHS "${_arch}")
-    endif()
-  endforeach()
-  set(_CUDA_ARCHS ${_FINAL_ARCHS})
-
-  set(${OUT_CUDA_ARCHS} ${_CUDA_ARCHS} PARENT_SCOPE)
-endfunction()
-
-#
-# For the given `SRC_ROCM_ARCHS` list of architecture versions in the form 
-# `<name>` compute the "loose intersection" with the `TGT_ROCM_ARCHS` list. 
-# The loose intersection is defined as:
-#   { max{ x \in tgt | x <= y } | y \in src, { x \in tgt | x <= y } != {} }
-#  where `<=` is the version comparison operator.
-# In other words, for each version in `TGT_ROCM_ARCHS` find the highest version
-#  in `SRC_ROCM_ARCHS` that is less or equal to the version in `TGT_ROCM_ARCHS`.
-# The result is stored in `OUT_ROCM_ARCHS`.
-#
-# Example:
-#   SRC_ROCM_ARCHS="gfx900;gfx906;gfx908;gfx90a"
-#   TGT_ROCM_ARCHS="gfx906;gfx908;gfx1030"
-#   hip_archs_loose_intersection(OUT_ROCM_ARCHS SRC_ROCM_ARCHS TGT_ROCM_ARCHS)
-#   OUT_ROCM_ARCHS="gfx906;gfx908"
-#
-function(hip_archs_loose_intersection OUT_ROCM_ARCHS SRC_ROCM_ARCHS TGT_ROCM_ARCHS)
-  list(REMOVE_DUPLICATES SRC_ROCM_ARCHS)
-  
-  # ROCm architectures are typically in format gfxNNN or gfxNNNx where N is a digit
-  # and x is a letter. We can sort them by string comparison which works for this format.
-  list(SORT SRC_ROCM_ARCHS COMPARE STRING ORDER ASCENDING)
-  
-  set(_ROCM_ARCHS)
-  
-  # Find the intersection of supported architectures
-  foreach(_SRC_ARCH ${SRC_ROCM_ARCHS})
-    if(_SRC_ARCH IN_LIST TGT_ROCM_ARCHS)
-      list(APPEND _ROCM_ARCHS ${_SRC_ARCH})
-    endif()
-  endforeach()
-  
-  list(REMOVE_DUPLICATES _ROCM_ARCHS)
-  set(${OUT_ROCM_ARCHS} ${_ROCM_ARCHS} PARENT_SCOPE)
-endfunction()
-
-#
-# Override the GPU architectures detected by cmake/torch and filter them by
-# `GPU_SUPPORTED_ARCHES`. Sets the final set of architectures in
-# `GPU_ARCHES`. This only applies to the HIP language since for CUDA we set 
-# the architectures on a per file basis.
-#
-# Note: this is defined as a macro since it updates `CMAKE_CUDA_FLAGS`.
-#
-macro(override_gpu_arches GPU_ARCHES GPU_LANG GPU_SUPPORTED_ARCHES)
-  set(_GPU_SUPPORTED_ARCHES_LIST ${GPU_SUPPORTED_ARCHES} ${ARGN})
-  message(STATUS "${GPU_LANG} supported arches: ${_GPU_SUPPORTED_ARCHES_LIST}")
-
-  if (${GPU_LANG} STREQUAL "HIP")
-    #
-    # `GPU_ARCHES` controls the `--offload-arch` flags.
-    #
-    # If PYTORCH_ROCM_ARCH env variable exists, then we take it as a list,
-    # if not, then we use CMAKE_HIP_ARCHITECTURES which was generated by calling
-    # "rocm_agent_enumerator" in "enable_language(HIP)"
-    # (in file Modules/CMakeDetermineHIPCompiler.cmake)
-    #
-    if(DEFINED ENV{PYTORCH_ROCM_ARCH})
-      set(HIP_ARCHITECTURES $ENV{PYTORCH_ROCM_ARCH})
-    else()
-      set(HIP_ARCHITECTURES ${CMAKE_HIP_ARCHITECTURES})
-    endif()
-    #
-    # Find the intersection of the supported + detected architectures to
-    # set the module architecture flags.
-    #
-    set(${GPU_ARCHES})
-    foreach (_ARCH ${HIP_ARCHITECTURES})
-      if (_ARCH IN_LIST _GPU_SUPPORTED_ARCHES_LIST)
-        list(APPEND ${GPU_ARCHES} ${_ARCH})
-      endif()
-    endforeach()
-
-    if(NOT ${GPU_ARCHES})
-      message(FATAL_ERROR
-        "None of the detected ROCm architectures: ${HIP_ARCHITECTURES} is"
-        " supported. Supported ROCm architectures are: ${_GPU_SUPPORTED_ARCHES_LIST}.")
-    endif()
-  endif()
-endmacro()
-
-#
-# Define a target named `GPU_MOD_NAME` for a single extension. The
-# arguments are:
-#
-# DESTINATION <dest>         - Module destination directory.
-# LANGUAGE <lang>            - The GPU language for this module, e.g CUDA, HIP,
-#                              etc.
-# SOURCES <sources>          - List of source files relative to CMakeLists.txt
-#                              directory.
-#
-# Optional arguments:
-#
-# ARCHITECTURES <arches>     - A list of target GPU architectures in cmake
-#                              format.
-#                              Refer `CMAKE_CUDA_ARCHITECTURES` documentation
-#                              and `CMAKE_HIP_ARCHITECTURES` for more info.
-#                              ARCHITECTURES will use cmake's defaults if
-#                              not provided.
-# COMPILE_FLAGS <flags>      - Extra compiler flags passed to NVCC/hip.
-# INCLUDE_DIRECTORIES <dirs> - Extra include directories.
-# LIBRARIES <libraries>      - Extra link libraries.
-# WITH_SOABI                 - Generate library with python SOABI suffix name.
-# USE_SABI <version>         - Use python stable api <version>
-#
-# Note: optimization level/debug info is set via cmake build type.
-#
-function (define_gpu_extension_target GPU_MOD_NAME)
-  cmake_parse_arguments(PARSE_ARGV 1
-    GPU
-    "WITH_SOABI"
-    "DESTINATION;LANGUAGE;USE_SABI"
-    "SOURCES;ARCHITECTURES;COMPILE_FLAGS;INCLUDE_DIRECTORIES;LIBRARIES")
-
-  # Add hipify preprocessing step when building with HIP/ROCm.
-  if (GPU_LANGUAGE STREQUAL "HIP")
-    hipify_sources_target(GPU_SOURCES ${GPU_MOD_NAME} "${GPU_SOURCES}")
-  endif()
-
-  if (GPU_WITH_SOABI)
-    set(GPU_WITH_SOABI WITH_SOABI)
-  else()
-    set(GPU_WITH_SOABI)
-  endif()
-
-  if (GPU_USE_SABI)
-    Python3_add_library(${GPU_MOD_NAME} MODULE USE_SABI ${GPU_USE_SABI} ${GPU_WITH_SOABI} "${GPU_SOURCES}")
-  else()
-    Python3_add_library(${GPU_MOD_NAME} MODULE ${GPU_WITH_SOABI} "${GPU_SOURCES}")
-  endif()
-
-  if (GPU_LANGUAGE STREQUAL "HIP")
-    # Make this target dependent on the hipify preprocessor step.
-    add_dependencies(${GPU_MOD_NAME} hipify${GPU_MOD_NAME})
-  endif()
-
-  if (GPU_ARCHITECTURES)
-    if (GPU_LANGUAGE STREQUAL "HIP")
-      # Clear target architectures, we are passing arch flags per source file.
-      set_property(TARGET ${GPU_MOD_NAME} PROPERTY HIP_ARCHITECTURES off)
-    else()
-      set_target_properties(${GPU_MOD_NAME} PROPERTIES
-        ${GPU_LANGUAGE}_ARCHITECTURES "${GPU_ARCHITECTURES}")
-    endif()
-  endif()
-
-  set_property(TARGET ${GPU_MOD_NAME} PROPERTY CXX_STANDARD 17)
-
-  target_compile_options(${GPU_MOD_NAME} PRIVATE
-    $<$<COMPILE_LANGUAGE:${GPU_LANGUAGE}>:${GPU_COMPILE_FLAGS}>)
-
-  target_compile_definitions(${GPU_MOD_NAME} PRIVATE
-    "-DTORCH_EXTENSION_NAME=${GPU_MOD_NAME}")
-
-  target_include_directories(${GPU_MOD_NAME} PRIVATE csrc
-    ${GPU_INCLUDE_DIRECTORIES})
-
-  target_link_libraries(${GPU_MOD_NAME} PRIVATE torch ${GPU_LIBRARIES})
-
-  # Don't use `TORCH_LIBRARIES` for CUDA since it pulls in a bunch of
-  # dependencies that are not necessary and may not be installed.
-  if (GPU_LANGUAGE STREQUAL "CUDA")
-    target_link_libraries(${GPU_MOD_NAME} PRIVATE CUDA::cudart)
-  else()
-    target_link_libraries(${GPU_MOD_NAME} PRIVATE ${TORCH_LIBRARIES})
-  endif()
-
-  install(TARGETS ${GPU_MOD_NAME} LIBRARY DESTINATION ${GPU_DESTINATION} COMPONENT ${GPU_MOD_NAME})
-endfunction()

pyproject.toml

@@ -1,10 +0,0 @@
-[build-system]
-requires = [
-  "cmake>=3.26",
-  "ninja",
-  "packaging",
-  "setuptools>=61",
-  "torch",
-  "wheel",
-]
-build-backend = "setuptools.build_meta"

setup.py

@@ -1,118 +0,0 @@
-import logging
-import os
-from shutil import which, move
-import subprocess
-import sys
-from pathlib import Path
-
-from setuptools import Extension, find_packages, setup
-from setuptools.command.build_ext import build_ext
-
-logger = logging.getLogger(__name__)
-
-
-def is_sccache_available() -> bool:
-    return which("sccache") is not None
-
-
-def is_ccache_available() -> bool:
-    return which("ccache") is not None
-
-
-def is_ninja_available() -> bool:
-    return which("ninja") is not None
-
-
-class CMakeExtension(Extension):
-    def __init__(self, name: str, sourcedir: str = "") -> None:
-        super().__init__(name, sources=[], py_limited_api=True)
-        self.sourcedir = os.fspath(Path(sourcedir).resolve())
-
-
-class CMakeBuild(build_ext):
-    def build_extension(self, ext: CMakeExtension) -> None:
-        ext_fullpath = Path.cwd() / self.get_ext_fullpath(ext.name)
-        extdir = ext_fullpath.parent.resolve()
-
-        debug = int(os.environ.get("DEBUG", 0)) if self.debug is None else self.debug
-        cfg = "Debug" if debug else "Release"
-
-        cmake_generator = os.environ.get("CMAKE_GENERATOR", "")
-
-        # Set Python3_EXECUTABLE instead if you use PYBIND11_FINDPYTHON
-        # EXAMPLE_VERSION_INFO shows you how to pass a value into the C++ code
-        # from Python.
-        cmake_args = [
-            f"-DCMAKE_LIBRARY_OUTPUT_DIRECTORY={extdir}{os.sep}",
-            f"-DPython3_EXECUTABLE={sys.executable}",
-            f"-DCMAKE_BUILD_TYPE={cfg}",  # not used on MSVC, but no harm
-        ]
-        build_args = []
-        if "CMAKE_ARGS" in os.environ:
-            cmake_args += [item for item in os.environ["CMAKE_ARGS"].split(" ") if item]
-
-        if not cmake_generator or cmake_generator == "Ninja":
-            try:
-                import ninja
-
-                ninja_executable_path = Path(ninja.BIN_DIR) / "ninja"
-                cmake_args += [
-                    "-GNinja",
-                    f"-DCMAKE_MAKE_PROGRAM:FILEPATH={ninja_executable_path}",
-                ]
-            except ImportError:
-                pass
-
-        if is_sccache_available():
-            cmake_args += [
-                "-DCMAKE_C_COMPILER_LAUNCHER=sccache",
-                "-DCMAKE_CXX_COMPILER_LAUNCHER=sccache",
-                "-DCMAKE_HIP_COMPILER_LAUNCHER=sccache",
-                "-DCMAKE_OBJC_COMPILER_LAUNCHER=sccache",
-                "-DCMAKE_OBJCXX_COMPILER_LAUNCHER=sccache",
-            ]
-        elif is_ccache_available():
-            cmake_args += [
-                "-DCMAKE_C_COMPILER_LAUNCHER=ccache",
-                "-DCMAKE_CXX_COMPILER_LAUNCHER=ccache",
-                "-DCMAKE_HIP_COMPILER_LAUNCHER=ccache",
-                "-DCMAKE_OBJC_COMPILER_LAUNCHER=ccache",
-                "-DCMAKE_OBJCXX_COMPILER_LAUNCHER=ccache",
-            ]
-
-        num_jobs = os.getenv("MAX_JOBS", None)
-        if num_jobs is not None:
-            num_jobs = int(num_jobs)
-            logger.info("Using MAX_JOBS=%d as the number of jobs.", num_jobs)
-        else:
-            try:
-                # os.sched_getaffinity() isn't universally available, so fall
-                #  back to os.cpu_count() if we get an error here.
-                num_jobs = len(os.sched_getaffinity(0))
-            except AttributeError:
-                num_jobs = os.cpu_count()
-
-        build_temp = Path(self.build_temp) / ext.name
-        if not build_temp.exists():
-            build_temp.mkdir(parents=True)
-
-        subprocess.run(
-            ["cmake", ext.sourcedir, *cmake_args], cwd=build_temp, check=True
-        )
-        subprocess.run(
-            ["cmake", "--build", ".", *build_args], cwd=build_temp, check=True
-        )
-
-
-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_3a886f8_dirty")],
-    cmdclass={"build_ext": CMakeBuild},
-    packages=find_packages(where="torch-ext", include=["gptoss_kernels*"]),
-    package_dir={"": "torch-ext"},
-    zip_safe=False,
-    install_requires=["torch"],
-    python_requires=">=3.9",
-)

torch-ext/gptoss_kernels/test.py

@@ -1,15 +0,0 @@
-import _gptoss_kernels_931bc1b_dirty 
-import torch
-
-print(dir(_gptoss_kernels_931bc1b_dirty))
-
-from gptoss_kernels import _gptoss_kernels_931bc1b_dirty
-
-print(dir(f32_bf16w_matmul))
-
-input = torch.randn(10, 10)
-weight_bf16 = torch.randn(10, 10)
-bias_bf16 = torch.randn(10)
-output = torch.randn(10, 10)
-f32_bf16w_matmul(input, weight_bf16, bias_bf16, output, 10, 10, 10, 10)
-print(output)

torch-ext/registration.h

@@ -1,30 +0,0 @@
-// Registration macros from vLLM:
-// https://github.com/vllm-project/vllm/blob/main/csrc/core/registration.h
-
-#pragma once
-
-#include <Python.h>
-
-#define _CONCAT(A, B) A##B
-#define CONCAT(A, B) _CONCAT(A, B)
-
-#define _STRINGIFY(A) #A
-#define STRINGIFY(A) _STRINGIFY(A)
-
-// A version of the TORCH_LIBRARY macro that expands the NAME, i.e. so NAME
-// could be a macro instead of a literal token.
-#define TORCH_LIBRARY_EXPAND(NAME, MODULE) TORCH_LIBRARY(NAME, MODULE)
-
-// A version of the TORCH_LIBRARY_IMPL macro that expands the NAME, i.e. so NAME
-// could be a macro instead of a literal token.
-#define TORCH_LIBRARY_IMPL_EXPAND(NAME, DEVICE, MODULE) \
-  TORCH_LIBRARY_IMPL(NAME, DEVICE, MODULE)
-
-// REGISTER_EXTENSION allows the shared library to be loaded and initialized
-// via python's import statement.
-#define REGISTER_EXTENSION(NAME)                                               \
-  PyMODINIT_FUNC CONCAT(PyInit_, NAME)() {                                     \
-    static struct PyModuleDef module = {PyModuleDef_HEAD_INIT,                 \
-                                        STRINGIFY(NAME), nullptr, 0, nullptr}; \
-    return PyModule_Create(&module);                                           \
-  }