commit evtn

CMakeLists.txt

@@ -0,0 +1,104 @@
+cmake_minimum_required(VERSION 3.26)
+project(gptoss_kernels LANGUAGES CXX)
+
+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/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"
+)
+
+# 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()
+# Include Metal shader compilation utilities
+include(${CMAKE_CURRENT_LIST_DIR}/cmake/compile-metal.cmake)
+
+define_gpu_extension_target(
+  _gptoss_kernels_931bc1b_dirty
+  DESTINATION _gptoss_kernels_931bc1b_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_931bc1b_dirty "${ALL_METAL_SOURCES}")
+endif()

README.md

@@ -8,3 +8,107 @@ tags:
 
 This is a build for some kernel released by OpenAI in the GPT-OSS repo : https://github.com/openai/gpt-oss
 
+```21:69:/Users/medmekk/projects/ai/kernels/gpt-oss/gpt_oss/metal/source/matmul.metal
+kernel void gptoss_f32_bf16w_matmul(
+    constant gptoss_matmul_args& args [[ buffer(0) ]],
+    const device float4* input [[ buffer(1) ]],
+    const device bfloat4* weight [[ buffer(2) ]],
+    const device bfloat* bias [[ buffer(3) ]],
+    device float* output [[ buffer(4) ]],
+    const device gptoss_control* control [[ buffer(5) ]],
+    uint2 gid [[threadgroup_position_in_grid]],
+    uint simdgroup_tid [[thread_index_in_simdgroup]],
+    uint simdgroup_idx [[simdgroup_index_in_threadgroup]],
+    uint num_simdgroups [[simdgroups_per_threadgroup]])
+{
+    const uint simdgroup_size = 32;
+    if (control->abort != 0) {
+        return;
+    }
+
+    const uint num_column_vecs = args.num_column_vecs;
+    const uint row = gid.x * num_simdgroups + simdgroup_idx;
+
+    input += gid.y * num_column_vecs + simdgroup_tid;
+    weight += num_column_vecs * row + simdgroup_tid;
+    bias += row;
+    output += gid.y * args.num_rows + row;
+
+    uint num_iter = 0;
+    num_iter = (num_column_vecs - simdgroup_tid + (simdgroup_size - 1)) / simdgroup_size;
+
+    float4 sum4 = 0.0f;
+    do {
+        const bfloat4 w = *weight;
+        const float4 i = *input;
+        sum4 = metal::fma(static_cast<float4>(w), i, sum4);
+
+        weight += simdgroup_size;
+        input += simdgroup_size;
+    } while (--num_iter != 0);
+    const float2 sum2 = sum4.xy + sum4.zw;
+    float sum = sum2.x + sum2.y;
+    sum = metal::simd_sum(sum);
+    if (metal::simd_is_first()) {
+        sum += static_cast<float>(*bias);
+        if (args.add) {
+            *output += sum;
+        } else {
+            *output = sum;
+        }
+    }
+}
+```
+
+### What it computes
+- Computes Y = X · W + b for a batch of tokens.
+- Types/layout:
+  - X is float32, shape [num_tokens, num_cols], viewed as `float4` vectors → num_column_vecs = num_cols/4.
+  - W is bfloat16, shape [num_rows, num_cols], viewed as `bfloat4` vectors per row (row-major).
+  - b is bfloat16, length num_rows.
+  - Y is float32, shape [num_tokens, num_rows].
+
+### Work decomposition
+- Grid Y (gid.y) = token index t in [0, num_tokens).
+- Grid X (gid.x) spans output rows in groups of `num_simdgroups`. Within a threadgroup:
+  - simdgroup_idx in [0, num_simdgroups) selects one output row r.
+  - Therefore row r = gid.x*num_simdgroups + simdgroup_idx.
+- Each simdgroup computes exactly one scalar Y[t, r].
+- Lanes inside a simdgroup (simdgroup_tid in [0, 31]) split the K dimension (num_column_vecs vectors) in a strided pattern: lane ℓ processes indices ℓ, ℓ+32, ℓ+64, ...
+
+### Pointer setup per simdgroup
+- input points to the start of token t’s vector-list, then lane offset: `input += t*num_column_vecs + lane`.
+- weight points to the start of row r’s vector-list, then lane offset: `weight += r*num_column_vecs + lane`.
+- bias points to `b[r]`. output points to `Y[t, r]`.
+
+### Inner loop
+- num_iter = ceil((num_column_vecs - lane)/32). Each lane loops over its share of the K/4 vectors.
+- On each iteration:
+  - Load one `float4` from input (4 consecutive columns) and one `bfloat4` from weight for row r.
+  - Fused multiply-add into `sum4` (vector-wise).
+  - Advance both pointers by 32 vectors (next stripe for this lane).
+- After loop, reduce the 4 lanes of `sum4` into a scalar: sum4.xy + sum4.zw → sum2, then sum2.x + sum2.y.
+- Reduce across all 32 lanes with `metal::simd_sum(sum)`.
+- Lane 0 adds bias[r] and writes `Y[t, r]` (add or overwrite depending on args.add).
+
+### Example mapping (num_tokens=2, num_cols=128, num_rows=4, threadgroup_size=32)
+- num_column_vecs = 128/4 = 32.
+- threadgroup_size=32 ⇒ num_simdgroups=1 ⇒ each threadgroup computes 1 row.
+- Grid:
+  - gid.y ∈ {0,1} (two tokens).
+  - gid.x ∈ {0,1,2,3} (four rows).
+- For a given (gid.x, gid.y), simdgroup_idx=0 computes one output scalar Y[t=gid.y, r=gid.x].
+- Per-lane work:
+  - lane ℓ loads X[t, cols 4ℓ..4ℓ+3] as float4 and W[r, cols 4ℓ..4ℓ+3] as bfloat4 (exactly one iteration since 32 vectors total).
+  - Each lane accumulates the dot over its 4 elements; lanes are then summed → full dot(X_row, W_row).
+  - Lane 0 writes Y[t, r] = dot + b[r].
+
+### Example mapping (same shapes, threadgroup_size=64)
+- num_simdgroups=64/32=2 ⇒ each threadgroup computes 2 rows at once.
+- For gid.x=k:
+  - simdgroup_idx=0 computes row r=2k, simdgroup_idx=1 computes row r=2k+1.
+- Lanes split K identically; two output scalars are produced per threadgroup (one per simdgroup).
+
+### Which piece each unit owns
+- Token t (grid y) × Row r (simdgroup within grid x) → one output scalar Y[t, r].
+- Lane ℓ within that simdgroup → partial dot over columns {4ℓ, 4ℓ+1, 4ℓ+2, 4ℓ+3}, plus any further stripes {4(ℓ+32), ...} if num_cols > 128.

build.toml

@@ -9,15 +9,29 @@ src = [
 ]
 
 [kernel.gptoss_kernels]
+
 depends = ["torch"]
-backend = "cuda"
+backend = "metal"
 
 src = [
-    "gptoss_kernels/attention_cuda_fwd.cu",
-    "gptoss_kernels/attention_cuda_bwd.cu",
-    "gptoss_kernels/attention_cuda_utils.cu",
-    "gptoss_kernels/attention_cuda_utils.cuh",
-    "gptoss_kernels/attention_cuda.cuh",
-    "gptoss_kernels/attention.h",
-    "gptoss_kernels/cudamacro.h",
+  "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",
 ]
+
+include = ["gptoss_kernels/source/include", "gptoss_kernels/include", "."]

cmake/compile-metal.cmake

@@ -0,0 +1,86 @@
+# Metal shader compilation function
+function(compile_metal_shaders TARGET_NAME METAL_SOURCES)
+    # Find the Metal compiler
+    find_program(METAL_COMPILER xcrun REQUIRED)
+    
+    # Set Metal compiler flags
+    set(METAL_FLAGS "-std=metal3.0" "-O2")
+    
+    # Output directory for compiled metallib
+    set(METALLIB_OUTPUT_DIR "${CMAKE_BINARY_DIR}/metallib")
+    file(MAKE_DIRECTORY ${METALLIB_OUTPUT_DIR})
+    
+    # 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

@@ -0,0 +1,73 @@
+#!/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

@@ -0,0 +1,557 @@
+# 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()

flake.lock

@@ -0,0 +1,169 @@
+{
+  "nodes": {
+    "flake-compat": {
+      "locked": {
+        "lastModified": 1761588595,
+        "narHash": "sha256-XKUZz9zewJNUj46b4AJdiRZJAvSZ0Dqj2BNfXvFlJC4=",
+        "owner": "edolstra",
+        "repo": "flake-compat",
+        "rev": "f387cd2afec9419c8ee37694406ca490c3f34ee5",
+        "type": "github"
+      },
+      "original": {
+        "owner": "edolstra",
+        "repo": "flake-compat",
+        "type": "github"
+      }
+    },
+    "flake-compat_2": {
+      "locked": {
+        "lastModified": 1747046372,
+        "narHash": "sha256-CIVLLkVgvHYbgI2UpXvIIBJ12HWgX+fjA8Xf8PUmqCY=",
+        "owner": "edolstra",
+        "repo": "flake-compat",
+        "rev": "9100a0f413b0c601e0533d1d94ffd501ce2e7885",
+        "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"
+      }
+    },
+    "flake-utils_2": {
+      "inputs": {
+        "systems": "systems_2"
+      },
+      "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"
+      }
+    },
+    "hf-nix": {
+      "inputs": {
+        "flake-compat": "flake-compat_2",
+        "flake-utils": "flake-utils_2",
+        "nixpkgs": "nixpkgs"
+      },
+      "locked": {
+        "lastModified": 1761756835,
+        "narHash": "sha256-Vjrv8ZIhkQRgQ3MHGVFaj/fUcE4yuGr+vnoKYRwWmYw=",
+        "owner": "huggingface",
+        "repo": "hf-nix",
+        "rev": "6839b6998be18679992978c2f3abddc902276280",
+        "type": "github"
+      },
+      "original": {
+        "owner": "huggingface",
+        "repo": "hf-nix",
+        "type": "github"
+      }
+    },
+    "kernel-builder": {
+      "inputs": {
+        "flake-compat": "flake-compat",
+        "flake-utils": "flake-utils",
+        "hf-nix": "hf-nix",
+        "nixpkgs": [
+          "kernel-builder",
+          "hf-nix",
+          "nixpkgs"
+        ]
+      },
+      "locked": {
+        "lastModified": 1761991868,
+        "narHash": "sha256-+csvkWC9jC4mwq1LNfK4O6m3Qg4dCCXjP5JGdPa3TEo=",
+        "owner": "huggingface",
+        "repo": "kernel-builder",
+        "rev": "79cbfcdfde82c8847551f67f4b951a410794a5c6",
+        "type": "github"
+      },
+      "original": {
+        "owner": "huggingface",
+        "ref": "metal_kernels",
+        "repo": "kernel-builder",
+        "type": "github"
+      }
+    },
+    "nixpkgs": {
+      "locked": {
+        "lastModified": 1755963616,
+        "narHash": "sha256-6yD0ww/S8n+U2uPYcJZ3DRURP8Kx036GRpR2uPNZroE=",
+        "owner": "nixos",
+        "repo": "nixpkgs",
+        "rev": "73e96df7cff5783f45e21342a75a1540c4eddce4",
+        "type": "github"
+      },
+      "original": {
+        "owner": "nixos",
+        "ref": "nixos-unstable-small",
+        "repo": "nixpkgs",
+        "type": "github"
+      }
+    },
+    "root": {
+      "inputs": {
+        "kernel-builder": "kernel-builder"
+      }
+    },
+    "systems": {
+      "locked": {
+        "lastModified": 1681028828,
+        "narHash": "sha256-Vy1rq5AaRuLzOxct8nz4T6wlgyUR7zLU309k9mBC768=",
+        "owner": "nix-systems",
+        "repo": "default",
+        "rev": "da67096a3b9bf56a91d16901293e51ba5b49a27e",
+        "type": "github"
+      },
+      "original": {
+        "owner": "nix-systems",
+        "repo": "default",
+        "type": "github"
+      }
+    },
+    "systems_2": {
+      "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
+}

flake.nix

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

gptoss_kernels/CMakeLists.txt

@@ -1,191 +0,0 @@
-cmake_minimum_required(VERSION 3.24)
-project(GPTOSS
-    VERSION 1.0
-    DESCRIPTION "Local GPT-OSS inference"
-    LANGUAGES C CXX OBJC)
-
-set(CMAKE_C_STANDARD 11)
-set(CMAKE_CXX_STANDARD 20)
-set(CMAKE_OBJC_STANDARD 11)
-set(CMAKE_OBJC_STANDARD_REQUIRED ON)
-
-find_library(FOUNDATION_FRAMEWORK Foundation REQUIRED)
-find_library(METAL_FRAMEWORK      Metal      REQUIRED)
-find_library(IOKIT_FRAMEWORK      IOKit      REQUIRED)
-
-set(METAL_SOURCES
-    ${CMAKE_CURRENT_SOURCE_DIR}/source/accumulate.metal
-    ${CMAKE_CURRENT_SOURCE_DIR}/source/convert.metal
-    ${CMAKE_CURRENT_SOURCE_DIR}/source/embeddings.metal
-    ${CMAKE_CURRENT_SOURCE_DIR}/source/expert_routing_metadata.metal
-    ${CMAKE_CURRENT_SOURCE_DIR}/source/gather_and_accumulate.metal
-    ${CMAKE_CURRENT_SOURCE_DIR}/source/matmul.metal
-    ${CMAKE_CURRENT_SOURCE_DIR}/source/moematmul.metal
-    ${CMAKE_CURRENT_SOURCE_DIR}/source/random.metal
-    ${CMAKE_CURRENT_SOURCE_DIR}/source/rmsnorm.metal
-    ${CMAKE_CURRENT_SOURCE_DIR}/source/rope.metal
-    ${CMAKE_CURRENT_SOURCE_DIR}/source/sample.metal
-    ${CMAKE_CURRENT_SOURCE_DIR}/source/scatter.metal
-    ${CMAKE_CURRENT_SOURCE_DIR}/source/sdpa.metal
-    ${CMAKE_CURRENT_SOURCE_DIR}/source/topk.metal
-)
-set(METAL_LIB default.metallib)
-
-include_directories(BEFORE include source/include)
-
-add_custom_command(
-    OUTPUT  ${CMAKE_CURRENT_BINARY_DIR}/${METAL_LIB}
-    COMMAND ${CMAKE_COMMAND} -E make_directory "${CMAKE_CURRENT_BINARY_DIR}/source/"
-    COMMAND xcrun -sdk macosx metal -g "-I${CMAKE_CURRENT_SOURCE_DIR}/source/include" -c "${CMAKE_CURRENT_SOURCE_DIR}/source/accumulate.metal" -o "${CMAKE_CURRENT_BINARY_DIR}/source/accumulate.air"
-    COMMAND xcrun -sdk macosx metal -g "-I${CMAKE_CURRENT_SOURCE_DIR}/source/include" -c "${CMAKE_CURRENT_SOURCE_DIR}/source/convert.metal" -o "${CMAKE_CURRENT_BINARY_DIR}/source/convert.air"
-    COMMAND xcrun -sdk macosx metal -g "-I${CMAKE_CURRENT_SOURCE_DIR}/source/include" -c "${CMAKE_CURRENT_SOURCE_DIR}/source/embeddings.metal" -o "${CMAKE_CURRENT_BINARY_DIR}/source/embeddings.air"
-    COMMAND xcrun -sdk macosx metal -g "-I${CMAKE_CURRENT_SOURCE_DIR}/source/include" -c "${CMAKE_CURRENT_SOURCE_DIR}/source/expert_routing_metadata.metal" -o "${CMAKE_CURRENT_BINARY_DIR}/source/expert_routing_metadata.air"
-    COMMAND xcrun -sdk macosx metal -g "-I${CMAKE_CURRENT_SOURCE_DIR}/source/include" -c "${CMAKE_CURRENT_SOURCE_DIR}/source/matmul.metal" -o "${CMAKE_CURRENT_BINARY_DIR}/source/matmul.air"
-    COMMAND xcrun -sdk macosx metal -g "-I${CMAKE_CURRENT_SOURCE_DIR}/source/include" -c "${CMAKE_CURRENT_SOURCE_DIR}/source/moematmul.metal" -o "${CMAKE_CURRENT_BINARY_DIR}/source/moematmul.air"
-    COMMAND xcrun -sdk macosx metal -g "-I${CMAKE_CURRENT_SOURCE_DIR}/source/include" -c "${CMAKE_CURRENT_SOURCE_DIR}/source/gather_and_accumulate.metal" -o "${CMAKE_CURRENT_BINARY_DIR}/source/gather_and_accumulate.air"
-    COMMAND xcrun -sdk macosx metal -g "-I${CMAKE_CURRENT_SOURCE_DIR}/source/include" -c "${CMAKE_CURRENT_SOURCE_DIR}/source/random.metal" -o "${CMAKE_CURRENT_BINARY_DIR}/source/random.air"
-    COMMAND xcrun -sdk macosx metal -g "-I${CMAKE_CURRENT_SOURCE_DIR}/source/include" -c "${CMAKE_CURRENT_SOURCE_DIR}/source/rmsnorm.metal" -o "${CMAKE_CURRENT_BINARY_DIR}/source/rmsnorm.air"
-    COMMAND xcrun -sdk macosx metal -g "-I${CMAKE_CURRENT_SOURCE_DIR}/source/include" -c "${CMAKE_CURRENT_SOURCE_DIR}/source/rope.metal" -o "${CMAKE_CURRENT_BINARY_DIR}/source/rope.air"
-    COMMAND xcrun -sdk macosx metal -g "-I${CMAKE_CURRENT_SOURCE_DIR}/source/include" -c "${CMAKE_CURRENT_SOURCE_DIR}/source/sample.metal" -o "${CMAKE_CURRENT_BINARY_DIR}/source/sample.air"
-    COMMAND xcrun -sdk macosx metal -g "-I${CMAKE_CURRENT_SOURCE_DIR}/source/include" -c "${CMAKE_CURRENT_SOURCE_DIR}/source/scatter.metal" -o "${CMAKE_CURRENT_BINARY_DIR}/source/scatter.air"
-    COMMAND xcrun -sdk macosx metal -g "-I${CMAKE_CURRENT_SOURCE_DIR}/source/include" -c "${CMAKE_CURRENT_SOURCE_DIR}/source/sdpa.metal" -o "${CMAKE_CURRENT_BINARY_DIR}/source/sdpa.air"
-    COMMAND xcrun -sdk macosx metal -g "-I${CMAKE_CURRENT_SOURCE_DIR}/source/include" -c "${CMAKE_CURRENT_SOURCE_DIR}/source/topk.metal" -o "${CMAKE_CURRENT_BINARY_DIR}/source/topk.air"
-    COMMAND xcrun -sdk macosx metallib "${CMAKE_CURRENT_BINARY_DIR}/source/accumulate.air" "${CMAKE_CURRENT_BINARY_DIR}/source/convert.air" "${CMAKE_CURRENT_BINARY_DIR}/source/embeddings.air" "${CMAKE_CURRENT_BINARY_DIR}/source/expert_routing_metadata.air" "${CMAKE_CURRENT_BINARY_DIR}/source/gather_and_accumulate.air" "${CMAKE_CURRENT_BINARY_DIR}/source/matmul.air" "${CMAKE_CURRENT_BINARY_DIR}/source/moematmul.air" "${CMAKE_CURRENT_BINARY_DIR}/source/random.air" "${CMAKE_CURRENT_BINARY_DIR}/source/rmsnorm.air" "${CMAKE_CURRENT_BINARY_DIR}/source/rope.air" "${CMAKE_CURRENT_BINARY_DIR}/source/sample.air" "${CMAKE_CURRENT_BINARY_DIR}/source/scatter.air" "${CMAKE_CURRENT_BINARY_DIR}/source/sdpa.air" "${CMAKE_CURRENT_BINARY_DIR}/source/topk.air" -o "${METAL_LIB}"
-    DEPENDS ${METAL_SOURCES}
-    COMMENT "Compiling Metal compute library"
-)
-
-add_custom_target(build_metallib ALL
-    DEPENDS ${CMAKE_CURRENT_BINARY_DIR}/${METAL_LIB})
-
-add_library(log OBJECT source/log.c)
-
-add_library(metal-kernels STATIC source/metal.m source/metal-kernels.c)
-target_link_libraries(metal-kernels PRIVATE log)
-
-add_dependencies(metal-kernels build_metallib)
-add_custom_command(TARGET metal-kernels POST_BUILD
-    COMMAND ${CMAKE_COMMAND} -E copy
-            ${CMAKE_CURRENT_BINARY_DIR}/${METAL_LIB}
-            $<TARGET_FILE_DIR:metal-kernels>)
-
-target_link_libraries(metal-kernels PRIVATE ${FOUNDATION_FRAMEWORK} ${METAL_FRAMEWORK} ${IOKIT_FRAMEWORK})
-
-add_library(gptoss STATIC source/model.c source/tokenizer.c source/context.c)
-target_link_libraries(gptoss PRIVATE log metal-kernels)
-
-add_executable(generate source/generate.c)
-target_link_libraries(generate gptoss)
-
-# --- [ Tests
-include(FetchContent)
-FetchContent_Declare(
-    googletest
-    URL https://github.com/google/googletest/archive/refs/tags/v1.17.0.zip
-    DOWNLOAD_EXTRACT_TIMESTAMP OFF
-)
-# For Windows: Prevent overriding the parent project's compiler/linker settings
-set(gtest_force_shared_crt ON CACHE BOOL "" FORCE)
-set(INSTALL_GTEST OFF CACHE BOOL "" FORCE)
-FetchContent_MakeAvailable(googletest)
-
-enable_testing()
-
-add_executable(u32-random-test test/u32-random.cc)
-target_link_libraries(u32-random-test PRIVATE GTest::gtest_main metal-kernels)
-target_include_directories(u32-random-test PRIVATE source/include)
-add_test(NAME u32-random-test COMMAND u32-random-test)
-
-add_executable(f32-random-test test/f32-random.cc)
-target_link_libraries(f32-random-test PRIVATE GTest::gtest_main metal-kernels)
-target_include_directories(f32-random-test PRIVATE source/include)
-add_test(NAME f32-random-test COMMAND f32-random-test)
-
-add_executable(mf4-f32-convert-test test/mf4-f32-convert.cc)
-target_link_libraries(mf4-f32-convert-test PRIVATE GTest::gtest_main metal-kernels)
-target_include_directories(mf4-f32-convert-test PRIVATE source/include)
-add_test(NAME mf4-f32-convert-test COMMAND mf4-f32-convert-test)
-
-add_executable(bf16-f32-embeddings-test test/bf16-f32-embeddings.cc)
-target_link_libraries(bf16-f32-embeddings-test PRIVATE GTest::gtest_main metal-kernels)
-target_include_directories(bf16-f32-embeddings-test PRIVATE source/include)
-add_test(NAME bf16-f32-embeddings-test COMMAND bf16-f32-embeddings-test)
-
-add_executable(f32-bf16w-rmsnorm-test test/f32-bf16w-rmsnorm.cc)
-target_link_libraries(f32-bf16w-rmsnorm-test PRIVATE GTest::gtest_main metal-kernels)
-target_include_directories(f32-bf16w-rmsnorm-test PRIVATE source/include)
-add_test(NAME f32-bf16w-rmsnorm-test COMMAND f32-bf16w-rmsnorm-test)
-
-add_executable(f32-bf16w-matmul-test test/f32-bf16w-matmul.cc)
-target_link_libraries(f32-bf16w-matmul-test PRIVATE GTest::gtest_main metal-kernels)
-target_include_directories(f32-bf16w-matmul-test PRIVATE source/include)
-add_test(NAME f32-bf16w-matmul-test COMMAND f32-bf16w-matmul-test)
-
-add_executable(f32-rope-test test/f32-rope.cc)
-target_link_libraries(f32-rope-test PRIVATE GTest::gtest_main metal-kernels)
-target_include_directories(f32-rope-test PRIVATE source/include)
-add_test(NAME f32-rope-test COMMAND f32-rope-test)
-
-# --- [ Benchmarks
-include(FetchContent)
-set(BENCHMARK_ENABLE_TESTING OFF CACHE BOOL "Disable self-tests in Google Benchmark" FORCE)
-set(BENCHMARK_ENABLE_INSTALL OFF CACHE BOOL "Disable installation of Google Benchmark" FORCE)
-FetchContent_Declare(
-    benchmark
-    URL https://github.com/google/benchmark/archive/refs/tags/v1.9.4.zip
-    DOWNLOAD_EXTRACT_TIMESTAMP OFF
-)
-FetchContent_MakeAvailable(benchmark)
-
-add_executable(f32-random-bench benchmark/f32-random.cc)
-target_link_libraries(f32-random-bench PRIVATE benchmark::benchmark metal-kernels)
-target_include_directories(f32-random-bench PRIVATE source/include)
-
-add_executable(u32-random-bench benchmark/u32-random.cc)
-target_link_libraries(u32-random-bench PRIVATE benchmark::benchmark metal-kernels)
-target_include_directories(u32-random-bench PRIVATE source/include)
-
-add_executable(mf4-f32-convert-bench benchmark/mf4-f32-convert.cc)
-target_link_libraries(mf4-f32-convert-bench PRIVATE benchmark::benchmark metal-kernels)
-target_include_directories(mf4-f32-convert-bench PRIVATE source/include)
-
-add_executable(f32-bf16w-rmsnorm-bench benchmark/f32-bf16w-rmsnorm.cc)
-target_link_libraries(f32-bf16w-rmsnorm-bench PRIVATE benchmark::benchmark metal-kernels)
-target_include_directories(f32-bf16w-rmsnorm-bench PRIVATE source/include)
-
-add_executable(end-to-end-bench benchmark/end-to-end.cc)
-target_link_libraries(end-to-end-bench PRIVATE benchmark::benchmark gptoss)
-target_include_directories(end-to-end-bench PRIVATE source/include)
-
-add_executable(end-to-end-threadgroup-bench benchmark/end-to-end-threadgroup.cc)
-target_link_libraries(end-to-end-threadgroup-bench PRIVATE benchmark::benchmark gptoss)
-target_include_directories(end-to-end-threadgroup-bench PRIVATE source/include)
-
-# --- [ Python extension ] -----------------------------------------------
-find_package(pybind11 CONFIG REQUIRED)          # provides pybind11_add_module
-
-pybind11_add_module(_metal
-    python/module.c
-    python/context.c
-    python/model.c
-    python/tokenizer.c
-)
-set_target_properties(_metal PROPERTIES PREFIX "")
-
-target_link_libraries(_metal PRIVATE gptoss)
-add_dependencies(_metal build_metallib)
-target_link_options(_metal PRIVATE
-    LINKER:-sectcreate,__METAL,__shaders,${CMAKE_CURRENT_BINARY_DIR}/${METAL_LIB}
-)
-add_custom_command(TARGET _metal POST_BUILD
-    COMMAND ${CMAKE_COMMAND} -E copy
-            ${CMAKE_CURRENT_BINARY_DIR}/${METAL_LIB}
-            $<TARGET_FILE_DIR:_metal>)
-
-# 1️⃣  install the extension module into the Python package
-install(TARGETS _metal LIBRARY DESTINATION gpt_oss/metal)
-
-# 2️⃣  make sure the Metal shader archive travels with it
-install(FILES ${CMAKE_CURRENT_BINARY_DIR}/${METAL_LIB}
-        DESTINATION gpt_oss/metal)
-# ------------------------------------------------------------------------

gptoss_kernels/__init__.py

@@ -1,6 +0,0 @@
-from importlib import import_module as _im
-
-# Load the compiled extension (gpt_oss.metal._metal)
-_ext = _im(f"{__name__}._metal")
-globals().update({k: v for k, v in _ext.__dict__.items() if not k.startswith("_")})
-del _im, _ext

gptoss_kernels/examples/chat.py

@@ -1,104 +0,0 @@
-#!/usr/bin/env python
-
-import argparse
-import sys
-
-from datetime import date
-from gpt_oss.metal import Context, Model
-
-
-DEFAULT_PROMPT = f"""You are ChatGPT, a large language model trained by OpenAI.
-Knowledge cutoff: 2024-06
-Current date: {date.today().isoformat()}
-
-reasoning effort high
-
-# Valid channels: analysis, final. Channel must be included for every message."""
-
-
-parser = argparse.ArgumentParser(description="Chat with gpt-oss", formatter_class=argparse.ArgumentDefaultsHelpFormatter)
-parser.add_argument("model", metavar="PATH", type=str, help="Path to gpt-oss model in Metal inference format")
-parser.add_argument("--prompt", type=str, default=DEFAULT_PROMPT, help="System prompt")
-parser.add_argument(
-    "--context-length", type=int, default=0, help="The maximum context length"
-)
-parser.add_argument(
-    "--temperature", type=float, default=1.0, help="Sampling temperature"
-)
-parser.add_argument(
-    "--seed", type=int, default=0, help="Sampling seed"
-)
-
-
-GREY = "\33[90m"
-BOLD = "\33[1m"
-RESET = "\33[0m"
-
-
-def main(args):
-    options = parser.parse_args(args)
-    model = Model(options.model)
-    tokenizer = model.tokenizer
-    start_token = tokenizer.encode_special_token("<|start|>")
-    message_token = tokenizer.encode_special_token("<|message|>")
-    end_token = tokenizer.encode_special_token("<|end|>")
-    return_token = tokenizer.encode_special_token("<|return|>")
-    channel_token = tokenizer.encode_special_token("<|channel|>")
-
-    context = Context(model, context_length=options.context_length)
-    context.append(start_token)
-    context.append("system")
-    context.append(message_token)
-    context.append(options.prompt)
-    context.append(end_token)
-
-    while True:
-        context.append(start_token)
-        context.append("user")
-        context.append(message_token)
-        message = input(f"{BOLD}User:{RESET} ").rstrip()
-        context.append(message)
-        context.append(end_token)
-        print(f"{BOLD}Assistant:{RESET} {GREY}", end="", flush=True)
-        context.append(start_token)
-        context.append("assistant")
-        context.append(channel_token)
-
-        inside_start_block = True
-        inside_channel_block = True
-        role = "assistant"
-        channel = ""
-        while True:
-            token = context.sample(
-                temperature=options.temperature,
-                seed=options.seed,
-            )
-            context.append(token)
-            if token == return_token:
-                print(flush=True)
-                break
-            elif token == start_token:
-                inside_start_block = True
-                role = ""
-                channel = ""
-            elif token == message_token:
-                inside_start_block = False
-                inside_channel_block = False
-                if channel == "analysis":
-                    print(f"{GREY}", end="", flush=True)
-            elif token == end_token:
-                print(f"{RESET}", flush=True)
-            elif token == channel_token:
-                inside_channel_block = True
-            elif token < tokenizer.num_text_tokens:
-                if inside_channel_block:
-                    channel += str(tokenizer.decode(token), encoding="utf-8")
-                elif inside_start_block:
-                    role += str(tokenizer.decode(token), encoding="utf-8")
-                else:
-                    sys.stdout.buffer.write(tokenizer.decode(token))
-                    sys.stdout.buffer.flush()
-
-
-if __name__ == "__main__":
-    main(sys.argv[1:])

gptoss_kernels/examples/generate.py

@@ -1,34 +0,0 @@
-#!/usr/bin/env python
-
-import argparse
-import sys
-
-from gpt_oss.metal import Context, Model
-
-
-parser = argparse.ArgumentParser(description='Chat with gpt-oss', formatter_class=argparse.ArgumentDefaultsHelpFormatter)
-parser.add_argument('model', metavar='PATH', type=str, help='Path to gpt-oss checkpoint')
-parser.add_argument('-p', '--prompt', type=str, required=True, help='Prompt')
-parser.add_argument('-l', '--limit', type=int, default=100, help='Number of tokens to generate')
-parser.add_argument('--context-length', type=int, default=0, help='The maximum context length')
-
-
-def main(args):
-    options = parser.parse_args(args)
-    model = Model(options.model)
-
-    context = Context(model, context_length=options.context_length)
-    context.append(options.prompt)
-    print(context.tokens)
-    prompt_tokens = context.num_tokens
-
-    tokenizer = model.tokenizer
-
-    while context.num_tokens - prompt_tokens < options.limit:
-        token = context.sample()
-        context.append(token)
-        print(str(tokenizer.decode(token), encoding="utf-8"), end='', flush=True)
-
-
-if __name__ == '__main__':
-    main(sys.argv[1:])

gptoss_kernels/source/context.c

@@ -1,1115 +0,0 @@
-#include <assert.h>
-#include <float.h>
-#include <inttypes.h>
-#include <stdbool.h>
-#include <stdint.h>
-#include <stdlib.h>
-#include <string.h>
-
-#include <gpt-oss.h>
-
-#include "internal/datatype.h"
-#include "internal/model.h"
-#include "internal/metal.h"
-#include "internal/metal-kernels.h"
-#include "internal/log.h"
-#include "internal/rng.h"
-
-
-enum gptoss_status GPTOSS_ABI gptoss_context_create(
-    gptoss_model_t model,
-    size_t context_length,
-    size_t max_batch_tokens,
-    gptoss_context_t* context_out)
-{
-    *context_out = NULL;
-
-    enum gptoss_status status = gptoss_status_success;
-    struct gptoss_context* context = NULL;
-
-    // Validate context_length
-    if (context_length == 0) {
-        context_length = model->context_length;
-    } else if (context_length > model->context_length) {
-        GPTOSS_LOG_ERROR("requested context length %zu exceeds model context length %" PRIu32,
-            context_length, model->context_length);
-        status = gptoss_status_invalid_argument;
-        goto cleanup;
-    }
-    assert(context_length != 0);
-    assert(context_length <= model->context_length);
-
-    // Validate max_batch_tokens
-    if (max_batch_tokens == 0) {
-        max_batch_tokens = GPTOSS_DEFAULT_BATCH_SIZE;
-    } else if (max_batch_tokens > context_length) {
-        GPTOSS_LOG_ERROR("requested max batch tokens %zu exceeds context length %zu",
-            max_batch_tokens, context_length);
-        status = gptoss_status_invalid_argument;
-        goto cleanup;
-    }
-    assert(max_batch_tokens != 0);
-    assert(max_batch_tokens <= context_length);
-
-    context = malloc(sizeof(struct gptoss_context));
-    if (context == NULL) {
-        GPTOSS_LOG_ERROR("failed to allocate %zu bytes for Context object",
-            sizeof(struct gptoss_context));
-        status = gptoss_status_insufficient_memory;
-        goto cleanup;
-    }
-    memset(context, 0, sizeof(struct gptoss_context));
-
-    atomic_store_explicit(&context->ref_count, 1, memory_order_relaxed);
-    context->max_tokens = context_length;
-    context->max_batch_tokens = max_batch_tokens;
-
-    // Activation buffers
-    status = gptoss_metal_buffer_create(&model->device, max_batch_tokens * model->embedding_dim * sizeof(float), NULL, &context->residual_activation_buffer);
-    if (status != gptoss_status_success) {
-        goto cleanup;
-    }
-    status = gptoss_metal_buffer_create(&model->device, max_batch_tokens * model->embedding_dim * sizeof(float), NULL, &context->rmsnorm_activation_buffer);
-    if (status != gptoss_status_success) {
-        goto cleanup;
-    }
-    status = gptoss_metal_buffer_create(&model->device, max_batch_tokens * model->head_dim * (model->num_heads + 2 * model->num_kv_heads) * sizeof(float), NULL, &context->qkv_activation_buffer);
-    if (status != gptoss_status_success) {
-        goto cleanup;
-    }
-    status = gptoss_metal_buffer_create(&model->device, max_batch_tokens * model->head_dim * model->num_heads * sizeof(float), NULL, &context->sdpa_activation_buffer);
-    if (status != gptoss_status_success) {
-        goto cleanup;
-    }
-    status = gptoss_metal_buffer_create(&model->device, max_batch_tokens * model->num_experts * sizeof(float), NULL, &context->gate_activation_buffer);
-    if (status != gptoss_status_success) {
-        goto cleanup;
-    }
-    status = gptoss_metal_buffer_create(&model->device, max_batch_tokens * model->num_experts * sizeof(struct gptoss_expert_prediction), NULL, &context->expert_activation_buffer);
-    if (status != gptoss_status_success) {
-        goto cleanup;
-    }
-    // The last entry will hold the total number of tokens.
-    status = gptoss_metal_buffer_create(&model->device, (1 + model->num_experts) * sizeof(uint32_t), NULL, &context->expert_offset_buffer);
-    if (status != gptoss_status_success) {
-        goto cleanup;
-    }
-    status = gptoss_metal_buffer_create(&model->device, max_batch_tokens * model->num_active_experts * sizeof(uint32_t), NULL, &context->token_to_expert_routing_buffer);
-    if (status != gptoss_status_success) {
-        goto cleanup;
-    }
-    status = gptoss_metal_buffer_create(&model->device, max_batch_tokens * model->num_active_experts * model->embedding_dim * sizeof(float), NULL, &context->swiglu_input_buffer);
-    if (status != gptoss_status_success) {
-        goto cleanup;
-    }
-    status = gptoss_metal_buffer_create(&model->device, max_batch_tokens * model->num_active_experts * model->mlp_dim * sizeof(float), NULL, &context->swiglu_activation_buffer);
-    if (status != gptoss_status_success) {
-        goto cleanup;
-    }
-    status = gptoss_metal_buffer_create(&model->device, max_batch_tokens * model->num_active_experts * model->embedding_dim * sizeof(float), NULL, &context->moe_activation_buffer);
-    if (status != gptoss_status_success) {
-        goto cleanup;
-    }
-
-    // Input/output buffers
-    status = gptoss_metal_buffer_create(&model->device, sizeof(struct gptoss_control), NULL, &context->control_buffer);
-    if (status != gptoss_status_success) {
-        goto cleanup;
-    }
-    status = gptoss_metal_buffer_create(&model->device, context_length * sizeof(uint32_t), NULL, &context->token_buffer);
-    if (status != gptoss_status_success) {
-        goto cleanup;
-    }
-    status = gptoss_metal_buffer_create(&model->device, max_batch_tokens * model->vocabulary_size * sizeof(float), NULL, &context->score_buffer);
-    if (status != gptoss_status_success) {
-        goto cleanup;
-    }
-    status = gptoss_metal_buffer_create(&model->device, max_batch_tokens * model->vocabulary_size * sizeof(float), NULL, &context->prob_buffer);
-    if (status != gptoss_status_success) {
-        goto cleanup;
-    }
-    status = gptoss_metal_buffer_create(&model->device, max_batch_tokens * model->max_threadgroups * sizeof(float), NULL, &context->sum_buffer);
-    if (status != gptoss_status_success) {
-        goto cleanup;
-    }
-    status = gptoss_metal_buffer_create(&model->device, max_batch_tokens * sizeof(uint64_t), NULL, &context->argmax_buffer);
-    if (status != gptoss_status_success) {
-        goto cleanup;
-    }
-    status = gptoss_metal_buffer_create(&model->device, model->num_blocks * context_length * 2 * model->num_kv_heads * model->head_dim * sizeof(float), NULL, &context->kvcache_buffer);
-    if (status != gptoss_status_success) {
-        goto cleanup;
-    }
-
-    context->kvcache_size = context->kvcache_buffer.size;
-    context->allocation_size = 
-        context->residual_activation_buffer.size + context->rmsnorm_activation_buffer.size +
-        context->qkv_activation_buffer.size + context->sdpa_activation_buffer.size +
-        context->gate_activation_buffer.size + context->expert_activation_buffer.size +
-        context->expert_offset_buffer.size + context->token_to_expert_routing_buffer.size + context->swiglu_input_buffer.size +
-        context->swiglu_activation_buffer.size + context->moe_activation_buffer.size +
-        context->token_buffer.size + context->kvcache_buffer.size + context->score_buffer.size + context->argmax_buffer.size;
-
-    context->model = model;
-    gptoss_model_retain(model);
-    *context_out = context;
-    context = NULL;
-
-cleanup:
-    gptoss_context_release(context);
-    return status;
-}
-
-enum gptoss_status GPTOSS_ABI gptoss_context_get_num_tokens(
-    gptoss_context_t context,
-    size_t* num_tokens_out)
-{
-    *num_tokens_out = context->num_tokens;
-    return gptoss_status_success;
-}
-
-enum gptoss_status GPTOSS_ABI gptoss_context_get_max_tokens(
-    gptoss_context_t context,
-    size_t* max_tokens_out)
-{
-    *max_tokens_out = context->max_tokens;
-    return gptoss_status_success;
-}
-
-enum gptoss_status GPTOSS_ABI gptoss_context_get_tokens(
-    gptoss_context_t context,
-    uint32_t* tokens_out,
-    size_t max_tokens,
-    size_t* num_tokens_out)
-{
-    *num_tokens_out = context->num_tokens;
-    if (max_tokens < context->num_tokens) {
-        return gptoss_status_insufficient_memory;
-    }
-
-    if (context->num_tokens != 0) {
-        memcpy(tokens_out, context->token_buffer.ptr, context->num_tokens * sizeof(uint32_t));
-    }
-    return gptoss_status_success;
-}
-
-// Prefill: input_tokens_offset = number of tokens in KV cache, num_input_tokens > 0, num_output_tokens = 0.
-// Sampling: input_tokens_offset = number of tokens in the context - 1, num_input_tokens = 1, num_output_tokens = 1.
-// Perplexity: input_tokens_offset = 0, num_input_tokens > 1, num_output_tokens = num_input_tokens.
-static enum gptoss_status process_tokens(
-    gptoss_context_t context,
-    struct gptoss_metal_command_buffer* command_buffer,
-    size_t input_tokens_offset,
-    size_t num_input_tokens,
-    size_t num_output_tokens)
-{
-    assert(num_input_tokens != 0);
-    assert(num_input_tokens <= context->max_batch_tokens);
-    assert(num_output_tokens <= context->max_batch_tokens);
-    assert(num_input_tokens >= num_output_tokens);
-    const size_t dense_matmul_kernel_token_multiple_constraint = 64;
-    const size_t min_tokens_for_dense_moe_kernels = 64;
-
-    enum gptoss_status status = gptoss_status_success;
-    const struct gptoss_model* model = context->model;
-
-    const size_t attn_qkv_dim = model->head_dim * (model->num_heads + 2 * model->num_kv_heads);
-
-    const size_t input_tokens_end = input_tokens_offset + num_input_tokens;
-    for (size_t input_batch_start = input_tokens_offset;
-        input_batch_start < input_tokens_end;
-        input_batch_start += context->max_batch_tokens)
-    {
-        const size_t input_batch_size = math_min(context->max_batch_tokens, input_tokens_end - input_batch_start);
-        const size_t input_batch_end = input_batch_start + input_batch_size;
-        const size_t output_batch_size = math_sub_sat(num_output_tokens, input_tokens_end - input_batch_end);
-
-        status = gptoss_metal_command_buffer_encode_launch_bf16_f32_embeddings(
-            command_buffer,
-            &model->bf16_f32_embeddings_fn,
-            model->embeddings_threadgroup_size,
-            &context->token_buffer,
-            input_batch_start * sizeof(uint32_t),
-            &model->shared_weight_buffer,
-            /*weight_offset=*/0,
-            &context->residual_activation_buffer,
-            /*output_offset=*/0,
-            &context->control_buffer,
-            /*control_offset=*/0,
-            /*num_tokens=*/input_batch_size,
-            /*num_channels=*/model->embedding_dim);
-        if (status != gptoss_status_success) {
-            GPTOSS_LOG_ERROR("failed to encode bf16_f32_embeddings kernel launch");
-            return status;
-        }
-        for (uint32_t n = 0; n < model->num_blocks; n++) {
-            const bool last_block = n + 1 == model->num_blocks;
-            const size_t num_block_output_tokens = last_block ? output_batch_size : input_batch_size;
-
-            status = gptoss_metal_command_buffer_encode_launch_f32_bf16w_rmsnorm(
-                command_buffer,
-                &model->f32_bf16w_rmsnorm_fn,
-                &context->residual_activation_buffer,
-                /*input_offset=*/0,
-                &model->shared_weight_buffer,
-                /*weight_offset=*/model->attn_rmsnorm_gain_offset + model->per_block_shared_weights_size * n,
-                &context->rmsnorm_activation_buffer,
-                /*output_offset=*/0,
-                &context->control_buffer,
-                /*control_offset=*/0,
-                /*num_tokens=*/input_batch_size,
-                /*num_channels=*/model->embedding_dim,
-                model->rmsnorm_epsilon);
-            if (status != gptoss_status_success) {
-                GPTOSS_LOG_ERROR("failed to encode f32_bf16w_rmsnorm kernel launch");
-                return status;
-            }
-
-            if (input_batch_size % dense_matmul_kernel_token_multiple_constraint == 0) {
-                status = gptoss_metal_command_buffer_encode_launch_f32_bf16w_dense_matmul_qkv(
-                    command_buffer,
-                    &model->f32_bf16w_dense_matmul_qkv_fn,
-                    &context->rmsnorm_activation_buffer,
-                    /*input_offset=*/0,
-                    &model->shared_weight_buffer,
-                    /*weight_offset=*/model->attn_qkv_weight_offset + model->per_block_shared_weights_size * n,
-                    &model->shared_weight_buffer,
-                    /*bias_offset=*/model->attn_qkv_bias_offset + model->per_block_shared_weights_size * n,
-                    &context->qkv_activation_buffer,
-                    /*output_offset=*/0,
-                    &context->control_buffer,
-                    /*control_offset=*/0,
-                    /*num_tokens=*/input_batch_size,
-                    /*num_cols=*/model->embedding_dim,
-                    /*num_rows=*/attn_qkv_dim);
-                if (status != gptoss_status_success) {
-                    GPTOSS_LOG_ERROR("failed to encode f32_bf16w_dense_matmul_qkv kernel launch");
-                    return status;
-                }
-
-                status = gptoss_metal_command_buffer_encode_launch_f32_rope(
-                    command_buffer,
-                    &model->f32_rope_fn,
-                    /*threadgroup_size=*/32,
-                    &context->qkv_activation_buffer,
-                    /*input_offset=*/0,
-                    &context->control_buffer,
-                    /*control_offset=*/0,
-                    model->rope_theta,
-                    model->interpolation_scale,
-                    model->yarn_offset,
-                    model->yarn_scale,
-                    model->yarn_multiplier,
-                    input_batch_size,
-                    model->num_heads,
-                    model->num_kv_heads,
-                    model->head_dim,
-                    /*token_offset=*/input_batch_start);
-                if (status != gptoss_status_success) {
-                    GPTOSS_LOG_ERROR("failed to encode f32_rope kernel launch");
-                    return status;
-                }
-
-                for (uint32_t t = 0; t < input_batch_size; t++) {
-                    for (uint32_t kv = 0; kv < 2; kv++) {
-                        for (uint32_t h = 0; h < model->num_kv_heads; h++) {
-                            status = gptoss_metal_command_buffer_encode_copy_buffer(
-                                command_buffer,
-                                &context->qkv_activation_buffer,
-                                /*input_offset=*/(t * attn_qkv_dim + (model->num_heads + kv * model->num_kv_heads + h) * model->head_dim) * sizeof(float),
-                                &context->kvcache_buffer,
-                                /*output_offset=*/(((n * model->num_kv_heads + h) * context->max_tokens + input_batch_start + t) * 2 + kv) * model->head_dim * sizeof(float),
-                                /*size=*/model->head_dim * sizeof(float));
-                            if (status != gptoss_status_success) {
-                                GPTOSS_LOG_ERROR("failed to encode copy of token %" PRIu32 " to KV cache", t);
-                                return status;
-                            }
-                        }
-                    }
-                }
-            } else {
-                status = gptoss_metal_command_buffer_encode_launch_f32_bf16w_matmul_qkv(
-                    command_buffer,
-                    &model->f32_bf16w_matmul_qkv_fn,
-                    model->attn_qkv_threadgroup_size,
-                    &context->rmsnorm_activation_buffer,
-                    /*input_offset=*/0,
-                    &model->shared_weight_buffer,
-                    /*weight_offset=*/model->attn_qkv_weight_offset + model->per_block_shared_weights_size * n,
-                    &model->shared_weight_buffer,
-                    /*bias_offset=*/model->attn_qkv_bias_offset + model->per_block_shared_weights_size * n,
-                    &context->qkv_activation_buffer,
-                    /*output_offset=*/0,
-                    &context->kvcache_buffer,
-                    /*kv_offset=*/n * model->num_kv_heads * context->max_tokens * 2 * model->head_dim * sizeof(float),
-                    &context->control_buffer,
-                    /*control_offset=*/0,
-                    /*num_tokens=*/input_batch_size,
-                    /*num_cols=*/model->embedding_dim,
-                    /*num_q_heads=*/model->num_heads,
-                    /*num_kv_heads=*/model->num_kv_heads,
-                    /*attn_head_dim=*/model->head_dim,
-                    /*token_offset=*/input_batch_start,
-                    /*max_tokens=*/context->max_tokens,
-                    /*rope_base=*/model->rope_theta,
-                    /*interpolation_scale=*/model->interpolation_scale,
-                    /*yarn_offset=*/model->yarn_offset,
-                    /*yarn_scale=*/model->yarn_scale,
-                    /*yarn_multiplier=*/model->yarn_multiplier);
-                if (status != gptoss_status_success) {
-                    GPTOSS_LOG_ERROR("failed to encode f32_bf16w_matmul_qkv kernel launch");
-                    return status;
-                }
-            }
-
-            if (num_block_output_tokens != 0) {
-                status = gptoss_metal_command_buffer_encode_launch_f32_sdpa(
-                    command_buffer,
-                    &model->f32_sdpa_q8_d64_fn,
-                    &context->qkv_activation_buffer,
-                    /*q_offset=*/attn_qkv_dim * (input_batch_size - num_block_output_tokens) * sizeof(float),
-                    &context->kvcache_buffer,
-                    /*kv_offset=*/n * model->num_kv_heads * context->max_tokens * 2 * model->head_dim * sizeof(float),
-                    &model->shared_weight_buffer,
-                    /*s_offset=*/model->attn_sdpa_sink_offset + model->per_block_shared_weights_size * n,
-                    &context->sdpa_activation_buffer,
-                    /*output_offset=*/0,
-                    &context->control_buffer,
-                    /*control_offset=*/0,
-                    /*window=*/n % 2 == 0 ? model->attention_window : UINT32_MAX,
-                    /*kv_stride=*/2 * context->max_tokens * model->head_dim,
-                    num_block_output_tokens,
-                    input_batch_start + input_batch_size - num_block_output_tokens,
-                    model->num_heads, model->num_kv_heads, model->head_dim);
-                if (status != gptoss_status_success) {
-                    GPTOSS_LOG_ERROR("failed to encode f32_sdpa kernel launch");
-                    return status;
-                }
-
-                if (input_batch_size % dense_matmul_kernel_token_multiple_constraint == 0) {
-                    status = gptoss_metal_command_buffer_encode_launch_f32_bf16w_dense_matmul_attn_output(
-                        command_buffer,
-                        &model->f32_bf16w_dense_matmul_attn_output_fn,
-                        &context->sdpa_activation_buffer,
-                        /*input_offset=*/0,
-                        &model->shared_weight_buffer,
-                        /*weight_offset=*/model->attn_out_weight_offset + model->per_block_shared_weights_size * n,
-                        &model->shared_weight_buffer,
-                        /*bias_offset=*/model->attn_out_bias_offset + model->per_block_shared_weights_size * n,
-                        &context->residual_activation_buffer,
-                        /*output_offset=*/model->embedding_dim * (input_batch_size - num_block_output_tokens) * sizeof(float),
-                        &context->control_buffer,
-                        /*control_offset=*/0,
-                        /*num_tokens=*/num_block_output_tokens,
-                        /*num_cols=*/model->num_heads * model->head_dim,
-                        /*num_rows=*/model->embedding_dim);
-                    if (status != gptoss_status_success) {
-                        GPTOSS_LOG_ERROR("failed to encode f32_bf16w_dense_matmul_attn_output kernel launch");
-                        return status;
-                    }
-                } else {
-                    status = gptoss_metal_command_buffer_encode_launch_f32_bf16w_matmul_add(
-                        command_buffer,
-                        &model->f32_bf16w_matmul_fn,
-                        model->attn_out_threadgroup_size,
-                        &context->sdpa_activation_buffer,
-                        /*input_offset=*/0,
-                        &model->shared_weight_buffer,
-                        /*weight_offset=*/model->attn_out_weight_offset + model->per_block_shared_weights_size * n,
-                        &model->shared_weight_buffer,
-                        /*bias_offset=*/model->attn_out_bias_offset + model->per_block_shared_weights_size * n,
-                        &context->residual_activation_buffer,
-                        /*output_offset=*/model->embedding_dim * (input_batch_size - num_block_output_tokens) * sizeof(float),
-                        &context->control_buffer,
-                        /*control_offset=*/0,
-                        /*num_tokens=*/num_block_output_tokens,
-                        /*num_cols=*/model->num_heads * model->head_dim,
-                        /*num_rows=*/model->embedding_dim);
-                    if (status != gptoss_status_success) {
-                        GPTOSS_LOG_ERROR("failed to encode f32_bf16w_matmul_add kernel launch");
-                        return status;
-                    }
-                }
-                status = gptoss_metal_command_buffer_encode_launch_f32_bf16w_rmsnorm(
-                    command_buffer,
-                    &model->f32_bf16w_rmsnorm_fn,
-                    &context->residual_activation_buffer,
-                    /*input_offset=*/model->embedding_dim * (input_batch_size - num_block_output_tokens) * sizeof(float),
-                    &model->shared_weight_buffer,
-                    /*weight_offset=*/model->mlp_rmsnorm_gain_offset + model->per_block_shared_weights_size * n,
-                    &context->rmsnorm_activation_buffer,
-                    /*output_offset=*/0,
-                    &context->control_buffer,
-                    /*control_offset=*/0,
-                    num_block_output_tokens,
-                    model->embedding_dim,
-                    model->rmsnorm_epsilon);
-                if (status != gptoss_status_success) {
-                    GPTOSS_LOG_ERROR("failed to encode f32_bf16w_rmsnorm kernel launch");
-                    return status;
-                }
-                if (input_batch_size % dense_matmul_kernel_token_multiple_constraint == 0) {
-                    status = gptoss_metal_command_buffer_encode_launch_f32_bf16w_dense_matmul_mlp_gate(
-                        command_buffer,
-                        &model->f32_bf16w_dense_matmul_mlp_gate_fn,
-                        &context->rmsnorm_activation_buffer,
-                        /*input_offset=*/0,
-                        &model->shared_weight_buffer,
-                        /*weight_offset=*/model->mlp_gate_weight_offset + model->per_block_shared_weights_size * n,
-                        &model->shared_weight_buffer,
-                        /*bias_offset=*/model->mlp_gate_bias_offset + model->per_block_shared_weights_size * n,
-                        &context->gate_activation_buffer,
-                        /*output_offset=*/0,
-                        &context->control_buffer,
-                        /*control_offset=*/0,
-                        num_block_output_tokens,
-                        model->embedding_dim,
-                        model->num_experts);
-                    if (status != gptoss_status_success) {
-                        GPTOSS_LOG_ERROR("failed to encode f32_bf16w_dense_matmul_mlp_gate kernel launch");
-                        return status;
-                    }
-                } else {
-                    status = gptoss_metal_command_buffer_encode_launch_f32_bf16w_matmul(
-                        command_buffer,
-                        &model->f32_bf16w_matmul_fn,
-                        model->mlp_gate_threadgroup_size,
-                        &context->rmsnorm_activation_buffer,
-                        /*input_offset=*/0,
-                        &model->shared_weight_buffer,
-                        /*weight_offset=*/model->mlp_gate_weight_offset + model->per_block_shared_weights_size * n,
-                        &model->shared_weight_buffer,
-                        /*bias_offset=*/model->mlp_gate_bias_offset + model->per_block_shared_weights_size * n,
-                        &context->gate_activation_buffer,
-                        /*output_offset=*/0,
-                        &context->control_buffer,
-                        /*control_offset=*/0,
-                        /*num_tokens=*/num_block_output_tokens,
-                        /*num_cols=*/model->embedding_dim,
-                        /*num_rows=*/model->num_experts);
-                    if (status != gptoss_status_success) {
-                        GPTOSS_LOG_ERROR("failed to encode f32_bf16w_matmul kernel launch");
-                        return status;
-                    }
-                }
-
-                const char* kernel_name = NULL;
-                switch (model->num_experts) {
-                    case 32:
-                        kernel_name = "f32_topk_softmax_e32_k4_fn";
-                        status = gptoss_metal_command_buffer_encode_launch_f32_topk(
-                            command_buffer,
-                            &model->f32_topk_softmax_e32_k4_fn,
-                            &context->gate_activation_buffer, /*input_offset=*/0,
-                            &context->expert_activation_buffer, /*output_offset=*/0,
-                            &context->control_buffer, /*control_offset=*/0,
-                            num_block_output_tokens,
-                            model->num_experts,
-                            model->num_active_experts);
-                        break;
-                    case 128:
-                        kernel_name = "f32_topk_softmax_e128_k4_fn";
-                        status = gptoss_metal_command_buffer_encode_launch_f32_topk(
-                            command_buffer,
-                            &model->f32_topk_softmax_e128_k4_fn,
-                            &context->gate_activation_buffer, /*input_offset=*/0,
-                            &context->expert_activation_buffer, /*output_offset=*/0,
-                            &context->control_buffer, /*control_offset=*/0,
-                            num_block_output_tokens,
-                            model->num_experts,
-                            model->num_active_experts);
-                        break;
-                    default:
-                        status = gptoss_status_unsupported_argument;
-                        GPTOSS_LOG_ERROR("missing Top-K kernel for %" PRIu32 " experts", model->num_experts);
-                        return status;
-                }
-                if (status != gptoss_status_success) {
-                    GPTOSS_LOG_ERROR("failed to encode %s kernel launch", kernel_name);
-                    return status;
-                }
-
-                // If we have enough tokens in prefill, we will pick the prefill-optimized kernels.
-                if (num_block_output_tokens >= min_tokens_for_dense_moe_kernels) {
-                    status = gptoss_metal_command_buffer_encode_launch_expert_routing_metadata(
-                        command_buffer,
-                        &model->f32_expert_routing_metadata_fn,
-                        &context->expert_activation_buffer,
-                        /*expert_predictions_offset=*/0,
-                        &context->expert_offset_buffer,
-                        /*expert_offsets_offset=*/0,
-                        &context->token_to_expert_routing_buffer,
-                        /*intra_expert_offsets_offset=*/0,
-                        num_block_output_tokens * model->num_active_experts,
-                        model->num_experts);
-                    if (status != gptoss_status_success) {
-                        GPTOSS_LOG_ERROR("failed to encode f32_expert_routing_metadata kernel launch");
-                        return status;
-                    }
-                    status = gptoss_metal_command_buffer_encode_launch_f32_scatter(
-                        command_buffer,
-                        &model->f32_scatter_e4_fn,
-                        &context->rmsnorm_activation_buffer,
-                        /*input_offset=*/0,
-                        &context->expert_activation_buffer,
-                        /*expert_predictions_offset=*/0,
-                        &context->expert_offset_buffer,
-                        /*expert_offsets_offset=*/0,
-                        &context->token_to_expert_routing_buffer,
-                        /*intra_expert_offsets_offset=*/0,
-                        &context->swiglu_input_buffer,
-                        /*output_offset=*/0,
-                        model->embedding_dim,
-                        num_block_output_tokens,
-                        model->num_active_experts);
-                    if (status != gptoss_status_success) {
-                        GPTOSS_LOG_ERROR("failed to encode f32_scatter kernel launch");
-                        return status;
-                    } 
-                    // Dense MoE SwiGLU matmul.
-                    status = gptoss_metal_command_buffer_encode_launch_f32_mf4w_moe_dense_matmul_swiglu(
-                        command_buffer,
-                        &model->f32_mf4w_moe_dense_matmul_swiglu_fn,
-                        &context->expert_offset_buffer,
-                        /*expert_offsets_offset=*/0,
-                        &context->swiglu_input_buffer,
-                        /*input_offset=*/0,
-                        &model->block_weight_buffers[n],
-                        /*weight_block_offset=*/0,
-                        &model->block_weight_buffers[n],
-                        /*weight_scale_offset=*/model->mlp_swiglu_scale_offset,
-                        &model->block_weight_buffers[n],
-                        /*bias_offset=*/model->mlp_swiglu_bias_offset,
-                        &context->swiglu_activation_buffer,
-                        /*output_offset=*/0,
-                        model->swiglu_limit,
-                        /*expert_stride_bytes=*/model->per_expert_block_weight_size,
-                        num_block_output_tokens,
-                        model->num_experts,
-                        model->embedding_dim,
-                        2 * model->mlp_dim);
-                    if (status != gptoss_status_success) {
-                        GPTOSS_LOG_ERROR("failed to encode f32_mf4w_moe_dense_matmul_swiglu kernel launch");
-                        return status;
-                    }
-
-                    // Dense MoE proj matmul.
-                    status = gptoss_metal_command_buffer_encode_launch_f32_mf4w_moe_dense_matmul(
-                        command_buffer,
-                        &model->f32_mf4w_moe_dense_matmul_fn,
-                        &context->expert_offset_buffer,
-                        /*expert_offsets_offset=*/0,
-                        &context->swiglu_activation_buffer,
-                        /*input_offset=*/0,
-                        &model->block_weight_buffers[n],
-                        /*weight_block_offset=*/model->mlp_out_block_offset,
-                        &model->block_weight_buffers[n],
-                        /*weight_scale_offset=*/model->mlp_out_scale_offset,
-                        &model->block_weight_buffers[n],
-                        /*bias_offset=*/model->mlp_out_bias_offset,
-                        &context->moe_activation_buffer,
-                        /*output_offset=*/0,
-                        /*expert_stride_bytes=*/model->per_expert_block_weight_size,
-                        num_block_output_tokens,
-                        model->num_experts,
-                        model->mlp_dim,
-                        model->embedding_dim);
-                    if (status != gptoss_status_success) {
-                        GPTOSS_LOG_ERROR("failed to encode f32_mf4w_moe_dense_matmul_swiglu kernel launch");
-                        return status;
-                    }
-                    // Gather and accumulate.
-                    status = gptoss_metal_command_buffer_encode_launch_f32_gather_and_accumulate_e4(
-                        command_buffer,
-                        &model->f32_gather_and_accumulate_e4_fn,
-                        &context->moe_activation_buffer,
-                        /*input_offset=*/0,
-                        &context->expert_activation_buffer,
-                        /*expert_predictions_offset=*/0,
-                        &context->expert_offset_buffer,
-                        /*expert_offsets_offset=*/0,
-                        &context->token_to_expert_routing_buffer,
-                        /*intra_expert_offsets_offset=*/0,
-                        &context->residual_activation_buffer, 
-                        /*output_offset=*/model->embedding_dim * (input_batch_size - num_block_output_tokens) * sizeof(float),
-                        model->embedding_dim,
-                        num_block_output_tokens,
-                        model->num_active_experts);
-                    if (status != gptoss_status_success) {
-                        GPTOSS_LOG_ERROR("failed to encode f32_gather_and_accumulate_e4 kernel launch");
-                        return status;
-                    }
-
-                } else {
-                    status = gptoss_metal_command_buffer_encode_launch_f32_mf4w_moe_matmul_swiglu(
-                        command_buffer,
-                        &model->f32_mf4w_moe_matmul_swiglu_fn,
-                        model->mlp_swiglu_threadgroup_size,
-                        &context->rmsnorm_activation_buffer,
-                        /*input_offset=*/0,
-                        &context->expert_activation_buffer,
-                        /*expert_offset=*/0,
-                        &model->block_weight_buffers[n],
-                        /*weight_block_offset=*/0,
-                        &model->block_weight_buffers[n],
-                        /*weight_scale_offset=*/model->mlp_swiglu_scale_offset,
-                        &model->block_weight_buffers[n],
-                        /*bias_offset=*/model->mlp_swiglu_bias_offset,
-                        &context->swiglu_activation_buffer,
-                        /*output_offset=*/0,
-                        &context->control_buffer,
-                        /*control_offset=*/0,
-                        model->swiglu_limit,
-                        model->per_expert_block_weight_size,
-                        num_block_output_tokens,
-                        model->num_active_experts,
-                        model->embedding_dim,
-                        model->mlp_dim);
-                    if (status != gptoss_status_success) {
-                        GPTOSS_LOG_ERROR("failed to encode f32_mf4w_moe_matmul_swiglu kernel launch");
-                        return status;
-                    }
-
-                    status = gptoss_metal_command_buffer_encode_launch_f32_mf4w_moe_matmul(
-                        command_buffer,
-                        &model->f32_mf4w_moe_matmul_fn,
-                        model->mlp_out_threadgroup_size,
-                        &context->swiglu_activation_buffer,
-                        /*input_offset=*/0,
-                        &context->expert_activation_buffer,
-                        /*expert_offset=*/0,
-                        &model->block_weight_buffers[n],
-                        /*weight_block_offset=*/model->mlp_out_block_offset,
-                        &model->block_weight_buffers[n],
-                        /*weight_scale_offset=*/model->mlp_out_scale_offset,
-                        &model->block_weight_buffers[n],
-                        /*bias_offset=*/model->mlp_out_bias_offset,
-                        &context->moe_activation_buffer,
-                        /*output_offset=*/0,
-                        &context->control_buffer,
-                        /*control_offset=*/0,
-                        model->per_expert_block_weight_size,
-                        num_block_output_tokens,
-                        model->num_active_experts,
-                        model->mlp_dim,
-                        model->embedding_dim);
-                    if (status != gptoss_status_success) {
-                        GPTOSS_LOG_ERROR("failed to encode f32_mf4w_moe_matmul kernel launch");
-                        return status;
-                    }
-
-                    status = gptoss_metal_command_buffer_encode_launch_f32_accumulate(
-                        command_buffer,
-                        &model->f32_accumulate_e4_fn,
-                        model->mlp_acc_threadgroup_size,
-                        model->max_threadgroups,
-                        &context->moe_activation_buffer,
-                        /*input_offset=*/0,
-                        &context->expert_activation_buffer,
-                        /*expert_offset=*/0,
-                        &context->residual_activation_buffer,
-                        /*output_offset=*/model->embedding_dim * (input_batch_size - num_block_output_tokens) * sizeof(float),
-                        &context->control_buffer,
-                        /*control_offset=*/0,
-                        model->embedding_dim,
-                        num_block_output_tokens,
-                        model->num_active_experts);
-                    if (status != gptoss_status_success) {
-                        GPTOSS_LOG_ERROR("failed to encode f32_accumulate kernel launch");
-                        return status;
-                    }
-                }
-            }
-        }
-
-        if (output_batch_size != 0) {
-            status = gptoss_metal_command_buffer_encode_launch_f32_bf16w_rmsnorm(
-                command_buffer,
-                &model->f32_bf16w_rmsnorm_fn,
-                &context->residual_activation_buffer,
-                /*input_offset=*/model->embedding_dim * (input_batch_size - output_batch_size) * sizeof(float),
-                &model->shared_weight_buffer,
-                /*weight_offset=*/model->rmsnorm_weight_offset,
-                &context->rmsnorm_activation_buffer,
-                /*output_offset=*/0,
-                &context->control_buffer,
-                /*control_offset=*/0,
-                /*num_tokens=*/output_batch_size,
-                /*num_channels=*/model->embedding_dim,
-                model->rmsnorm_epsilon);
-            if (status != gptoss_status_success) {
-                GPTOSS_LOG_ERROR("failed to encode f32_bf16w_rmsnorm kernel launch");
-                return status;
-            }
-
-            status = gptoss_metal_command_buffer_encode_fill_buffer(
-                command_buffer,
-                &context->argmax_buffer,
-                /*offset=*/0,
-                /*size=*/sizeof(uint64_t) * output_batch_size,
-                /*fill_value=*/0xFF);
-            if (status != gptoss_status_success) {
-                GPTOSS_LOG_ERROR("failed to encode fill buffer command");
-                return status;
-            }
-
-            status = gptoss_metal_command_buffer_encode_launch_f32_bf16w_unembedding(
-                command_buffer,
-                &model->f32_bf16w_unembedding_fn,
-                model->unembedding_threadgroup_size,
-                model->max_threadgroups,
-                &context->rmsnorm_activation_buffer,
-                /*input_offset=*/0,
-                &model->shared_weight_buffer,
-                /*weight_offset=*/model->unembedding_weight_offset,
-                &context->score_buffer,
-                /*output_offset=*/0,
-                &context->argmax_buffer,
-                /*argmax_offset=*/0,
-                &context->control_buffer,
-                /*control_offset=*/0,
-                /*num_tokens=*/output_batch_size,
-                /*num_cols=*/model->embedding_dim,
-                /*num_rows=*/model->vocabulary_size);
-            if (status != gptoss_status_success) {
-                GPTOSS_LOG_ERROR("failed to encode f32_bf16w_unembedding kernel launch");
-                return status;
-            }
-        }
-    }
-    return gptoss_status_success;
-}
-
-enum gptoss_status GPTOSS_ABI gptoss_context_append_chars(
-    gptoss_context_t context,
-    const char* text,
-    size_t text_length,
-    size_t* num_tokens_out)
-{
-    enum gptoss_status status = gptoss_status_success;
-    const struct gptoss_model* model = context->model;
-    const struct gptoss_tokenizer* tokenizer = model->tokenizer;
-    size_t num_appended_tokens = 0;
-    while (text_length != 0) {
-        if (context->num_tokens == context->max_tokens) {
-            status = gptoss_status_context_overflow;
-            break;
-        }
-        const char* tokens = tokenizer->tokens_ptr;
-        uint32_t best_token = UINT32_MAX;
-        uint32_t best_token_length = 0;
-        for (size_t t = 0; t < tokenizer->num_text_tokens; t++) {
-            uint16_t token_length;
-            memcpy(&token_length, tokens, sizeof(uint16_t));
-            tokens += sizeof(uint16_t);
-            if (token_length <= text_length && token_length > best_token_length) {
-                if (memcmp(text, tokens, token_length) == 0) {
-                    if (token_length > best_token_length) {
-                        best_token = (uint32_t) t;
-                        best_token_length = token_length;
-                    }
-                }
-            }
-            tokens += token_length;
-        }
-
-        if (best_token == UINT32_MAX) {
-            GPTOSS_LOG_ERROR("failed to tokenize text \"%.*s\"", (int) text_length, text);
-            return gptoss_status_invalid_argument;
-        }
-
-        uint32_t* input_tokens = (uint32_t*) context->token_buffer.ptr;
-        if (context->num_kv_tokens > context->num_tokens) {
-            if (input_tokens[context->num_tokens] != best_token) {
-                input_tokens[context->num_tokens] = best_token;
-
-                // Invalidate the KV cache starting with the newly added token.
-                context->num_kv_tokens = context->num_tokens;
-            }
-            context->num_tokens++;
-        } else {
-            input_tokens[context->num_tokens++] = best_token;
-        }
-        num_appended_tokens++;
-        text += best_token_length;
-        text_length -= best_token_length;
-    }
-    if (num_tokens_out != NULL) {
-        *num_tokens_out = num_appended_tokens;
-    }
-    return status;
-}
-
-enum gptoss_status GPTOSS_ABI gptoss_context_append_tokens(
-    gptoss_context_t context,
-    size_t num_tokens,
-    const uint32_t* tokens)
-{
-    const struct gptoss_model* model = context->model;
-
-    // Validate all tokens
-    for (size_t t = 0; t < num_tokens; t++) {
-        const uint32_t token = tokens[t];
-        if (token >= model->vocabulary_size) {
-            GPTOSS_LOG_ERROR("token %" PRIu32 " at index %zu is out of bounds for vocabulary size %" PRIu32,
-                token, t, context->model->vocabulary_size);
-            return gptoss_status_invalid_argument;
-        }
-    }
-
-    enum gptoss_status status = gptoss_status_success;
-    uint32_t* input_tokens = (uint32_t*) context->token_buffer.ptr;
-    while (num_tokens != 0) {
-        if (context->num_tokens == context->max_tokens) {
-            status = gptoss_status_context_overflow;
-            break;
-        }
-
-        if (context->num_kv_tokens > context->num_tokens) {
-            const size_t num_tokens_to_verify = math_min(context->num_kv_tokens - context->num_tokens, num_tokens);
-            size_t num_verified_tokens = 0;
-            for (; num_verified_tokens < num_tokens_to_verify; num_verified_tokens++) {
-                if (input_tokens[context->num_tokens + num_verified_tokens] != tokens[num_verified_tokens]) {
-                    // Invalidate the KV cache starting with the newly added tokens.
-                    context->num_kv_tokens = context->num_tokens + num_verified_tokens;
-                    break;
-                }
-            }
-
-            context->num_tokens += num_verified_tokens;
-            tokens += num_verified_tokens;
-            num_tokens -= num_verified_tokens;
-        } else {
-            const size_t num_tokens_to_copy = math_min(context->max_tokens - context->num_tokens, num_tokens);
-            memcpy(input_tokens + context->num_tokens, tokens, num_tokens_to_copy * sizeof(uint32_t));
-            context->num_tokens += num_tokens_to_copy;
-            tokens += num_tokens_to_copy;
-            num_tokens -= num_tokens_to_copy;
-        }
-    }
-
-    return status;
-}
-
-enum gptoss_status GPTOSS_ABI gptoss_context_process(
-    gptoss_context_t context)
-{
-    if (context->num_tokens > context->num_kv_tokens) {
-        struct gptoss_metal_command_buffer command_buffer = {0};
-
-        enum gptoss_status status = gptoss_metal_command_buffer_create(&context->model->command_queue, &command_buffer);
-        if (status != gptoss_status_success) {
-            goto cleanup;
-        }
-
-        struct gptoss_control* control = (struct gptoss_control*) context->control_buffer.ptr;
-        control->abort = 0;
-
-        status = process_tokens(
-            context,
-            &command_buffer,
-            /*input_tokens_offset=*/context->num_kv_tokens,
-            /*num_input_tokens=*/context->num_tokens - context->num_kv_tokens,
-            /*num_output_tokens=*/0);
-        if (status != gptoss_status_success) {
-            goto cleanup;
-        }
-
-        status = gptoss_metal_command_buffer_commit(&command_buffer);
-        if (status != gptoss_status_success) {
-            goto cleanup;
-        }
-
-        status = gptoss_metal_command_buffer_wait_completion(&command_buffer, NULL);
-        if (status != gptoss_status_success) {
-            goto cleanup;
-        }
-
-        context->num_kv_tokens = context->num_tokens;
-
-cleanup:
-        gptoss_metal_command_buffer_release(&command_buffer);
-        return status;
-    }
-    
-    return gptoss_status_success;
-}
-
-enum gptoss_status GPTOSS_ABI gptoss_context_sample(
-    gptoss_context_t context,
-    float temperature,
-    uint64_t seed,
-    size_t max_tokens,
-    uint32_t* tokens_out,
-    size_t* num_tokens_out)
-{
-    enum gptoss_status status = gptoss_status_success;
-    const struct gptoss_model* model = context->model;
-    struct gptoss_metal_command_buffer command_buffer = {0};
-
-    *num_tokens_out = 0;
-
-    const uint32_t num_original_tokens = context->num_tokens;
-
-    status = gptoss_metal_command_buffer_create(&context->model->command_queue, &command_buffer);
-    if (status != gptoss_status_success) {
-        goto cleanup;
-    }
-
-    struct gptoss_control* control = (struct gptoss_control*) context->control_buffer.ptr;
-    control->abort = 0;
-
-    for (size_t t = 0; t < max_tokens; t++) {
-        if (context->num_kv_tokens < context->num_tokens) {
-            status = process_tokens(
-                context,
-                &command_buffer,
-                /*input_tokens_offset=*/context->num_kv_tokens,
-                /*num_input_tokens=*/context->num_tokens - context->num_kv_tokens,
-                /*num_output_tokens=*/1);
-            context->num_kv_tokens = context->num_tokens;
-        } else {
-            status = process_tokens(
-                context,
-                &command_buffer,
-                /*input_tokens_offset=*/context->num_tokens - 1,
-                /*num_input_tokens=*/1,
-                /*num_output_tokens=*/1);
-        }
-        if (status != gptoss_status_success) {
-            goto cleanup;
-        }
-
-        if (temperature != 0.0f) {
-            assert(context->num_processed_tokens != 0);
-            uint32_t num_threadgroups = 0;
-            uint32_t num_dims_per_threadgroup = 0;
-            status = gptoss_metal_command_buffer_encode_launch_f32_softmax(
-                &command_buffer,
-                &model->f32_softmax_fn,
-                /*threadgroup_size=*/512,
-                model->max_threadgroups,
-                &context->score_buffer,
-                /*score_offset=*/0,
-                &context->argmax_buffer,
-                /*argmax_offset=*/0,
-                &context->prob_buffer,
-                /*prob_offset=*/0,
-                &context->sum_buffer,
-                /*sum_offset=*/0,
-                &context->control_buffer,
-                /*control_offset=*/0,
-                model->vocabulary_size,
-                /*num_tokens=*/1,
-                temperature,
-                &num_threadgroups,
-                &num_dims_per_threadgroup);
-            if (status != gptoss_status_success) {
-                GPTOSS_LOG_ERROR("failed to encode f32_softmax kernel launch");
-                goto cleanup;
-            }
-
-            status = gptoss_metal_command_buffer_encode_launch_f32_sample(
-                &command_buffer,
-                &model->f32_sample_fn,
-                /*min_threadgroup_size=*/512,
-                &context->prob_buffer,
-                /*prob_offset=*/0,
-                &context->sum_buffer,
-                /*sum_offset=*/0,
-                &context->token_buffer,
-                /*token_offset=*/context->num_tokens * sizeof(uint32_t),
-                &context->control_buffer,
-                /*control_offset=*/0,
-                /*rng_seed=*/seed + UINT64_C(0x123456789ABCDEF),
-                /*rng_offset=*/context->num_tokens,
-                /*num_blocks=*/num_threadgroups,
-                /*num_channels=*/model->vocabulary_size,
-                /*num_channels_per_block=*/num_dims_per_threadgroup);
-            if (status != gptoss_status_success) {
-                GPTOSS_LOG_ERROR("failed to encode f32_sample kernel launch");
-                goto cleanup;
-            }
-        } else {
-            status = gptoss_metal_command_buffer_encode_copy_buffer(
-                &command_buffer,
-                &context->argmax_buffer,
-                /*input_offset=*/0,
-                &context->token_buffer,
-                /*output_offset=*/context->num_tokens * sizeof(uint32_t),
-                /*size=*/sizeof(uint32_t));
-            if (status != gptoss_status_success) {
-                GPTOSS_LOG_ERROR("failed to encode copy buffer");
-                goto cleanup;
-            }
-        }
-        context->num_tokens += 1;
-        context->num_kv_tokens = context->num_tokens;
-    }
-
-    gptoss_metal_command_buffer_commit(&command_buffer);
-    gptoss_metal_command_buffer_wait_completion(&command_buffer, NULL);
-
-    const uint32_t* token_ptr = (const uint32_t*) context->token_buffer.ptr;
-    const uint32_t num_generated_tokens = context->num_tokens - num_original_tokens;
-    memcpy(tokens_out, token_ptr + num_original_tokens, num_generated_tokens * sizeof(uint32_t));
-    *num_tokens_out = num_generated_tokens;
-
-cleanup:
-    gptoss_metal_command_buffer_release(&command_buffer);
-    return status;
-}
-
-enum gptoss_status GPTOSS_ABI gptoss_context_reset(
-    gptoss_context_t context)
-{
-    context->num_tokens = 0;
-
-    // Note: context->num_kv_tokens is not reset and context->input_tokens_buffer is not cleared.
-    // If the subsequently added tokens match the tokens already in the KV cache, we reuse the KV cache.
-
-    return gptoss_status_success;
-}
-
-enum gptoss_status GPTOSS_ABI gptoss_context_retain(
-    gptoss_context_t context)
-{
-    atomic_fetch_add_explicit(&context->ref_count, 1, memory_order_relaxed);
-    return gptoss_status_success;
-}
-
-enum gptoss_status GPTOSS_ABI gptoss_context_release(
-    gptoss_context_t context)
-{
-    if (context != NULL) {
-        if (atomic_fetch_sub_explicit(&context->ref_count, 1, memory_order_acq_rel) == 1) {
-            // Activation buffers
-            gptoss_metal_buffer_release(&context->residual_activation_buffer);
-            gptoss_metal_buffer_release(&context->rmsnorm_activation_buffer);
-            gptoss_metal_buffer_release(&context->qkv_activation_buffer);
-            gptoss_metal_buffer_release(&context->sdpa_activation_buffer);
-            gptoss_metal_buffer_release(&context->gate_activation_buffer);
-            gptoss_metal_buffer_release(&context->expert_activation_buffer);
-            gptoss_metal_buffer_release(&context->swiglu_activation_buffer);
-            gptoss_metal_buffer_release(&context->moe_activation_buffer);
-            gptoss_metal_buffer_release(&context->expert_offset_buffer);
-            gptoss_metal_buffer_release(&context->token_to_expert_routing_buffer);
-            gptoss_metal_buffer_release(&context->swiglu_input_buffer);
-
-            // Input/output buffers
-            gptoss_metal_buffer_release(&context->control_buffer);
-            gptoss_metal_buffer_release(&context->token_buffer);
-            gptoss_metal_buffer_release(&context->score_buffer);
-            gptoss_metal_buffer_release(&context->prob_buffer);
-            gptoss_metal_buffer_release(&context->sum_buffer);
-            gptoss_metal_buffer_release(&context->argmax_buffer);
-            gptoss_metal_buffer_release(&context->kvcache_buffer);
-
-            gptoss_model_release(context->model);
-
-            memset(context, 0, sizeof(struct gptoss_context));
-            free(context);
-        }
-    }
-    return gptoss_status_success;
-}

gptoss_kernels/source/generate.c

@@ -1,317 +0,0 @@
-#include <assert.h>
-#include <inttypes.h>
-#include <math.h>
-#include <signal.h>
-#include <stdatomic.h>
-#include <stdbool.h>
-#include <stdio.h>
-#include <stdint.h>
-#include <stdlib.h>
-#include <string.h>
-
-#include <mach/mach_time.h>
-
-#include <gpt-oss.h>
-
-#include "internal/model.h"
-
-struct {
-    atomic_uint_least64_t inference_bytes;
-    atomic_size_t num_prefill_tokens;
-    atomic_uint_least64_t prefill_microseconds;
-    atomic_size_t num_generated_tokens;
-    atomic_uint_least64_t generation_microseconds;
-} globals = {
-    .inference_bytes = 0,
-    .num_prefill_tokens = 0,
-    .prefill_microseconds = 0,
-    .num_generated_tokens = 0,
-    .generation_microseconds = 0,
-};
-
-struct options {
-    const char* model;
-    const char* prompt;
-    size_t context_length;
-    size_t max_tokens;
-    float temperature;
-    bool verbose;
-};
-
-static inline double mach_timestamp_diff_to_seconds(uint64_t start_timestamp, uint64_t end_timestamp) {
-    static mach_timebase_info_data_t timebase_info = {0};
-    if (timebase_info.denom == 0) {
-        mach_timebase_info(&timebase_info);
-    }
-    const uint64_t elapsed_mach_time = end_timestamp - start_timestamp;
-    return ((double) elapsed_mach_time * (double) timebase_info.numer) / ((double) timebase_info.denom * 1.0e+9);
-}
-
-static inline uint64_t mach_timestamp_diff_to_microseconds(uint64_t start_timestamp, uint64_t end_timestamp) {
-    static mach_timebase_info_data_t timebase_info = {0};
-    if (timebase_info.denom == 0) {
-        mach_timebase_info(&timebase_info);
-    }
-    const uint64_t elapsed_mach_time = end_timestamp - start_timestamp;
-    const uint64_t denominator = timebase_info.denom * UINT64_C(1000);
-    return (elapsed_mach_time * timebase_info.numer + denominator / 2) / denominator;
-}
-
-static void print_usage(const char* program_name) {
-    printf("Usage: %s <model-path> [-p <prompt>] [-n <tokens>]\n", program_name);
-}
-
-struct options parse_options(int argc, char** argv) {
-    struct options options = (struct options) {
-        .model = NULL,
-        .prompt = NULL,
-        .context_length = 0,
-        .max_tokens = 0,
-        .temperature = 0.0f,
-        .verbose = false,
-    };
-    if (argc < 2) {
-        fprintf(stderr, "Error: missing required command-line argument\n");
-        print_usage(argv[0]);
-        exit(EXIT_FAILURE);
-    }
-    for (int i = 1; i < argc; i++) {
-        if (strcmp(argv[i], "--help") == 0) {
-            print_usage(argv[0]);
-            exit(EXIT_SUCCESS);
-        } else if (strcmp(argv[i], "-p") == 0 || strcmp(argv[i], "--prompt") == 0) {
-            if (i + 1 >= argc) {
-                fprintf(stderr, "Error: missing argument for %s\n", argv[i]);
-                print_usage(argv[0]);
-                exit(EXIT_FAILURE);
-            }
-            options.prompt = argv[++i];
-        } else if (strcmp(argv[i], "--context-length") == 0) {
-            if (i + 1 >= argc) {
-                fprintf(stderr, "Error: missing argument for --context-length\n");
-                print_usage(argv[0]);
-                exit(EXIT_FAILURE);
-            }
-            char* context_length_start = argv[++i];
-            char* context_length_end = context_length_start;
-            options.context_length = strtoul(context_length_start, &context_length_end, 10);
-            if (context_length_end == context_length_start || *context_length_end != 0) {
-                fprintf(stderr, "Error: failed to parse context length value \"%s\"\n", context_length_start);
-                exit(EXIT_FAILURE);
-            }
-        } else if (strcmp(argv[i], "-n") == 0 || strcmp(argv[i], "--max-tokens") == 0) {
-            if (i + 1 >= argc) {
-                fprintf(stderr, "Error: missing argument for %s\n", argv[i]);
-                print_usage(argv[0]);
-                exit(EXIT_FAILURE);
-            }
-            char* max_tokens_start = argv[++i];
-            char* max_tokens_end = max_tokens_start;
-            options.max_tokens = strtoul(max_tokens_start, &max_tokens_end, 10);
-            if (max_tokens_end == max_tokens_start || *max_tokens_end != 0) {
-                fprintf(stderr, "Error: failed to max tokens value \"%s\"\n", max_tokens_start);
-                exit(EXIT_FAILURE);
-            }
-            if (options.max_tokens == 0) {
-                fprintf(stderr, "Error: invalid max tokens value %zu\n", options.max_tokens);
-                exit(EXIT_FAILURE);
-            }
-        } else if (strcmp(argv[i], "-t") == 0 || strcmp(argv[i], "--temperature") == 0) {
-            if (i + 1 >= argc) {
-                fprintf(stderr, "Error: missing argument for %s\n", argv[i]);
-                print_usage(argv[0]);
-                exit(EXIT_FAILURE);
-            }
-            char* temperature_start = argv[++i];
-            char* temperature_end = temperature_start;
-            options.temperature = strtof(temperature_start, &temperature_end);
-            if (temperature_end == temperature_start || *temperature_end != 0) {
-                fprintf(stderr, "Error: failed to parse temperature value \"%s\"\n", temperature_start);
-                exit(EXIT_FAILURE);
-            }
-            if (signbit(options.temperature) != 0 || !(options.temperature <= 2.0f)) {
-                fprintf(stderr, "Error: invalid temperature value %f\n", options.temperature);
-                exit(EXIT_FAILURE);
-            }
-        } else if (strcmp(argv[i], "-v") == 0 || strcmp(argv[i], "--verbose") == 0) {
-            options.verbose = true;
-        } else {
-            if (options.model == NULL) {
-                options.model = argv[i];
-            } else {
-                fprintf(stderr, "Error: unexpected command-line argument %s\n", argv[i]);
-                print_usage(argv[0]);
-                exit(EXIT_FAILURE);
-            }
-        }
-    }
-    if (options.model == NULL) {
-        fprintf(stderr, "Error: missing required model argument\n");
-        print_usage(argv[0]);
-        exit(EXIT_FAILURE);
-    }
-    if (options.prompt == NULL) {
-        fprintf(stderr, "Error: missing required prompt argument\n");
-        print_usage(argv[0]);
-        exit(EXIT_FAILURE);
-    }
-    return options;
-}
-
-
-static void print_profile() {
-    const size_t num_prefill_tokens = atomic_load(&globals.num_prefill_tokens);
-    const uint64_t prefill_microseconds = atomic_load(&globals.prefill_microseconds);
-    const size_t num_generated_tokens = atomic_load(&globals.num_generated_tokens);
-    const uint64_t generation_microseconds = atomic_load(&globals.generation_microseconds);
-    const uint64_t inference_bytes = atomic_load(&globals.inference_bytes);
-    if (num_prefill_tokens != 0 || num_generated_tokens != 0) {
-        printf("\n");
-    }
-    if (num_prefill_tokens != 0) {
-        printf("Prefill speed (%zu tokens): %.1f tokens/second\n",
-            num_prefill_tokens,
-            (double) num_prefill_tokens / (double) prefill_microseconds * 1.0e+6);
-    }
-    if (num_generated_tokens != 0) {
-        printf("Generation speed (%zu tokens): %.1f tokens/second\n",
-            num_generated_tokens,
-            (double) num_generated_tokens / (double) generation_microseconds * 1.0e+6);
-    }
-}
-
-static void ctrl_c_handler(int signum) {
-    print_profile();
-    exit(EXIT_SUCCESS);
-}
-
-int main(int argc, char *argv[]) {
-    enum gptoss_status status;
-    gptoss_model_t model = NULL;
-    gptoss_tokenizer_t tokenizer = NULL;
-    gptoss_context_t context = NULL;
-
-    struct sigaction act;
-    act.sa_handler = ctrl_c_handler;
-    sigaction(SIGINT, &act, NULL);
-
-    setvbuf(stdout, NULL, _IONBF, 0);
-
-    struct options options = parse_options(argc, argv);
-
-    const uint64_t load_start_time = mach_continuous_time();
-    status = gptoss_model_create_from_file(options.model, &model);
-    if (status != gptoss_status_success) {
-        fprintf(stderr, "Error: failed to load model from file %s\n", options.model);
-        goto error;
-    }
-    size_t max_model_context_length = 0;
-    status = gptoss_model_get_max_context_length(model, &max_model_context_length);
-    if (status != gptoss_status_success) {
-        fprintf(stderr, "Error: failed to query maximum context length\n");
-        goto error;
-    }
-    assert(max_model_context_length != 0);
-    if (options.context_length == 0) {
-        options.context_length = max_model_context_length;
-    } else if (options.context_length > max_model_context_length) {
-        fprintf(stderr, "Error: context length %zu exceeds maximum context length %zu supported by the model\n", options.context_length, max_model_context_length);
-        goto error;
-    }
-
-    status = gptoss_model_get_tokenizer(model, &tokenizer);
-    if (status != gptoss_status_success) {
-        fprintf(stderr, "Error: failed to retrieve Tokenizer\n");
-        goto error;
-    }
-
-    uint32_t return_token_id = UINT32_MAX;
-    status = gptoss_tokenizer_get_special_token_id(tokenizer, gptoss_special_token_return, &return_token_id);
-    if (status != gptoss_status_success) {
-        fprintf(stderr, "Error: failed to query end-of-text token ID\n");
-        goto error;
-    }
-
-    status = gptoss_context_create(model, options.context_length, /*max_batch_tokens=*/0, &context);
-    if (status != gptoss_status_success) {
-        fprintf(stderr, "Error: failed to create Context object\n");
-        goto error;
-    }
-    if (options.verbose) {
-        printf("Model weights size: %.2lf MB\n", (double) model->weights_size * 0x1.0p-20);
-        printf("Model allocation size: %.2lf MB\n", (double) model->allocation_size * 0x1.0p-20);
-        printf("Context allocation size: %.2lf MB\n", (double) context->allocation_size * 0x1.0p-20);
-        printf("  Including KV cache: %.2lf MB\n", (double) context->kvcache_size * 0x1.0p-20);
-    }
-
-    const uint64_t load_end_time = mach_continuous_time();
-    const double load_elapsed_seconds = mach_timestamp_diff_to_seconds(load_start_time, load_end_time);
-    if (options.verbose) {
-        printf("Loaded model in %.3f seconds\n", load_elapsed_seconds);
-    }
-
-    const uint64_t prefill_start_time = mach_continuous_time();
-    size_t num_prefill_tokens = 0;
-    status = gptoss_context_append_chars(context, options.prompt, strlen(options.prompt), &num_prefill_tokens);
-    if (status != gptoss_status_success) {
-        fprintf(stderr, "Error: failed to tokenize prompt \"%s\"\n", options.prompt);
-        goto error;
-    }
-    atomic_store(&globals.num_prefill_tokens, num_prefill_tokens);
-    status = gptoss_context_process(context);
-    if (status != gptoss_status_success) {
-        fprintf(stderr, "Error: failed to process Context object\n");
-        goto error;
-    }
-    const uint64_t prefill_end_time = mach_continuous_time();
-
-    while (options.max_tokens == 0 || atomic_load(&globals.num_generated_tokens) < options.max_tokens) {
-
-        uint32_t predicted_token = UINT32_MAX;
-        size_t num_predicted_tokens = 0;
-        const uint64_t inference_start_timestamp = mach_continuous_time();
-        status = gptoss_context_sample(context, options.temperature, /*rng_state=*/0, /*num_tokens=*/1, &predicted_token, &num_predicted_tokens);
-        if (status != gptoss_status_success) {
-            fprintf(stderr, "Error: failed to sample from the Context object\n");
-            goto error;
-        }
-        const uint64_t inference_end_timestamp = mach_continuous_time();
-
-        if (predicted_token == return_token_id) {
-            // Yield token -> stop generation
-            break;
-        }
-
-        // Unembedding: detokenize
-        size_t token_size = 0;
-        const void* token_ptr = NULL;
-        status = gptoss_tokenizer_decode(tokenizer, predicted_token, &token_ptr, &token_size);
-        if (status != gptoss_status_success) {
-            fprintf(stderr, "Error: failed to detokenize predicted token %" PRIu32 "\n", predicted_token);
-            goto error;
-        }
-        const size_t previous_num_generated_tokens = atomic_fetch_add(&globals.num_generated_tokens, 1);
-        if (previous_num_generated_tokens == 0) {
-            atomic_fetch_add(&globals.prefill_microseconds, mach_timestamp_diff_to_microseconds(prefill_start_time, prefill_end_time));
-        } else {
-            atomic_fetch_add(&globals.generation_microseconds, mach_timestamp_diff_to_microseconds(inference_start_timestamp, inference_end_timestamp));
-        }
-        printf("%.*s", (int) token_size, (const char*) token_ptr);
-
-        status = gptoss_context_append_tokens(context, 1, &predicted_token);
-        if (status != gptoss_status_success) {
-            fprintf(stderr, "Error: failed to append predicted token %" PRIu32 " to context\n", predicted_token);
-            goto error;
-        }
-    }
-
-    print_profile();
-
-    return EXIT_SUCCESS;
-
-error:
-    gptoss_context_release(context);
-    gptoss_tokenizer_release(tokenizer);
-    gptoss_model_release(model);
-    return EXIT_FAILURE;
-}

gptoss_kernels/source/include/internal/log.h

@@ -2,6 +2,9 @@
 
 #include <stdarg.h>
 
+#ifdef __cplusplus
+extern "C" {
+#endif
 
 void gptoss_format_log(const char* format, va_list args);
 
@@ -13,6 +16,10 @@ inline static void gptoss_log(const char* format, ...) {
     va_end(args);
 }
 
+#ifdef __cplusplus
+} // extern "C"
+#endif
+
 #define GPTOSS_LOG_ERROR(message, ...) \
     gptoss_log("Error: " message "\n", ##__VA_ARGS__)
 

gptoss_kernels/source/include/internal/metal.h

@@ -1,7 +1,6 @@
 #pragma once
 
 #include <stddef.h>
-
 #include <gpt-oss/types.h>
 
 #ifdef __cplusplus

gptoss_kernels/source/matmul.metal

@@ -43,7 +43,10 @@ kernel void gptoss_f32_bf16w_matmul(
     bias += row;
     output += gid.y * args.num_rows + row;
 
-    uint num_iter = (num_column_vecs - simdgroup_tid + (simdgroup_size - 1)) / simdgroup_size;
+    uint num_iter = 0;
+    if (simdgroup_tid < num_column_vecs) {
+        num_iter = (num_column_vecs - simdgroup_tid + (simdgroup_size - 1)) / simdgroup_size;
+    }
 
     float4 sum4 = 0.0f;
     do {
@@ -97,7 +100,10 @@ kernel void gptoss_f32_bf16w_matmul_qkv(
     bias += row;
     q += gid.y * args.num_rows;
 
-    uint num_iter = (num_column_vecs - simdgroup_tid + (simdgroup_size - 1)) / simdgroup_size;
+    uint num_iter = 0;
+    if (simdgroup_tid < num_column_vecs) {
+        num_iter = (num_column_vecs - simdgroup_tid + (simdgroup_size - 1)) / simdgroup_size;
+    }
 
     float4 sum4 = 0.0f;
     do {

gptoss_kernels/source/metal.m

@@ -9,7 +9,6 @@
 #include <internal/log.h>
 #include <internal/metal.h>
 
-
 static size_t gptoss_metal_device_get_core_count(id<MTLDevice> device) {
     if (!device) {
         return 0;

gptoss_kernels/source/model.c

@@ -1,581 +0,0 @@
-#include <assert.h>
-#include <inttypes.h>
-#include <stdatomic.h>
-#include <stdint.h>
-#include <stdlib.h>
-#include <string.h>
-
-#include <errno.h>  // errno, EISDIR, ENOENT, ENOTDIR
-#include <fcntl.h>  // open
-#include <mach/vm_page_size.h>  // vm_page_size
-#include <sys/mman.h>  // mmap, PROT_READ, MAP_PRIVATE
-#include <sys/stat.h>  // fstat, stat
-#include <sys/types.h>  // off_t, ssize_t
-#include <unistd.h>  // close
-
-#include <gpt-oss.h>
-
-#include "internal/datatype.h"
-#include "internal/kernel-args.h"  // gptoss_expert_prediction
-#include "internal/log.h"
-#include "internal/uuid.h"
-#include "internal/storage.h"
-#include "internal/math.h"
-#include "internal/model.h"
-
-
-static size_t round_up_to_page_size(size_t bytes) {
-    const size_t page_size_mask = (size_t) vm_page_size - 1;
-    if ((bytes & page_size_mask) != 0) {
-        bytes |= page_size_mask;
-        bytes += 1;
-    }
-    return bytes;
-}
-
-static size_t round_down_to_page_size(size_t bytes) {
-    const size_t page_size_mask = (size_t) vm_page_size - 1;
-    return bytes & ~page_size_mask;
-}
-
-static enum gptoss_status read_fd(int fd, void* data, size_t size, const char* path) {
-    assert(fd != -1);
-    assert(data != NULL);
-    assert(size != 0);
-
-    size_t bytes_to_read = size;
-    char* current_byte = (char*) data;
-    do {
-        const ssize_t read_result = read(fd, current_byte, bytes_to_read);
-        if (read_result < 0) {
-            GPTOSS_LOG_ERROR("reading %zu bytes from file %s failed with error %d",
-                size, path, errno);
-            return gptoss_status_io_error;
-        }
-        current_byte += (size_t) read_result;
-        bytes_to_read -= (size_t) read_result;
-    } while (bytes_to_read != 0);
-    return gptoss_status_success;
-}
-
-static void prefetch_fd(int fd, size_t offset, size_t size, const char* path) {
-    // radvisory.ra_count is int, so we can't prefetch 2GB+ at once
-    const size_t prefetch_max = round_down_to_page_size((size_t) INT_MAX);
-    do {
-        const size_t prefetch_size = math_min(size, prefetch_max);
-        const struct radvisory ra = {
-            .ra_offset = offset,
-            .ra_count = (int) prefetch_size,
-        };
-        if (fcntl(fd, F_RDADVISE, &ra) == -1) {
-            GPTOSS_LOG_WARNING("fcntl(%s, F_RDADVISE, .ra_offset=%zu, .ra_count=%d) failed with error %d\n",
-                path, (size_t) ra.ra_offset, ra.ra_count, errno);
-            return;
-        }
-        offset += prefetch_size;
-        size -= prefetch_size;
-    } while (size != 0);
-}
-
-enum gptoss_status GPTOSS_ABI gptoss_model_create_from_file(
-    const char* path,
-    gptoss_model_t* model_out)
-{
-    *model_out = NULL;
-
-    enum gptoss_status status = gptoss_status_success;
-    struct gptoss_model* model = NULL;
-    struct gptoss_tokenizer* tokenizer = NULL;
-    int fd = -1;
-    size_t file_offset = 0;
-
-    fd = open(path, O_RDONLY);
-    if (fd == -1) {
-        GPTOSS_LOG_ERROR("open(%s) failed with error %d", path, errno);
-        switch (errno) {
-            case EISDIR:
-            case ENOENT:
-            case ENOTDIR:
-                status = gptoss_status_invalid_argument;
-                break;
-            default:
-                status = gptoss_status_io_error;
-                break;
-        }
-        goto cleanup;
-    }
-
-    struct gptoss_file_header file_header;
-    status = read_fd(fd, &file_header, sizeof(file_header), path);
-    if (status != gptoss_status_success) {
-        goto cleanup;
-    }
-    file_offset += sizeof(file_header);
-
-    if (file_header.magic[0] != 'G' ||
-        file_header.magic[1] != 'P' ||
-        file_header.magic[2] != 'T' ||
-        file_header.magic[3] != '-' ||
-        file_header.magic[4] != 'O' ||
-        file_header.magic[5] != 'S' ||
-        file_header.magic[6] != 'S' ||
-        file_header.magic[7] != ' ' ||
-        file_header.magic[8] != 'v' ||
-        file_header.magic[9] != '1' ||
-        file_header.magic[10] != '.' ||
-        file_header.magic[11] != '0' ||
-        file_header.zero != 0)
-    {
-        GPTOSS_LOG_ERROR("invalid magic in file %s", path);
-        status = gptoss_status_invalid_argument;
-        goto cleanup;
-    }
-
-    struct gptoss_uuid model_uuid;
-    status = read_fd(fd, &model_uuid, sizeof(model_uuid), path);
-    if (status != gptoss_status_success) {
-        goto cleanup;
-    }
-    file_offset += sizeof(model_uuid);
-
-    if (!gptoss_is_gptoss_model_uuid(&model_uuid)) {
-        GPTOSS_LOG_ERROR("unsupported model UUID " UUID_FORMAT, UUID_ARGS(model_uuid));
-        status = gptoss_status_invalid_argument;
-        goto cleanup;
-    }
-
-    struct gptoss_gptoss_model_header model_header;
-    status = read_fd(fd, &model_header, sizeof(model_header), path);
-    if (status != gptoss_status_success) {
-        goto cleanup;
-    }
-    file_offset += sizeof(model_header);
-
-    struct gptoss_uuid layout_uuid;
-    status = read_fd(fd, &layout_uuid, sizeof(layout_uuid), path);
-    if (status != gptoss_status_success) {
-        goto cleanup;
-    }
-    file_offset += sizeof(layout_uuid);
-
-    if (!gptoss_is_applegpu_layout_uuid(&layout_uuid)) {
-        GPTOSS_LOG_ERROR("unsupported layout UUID " UUID_FORMAT, UUID_ARGS(layout_uuid));
-        status = gptoss_status_invalid_argument;
-        goto cleanup;
-    }
-
-    const size_t model_size = sizeof(struct gptoss_model) + model_header.num_blocks * sizeof(struct gptoss_metal_buffer);
-    model = malloc(model_size);
-    if (model == NULL) {
-        GPTOSS_LOG_ERROR("failed to allocate %zu bytes for model descriptor", model_size);
-        status = gptoss_status_insufficient_memory;
-        goto cleanup;
-    }
-    memset(model, 0, model_size);
-
-    atomic_store_explicit(&model->ref_count, 1, memory_order_relaxed);
-    model->context_length = model_header.context_length;
-    model->num_blocks = model_header.num_blocks;
-    model->num_experts = model_header.num_experts;
-    model->num_active_experts = model_header.num_active_experts;
-    model->embedding_dim = model_header.embedding_dim;
-    model->mlp_dim = model_header.mlp_dim;
-    model->swiglu_limit = model_header.swiglu_limit;
-    model->head_dim = model_header.head_dim;
-    model->num_heads = model_header.num_heads;
-    model->num_kv_heads = model_header.num_kv_heads;
-    model->attention_window = model_header.attention_window;
-    model->rope_theta = model_header.rope_theta;
-    model->interpolation_scale = model_header.interpolation_scale;
-    model->yarn_offset = model_header.yarn_offset;
-    model->yarn_scale = model_header.yarn_scale;
-    model->yarn_multiplier = model_header.yarn_multiplier;
-    model->rmsnorm_epsilon = model_header.rmsnorm_epsilon;
-
-    struct gptoss_uuid tokenizer_uuid;
-    status = read_fd(fd, &tokenizer_uuid, sizeof(tokenizer_uuid), path);
-    if (status != gptoss_status_success) {
-        goto cleanup;
-    }
-    file_offset += sizeof(tokenizer_uuid);
-
-    if (!gptoss_is_tiktoken_tokenizer_uuid(&tokenizer_uuid)) {
-        GPTOSS_LOG_ERROR("unsupported tokenizer UUID " UUID_FORMAT, UUID_ARGS(tokenizer_uuid));
-        status = gptoss_status_invalid_argument;
-        goto cleanup;
-    }
-
-    struct gptoss_tiktoken_tokenizer_header tokenizer_header;
-    status = read_fd(fd, &tokenizer_header, sizeof(tokenizer_header), path);
-    if (status != gptoss_status_success) {
-        goto cleanup;
-    }
-    file_offset += sizeof(tokenizer_header);
-
-    tokenizer = malloc(sizeof(struct gptoss_tokenizer));
-    if (tokenizer == NULL) {
-        GPTOSS_LOG_ERROR("failed to allocate %zu bytes for tokenizer descriptor", sizeof(struct gptoss_tokenizer));
-        status = gptoss_status_insufficient_memory;
-        goto cleanup;
-    }
-    memset(tokenizer, 0, sizeof(struct gptoss_tokenizer));
-    // Initialize all special token IDs to UINT32_MAX (0xFF in all bytes)
-    memset(tokenizer->special_token_id, 0xFF, sizeof(tokenizer->special_token_id));
-
-    atomic_store_explicit(&tokenizer->ref_count, 1, memory_order_relaxed);
-    tokenizer->num_special_tokens = tokenizer_header.num_special_tokens;
-    tokenizer->num_text_tokens = tokenizer_header.num_text_tokens;
-    model->vocabulary_size = tokenizer_header.num_special_tokens + tokenizer_header.num_text_tokens;
-    for (uint32_t t = 0; t < tokenizer_header.num_special_tokens; t++) {
-        struct gptoss_uuid token_uuid;
-        status = read_fd(fd, &token_uuid, sizeof(token_uuid), path);
-        if (status != gptoss_status_success) {
-            goto cleanup;
-        }
-        file_offset += sizeof(token_uuid);
-
-        const enum gptoss_special_token token = gptoss_special_token_decode_uuid(&token_uuid);
-        if (token != gptoss_special_token_invalid) {
-            tokenizer->special_token_id[token - 1] = tokenizer_header.num_text_tokens + t;
-        }
-    }
-
-    const size_t tokenizer_start_offset = file_offset;
-    const size_t tokenizer_end_offset = tokenizer_start_offset + tokenizer_header.regex_size + tokenizer_header.tokens_size;
-    const size_t tokenizer_mapping_start = round_down_to_page_size(tokenizer_start_offset);
-    const size_t tokenizer_mapping_size = round_up_to_page_size(tokenizer_end_offset) - tokenizer_mapping_start;
-    void* tokenizer_mapping_ptr = mmap(NULL, tokenizer_mapping_size, PROT_READ, MAP_PRIVATE, fd, tokenizer_mapping_start);
-    if (tokenizer_mapping_ptr == (void*) -1) {
-        GPTOSS_LOG_ERROR("failed to mmap(%s) tokenizer at offset %zu size %zu",
-            path, tokenizer_mapping_start, tokenizer_mapping_size);
-        status = gptoss_status_io_error;
-        goto cleanup;
-    }
-    tokenizer->mapping_ptr = tokenizer_mapping_ptr;
-    tokenizer->mapping_size = tokenizer_mapping_size;
-    tokenizer->regex_ptr = (const char*) tokenizer_mapping_ptr + (tokenizer_start_offset - tokenizer_mapping_start);
-    tokenizer->tokens_ptr = tokenizer->regex_ptr + tokenizer_header.regex_size;
-
-    if (madvise(tokenizer_mapping_ptr, tokenizer_mapping_size, MADV_RANDOM | MADV_WILLNEED) != 0) {
-        GPTOSS_LOG_WARNING("madvise(%s, size=%zu) failed with error %d", path, tokenizer_mapping_size, errno);
-    }
-
-    prefetch_fd(fd, tokenizer_mapping_start, tokenizer_mapping_size, path);
-
-    struct stat model_stat = {0};
-    int stat_result = fstat(fd, &model_stat);
-    if (stat_result != 0) {
-        GPTOSS_LOG_ERROR("stat(%s) failed with error %d", path, errno);
-        status = gptoss_status_io_error;
-        goto cleanup;
-    }
-
-    const size_t model_mapping_start = round_up_to_page_size(tokenizer_end_offset);
-    const size_t model_mapping_size = round_up_to_page_size((size_t) model_stat.st_size) - model_mapping_start;
-    void* model_mapping_ptr = mmap(NULL, model_mapping_size, PROT_READ, MAP_PRIVATE, fd, model_mapping_start);
-    if (model_mapping_ptr == (void*) -1) {
-        GPTOSS_LOG_ERROR("failed to mmap(%s) model weights at offset %zu size %zu",
-            path, model_mapping_start, model_mapping_size);
-        status = gptoss_status_io_error;
-        goto cleanup;
-    }
-    model->mapping_ptr = model_mapping_ptr;
-    model->mapping_size = model_mapping_size;
-
-    if (madvise(model_mapping_ptr, model_mapping_size, MADV_SEQUENTIAL | MADV_WILLNEED) != 0) {
-        GPTOSS_LOG_WARNING("madvise(%s, size=%zu) failed with error %d", path, model_mapping_size, errno);
-    }
-
-    prefetch_fd(fd, model_mapping_start, model_mapping_size, path);
-
-    if (mlock(model_mapping_ptr, model_mapping_size) != 0) {
-        GPTOSS_LOG_WARNING("mlock(%s, size=%zu) failed with error %d", path, model_mapping_size, errno);
-    } else {
-        model->lock_memory = true;
-    }
-
-    // Initialize Metal
-    status = gptoss_metal_device_create_system_default(&model->device);
-    if (status != gptoss_status_success) {
-        goto cleanup;
-    }
-    model->max_threadgroups = model->device.num_cores * 3;
-    status = gptoss_metal_command_queue_create(&model->device, &model->command_queue);
-    if (status != gptoss_status_success) {
-        goto cleanup;
-    }
-
-    // Metal kernels
-    status = gptoss_metal_library_create_default(&model->device, &model->library);
-    if (status != gptoss_status_success) {
-        goto cleanup;
-    }
-    status = gptoss_metal_function_create(&model->library, "gptoss_bf16_f32_embeddings", &model->bf16_f32_embeddings_fn);
-    if (status != gptoss_status_success) {
-        goto cleanup;
-    }
-    status = gptoss_metal_function_create(&model->library, "gptoss_f32_bf16w_rmsnorm", &model->f32_bf16w_rmsnorm_fn);
-    if (status != gptoss_status_success) {
-        goto cleanup;
-    }
-    status = gptoss_metal_function_create(&model->library, "gptoss_f32_bf16w_matmul", &model->f32_bf16w_matmul_fn);
-    if (status != gptoss_status_success) {
-        goto cleanup;
-    }
-    status = gptoss_metal_function_create(&model->library, "gptoss_f32_bf16w_matmul_qkv", &model->f32_bf16w_matmul_qkv_fn);
-    if (status != gptoss_status_success) {
-        goto cleanup;
-    }
-    status = gptoss_metal_function_create(&model->library, "gptoss_f32_bf16w_dense_matmul_qkv", &model->f32_bf16w_dense_matmul_qkv_fn);
-    if (status != gptoss_status_success) {
-        goto cleanup;
-    }
-    status = gptoss_metal_function_create(&model->library, "gptoss_f32_bf16w_dense_matmul_attn_output", &model->f32_bf16w_dense_matmul_attn_output_fn);
-    if (status != gptoss_status_success) {
-        goto cleanup;
-    }
-    status = gptoss_metal_function_create(&model->library, "gptoss_f32_bf16w_dense_matmul_mlp_gate", &model->f32_bf16w_dense_matmul_mlp_gate_fn);
-    if (status != gptoss_status_success) {
-        goto cleanup;
-    }
-    status = gptoss_metal_function_create(&model->library, "gptoss_f32_bf16w_unembedding", &model->f32_bf16w_unembedding_fn);
-    if (status != gptoss_status_success) {
-        goto cleanup;
-    }
-    status = gptoss_metal_function_create(&model->library, "gptoss_f32_rope", &model->f32_rope_fn);
-    if (status != gptoss_status_success) {
-        goto cleanup;
-    }
-    status = gptoss_metal_function_create(&model->library, "gptoss_f32_expert_routing_metadata", &model->f32_expert_routing_metadata_fn);
-    if (status != gptoss_status_success) {
-        goto cleanup;
-    }
-    status = gptoss_metal_function_create(&model->library, "gptoss_f32_scatter_e4", &model->f32_scatter_e4_fn);
-    if (status != gptoss_status_success) {
-        goto cleanup;
-    }
-    status = gptoss_metal_function_create(&model->library, "gptoss_f32_mf4w_moe_dense_matmul_swiglu", &model->f32_mf4w_moe_dense_matmul_swiglu_fn);
-    if (status != gptoss_status_success) {
-        goto cleanup;
-    }
-    status = gptoss_metal_function_create(&model->library, "gptoss_f32_mf4w_moe_dense_matmul", &model->f32_mf4w_moe_dense_matmul_fn);
-    if (status != gptoss_status_success) {
-        goto cleanup;
-    }
-    status = gptoss_metal_function_create(&model->library, "gptoss_f32_gather_and_accumulate_e4", &model->f32_gather_and_accumulate_e4_fn);
-    if (status != gptoss_status_success) {
-        goto cleanup;
-    }
-    status = gptoss_metal_function_create(&model->library, "gptoss_f32_mf4w_moe_matmul_swiglu", &model->f32_mf4w_moe_matmul_swiglu_fn);
-    if (status != gptoss_status_success) {
-        goto cleanup;
-    }
-    status = gptoss_metal_function_create(&model->library, "gptoss_f32_mf4w_moe_matmul", &model->f32_mf4w_moe_matmul_fn);
-    if (status != gptoss_status_success) {
-        goto cleanup;
-    }
-    status = gptoss_metal_function_create(&model->library, "gptoss_f32_accumulate_e4", &model->f32_accumulate_e4_fn);
-    if (status != gptoss_status_success) {
-        goto cleanup;
-    }
-    status = gptoss_metal_function_create(&model->library, "gptoss_f32_topk_softmax_e32_k4", &model->f32_topk_softmax_e32_k4_fn);
-    if (status != gptoss_status_success) {
-        goto cleanup;
-    }
-    status = gptoss_metal_function_create(&model->library, "gptoss_f32_topk_softmax_e128_k4", &model->f32_topk_softmax_e128_k4_fn);
-    if (status != gptoss_status_success) {
-        goto cleanup;
-    }
-    status = gptoss_metal_function_create(&model->library, "gptoss_f32_softmax", &model->f32_softmax_fn);
-    if (status != gptoss_status_success) {
-        goto cleanup;
-    }
-    status = gptoss_metal_function_create(&model->library, "gptoss_f32_sample", &model->f32_sample_fn);
-    if (status != gptoss_status_success) {
-        goto cleanup;
-    }
-    status = gptoss_metal_function_create(&model->library, "gptoss_f32_sdpa_q8_d64", &model->f32_sdpa_q8_d64_fn);
-    if (status != gptoss_status_success) {
-        goto cleanup;
-    }
-
-    // Kernel launch parameters
-    model->embeddings_threadgroup_size = 512;
-    model->attn_qkv_threadgroup_size = 1024;
-    model->attn_out_threadgroup_size = 768;
-    model->mlp_gate_threadgroup_size = 256;
-    model->mlp_swiglu_threadgroup_size = 192;
-    model->mlp_out_threadgroup_size = 192;
-    model->mlp_acc_threadgroup_size = 768;
-    model->unembedding_threadgroup_size = 416;
-
-    // Weight buffers
-    const char* current_ptr = (const char*) model->mapping_ptr;
-
-    const size_t embedding_weight_size = math_round_up_po2(model->vocabulary_size * model->embedding_dim * sizeof(gptoss_bfloat16), 16);
-    model->attn_rmsnorm_gain_offset = embedding_weight_size;
-    const size_t rmsnorm_weight_size = math_round_up_po2(model->embedding_dim * sizeof(gptoss_bfloat16), 16);
-    model->attn_qkv_weight_offset = model->attn_rmsnorm_gain_offset + rmsnorm_weight_size;
-    const size_t attn_qkv_dim = model->head_dim * (model->num_heads + 2 * model->num_kv_heads);
-    const size_t attn_qkv_weight_size = math_round_up_po2(attn_qkv_dim * model->embedding_dim * sizeof(gptoss_bfloat16), 16);
-    model->attn_qkv_bias_offset = model->attn_qkv_weight_offset + attn_qkv_weight_size;
-    const size_t attn_qkv_bias_size = math_round_up_po2(attn_qkv_dim * sizeof(gptoss_bfloat16), 16);
-    model->attn_sdpa_sink_offset = model->attn_qkv_bias_offset + attn_qkv_bias_size;
-    const size_t attn_sink_weight_size = math_round_up_po2(model->num_heads * sizeof(gptoss_bfloat16), 16);
-    model->attn_out_weight_offset = model->attn_sdpa_sink_offset + attn_sink_weight_size;
-    const size_t attn_out_weight_size = math_round_up_po2(model->embedding_dim * model->num_heads * model->head_dim * sizeof(gptoss_bfloat16), 16);
-    model->attn_out_bias_offset = model->attn_out_weight_offset + attn_out_weight_size;
-    const size_t attn_out_bias_size = math_round_up_po2(model->embedding_dim * sizeof(gptoss_bfloat16), 16);
-    model->mlp_rmsnorm_gain_offset = model->attn_out_bias_offset + attn_out_bias_size;
-    model->mlp_gate_weight_offset = model->mlp_rmsnorm_gain_offset + rmsnorm_weight_size;
-    const size_t mlp_gate_weight_size = math_round_up_po2(model->num_experts * model->embedding_dim * sizeof(gptoss_bfloat16), 16);
-    model->mlp_gate_bias_offset = model->mlp_gate_weight_offset + mlp_gate_weight_size;
-    const size_t mlp_gate_bias_size = math_round_up_po2(model->num_experts * sizeof(gptoss_bfloat16), 16);
-    const size_t per_block_shared_weights_size =
-        rmsnorm_weight_size + attn_qkv_weight_size + attn_qkv_bias_size + attn_sink_weight_size + attn_out_weight_size + attn_out_bias_size +
-        rmsnorm_weight_size + mlp_gate_weight_size + mlp_gate_bias_size;
-    model->rmsnorm_weight_offset = embedding_weight_size + model->num_blocks * per_block_shared_weights_size;
-    model->unembedding_weight_offset = model->rmsnorm_weight_offset + rmsnorm_weight_size;
-    const size_t unembedding_weight_size = math_round_up_po2(model->vocabulary_size * model->embedding_dim * sizeof(gptoss_bfloat16), 16);
-
-    model->per_block_shared_weights_size = per_block_shared_weights_size;
-    const size_t shared_weights_size =
-        round_up_to_page_size(embedding_weight_size + rmsnorm_weight_size + unembedding_weight_size + model->num_blocks * per_block_shared_weights_size);
-
-    status = gptoss_metal_buffer_wrap(&model->device, shared_weights_size, current_ptr, &model->shared_weight_buffer);
-    if (status != gptoss_status_success) {
-        GPTOSS_LOG_ERROR("failed to map expert-shared weight of size %zu onto a Metal buffer", shared_weights_size);
-        goto cleanup;
-    }
-    current_ptr += shared_weights_size;
-    model->weights_size += shared_weights_size;
-
-    const size_t mlp_swiglu_weight_block_size = math_round_up_po2(2 * model->mlp_dim * model->embedding_dim / 2, 16);
-    model->mlp_swiglu_scale_offset = mlp_swiglu_weight_block_size;
-    const size_t mlp_swiglu_weight_scale_size = math_round_up_po2(2 * model->mlp_dim * model->embedding_dim / 32, 16);
-    model->mlp_swiglu_bias_offset = model->mlp_swiglu_scale_offset + mlp_swiglu_weight_scale_size;
-    const size_t mlp_swiglu_bias_size = math_round_up_po2(2 * model->mlp_dim * sizeof(gptoss_bfloat16), 16);
-    model->mlp_out_block_offset = model->mlp_swiglu_bias_offset + mlp_swiglu_bias_size;
-    const size_t mlp_out_weight_block_size = math_round_up_po2(model->embedding_dim * model->mlp_dim / 2, 16);
-    model->mlp_out_scale_offset = model->mlp_out_block_offset + mlp_out_weight_block_size;
-    const size_t mlp_out_weight_scale_size = math_round_up_po2(model->embedding_dim * model->mlp_dim / 32, 16);
-    model->mlp_out_bias_offset = model->mlp_out_scale_offset + mlp_out_weight_scale_size;
-    const size_t mlp_out_bias_size = math_round_up_po2(model->embedding_dim * sizeof(gptoss_bfloat16), 16);
-    model->per_expert_block_weight_size =
-        mlp_swiglu_weight_block_size + mlp_swiglu_weight_scale_size + mlp_swiglu_bias_size + mlp_out_weight_block_size + mlp_out_weight_scale_size + mlp_out_bias_size;
-    const size_t moe_block_weight_size = round_up_to_page_size(model->num_experts * model->per_expert_block_weight_size);
-    for (uint32_t n = 0; n < model->num_blocks; n++) {
-        status = gptoss_metal_buffer_wrap(&model->device, moe_block_weight_size, current_ptr, &model->block_weight_buffers[n]);
-        if (status != gptoss_status_success) {
-            GPTOSS_LOG_ERROR("failed to map block #%" PRIu32 " MoE weight of size %zu onto a Metal buffer",
-                n, moe_block_weight_size);
-            goto cleanup;
-        }
-        current_ptr += moe_block_weight_size;
-        model->weights_size += moe_block_weight_size;
-    }
-
-    // Commit tokenizer
-    model->tokenizer = tokenizer;
-    tokenizer = NULL;
-
-    // Commit model
-    *model_out = model;
-    model = NULL;
-
-cleanup:
-    if (fd != -1) {
-        close(fd);
-        fd = -1;
-    }
-    gptoss_model_release(model);  // does nothing if model is NULL
-    gptoss_tokenizer_release(tokenizer);  // does nothing if tokenizer is NULL
-    return status;
-}
-
-enum gptoss_status GPTOSS_ABI gptoss_model_get_tokenizer(
-    gptoss_model_t model,
-    gptoss_tokenizer_t* tokenizer_out)
-{
-    gptoss_tokenizer_t tokenizer = model->tokenizer;
-    atomic_fetch_add_explicit(&tokenizer->ref_count, 1, memory_order_relaxed);
-    *tokenizer_out = tokenizer;
-    return gptoss_status_success;
-}
-
-enum gptoss_status GPTOSS_ABI gptoss_model_get_max_context_length(
-    gptoss_model_t model,
-    size_t* max_context_length_out)
-{
-    *max_context_length_out = model->context_length;
-    return gptoss_status_success;
-}
-
-enum gptoss_status GPTOSS_ABI gptoss_model_retain(
-    gptoss_model_t model)
-{
-    atomic_fetch_add_explicit(&model->ref_count, 1, memory_order_relaxed);
-    return gptoss_status_success;
-}
-
-enum gptoss_status GPTOSS_ABI gptoss_model_release(
-    gptoss_model_t model)
-{
-    if (model != NULL) {
-        if (atomic_fetch_sub_explicit(&model->ref_count, 1, memory_order_acq_rel) == 1) {
-            gptoss_tokenizer_release(model->tokenizer);
-
-            // Weight buffers
-            gptoss_metal_buffer_release(&model->shared_weight_buffer);
-            for (uint32_t n = 0; n < model->num_blocks; n++) {
-                gptoss_metal_buffer_release(&model->block_weight_buffers[n]);
-            }
-
-            // Metal kernels
-            gptoss_metal_function_release(&model->bf16_f32_embeddings_fn);
-            gptoss_metal_function_release(&model->f32_bf16w_rmsnorm_fn);
-            gptoss_metal_function_release(&model->f32_bf16w_matmul_fn);
-            gptoss_metal_function_release(&model->f32_bf16w_matmul_qkv_fn);
-            gptoss_metal_function_release(&model->f32_bf16w_dense_matmul_qkv_fn);
-            gptoss_metal_function_release(&model->f32_bf16w_dense_matmul_attn_output_fn);
-            gptoss_metal_function_release(&model->f32_bf16w_dense_matmul_mlp_gate_fn);
-            gptoss_metal_function_release(&model->f32_bf16w_unembedding_fn);
-            gptoss_metal_function_release(&model->f32_rope_fn);
-            gptoss_metal_function_release(&model->f32_expert_routing_metadata_fn);
-            gptoss_metal_function_release(&model->f32_scatter_e4_fn);
-            gptoss_metal_function_release(&model->f32_mf4w_moe_dense_matmul_swiglu_fn);
-            gptoss_metal_function_release(&model->f32_mf4w_moe_dense_matmul_fn);
-            gptoss_metal_function_release(&model->f32_gather_and_accumulate_e4_fn);
-            gptoss_metal_function_release(&model->f32_mf4w_moe_matmul_swiglu_fn);
-            gptoss_metal_function_release(&model->f32_mf4w_moe_matmul_fn);
-            gptoss_metal_function_release(&model->f32_accumulate_e4_fn);
-            gptoss_metal_function_release(&model->f32_topk_softmax_e32_k4_fn);
-            gptoss_metal_function_release(&model->f32_topk_softmax_e128_k4_fn);
-            gptoss_metal_function_release(&model->f32_softmax_fn);
-            gptoss_metal_function_release(&model->f32_sample_fn);
-            gptoss_metal_function_release(&model->f32_sdpa_q8_d64_fn);
-            gptoss_metal_library_release(&model->library);
-
-            gptoss_metal_command_queue_release(&model->command_queue);
-            gptoss_metal_device_release(&model->device);
-            // Weight buffers
-
-            if (model->mapping_ptr != NULL && model->mapping_size != 0) {
-                if (model->lock_memory) {
-                    if (munlock(model->mapping_ptr, model->mapping_size) != 0) {
-                        GPTOSS_LOG_WARNING("munlock for model weight mapping failed with error %d", errno);
-                    }
-                }
-
-                if (munmap(model->mapping_ptr, model->mapping_size) != 0) {
-                    GPTOSS_LOG_WARNING("munmap for model weight mapping failed with error %d", errno);
-                }
-            }
-
-            const size_t model_size = sizeof(struct gptoss_model) + model->num_blocks * sizeof(struct gptoss_metal_buffer);
-            memset(model, 0, model_size);
-            free(model);
-        }
-    }
-    return gptoss_status_success;
-}

gptoss_kernels/source/tensor_wrappers.cpp

@@ -0,0 +1,77 @@
+#include <internal/metal-kernels.h>
+#include <internal/metal.h>
+#include <ATen/Tensor.h>
+
+void f32_bf16w_matmul_torch(const at::Tensor &input,
+    const at::Tensor &weight_bf16,
+    const at::Tensor &bias_bf16,
+    at::Tensor &output,
+    int64_t num_tokens, int64_t num_cols, int64_t num_rows, int64_t threadgroup_size)
+{
+    TORCH_CHECK(input.dtype() == at::kFloat, "input must be float32");
+    TORCH_CHECK(weight_bf16.dtype() == at::kBFloat16, "weight must be bfloat16");
+    TORCH_CHECK(bias_bf16.dtype() == at::kBFloat16, "bias must be bfloat16");
+    TORCH_CHECK(output.dtype() == at::kFloat, "output must be float32");
+
+    TORCH_CHECK(input.dim() == 2, "input must be 2D");
+    TORCH_CHECK(weight_bf16.dim() == 2, "weight must be 2D");
+    TORCH_CHECK(bias_bf16.dim() == 1, "bias must be 1D");
+    TORCH_CHECK(output.dim() == 2, "output must be 2D");
+    TORCH_CHECK(input.size(0) == num_tokens && input.size(1) == num_cols,
+                "input shape must be [num_tokens, num_cols]");
+    TORCH_CHECK(weight_bf16.size(0) == num_cols && weight_bf16.size(1) == num_rows,
+                "weight shape must be [num_cols, num_rows]");
+    TORCH_CHECK(bias_bf16.size(0) == num_rows, "bias length must be num_rows");
+    TORCH_CHECK(output.size(0) == num_tokens && output.size(1) == num_rows,
+                "output shape must be [num_tokens, num_rows]");
+
+    auto input_cpu  = input.contiguous().to(at::kCPU);
+    auto weight_cpu = weight_bf16.transpose(0, 1).contiguous().to(at::kCPU);
+    auto bias_cpu   = bias_bf16.contiguous().to(at::kCPU);
+    auto out_cpu = output.detach().to(at::kCPU).contiguous().clone();
+
+    gptoss_metal_device device{}; gptoss_metal_library library{};
+    gptoss_metal_function fn{};   gptoss_metal_command_queue cq{};
+    gptoss_metal_command_buffer cb{};
+
+    TORCH_CHECK(gptoss_metal_device_create_system_default(&device) == gptoss_status_success, "device_create failed");
+    TORCH_CHECK(gptoss_metal_library_create_default(&device, &library) == gptoss_status_success, "library_create failed");
+    TORCH_CHECK(gptoss_metal_function_create(&library, "gptoss_f32_bf16w_matmul", &fn) == gptoss_status_success, "function_create failed");
+    TORCH_CHECK(gptoss_metal_command_queue_create(&device, &cq) == gptoss_status_success, "cq_create failed");
+    TORCH_CHECK(gptoss_metal_command_buffer_create(&cq, &cb) == gptoss_status_success, "cb_create failed");
+
+    const size_t in_bytes  = (size_t)num_tokens * (size_t)num_cols * sizeof(float);
+    const size_t wt_bytes  = (size_t)num_rows   * (size_t)num_cols * sizeof(uint16_t);
+    const size_t bs_bytes  = (size_t)num_rows * sizeof(uint16_t);
+    const size_t out_bytes = (size_t)num_tokens * (size_t)num_rows * sizeof(float);
+
+    gptoss_metal_buffer in_buf{}, wt_buf{}, bs_buf{}, out_buf{}, ctrl_buf{};
+    TORCH_CHECK(gptoss_metal_buffer_wrap(&device, in_bytes,  input_cpu.data_ptr(),  &in_buf)  == gptoss_status_success, "wrap input failed");
+    TORCH_CHECK(gptoss_metal_buffer_wrap(&device, wt_bytes,  weight_cpu.data_ptr(), &wt_buf)  == gptoss_status_success, "wrap weight failed");
+    TORCH_CHECK(gptoss_metal_buffer_wrap(&device, bs_bytes,  bias_cpu.data_ptr(),   &bs_buf)  == gptoss_status_success, "wrap bias failed");
+    TORCH_CHECK(gptoss_metal_buffer_create(&device, out_bytes, nullptr, &out_buf)   == gptoss_status_success, "alloc out failed");
+    uint32_t ctrl_zero = 0;
+    TORCH_CHECK(gptoss_metal_buffer_create(&device, sizeof(uint32_t), &ctrl_zero, &ctrl_buf) == gptoss_status_success, "alloc ctrl failed");
+
+    TORCH_CHECK(gptoss_metal_command_buffer_encode_launch_f32_bf16w_matmul(
+    &cb, &fn, (size_t)threadgroup_size,
+    &in_buf, 0, &wt_buf, 0, &bs_buf, 0, &out_buf, 0, &ctrl_buf, 0,
+    (uint32_t)num_tokens, (uint32_t)num_cols, (uint32_t)num_rows) == gptoss_status_success, "encode failed");
+
+    TORCH_CHECK(gptoss_metal_command_buffer_commit(&cb) == gptoss_status_success, "commit failed");
+    TORCH_CHECK(gptoss_metal_command_buffer_wait_completion(&cb, nullptr) == gptoss_status_success, "wait failed");
+
+    std::memcpy(out_cpu.data_ptr(), out_buf.ptr, out_bytes);
+    output.copy_(out_cpu.to(output.device(), /*non_blocking=*/false, /*copy=*/true));
+
+    (void) gptoss_metal_command_buffer_release(&cb);
+    (void) gptoss_metal_command_queue_release(&cq);
+    (void) gptoss_metal_function_release(&fn);
+    (void) gptoss_metal_library_release(&library);
+    (void) gptoss_metal_device_release(&device);
+    (void) gptoss_metal_buffer_release(&ctrl_buf);
+    (void) gptoss_metal_buffer_release(&out_buf);
+    (void) gptoss_metal_buffer_release(&bs_buf);
+    (void) gptoss_metal_buffer_release(&wt_buf);
+    (void) gptoss_metal_buffer_release(&in_buf);
+}

gptoss_kernels/source/tokenizer.c

@@ -1,106 +0,0 @@
-#include <assert.h>
-#include <stdatomic.h>
-#include <stddef.h>
-#include <stdint.h>
-#include <stdlib.h>
-#include <string.h>
-
-#include <errno.h>
-#include <sys/mman.h>
-
-#include <gpt-oss.h>
-
-#include "internal/log.h"
-#include "internal/model.h"
-
-
-enum gptoss_status GPTOSS_ABI gptoss_tokenizer_get_special_token_id(
-    gptoss_tokenizer_t tokenizer,
-    enum gptoss_special_token token_type,
-    uint32_t* token_id_out)
-{
-    uint32_t token_id = UINT32_MAX;
-    if (token_type != gptoss_special_token_invalid && token_type < gptoss_special_token_max)
-    {
-        token_id = tokenizer->special_token_id[(uint32_t) token_type - 1];
-    }
-    if (token_id == UINT32_MAX) {
-        return gptoss_status_invalid_argument;
-    }
-
-    *token_id_out = token_id;
-    return gptoss_status_success;
-}
-
-enum gptoss_status GPTOSS_ABI gptoss_tokenizer_get_num_text_tokens(
-    gptoss_tokenizer_t tokenizer,
-    uint32_t* num_text_tokens_out)
-{
-    *num_text_tokens_out = tokenizer->num_text_tokens;
-    return gptoss_status_success;
-}
-
-enum gptoss_status GPTOSS_ABI gptoss_tokenizer_get_num_special_tokens(
-    gptoss_tokenizer_t tokenizer,
-    uint32_t* num_special_tokens_out)
-{
-    *num_special_tokens_out = tokenizer->num_special_tokens;
-    return gptoss_status_success;
-}
-
-enum gptoss_status GPTOSS_ABI gptoss_tokenizer_get_num_tokens(
-    gptoss_tokenizer_t tokenizer,
-    uint32_t* num_tokens_out)
-{
-    *num_tokens_out = tokenizer->num_text_tokens + tokenizer->num_special_tokens;
-    return gptoss_status_success;
-}
-
-enum gptoss_status GPTOSS_ABI gptoss_tokenizer_decode(
-    gptoss_tokenizer_t tokenizer,
-    uint32_t token_id,
-    const void** token_ptr_out,
-    size_t* token_size_out)
-{
-    if (token_id >= tokenizer->num_text_tokens) {
-        return gptoss_status_invalid_argument;
-    }
-
-    const char* token_ptr = (const char*) tokenizer->tokens_ptr;
-    for (uint32_t t = 0; t < token_id; t++) {
-        // Reading unaligned uint16_t
-        uint16_t token_length;
-        memcpy(&token_length, token_ptr, sizeof(token_length));
-
-        token_ptr += (size_t) token_length + sizeof(uint16_t);
-    }
-
-    *token_ptr_out = (const void*) (token_ptr + sizeof(uint16_t));
-    *token_size_out = (size_t) *token_ptr;
-    return gptoss_status_success;
-}
-
-enum gptoss_status GPTOSS_ABI gptoss_tokenizer_retain(
-    gptoss_tokenizer_t tokenizer)
-{
-    atomic_fetch_add_explicit(&tokenizer->ref_count, 1, memory_order_relaxed);
-    return gptoss_status_success;
-}
-
-enum gptoss_status GPTOSS_ABI gptoss_tokenizer_release(
-    gptoss_tokenizer_t tokenizer)
-{
-    if (tokenizer != NULL) {
-        if (atomic_fetch_sub_explicit(&tokenizer->ref_count, 1, memory_order_acquire) == 1) {
-            if (tokenizer->mapping_ptr != NULL && tokenizer->mapping_size != 0) {
-                if (munmap(tokenizer->mapping_ptr, tokenizer->mapping_size) != 0) {
-                    GPTOSS_LOG_WARNING("munmap for tokenizer mapping failed with error %d", errno);
-                }
-            }
-
-            memset(tokenizer, 0, sizeof(struct gptoss_tokenizer));
-            free(tokenizer);
-        }
-    }
-    return gptoss_status_success;
-}

pyproject.toml

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

setup.py

@@ -0,0 +1,118 @@
+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_931bc1b_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/__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

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

torch-ext/gptoss_kernels/_ops.py

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

torch-ext/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

torch-ext/registration.h

@@ -0,0 +1,30 @@
+// 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);                                           \
+  }

torch-ext/torch_binding.cpp

@@ -0,0 +1,10 @@
+#include <ATen/Tensor.h>
+#include "torch_binding.h"
+#include "registration.h"
+
+TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
+  ops.def("f32_bf16w_matmul(Tensor input, Tensor weight_bf16, Tensor bias_bf16, Tensor output, int num_tokens, int num_cols, int num_rows, int threadgroup_size) -> ()");
+  ops.impl("f32_bf16w_matmul", torch::kMPS, &f32_bf16w_matmul_torch);
+}
+
+REGISTER_EXTENSION(TORCH_EXTENSION_NAME)

torch-ext/torch_binding.h

@@ -0,0 +1,5 @@
+#pragma once
+
+#include <torch/torch.h>
+
+void f32_bf16w_matmul_torch(const at::Tensor &input, const at::Tensor &weight_bf16, const at::Tensor &bias_bf16, at::Tensor &output, int64_t num_tokens, int64_t num_cols, int64_t num_rows, int64_t threadgroup_size);