Upload folder using huggingface_hub

README.md

@@ -0,0 +1,10 @@
+---
+tags:
+- kernels
+- gptoss
+---
+
+# gptoss_kernels
+
+This is a build for some kernel released by OpenAI in the GPT-OSS repo : https://github.com/openai/gpt-oss
+

build.toml

@@ -0,0 +1,23 @@
+[general]
+name = "gptoss_kernels"
+universal = false
+
+[torch]
+src = [
+  "torch-ext/torch_binding.cpp",
+  "torch-ext/torch_binding.h",
+]
+
+[kernel.gptoss_kernels]
+depends = ["torch"]
+backend = "cuda"
+
+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",
+]

flake.nix

@@ -0,0 +1,13 @@
+{
+  description = "Flake for Torch kernel extension";
+
+  inputs = {
+    kernel-builder.url = "github:huggingface/kernel-builder";
+  };
+
+  outputs = { self, kernel-builder, }:
+    kernel-builder.lib.genFlakeOutputs {
+        inherit self;
+        path = ./.;
+    };
+}

gptoss_kernels/CMakeLists.txt

@@ -0,0 +1,191 @@
+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

@@ -0,0 +1,6 @@
+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

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

@@ -0,0 +1,34 @@
+#!/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/include/gpt-oss.h

@@ -0,0 +1,5 @@
+#pragma once
+
+#include <gpt-oss/macros.h>
+#include <gpt-oss/types.h>
+#include <gpt-oss/functions.h>

gptoss_kernels/include/gpt-oss/functions.h

@@ -0,0 +1,401 @@
+#pragma once
+
+#include <stddef.h>
+#include <stdint.h>
+
+#include <gpt-oss/macros.h>
+#include <gpt-oss/types.h>
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+/*
+ * Creates a Model object from a file in the filesystem.
+ *
+ * @param path Path to the file containing the model in GPT-OSS format.
+ * @param model_out Pointer to the Model object that will be created. Must be released with gptoss_release_model.
+ *
+ * On success, returns gptoss_status_success and saves a pointer to the created Model in the model_out argument.
+ * On failure, returns an error code and stores null pointer in the model_out argument.
+ */
+enum gptoss_status GPTOSS_ABI gptoss_model_create_from_file(
+    const char* path,
+    gptoss_model_t* model_out);
+
+/*
+ * Query the Tokenizer object associated with the Model.
+ *
+ * @param model Pointer to the Model object created by gptoss_model_create_from_file.
+ * @param tokenizer_out Pointer to the variable where the Tokenizer reference will be stored.
+ *
+ * On success, returns gptoss_status_success and stores reference to the Tokenizer object in the tokenizer_out argument.
+ * On failure, returns an error code and stores NULL in the tokenizer_out argument.
+ */
+enum gptoss_status GPTOSS_ABI gptoss_model_get_tokenizer(
+    gptoss_model_t model,
+    gptoss_tokenizer_t* tokenizer_out);
+
+/*
+ * Query the maximum context length supported by the Model.
+ *
+ * @param model Pointer to the Model object created by gptoss_model_create_from_file.
+ * @param max_context_length_out Pointer to the variable where the maximum context length will be stored.
+ *
+ * On success, returns gptoss_status_success and stores maximum context length in the max_context_length_out argument.
+ * On failure, returns an error code and leaves the value specified by max_context_length_out unchanged.
+ */
+enum gptoss_status GPTOSS_ABI gptoss_model_get_max_context_length(
+    gptoss_model_t model,
+    size_t* max_context_length_out);
+
+/*
+ * Increments a Model object's reference count.
+ *
+ * @param model Pointer to the Model object created by gptoss_model_create_from_file.
+ *
+ * On success, returns gptoss_status_success, otherwise returns an error code.
+ */
+enum gptoss_status GPTOSS_ABI gptoss_model_retain(
+    gptoss_model_t model);
+
+/*
+ * Decrements a Model object's reference count and possibly release associated resources.
+ *
+ * @param model Pointer to the Model object created by gptoss_model_create_from_file.
+ *
+ * On success, returns gptoss_status_success, otherwise returns an error code.
+ */
+enum gptoss_status GPTOSS_ABI gptoss_model_release(
+    gptoss_model_t model);
+
+/*
+ * Query the token ID for a special token in the Tokenizer vocabulary.
+ *
+ * @param tokenizer Pointer to the Tokenizer object created by gptoss_model_get_tokenizer.
+ * @param token_type Type of the special token to query an ID for.
+ * @param token_id_out Pointer to the variable where the token ID will be stored.
+ *
+ * On success, returns gptoss_status_success and stores the token ID in the token_id_out argument.
+ * On failure, returns an error code and leaves the value specified by token_id_out unchanged.
+ */
+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);
+
+/*
+ * Query the number of text tokens in the Tokenizer vocabulary.
+ *
+ * @param tokenizer Pointer to the Tokenizer object created by gptoss_model_get_tokenizer.
+ * @param num_text_tokens_out Pointer to the variable where the number of text tokens will be stored.
+ *
+ * On success, returns gptoss_status_success and stores the number of text tokens in the num_text_tokens_out argument.
+ * On failure, returns an error code and leaves the value specified by num_text_tokens_out unchanged.
+ */
+enum gptoss_status GPTOSS_ABI gptoss_tokenizer_get_num_text_tokens(
+    gptoss_tokenizer_t tokenizer,
+    uint32_t* num_text_tokens_out);
+
+/*
+ * Query the number of special tokens in the Tokenizer vocabulary.
+ *
+ * @param tokenizer Pointer to the Tokenizer object created by gptoss_model_get_tokenizer.
+ * @param num_special_tokens_out Pointer to the variable where the number of special tokens will be stored.
+ *
+ * On success, returns gptoss_status_success and stores the number of text tokens in the num_special_tokens_out argument.
+ * On failure, returns an error code and leaves the value specified by num_special_tokens_out unchanged.
+ */
+enum gptoss_status GPTOSS_ABI gptoss_tokenizer_get_num_special_tokens(
+    gptoss_tokenizer_t tokenizer,
+    uint32_t* num_special_tokens_out);
+
+/*
+ * Query the total number of tokens in the Tokenizer vocabulary.
+ *
+ * @param tokenizer Pointer to the Tokenizer object created by gptoss_model_get_tokenizer.
+ * @param num_tokens_out Pointer to the variable where the total number of tokens will be stored.
+ *
+ * On success, returns gptoss_status_success and stores the total number of tokens in the num_special_tokens_out argument.
+ * On failure, returns an error code and leaves the value specified by num_special_tokens_out unchanged.
+ */
+enum gptoss_status GPTOSS_ABI gptoss_tokenizer_get_num_tokens(
+    gptoss_tokenizer_t tokenizer,
+    uint32_t* num_tokens_out);
+
+/*
+ * Convert a text token ID to byte representation.
+ *
+ * @param tokenizer Pointer to the Tokenizer object returned by gptoss_model_get_tokenizer. The lifetime of the returned
+ *                  byte representation would match the lifetime of this Tokenizer object.
+ * @param token_ptr_out Pointer to the variable where the pointer to the byte representation of the token will be
+ *                      stored.
+ * @param token_size_out Pointer to the variable where the size of the byte representation of the token will be stored.
+ *
+ * On success, returns gptoss_status_success and stores pointer and size of the byte representation of the token in the
+ *                     token_ptr_out and token_size_out arguments.
+ * On failure, returns an error code and leaves the values specified in token_ptr_out and token_size_out unchanged.
+ */
+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);
+
+/*
+ * Increments a Tokenizer object's reference count.
+ *
+ * @param tokenizer Pointer to the Tokenizer object returned by gptoss_model_get_tokenizer.
+ *
+ * On success, returns gptoss_status_success, otherwise returns an error code.
+ */
+enum gptoss_status GPTOSS_ABI gptoss_tokenizer_retain(
+    gptoss_tokenizer_t tokenizer);
+
+/*
+ * Decrements a Tokenizer object's reference count and possibly release associated resources.
+ *
+ * @param tokenizer Pointer to the Tokenizer object returned by gptoss_model_get_tokenizer.
+ *
+ * On success, returns gptoss_status_success, otherwise returns an error code.
+ */
+enum gptoss_status GPTOSS_ABI gptoss_tokenizer_release(
+    gptoss_tokenizer_t tokenizer);
+
+/*
+ * Creates a Context object for use with the particular Model object.
+ *
+ * @param model Model object to create a context for.
+ * @param context_length Maximum number of tokens in the context.
+ *                       Specify 0 to use the maximum context length supported by the model.
+ * @param max_batch_size Maximum number of tokens that can be processed in a single batch.
+ *                       Larger values may improve prefill performance, but require more memory.
+ *                       Specify 0 to use the default value.
+ * @param context_out Pointer to the Context object that will be created.
+ *                    Must be released with gptoss_release_context.
+ *
+ * On success, returns gptoss_status_success and saves a pointer to the created Context in the context_out argument.
+ * On failure, returns an error code and stores null pointer in the context_out argument.
+ */
+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);
+
+/*
+ * Query the current number of tokens cached in the Context.
+ *
+ * @param context Pointer to the Context object created by gptoss_context_create.
+ * @param num_tokens_out Pointer to the variable where the current number of cached tokens will be stored.
+ *
+ * On success, returns gptoss_status_success and stores current number of cached tokens in the num_tokens_out argument.
+ * On failure, returns an error code and leaves the value specified by num_tokens_out unchanged.
+ */
+enum gptoss_status GPTOSS_ABI gptoss_context_get_num_tokens(
+    gptoss_context_t context,
+    size_t* num_tokens_out);
+
+/*
+ * Query the maximum number of tokens cached in the Context.
+ *
+ * @param context Pointer to the Context object created by gptoss_context_create.
+ * @param max_tokens_out Pointer to the variable where the maximum number of cached tokens will be stored.
+ *
+ * On success, returns gptoss_status_success and stores maximum number of cached tokens in the max_tokens_out argument.
+ * On failure, returns an error code and leaves the value specified by max_tokens_out unchanged.
+ */
+enum gptoss_status GPTOSS_ABI gptoss_context_get_max_tokens(
+    gptoss_context_t context,
+    size_t* max_tokens_out);
+
+/*
+ * Query the list of token IDs cached in the Context.
+ *
+ * @param context Pointer to the Context object created by gptoss_context_create.
+ * @param tokens_out Pointer to the array where up to max_tokens_out of cached tokens will be stored.
+ * @param max_tokens Maximum capacity of the buffer specified by tokens_out.
+ * @param num_tokens_out Pointer to the variable where the actual number of cached tokens will be stored.
+ *                       This value can exceed max_tokens if the buffer capacity is insufficient.
+ *
+ * On success, returns gptoss_status_success and stores cached token IDs in the tokens_out argument and the number of
+ * cached tokens in the num_tokens_out argument.
+ * On failure, returns an error code and leaves the values specified by tokens_out and num_tokens_out unchanged.
+ */
+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);
+
+/*
+ * Tokenize and appends a character string to the Context object.
+ *
+ * @param context Context object created by gptoss_context_create.
+ * @param text Pointer to the character string to tokenizer and append.
+ * @param text_length Length of the string, in chars.
+ * @param num_tokens_out Optional pointer to the variable where the number of appended tokens will be stored. Ignored if a null pointer is provided.
+ *
+ * On success, returns gptoss_status_success, otherwise returns an error code.
+ */
+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);
+
+/*
+ * Appends a list of tokens to the context.
+ *
+ * @param context Context object created by gptoss_context_create.
+ * @param num_tokens Number of tokens to be appended.
+ * @param tokens Pointer to the array of tokens to be appended.
+ *
+ * On success, returns gptoss_status_success, otherwise returns an error code.
+ */
+enum gptoss_status GPTOSS_ABI gptoss_context_append_tokens(
+    gptoss_context_t context,
+    size_t num_tokens,
+    const uint32_t* tokens);
+
+/*
+ * Resets the context, clearing its state.
+ *
+ * @param context Context object created by gptoss_context_create.
+ *
+ * On success, returns gptoss_status_success, otherwise returns an error code.
+ */
+enum gptoss_status GPTOSS_ABI gptoss_context_reset(
+    gptoss_context_t context);
+
+/*
+ * Pre-process the tokens in the Context and generate probability distribution over the next token.
+ *
+ * @param context Context object created by gptoss_context_create.
+ *
+ * On success, returns gptoss_status_success, otherwise returns an error code.
+ */
+enum gptoss_status GPTOSS_ABI gptoss_context_process(
+    gptoss_context_t context);
+
+/*
+ * Generate a token probability distribution over the next token conditioned on the Context.
+ *
+ * @param context Context object created by gptoss_context_create.
+ * @param temperature Sampling temperature. Must be non-negative.
+ * @param seed Random number generator seed to use for sampling.
+ * @param token_out Pointer to the variable where the token ID will be stored.
+ *
+ * On success, returns gptoss_status_success, otherwise returns an error code.
+ */
+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);
+
+/*
+ * Increments a Context object's reference count.
+ *
+ * @param context Pointer to the Context object created by gptoss_create_context.
+ *
+ * On success, returns gptoss_status_success, otherwise returns an error code.
+ */
+enum gptoss_status GPTOSS_ABI gptoss_context_retain(
+    gptoss_context_t context);
+
+/*
+ * Decrements a Context object's reference count and possibly release associated resources.
+ *
+ * @param context Pointer to the Context object created by gptoss_create_context.
+ *
+ * On success, returns gptoss_status_success, otherwise returns an error code.
+ */
+enum gptoss_status GPTOSS_ABI gptoss_context_release(
+    gptoss_context_t context);
+
+/*
+ * Creates a Sampler object.
+ *
+ * @param sampler_out Pointer to the Sampler object that will be created.
+ *                    Must be released with gptoss_sampler_release.
+ *
+ * On success, returns gptoss_status_success and saves a pointer to the created Sampler in the sampler_out argument.
+ * On failure, returns an error code and stores a null pointer in the sampler_out argument.
+ */
+enum gptoss_status GPTOSS_ABI gptoss_sampler_create(
+    gptoss_sampler_t* sampler_out);
+
+/*
+ * Sets the sampling temperature for the Sampler.
+ *
+ * @param sampler Sampler object created by gptoss_sampler_create.
+ * @param temperature Temperature value to be set. Must be in the [0.0, 1.0] range.
+ *
+ * On success, returns gptoss_status_success, otherwise returns an error code.
+ */
+enum gptoss_status GPTOSS_ABI gptoss_sampler_set_temperature(
+    gptoss_sampler_t sampler,
+    float temperature);
+
+/*
+ * Sets the Top-P nucleus sampling parameter for the Sampler.
+ *
+ * @param sampler Sampler object created by gptoss_sampler_create.
+ * @param top_p Top-P value to be set. Must be in the (0.0, 1.0] range.
+ *
+ * On success, returns gptoss_status_success, otherwise returns an error code.
+ */
+enum gptoss_status GPTOSS_ABI gptoss_sampler_set_top_p(
+    gptoss_sampler_t sampler,
+    float top_p);
+
+/*
+ * Sets the presence penalty for the Sampler.
+ *
+ * @param sampler Sampler object created by gptoss_sampler_create.
+ * @param presence_penalty Presence penalty value to be set. Must be in the [-2.0, 2.0] range.
+ *
+ * On success, returns gptoss_status_success, otherwise returns an error code.
+ */
+enum gptoss_status GPTOSS_ABI gptoss_sampler_set_presence_penalty(
+    gptoss_sampler_t sampler,
+    float presence_penalty);
+
+/*
+ * Sets the frequency penalty for the Sampler.
+ *
+ * @param sampler Sampler object created by gptoss_sampler_create.
+ * @param frequency_penalty Frequency penalty value to be set. Must be in the [-2.0, 2.0] range.
+ *
+ * On success, returns gptoss_status_success, otherwise returns an error code.
+ */
+enum gptoss_status GPTOSS_ABI gptoss_sampler_set_frequency_penalty(
+    gptoss_sampler_t sampler,
+    float frequency_penalty);
+
+/*
+ * Increments a Sampler object's reference count.
+ *
+ * @param sampler Pointer to the Sampler object created by gptoss_sampler_create.
+ *
+ * On success, returns gptoss_status_success, otherwise returns an error code.
+ */
+enum gptoss_status GPTOSS_ABI gptoss_sampler_retain(
+    gptoss_sampler_t sampler);
+
+/*
+ * Decrements a Sampler object's reference count and possibly releases associated resources.
+ *
+ * @param sampler Pointer to the Sampler object created by gptoss_sampler_create.
+ *
+ * On success, returns gptoss_status_success, otherwise returns an error code.
+ */
+enum gptoss_status GPTOSS_ABI gptoss_sampler_release(
+    gptoss_sampler_t sampler);
+
+#ifdef __cplusplus
+}  // extern "C"
+#endif

gptoss_kernels/include/gpt-oss/macros.h

@@ -0,0 +1,5 @@
+#pragma once
+
+#ifndef GPTOSS_ABI
+    #define GPTOSS_ABI
+#endif  // GPTOSS_ABI

gptoss_kernels/include/gpt-oss/types.h

@@ -0,0 +1,62 @@
+#pragma once
+
+/*
+ * Status codes returned by GPT-OSS API functions.
+ */
+enum gptoss_status {
+    gptoss_status_success = 0,
+    gptoss_status_invalid_argument = 1,
+    gptoss_status_unsupported_argument = 2,
+    gptoss_status_invalid_state = 3,
+    gptoss_status_io_error = 4,
+    gptoss_status_insufficient_memory = 5,
+    gptoss_status_insufficient_resources = 6,
+    gptoss_status_unsupported_system = 7,
+    gptoss_status_context_overflow = 8,
+};
+
+enum gptoss_special_token {
+    gptoss_special_token_invalid = 0,
+    gptoss_special_token_return = 1,
+    gptoss_special_token_start = 2,
+    gptoss_special_token_message = 3,
+    gptoss_special_token_end = 4,
+    gptoss_special_token_refusal = 5,
+    gptoss_special_token_constrain = 6,
+    gptoss_special_token_channel = 7,
+    gptoss_special_token_call = 8,
+    gptoss_special_token_untrusted = 9,
+    gptoss_special_token_end_untrusted = 10,
+    gptoss_special_token_max,
+};
+
+/*
+ * Model object is an opaque container comprised of:
+ * - Weights
+ * - Temporary buffers required to run the model
+ * - Any other resources requires to run the model
+ */
+typedef struct gptoss_model* gptoss_model_t;
+
+typedef struct gptoss_tokenizer* gptoss_tokenizer_t;
+
+/*
+ * Context is an opaque container comprised of:
+ * - Input tokens
+ * - Distribution over the output tokens
+ * - KV cache
+ * 
+ * Multiple contexts can be created and used with the same model.
+ */
+typedef struct gptoss_context* gptoss_context_t;
+
+/*
+ * Sampler is an opaque container for sampling parameters:
+ * - Temperature
+ * - Top-p (nucleus sampling)
+ * - Frequency penalty
+ * - Presence penalty
+ *
+ * Multiple samplers can be created and used with the same context.
+ */
+typedef struct gptoss_sampler* gptoss_sampler_t;

gptoss_kernels/source/accumulate.metal

@@ -0,0 +1,59 @@
+#include <metal_integer>
+#include <metal_math>
+
+#include <internal/kernel-args.h>
+
+#pragma METAL fp math_mode(safe)
+#pragma METAL fp contract(off)
+
+
+kernel void gptoss_f32_accumulate_e4(
+    constant gptoss_accumulate_args& args [[ buffer(0) ]],
+    const device float4* input [[ buffer(1) ]],
+    const device gptoss_expert_prediction* expert [[ buffer(2) ]],
+    device float4* output [[ buffer(3) ]],
+    const device gptoss_control* control [[ buffer(4) ]],
+    uint2 gid [[threadgroup_position_in_grid]],
+    uint tid [[thread_index_in_threadgroup]],
+    uint2 threadgroup_size [[ threads_per_threadgroup ]])
+{
+    const uint num_active_experts = 4;
+    if (control->abort != 0) {
+        return;
+    }
+
+    const uint num_vecs_per_threadgroup = args.num_vecs_per_threadgroup;
+    const uint threadgroup_start = gid.x * num_vecs_per_threadgroup;
+    const uint num_vecs = args.num_vecs;
+    const uint threadgroup_end = metal::min(threadgroup_start + num_vecs_per_threadgroup, num_vecs);
+    const uint thread_start = threadgroup_start + tid;
+    uint num_iter = static_cast<uint>((threadgroup_end - thread_start + (threadgroup_size.x - 1)) / threadgroup_size.x);
+
+    const uint num_vecs_per_expert = args.num_vecs_per_expert;
+    const float scale0 = expert[gid.y * num_active_experts + 0].score;
+    const device float4* input0 = input + gid.y * num_vecs + thread_start;
+    const float scale1 = expert[gid.y * num_active_experts + 1].score;
+    const device float4* input1 = input0 + num_vecs_per_expert;
+    const float scale2 = expert[gid.y * num_active_experts + 2].score;
+    const device float4* input2 = input1 + num_vecs_per_expert;
+    const float scale3 = expert[gid.y * num_active_experts + 3].score;
+    const device float4* input3 = input2 + num_vecs_per_expert;
+    output += gid.y * num_vecs + thread_start;
+    for (; num_iter != 0; num_iter--) {
+        float4 acc = *output;
+        const float4 val0 = *input0;
+        const float4 val1 = *input1;
+        const float4 val2 = *input2;
+        const float4 val3 = *input3;
+        input0 += threadgroup_size.x;
+        acc = metal::fma(val0, scale0, acc);
+        input1 += threadgroup_size.x;
+        acc = metal::fma(val1, scale1, acc);
+        input2 += threadgroup_size.x;
+        acc = metal::fma(val2, scale2, acc);
+        input3 += threadgroup_size.x;
+        acc = metal::fma(val3, scale3, acc);
+        *output = acc;
+        output += threadgroup_size.x;
+    }
+}

gptoss_kernels/source/context.c

@@ -0,0 +1,1115 @@
+#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/convert.metal

@@ -0,0 +1,64 @@
+#include <metal_integer>
+
+#include <internal/kernel-args.h>
+
+#pragma METAL fp math_mode(safe)
+#pragma METAL fp contract(off)
+
+
+kernel void gptoss_mf4_f32_convert(
+    constant gptoss_convert_args& args [[ buffer(0) ]],
+    const device uint4* blocks [[ buffer(1) ]],
+    const device uchar* scales [[ buffer(2) ]],
+    device float4* output [[ buffer(3) ]],
+    uint gid [[threadgroup_position_in_grid]],
+    uint tid [[thread_position_in_threadgroup]],
+    uint threadgroup_size [[ threads_per_threadgroup ]])
+{
+    const ulong num_vecs_per_threadgroup = args.num_vecs_per_threadgroup;
+    const ulong threadgroup_start = gid * num_vecs_per_threadgroup;
+    const ulong threadgroup_end = metal::min(threadgroup_start + num_vecs_per_threadgroup, args.num_vecs);
+    const ulong thread_start = threadgroup_start + tid;
+    uint num_iter = static_cast<uint>((threadgroup_end - thread_start + (threadgroup_size - 1)) / threadgroup_size);
+
+    blocks += thread_start;
+    scales += thread_start;
+    output += 8 * thread_start;
+    for (; num_iter != 0; num_iter--) {
+        const uint4 block = *blocks;
+        const float scale = as_type<float>((static_cast<uint>(*scales) + 14) << 23);
+        uint4 block02468ACEGIKMOQSU = block + block;
+        uint4 block13579BDFHJLNPRTV = block >> 3;
+        block02468ACEGIKMOQSU &= 0x1E1E1E1Eu;
+        block13579BDFHJLNPRTV &= 0x1E1E1E1Eu;
+        block02468ACEGIKMOQSU += 0x70707070u;
+        block13579BDFHJLNPRTV += 0x70707070u;
+        block02468ACEGIKMOQSU &= 0x8E8E8E8Eu;
+        block13579BDFHJLNPRTV &= 0x8E8E8E8Eu;
+        const uint4 block26AEIMQU = block02468ACEGIKMOQSU & 0xFF00FF00u;
+        const uint4 block048CGKOS = (block02468ACEGIKMOQSU << 8) & 0xFF00FF00u;
+        const uint4 block37BFJNRV = block13579BDFHJLNPRTV & 0xFF00FF00u;
+        const uint4 block159DHLPT = (block13579BDFHJLNPRTV << 8) & 0xFF00FF00u;
+        const float4 block048C = static_cast<float4>(as_type<half4>(block048CGKOS.xy)) * scale;
+        const float4 blockGKOS = static_cast<float4>(as_type<half4>(block048CGKOS.zw)) * scale;
+        const float4 block26AE = static_cast<float4>(as_type<half4>(block26AEIMQU.xy)) * scale;
+        const float4 blockIMQU = static_cast<float4>(as_type<half4>(block26AEIMQU.zw)) * scale;
+        const float4 block159D = static_cast<float4>(as_type<half4>(block159DHLPT.xy)) * scale;
+        const float4 blockHLPT = static_cast<float4>(as_type<half4>(block159DHLPT.zw)) * scale;
+        const float4 block37BF = static_cast<float4>(as_type<half4>(block37BFJNRV.xy)) * scale;
+        const float4 blockJNRV = static_cast<float4>(as_type<half4>(block37BFJNRV.zw)) * scale;
+
+        output[0] = (float4) { block048C.x, block159D.x, block26AE.x, block37BF.x };
+        output[1] = (float4) { block048C.y, block159D.y, block26AE.y, block37BF.y };
+        output[2] = (float4) { block048C.z, block159D.z, block26AE.z, block37BF.z };
+        output[3] = (float4) { block048C.w, block159D.w, block26AE.w, block37BF.w };
+        output[4] = (float4) { blockGKOS.x, blockHLPT.x, blockIMQU.x, blockJNRV.x };
+        output[5] = (float4) { blockGKOS.y, blockHLPT.y, blockIMQU.y, blockJNRV.y };
+        output[6] = (float4) { blockGKOS.z, blockHLPT.z, blockIMQU.z, blockJNRV.z };
+        output[7] = (float4) { blockGKOS.w, blockHLPT.w, blockIMQU.w, blockJNRV.w };
+
+        blocks += threadgroup_size;
+        scales += threadgroup_size;
+        output += 8 * threadgroup_size;
+    }
+}

gptoss_kernels/source/embeddings.metal

@@ -0,0 +1,29 @@
+#include <internal/kernel-args.h>
+
+#pragma METAL fp math_mode(safe)
+#pragma METAL fp contract(off)
+
+
+kernel void gptoss_bf16_f32_embeddings(
+    constant gptoss_embeddings_args& args [[ buffer(0) ]],
+    const device uint* tokens [[ buffer(1) ]],
+    const device bfloat4* weights [[ buffer(2) ]],
+    device float4* output [[ buffer(3) ]],
+    const device gptoss_control* control [[ buffer(4) ]],
+    uint gid [[threadgroup_position_in_grid]],
+    uint tid [[thread_position_in_threadgroup]],
+    uint threadgroup_size [[ threads_per_threadgroup ]])
+{
+    if (control->abort != 0) {
+        return;
+    }
+
+    const uint t = tokens[gid];
+
+    weights += t * args.num_vecs;
+    output += gid * args.num_vecs;
+    for (uint i = tid; i < args.num_vecs; i += threadgroup_size) {
+        const bfloat4 w = weights[i];
+        output[i] = static_cast<float4>(w);
+    }
+}

gptoss_kernels/source/expert_routing_metadata.metal

@@ -0,0 +1,41 @@
+#include <internal/kernel-args.h>
+#include <metal_integer>
+#include <metal_math>
+#include <metal_stdlib>
+
+constant uint kMaxExperts = 128;
+
+kernel void gptoss_f32_expert_routing_metadata(
+    constant gptoss_expert_routing_metadata_args& args [[ buffer(0) ]],
+    const device gptoss_expert_prediction* __restrict__ expert_predictions [[ buffer(1) ]],
+    device uint* __restrict__ expert_offsets [[ buffer(2) ]],
+    device uint* __restrict__ intra_expert_offsets [[ buffer(3) ]],
+    uint tg_size [[threads_per_threadgroup]],
+    uint tid [[thread_position_in_threadgroup]]) 
+{
+    assert(args.num_experts <= kMaxExperts);
+    // Create threadgroup mem and initialize it to 0.
+    threadgroup metal::atomic_uint tg_counts[kMaxExperts];
+    for (uint e = tid; e < args.num_experts; e += tg_size) {
+        metal::atomic_store_explicit(&tg_counts[e], 0u, metal::memory_order_relaxed);
+    }
+
+    threadgroup_barrier(metal::mem_flags::mem_threadgroup);
+
+    for (uint i = tid; i < args.tokens; i += tg_size) {
+        const uint e = expert_predictions[i].expert_id;
+        const uint r = metal::atomic_fetch_add_explicit(&tg_counts[e], 1u, metal::memory_order_relaxed);
+        intra_expert_offsets[i] = r;
+    }
+    threadgroup_barrier(metal::mem_flags::mem_threadgroup);
+
+    if (tid == 0) {
+        uint total = 0;
+        for (uint e = 0; e < args.num_experts; ++e) {
+            const uint bin = metal::atomic_load_explicit(&tg_counts[e], metal::memory_order_relaxed);
+            expert_offsets[e] = total;
+            total += bin;
+        }
+        expert_offsets[args.num_experts] = total;
+    }
+}

gptoss_kernels/source/gather_and_accumulate.metal

@@ -0,0 +1,74 @@
+#include <internal/kernel-args.h>
+#include <metal_integer>
+#include <metal_math>
+#include <metal_stdlib>
+
+// TODO(ibrahim): This is not optimal as each thread only gathers and accumulates a single float4. To amortize the
+// cost of reading the expert, offset and scales for a token, we should let each thread gather and accumulate several
+// float4s.
+kernel void gptoss_f32_gather_and_accumulate_e4(
+    constant gptoss_gather_args& args [[ buffer(0) ]],
+    const device float* in [[ buffer(1) ]],
+    const device gptoss_expert_prediction* __restrict__ expert_predictions [[ buffer(2) ]],
+    const device uint* expert_offsets [[ buffer(3) ]],
+    const device uint* intra_expert_offsets [[ buffer(4) ]],
+    device float* out [[ buffer(5) ]],
+    uint3 gid [[thread_position_in_grid]]) 
+{
+    const uint T = args.tokens;
+    const uint k = args.active_experts_per_token;
+    const uint D = args.token_stride;
+
+    assert((D & 3u) == 0);
+    assert(k == 4);
+
+    const uint row = gid.y;
+    if (row >= T) {
+        return;
+    }
+
+    const uint col_vec4 = gid.x;
+    const uint col = col_vec4 * 4u;
+    if (col >= D) {
+        return;
+    }
+
+    device float4* dst4 = reinterpret_cast<device float4*>(out + row * D + col);
+
+    const uint base = row * k;
+    const gptoss_expert_prediction expert0 = expert_predictions[base];
+    const gptoss_expert_prediction expert1 = expert_predictions[base + 1];
+    const gptoss_expert_prediction expert2 = expert_predictions[base + 2];
+    const gptoss_expert_prediction expert3 = expert_predictions[base + 3];
+    const uint expert0_id = expert0.expert_id;
+    const uint expert1_id = expert1.expert_id;
+    const uint expert2_id = expert2.expert_id;
+    const uint expert3_id = expert3.expert_id;
+    const float scale0 = expert0.score;
+    const float scale1 = expert1.score;
+    const float scale2 = expert2.score;
+    const float scale3 = expert3.score;
+    const uint4 current_intra_expert_offsets =
+        *reinterpret_cast<const device uint4*>(&intra_expert_offsets[base]);
+    // Get the row indices for the current expert ids
+    const uint r0 = expert_offsets[expert0_id] + current_intra_expert_offsets.x;
+    const uint r1 = expert_offsets[expert1_id] + current_intra_expert_offsets.y;
+    const uint r2 = expert_offsets[expert2_id] + current_intra_expert_offsets.z;
+    const uint r3 = expert_offsets[expert3_id] + current_intra_expert_offsets.w;
+
+    const device float4* src0 =
+        reinterpret_cast<const device float4*>(in + r0 * D + col);
+    const device float4* src1 =
+        reinterpret_cast<const device float4*>(in + r1 * D + col);
+    const device float4* src2 =
+        reinterpret_cast<const device float4*>(in + r2 * D + col);
+    const device float4* src3 =
+        reinterpret_cast<const device float4*>(in + r3 * D + col);
+
+    float4 acc = *dst4;
+    acc = metal::fma(*src0, scale0, acc);
+    acc = metal::fma(*src1, scale1, acc);
+    acc = metal::fma(*src2, scale2, acc);
+    acc = metal::fma(*src3, scale3, acc);
+    *dst4 = acc;
+}

gptoss_kernels/source/generate.c

@@ -0,0 +1,317 @@
+#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/datatype.h

@@ -0,0 +1,41 @@
+#pragma once
+
+#include <stdint.h>
+
+#include <internal/macros.h>
+
+
+typedef struct GPTOSS_DENSELY_PACKED_STRUCTURE {
+    GPTOSS_ALIGN(2) uint16_t bits;
+} gptoss_bfloat16;
+static_assert(sizeof(gptoss_bfloat16) == 2, "bfloat16 size is not 2 bytes");
+
+
+typedef struct GPTOSS_DENSELY_PACKED_STRUCTURE {
+    GPTOSS_ALIGN(2) uint16_t bits;
+} gptoss_float16;
+static_assert(sizeof(gptoss_float16) == 2, "float16 size is not 2 bytes");
+
+
+typedef struct GPTOSS_DENSELY_PACKED_STRUCTURE {
+    GPTOSS_ALIGN(1) uint8_t bits;
+} gptoss_float8ue8m0;
+static_assert(sizeof(gptoss_float8ue8m0) == 1, "gptoss_float8ue8m0 size is not 1 bytes");
+
+
+typedef struct GPTOSS_DENSELY_PACKED_STRUCTURE {
+    GPTOSS_ALIGN(1) uint8_t bits;
+} gptoss_float8e5m2;
+static_assert(sizeof(gptoss_float8e5m2) == 1, "float8e5m2 size is not 1 bytes");
+
+
+typedef struct GPTOSS_DENSELY_PACKED_STRUCTURE {
+    GPTOSS_ALIGN(1) uint8_t bits;
+} gptoss_float8e4m3;
+static_assert(sizeof(gptoss_float8e4m3) == 1, "gptoss_float8e4m3 size is not 1 bytes");
+
+
+typedef struct GPTOSS_DENSELY_PACKED_STRUCTURE {
+    GPTOSS_ALIGN(1) uint8_t bits;
+} gptoss_float4e2m1x2;
+static_assert(sizeof(gptoss_float4e2m1x2) == 1, "gptoss_float4e2m1x2 size is not 1 bytes");

gptoss_kernels/source/include/internal/datatype.hpp

@@ -0,0 +1,87 @@
+#pragma once
+
+#include <bit>
+
+#include <internal/datatype.h>
+
+
+namespace gptoss {
+
+template <typename WideT, typename NarrowT>
+WideT upcast(NarrowT);
+
+template <>
+inline float upcast<float>(gptoss_bfloat16 bf16_value) {
+    const uint32_t bits = static_cast<uint32_t>(bf16_value.bits) << 16;
+    return std::bit_cast<float>(bits);
+}
+
+template <>
+inline float upcast<float>(gptoss_float16 fp16_value) {
+    return static_cast<float>(std::bit_cast<_Float16>(fp16_value.bits));
+}
+
+template <>
+inline float upcast<float>(gptoss_float8e4m3 fp8_value) {
+    static constexpr uint16_t fp8e4m3_to_fp32[256] = {
+        0x0000, 0x3B00, 0x3B80, 0x3BC0, 0x3C00, 0x3C20, 0x3C40, 0x3C60,
+        0x3C80, 0x3C90, 0x3CA0, 0x3CB0, 0x3CC0, 0x3CD0, 0x3CE0, 0x3CF0,
+        0x3D00, 0x3D10, 0x3D20, 0x3D30, 0x3D40, 0x3D50, 0x3D60, 0x3D70,
+        0x3D80, 0x3D90, 0x3DA0, 0x3DB0, 0x3DC0, 0x3DD0, 0x3DE0, 0x3DF0,
+        0x3E00, 0x3E10, 0x3E20, 0x3E30, 0x3E40, 0x3E50, 0x3E60, 0x3E70,
+        0x3E80, 0x3E90, 0x3EA0, 0x3EB0, 0x3EC0, 0x3ED0, 0x3EE0, 0x3EF0,
+        0x3F00, 0x3F10, 0x3F20, 0x3F30, 0x3F40, 0x3F50, 0x3F60, 0x3F70,
+        0x3F80, 0x3F90, 0x3FA0, 0x3FB0, 0x3FC0, 0x3FD0, 0x3FE0, 0x3FF0,
+        0x4000, 0x4010, 0x4020, 0x4030, 0x4040, 0x4050, 0x4060, 0x4070,
+        0x4080, 0x4090, 0x40A0, 0x40B0, 0x40C0, 0x40D0, 0x40E0, 0x40F0,
+        0x4100, 0x4110, 0x4120, 0x4130, 0x4140, 0x4150, 0x4160, 0x4170,
+        0x4180, 0x4190, 0x41A0, 0x41B0, 0x41C0, 0x41D0, 0x41E0, 0x41F0,
+        0x4200, 0x4210, 0x4220, 0x4230, 0x4240, 0x4250, 0x4260, 0x4270,
+        0x4280, 0x4290, 0x42A0, 0x42B0, 0x42C0, 0x42D0, 0x42E0, 0x42F0,
+        0x4300, 0x4310, 0x4320, 0x4330, 0x4340, 0x4350, 0x4360, 0x4370,
+        0x4380, 0x4390, 0x43A0, 0x43B0, 0x43C0, 0x43D0, 0x43E0, 0x7FF0,
+        0x8000, 0xBB00, 0xBB80, 0xBBC0, 0xBC00, 0xBC20, 0xBC40, 0xBC60,
+        0xBC80, 0xBC90, 0xBCA0, 0xBCB0, 0xBCC0, 0xBCD0, 0xBCE0, 0xBCF0,
+        0xBD00, 0xBD10, 0xBD20, 0xBD30, 0xBD40, 0xBD50, 0xBD60, 0xBD70,
+        0xBD80, 0xBD90, 0xBDA0, 0xBDB0, 0xBDC0, 0xBDD0, 0xBDE0, 0xBDF0,
+        0xBE00, 0xBE10, 0xBE20, 0xBE30, 0xBE40, 0xBE50, 0xBE60, 0xBE70,
+        0xBE80, 0xBE90, 0xBEA0, 0xBEB0, 0xBEC0, 0xBED0, 0xBEE0, 0xBEF0,
+        0xBF00, 0xBF10, 0xBF20, 0xBF30, 0xBF40, 0xBF50, 0xBF60, 0xBF70,
+        0xBF80, 0xBF90, 0xBFA0, 0xBFB0, 0xBFC0, 0xBFD0, 0xBFE0, 0xBFF0,
+        0xC000, 0xC010, 0xC020, 0xC030, 0xC040, 0xC050, 0xC060, 0xC070,
+        0xC080, 0xC090, 0xC0A0, 0xC0B0, 0xC0C0, 0xC0D0, 0xC0E0, 0xC0F0,
+        0xC100, 0xC110, 0xC120, 0xC130, 0xC140, 0xC150, 0xC160, 0xC170,
+        0xC180, 0xC190, 0xC1A0, 0xC1B0, 0xC1C0, 0xC1D0, 0xC1E0, 0xC1F0,
+        0xC200, 0xC210, 0xC220, 0xC230, 0xC240, 0xC250, 0xC260, 0xC270,
+        0xC280, 0xC290, 0xC2A0, 0xC2B0, 0xC2C0, 0xC2D0, 0xC2E0, 0xC2F0,
+        0xC300, 0xC310, 0xC320, 0xC330, 0xC340, 0xC350, 0xC360, 0xC370,
+        0xC380, 0xC390, 0xC3A0, 0xC3B0, 0xC3C0, 0xC3D0, 0xC3E0, 0xFFF0,
+    };
+    const gptoss_bfloat16 bf16_value{.bits = fp8e4m3_to_fp32[fp8_value.bits]};
+    return upcast<float>(bf16_value);
+}
+
+template <>
+inline double upcast<double>(float fp32_value) {
+    return static_cast<double>(fp32_value);
+}
+
+template <>
+inline double upcast<double>(gptoss_bfloat16 bf16_value) {
+    const float fp32_value = upcast<float>(bf16_value);
+    return upcast<double>(fp32_value);
+}
+
+template <>
+inline double upcast<double>(gptoss_float16 fp16_value) {
+    const float fp32_value = upcast<float>(fp16_value);
+    return upcast<double>(fp32_value);
+}
+
+template <>
+inline double upcast<double>(gptoss_float8e4m3 fp8_value) {
+    const float fp32_value = upcast<float>(fp8_value);
+    return upcast<double>(fp32_value);
+}
+
+}  // namespace gptoss

gptoss_kernels/source/include/internal/kernel-args.h

@@ -0,0 +1,201 @@
+#pragma once
+
+#if !defined(__METAL_VERSION__)
+#include <stdint.h>
+#endif
+
+// TODO(ibahmed): specalize using metal function constants.
+#define QKV_Bm 64
+#define QKV_Bn 64
+#define QKV_Bk 32
+#define QKV_Sg_Bm 32
+#define QKV_Sg_Bn 32
+
+#define ATTN_OUTPUT_Bm 32
+#define ATTN_OUTPUT_Bn 64
+#define ATTN_OUTPUT_Bk 64
+#define ATTN_OUTPUT_Sg_Bm 32
+#define ATTN_OUTPUT_Sg_Bn 16
+
+#define MLP_GATE_Bm 64
+#define MLP_GATE_Bn 16
+#define MLP_GATE_Bk 64
+#define MLP_GATE_Sg_Bm 16
+#define MLP_GATE_Sg_Bn 16
+
+#define MOE_DENSE_MATMUL_SWIGLU_Bm 32
+#define MOE_DENSE_MATMUL_SWIGLU_Bn 64
+#define MOE_DENSE_MATMUL_SWIGLU_Bk 16
+#define MOE_DENSE_MATMUL_SWIGLU_Sg_Bm 32
+#define MOE_DENSE_MATMUL_SWIGLU_Sg_Bn 16
+
+#define MOE_DENSE_MATMUL_Bm 32
+#define MOE_DENSE_MATMUL_Bn 64
+#define MOE_DENSE_MATMUL_Bk 16
+#define MOE_DENSE_MATMUL_Sg_Bm 32
+#define MOE_DENSE_MATMUL_Sg_Bn 16
+
+struct gptoss_expert_prediction {
+    uint32_t expert_id;
+    float score;
+};
+
+struct gptoss_control {
+    uint32_t abort;
+};
+
+struct gptoss_topk_args {
+    uint32_t num_vecs_per_token;
+};
+
+struct gptoss_sdpa_args {
+    uint32_t qkv_dim;
+    uint32_t num_kv_tokens;
+    uint32_t kv_stride;
+    uint32_t window;
+};
+
+struct gptoss_u32_fill_random_args {
+    uint64_t num_vecs_per_threadgroup;
+    uint64_t num_vecs;
+    uint64_t offset;
+    uint64_t seed;
+};
+
+struct gptoss_f32_fill_random_args {
+    uint64_t num_vecs_per_threadgroup;
+    uint64_t num_vecs;
+    uint64_t offset;
+    uint64_t seed;
+    float scale;
+    float bias;
+};
+
+struct gptoss_accumulate_args {
+    uint32_t num_vecs_per_expert;
+    uint32_t num_vecs_per_threadgroup;
+    uint32_t num_vecs;
+};
+
+struct gptoss_convert_args {
+    uint64_t num_vecs_per_threadgroup;
+    uint64_t num_vecs;
+};
+
+struct gptoss_embeddings_args {
+    uint32_t num_vecs;
+};
+
+struct gptoss_rmsnorm_args {
+    uint32_t num_vecs;
+    float num_channels;
+    float epsilon;
+};
+
+struct gptoss_matmul_args {
+    uint32_t num_column_vecs;
+    uint32_t num_rows;
+    uint32_t add;
+};
+
+struct gptoss_dense_matmul_args {
+    uint32_t m;
+    uint32_t n;
+    uint32_t k;
+};
+
+struct gptoss_scatter_args {
+    uint32_t tokens;
+    uint32_t active_experts_per_token;
+    uint32_t token_stride;
+};
+
+struct gptoss_moe_dense_matmul_swiglu_args {
+    uint32_t k;
+    uint32_t n;
+    uint32_t weight_blocks_expert_stride_bytes;
+    uint32_t weight_scales_expert_stride_bytes;
+    uint32_t bias_expert_stride_bytes;
+    float swiglu_min;
+    float swiglu_max;
+};
+struct gptoss_moe_dense_matmul_args {
+    uint32_t k;
+    uint32_t n;
+    uint32_t weight_blocks_expert_stride_bytes;
+    uint32_t weight_scales_expert_stride_bytes;
+    uint32_t bias_expert_stride_bytes;
+};
+
+struct gptoss_expert_routing_metadata_args {
+uint32_t tokens;
+    uint32_t num_experts;
+};
+
+struct gptoss_gather_args {
+    uint32_t tokens;
+    uint32_t active_experts_per_token;
+    uint32_t token_stride;
+};
+
+struct gptoss_unembedding_args {
+    uint32_t num_column_vecs;
+    uint32_t num_rows_per_threadgroup;
+    uint32_t num_rows;
+};
+
+struct gptoss_moe_matmul_swiglu_args {
+    uint32_t num_column_vecs;
+    uint32_t num_rows;
+    uint32_t num_active_experts;
+    uint32_t weight_expert_stride;  // in bytes
+    uint32_t output_expert_stride;  // in elements
+    float swiglu_min;
+    float swiglu_max;
+};
+
+struct gptoss_moe_matmul_args {
+    uint32_t num_column_vecs;
+    uint32_t num_rows;
+    uint32_t num_active_experts;
+    uint32_t input_expert_stride;  // in blocks of 32 elements
+    uint32_t weight_expert_stride;  // in bytes
+    uint32_t output_expert_stride;  // in elements
+};
+
+struct gptoss_rope_args {
+    uint32_t token_stride;
+    uint32_t token_offset;
+    float freq_scale;
+    float interpolation_scale;
+    float yarn_offset;
+    float yarn_scale;
+    float yarn_multiplier;
+};
+
+struct gptoss_qkv_args {
+    uint32_t num_column_vecs;
+    uint32_t num_rows;
+    uint32_t token_offset;
+    float freq_scale;
+    float interpolation_scale;
+    float yarn_offset;
+    float yarn_scale;
+    float yarn_multiplier;
+    uint32_t max_tokens;
+};
+
+struct gptoss_softmax_args {
+    uint32_t num_vecs;
+    uint32_t num_vecs_per_threadgroup;
+    uint32_t max_threadgroups;
+    float temperature;
+};
+
+struct gptoss_sample_args {
+    uint64_t rng_seed;
+    uint32_t rng_offset;
+    uint32_t num_blocks;
+    uint32_t num_dims;
+    uint32_t num_dims_per_block;
+};

gptoss_kernels/source/include/internal/log.h

@@ -0,0 +1,20 @@
+#pragma once
+
+#include <stdarg.h>
+
+
+void gptoss_format_log(const char* format, va_list args);
+
+__attribute__((__format__(__printf__, 1, 2)))
+inline static void gptoss_log(const char* format, ...) {
+    va_list args;
+    va_start(args, format);
+    gptoss_format_log(format, args);
+    va_end(args);
+}
+
+#define GPTOSS_LOG_ERROR(message, ...) \
+    gptoss_log("Error: " message "\n", ##__VA_ARGS__)
+
+#define GPTOSS_LOG_WARNING(message, ...) \
+    gptoss_log("Warning: " message "\n", ##__VA_ARGS__)

gptoss_kernels/source/include/internal/macros.h

@@ -0,0 +1,107 @@
+#pragma once
+
+/***** Architecture detection macros *****/
+
+#ifdef GPTOSS_ARCH_X86_64
+    #if GPTOSS_ARCH_X86_64 != 0 && GPTOSS_ARCH_X86_64 != 1
+        #error "Invalid GPTOSS_ARCH_X86_64 value: must be either 0 or 1"
+    #endif
+#else
+    #if defined(__x86_64__) || defined(_M_X64) && !defined(_M_ARM64EC)
+        #define GPTOSS_ARCH_X86_64 1
+    #else
+        #define GPTOSS_ARCH_X86_64 0
+    #endif
+#endif
+
+#ifdef GPTOSS_ARCH_ARM64
+    #if GPTOSS_ARCH_ARM64 != 0 && GPTOSS_ARCH_ARM64 != 1
+        #error "Invalid GPTOSS_ARCH_ARM64 value: must be either 0 or 1"
+    #endif
+#else
+    #if defined(__aarch64__) || defined(_M_ARM64) || defined(_M_ARM64EC)
+        #define GPTOSS_ARCH_ARM64 1
+    #else
+        #define GPTOSS_ARCH_ARM64 0
+    #endif
+#endif
+
+#if GPTOSS_ARCH_X86_64 + GPTOSS_ARCH_ARM64 == 0
+    #error "Unsupported architecture: neither x86-64 nor ARM64 detected"
+#elif GPTOSS_ARCH_X86_64 + GPTOSS_ARCH_ARM64 != 1
+    #error "Inconsistent architecture detection: both x86-64 and ARM64 detection macros are specified"
+#endif
+
+/***** Compiler portability macros *****/
+
+#ifndef GPTOSS_LIKELY
+    #if defined(__GNUC__)
+        #define GPTOSS_LIKELY(condition) (__builtin_expect(!!(condition), 1))
+    #else
+        #define GPTOSS_LIKELY(condition) (!!(condition))
+    #endif
+#endif
+
+#ifndef GPTOSS_UNLIKELY
+    #if defined(__GNUC__)
+        #define GPTOSS_UNLIKELY(condition) (__builtin_expect(!!(condition), 0))
+    #else
+        #define GPTOSS_UNLIKELY(condition) (!!(condition))
+    #endif
+#endif
+
+#ifndef GPTOSS_UNPREDICTABLE
+    #if defined(__has_builtin)
+        #if __has_builtin(__builtin_unpredictable)
+            #define GPTOSS_UNPREDICTABLE(condition) (__builtin_unpredictable(!!(condition)))
+        #endif
+    #endif
+#endif
+#ifndef GPTOSS_UNPREDICTABLE
+    #if defined(__GNUC__) && (__GNUC__ >= 9) && !defined(__INTEL_COMPILER)
+        #define GPTOSS_UNPREDICTABLE(condition) (__builtin_expect_with_probability(!!(condition), 0, 0.5))
+    #else
+        #define GPTOSS_UNPREDICTABLE(condition) (!!(condition))
+    #endif
+#endif
+
+// Disable padding for structure members.
+#ifndef GPTOSS_DENSELY_PACKED_STRUCTURE
+    #if defined(__GNUC__)
+        #define GPTOSS_DENSELY_PACKED_STRUCTURE __attribute__((__packed__))
+    #else
+        #error "Compiler-specific implementation of GPTOSS_DENSELY_PACKED_STRUCTURE required"
+    #endif
+#endif
+
+#ifndef GPTOSS_ALIGN
+    #if defined(__GNUC__)
+        #define GPTOSS_ALIGN(alignment) __attribute__((__aligned__(alignment)))
+    #elif defined(_MSC_VER)
+        #define GPTOSS_ALIGN(alignment) __declspec(align(alignment))
+    #else
+        #error "Compiler-specific implementation of GPTOSS_ALIGN required"
+    #endif
+#endif
+
+#ifndef GPTOSS_FORCE_INLINE
+    #if defined(__GNUC__)
+        #define GPTOSS_FORCE_INLINE inline __attribute__((__always_inline__))
+    #elif defined(_MSC_VER)
+        #define GPTOSS_FORCE_INLINE __forceinline
+    #else
+        #define GPTOSS_FORCE_INLINE inline
+    #endif
+#endif
+
+/***** Symbol visibility macros *****/
+
+#ifndef GPTOSS_INTERNAL_SYMBOL
+    #if defined(__ELF__)
+        #define GPTOSS_INTERNAL_SYMBOL __attribute__((__visibility__("internal")))
+    #elif defined(__MACH__)
+        #define GPTOSS_INTERNAL_SYMBOL __attribute__((__visibility__("hidden")))
+    #else
+        #define GPTOSS_INTERNAL_SYMBOL
+    #endif
+#endif

gptoss_kernels/source/include/internal/math.h

@@ -0,0 +1,40 @@
+#pragma once
+
+#include <assert.h>
+#include <stddef.h>
+#include <stdint.h>
+
+inline static size_t math_ceil_div(size_t numer, size_t denom) {
+    return (numer + denom - 1) / denom;
+}
+
+inline static size_t math_max(size_t a, size_t b) {
+    return a >= b ? a : b;
+}
+
+inline static size_t math_min(size_t a, size_t b) {
+    return a < b ? a : b;
+}
+
+inline static size_t math_sub_sat(size_t a, size_t b) {
+    return a > b ? a - b : 0;
+}
+
+static size_t math_round_down_po2(size_t number, size_t multiple) {
+    assert(multiple != 0);
+    assert((multiple & (multiple - 1)) == 0);
+
+    return number & -multiple;
+}
+
+static size_t math_round_up_po2(size_t number, size_t multiple) {
+    assert(multiple != 0);
+    assert((multiple & (multiple - 1)) == 0);
+
+    const size_t multiple_mask = multiple - 1;
+    if ((number & multiple_mask) != 0) {
+        number |= multiple_mask;
+        number += 1;
+    }
+    return number;
+}

gptoss_kernels/source/include/internal/metal-kernels.h

@@ -0,0 +1,486 @@
+#pragma once
+
+#include <stddef.h>
+#include <stdint.h>
+
+#include <internal/metal.h>
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+#include <stddef.h>
+#include <stdint.h>
+
+#include <internal/kernel-args.h>
+#include <internal/math.h>
+#include <internal/metal.h>
+
+
+enum gptoss_status gptoss_metal_command_buffer_encode_launch_u32_fill_random(
+    const struct gptoss_metal_command_buffer* command_buffer,
+    const struct gptoss_metal_function* u32_fill_random_fn,
+    size_t threadgroup_size,
+    size_t max_threadgroups,
+    const struct gptoss_metal_buffer* output_buffer,
+    size_t output_offset,
+    uint64_t num_elements,
+    uint64_t rng_seed,
+    uint64_t rng_offset);
+
+enum gptoss_status gptoss_metal_command_buffer_encode_launch_f32_fill_random(
+    const struct gptoss_metal_command_buffer* command_buffer,
+    const struct gptoss_metal_function* f32_fill_random_fn,
+    size_t threadgroup_size,
+    size_t max_threadgroups,
+    const struct gptoss_metal_buffer* output_buffer,
+    size_t output_offset,
+    uint64_t num_elements,
+    uint64_t rng_seed,
+    uint64_t rng_offset,
+    float rng_min,
+    float rng_max);
+
+enum gptoss_status gptoss_metal_command_buffer_encode_launch_bf16_fill_random(
+    const struct gptoss_metal_command_buffer* command_buffer,
+    const struct gptoss_metal_function* bf16_fill_random_fn,
+    size_t threadgroup_size,
+    size_t max_threadgroups,
+    const struct gptoss_metal_buffer* output_buffer,
+    size_t output_offset,
+    uint64_t num_elements,
+    uint64_t rng_seed,
+    uint64_t rng_offset,
+    float rng_min,
+    float rng_max);
+
+enum gptoss_status gptoss_metal_command_buffer_encode_launch_mf4_f32_convert(
+    const struct gptoss_metal_command_buffer* command_buffer,
+    const struct gptoss_metal_function* mf4_f32_convert_fn,
+    size_t threadgroup_size,
+    size_t max_threadgroups,
+    const struct gptoss_metal_buffer* block_buffer,
+    const struct gptoss_metal_buffer* scale_buffer,
+    const struct gptoss_metal_buffer* output_buffer,
+    uint64_t num_elements);
+
+enum gptoss_status gptoss_metal_command_buffer_encode_launch_bf16_f32_embeddings(
+    const struct gptoss_metal_command_buffer* command_buffer,
+    const struct gptoss_metal_function* bf16_f32_embeddings_fn,
+    size_t threadgroup_size,
+    const struct gptoss_metal_buffer* token_buffer,
+    size_t token_offset,
+    const struct gptoss_metal_buffer* weight_buffer,
+    size_t weight_offset,
+    const struct gptoss_metal_buffer* output_buffer,
+    size_t output_offset,
+    const struct gptoss_metal_buffer* control_buffer,
+    size_t control_offset,
+    uint32_t num_tokens,
+    uint32_t num_channels);
+
+enum gptoss_status gptoss_metal_command_buffer_encode_launch_f32_bf16w_rmsnorm(
+    const struct gptoss_metal_command_buffer* command_buffer,
+    const struct gptoss_metal_function* f32_bf16w_rmsnorm_fn,
+    const struct gptoss_metal_buffer* input_buffer,
+    size_t input_offset,
+    const struct gptoss_metal_buffer* weight_buffer,
+    size_t weight_offset,
+    const struct gptoss_metal_buffer* output_buffer,
+    size_t output_offset,
+    const struct gptoss_metal_buffer* control_buffer,
+    size_t control_offset,
+    uint32_t num_tokens,
+    uint32_t num_channels,
+    float epsilon);
+
+enum gptoss_status gptoss_metal_command_buffer_encode_launch_f32_bf16w_matmul(
+    const struct gptoss_metal_command_buffer* command_buffer,
+    const struct gptoss_metal_function* f32_bf16w_matmul_fn,
+    size_t threadgroup_size,
+    const struct gptoss_metal_buffer* input_buffer,
+    size_t input_offset,
+    const struct gptoss_metal_buffer* weight_buffer,
+    size_t weight_offset,
+    const struct gptoss_metal_buffer* bias_buffer,
+    size_t bias_offset,
+    const struct gptoss_metal_buffer* output_buffer,
+    size_t output_offset,
+    const struct gptoss_metal_buffer* control_buffer,
+    size_t control_offset,
+    uint32_t num_tokens,
+    uint32_t num_cols,
+    uint32_t num_rows);
+
+enum gptoss_status gptoss_metal_command_buffer_encode_launch_f32_bf16w_matmul_qkv(
+    const struct gptoss_metal_command_buffer* command_buffer,
+    const struct gptoss_metal_function* f32_bf16w_matmul_qkv_fn,
+    size_t threadgroup_size,
+    const struct gptoss_metal_buffer* input_buffer,
+    size_t input_offset,
+    const struct gptoss_metal_buffer* weight_buffer,
+    size_t weight_offset,
+    const struct gptoss_metal_buffer* bias_buffer,
+    size_t bias_offset,
+    const struct gptoss_metal_buffer* output_buffer,
+    size_t output_offset,
+    const struct gptoss_metal_buffer* kv_buffer,
+    size_t kv_offset,
+    const struct gptoss_metal_buffer* control_buffer,
+    size_t control_offset,
+    uint32_t num_tokens,
+    uint32_t num_cols,
+    uint32_t num_q_heads,
+    uint32_t num_kv_heads,
+    uint32_t attn_head_dim,
+    uint32_t token_offset,
+    uint32_t max_tokens,
+    float rope_base,
+    float interpolation_scale,
+    float yarn_offset,
+    float yarn_scale,
+    float yarn_multiplier);
+
+enum gptoss_status gptoss_metal_command_buffer_encode_launch_f32_bf16w_matmul_add(
+    const struct gptoss_metal_command_buffer* command_buffer,
+    const struct gptoss_metal_function* f32_bf16w_matmul_fn,
+    size_t threadgroup_size,
+    const struct gptoss_metal_buffer* input_buffer,
+    size_t input_offset,
+    const struct gptoss_metal_buffer* weight_buffer,
+    size_t weight_offset,
+    const struct gptoss_metal_buffer* bias_buffer,
+    size_t bias_offset,
+    const struct gptoss_metal_buffer* output_buffer,
+    size_t output_offset,
+    const struct gptoss_metal_buffer* control_buffer,
+    size_t control_offset,
+    uint32_t num_tokens,
+    uint32_t num_cols,
+    uint32_t num_rows);
+
+enum gptoss_status
+gptoss_metal_command_buffer_encode_launch_f32_bf16w_dense_matmul_qkv(
+    const struct gptoss_metal_command_buffer* command_buffer,
+    const struct gptoss_metal_function* f32_bf16w_dense_matmul_fn,
+    const struct gptoss_metal_buffer* input_buffer,
+    size_t input_offset,
+    const struct gptoss_metal_buffer* weight_buffer,
+    size_t weight_offset,
+    const struct gptoss_metal_buffer* bias_buffer,
+    size_t bias_offset,
+    const struct gptoss_metal_buffer* output_buffer,
+    size_t output_offset,
+    const struct gptoss_metal_buffer* control_buffer,
+    size_t control_offset,
+    uint32_t num_tokens,
+    uint32_t num_cols,
+    uint32_t num_rows);
+
+enum gptoss_status
+gptoss_metal_command_buffer_encode_launch_f32_bf16w_dense_matmul_attn_output(
+    const struct gptoss_metal_command_buffer* command_buffer,
+    const struct gptoss_metal_function* f32_bf16w_dense_matmul_fn,
+    const struct gptoss_metal_buffer* input_buffer,
+    size_t input_offset,
+    const struct gptoss_metal_buffer* weight_buffer,
+    size_t weight_offset,
+    const struct gptoss_metal_buffer* bias_buffer,
+    size_t bias_offset,
+    const struct gptoss_metal_buffer* output_buffer,
+    size_t output_offset,
+    const struct gptoss_metal_buffer* control_buffer,
+    size_t control_offset,
+    uint32_t num_tokens,
+    uint32_t num_cols,
+    uint32_t num_rows);
+
+enum gptoss_status
+gptoss_metal_command_buffer_encode_launch_f32_bf16w_dense_matmul_mlp_gate(
+    const struct gptoss_metal_command_buffer* command_buffer,
+    const struct gptoss_metal_function* f32_bf16w_dense_matmul_fn,
+    const struct gptoss_metal_buffer* input_buffer,
+    size_t input_offset,
+    const struct gptoss_metal_buffer* weight_buffer,
+    size_t weight_offset,
+    const struct gptoss_metal_buffer* bias_buffer,
+    size_t bias_offset,
+    const struct gptoss_metal_buffer* output_buffer,
+    size_t output_offset,
+    const struct gptoss_metal_buffer* control_buffer,
+    size_t control_offset,
+    uint32_t num_tokens,
+    uint32_t num_cols,
+    uint32_t num_rows);
+
+enum gptoss_status gptoss_metal_command_buffer_encode_launch_f32_bf16w_unembedding(
+    const struct gptoss_metal_command_buffer* command_buffer,
+    const struct gptoss_metal_function* f32_bf16w_matmul_fn,
+    size_t threadgroup_size,
+    size_t max_threadgroups,
+    const struct gptoss_metal_buffer* input_buffer,
+    size_t input_offset,
+    const struct gptoss_metal_buffer* weight_buffer,
+    size_t weight_offset,
+    const struct gptoss_metal_buffer* output_buffer,
+    size_t output_offset,
+    const struct gptoss_metal_buffer* argmax_buffer,
+    size_t argmax_offset,
+    const struct gptoss_metal_buffer* control_buffer,
+    size_t control_offset,
+    uint32_t num_tokens,
+    uint32_t num_cols,
+    uint32_t num_rows);
+
+enum gptoss_status gptoss_metal_command_buffer_encode_launch_f32_mf4w_moe_matmul_swiglu(
+    const struct gptoss_metal_command_buffer* command_buffer,
+    const struct gptoss_metal_function* f32_mf4w_moe_matmul_swiglu_fn,
+    size_t threadgroup_size,
+    const struct gptoss_metal_buffer* input_buffer,
+    size_t input_offset,
+    const struct gptoss_metal_buffer* expert_buffer,
+    size_t expert_offset,
+    const struct gptoss_metal_buffer* weight_block_buffer,
+    size_t weight_block_offset,
+    const struct gptoss_metal_buffer* weight_scale_buffer,
+    size_t weight_scale_offset,
+    const struct gptoss_metal_buffer* bias_buffer,
+    size_t bias_offset,
+    const struct gptoss_metal_buffer* output_buffer,
+    size_t output_offset,
+    const struct gptoss_metal_buffer* control_buffer,
+    size_t control_offset,
+    float swiglu_limit,
+    uint32_t expert_stride,
+    uint32_t num_tokens,
+    uint32_t num_active_experts,
+    uint32_t num_cols,
+    uint32_t num_rows);
+
+enum gptoss_status gptoss_metal_command_buffer_encode_launch_f32_mf4w_moe_matmul(
+    const struct gptoss_metal_command_buffer* command_buffer,
+    const struct gptoss_metal_function* f32_mf4w_moe_matmul_fn,
+    size_t threadgroup_size,
+    const struct gptoss_metal_buffer* input_buffer,
+    size_t input_offset,
+    const struct gptoss_metal_buffer* expert_buffer,
+    size_t expert_offset,
+    const struct gptoss_metal_buffer* weight_block_buffer,
+    size_t weight_block_offset,
+    const struct gptoss_metal_buffer* weight_scale_buffer,
+    size_t weight_scale_offset,
+    const struct gptoss_metal_buffer* bias_buffer,
+    size_t bias_offset,
+    const struct gptoss_metal_buffer* output_buffer,
+    size_t output_offset,
+    const struct gptoss_metal_buffer* control_buffer,
+    size_t control_offset,
+    uint32_t expert_stride,
+    uint32_t num_tokens,
+    uint32_t num_active_experts,
+    uint32_t num_cols,
+    uint32_t num_rows);
+
+enum gptoss_status gptoss_metal_command_buffer_encode_launch_f32_rope(
+    const struct gptoss_metal_command_buffer* command_buffer,
+    const struct gptoss_metal_function* f32_rope_fn,
+    size_t threadgroup_size,
+    const struct gptoss_metal_buffer* activations_buffer,
+    size_t activations_offset,
+    const struct gptoss_metal_buffer* control_buffer,
+    size_t control_offset,
+    float rope_base,
+    float interpolation_scale,
+    float yarn_offset,
+    float yarn_scale,
+    float yarn_multiplier,
+    uint32_t num_tokens,
+    uint32_t num_q_heads,
+    uint32_t num_kv_heads,
+    uint32_t attn_head_dim,
+    uint32_t token_offset);
+
+enum gptoss_status gptoss_metal_command_buffer_encode_launch_f32_accumulate(
+    const struct gptoss_metal_command_buffer* command_buffer,
+    const struct gptoss_metal_function* f32_accumulate_fn,
+    size_t threadgroup_size,
+    size_t max_threadgroups,
+    const struct gptoss_metal_buffer* input_buffer,
+    size_t input_offset,
+    const struct gptoss_metal_buffer* expert_buffer,
+    size_t expert_offset,
+    const struct gptoss_metal_buffer* output_buffer,
+    size_t output_offset,
+    const struct gptoss_metal_buffer* control_buffer,
+    size_t control_offset,
+    uint32_t num_channels,
+    uint32_t num_tokens,
+    uint32_t num_experts);
+
+enum gptoss_status gptoss_metal_command_buffer_encode_launch_expert_routing_metadata(
+    const struct gptoss_metal_command_buffer* command_buffer,
+    const struct gptoss_metal_function* expert_routing_metadata_fn,
+    const struct gptoss_metal_buffer* expert_predictions_buffer,
+    size_t expert_predictions_offset,
+    const struct gptoss_metal_buffer* expert_offsets_buffer,
+    size_t expert_offsets_offset,
+    const struct gptoss_metal_buffer* intra_expert_offsets_buffer,
+    size_t intra_expert_offsets_offset,
+    uint32_t num_tokens,
+    uint32_t num_experts);
+
+
+enum gptoss_status gptoss_metal_command_buffer_encode_launch_f32_scatter(
+    const struct gptoss_metal_command_buffer* command_buffer,
+    const struct gptoss_metal_function* f32_scatter_fn,
+    const struct gptoss_metal_buffer* input_buffer,
+    size_t input_offset,
+    const struct gptoss_metal_buffer* expert_predictions_buffer,
+    size_t expert_predictions_offset,
+    const struct gptoss_metal_buffer* expert_offsets_buffer,
+    size_t expert_offsets_offset,
+    const struct gptoss_metal_buffer* intra_expert_offsets_buffer,
+    size_t intra_expert_offsets_offset,
+    const struct gptoss_metal_buffer* output_buffer,
+    size_t output_offset,
+    uint32_t num_channels,
+    uint32_t num_tokens,
+    uint32_t num_active_experts);
+    
+enum gptoss_status
+gptoss_metal_command_buffer_encode_launch_f32_mf4w_moe_dense_matmul_swiglu(
+    const struct gptoss_metal_command_buffer* command_buffer,
+    const struct gptoss_metal_function* f32_mf4w_moe_dense_matmul_swiglu_fn,
+    const struct gptoss_metal_buffer* expert_offsets_buffer,
+    size_t expert_offsets_offset,
+    const struct gptoss_metal_buffer* input_buffer,
+    size_t input_offset,
+    const struct gptoss_metal_buffer* weight_block_buffer,
+    size_t weight_block_offset,
+    const struct gptoss_metal_buffer* weight_scale_buffer,
+    size_t weight_scale_offset,
+    const struct gptoss_metal_buffer* bias_buffer,
+    size_t bias_offset,
+    const struct gptoss_metal_buffer* output_buffer,
+    size_t output_offset,
+    float swiglu_limit,
+    uint32_t expert_stride_bytes,
+    uint32_t num_tokens,
+    uint32_t num_experts,
+    uint32_t num_cols,
+    uint32_t num_rows);
+    
+enum gptoss_status gptoss_metal_command_buffer_encode_launch_f32_mf4w_moe_dense_matmul(
+    const struct gptoss_metal_command_buffer* command_buffer,
+    const struct gptoss_metal_function* f32_mf4w_moe_dense_matmul_fn,
+    const struct gptoss_metal_buffer* expert_offsets_buffer,
+    size_t expert_offsets_offset,
+    const struct gptoss_metal_buffer* input_buffer,
+    size_t input_offset,
+    const struct gptoss_metal_buffer* weight_block_buffer,
+    size_t weight_block_offset,
+    const struct gptoss_metal_buffer* weight_scale_buffer,
+    size_t weight_scale_offset,
+    const struct gptoss_metal_buffer* bias_buffer,
+    size_t bias_offset,
+    const struct gptoss_metal_buffer* output_buffer,
+    size_t output_offset,
+    uint32_t expert_stride_bytes,
+    uint32_t num_tokens,
+    uint32_t num_experts,
+    uint32_t num_cols,
+    uint32_t num_rows);
+    
+enum gptoss_status gptoss_metal_command_buffer_encode_launch_f32_gather_and_accumulate_e4(
+    const struct gptoss_metal_command_buffer* command_buffer,
+    const struct gptoss_metal_function* f32_gather_and_accumulate_e4_fn,
+    const struct gptoss_metal_buffer* input_buffer,
+    size_t input_offset,
+    const struct gptoss_metal_buffer* expert_predictions_buffer,
+    size_t expert_predictions_offset,
+    const struct gptoss_metal_buffer* expert_offsets_buffer,
+    size_t expert_offsets_offset,
+    const struct gptoss_metal_buffer* intra_expert_offsets_buffer,
+    size_t intra_expert_offsets_offset,
+    const struct gptoss_metal_buffer* output_buffer,
+    size_t output_offset,
+    uint32_t num_channels,
+    uint32_t num_tokens,
+    uint32_t num_active_experts);
+
+enum gptoss_status gptoss_metal_command_buffer_encode_launch_f32_topk(
+    const struct gptoss_metal_command_buffer* command_buffer,
+    const struct gptoss_metal_function* f32_topk_fn,
+    const struct gptoss_metal_buffer* input_buffer,
+    size_t input_offset,
+    const struct gptoss_metal_buffer* output_buffer,
+    size_t output_offset,
+    const struct gptoss_metal_buffer* control_buffer,
+    size_t control_offset,
+    uint32_t num_tokens,
+    uint32_t num_experts,
+    uint32_t num_active_experts);
+
+enum gptoss_status gptoss_metal_command_buffer_encode_launch_f32_sdpa(
+    const struct gptoss_metal_command_buffer* command_buffer,
+    const struct gptoss_metal_function* f32_sdpa_fn,
+    const struct gptoss_metal_buffer* q_buffer,
+    size_t q_offset,
+    const struct gptoss_metal_buffer* kv_buffer,
+    size_t kv_offset,
+    const struct gptoss_metal_buffer* s_buffer,
+    size_t s_offset,
+    const struct gptoss_metal_buffer* output_buffer,
+    size_t output_offset,
+    const struct gptoss_metal_buffer* control_buffer,
+    size_t control_offset,
+    uint32_t window,
+    uint32_t kv_stride,
+    uint32_t num_q_tokens,
+    uint32_t num_kv_tokens,
+    uint32_t num_q_heads,
+    uint32_t num_kv_heads,
+    uint32_t head_dim);
+
+enum gptoss_status gptoss_metal_command_buffer_encode_launch_f32_softmax(
+    const struct gptoss_metal_command_buffer* command_buffer,
+    const struct gptoss_metal_function* f32_softmax_fn,
+    size_t threadgroup_size,
+    size_t max_threadgroups,
+    const struct gptoss_metal_buffer* score_buffer,
+    size_t score_offset,
+    const struct gptoss_metal_buffer* argmax_buffer,
+    size_t argmax_offset,
+    const struct gptoss_metal_buffer* prob_buffer,
+    size_t prob_offset,
+    const struct gptoss_metal_buffer* sum_buffer,
+    size_t sum_offset,
+    const struct gptoss_metal_buffer* control_buffer,
+    size_t control_offset,
+    uint32_t num_channels,
+    uint32_t num_tokens,
+    float temperature,
+    uint32_t* num_threadgroups_out,
+    uint32_t* num_channels_per_threadgroup_out);
+
+enum gptoss_status gptoss_metal_command_buffer_encode_launch_f32_sample(
+    const struct gptoss_metal_command_buffer* command_buffer,
+    const struct gptoss_metal_function* f32_sample_fn,
+    size_t min_threadgroup_size,
+    const struct gptoss_metal_buffer* prob_buffer,
+    size_t prob_offset,
+    const struct gptoss_metal_buffer* sum_buffer,
+    size_t sum_offset,
+    const struct gptoss_metal_buffer* token_buffer,
+    size_t token_offset,
+    const struct gptoss_metal_buffer* control_buffer,
+    size_t control_offset,
+    uint64_t rng_seed,
+    uint32_t rng_offset,
+    uint32_t num_blocks,
+    uint32_t num_channels,
+    uint32_t num_channels_per_block);
+
+#ifdef __cplusplus
+}  // extern "C"
+#endif

gptoss_kernels/source/include/internal/metal.h

@@ -0,0 +1,138 @@
+#pragma once
+
+#include <stddef.h>
+
+#include <gpt-oss/types.h>
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+struct gptoss_metal_device {
+    void* object; // id<MTLDevice>
+    size_t num_cores;
+    size_t max_buffer_size;
+    size_t max_threadgroup_memory;
+    size_t max_threadgroup_threads_x;
+    size_t max_threadgroup_threads_y;
+    size_t max_threadgroup_threads_z;
+};
+
+enum gptoss_status gptoss_metal_device_create_system_default(
+    struct gptoss_metal_device* device_out);
+
+enum gptoss_status gptoss_metal_device_release(
+    struct gptoss_metal_device* device);
+
+
+struct gptoss_metal_library {
+    void* object; // id<MTLLibrary>
+};
+
+enum gptoss_status gptoss_metal_library_create_default(
+    const struct gptoss_metal_device* device,
+    struct gptoss_metal_library* library_out);
+
+enum gptoss_status gptoss_metal_library_release(
+    struct gptoss_metal_library* library);
+
+struct gptoss_metal_function {
+    void* function_object; // id<MTLFunction>
+    void* pipeline_state_object; // id<MTLComputePipelineState>
+    size_t max_threadgroup_threads;
+    size_t simdgroup_threads;
+    size_t static_threadgroup_memory;
+};
+
+enum gptoss_status gptoss_metal_function_create(
+    const struct gptoss_metal_library* library,
+    const char* name,
+    struct gptoss_metal_function* function_out);
+
+enum gptoss_status gptoss_metal_function_release(
+    struct gptoss_metal_function* function);
+
+struct gptoss_metal_buffer {
+    void* object; // id<MTLBuffer>
+    size_t size;
+    void* ptr;
+};
+
+enum gptoss_status gptoss_metal_buffer_create(
+    const struct gptoss_metal_device* device,
+    size_t size,
+    const void* data,
+    struct gptoss_metal_buffer* buffer_out);
+
+enum gptoss_status gptoss_metal_buffer_wrap(
+    const struct gptoss_metal_device* device,
+    size_t size,
+    const void* data,
+    struct gptoss_metal_buffer* buffer_out);
+
+enum gptoss_status gptoss_metal_buffer_release(
+    struct gptoss_metal_buffer* buffer);
+
+struct gptoss_metal_command_queue {
+    void* object; // id<MTLCommandQueue>
+};
+
+enum gptoss_status gptoss_metal_command_queue_create(
+    const struct gptoss_metal_device* device,
+    struct gptoss_metal_command_queue* command_queue_out);
+
+enum gptoss_status gptoss_metal_command_queue_release(
+    struct gptoss_metal_command_queue* command_queue);
+
+struct gptoss_metal_command_buffer {
+    void* object; // id<MTLCommandBuffer>
+};
+
+enum gptoss_status gptoss_metal_command_buffer_create(
+    const struct gptoss_metal_command_queue* command_queue,
+    struct gptoss_metal_command_buffer* command_buffer_out);
+
+enum gptoss_status gptoss_metal_command_buffer_encode_fill_buffer(
+    const struct gptoss_metal_command_buffer* command_buffer,
+    const struct gptoss_metal_buffer* buffer,
+    size_t offset,
+    size_t size,
+    uint8_t fill_value);
+
+enum gptoss_status gptoss_metal_command_buffer_encode_copy_buffer(
+    const struct gptoss_metal_command_buffer* command_buffer,
+    const struct gptoss_metal_buffer* input_buffer,
+    size_t input_offset,
+    const struct gptoss_metal_buffer* output_buffer,
+    size_t output_offset,
+    size_t size);
+
+enum gptoss_status gptoss_metal_command_buffer_encode_launch_kernel(
+    const struct gptoss_metal_command_buffer* command_buffer,
+    const struct gptoss_metal_function* function,
+    size_t threadgroup_size_x,
+    size_t threadgroup_size_y,
+    size_t threadgroup_size_z,
+    size_t num_threadgroups_x,
+    size_t num_threadgroups_y,
+    size_t num_threadgroups_z,
+    size_t params_size,
+    const void* params,
+    size_t num_device_buffers,
+    const struct gptoss_metal_buffer** device_buffers,
+    const size_t* device_buffer_offsets,
+    size_t threadgroup_buffer_size);
+
+enum gptoss_status gptoss_metal_command_buffer_commit(
+    const struct gptoss_metal_command_buffer* command_buffer);
+
+enum gptoss_status gptoss_metal_command_buffer_wait_completion(
+    const struct gptoss_metal_command_buffer* command_buffer,
+    double* elapsed_seconds);
+
+enum gptoss_status gptoss_metal_command_buffer_release(
+    struct gptoss_metal_command_buffer* command_buffer);
+
+#ifdef __cplusplus
+}  // extern "C"
+#endif

gptoss_kernels/source/include/internal/metal.hpp

@@ -0,0 +1,342 @@
+#pragma once
+
+#include <array>
+#include <initializer_list>
+#include <cstring>
+#include <stdexcept>
+#include <vector>
+
+#include <gpt-oss/types.h>
+#include <internal/metal.h>
+#include <internal/metal-kernels.h>
+
+
+namespace gptoss {
+
+inline void Check(gptoss_status s, const char* what) {
+    if (s != gptoss_status_success) {
+        throw std::runtime_error(what);
+    }
+}
+
+inline std::size_t round_up(std::size_t p, std::size_t q) {
+    const std::size_t r = p % q;
+    if (r == 0) {
+        return p;
+    } else {
+        return p - r + q;
+    }
+}
+
+namespace metal {
+
+class Device {
+public:
+    inline Device() {
+        Check(gptoss_metal_device_create_system_default(&device_), "create Device");
+    }
+
+    inline ~Device() {
+        gptoss_metal_device_release(&device_);
+    }
+
+    Device(const Device&) = delete;
+    Device& operator=(const Device&) = delete;
+
+    inline Device(Device&& other) noexcept {
+        device_ = other.device_;
+        std::memset(&other.device_, 0, sizeof(other.device_));
+    }
+
+    inline Device& operator=(Device&& other) noexcept {
+        if (this != &other) {
+            gptoss_metal_device_release(&device_);
+            device_ = other.device_;
+            std::memset(&other.device_, 0, sizeof(other.device_));
+        }
+        return *this;
+    }
+
+    inline const gptoss_metal_device* handle() const noexcept { return &device_; }
+
+    inline size_t max_buffer_size() const noexcept { return device_.max_buffer_size; }
+    inline size_t max_threadgroup_memory() const noexcept { return device_.max_threadgroup_memory; }
+    inline size_t max_threadgroup_threads_x() const noexcept { return device_.max_threadgroup_threads_x; }
+    inline size_t max_threadgroup_threads_y() const noexcept { return device_.max_threadgroup_threads_y; }
+    inline size_t max_threadgroup_threads_z() const noexcept { return device_.max_threadgroup_threads_z; }
+
+private:
+    gptoss_metal_device device_{};
+};
+
+class Library {
+public:
+    inline explicit Library(const Device& dev) {
+        Check(gptoss_metal_library_create_default(dev.handle(), &library_),
+            "gptoss_metal_library_create_default");
+    }
+
+    inline ~Library() {
+        gptoss_metal_library_release(&library_);
+    }
+
+    Library(const Library&) = delete;
+    Library& operator=(const Library&) = delete;
+
+    inline Library(Library&& other) noexcept {
+        library_ = other.library_;
+        std::memset(&other.library_, 0, sizeof(other.library_));
+    }
+
+    inline Library& operator=(Library&& other) noexcept {
+        if (this != &other) {
+            gptoss_metal_library_release(&library_);
+            library_ = other.library_;
+            std::memset(&other.library_, 0, sizeof(other.library_));
+        }
+        return *this;
+    }
+
+    inline const gptoss_metal_library* handle() const noexcept {
+        return &library_;
+    }
+
+private:
+    gptoss_metal_library library_{};
+};
+
+class Function {
+public:
+    inline Function(const Library& library, const char* name) {
+        Check(gptoss_metal_function_create(library.handle(), name, &function_),
+            "gptoss_metal_function_create");
+    }
+
+    inline ~Function() {
+        gptoss_metal_function_release(&function_);
+    }
+
+    Function(const Function&) = delete;
+    Function& operator=(const Function&) = delete;
+
+    inline Function(Function&& other) noexcept {
+        function_ = other.function_;
+        std::memset(&other.function_, 0, sizeof(other.function_));
+    }
+
+    inline Function& operator=(Function&& other) noexcept {
+        if (this != &other) {
+            gptoss_metal_function_release(&function_);
+            function_ = other.function_;
+            std::memset(&other.function_, 0, sizeof(other.function_));
+        }
+        return *this;
+    }
+
+    inline const gptoss_metal_function* handle() const noexcept { return &function_; }
+
+    inline size_t max_threadgroup_threads() const noexcept { return function_.max_threadgroup_threads; }
+    inline size_t simdgroup_threads() const noexcept { return function_.simdgroup_threads; }
+    inline size_t static_threadgroup_memory() const noexcept { return function_.static_threadgroup_memory; }
+
+private:
+    gptoss_metal_function function_{};
+};
+
+class Buffer {
+public:
+    inline Buffer(const Device& dev, size_t size, const void* data = nullptr) {
+        Check(gptoss_metal_buffer_create(dev.handle(), size, data, &buffer_), "create buffer");
+    }
+
+    inline ~Buffer() {
+        gptoss_metal_buffer_release(&buffer_);
+    }
+
+    Buffer(const Buffer&) = delete;
+    Buffer& operator=(const Buffer&) = delete;
+
+    inline Buffer(Buffer&& other) noexcept {
+        buffer_ = other.buffer_;
+        std::memset(&other.buffer_, 0, sizeof(other.buffer_));
+    }
+
+    inline Buffer& operator=(Buffer&& other) noexcept {
+        if (this != &other) {
+            gptoss_metal_buffer_release(&buffer_);
+            buffer_ = other.buffer_;
+            std::memset(&other.buffer_, 0, sizeof(other.buffer_));
+        }
+        return *this;
+    }
+
+    inline size_t size() const noexcept { return buffer_.size; }
+    inline void* ptr() const noexcept { return buffer_.ptr; }
+
+    inline const gptoss_metal_buffer* handle() const noexcept { return &buffer_; }
+
+private:
+    gptoss_metal_buffer buffer_{};
+};
+
+class CommandQueue {
+public:
+    inline explicit CommandQueue(const Device& dev) {
+        Check(gptoss_metal_command_queue_create(dev.handle(), &command_queue_),
+            "gptoss_metal_command_queue_create");
+    }
+
+    inline ~CommandQueue() {
+        gptoss_metal_command_queue_release(&command_queue_);
+    }
+
+    CommandQueue(const CommandQueue&) = delete;
+    CommandQueue& operator=(const CommandQueue&) = delete;
+
+    inline CommandQueue(CommandQueue&& other) noexcept {
+        command_queue_ = other.command_queue_;
+        std::memset(&other.command_queue_, 0, sizeof(other.command_queue_));
+    }
+
+    inline CommandQueue& operator=(CommandQueue&& other) noexcept {
+        if (this != &other) {
+            gptoss_metal_command_queue_release(&command_queue_);
+            command_queue_ = other.command_queue_;
+            std::memset(&other.command_queue_, 0, sizeof(other.command_queue_));
+        }
+        return *this;
+    }
+
+    inline const gptoss_metal_command_queue* handle() const noexcept {
+        return &command_queue_;
+    }
+
+private:
+    gptoss_metal_command_queue command_queue_{};
+};
+
+class CommandBuffer {
+public:
+    inline explicit CommandBuffer(const CommandQueue& command_queue) {
+        Check(gptoss_metal_command_buffer_create(command_queue.handle(), &command_buffer_),
+            "gptoss_metal_command_buffer_create");
+    }
+    inline ~CommandBuffer() {
+        gptoss_metal_command_buffer_release(&command_buffer_);
+    }
+
+    CommandBuffer(const CommandBuffer&)            = delete;
+    CommandBuffer& operator=(const CommandBuffer&) = delete;
+
+    inline CommandBuffer(CommandBuffer&& other) noexcept  {
+        command_buffer_ = other.command_buffer_;
+        std::memset(&other.command_buffer_, 0, sizeof(other.command_buffer_));
+    }
+
+    inline CommandBuffer& operator=(CommandBuffer&& other) noexcept {
+        if (this != &other) {
+            gptoss_metal_command_buffer_release(&command_buffer_);
+            command_buffer_ = other.command_buffer_;
+            std::memset(&other.command_buffer_, 0, sizeof(other.command_buffer_));
+        }
+        return *this;
+    }
+
+    inline void encode_launch_kernel(const Function& function,
+                                     const std::array<size_t, 3>& threadgroup_size,
+                                     const std::array<size_t, 3>& num_threadgroups,
+                                     size_t params_size, const void* params,
+                                     std::initializer_list<const Buffer*> device_buffers = {},
+                                     size_t threadgroup_buffer_size = 0)
+    {
+        std::vector<const gptoss_metal_buffer*> buffer_handles(device_buffers.size());
+        std::transform(device_buffers.begin(), device_buffers.end(), buffer_handles.begin(),
+            [](const Buffer* buffer) -> const gptoss_metal_buffer* { return buffer->handle(); });
+        Check(gptoss_metal_command_buffer_encode_launch_kernel(
+                &command_buffer_, function.handle(),
+                threadgroup_size[0], threadgroup_size[1], threadgroup_size[2],
+                num_threadgroups[0], num_threadgroups[1], num_threadgroups[2],
+                params_size, params,
+                buffer_handles.size(),
+                buffer_handles.data(),
+                /*buffer_offsets=*/nullptr,
+                threadgroup_buffer_size),
+            "gptoss_metal_command_buffer_encode_launch_kernel");
+    }
+
+    inline void encode_launch_f32_fill_random(const Function& f32_fill_random_fn,
+                                              size_t threadgroup_size,
+                                              size_t num_threadgroups,
+                                              const Buffer& output_buffer,
+                                              size_t output_offset,
+                                              size_t num_channels,
+                                              uint64_t rng_seed,
+                                              uint64_t rng_offset,
+                                              float rng_min,
+                                              float rng_max)
+    {
+        Check(gptoss_metal_command_buffer_encode_launch_f32_fill_random(
+                &command_buffer_, f32_fill_random_fn.handle(),
+                threadgroup_size, num_threadgroups,
+                output_buffer.handle(), output_offset,
+                num_channels,
+                rng_seed, rng_offset, rng_min, rng_max),
+            "gptoss_metal_command_buffer_encode_launch_f32_fill_random");
+    }
+
+    inline void encode_launch_bf16_fill_random(const Function& bf16_fill_random_fn,
+                                               size_t threadgroup_size,
+                                               size_t num_threadgroups,
+                                               const Buffer& output_buffer,
+                                               size_t output_offset,
+                                               size_t num_channels,
+                                               uint64_t rng_seed,
+                                               uint64_t rng_offset,
+                                               float rng_min,
+                                               float rng_max)
+    {
+        Check(gptoss_metal_command_buffer_encode_launch_bf16_fill_random(
+                &command_buffer_, bf16_fill_random_fn.handle(),
+                threadgroup_size, num_threadgroups,
+                output_buffer.handle(), output_offset,
+                num_channels,
+                rng_seed, rng_offset, rng_min, rng_max),
+            "gptoss_metal_command_buffer_encode_launch_bf16_fill_random");
+    }
+
+    inline void encode_launch_u32_fill_random(const Function& u32_fill_random_fn,
+                                              size_t threadgroup_size,
+                                              size_t num_threadgroups,
+                                              const Buffer& output_buffer,
+                                              size_t output_offset,
+                                              size_t num_channels,
+                                              uint64_t rng_seed,
+                                              uint64_t rng_offset)
+    {
+        Check(gptoss_metal_command_buffer_encode_launch_u32_fill_random(
+                &command_buffer_, u32_fill_random_fn.handle(),
+                threadgroup_size, num_threadgroups,
+                output_buffer.handle(), output_offset,
+                num_channels,
+                rng_seed, rng_offset),
+            "gptoss_metal_command_buffer_encode_launch_u32_fill_random");
+    }
+
+    inline void commit() {
+        Check(gptoss_metal_command_buffer_commit(&command_buffer_), "commit");
+    }
+
+    inline double wait_completion() {
+        double secs = 0.0;
+        Check(gptoss_metal_command_buffer_wait_completion(&command_buffer_, &secs), "wait completion");
+        return secs;
+    }
+
+    inline const gptoss_metal_command_buffer* handle() const noexcept { return &command_buffer_; }
+
+private:
+    gptoss_metal_command_buffer command_buffer_{};
+};
+
+} // namespace metal
+} // namespace gptoss

gptoss_kernels/source/include/internal/model.h

@@ -0,0 +1,178 @@
+#pragma once
+
+#ifndef __cplusplus
+    #include <stdatomic.h>
+#endif
+#include <stdbool.h>
+#include <stddef.h>
+#include <stdint.h>
+
+#include "internal/metal.h"
+
+
+struct gptoss_tokenizer {
+#ifndef __cplusplus
+    atomic_uint_least64_t ref_count;
+#else
+    uint_least64_t ref_count;
+#endif
+
+    void* mapping_ptr;
+    size_t mapping_size;
+
+    const char* regex_ptr;
+    const char* tokens_ptr;
+
+    uint32_t num_text_tokens;
+    uint32_t num_special_tokens;
+
+    uint32_t special_token_id[gptoss_special_token_max - 1];
+};
+
+struct gptoss_model {
+#ifndef __cplusplus
+    atomic_uint_least64_t ref_count;
+#else
+    uint_least64_t ref_count;
+#endif
+
+    struct gptoss_tokenizer* tokenizer;
+
+    void* mapping_ptr;
+    size_t mapping_size;
+
+    uint32_t context_length;
+    uint32_t num_blocks;
+    uint32_t num_experts;
+    uint32_t num_active_experts;
+    uint32_t embedding_dim;
+    uint32_t mlp_dim;
+    float swiglu_limit;
+    uint32_t head_dim;
+    uint32_t num_heads;
+    uint32_t num_kv_heads;
+    uint32_t attention_window;
+    float rope_theta;
+    float interpolation_scale;
+    float yarn_offset;
+    float yarn_scale;
+    float yarn_multiplier;
+    float rmsnorm_epsilon;
+
+    uint32_t vocabulary_size;
+
+    bool lock_memory;
+
+    size_t weights_size;
+    size_t allocation_size;
+
+    // Metal objects
+    struct gptoss_metal_device device;
+    size_t max_threadgroups;
+    struct gptoss_metal_command_queue command_queue;
+    struct gptoss_metal_library library;
+    struct gptoss_metal_function bf16_f32_embeddings_fn;
+    struct gptoss_metal_function f32_bf16w_rmsnorm_fn;
+    struct gptoss_metal_function f32_bf16w_matmul_fn;
+    struct gptoss_metal_function f32_bf16w_matmul_qkv_fn;
+    struct gptoss_metal_function f32_bf16w_dense_matmul_qkv_fn;
+    struct gptoss_metal_function f32_bf16w_dense_matmul_attn_output_fn;
+    struct gptoss_metal_function f32_bf16w_dense_matmul_mlp_gate_fn;
+    struct gptoss_metal_function f32_bf16w_unembedding_fn;
+    struct gptoss_metal_function f32_rope_fn;
+    struct gptoss_metal_function f32_mf4w_moe_matmul_swiglu_fn;
+    struct gptoss_metal_function f32_mf4w_moe_matmul_fn;
+    struct gptoss_metal_function f32_accumulate_e4_fn;
+    struct gptoss_metal_function f32_scatter_e4_fn;
+    struct gptoss_metal_function f32_mf4w_moe_dense_matmul_swiglu_fn;
+    struct gptoss_metal_function f32_mf4w_moe_dense_matmul_fn;
+    struct gptoss_metal_function f32_gather_and_accumulate_e4_fn;
+    struct gptoss_metal_function f32_expert_routing_metadata_fn;
+    struct gptoss_metal_function f32_topk_softmax_e32_k4_fn;
+    struct gptoss_metal_function f32_topk_softmax_e128_k4_fn;
+    struct gptoss_metal_function f32_sdpa_q8_d64_fn;
+    struct gptoss_metal_function f32_softmax_fn;
+    struct gptoss_metal_function f32_sample_fn;
+
+    size_t per_block_shared_weights_size;
+    size_t per_expert_block_weight_size;
+
+    size_t embeddings_threadgroup_size;
+    size_t attn_qkv_threadgroup_size;
+    size_t attn_out_threadgroup_size;
+    size_t mlp_gate_threadgroup_size;
+    size_t mlp_swiglu_threadgroup_size;
+    size_t mlp_out_threadgroup_size;
+    size_t mlp_acc_threadgroup_size;
+    size_t unembedding_threadgroup_size;
+
+    size_t attn_rmsnorm_gain_offset;
+    size_t attn_qkv_weight_offset;
+    size_t attn_qkv_bias_offset;
+    size_t attn_sdpa_sink_offset;
+    size_t attn_out_weight_offset;
+    size_t attn_out_bias_offset;
+    size_t mlp_rmsnorm_gain_offset;
+    size_t mlp_gate_weight_offset;
+    size_t mlp_gate_bias_offset;
+    size_t mlp_swiglu_scale_offset;
+    size_t mlp_swiglu_bias_offset;
+    size_t mlp_out_block_offset;
+    size_t mlp_out_scale_offset;
+    size_t mlp_out_bias_offset;
+    size_t rmsnorm_weight_offset;
+    size_t unembedding_weight_offset;
+
+    // Buffer with non-MoE weights. Includes MoE gates, embeddings/unembeddings.
+    struct gptoss_metal_buffer shared_weight_buffer;
+    // num_blocks per-block buffers with MoE weights to follow.
+    struct gptoss_metal_buffer block_weight_buffers[];
+};
+
+#define GPTOSS_DEFAULT_BATCH_SIZE 128
+
+struct gptoss_context {
+#ifndef __cplusplus
+    atomic_uint_least64_t ref_count;
+#else
+    uint_least64_t ref_count;
+#endif
+
+    struct gptoss_model* model;
+    // Number of tokens processed in the context.
+    size_t num_tokens;
+    // Number of tokens in the KV cache.
+    size_t num_kv_tokens;
+    // Length of the context.
+    size_t max_tokens;
+    // Maximum number of tokens that can be processed in a single batch.
+    // Activation buffers are allocated with this size.
+    size_t max_batch_tokens;
+
+
+    size_t kvcache_size;
+    size_t allocation_size;
+
+    // Activation buffers.
+    // TODO: merge into a single buffer.
+    struct gptoss_metal_buffer residual_activation_buffer;  // Residual stream
+    struct gptoss_metal_buffer rmsnorm_activation_buffer;  // Both attention & MLP RMSNorm output
+    struct gptoss_metal_buffer qkv_activation_buffer;  // QKV projection output
+    struct gptoss_metal_buffer sdpa_activation_buffer;  // SDPA output
+    struct gptoss_metal_buffer gate_activation_buffer;  // MoE gating output
+    struct gptoss_metal_buffer expert_activation_buffer;  // MoE expert predictions
+    struct gptoss_metal_buffer expert_offset_buffer; // MoE expert histograms cumsum
+    struct gptoss_metal_buffer token_to_expert_routing_buffer; // MoE token to expert routing
+    struct gptoss_metal_buffer swiglu_input_buffer; // MLP+SwiGLU input for prefill.
+    struct gptoss_metal_buffer swiglu_activation_buffer;  // MLP+SwiGLU output
+    struct gptoss_metal_buffer moe_activation_buffer;  // MoE MLP output (per-active expert)
+
+    // Input/output buffers.
+    struct gptoss_metal_buffer control_buffer;
+    struct gptoss_metal_buffer token_buffer;  // uint32 token IDs
+    struct gptoss_metal_buffer score_buffer;  // unembedding outputs
+    struct gptoss_metal_buffer prob_buffer;
+    struct gptoss_metal_buffer sum_buffer;
+    struct gptoss_metal_buffer argmax_buffer;
+    struct gptoss_metal_buffer kvcache_buffer;
+};

gptoss_kernels/source/include/internal/rng.h

@@ -0,0 +1,24 @@
+#pragma once
+
+#include <stdint.h>
+
+inline static uint32_t rng_squares32(uint64_t offset, uint64_t seed) {
+    const uint64_t y = offset * seed;
+    const uint64_t z = y + seed;
+
+    /* Round 1 */
+    uint64_t x = y * y + y;
+    x = (x >> 32) | (x << 32);
+
+    /* Round 2 */
+    x = x * x + z;
+    x = (x >> 32) | (x << 32);
+
+    /* Round 3 */
+    x = x * x + y;
+    x = (x >> 32) | (x << 32);
+
+    /* Round 4 */
+    x = x * x + z;
+    return (uint32_t) (x >> 32);
+}

gptoss_kernels/source/include/internal/rng.hpp

@@ -0,0 +1,32 @@
+#pragma once
+
+#include <cstdint>
+
+namespace gptoss {
+
+namespace rng {
+
+inline static std::uint32_t squares32(std::uint64_t offset, std::uint64_t seed) {
+    const std::uint64_t y = offset * seed;
+    const std::uint64_t z = y + seed;
+
+    /* Round 1 */
+    std::uint64_t x = y * y + y;
+    x = (x >> 32) | (x << 32);
+
+    /* Round 2 */
+    x = x * x + z;
+    x = (x >> 32) | (x << 32);
+
+    /* Round 3 */
+    x = x * x + y;
+    x = (x >> 32) | (x << 32);
+
+    /* Round 4 */
+    x = x * x + z;
+    return static_cast<uint32_t>(x >> 32);
+}
+
+}  // namespace rng
+
+}  // namespace gptoss

gptoss_kernels/source/include/internal/storage.h

@@ -0,0 +1,36 @@
+#pragma once
+
+#include <stdbool.h>
+#include <stdint.h>
+
+struct gptoss_file_header {
+    char magic[12];
+    uint32_t zero;
+};
+
+struct gptoss_gptoss_model_header {
+    uint32_t context_length;
+    uint32_t num_blocks;
+    uint32_t num_experts;
+    uint32_t num_active_experts;
+    uint32_t embedding_dim;
+    uint32_t mlp_dim;
+    float swiglu_limit;
+    uint32_t head_dim;
+    uint32_t num_heads;
+    uint32_t num_kv_heads;
+    uint32_t attention_window;
+    float rope_theta;
+    float interpolation_scale;
+    float yarn_offset;
+    float yarn_scale;
+    float yarn_multiplier;
+    float rmsnorm_epsilon;
+};
+
+struct gptoss_tiktoken_tokenizer_header {
+    uint32_t num_special_tokens;
+    uint32_t num_text_tokens;
+    uint32_t regex_size;
+    uint32_t tokens_size;
+};

gptoss_kernels/source/include/internal/uuid.h

@@ -0,0 +1,114 @@
+#pragma once
+
+#include <stdbool.h>
+#include <stdint.h>
+#include <string.h>
+
+#include "internal/macros.h"
+
+
+struct GPTOSS_DENSELY_PACKED_STRUCTURE gptoss_uuid {
+    uint8_t bytes[16];
+};
+static_assert(sizeof(struct gptoss_uuid) == 16, "UUID size is not 16 bytes");
+
+
+#define UUID_FORMAT "%02X%02X%02X%02X-%02X%02X-%02X%02X-%02X%02X-%02X%02X%02X%02X%02X%02X"
+#define UUID_ARGS(uuid) (uuid).bytes[0], (uuid).bytes[1], (uuid).bytes[2], (uuid).bytes[3], \
+    (uuid).bytes[4], (uuid).bytes[5], (uuid).bytes[6], (uuid).bytes[7], (uuid).bytes[8], (uuid).bytes[9], \
+    (uuid).bytes[10], (uuid).bytes[11], (uuid).bytes[12], (uuid).bytes[13], (uuid).bytes[14], (uuid).bytes[15]
+
+static inline bool gptoss_is_gptoss_model_uuid(const struct gptoss_uuid* uuid) {
+    return memcmp(
+        &(struct gptoss_uuid) {0xDF, 0x52, 0xDC, 0x86, 0x17, 0x89, 0x4E, 0xD0, 0xA2, 0x95, 0x66, 0xF1, 0x05, 0x08, 0x14, 0x5B},
+        uuid,
+        sizeof(struct gptoss_uuid)) == 0;
+}
+
+static inline bool gptoss_is_applegpu_layout_uuid(const struct gptoss_uuid* uuid) {
+    return memcmp(
+        &(struct gptoss_uuid) {0x22, 0x91, 0x77, 0xA8, 0x57, 0x75, 0x42, 0x68, 0xBF, 0xD8, 0xD5, 0x88, 0xB3, 0x51, 0xC5, 0x6D},
+        uuid,
+        sizeof(struct gptoss_uuid)) == 0;
+}
+
+static inline bool gptoss_is_tiktoken_tokenizer_uuid(const struct gptoss_uuid* uuid) {
+    return memcmp(
+        &(struct gptoss_uuid) {0x74, 0x01, 0xAD, 0xED, 0x2A, 0x95, 0x40, 0xCB, 0xB7, 0x82, 0x9C, 0xCE, 0xBA, 0xAF, 0xE7, 0x2B},
+        uuid,
+        sizeof(struct gptoss_uuid)) == 0;
+}
+
+static inline enum gptoss_special_token gptoss_special_token_decode_uuid(const struct gptoss_uuid* uuid) {
+    if (memcmp(
+        &(struct gptoss_uuid) {0x55, 0xA7, 0x7C, 0x2F, 0x8A, 0x01, 0x4C, 0x54, 0x8A, 0xC2, 0x31, 0x3B, 0xFC, 0x7E, 0x20, 0x8D},
+        uuid,
+        sizeof(struct gptoss_uuid)) == 0)
+    {
+        return gptoss_special_token_start;
+    } else if (memcmp(
+        &(struct gptoss_uuid) {0x16, 0xE4, 0x04, 0x31, 0xF4, 0x7F, 0x4B, 0x22, 0xB5, 0x9B, 0x8B, 0x27, 0x8F, 0xC3, 0x0A, 0x54},
+        uuid,
+        sizeof(struct gptoss_uuid)) == 0)
+    {
+        return gptoss_special_token_message;
+    } else if (memcmp(
+        &(struct gptoss_uuid) {0xFC, 0xAC, 0x2F, 0x6D, 0x47, 0x05, 0x4F, 0x6B, 0xB2, 0x28, 0x64, 0x2A, 0xCC, 0xAC, 0x72, 0x38},
+        uuid,
+        sizeof(struct gptoss_uuid)) == 0)
+    {
+        return gptoss_special_token_end;
+    } else if (memcmp(
+        &(struct gptoss_uuid) {0xF7, 0x99, 0xFF, 0x69, 0x19, 0x92, 0x43, 0xC4, 0xA3, 0xD8, 0xD8, 0x31, 0xF4, 0x75, 0xDC, 0x75},
+        uuid,
+        sizeof(struct gptoss_uuid)) == 0)
+    {
+        return gptoss_special_token_return;
+    } else if (memcmp(
+        &(struct gptoss_uuid) {0xE1, 0x5B, 0xA7, 0x02, 0x28, 0xC4, 0x42, 0x92, 0xAB, 0x8F, 0xFF, 0xA4, 0x34, 0x70, 0x91, 0x28},
+        uuid,
+        sizeof(struct gptoss_uuid)) == 0)
+    {
+        return gptoss_special_token_refusal;
+    } else if (memcmp(
+        &(struct gptoss_uuid) {0xC0, 0xBB, 0x14, 0xC7, 0x60, 0x22, 0x49, 0xDA, 0xAD, 0x08, 0x79, 0x2D, 0x67, 0xE8, 0xB4, 0x70},
+        uuid,
+        sizeof(struct gptoss_uuid)) == 0)
+    {
+        return gptoss_special_token_constrain;
+    } else if (memcmp(
+        &(struct gptoss_uuid) {0xFD, 0x3D, 0xDA, 0x11, 0xC8, 0xAB, 0x40, 0x33, 0x87, 0x6E, 0xD9, 0x3D, 0xEB, 0x17, 0x2C, 0x93},
+        uuid,
+        sizeof(struct gptoss_uuid)) == 0)
+    {
+        return gptoss_special_token_channel;
+    } else if (memcmp(
+        &(struct gptoss_uuid) {0x12, 0x20, 0xF7, 0x96, 0xE3, 0x88, 0x4D, 0xE5, 0xB4, 0x87, 0xFE, 0x2E, 0xB5, 0xFE, 0x03, 0xC0},
+        uuid,
+        sizeof(struct gptoss_uuid)) == 0)
+    {
+        return gptoss_special_token_call;
+    } else if (memcmp(
+        &(struct gptoss_uuid) {0x07, 0xD7, 0xDA, 0x55, 0xB3, 0x46, 0x4C, 0xFF, 0x8B, 0x37, 0x7C, 0xEF, 0xAC, 0xF8, 0xA3, 0xE8},
+        uuid,
+        sizeof(struct gptoss_uuid)) == 0)
+    {
+        return gptoss_special_token_untrusted;
+    } else if (memcmp(
+        &(struct gptoss_uuid) {0xF2, 0x65, 0xBD, 0x9C, 0xC7, 0x17, 0x46, 0x9E, 0xA4, 0x47, 0x92, 0x06, 0x87, 0xD6, 0x5D, 0x90},
+        uuid,
+        sizeof(struct gptoss_uuid)) == 0)
+    {
+        return gptoss_special_token_end_untrusted;
+    } else if (memcmp(
+        &(struct gptoss_uuid) {0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00},
+        uuid,
+        sizeof(struct gptoss_uuid)) == 0)
+    {
+        // Suppress warning
+        return gptoss_special_token_invalid;
+    } else {
+        GPTOSS_LOG_WARNING("unsupported special token " UUID_FORMAT, UUID_ARGS(*uuid));
+        return gptoss_special_token_invalid;
+    }
+}

gptoss_kernels/source/log.c

@@ -0,0 +1,50 @@
+#include <assert.h>  // assert
+#include <stdarg.h>  // va_list, va_copy, va_end
+#include <stdio.h>  // vsnprintf
+#include <stdlib.h>  // malloc, free
+
+#include <unistd.h>  // STDERR_FILENO
+
+
+
+#define GPTOSS_ON_STACK_FORMAT_BUFFER_SIZE 16384
+
+void gptoss_format_log(const char* format, va_list args) {
+    char stack_buffer[GPTOSS_ON_STACK_FORMAT_BUFFER_SIZE];
+    char* heap_buffer = NULL;
+
+    va_list args_copy;
+    va_copy(args_copy, args);
+
+    const int vsnprintf_result = vsnprintf(stack_buffer, GPTOSS_ON_STACK_FORMAT_BUFFER_SIZE, format, args);
+    assert(vsnprintf_result >= 0);
+
+    // At least a partially formatted buffer is ready.
+    char* message_buffer = &stack_buffer[0];
+    size_t message_size = (size_t) vsnprintf_result;
+    if (message_size > GPTOSS_ON_STACK_FORMAT_BUFFER_SIZE) {
+        heap_buffer = malloc(message_size);
+        if (heap_buffer == NULL) {
+            // Fall back to the truncated message in the on-stack buffer.
+            message_size = GPTOSS_ON_STACK_FORMAT_BUFFER_SIZE;
+        } else {
+            // Use the full message in the in-heap buffer.
+            vsnprintf(heap_buffer, message_size, format, args_copy);
+            message_buffer = heap_buffer;
+        }
+    }
+
+    ssize_t bytes_written;
+    do {
+        bytes_written = write(STDERR_FILENO, message_buffer, message_size);
+        if (bytes_written > 0) {
+            assert((size_t) bytes_written <= message_size);
+            message_buffer += bytes_written;
+            message_size -= bytes_written;
+        }
+    } while (bytes_written >= 0 && message_size != 0);
+
+cleanup:
+    free(heap_buffer);
+    va_end(args_copy);
+}

gptoss_kernels/source/matmul.metal

@@ -0,0 +1,422 @@
+#include <metal_atomic>
+#include <metal_compute>
+#include <metal_integer>
+#include <metal_math>
+#include <metal_simdgroup>
+#include <metal_stdlib>
+
+#include <internal/kernel-args.h>
+
+#pragma METAL fp math_mode(safe)
+#pragma METAL fp contract(off)
+
+
+// Each simdgroup reduces all channels of the input and computes a single channel of the output
+// + Efficient synchronization
+// + Sequential memory access within a warp
+// Each threadgroup computes (simdgroups_per_threadgroup) consecutive output channels
+// + Reuse input vector from threadgroup memory
+// + Avoid synchronization across warps when doing reduction
+
+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 = (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;
+        }
+    }
+}
+
+kernel void gptoss_f32_bf16w_matmul_qkv(
+    constant gptoss_qkv_args& args [[ buffer(0) ]],
+    const device float4* input [[ buffer(1) ]],
+    const device bfloat4* weight [[ buffer(2) ]],
+    const device bfloat* bias [[ buffer(3) ]],
+    device float* q [[ buffer(4) ]],
+    device float* kv [[ buffer(5) ]],
+    const device gptoss_control* control [[ buffer(6) ]],
+    threadgroup void* scratch [[ threadgroup(0) ]],
+    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;
+    const uint head_dim = 64;
+    const uint num_q_heads = 64;
+    const uint num_kv_heads = 8;
+    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;
+    q += gid.y * args.num_rows;
+
+    uint 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);
+        static_cast<threadgroup float*>(scratch)[simdgroup_idx] = sum;
+    }
+    metal::threadgroup_barrier(metal::mem_flags::mem_threadgroup);
+    if (simdgroup_idx == 0) {
+        const uint num_half_simdgroups = num_simdgroups / 2;
+        if (simdgroup_tid < num_half_simdgroups) {
+            float2 vals = static_cast<const threadgroup float2*>(scratch)[simdgroup_tid];
+            const uint idx = gid.x * num_half_simdgroups + simdgroup_tid;
+            const uint head_idx = idx / (head_dim / 2);
+            const uint token_idx = args.token_offset + gid.y;
+            const uint dim_idx = idx % (head_dim / 2);
+            if (head_idx < num_q_heads + num_kv_heads) {
+                const float dim_idx_val = static_cast<float>(dim_idx);
+                const float inv_extrapolation_freq = metal::precise::exp(dim_idx_val * args.freq_scale);
+                const float inv_interpolation_freq = inv_extrapolation_freq * args.interpolation_scale;
+                const float alpha = metal::saturate(metal::fma(dim_idx_val, args.yarn_scale, args.yarn_offset));
+                const float inv_freq = metal::mix(inv_extrapolation_freq, inv_interpolation_freq, alpha);
+
+                const float phi = static_cast<float>(token_idx) * inv_freq;
+                const float yarn_multiplier = args.yarn_multiplier;
+                float cosphi;
+                const float sinphi = metal::precise::sincos(phi, cosphi) * yarn_multiplier;
+                cosphi *= yarn_multiplier;
+
+                vals = (float2) {
+                    vals.x * cosphi - vals.y * sinphi,
+                    vals.x * sinphi + vals.y * cosphi,
+                };
+            }
+            if (head_idx < num_q_heads) {
+                reinterpret_cast<device float2*>(q)[idx] = vals;
+            } else if (head_idx < num_q_heads + num_kv_heads) {
+                const uint h = head_idx - num_q_heads;
+                reinterpret_cast<device float2*>(kv + (h * args.max_tokens + token_idx) * 2 * head_dim)[dim_idx] = vals;
+            } else {
+                const uint h = head_idx - num_q_heads - num_kv_heads;
+                reinterpret_cast<device float2*>(kv + (h * args.max_tokens + token_idx) * 2 * head_dim + head_dim)[dim_idx] = vals;
+            }
+        }
+    }
+}
+
+kernel void gptoss_f32_bf16w_unembedding(
+    constant gptoss_unembedding_args& args [[ buffer(0) ]],
+    const device float4* input [[ buffer(1) ]],
+    const device bfloat4* weight [[ buffer(2) ]],
+    device float* output [[ buffer(3) ]],
+    device metal::atomic_ulong* argmax [[ 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;
+    threadgroup uint2 threadgroup_buffer[32];
+    if (control->abort != 0) {
+        return;
+    }
+
+    const uint num_column_vecs = args.num_column_vecs;
+    const uint row_start = gid.x * args.num_rows_per_threadgroup + simdgroup_idx;
+    const uint row_end = metal::min(gid.x * args.num_rows_per_threadgroup + args.num_rows_per_threadgroup, args.num_rows);
+    const uint num_iter = (num_column_vecs - simdgroup_tid + (simdgroup_size - 1)) / simdgroup_size;
+
+    input += gid.y * num_column_vecs + simdgroup_tid;
+    weight += num_column_vecs * row_start + simdgroup_tid;
+    output += gid.y * args.num_rows + row_start;
+
+    uint2 row_sum{0xFFFFFFFFul, 0xFFFFFFFFul};
+    for (uint row = row_start; row < row_end; row += num_simdgroups) {
+        uint n = num_iter;
+
+        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 (--n != 0);
+        input -= num_iter * simdgroup_size;
+        weight -= num_iter * simdgroup_size;
+
+        const float2 sum2 = sum4.xy + sum4.zw;
+        float sum = sum2.x + sum2.y;
+        sum = metal::simd_sum(sum);
+        uint sum_bits = as_type<uint>(sum);
+        if (static_cast<int>(sum_bits) >= 0) {
+            sum_bits ^= 0x7FFFFFFFu;
+        }
+        row_sum = as_type<uint2>(metal::min(as_type<ulong>(row_sum), as_type<ulong>(uint2{row, sum_bits})));
+        if (metal::simd_is_first()) {
+            *output = sum;
+        }
+
+        weight += num_column_vecs * num_simdgroups;
+        output += num_simdgroups;
+    }
+    if (metal::simd_is_first()) {
+        threadgroup_buffer[simdgroup_idx] = row_sum;
+    }
+    metal::threadgroup_barrier(metal::mem_flags::mem_threadgroup);
+    if (simdgroup_idx == 0) {
+        // Min-Reduce threadgroup_buffer
+        if (simdgroup_tid < num_simdgroups) {
+            row_sum = threadgroup_buffer[simdgroup_tid];
+        }
+        const uint sum_bits = row_sum.y;
+        const uint sum_bits_min = metal::simd_min(sum_bits);
+        const uint row_min = metal::simd_min(sum_bits == sum_bits_min ? row_sum.x : 0xFFFFFFFFu);
+        if (metal::simd_is_first()) {
+            const uint2 threadgroup_output{row_min, sum_bits_min};
+            atomic_min_explicit(&argmax[gid.y], as_type<ulong>(threadgroup_output), metal::memory_order_relaxed);
+        }
+    }
+}
+
+// Current constraints for the dense matmul kernel:
+//  1- All B* and Sg_* are a multiple of 8.
+//  2- Bm is divisible by Sg_n and Bn is divisible by Sg_n.
+//  3- M, N and K are all divisible by 8..
+template <uint Bm, uint Bn, uint Bk, uint Sg_Bm, uint Sg_Bn, uint add = 0>
+inline void _gptoss_f32_bf16w_dense_matmul_impl(
+    constant gptoss_dense_matmul_args& args, const device float* lhs,
+    const device bfloat* rhs, const device bfloat* __restrict__ bias,
+    device float* out, const device gptoss_control* control, threadgroup float* scratch, threadgroup float* bias_tile,
+    uint sg_id, uint sg_count_per_tg, uint3 gid, uint3 tg_id, uint3 local_tid,
+    uint3 threadgroup_size) {
+
+    if (control->abort != 0) {
+        return;
+    }
+
+    // The kernel assumes that M, K, and N are divisible by 8.
+    const uint M = args.m;
+    const uint K = args.k;
+    const uint N = args.n;
+    static_assert((Bm % 8u) == 0u, "Bm must be a multiple of 8");
+    static_assert((Bn % 8u) == 0u, "Bn must be a multiple of 8");
+    static_assert((Bk % 8u) == 0u, "Bk must be a multiple of 8");
+    static_assert((Sg_Bm % 8u) == 0u, "Bk must be a multiple of 8");
+    static_assert((Sg_Bn % 8u) == 0u, "Bk must be a multiple of 8");
+    static_assert((Bn % Sg_Bn) == 0u, "Bn must be a multiple of Sg_Bn");
+    static_assert((Bm % Sg_Bm) == 0u, "Bm must be a multiple of Sg_Bm");
+
+    // Get row and col tg.
+    const uint row_tg = tg_id.y;
+    const uint col_tg = tg_id.x;
+    // Get row and col local tid.
+    const uint row_tg_offset = row_tg * Bm;
+    const uint col_tg_offset = col_tg * Bn;
+
+    const uint sg_col_count = Bn / Sg_Bn;
+    const uint row_sg = sg_id / sg_col_count;
+    const uint col_sg = sg_id % sg_col_count;
+
+    const uint row_sg_offset = row_sg * Sg_Bm;
+    const uint col_sg_offset = col_sg * Sg_Bn;
+    constexpr uint temp_result_size = (Sg_Bm / 8) * (Sg_Bn / 8);
+    // Create an array of simdgroup_float8x8 to hold temp results.
+    metal::simdgroup_float8x8 OutTiles[temp_result_size];
+#pragma clang loop unroll(full)
+    for (uint i = 0; i < temp_result_size; i++) {
+        OutTiles[i] = metal::make_filled_simdgroup_matrix<float, 8, 8>(
+            static_cast<float>(0.0));
+    }
+
+    for (uint k_offset = 0; k_offset < K; k_offset += Bk) {
+#pragma clang loop unroll(full)
+        for (uint k = 0; k < Bk; k += 8) {
+#pragma clang loop unroll(full)
+            for (uint m_subtile_ = 0; m_subtile_ < Sg_Bm; m_subtile_ += 8) {
+                // const uint m_subtile = row_sg_offset + m_subtile_;
+                // const uint row_index_in_out_tile = (m_subtile - row_sg_offset) / 8;
+                const uint row_index_in_out_tile = m_subtile_ / 8;
+                metal::simdgroup_float8x8 LHStile;
+                const uint k_id = k + k_offset;
+                const uint row_offset = row_tg_offset + row_sg_offset + m_subtile_;
+                metal::simdgroup_load(LHStile, lhs, K, ulong2(k_id, row_offset));
+                metal::simdgroup_bfloat8x8 RHStile;
+#pragma clang loop unroll(full)
+                for (uint n_subtile_ = 0; n_subtile_ < Sg_Bn; n_subtile_ += 8) {
+                    const uint col_index_in_out_tile = n_subtile_ / 8;
+                    const uint current_index_out_tile =
+                        row_index_in_out_tile * (Sg_Bn / 8) + col_index_in_out_tile;
+                    const uint col_offset = col_tg_offset + col_sg_offset + n_subtile_;
+                    simdgroup_load(RHStile, rhs, K, ulong2(k_id, col_offset), /*transpose=*/true);
+                    // If rhs was not transposed, use the following instead:
+                    // simdgroup_load(RHStile, rhs, N, ulong2(col_offset, k_id));
+                    simdgroup_multiply_accumulate(OutTiles[current_index_out_tile],
+                                                  LHStile, RHStile,
+                                                  OutTiles[current_index_out_tile]);
+                }
+            }
+        }
+    }
+    // Epilogue.
+#pragma clang loop unroll(full)
+    for (uint n_subtile_ = 0; n_subtile_ < Sg_Bn; n_subtile_ += 8) {
+        const uint col_index_in_out_tile = n_subtile_ / 8;
+        const uint local_col_offset = col_sg_offset + n_subtile_;
+#pragma clang loop unroll(full)
+        for (uint m_subtile_ = 0; m_subtile_ < Sg_Bm; m_subtile_ += 8) {
+            const uint row_index_in_out_tile = m_subtile_ / 8;
+            const uint local_row_offset = row_sg_offset + m_subtile_;
+            const uint current_index_out_tile =
+                row_index_in_out_tile * (Sg_Bn / 8) + col_index_in_out_tile;
+            simdgroup_store(OutTiles[current_index_out_tile], scratch, Bn,
+                            ulong2(local_col_offset, local_row_offset));
+        }
+    }
+    // TODO(ibahmed): vectorize these loads an maybe unroll the loop.
+    const uint thread_count_per_tg =
+        threadgroup_size.x * threadgroup_size.y * threadgroup_size.z;
+    for (uint c_local = local_tid.x; c_local < Bn;
+         c_local += thread_count_per_tg) {
+        const uint c_global = col_tg_offset + c_local;
+        bias_tile[c_local] =
+            (c_global < N) ? static_cast<float>(bias[c_global]) : 0.0f;
+    }
+
+    metal::threadgroup_barrier(metal::mem_flags::mem_threadgroup);
+
+    // TODO(ibahmed): vectorize these stores and maybe unroll the loop.
+    for (uint idx = local_tid.x; idx < Bm * Bn; idx += thread_count_per_tg) {
+        const uint r = idx / Bn;
+        const uint c = idx % Bn;
+
+        const uint out_row = row_tg_offset + r;
+        const uint out_col = col_tg_offset + c;
+
+        if (out_row < M && out_col < N) {
+            float acc = scratch[idx] + bias_tile[c];
+            if (add) {
+                acc += out[out_row * N + out_col];
+            }
+            out[out_row * N + out_col] = acc;
+        }
+    }
+}
+
+kernel void gptoss_f32_bf16w_dense_matmul_qkv(
+    constant gptoss_dense_matmul_args& args [[buffer(0)]],
+    const device float* lhs [[buffer(1)]],
+    const device bfloat* rhs [[buffer(2)]],
+    const device bfloat* __restrict__ bias [[buffer(3)]],
+    device float* out [[buffer(4)]],
+    const device gptoss_control* control [[buffer(5)]],
+    uint sg_id [[simdgroup_index_in_threadgroup]],
+    uint sg_count_per_tg [[dispatch_simdgroups_per_threadgroup]],
+    uint3 gid [[thread_position_in_grid]],
+    uint3 tg_id [[threadgroup_position_in_grid]],
+    uint3 local_tid [[thread_position_in_threadgroup]],
+    uint3 threadgroup_size [[threads_per_threadgroup]]) {
+    threadgroup float scratch[QKV_Bm * QKV_Bn];
+    threadgroup float bias_tile[QKV_Bn];
+    _gptoss_f32_bf16w_dense_matmul_impl<QKV_Bm, QKV_Bn, QKV_Bk, QKV_Sg_Bm,
+                                        QKV_Sg_Bn>(
+        args, lhs, rhs, bias, out, control, scratch, bias_tile, sg_id, sg_count_per_tg,
+        gid, tg_id, local_tid, threadgroup_size);
+}
+
+kernel void gptoss_f32_bf16w_dense_matmul_attn_output(
+    constant gptoss_dense_matmul_args& args [[buffer(0)]],
+    const device float* lhs [[buffer(1)]],
+    const device bfloat* rhs [[buffer(2)]],
+    const device bfloat* __restrict__ bias [[buffer(3)]],
+    device float* out [[buffer(4)]],
+    const device gptoss_control* control [[buffer(5)]],
+    uint sg_id [[simdgroup_index_in_threadgroup]],
+    uint sg_count_per_tg [[dispatch_simdgroups_per_threadgroup]],
+    uint3 gid [[thread_position_in_grid]],
+    uint3 tg_id [[threadgroup_position_in_grid]],
+    uint3 local_tid [[thread_position_in_threadgroup]],
+    uint3 threadgroup_size [[threads_per_threadgroup]]) {
+    threadgroup float scratch[ATTN_OUTPUT_Bm * ATTN_OUTPUT_Bn];
+    threadgroup float bias_tile[ATTN_OUTPUT_Bn];
+    _gptoss_f32_bf16w_dense_matmul_impl<ATTN_OUTPUT_Bm, ATTN_OUTPUT_Bn,
+                                        ATTN_OUTPUT_Bk, ATTN_OUTPUT_Sg_Bm,
+                                        ATTN_OUTPUT_Sg_Bn, /*add=*/1>(
+        args, lhs, rhs, bias, out, control, scratch, bias_tile, sg_id, sg_count_per_tg,
+        gid, tg_id, local_tid, threadgroup_size);
+}
+
+kernel void gptoss_f32_bf16w_dense_matmul_mlp_gate(
+    constant gptoss_dense_matmul_args& args [[buffer(0)]],
+    const device float* lhs [[buffer(1)]],
+    const device bfloat* rhs [[buffer(2)]],
+    const device bfloat* __restrict__ bias [[buffer(3)]],
+    device float* out [[buffer(4)]],
+    const device gptoss_control* control [[buffer(5)]],
+    uint sg_id [[simdgroup_index_in_threadgroup]],
+    uint sg_count_per_tg [[dispatch_simdgroups_per_threadgroup]],
+    uint3 gid [[thread_position_in_grid]],
+    uint3 tg_id [[threadgroup_position_in_grid]],
+    uint3 local_tid [[thread_position_in_threadgroup]],
+    uint3 threadgroup_size [[threads_per_threadgroup]]) {
+    threadgroup float scratch[MLP_GATE_Bm * MLP_GATE_Bn];
+    threadgroup float bias_tile[MLP_GATE_Bn];
+    _gptoss_f32_bf16w_dense_matmul_impl<MLP_GATE_Bm, MLP_GATE_Bn, MLP_GATE_Bk,
+                                        MLP_GATE_Sg_Bm, MLP_GATE_Sg_Bn>(
+        args, lhs, rhs, bias, out, control, scratch, bias_tile, sg_id, sg_count_per_tg,
+        gid, tg_id, local_tid, threadgroup_size);
+}

gptoss_kernels/source/metal-kernels.c

@@ -0,0 +1,1518 @@
+#include <inttypes.h>
+#include <stddef.h>
+#include <stdint.h>
+#include <math.h>
+
+#include <internal/kernel-args.h>
+#include <internal/log.h>
+#include <internal/math.h>
+#include <internal/metal.h>
+#include <internal/metal-kernels.h>
+
+
+enum gptoss_status gptoss_metal_command_buffer_encode_launch_u32_fill_random(
+    const struct gptoss_metal_command_buffer* command_buffer,
+    const struct gptoss_metal_function* u32_fill_random_fn,
+    size_t threadgroup_size,
+    size_t max_threadgroups,
+    const struct gptoss_metal_buffer* output_buffer,
+    size_t output_offset,
+    uint64_t num_elements,
+    uint64_t rng_seed,
+    uint64_t rng_offset)
+{
+    if (command_buffer->object == NULL || u32_fill_random_fn->pipeline_state_object == NULL) {
+        return gptoss_status_invalid_state;
+    }
+
+    if (threadgroup_size == 0) {
+        threadgroup_size = u32_fill_random_fn->max_threadgroup_threads;
+    } else if (threadgroup_size > u32_fill_random_fn->max_threadgroup_threads) {
+        return gptoss_status_invalid_argument;
+    }
+
+    const size_t num_vecs = num_elements;
+    const size_t num_vecs_per_threadgroup = math_ceil_div(num_vecs, max_threadgroups * threadgroup_size) * threadgroup_size;
+    const size_t num_threadgroups = math_min(max_threadgroups, math_ceil_div(num_vecs, num_vecs_per_threadgroup));
+    const struct gptoss_u32_fill_random_args args = {
+        .num_vecs = num_vecs,
+        .num_vecs_per_threadgroup = num_vecs_per_threadgroup,
+        .seed = rng_seed,
+        .offset = rng_offset,
+    };
+
+    return gptoss_metal_command_buffer_encode_launch_kernel(
+        command_buffer, u32_fill_random_fn,
+        threadgroup_size, 1, 1,
+        num_threadgroups, 1, 1,
+        sizeof(args), &args,
+        1, &output_buffer, &output_offset,
+        /*threadgroup_buffer_size=*/0);
+}
+
+enum gptoss_status gptoss_metal_command_buffer_encode_launch_f32_fill_random(
+    const struct gptoss_metal_command_buffer* command_buffer,
+    const struct gptoss_metal_function* f32_fill_random_fn,
+    size_t threadgroup_size,
+    size_t max_threadgroups,
+    const struct gptoss_metal_buffer* output_buffer,
+    size_t output_offset,
+    uint64_t num_elements,
+    uint64_t rng_seed,
+    uint64_t rng_offset,
+    float rng_min,
+    float rng_max)
+{
+    if (command_buffer->object == NULL || f32_fill_random_fn->pipeline_state_object == NULL) {
+        return gptoss_status_invalid_state;
+    }
+
+    if (threadgroup_size == 0) {
+        threadgroup_size = f32_fill_random_fn->max_threadgroup_threads;
+    } else if (threadgroup_size > f32_fill_random_fn->max_threadgroup_threads) {
+        return gptoss_status_invalid_argument;
+    }
+
+    if (rng_min >= rng_max) {
+        return gptoss_status_invalid_argument;
+    }
+
+    const size_t num_vecs = num_elements;
+    const size_t num_vecs_per_threadgroup = math_ceil_div(num_vecs, max_threadgroups * threadgroup_size) * threadgroup_size;
+    const size_t num_threadgroups = math_min(max_threadgroups, math_ceil_div(num_vecs, num_vecs_per_threadgroup));
+    const struct gptoss_f32_fill_random_args args = {
+        .num_vecs = num_vecs,
+        .num_vecs_per_threadgroup = num_vecs_per_threadgroup,
+        .seed = rng_seed,
+        .offset = rng_offset,
+        .scale = (rng_max - rng_min) * 0x1.0p-32f,
+        .bias = (rng_min + rng_max) * 0.5f,
+    };
+
+    return gptoss_metal_command_buffer_encode_launch_kernel(
+        command_buffer, f32_fill_random_fn,
+        threadgroup_size, 1, 1,
+        num_threadgroups, 1, 1,
+        sizeof(args), &args,
+        1, &output_buffer, &output_offset,
+        /*threadgroup_buffer_size=*/0);
+}
+
+enum gptoss_status gptoss_metal_command_buffer_encode_launch_bf16_fill_random(
+    const struct gptoss_metal_command_buffer* command_buffer,
+    const struct gptoss_metal_function* bf16_fill_random_fn,
+    size_t threadgroup_size,
+    size_t max_threadgroups,
+    const struct gptoss_metal_buffer* output_buffer,
+    size_t output_offset,
+    uint64_t num_elements,
+    uint64_t rng_seed,
+    uint64_t rng_offset,
+    float rng_min,
+    float rng_max)
+{
+    if (command_buffer->object == NULL || bf16_fill_random_fn->pipeline_state_object == NULL) {
+        return gptoss_status_invalid_state;
+    }
+
+    if (threadgroup_size == 0) {
+        threadgroup_size = bf16_fill_random_fn->max_threadgroup_threads;
+    } else if (threadgroup_size > bf16_fill_random_fn->max_threadgroup_threads) {
+        return gptoss_status_invalid_argument;
+    }
+
+    if (rng_min >= rng_max) {
+        return gptoss_status_invalid_argument;
+    }
+
+    const size_t num_vecs = num_elements;
+    const size_t num_vecs_per_threadgroup = math_ceil_div(num_vecs, max_threadgroups * threadgroup_size) * threadgroup_size;
+    const size_t num_threadgroups = math_min(max_threadgroups, math_ceil_div(num_vecs, num_vecs_per_threadgroup));
+    const struct gptoss_f32_fill_random_args args = {
+        .num_vecs = num_vecs,
+        .num_vecs_per_threadgroup = num_vecs_per_threadgroup,
+        .seed = rng_seed,
+        .offset = rng_offset,
+        .scale = (rng_max - rng_min) * 0x1.0p-32f,
+        .bias = (rng_min + rng_max) * 0.5f,
+    };
+
+    return gptoss_metal_command_buffer_encode_launch_kernel(
+        command_buffer, bf16_fill_random_fn,
+        threadgroup_size, 1, 1,
+        num_threadgroups, 1, 1,
+        sizeof(args), &args,
+        1, &output_buffer, &output_offset,
+        /*threadgroup_buffer_size=*/0);
+}
+
+enum gptoss_status gptoss_metal_command_buffer_encode_launch_mf4_f32_convert(
+    const struct gptoss_metal_command_buffer* command_buffer,
+    const struct gptoss_metal_function* mf4_f32_convert_fn,
+    size_t threadgroup_size,
+    size_t max_threadgroups,
+    const struct gptoss_metal_buffer* block_buffer,
+    const struct gptoss_metal_buffer* scale_buffer,
+    const struct gptoss_metal_buffer* output_buffer,
+    uint64_t num_elements)
+{
+    if (command_buffer->object == NULL || mf4_f32_convert_fn->pipeline_state_object == NULL) {
+        return gptoss_status_invalid_state;
+    }
+
+    if (num_elements % 32 != 0) {
+        return gptoss_status_invalid_argument;
+    }
+
+    if (threadgroup_size == 0) {
+        threadgroup_size = mf4_f32_convert_fn->max_threadgroup_threads;
+    } else if (threadgroup_size > mf4_f32_convert_fn->max_threadgroup_threads) {
+        return gptoss_status_invalid_argument;
+    }
+
+    const size_t num_vecs = num_elements / 32;
+    const size_t num_vecs_per_threadgroup = math_ceil_div(num_vecs, max_threadgroups * threadgroup_size) * threadgroup_size;
+    const size_t num_threadgroups = math_min(max_threadgroups, math_ceil_div(num_vecs, num_vecs_per_threadgroup));
+    const struct gptoss_convert_args args = {
+        .num_vecs = num_vecs,
+        .num_vecs_per_threadgroup = num_vecs_per_threadgroup,
+    };
+
+    return gptoss_metal_command_buffer_encode_launch_kernel(
+        command_buffer, mf4_f32_convert_fn,
+        threadgroup_size, 1, 1,
+        num_threadgroups, 1, 1,
+        sizeof(args), &args,
+        3, (const struct gptoss_metal_buffer *[]) {block_buffer, scale_buffer, output_buffer}, NULL,
+        /*threadgroup_buffer_size=*/0);
+}
+
+enum gptoss_status gptoss_metal_command_buffer_encode_launch_bf16_f32_embeddings(
+    const struct gptoss_metal_command_buffer* command_buffer,
+    const struct gptoss_metal_function* bf16_f32_embeddings_fn,
+    size_t threadgroup_size,
+    const struct gptoss_metal_buffer* token_buffer,
+    size_t token_offset,
+    const struct gptoss_metal_buffer* weight_buffer,
+    size_t weight_offset,
+    const struct gptoss_metal_buffer* output_buffer,
+    size_t output_offset,
+    const struct gptoss_metal_buffer* control_buffer,
+    size_t control_offset,
+    uint32_t num_tokens,
+    uint32_t num_channels)
+{
+    if (command_buffer->object == NULL || bf16_f32_embeddings_fn->pipeline_state_object == NULL) {
+        return gptoss_status_invalid_state;
+    }
+
+    if (num_channels % 4 != 0) {
+        return gptoss_status_invalid_argument;
+    }
+
+    if (threadgroup_size == 0) {
+        threadgroup_size = bf16_f32_embeddings_fn->max_threadgroup_threads;
+    } else if (threadgroup_size > bf16_f32_embeddings_fn->max_threadgroup_threads) {
+        return gptoss_status_invalid_argument;
+    }
+
+    const uint32_t num_vecs = num_channels / 4;
+    const struct gptoss_embeddings_args args = {
+        .num_vecs = num_vecs,
+    };
+
+    return gptoss_metal_command_buffer_encode_launch_kernel(
+        command_buffer, bf16_f32_embeddings_fn,
+        threadgroup_size, 1, 1,
+        num_tokens, 1, 1,
+        sizeof(args), &args,
+        4,
+        (const struct gptoss_metal_buffer *[]) {token_buffer, weight_buffer, output_buffer, control_buffer},
+        (const size_t[]) {token_offset, weight_offset, output_offset, control_offset},
+        /*threadgroup_buffer_size=*/0);
+}
+
+enum gptoss_status gptoss_metal_command_buffer_encode_launch_f32_bf16w_rmsnorm(
+    const struct gptoss_metal_command_buffer* command_buffer,
+    const struct gptoss_metal_function* f32_bf16w_rmsnorm_fn,
+    const struct gptoss_metal_buffer* input_buffer,
+    size_t input_offset,
+    const struct gptoss_metal_buffer* weight_buffer,
+    size_t weight_offset,
+    const struct gptoss_metal_buffer* output_buffer,
+    size_t output_offset,
+    const struct gptoss_metal_buffer* control_buffer,
+    size_t control_offset,
+    uint32_t num_tokens,
+    uint32_t num_channels,
+    float epsilon)
+{
+    if (command_buffer->object == NULL || f32_bf16w_rmsnorm_fn->pipeline_state_object == NULL) {
+        return gptoss_status_invalid_state;
+    }
+
+    if (num_channels % 4 != 0) {
+        return gptoss_status_invalid_argument;
+    }
+
+    if (f32_bf16w_rmsnorm_fn->max_threadgroup_threads < 1024) {
+        return gptoss_status_unsupported_system;
+    }
+
+    if (f32_bf16w_rmsnorm_fn->simdgroup_threads != 32) {
+        return gptoss_status_unsupported_system;
+    }
+
+    const uint32_t num_vecs = num_channels / 4;
+    const struct gptoss_rmsnorm_args args = {
+        .num_vecs = num_vecs,
+        .num_channels = (float) num_channels,
+        .epsilon = epsilon,
+    };
+
+    return gptoss_metal_command_buffer_encode_launch_kernel(
+        command_buffer, f32_bf16w_rmsnorm_fn,
+        /*threadgroup_size=*/1024, 1, 1,
+        num_tokens, 1, 1,
+        sizeof(args), &args,
+        4,
+        (const struct gptoss_metal_buffer *[]) {input_buffer, weight_buffer, output_buffer, control_buffer},
+        (const size_t[]) {input_offset, weight_offset, output_offset, control_offset},
+        /*threadgroup_buffer_size=*/0);
+}
+
+enum gptoss_status gptoss_metal_command_buffer_encode_launch_f32_bf16w_matmul(
+    const struct gptoss_metal_command_buffer* command_buffer,
+    const struct gptoss_metal_function* f32_bf16w_matmul_fn,
+    size_t threadgroup_size,
+    const struct gptoss_metal_buffer* input_buffer,
+    size_t input_offset,
+    const struct gptoss_metal_buffer* weight_buffer,
+    size_t weight_offset,
+    const struct gptoss_metal_buffer* bias_buffer,
+    size_t bias_offset,
+    const struct gptoss_metal_buffer* output_buffer,
+    size_t output_offset,
+    const struct gptoss_metal_buffer* control_buffer,
+    size_t control_offset,
+    uint32_t num_tokens,
+    uint32_t num_cols,
+    uint32_t num_rows)
+{
+    if (command_buffer->object == NULL || f32_bf16w_matmul_fn->pipeline_state_object == NULL) {
+        GPTOSS_LOG_ERROR("failed to encode f32_bf16w_matmul kernel launch: invalid command buffer or pipeline state object");
+        return gptoss_status_invalid_state;
+    }
+
+    if (threadgroup_size == 0) {
+        threadgroup_size = f32_bf16w_matmul_fn->simdgroup_threads;
+    } else if (threadgroup_size > f32_bf16w_matmul_fn->max_threadgroup_threads) {
+        GPTOSS_LOG_ERROR("failed to encode f32_bf16w_matmul kernel launch: threadgroup size (%zu) exceeds supported maximum (%zu)",
+            threadgroup_size, f32_bf16w_matmul_fn->max_threadgroup_threads);
+        return gptoss_status_invalid_argument;
+    }
+
+    if (num_cols % 4 != 0) {
+        GPTOSS_LOG_ERROR("failed to encode f32_bf16w_matmul kernel launch: number of columns (%" PRIu32 ") is not divisible by 4",
+            num_cols);
+        return gptoss_status_invalid_argument;
+    }
+    const size_t num_simdgroups = threadgroup_size / f32_bf16w_matmul_fn->simdgroup_threads;
+    if (num_rows % num_simdgroups != 0) {
+        GPTOSS_LOG_ERROR("failed to encode f32_bf16w_matmul kernel launch: number of rows (%" PRIu32 ") is not divisible by the number of simdgroups (%zu)",
+            num_rows, num_simdgroups);
+        return gptoss_status_invalid_argument;
+    }
+
+    const struct gptoss_matmul_args args = {
+        .num_column_vecs = num_cols / 4,
+        .num_rows = num_rows,
+        .add = 0,
+    };
+
+    return gptoss_metal_command_buffer_encode_launch_kernel(
+        command_buffer, f32_bf16w_matmul_fn,
+        threadgroup_size, 1, 1,
+        num_rows / num_simdgroups, num_tokens, 1,
+        sizeof(args), &args,
+        5,
+        (const struct gptoss_metal_buffer *[]) {input_buffer, weight_buffer, bias_buffer, output_buffer, control_buffer},
+        (const size_t[]) {input_offset, weight_offset, bias_offset, output_offset, control_offset},
+        /*threadgroup_buffer_size=*/0);
+}
+
+enum gptoss_status gptoss_metal_command_buffer_encode_launch_f32_bf16w_matmul_qkv(
+    const struct gptoss_metal_command_buffer* command_buffer,
+    const struct gptoss_metal_function* f32_bf16w_matmul_qkv_fn,
+    size_t threadgroup_size,
+    const struct gptoss_metal_buffer* input_buffer,
+    size_t input_offset,
+    const struct gptoss_metal_buffer* weight_buffer,
+    size_t weight_offset,
+    const struct gptoss_metal_buffer* bias_buffer,
+    size_t bias_offset,
+    const struct gptoss_metal_buffer* output_buffer,
+    size_t output_offset,
+    const struct gptoss_metal_buffer* kv_buffer,
+    size_t kv_offset,
+    const struct gptoss_metal_buffer* control_buffer,
+    size_t control_offset,
+    uint32_t num_tokens,
+    uint32_t num_cols,
+    uint32_t num_q_heads,
+    uint32_t num_kv_heads,
+    uint32_t attn_head_dim,
+    uint32_t token_offset,
+    uint32_t max_tokens,
+    float rope_base,
+    float interpolation_scale,
+    float yarn_offset,
+    float yarn_scale,
+    float yarn_multiplier)
+{
+    if (command_buffer->object == NULL || f32_bf16w_matmul_qkv_fn->pipeline_state_object == NULL) {
+        GPTOSS_LOG_ERROR("failed to encode f32_bf16w_matmul_qkv kernel launch: invalid command buffer or pipeline state object");
+        return gptoss_status_invalid_state;
+    }
+
+    if (threadgroup_size == 0) {
+        threadgroup_size = f32_bf16w_matmul_qkv_fn->simdgroup_threads;
+    } else if (threadgroup_size > f32_bf16w_matmul_qkv_fn->max_threadgroup_threads) {
+        GPTOSS_LOG_ERROR("failed to encode f32_bf16w_matmul_qkv kernel launch: threadgroup size (%zu) exceeds supported maximum (%zu)",
+            threadgroup_size, f32_bf16w_matmul_qkv_fn->max_threadgroup_threads);
+        return gptoss_status_invalid_argument;
+    }
+
+    if (num_cols % 4 != 0) {
+        GPTOSS_LOG_ERROR("failed to encode f32_bf16w_matmul_qkv kernel launch: number of columns (%" PRIu32 ") is not divisible by 4",
+            num_cols);
+        return gptoss_status_invalid_argument;
+    }
+
+    if (num_q_heads != 64) {
+        GPTOSS_LOG_ERROR("failed to encode f32_bf16w_matmul_qkv kernel launch: number of Q heads (%" PRIu32 ") must be 64",
+            num_q_heads);
+        return gptoss_status_invalid_argument;
+    }
+    if (num_kv_heads != 8) {
+        GPTOSS_LOG_ERROR("failed to encode f32_bf16w_matmul_qkv kernel launch: number of KV heads (%" PRIu32 ") must be 8",
+            num_kv_heads);
+        return gptoss_status_invalid_argument;
+    }
+    if (attn_head_dim != 64) {
+        GPTOSS_LOG_ERROR("failed to encode f32_bf16w_matmul_qkv kernel launch: attention head dimension (%" PRIu32 ") must be 64",
+            attn_head_dim);
+        return gptoss_status_invalid_argument;
+    }
+
+    const size_t num_simdgroups = threadgroup_size / f32_bf16w_matmul_qkv_fn->simdgroup_threads;
+    const uint32_t num_rows = (num_q_heads + 2 * num_kv_heads) * attn_head_dim;
+    if (num_rows % num_simdgroups != 0) {
+        GPTOSS_LOG_ERROR("failed to encode f32_bf16w_matmul_qkv kernel launch: number of rows (%" PRIu32 ") is not divisible by the number of simdgroups (%zu)",
+            num_rows, num_simdgroups);
+        return gptoss_status_invalid_argument;
+    }
+
+    const struct gptoss_qkv_args args = {
+        .num_column_vecs = num_cols / 4,
+        .num_rows = num_rows,
+        .token_offset = token_offset,
+        .freq_scale = -logf(rope_base) / (float) (int32_t) (attn_head_dim / 2),
+        .interpolation_scale = interpolation_scale,
+        .yarn_offset = yarn_offset,
+        .yarn_scale = yarn_scale,
+        .yarn_multiplier = yarn_multiplier,
+        .max_tokens = max_tokens,
+    };
+
+    return gptoss_metal_command_buffer_encode_launch_kernel(
+        command_buffer, f32_bf16w_matmul_qkv_fn,
+        threadgroup_size, 1, 1,
+        num_rows / num_simdgroups, num_tokens, 1,
+        sizeof(args), &args,
+        6,
+        (const struct gptoss_metal_buffer *[]) {input_buffer, weight_buffer, bias_buffer, output_buffer, kv_buffer, control_buffer},
+        (const size_t[]) {input_offset, weight_offset, bias_offset, output_offset, kv_offset, control_offset},
+        /*threadgroup_buffer_size=*/num_simdgroups * sizeof(float));
+}
+
+enum gptoss_status gptoss_metal_command_buffer_encode_launch_f32_bf16w_matmul_add(
+    const struct gptoss_metal_command_buffer* command_buffer,
+    const struct gptoss_metal_function* f32_bf16w_matmul_fn,
+    size_t threadgroup_size,
+    const struct gptoss_metal_buffer* input_buffer,
+    size_t input_offset,
+    const struct gptoss_metal_buffer* weight_buffer,
+    size_t weight_offset,
+    const struct gptoss_metal_buffer* bias_buffer,
+    size_t bias_offset,
+    const struct gptoss_metal_buffer* output_buffer,
+    size_t output_offset,
+    const struct gptoss_metal_buffer* control_buffer,
+    size_t control_offset,
+    uint32_t num_tokens,
+    uint32_t num_cols,
+    uint32_t num_rows)
+{
+    if (command_buffer->object == NULL || f32_bf16w_matmul_fn->pipeline_state_object == NULL) {
+        GPTOSS_LOG_ERROR("failed to encode f32_bf16w_matmul_add kernel launch: invalid command buffer or pipeline state object");
+        return gptoss_status_invalid_state;
+    }
+
+    if (threadgroup_size == 0) {
+        threadgroup_size = f32_bf16w_matmul_fn->simdgroup_threads;
+    } else if (threadgroup_size > f32_bf16w_matmul_fn->max_threadgroup_threads) {
+        GPTOSS_LOG_ERROR("failed to encode f32_bf16w_matmul_add kernel launch: threadgroup size (%zu) exceeds supported maximum (%zu)",
+            threadgroup_size, f32_bf16w_matmul_fn->max_threadgroup_threads);
+        return gptoss_status_invalid_argument;
+    }
+
+    if (num_cols % 4 != 0) {
+        GPTOSS_LOG_ERROR("failed to encode f32_bf16w_matmul_add kernel launch: number of columns (%" PRIu32 ") is not divisible by 4",
+            num_cols);
+        return gptoss_status_invalid_argument;
+    }
+    const size_t num_simdgroups = threadgroup_size / f32_bf16w_matmul_fn->simdgroup_threads;
+    if (num_rows % num_simdgroups != 0) {
+        GPTOSS_LOG_ERROR("failed to encode f32_bf16w_matmul_add kernel launch: number of rows (%" PRIu32 ") is not divisible by the number of simdgroups (%zu)",
+            num_rows, num_simdgroups);
+        return gptoss_status_invalid_argument;
+    }
+
+    const struct gptoss_matmul_args args = {
+        .num_column_vecs = num_cols / 4,
+        .num_rows = num_rows,
+        .add = 1,
+    };
+
+    return gptoss_metal_command_buffer_encode_launch_kernel(
+        command_buffer, f32_bf16w_matmul_fn,
+        threadgroup_size, 1, 1,
+        num_rows / num_simdgroups, num_tokens, 1,
+        sizeof(args), &args,
+        5,
+        (const struct gptoss_metal_buffer *[]) {input_buffer, weight_buffer, bias_buffer, output_buffer, control_buffer},
+        (const size_t[]) {input_offset, weight_offset, bias_offset, output_offset, control_offset},
+        /*threadgroup_buffer_size=*/0);
+}
+
+enum gptoss_status _gptoss_metal_command_buffer_encode_launch_f32_bf16w_dense_matmul_impl(
+    const struct gptoss_metal_command_buffer* command_buffer,
+    const struct gptoss_metal_function* f32_bf16w_dense_matmul_fn,
+    const struct gptoss_metal_buffer* input_buffer,
+    size_t input_offset,
+    const struct gptoss_metal_buffer* weight_buffer,
+    size_t weight_offset,
+    const struct gptoss_metal_buffer* bias_buffer,
+    size_t bias_offset,
+    const struct gptoss_metal_buffer* output_buffer,
+    size_t output_offset, 
+    const struct gptoss_metal_buffer* control_buffer,
+    size_t control_offset,
+    uint32_t num_tokens,
+    uint32_t num_cols,
+    uint32_t num_rows,
+    uint32_t Bm,
+    uint32_t Bn,
+    uint32_t Bk,
+    uint32_t Sg_Bm,
+    uint32_t Sg_Bn)
+{
+
+    if (command_buffer->object == NULL || f32_bf16w_dense_matmul_fn->pipeline_state_object == NULL) {
+        GPTOSS_LOG_ERROR("failed to encode f32_bf16w_dense_matmul kernel launch: invalid command buffer or pipeline state object");
+        return gptoss_status_invalid_state;
+    }
+
+    if (num_cols % 8 != 0) {
+        GPTOSS_LOG_ERROR("failed to encode f32_bf16w_dense_matmul kernel launch: number of columns (%" PRIu32 ") is not divisible by 8",
+                         num_cols);
+        return gptoss_status_invalid_argument;
+    }
+    if (num_rows % 8 != 0) {
+        GPTOSS_LOG_ERROR("failed to encode f32_bf16w_dense_matmul kernel launch: number of rows (%" PRIu32 ") is not divisible by 8",
+                         num_rows);
+        return gptoss_status_invalid_argument;
+    }
+    if (num_tokens % 8 != 0) {
+        GPTOSS_LOG_ERROR("failed to encode f32_bf16w_dense_matmul kernel launch: number of tokens (%" PRIu32 ") is not divisible by 8",
+                         num_tokens);
+        return gptoss_status_invalid_argument;
+    }
+
+    const struct gptoss_dense_matmul_args args = {
+        .m = num_tokens,
+        .n = num_rows,
+        .k = num_cols,
+    };
+    const size_t threads_per_simdgroup = f32_bf16w_dense_matmul_fn->simdgroup_threads;
+    const uint32_t m = args.m;
+    const uint32_t n = args.n;
+    const uint32_t k = args.k;
+    if (Bm % Sg_Bm != 0) {
+        GPTOSS_LOG_ERROR("failed to encode f32_bf16w_dense_matmul kernel launch: Bm (%" PRIu32 ") is not divisible by Sg_Bm (%" PRIu32 ")",
+                         Bm, Sg_Bm);
+        return gptoss_status_invalid_argument;
+    }
+    if (Bn % Sg_Bn != 0) {
+        GPTOSS_LOG_ERROR("failed to encode f32_bf16w_dense_matmul kernel launch: Bn (%" PRIu32 ") is not divisible by Sg_Bn (%" PRIu32 ")",
+                         Bn, Sg_Bn);
+        return gptoss_status_invalid_argument;
+    }
+    const size_t threadgroup_size_x = (Bm / Sg_Bm) * (Bn / Sg_Bn) * threads_per_simdgroup;
+    const size_t threadgroup_size_y = 1;
+    const size_t threadgroup_size_z = 1;
+    const size_t total_threadgroup_size = threadgroup_size_x * threadgroup_size_y * threadgroup_size_z;
+    if (total_threadgroup_size > f32_bf16w_dense_matmul_fn->max_threadgroup_threads) {
+        GPTOSS_LOG_ERROR("failed to encode f32_bf16w_dense_matmul kernel launch: total threadgroup size (%zu) exceeds supported maximum (%zu)",
+                         total_threadgroup_size, f32_bf16w_dense_matmul_fn->max_threadgroup_threads);
+        return gptoss_status_invalid_argument;
+    }
+    if (m % Bm != 0) {
+        GPTOSS_LOG_ERROR("failed to encode f32_bf16w_dense_matmul kernel launch: m (%" PRIu32 ") is not divisible by Bm (%" PRIu32 ")",
+                         m, Bm);
+        return gptoss_status_invalid_argument;
+    }
+    if (n % Bn != 0) {
+        GPTOSS_LOG_ERROR("failed to encode f32_bf16w_dense_matmul kernel launch: n (%" PRIu32 ") is not divisible by Bn (%" PRIu32 ")",
+                         n, Bn);
+        return gptoss_status_invalid_argument;
+    }
+    if (k % Bk != 0) {
+        GPTOSS_LOG_ERROR("failed to encode f32_bf16w_dense_matmul kernel launch: k (%" PRIu32 ") is not divisible by Bk (%" PRIu32 ")",
+                         k, Bk);
+        return gptoss_status_invalid_argument;
+    }
+    const size_t grid_x = n / Bn;
+    const size_t grid_y = m / Bm;
+    const size_t grid_z = 1;
+
+    return gptoss_metal_command_buffer_encode_launch_kernel(
+        command_buffer, f32_bf16w_dense_matmul_fn,
+        threadgroup_size_x, threadgroup_size_y, threadgroup_size_z,
+        grid_x, grid_y, grid_z,
+        sizeof(args), &args,
+        5,
+        (const struct gptoss_metal_buffer *[]){input_buffer, weight_buffer, bias_buffer, output_buffer, control_buffer},
+        (const size_t[]){input_offset, weight_offset, bias_offset, output_offset, control_offset},
+        /*threadgroup_buffer_size=*/0);
+    return gptoss_status_success;
+}
+
+enum gptoss_status gptoss_metal_command_buffer_encode_launch_f32_bf16w_dense_matmul_qkv(
+    const struct gptoss_metal_command_buffer* command_buffer,
+    const struct gptoss_metal_function* f32_bf16w_dense_matmul_fn,
+    const struct gptoss_metal_buffer* input_buffer,
+    size_t input_offset,
+    const struct gptoss_metal_buffer* weight_buffer,
+    size_t weight_offset,
+    const struct gptoss_metal_buffer* bias_buffer,
+    size_t bias_offset,
+    const struct gptoss_metal_buffer* output_buffer,
+    size_t output_offset,
+    const struct gptoss_metal_buffer* control_buffer,
+    size_t control_offset,
+    uint32_t num_tokens,
+    uint32_t num_cols,
+    uint32_t num_rows)
+{
+    return _gptoss_metal_command_buffer_encode_launch_f32_bf16w_dense_matmul_impl(
+        command_buffer, f32_bf16w_dense_matmul_fn, input_buffer, input_offset, 
+        weight_buffer, weight_offset, bias_buffer, bias_offset, output_buffer,
+        output_offset, control_buffer, control_offset, num_tokens, num_cols, num_rows, QKV_Bm, QKV_Bn, QKV_Bk,
+        QKV_Sg_Bm, QKV_Sg_Bn);
+}
+
+enum gptoss_status gptoss_metal_command_buffer_encode_launch_f32_bf16w_dense_matmul_attn_output(
+    const struct gptoss_metal_command_buffer* command_buffer,
+    const struct gptoss_metal_function* f32_bf16w_dense_matmul_fn,
+    const struct gptoss_metal_buffer* input_buffer,
+    size_t input_offset,
+    const struct gptoss_metal_buffer* weight_buffer,
+    size_t weight_offset,
+    const struct gptoss_metal_buffer* bias_buffer,
+    size_t bias_offset,
+    const struct gptoss_metal_buffer* output_buffer,
+    size_t output_offset,
+    const struct gptoss_metal_buffer* control_buffer,
+    size_t control_offset,
+    uint32_t num_tokens,
+    uint32_t num_cols,
+    uint32_t num_rows)
+{
+    return _gptoss_metal_command_buffer_encode_launch_f32_bf16w_dense_matmul_impl(
+        command_buffer, f32_bf16w_dense_matmul_fn, input_buffer, input_offset,
+        weight_buffer, weight_offset, bias_buffer, bias_offset, output_buffer,
+        output_offset, control_buffer, control_offset, num_tokens, num_cols, num_rows, ATTN_OUTPUT_Bm,
+        ATTN_OUTPUT_Bn, ATTN_OUTPUT_Bk, ATTN_OUTPUT_Sg_Bm, ATTN_OUTPUT_Sg_Bn);
+}
+
+enum gptoss_status gptoss_metal_command_buffer_encode_launch_f32_bf16w_dense_matmul_mlp_gate(
+    const struct gptoss_metal_command_buffer* command_buffer,
+    const struct gptoss_metal_function* f32_bf16w_dense_matmul_fn,
+    const struct gptoss_metal_buffer* input_buffer,
+    size_t input_offset,
+    const struct gptoss_metal_buffer* weight_buffer,
+    size_t weight_offset,
+    const struct gptoss_metal_buffer* bias_buffer,
+    size_t bias_offset,
+    const struct gptoss_metal_buffer* output_buffer,
+    size_t output_offset,
+    const struct gptoss_metal_buffer* control_buffer,
+    size_t control_offset,
+    uint32_t num_tokens,
+    uint32_t num_cols,
+    uint32_t num_rows)
+{
+    return _gptoss_metal_command_buffer_encode_launch_f32_bf16w_dense_matmul_impl(
+        command_buffer, f32_bf16w_dense_matmul_fn, input_buffer, input_offset,
+        weight_buffer, weight_offset, bias_buffer, bias_offset, output_buffer,
+        output_offset, control_buffer, control_offset, num_tokens, num_cols,
+        num_rows, MLP_GATE_Bm, MLP_GATE_Bn, MLP_GATE_Bk, MLP_GATE_Sg_Bm,
+        MLP_GATE_Sg_Bn);
+}
+
+enum gptoss_status gptoss_metal_command_buffer_encode_launch_f32_bf16w_unembedding(
+    const struct gptoss_metal_command_buffer* command_buffer,
+    const struct gptoss_metal_function* f32_bf16w_unembedding_fn,
+    size_t threadgroup_size,
+    size_t max_threadgroups,
+    const struct gptoss_metal_buffer* input_buffer,
+    size_t input_offset,
+    const struct gptoss_metal_buffer* weight_buffer,
+    size_t weight_offset,
+    const struct gptoss_metal_buffer* output_buffer,
+    size_t output_offset,
+    const struct gptoss_metal_buffer* argmax_buffer,
+    size_t argmax_offset,
+    const struct gptoss_metal_buffer* control_buffer,
+    size_t control_offset,
+    uint32_t num_tokens,
+    uint32_t num_cols,
+    uint32_t num_rows)
+{
+    if (command_buffer->object == NULL || f32_bf16w_unembedding_fn->pipeline_state_object == NULL) {
+        GPTOSS_LOG_ERROR("failed to encode f32_bf16w_unembedding kernel launch: invalid command buffer or pipeline state object");
+        return gptoss_status_invalid_state;
+    }
+
+    if (threadgroup_size == 0) {
+        threadgroup_size = f32_bf16w_unembedding_fn->simdgroup_threads;
+    } else if (threadgroup_size > f32_bf16w_unembedding_fn->max_threadgroup_threads) {
+        GPTOSS_LOG_ERROR("failed to encode f32_bf16w_unembedding kernel launch: threadgroup size (%zu) exceeds supported maximum (%zu)",
+            threadgroup_size, f32_bf16w_unembedding_fn->max_threadgroup_threads);
+        return gptoss_status_invalid_argument;
+    }
+
+    if (num_cols % 4 != 0) {
+        GPTOSS_LOG_ERROR("failed to encode f32_bf16w_unembedding kernel launch: number of columns (%" PRIu32 ") is not divisible by 4",
+            num_cols);
+        return gptoss_status_invalid_argument;
+    }
+
+    const size_t num_simdgroups = threadgroup_size / f32_bf16w_unembedding_fn->simdgroup_threads;
+    const size_t num_rows_per_threadgroup = math_ceil_div(num_rows, max_threadgroups * num_simdgroups) * num_simdgroups;
+    const size_t num_threadgroups = math_min(max_threadgroups, math_ceil_div(num_rows, num_rows_per_threadgroup));
+    const struct gptoss_unembedding_args args = {
+        .num_column_vecs = num_cols / 4,
+        .num_rows_per_threadgroup = num_rows_per_threadgroup,
+        .num_rows = num_rows,
+    };
+
+    return gptoss_metal_command_buffer_encode_launch_kernel(
+        command_buffer, f32_bf16w_unembedding_fn,
+        threadgroup_size, 1, 1,
+        num_threadgroups, num_tokens, 1,
+        sizeof(args), &args,
+        5,
+        (const struct gptoss_metal_buffer *[]) {input_buffer, weight_buffer, output_buffer, argmax_buffer, control_buffer},
+        (const size_t[]) {input_offset, weight_offset, output_offset, argmax_offset, control_offset},
+        /*threadgroup_buffer_size=*/0);
+}
+
+enum gptoss_status gptoss_metal_command_buffer_encode_launch_f32_mf4w_moe_matmul_swiglu(
+    const struct gptoss_metal_command_buffer* command_buffer,
+    const struct gptoss_metal_function* f32_mf4w_moe_matmul_swiglu_fn,
+    size_t threadgroup_size,
+    const struct gptoss_metal_buffer* input_buffer,
+    size_t input_offset,
+    const struct gptoss_metal_buffer* expert_buffer,
+    size_t expert_offset,
+    const struct gptoss_metal_buffer* weight_block_buffer,
+    size_t weight_block_offset,
+    const struct gptoss_metal_buffer* weight_scale_buffer,
+    size_t weight_scale_offset,
+    const struct gptoss_metal_buffer* bias_buffer,
+    size_t bias_offset,
+    const struct gptoss_metal_buffer* output_buffer,
+    size_t output_offset,
+    const struct gptoss_metal_buffer* control_buffer,
+    size_t control_offset,
+    float swiglu_limit,
+    uint32_t expert_stride,
+    uint32_t num_tokens,
+    uint32_t num_active_experts,
+    uint32_t num_cols,
+    uint32_t num_rows)
+{
+    if (command_buffer->object == NULL || f32_mf4w_moe_matmul_swiglu_fn->pipeline_state_object == NULL) {
+        GPTOSS_LOG_ERROR("failed to encode f32_mf4w_moe_matmul_swiglu kernel launch: invalid command buffer or pipeline state object");
+        return gptoss_status_invalid_state;
+    }
+
+    if (threadgroup_size == 0) {
+        threadgroup_size = 2 * f32_mf4w_moe_matmul_swiglu_fn->simdgroup_threads;
+    } else if (threadgroup_size > f32_mf4w_moe_matmul_swiglu_fn->max_threadgroup_threads) {
+        GPTOSS_LOG_ERROR("failed to encode f32_mf4w_moe_matmul_swiglu kernel launch: threadgroup size (%zu) exceeds supported maximum (%zu)",
+            threadgroup_size, f32_mf4w_moe_matmul_swiglu_fn->max_threadgroup_threads);
+        return gptoss_status_invalid_argument;
+    } else if (threadgroup_size % (2 * f32_mf4w_moe_matmul_swiglu_fn->simdgroup_threads)) {
+        GPTOSS_LOG_ERROR("failed to encode f32_mf4w_moe_matmul_swiglu kernel launch: threadgroup size (%zu) is not divisible by simdgroup size (%zu) multiplied by 2X",
+            threadgroup_size, f32_mf4w_moe_matmul_swiglu_fn->simdgroup_threads);
+        return gptoss_status_invalid_argument;
+    }
+
+    if (num_cols % 32 != 0) {
+        GPTOSS_LOG_ERROR("failed to encode f32_mf4w_moe_matmul_swiglu kernel launch: number of columns (%" PRIu32 ") is not divisible by 32",
+            num_cols);
+        return gptoss_status_invalid_argument;
+    }
+    const size_t num_simdgroups = threadgroup_size / f32_mf4w_moe_matmul_swiglu_fn->simdgroup_threads;
+    if ((2 * num_rows) % num_simdgroups != 0) {
+        GPTOSS_LOG_ERROR("failed to encode f32_bf16w_matmul_add kernel launch: "
+            "the number of rows (%" PRIu32 ") multiplied by 2X is not divisible by the number of simdgroups (%zu)",
+            num_rows, num_simdgroups);
+        return gptoss_status_invalid_argument;
+    }
+
+    const struct gptoss_moe_matmul_swiglu_args args = {
+        .num_column_vecs = num_cols / 32,
+        .num_rows = num_rows,
+        .num_active_experts = num_active_experts,
+        .weight_expert_stride = expert_stride,
+        .output_expert_stride = num_rows * num_tokens,
+        .swiglu_min = -swiglu_limit,
+        .swiglu_max = swiglu_limit,
+    };
+
+    return gptoss_metal_command_buffer_encode_launch_kernel(
+        command_buffer, f32_mf4w_moe_matmul_swiglu_fn,
+        threadgroup_size, 1, 1,
+        (2 * num_rows) / num_simdgroups, num_tokens, num_active_experts,
+        sizeof(args), &args,
+        7,
+        (const struct gptoss_metal_buffer *[]) {input_buffer, expert_buffer, weight_block_buffer, weight_scale_buffer, bias_buffer, output_buffer, control_buffer},
+        (const size_t[]) {input_offset, expert_offset, weight_block_offset, weight_scale_offset, bias_offset, output_offset, control_offset},
+        /*threadgroup_buffer_size=*/0);
+}
+
+enum gptoss_status gptoss_metal_command_buffer_encode_launch_f32_mf4w_moe_matmul(
+    const struct gptoss_metal_command_buffer* command_buffer,
+    const struct gptoss_metal_function* f32_mf4w_moe_matmul_fn,
+    size_t threadgroup_size,
+    const struct gptoss_metal_buffer* input_buffer,
+    size_t input_offset,
+    const struct gptoss_metal_buffer* expert_buffer,
+    size_t expert_offset,
+    const struct gptoss_metal_buffer* weight_block_buffer,
+    size_t weight_block_offset,
+    const struct gptoss_metal_buffer* weight_scale_buffer,
+    size_t weight_scale_offset,
+    const struct gptoss_metal_buffer* bias_buffer,
+    size_t bias_offset,
+    const struct gptoss_metal_buffer* output_buffer,
+    size_t output_offset,
+    const struct gptoss_metal_buffer* control_buffer,
+    size_t control_offset,
+    uint32_t expert_stride,
+    uint32_t num_tokens,
+    uint32_t num_active_experts,
+    uint32_t num_cols,
+    uint32_t num_rows)
+{
+    if (command_buffer->object == NULL || f32_mf4w_moe_matmul_fn->pipeline_state_object == NULL) {
+        GPTOSS_LOG_ERROR("failed to encode f32_mf4w_moe_matmul kernel launch: invalid command buffer or pipeline state object");
+        return gptoss_status_invalid_state;
+    }
+
+    if (threadgroup_size == 0) {
+        threadgroup_size = f32_mf4w_moe_matmul_fn->simdgroup_threads;
+    } else if (threadgroup_size > f32_mf4w_moe_matmul_fn->max_threadgroup_threads) {
+        GPTOSS_LOG_ERROR("failed to encode f32_mf4w_moe_matmul kernel launch: threadgroup size (%zu) exceeds supported maximum (%zu)",
+            threadgroup_size, f32_mf4w_moe_matmul_fn->max_threadgroup_threads);
+        return gptoss_status_invalid_argument;
+    } else if (threadgroup_size % f32_mf4w_moe_matmul_fn->simdgroup_threads) {
+        GPTOSS_LOG_ERROR("failed to encode f32_mf4w_moe_matmul kernel launch: threadgroup size (%zu) is not divisible by simdgroup size (%zu)",
+            threadgroup_size, f32_mf4w_moe_matmul_fn->simdgroup_threads);
+        return gptoss_status_invalid_argument;
+    }
+
+    if (num_cols % 32 != 0) {
+        GPTOSS_LOG_ERROR("failed to encode f32_mf4w_moe_matmul kernel launch: number of columns (%" PRIu32 ") is not divisible by 32",
+            num_cols);
+        return gptoss_status_invalid_argument;
+    }
+    const size_t num_simdgroups = threadgroup_size / f32_mf4w_moe_matmul_fn->simdgroup_threads;
+    if (num_rows % num_simdgroups != 0) {
+        GPTOSS_LOG_ERROR("failed to encode f32_mf4w_moe_matmul kernel launch: "
+            "the number of rows (%" PRIu32 ") is not divisible by the number of simdgroups (%zu)",
+            num_rows, num_simdgroups);
+        return gptoss_status_invalid_argument;
+    }
+
+    const struct gptoss_moe_matmul_args args = {
+        .num_column_vecs = num_cols / 32,
+        .num_rows = num_rows,
+        .num_active_experts = num_active_experts,
+        .input_expert_stride = num_tokens * (num_cols / 32),
+        .weight_expert_stride = expert_stride,
+        .output_expert_stride = num_rows * num_tokens,
+    };
+
+    return gptoss_metal_command_buffer_encode_launch_kernel(
+        command_buffer, f32_mf4w_moe_matmul_fn,
+        threadgroup_size, 1, 1,
+        num_rows / num_simdgroups, num_tokens, num_active_experts,
+        sizeof(args), &args,
+        7,
+        (const struct gptoss_metal_buffer *[]) {input_buffer, expert_buffer, weight_block_buffer, weight_scale_buffer, bias_buffer, output_buffer, control_buffer},
+        (const size_t[]) {input_offset, expert_offset, weight_block_offset, weight_scale_offset, bias_offset, output_offset, control_offset},
+        /*threadgroup_buffer_size=*/0);
+}
+
+enum gptoss_status gptoss_metal_command_buffer_encode_launch_f32_rope(
+    const struct gptoss_metal_command_buffer* command_buffer,
+    const struct gptoss_metal_function* f32_rope_fn,
+    size_t threadgroup_size,
+    const struct gptoss_metal_buffer* activations_buffer,
+    size_t activations_offset,
+    const struct gptoss_metal_buffer* control_buffer,
+    size_t control_offset,
+    float rope_base,
+    float interpolation_scale,
+    float yarn_offset,
+    float yarn_scale,
+    float yarn_multiplier,
+    uint32_t num_tokens,
+    uint32_t num_q_heads,
+    uint32_t num_kv_heads,
+    uint32_t attn_head_dim,
+    uint32_t token_offset)
+{
+    if (command_buffer->object == NULL || f32_rope_fn->pipeline_state_object == NULL) {
+        return gptoss_status_invalid_state;
+    }
+
+    if (threadgroup_size == 0) {
+        threadgroup_size = f32_rope_fn->max_threadgroup_threads;
+    } else if (threadgroup_size > f32_rope_fn->max_threadgroup_threads) {
+        return gptoss_status_invalid_argument;
+    }
+
+    const size_t num_simdgroups = threadgroup_size / f32_rope_fn->simdgroup_threads;
+    const uint32_t num_qk_heads = num_q_heads + num_kv_heads;
+    if (num_qk_heads % num_simdgroups != 0) {
+        return gptoss_status_invalid_argument;
+    }
+
+    const struct gptoss_rope_args args = {
+        .token_stride = (num_q_heads + 2 * num_kv_heads) * (attn_head_dim / 2),
+        .token_offset = token_offset,
+        .freq_scale = -logf(rope_base) / (float) (int32_t) (attn_head_dim / 2),
+        .interpolation_scale = interpolation_scale,
+        .yarn_offset = yarn_offset,
+        .yarn_scale = yarn_scale,
+        .yarn_multiplier = yarn_multiplier,
+    };
+
+    return gptoss_metal_command_buffer_encode_launch_kernel(
+        command_buffer, f32_rope_fn,
+        threadgroup_size, 1, 1,
+        num_qk_heads / num_simdgroups, num_tokens, 1,
+        sizeof(args), &args,
+        2,
+        (const struct gptoss_metal_buffer *[]) {activations_buffer, control_buffer},
+        (const size_t[]) {activations_offset, control_offset},
+        /*threadgroup_buffer_size=*/0);
+}
+
+enum gptoss_status gptoss_metal_command_buffer_encode_launch_expert_routing_metadata(
+    const struct gptoss_metal_command_buffer* command_buffer,
+    const struct gptoss_metal_function* expert_routing_metadata_fn,
+    const struct gptoss_metal_buffer* expert_predictions_buffer,
+    size_t expert_predictions_offset,
+    const struct gptoss_metal_buffer* expert_offsets_buffer,
+    size_t expert_offsets_offset,
+    const struct gptoss_metal_buffer* intra_expert_offsets_buffer,
+    size_t intra_expert_offsets_offset,
+    uint32_t num_tokens,
+    uint32_t num_experts)
+{
+    if (command_buffer->object == NULL || expert_routing_metadata_fn->pipeline_state_object == NULL) {
+        return gptoss_status_invalid_state;
+    }
+    
+    const struct gptoss_expert_routing_metadata_args args = {
+        .tokens = num_tokens,
+        .num_experts = num_experts,
+    };
+    const uint32_t threadgroup_size = 256;
+    return gptoss_metal_command_buffer_encode_launch_kernel(
+        command_buffer, expert_routing_metadata_fn,
+        threadgroup_size, 1, 1,
+        /*num_threadgroups_x=*/1, /*num_threadgroups_y=*/1, /*num_threadgroups_z=*/1,
+        sizeof(args), &args,
+        3,
+        (const struct gptoss_metal_buffer *[]) {expert_predictions_buffer, expert_offsets_buffer, intra_expert_offsets_buffer},
+        (const size_t[]) {expert_predictions_offset, expert_offsets_offset, intra_expert_offsets_offset},
+        /*threadgroup_buffer_size=*/0);
+}
+
+enum gptoss_status gptoss_metal_command_buffer_encode_launch_f32_scatter(
+    const struct gptoss_metal_command_buffer* command_buffer,
+    const struct gptoss_metal_function* f32_scatter_fn,
+    const struct gptoss_metal_buffer* input_buffer,
+    size_t input_offset,
+    const struct gptoss_metal_buffer* expert_predictions_buffer,
+    size_t expert_predictions_offset,
+    const struct gptoss_metal_buffer* expert_offsets_buffer,
+    size_t expert_offsets_offset,
+    const struct gptoss_metal_buffer* intra_expert_offsets_buffer,
+    size_t intra_expert_offsets_offset,
+    const struct gptoss_metal_buffer* output_buffer,
+    size_t output_offset,
+    uint32_t num_channels,
+    uint32_t num_tokens,
+    uint32_t num_active_experts)
+{
+    if (command_buffer->object == NULL || f32_scatter_fn->pipeline_state_object == NULL) {
+        return gptoss_status_invalid_state;
+    }
+
+    if (num_channels % 4 != 0) {
+        return gptoss_status_invalid_argument;
+    }
+
+    const size_t num_vecs = num_channels / 4;
+    const size_t tgx = math_min(num_vecs, 64);
+    const size_t tgy = 1;
+    const size_t tgz = 1;
+    const size_t grid_x = math_ceil_div(num_vecs, tgx);
+    const size_t grid_y = num_tokens;
+    const size_t grid_z = 1;
+    const size_t total_threadgroup_size = tgx * tgy * tgz;
+    if (total_threadgroup_size > f32_scatter_fn->max_threadgroup_threads) {
+        return gptoss_status_invalid_argument;
+    }
+    const struct gptoss_scatter_args args = {
+        .tokens = num_tokens,
+        .active_experts_per_token = num_active_experts,
+        .token_stride = num_channels,
+    };
+
+    return gptoss_metal_command_buffer_encode_launch_kernel(
+        command_buffer, f32_scatter_fn,
+        tgx, tgy, tgz,
+        grid_x, grid_y, grid_z,
+        sizeof(args), &args,
+        5,
+        (const struct gptoss_metal_buffer *[]) {input_buffer, expert_predictions_buffer, expert_offsets_buffer, intra_expert_offsets_buffer, output_buffer},
+        (const size_t[]) {input_offset, expert_predictions_offset, expert_offsets_offset, intra_expert_offsets_offset, output_offset},
+        /*threadgroup_buffer_size=*/0);
+}
+
+enum gptoss_status gptoss_metal_command_buffer_encode_launch_f32_gather_and_accumulate_e4(
+    const struct gptoss_metal_command_buffer* command_buffer,
+    const struct gptoss_metal_function* f32_gather_and_accumulate_e4_fn,
+    const struct gptoss_metal_buffer* input_buffer,
+    size_t input_offset,
+    const struct gptoss_metal_buffer* expert_predictions_buffer,
+    size_t expert_predictions_offset,
+    const struct gptoss_metal_buffer* expert_offsets_buffer,
+    size_t expert_offsets_offset,
+    const struct gptoss_metal_buffer* intra_expert_offsets_buffer,
+    size_t intra_expert_offsets_offset,
+    const struct gptoss_metal_buffer* output_buffer,
+    size_t output_offset,
+    uint32_t num_channels,
+    uint32_t num_tokens,
+    uint32_t num_active_experts) 
+{
+        if (command_buffer->object == NULL || f32_gather_and_accumulate_e4_fn->pipeline_state_object == NULL) {
+        return gptoss_status_invalid_state;
+    }
+
+    if (num_channels % 4 != 0) {
+        return gptoss_status_invalid_argument;
+    }
+
+    const size_t num_vecs = num_channels / 4;
+    const size_t tgx = math_min(num_vecs, 64);
+    const size_t tgy = 1;
+    const size_t tgz = 1;
+    const size_t grid_x = math_ceil_div(num_vecs, tgx);
+    const size_t grid_y = num_tokens;
+    const size_t grid_z = 1;
+    const size_t total_threadgroup_size = tgx * tgy * tgz;
+    if (total_threadgroup_size > f32_gather_and_accumulate_e4_fn->max_threadgroup_threads) {
+        return gptoss_status_invalid_argument;
+    }
+    const struct gptoss_gather_args args = {
+        .tokens = num_tokens,
+        .active_experts_per_token = num_active_experts,
+        .token_stride = num_channels,
+    };
+    
+    return gptoss_metal_command_buffer_encode_launch_kernel(
+        command_buffer, f32_gather_and_accumulate_e4_fn,
+        tgx, tgy, tgz,
+        grid_x, grid_y, grid_z,
+        sizeof(args), &args,
+        5,
+        (const struct gptoss_metal_buffer *[]) {input_buffer, expert_predictions_buffer, expert_offsets_buffer, intra_expert_offsets_buffer, output_buffer},
+        (const size_t[]) {input_offset, expert_predictions_offset, expert_offsets_offset, intra_expert_offsets_offset, output_offset},
+        /*threadgroup_buffer_size=*/0);
+}
+
+enum gptoss_status gptoss_metal_command_buffer_encode_launch_f32_mf4w_moe_dense_matmul_swiglu(
+    const struct gptoss_metal_command_buffer* command_buffer,
+    const struct gptoss_metal_function* f32_mf4w_moe_dense_matmul_swiglu_fn,
+    const struct gptoss_metal_buffer* expert_offsets_buffer,
+    size_t expert_offsets_offset,
+    const struct gptoss_metal_buffer* input_buffer,
+    size_t input_offset,
+    const struct gptoss_metal_buffer* weight_block_buffer,
+    size_t weight_block_offset,
+    const struct gptoss_metal_buffer* weight_scale_buffer,
+    size_t weight_scale_offset,
+    const struct gptoss_metal_buffer* bias_buffer,
+    size_t bias_offset,
+    const struct gptoss_metal_buffer* output_buffer,
+    size_t output_offset,
+    float swiglu_limit,
+    uint32_t expert_stride_bytes,
+    uint32_t num_tokens,
+    uint32_t num_experts,
+    uint32_t num_cols,
+    uint32_t num_rows)
+{
+    if (command_buffer->object == NULL || f32_mf4w_moe_dense_matmul_swiglu_fn->pipeline_state_object == NULL) {
+        return gptoss_status_invalid_state;
+    }
+
+    if (num_cols % 32 != 0) {
+        GPTOSS_LOG_ERROR("failed to encode f32_mf4w_moe_dense_matmul_swiglu kernel launch: number of columns (%" PRIu32 ") is not divisible by 32",
+            num_cols);
+        return gptoss_status_invalid_argument;
+    }
+
+    const struct gptoss_moe_dense_matmul_swiglu_args args = {
+        .n = num_rows,
+        .k = num_cols,
+        .weight_blocks_expert_stride_bytes = expert_stride_bytes,
+        .weight_scales_expert_stride_bytes = expert_stride_bytes,
+        .bias_expert_stride_bytes = expert_stride_bytes,
+        .swiglu_min = -swiglu_limit,
+        .swiglu_max = swiglu_limit,
+    };
+    const size_t threads_per_simdgroup = f32_mf4w_moe_dense_matmul_swiglu_fn->simdgroup_threads;
+    const uint32_t m = num_tokens;
+    const uint32_t n = args.n;
+    const uint32_t k = args.k;
+    const uint32_t Bm = MOE_DENSE_MATMUL_SWIGLU_Bm;
+    const uint32_t Bn = MOE_DENSE_MATMUL_SWIGLU_Bn;
+    const uint32_t Bk = MOE_DENSE_MATMUL_SWIGLU_Bk;
+    const uint32_t Sg_Bm = MOE_DENSE_MATMUL_SWIGLU_Sg_Bm;
+    const uint32_t Sg_Bn = MOE_DENSE_MATMUL_SWIGLU_Sg_Bn;
+    if (Bm % Sg_Bm != 0) {
+        GPTOSS_LOG_ERROR("failed to encode f32_mf4w_moe_dense_matmul_swiglu kernel launch: Bm (%" PRIu32 ") is not divisible by Sg_Bm (%" PRIu32 ")",
+            Bm, Sg_Bm);
+        return gptoss_status_invalid_argument;
+    }
+    if (Bn % Sg_Bn != 0) {
+        GPTOSS_LOG_ERROR("failed to encode f32_mf4w_moe_dense_matmul_swiglu kernel launch: Bn (%" PRIu32 ") is not divisible by Sg_Bn (%" PRIu32 ")",
+            Bn, Sg_Bn);
+        return gptoss_status_invalid_argument;
+    }
+
+    const size_t threadgroup_size_x = (Bm / Sg_Bm) * (Bn / Sg_Bn) * threads_per_simdgroup;
+    const size_t threadgroup_size_y = 1;
+    const size_t threadgroup_size_z = 1;
+    const size_t total_threadgroup_size = threadgroup_size_x * threadgroup_size_y * threadgroup_size_z;
+    if (total_threadgroup_size > f32_mf4w_moe_dense_matmul_swiglu_fn->max_threadgroup_threads) {
+        GPTOSS_LOG_ERROR("failed to encode f32_mf4w_moe_dense_matmul_swiglu kernel launch: total threadgroup size (%zu) exceeds supported maximum (%zu)",
+            total_threadgroup_size, f32_mf4w_moe_dense_matmul_swiglu_fn->max_threadgroup_threads);
+        return gptoss_status_invalid_argument;
+    }
+    if (n % Bn != 0) {
+        GPTOSS_LOG_ERROR("failed to encode f32_mf4w_moe_dense_matmul_swiglu kernel launch: n (%" PRIu32 ") is not divisible by Bn (%" PRIu32 ")",
+            n, Bn);
+        return gptoss_status_invalid_argument;
+    }
+    if (k % Bk != 0) {
+        GPTOSS_LOG_ERROR("failed to encode f32_mf4w_moe_dense_matmul_swiglu kernel launch: k (%" PRIu32 ") is not divisible by Bk (%" PRIu32 ")",
+            k, Bk);
+        return gptoss_status_invalid_argument;
+    }
+    const size_t grid_x = n / Bn;
+    const size_t grid_y = math_ceil_div(m, Bm);
+    const size_t grid_z = num_experts;
+
+    return gptoss_metal_command_buffer_encode_launch_kernel(
+        command_buffer, f32_mf4w_moe_dense_matmul_swiglu_fn,
+        threadgroup_size_x, threadgroup_size_y, threadgroup_size_z,
+        grid_x, grid_y, grid_z,
+        sizeof(args), &args,
+        6,
+        (const struct gptoss_metal_buffer *[]) {expert_offsets_buffer, input_buffer, weight_block_buffer, weight_scale_buffer, bias_buffer, output_buffer},
+        (const size_t[]) {expert_offsets_offset, input_offset, weight_block_offset, weight_scale_offset, bias_offset, output_offset},
+        /*threadgroup_buffer_size=*/0);
+
+    }
+
+enum gptoss_status gptoss_metal_command_buffer_encode_launch_f32_mf4w_moe_dense_matmul(
+    const struct gptoss_metal_command_buffer* command_buffer,
+    const struct gptoss_metal_function* f32_mf4w_moe_dense_matmul_fn,
+    const struct gptoss_metal_buffer* expert_offsets_buffer,
+    size_t expert_offsets_offset,
+    const struct gptoss_metal_buffer* input_buffer,
+    size_t input_offset,
+    const struct gptoss_metal_buffer* weight_block_buffer,
+    size_t weight_block_offset,
+    const struct gptoss_metal_buffer* weight_scale_buffer,
+    size_t weight_scale_offset,
+    const struct gptoss_metal_buffer* bias_buffer,
+    size_t bias_offset,
+    const struct gptoss_metal_buffer* output_buffer,
+    size_t output_offset,
+    uint32_t expert_stride_bytes,
+    uint32_t num_tokens,
+    uint32_t num_experts,
+    uint32_t num_cols,
+    uint32_t num_rows)
+{
+    if (command_buffer->object == NULL || f32_mf4w_moe_dense_matmul_fn->pipeline_state_object == NULL) {
+        return gptoss_status_invalid_state;
+    }
+
+    if (num_cols % 32 != 0) {
+        GPTOSS_LOG_ERROR("failed to encode f32_mf4w_moe_dense_matmul kernel launch: number of columns (%" PRIu32 ") is not divisible by 32",
+            num_cols);
+        return gptoss_status_invalid_argument;
+    }
+    const struct gptoss_moe_dense_matmul_args args = {
+        .k = num_cols,
+        .n = num_rows,
+        .weight_blocks_expert_stride_bytes = expert_stride_bytes,
+        .weight_scales_expert_stride_bytes = expert_stride_bytes,
+        .bias_expert_stride_bytes = expert_stride_bytes,
+    };
+
+    const size_t threads_per_simdgroup = f32_mf4w_moe_dense_matmul_fn->simdgroup_threads;
+    const uint32_t m = num_tokens;
+    const uint32_t n = args.n;
+    const uint32_t k = args.k;
+    const uint32_t Bm = MOE_DENSE_MATMUL_Bm;
+    const uint32_t Bn = MOE_DENSE_MATMUL_Bn;
+    const uint32_t Bk = MOE_DENSE_MATMUL_Bk;
+    const uint32_t Sg_Bm = MOE_DENSE_MATMUL_Sg_Bm;
+    const uint32_t Sg_Bn = MOE_DENSE_MATMUL_Sg_Bn;
+    if (Bm % Sg_Bm != 0) {
+        GPTOSS_LOG_ERROR("failed to encode f32_mf4w_moe_dense_matmul kernel launch: Bm (%" PRIu32 ") is not divisible by Sg_Bm (%" PRIu32 ")",
+            Bm, Sg_Bm);
+        return gptoss_status_invalid_argument;
+    }
+    if (Bn % Sg_Bn != 0) {
+        GPTOSS_LOG_ERROR("failed to encode f32_mf4w_moe_dense_matmul kernel launch: Bn (%" PRIu32 ") is not divisible by Sg_Bn (%" PRIu32 ")",
+            Bn, Sg_Bn);
+        return gptoss_status_invalid_argument;
+    }
+
+    const size_t threadgroup_size_x = (Bm / Sg_Bm) * (Bn / Sg_Bn) * threads_per_simdgroup;
+    const size_t threadgroup_size_y = 1;
+    const size_t threadgroup_size_z = 1;
+    const size_t total_threadgroup_size = threadgroup_size_x * threadgroup_size_y * threadgroup_size_z;
+    if (total_threadgroup_size > f32_mf4w_moe_dense_matmul_fn->max_threadgroup_threads) {
+        GPTOSS_LOG_ERROR("failed to encode f32_mf4w_moe_dense_matmul kernel launch: total threadgroup size (%zu) exceeds supported maximum (%zu)",
+            total_threadgroup_size, f32_mf4w_moe_dense_matmul_fn->max_threadgroup_threads);
+        return gptoss_status_invalid_argument;
+    }
+    if (n % Bn != 0) {
+        GPTOSS_LOG_ERROR("failed to encode f32_mf4w_moe_dense_matmul kernel launch: n (%" PRIu32 ") is not divisible by Bn (%" PRIu32 ")",
+            n, Bn);
+        return gptoss_status_invalid_argument;
+    }
+    if (k % Bk != 0) {
+        GPTOSS_LOG_ERROR("failed to encode f32_mf4w_moe_dense_matmul kernel launch: k (%" PRIu32 ") is not divisible by Bk (%" PRIu32 ")",
+            k, Bk);
+        return gptoss_status_invalid_argument;
+    }
+
+    const size_t grid_y = math_ceil_div(m, Bm);
+    const size_t grid_x = n / Bn;
+    const size_t grid_z = num_experts;
+
+    return gptoss_metal_command_buffer_encode_launch_kernel(
+        command_buffer, f32_mf4w_moe_dense_matmul_fn,
+        threadgroup_size_x, threadgroup_size_y, threadgroup_size_z,
+        grid_x, grid_y, grid_z,
+        sizeof(args), &args,
+        6,
+        (const struct gptoss_metal_buffer *[]) {expert_offsets_buffer, input_buffer, weight_block_buffer, weight_scale_buffer, bias_buffer, output_buffer},
+        (const size_t[]) {expert_offsets_offset, input_offset, weight_block_offset, weight_scale_offset, bias_offset, output_offset},
+        /*threadgroup_buffer_size=*/0);
+}
+
+enum gptoss_status gptoss_metal_command_buffer_encode_launch_f32_accumulate(
+    const struct gptoss_metal_command_buffer* command_buffer,
+    const struct gptoss_metal_function* f32_accumulate_fn,
+    size_t threadgroup_size,
+    size_t max_threadgroups,
+    const struct gptoss_metal_buffer* input_buffer,
+    size_t input_offset,
+    const struct gptoss_metal_buffer* expert_buffer,
+    size_t expert_offset,
+    const struct gptoss_metal_buffer* output_buffer,
+    size_t output_offset,
+    const struct gptoss_metal_buffer* control_buffer,
+    size_t control_offset,
+    uint32_t num_channels,
+    uint32_t num_tokens,
+    uint32_t num_experts)
+{
+    if (command_buffer->object == NULL || f32_accumulate_fn->pipeline_state_object == NULL) {
+        return gptoss_status_invalid_state;
+    }
+
+    if (num_channels% 4 != 0) {
+        return gptoss_status_invalid_argument;
+    }
+
+    if (threadgroup_size == 0) {
+        threadgroup_size = f32_accumulate_fn->max_threadgroup_threads;
+    } else if (threadgroup_size > f32_accumulate_fn->max_threadgroup_threads) {
+        return gptoss_status_invalid_argument;
+    }
+
+    const size_t num_vecs = num_channels / 4;
+    const size_t num_vecs_per_expert = num_vecs * num_tokens;
+    const size_t num_vecs_per_threadgroup = math_ceil_div(num_vecs, max_threadgroups * threadgroup_size) * threadgroup_size;
+    const size_t num_threadgroups = math_min(max_threadgroups, math_ceil_div(num_vecs, num_vecs_per_threadgroup));
+    const struct gptoss_accumulate_args args = {
+        .num_vecs_per_expert = num_vecs_per_expert,
+        .num_vecs_per_threadgroup = num_vecs_per_threadgroup,
+        .num_vecs = num_vecs,
+    };
+
+    return gptoss_metal_command_buffer_encode_launch_kernel(
+        command_buffer, f32_accumulate_fn,
+        threadgroup_size, 1, 1,
+        num_threadgroups, num_tokens, 1,
+        sizeof(args), &args,
+        4,
+        (const struct gptoss_metal_buffer *[]) {input_buffer, expert_buffer, output_buffer, control_buffer},
+        (const size_t[]) {input_offset, expert_offset, output_offset, control_offset},
+        /*threadgroup_buffer_size=*/0);
+}
+
+enum gptoss_status gptoss_metal_command_buffer_encode_launch_f32_topk(
+    const struct gptoss_metal_command_buffer* command_buffer,
+    const struct gptoss_metal_function* f32_topk_fn,
+    const struct gptoss_metal_buffer* input_buffer,
+    size_t input_offset,
+    const struct gptoss_metal_buffer* output_buffer,
+    size_t output_offset,
+    const struct gptoss_metal_buffer* control_buffer,
+    size_t control_offset,
+    uint32_t num_tokens,
+    uint32_t num_experts,
+    uint32_t num_active_experts)
+{
+    if (command_buffer->object == NULL || f32_topk_fn->pipeline_state_object == NULL) {
+        return gptoss_status_invalid_state;
+    }
+
+    if (num_experts != 32  && num_experts != 128) {
+        return gptoss_status_invalid_argument;
+    }
+
+    if (num_active_experts != 4) {
+        return gptoss_status_invalid_argument;
+    }
+
+    const struct gptoss_topk_args args = { 0 };
+
+    return gptoss_metal_command_buffer_encode_launch_kernel(
+        command_buffer, f32_topk_fn,
+        /*threadgroup_size=*/32, 1, 1,
+        num_tokens, 1, 1,
+        sizeof(args), &args,
+        3,
+        (const struct gptoss_metal_buffer *[]) {input_buffer, output_buffer, control_buffer},
+        (const size_t[]) {input_offset, output_offset, control_offset},
+        /*threadgroup_buffer_size=*/0);
+}
+
+enum gptoss_status gptoss_metal_command_buffer_encode_launch_f32_sdpa(
+    const struct gptoss_metal_command_buffer* command_buffer,
+    const struct gptoss_metal_function* f32_sdpa_fn,
+    const struct gptoss_metal_buffer* q_buffer,
+    size_t q_offset,
+    const struct gptoss_metal_buffer* kv_buffer,
+    size_t kv_offset,
+    const struct gptoss_metal_buffer* s_buffer,
+    size_t s_offset,
+    const struct gptoss_metal_buffer* output_buffer,
+    size_t output_offset,
+    const struct gptoss_metal_buffer* control_buffer,
+    size_t control_offset,
+    uint32_t window,
+    uint32_t kv_stride,
+    uint32_t num_q_tokens,
+    uint32_t num_kv_tokens,
+    uint32_t num_q_heads,
+    uint32_t num_kv_heads,
+    uint32_t head_dim)
+{
+    if (command_buffer->object == NULL || f32_sdpa_fn->pipeline_state_object == NULL) {
+        return gptoss_status_invalid_state;
+    }
+
+    if (num_q_heads != num_kv_heads * 8) {
+        GPTOSS_LOG_ERROR("number of Q heads (%" PRIu32 ") must be 8 times the number of KV heads (%" PRIu32 ")",
+            num_q_heads, num_kv_heads);
+        return gptoss_status_invalid_argument;
+    }
+
+    if (head_dim != 64) {
+        GPTOSS_LOG_ERROR("attention head dimension (%" PRIu32 ") must be 64", head_dim);
+        return gptoss_status_invalid_argument;
+    }
+
+    const size_t max_context_tokens = math_min(num_q_tokens + num_kv_tokens + 1, window);
+    const size_t threadgroup_size = math_min(f32_sdpa_fn->max_threadgroup_threads,
+        max_context_tokens * f32_sdpa_fn->simdgroup_threads);
+    const size_t half_threadgroup_size = math_round_down_po2(threadgroup_size / 2, f32_sdpa_fn->simdgroup_threads);
+
+    const struct gptoss_sdpa_args args = {
+        .qkv_dim = head_dim * (num_q_heads + 2 * num_kv_heads),
+        .num_kv_tokens = num_kv_tokens,
+        .kv_stride = kv_stride,
+        .window = window,
+    };
+
+    return gptoss_metal_command_buffer_encode_launch_kernel(
+        command_buffer, f32_sdpa_fn,
+        threadgroup_size, 1, 1,
+        num_q_tokens, num_kv_heads, 1,
+        sizeof(args), &args,
+        5,
+        (const struct gptoss_metal_buffer *[]) {q_buffer, kv_buffer, s_buffer, output_buffer, control_buffer},
+        (const size_t[]) {q_offset, kv_offset, s_offset, output_offset, control_offset},
+        /*threadgroup_buffer_size=*/half_threadgroup_size * 8 * 4 * sizeof(float));
+}
+
+enum gptoss_status gptoss_metal_command_buffer_encode_launch_f32_softmax(
+    const struct gptoss_metal_command_buffer* command_buffer,
+    const struct gptoss_metal_function* f32_softmax_fn,
+    size_t threadgroup_size,
+    size_t max_threadgroups,
+    const struct gptoss_metal_buffer* score_buffer,
+    size_t score_offset,
+    const struct gptoss_metal_buffer* argmax_buffer,
+    size_t argmax_offset,
+    const struct gptoss_metal_buffer* prob_buffer,
+    size_t prob_offset,
+    const struct gptoss_metal_buffer* sum_buffer,
+    size_t sum_offset,
+    const struct gptoss_metal_buffer* control_buffer,
+    size_t control_offset,
+    uint32_t num_channels,
+    uint32_t num_tokens,
+    float temperature,
+    uint32_t* num_threadgroups_out,
+    uint32_t* num_channels_per_threadgroup_out)
+{
+    *num_threadgroups_out = 0;
+    *num_channels_per_threadgroup_out = 0;
+    if (command_buffer->object == NULL || f32_softmax_fn->pipeline_state_object == NULL) {
+        return gptoss_status_invalid_state;
+    }
+
+    const size_t num_vecs = num_channels;
+    const size_t num_vecs_per_threadgroup = math_ceil_div(num_vecs, max_threadgroups * threadgroup_size) * threadgroup_size;
+    const size_t num_threadgroups = math_min(max_threadgroups, math_ceil_div(num_vecs, num_vecs_per_threadgroup));
+    const struct gptoss_softmax_args args = {
+        .num_vecs = num_vecs,
+        .num_vecs_per_threadgroup = num_vecs_per_threadgroup,
+        .max_threadgroups = max_threadgroups,
+        .temperature = temperature,
+    };
+
+    *num_threadgroups_out = num_threadgroups;
+    *num_channels_per_threadgroup_out = num_vecs_per_threadgroup;
+    return gptoss_metal_command_buffer_encode_launch_kernel(
+        command_buffer, f32_softmax_fn,
+        threadgroup_size, 1, 1,
+        num_threadgroups, num_tokens, 1,
+        sizeof(args), &args,
+        5,
+        (const struct gptoss_metal_buffer *[]) {score_buffer, argmax_buffer, prob_buffer, sum_buffer, control_buffer},
+        (const size_t[]) {score_offset, argmax_offset, prob_offset, sum_offset, control_offset},
+        /*threadgroup_buffer_size=*/0);
+}
+
+enum gptoss_status gptoss_metal_command_buffer_encode_launch_f32_sample(
+    const struct gptoss_metal_command_buffer* command_buffer,
+    const struct gptoss_metal_function* f32_sample_fn,
+    size_t min_threadgroup_size,
+    const struct gptoss_metal_buffer* prob_buffer,
+    size_t prob_offset,
+    const struct gptoss_metal_buffer* sum_buffer,
+    size_t sum_offset,
+    const struct gptoss_metal_buffer* token_buffer,
+    size_t token_offset,
+    const struct gptoss_metal_buffer* control_buffer,
+    size_t control_offset,
+    uint64_t rng_seed,
+    uint32_t rng_offset,
+    uint32_t num_blocks,
+    uint32_t num_channels,
+    uint32_t num_channels_per_block)
+{
+    if (command_buffer->object == NULL || f32_sample_fn->pipeline_state_object == NULL) {
+        return gptoss_status_invalid_state;
+    }
+
+    if (min_threadgroup_size > f32_sample_fn->max_threadgroup_threads) {
+        return gptoss_status_invalid_argument;
+    }
+
+    if (min_threadgroup_size % f32_sample_fn->simdgroup_threads != 0) {
+        return gptoss_status_invalid_argument;
+    }
+
+    if (num_blocks > f32_sample_fn->max_threadgroup_threads) {
+        return gptoss_status_invalid_argument;
+    }
+
+    const struct gptoss_sample_args args = {
+        .rng_seed = rng_seed,
+        .rng_offset = rng_offset,
+        .num_blocks = num_blocks,
+        .num_dims = num_channels,
+        .num_dims_per_block = num_channels_per_block,
+    };
+
+    const size_t threadgroup_size = math_max(min_threadgroup_size,
+        math_round_up_po2(num_blocks, f32_sample_fn->simdgroup_threads));
+    return gptoss_metal_command_buffer_encode_launch_kernel(
+        command_buffer, f32_sample_fn,
+        threadgroup_size, 1, 1,
+        1, 1, 1,
+        sizeof(args), &args,
+        4,
+        (const struct gptoss_metal_buffer *[]) {prob_buffer, sum_buffer, token_buffer, control_buffer},
+        (const size_t[]) {prob_offset, sum_offset, token_offset, control_offset},
+        /*threadgroup_buffer_size=*/0);
+}

gptoss_kernels/source/metal.m

@@ -0,0 +1,482 @@
+#import <Foundation/Foundation.h>
+#import <Metal/Metal.h>
+
+#include <dispatch/dispatch.h>
+#include <mach-o/getsect.h>
+
+#include <gpt-oss/types.h>
+
+#include <internal/log.h>
+#include <internal/metal.h>
+
+
+static size_t gptoss_metal_device_get_core_count(id<MTLDevice> device) {
+    if (!device) {
+        return 0;
+    }
+
+    const uint64_t target_registry_id = [device registryID];
+
+    io_iterator_t it = IO_OBJECT_NULL;
+    const kern_return_t kr = IOServiceGetMatchingServices(
+        kIOMainPortDefault,
+        IOServiceMatching("IOAccelerator"),
+        &it
+    );
+    if (kr != KERN_SUCCESS) {
+        GPTOSS_LOG_ERROR("failed to find IOAccelerator objects: error %d", kr);
+        return 0;
+    }
+
+    size_t result = 0;
+    for (io_object_t obj = IOIteratorNext(it); obj != IO_OBJECT_NULL; obj = IOIteratorNext(it)) {
+        uint64_t registry_id = 0;
+        if (IORegistryEntryGetRegistryEntryID(obj, &registry_id) == KERN_SUCCESS &&
+            registry_id == target_registry_id)
+        {
+            // Read "gpu-core-count" from this accelerator node
+            const CFTypeRef value = IORegistryEntryCreateCFProperty(
+                obj, CFSTR("gpu-core-count"), kCFAllocatorDefault, 0);
+            if (value != NULL) {
+                if (CFGetTypeID(value) == CFNumberGetTypeID()) {
+                    int32_t n = -1;
+                    if (CFNumberGetValue((CFNumberRef) value, kCFNumberSInt32Type, &n) && n > 0) {
+                        result = (size_t) n;
+                    }
+                }
+                CFRelease(value);
+            }
+            IOObjectRelease(obj);
+            break;
+        }
+        IOObjectRelease(obj);
+    }
+
+    IOObjectRelease(it);
+    return result;
+}
+
+enum gptoss_status gptoss_metal_device_create_system_default(
+    struct gptoss_metal_device* device_out)
+{
+    id<MTLDevice> device_obj = MTLCreateSystemDefaultDevice();
+    if (device_obj == nil) {
+        GPTOSS_LOG_ERROR("failed to create Metal device");
+        return gptoss_status_unsupported_system;
+    }
+
+    device_out->object = (void*) device_obj;
+    device_out->num_cores = gptoss_metal_device_get_core_count(device_obj);
+    device_out->max_buffer_size = (size_t) [device_obj maxBufferLength];
+    device_out->max_threadgroup_memory = (size_t) [device_obj maxThreadgroupMemoryLength];
+    const MTLSize max_threadgroup_threads = [device_obj maxThreadsPerThreadgroup];
+    device_out->max_threadgroup_threads_x = (size_t) max_threadgroup_threads.width;
+    device_out->max_threadgroup_threads_y = (size_t) max_threadgroup_threads.height;
+    device_out->max_threadgroup_threads_z = (size_t) max_threadgroup_threads.depth;
+    return gptoss_status_success;
+}
+
+enum gptoss_status gptoss_metal_device_release(
+    struct gptoss_metal_device* device)
+{
+    if (device->object != NULL) {
+        id<MTLDevice> device_obj = (id<MTLDevice>) device->object;
+        [device_obj release];
+    }
+    memset(device, 0, sizeof(struct gptoss_metal_device));
+    return gptoss_status_success;
+}
+
+extern const struct mach_header_64 __dso_handle;
+
+enum gptoss_status gptoss_metal_library_create_default(
+    const struct gptoss_metal_device* device,
+    struct gptoss_metal_library* library_out)
+{
+    enum gptoss_status status = gptoss_status_success;
+    id<MTLDevice> device_obj = (id<MTLDevice>) device->object;
+    id<MTLLibrary> library_obj = nil;
+    NSAutoreleasePool* autorelease_pool = nil;
+    dispatch_data_t library_blob = NULL;
+
+    unsigned long library_size = 0;
+    uint8_t* library_data = getsectiondata(&__dso_handle, "__METAL", "__shaders", &library_size);
+    if (library_data != NULL) {
+        library_blob = dispatch_data_create(library_data, library_size, NULL, DISPATCH_DATA_DESTRUCTOR_DEFAULT);
+
+        autorelease_pool = [[NSAutoreleasePool alloc] init];
+        NSError* error_obj = nil;
+        library_obj = [device_obj newLibraryWithData:library_blob error:&error_obj];
+        if (library_obj == nil) {
+            GPTOSS_LOG_ERROR("failed to create Metal library: %s", [[error_obj localizedDescription] UTF8String]);
+            status = gptoss_status_unsupported_system;
+            goto cleanup;
+        }
+    } else {
+        // Fall-back to loading from the bundle
+        library_obj = [device_obj newDefaultLibrary];
+        if (library_obj == nil) {
+            GPTOSS_LOG_ERROR("failed to create Metal default library");
+            status = gptoss_status_unsupported_system;
+            goto cleanup;
+        }
+    }
+
+    *library_out = (struct gptoss_metal_library) {
+        .object = (void*) library_obj,
+    };
+
+cleanup:
+    if (library_blob != NULL) {
+        dispatch_release(library_blob);
+    }
+    if (autorelease_pool != nil) {
+        [autorelease_pool drain];
+    }
+    return status;
+}
+
+enum gptoss_status gptoss_metal_library_release(
+    struct gptoss_metal_library* library)
+{
+    if (library->object != NULL) {
+        id<MTLLibrary> library_obj = (id<MTLLibrary>) library->object;
+        [library_obj release];
+    }
+    memset(library, 0, sizeof(struct gptoss_metal_library));
+    return gptoss_status_success;
+}
+
+enum gptoss_status gptoss_metal_function_create(
+    const struct gptoss_metal_library* library,
+    const char* name,
+    struct gptoss_metal_function* function_out)
+{
+    __block NSString* error_string_obj = nil;
+    id<MTLFunction> function_obj = nil;
+    MTLComputePipelineDescriptor* pipeline_descriptor_obj = nil;
+    __block id<MTLComputePipelineState> pipeline_state_obj = nil;
+    dispatch_semaphore_t pipeline_build_semaphore = NULL;
+    enum gptoss_status status = gptoss_status_success;
+
+    NSAutoreleasePool* autorelease_pool = [[NSAutoreleasePool alloc] init];
+    id<MTLLibrary> library_obj = (id<MTLLibrary>) library->object;
+    NSString* name_obj = [NSString stringWithUTF8String:name];
+    function_obj = [library_obj newFunctionWithName:name_obj];
+    if (function_obj == nil) {
+        GPTOSS_LOG_ERROR("failed to create Metal function %s", name);
+        status = gptoss_status_unsupported_system;
+        goto cleanup;
+    }
+    id<MTLDevice> device_obj = [library_obj device];
+    pipeline_descriptor_obj = [[MTLComputePipelineDescriptor alloc] init];
+    [pipeline_descriptor_obj setComputeFunction:function_obj];
+    [pipeline_descriptor_obj setThreadGroupSizeIsMultipleOfThreadExecutionWidth:YES];
+
+    pipeline_build_semaphore = dispatch_semaphore_create(/*value=*/0);
+    [device_obj newComputePipelineStateWithDescriptor:pipeline_descriptor_obj
+                                              options:MTLPipelineOptionNone
+                                    completionHandler:^(id<MTLComputePipelineState> _Nullable new_state,
+                                                        MTLComputePipelineReflection* _Nullable reflection,
+                                                        NSError* _Nullable error_obj) {
+        if (new_state != nil) {
+            pipeline_state_obj = [new_state retain];
+        }
+        if (error_obj != nil) {
+            error_string_obj = [[error_obj localizedDescription] copy];
+        }
+        dispatch_semaphore_signal(pipeline_build_semaphore);
+    }];
+    dispatch_semaphore_wait(pipeline_build_semaphore, DISPATCH_TIME_FOREVER);
+
+    if (pipeline_state_obj == nil) {
+        const char* error_string = "unknown error";
+        if (error_string_obj != nil) {
+            error_string = [error_string_obj UTF8String];
+        }
+        GPTOSS_LOG_ERROR("failed to create Metal compute pipeline state for function %s: %s",
+            name, error_string);
+        status = gptoss_status_unsupported_system;
+        goto cleanup;
+    }
+
+    // Commit
+    function_out->function_object = function_obj;
+    function_out->pipeline_state_object = pipeline_state_obj;
+    function_out->max_threadgroup_threads = (size_t) [pipeline_state_obj maxTotalThreadsPerThreadgroup];
+    function_out->simdgroup_threads = (size_t) [pipeline_state_obj threadExecutionWidth];
+    function_out->static_threadgroup_memory = (size_t) [pipeline_state_obj staticThreadgroupMemoryLength];
+
+    function_obj = nil;
+    pipeline_state_obj = nil;
+
+cleanup:
+    if (function_obj != nil) {
+        [function_obj release];
+    }
+    if (pipeline_descriptor_obj != nil) {
+        [pipeline_descriptor_obj release];
+    }
+    if (error_string_obj != nil) {
+        [error_string_obj release];
+    }
+    if (pipeline_build_semaphore != NULL) {
+        dispatch_release(pipeline_build_semaphore);
+    }
+    if (autorelease_pool != nil) {
+        [autorelease_pool drain];
+    }
+    return status;
+}
+
+enum gptoss_status gptoss_metal_function_release(
+    struct gptoss_metal_function* function)
+{
+    if (function->pipeline_state_object != NULL) {
+        id<MTLComputePipelineState> pipeline_state_obj = (id<MTLComputePipelineState>) function->pipeline_state_object;
+        [pipeline_state_obj release];
+    }
+    if (function->function_object != NULL) {
+        id<MTLFunction> function_obj = (id<MTLFunction>) function->function_object;
+        [function_obj release];
+    }
+    memset(function, 0, sizeof(struct gptoss_metal_function));
+    return gptoss_status_success;
+}
+
+enum gptoss_status gptoss_metal_buffer_create(
+    const struct gptoss_metal_device* device,
+    size_t size,
+    const void* data,
+    struct gptoss_metal_buffer* buffer_out)
+{
+    id<MTLDevice> device_obj = (id<MTLDevice>) device->object;
+    id<MTLBuffer> buffer_obj = nil;
+    if (data != NULL) {
+        buffer_obj = [device_obj newBufferWithBytes:data length:size options:MTLResourceStorageModeShared];
+    } else {
+        buffer_obj = [device_obj newBufferWithLength:size options:MTLResourceStorageModeShared];
+    }
+    if (buffer_obj == nil) {
+        GPTOSS_LOG_ERROR("failed to create Metal buffer of size %zu", size);
+        return gptoss_status_unsupported_system;
+    }
+    buffer_out->object = (void*) buffer_obj;
+    buffer_out->size = size;
+    buffer_out->ptr = [buffer_obj contents];
+    return gptoss_status_success;
+}
+
+enum gptoss_status gptoss_metal_buffer_wrap(
+    const struct gptoss_metal_device* device,
+    size_t size,
+    const void* data,
+    struct gptoss_metal_buffer* buffer_out)
+{
+    id<MTLDevice> device_obj = (id<MTLDevice>) device->object;
+    id<MTLBuffer> buffer_obj = [device_obj newBufferWithBytesNoCopy:(void*) data length:size options:MTLResourceStorageModeShared deallocator:nil];
+    if (buffer_obj == nil) {
+        GPTOSS_LOG_ERROR("failed to wrap Metal buffer of size %zu", size);
+        return gptoss_status_unsupported_system;
+    }
+    buffer_out->object = (void*) buffer_obj;
+    buffer_out->size = size;
+    buffer_out->ptr = (void*) data;
+    return gptoss_status_success;
+}
+
+enum gptoss_status gptoss_metal_buffer_release(
+    struct gptoss_metal_buffer* buffer)
+{
+    if (buffer->object != NULL) {
+        id<MTLBuffer> buffer_obj = (id<MTLBuffer>) buffer->object;
+        [buffer_obj release];
+    }
+    memset(buffer, 0, sizeof(struct gptoss_metal_buffer));
+    return gptoss_status_success;
+}
+
+enum gptoss_status gptoss_metal_command_queue_create(
+    const struct gptoss_metal_device* device,
+    struct gptoss_metal_command_queue* command_queue_out)
+{
+    id<MTLDevice> device_obj = (id<MTLDevice>) device->object;
+    id<MTLCommandQueue> command_queue_obj = [device_obj newCommandQueue];
+    if (command_queue_obj == nil) {
+        GPTOSS_LOG_ERROR("failed to create Metal command queue");
+        return gptoss_status_unsupported_system;
+    }
+    command_queue_out->object = (void*) command_queue_obj;
+    return gptoss_status_success;
+}
+
+enum gptoss_status gptoss_metal_command_queue_release(
+    struct gptoss_metal_command_queue* command_queue)
+{
+    if (command_queue->object != NULL) {
+        id<MTLCommandQueue> command_queue_obj = (id<MTLCommandQueue>) command_queue->object;
+        [command_queue_obj release];
+    }
+    memset(command_queue, 0, sizeof(struct gptoss_metal_command_queue));
+    return gptoss_status_success;
+}
+
+enum gptoss_status gptoss_metal_command_buffer_create(
+    const struct gptoss_metal_command_queue* command_queue,
+    struct gptoss_metal_command_buffer* command_buffer_out)
+{
+    id<MTLCommandQueue> command_queue_obj = (id<MTLCommandQueue>) command_queue->object;
+    id<MTLCommandBuffer> command_buffer_obj = [command_queue_obj commandBuffer];
+    if (command_buffer_obj == nil) {
+        GPTOSS_LOG_ERROR("failed to create Metal command buffer");
+        return gptoss_status_unsupported_system;
+    }
+    [command_buffer_obj retain];
+    command_buffer_out->object = (void*) command_buffer_obj;
+    return gptoss_status_success;
+}
+
+enum gptoss_status gptoss_metal_command_buffer_encode_fill_buffer(
+    const struct gptoss_metal_command_buffer* command_buffer,
+    const struct gptoss_metal_buffer* buffer,
+    size_t offset,
+    size_t size,
+    uint8_t fill_value)
+{
+    if (command_buffer->object == NULL) {
+        return gptoss_status_invalid_state;
+    }
+    if (buffer->object == NULL) {
+        return gptoss_status_invalid_argument;
+    }
+
+    id<MTLCommandBuffer> command_buffer_obj = (id<MTLCommandBuffer>) command_buffer->object;
+    id<MTLBuffer> buffer_obj = (id<MTLBuffer>) buffer->object;
+
+    id<MTLBlitCommandEncoder> command_encoder_obj = [command_buffer_obj blitCommandEncoder];
+
+    const NSRange range = NSMakeRange((NSUInteger) offset, (NSUInteger) size);
+    [command_encoder_obj fillBuffer:buffer_obj range:range value:fill_value];
+    [command_encoder_obj endEncoding];
+
+    return gptoss_status_success;
+}
+
+enum gptoss_status gptoss_metal_command_buffer_encode_copy_buffer(
+    const struct gptoss_metal_command_buffer* command_buffer,
+    const struct gptoss_metal_buffer* input_buffer,
+    size_t input_offset,
+    const struct gptoss_metal_buffer* output_buffer,
+    size_t output_offset,
+    size_t size)
+{
+    if (command_buffer->object == NULL) {
+        return gptoss_status_invalid_state;
+    }
+    if (input_buffer->object == NULL) {
+        return gptoss_status_invalid_argument;
+    }
+    if (output_buffer->object == NULL) {
+        return gptoss_status_invalid_argument;
+    }
+
+    id<MTLCommandBuffer> command_buffer_obj = (id<MTLCommandBuffer>) command_buffer->object;
+    id<MTLBuffer> input_buffer_obj = (id<MTLBuffer>) input_buffer->object;
+    id<MTLBuffer> output_buffer_obj = (id<MTLBuffer>) output_buffer->object;
+
+    id<MTLBlitCommandEncoder> command_encoder_obj = [command_buffer_obj blitCommandEncoder];
+
+    [command_encoder_obj copyFromBuffer:input_buffer_obj sourceOffset:(NSUInteger) input_offset
+                         toBuffer:output_buffer_obj destinationOffset:(NSUInteger) output_offset
+                         size:(NSUInteger) size];
+    [command_encoder_obj endEncoding];
+
+    return gptoss_status_success;
+}
+
+enum gptoss_status gptoss_metal_command_buffer_encode_launch_kernel(
+    const struct gptoss_metal_command_buffer* command_buffer,
+    const struct gptoss_metal_function* function,
+    size_t threadgroup_size_x,
+    size_t threadgroup_size_y,
+    size_t threadgroup_size_z,
+    size_t num_threadgroups_x,
+    size_t num_threadgroups_y,
+    size_t num_threadgroups_z,
+    size_t params_size,
+    const void* params,
+    size_t num_device_buffers,
+    const struct gptoss_metal_buffer** device_buffers,
+    const size_t* device_buffer_offsets,
+    size_t threadgroup_buffer_size)
+{
+    if (command_buffer->object == NULL || function->pipeline_state_object == NULL) {
+        return gptoss_status_invalid_state;
+    }
+
+    id<MTLCommandBuffer> command_buffer_obj = (id<MTLCommandBuffer>) command_buffer->object;
+    id<MTLComputePipelineState> pipeline_state_obj = (id<MTLComputePipelineState>) function->pipeline_state_object;
+
+    id<MTLComputeCommandEncoder> command_encoder_obj = [command_buffer_obj computeCommandEncoder];
+
+    // Set kernel arguments
+    [command_encoder_obj setComputePipelineState:pipeline_state_obj];
+    [command_encoder_obj setBytes:params length:params_size atIndex:0];
+    for (size_t i = 0; i < num_device_buffers; ++i) {
+        id<MTLBuffer> buffer_obj = (id<MTLBuffer>) device_buffers[i]->object;
+        const NSUInteger offset = device_buffer_offsets == NULL ? 0 : (NSUInteger) device_buffer_offsets[i];
+        [command_encoder_obj setBuffer:buffer_obj offset:offset atIndex:i + 1];
+    }
+    if (threadgroup_buffer_size != 0) {
+        [command_encoder_obj setThreadgroupMemoryLength:threadgroup_buffer_size atIndex:0];
+    }
+
+    // Dispatch kernel
+    const MTLSize threadgroup_size = MTLSizeMake(threadgroup_size_x, threadgroup_size_y, threadgroup_size_z);
+    const MTLSize num_threadgroups = MTLSizeMake(num_threadgroups_x, num_threadgroups_y, num_threadgroups_z);
+    [command_encoder_obj dispatchThreadgroups:num_threadgroups threadsPerThreadgroup:threadgroup_size];
+    [command_encoder_obj endEncoding];
+
+    return gptoss_status_success;
+}
+
+enum gptoss_status gptoss_metal_command_buffer_commit(
+    const struct gptoss_metal_command_buffer* command_buffer)
+{
+    if (command_buffer->object == NULL) {
+        return gptoss_status_invalid_state;
+    }
+
+    id<MTLCommandBuffer> command_buffer_obj = (id<MTLCommandBuffer>) command_buffer->object;
+    [command_buffer_obj commit];
+    return gptoss_status_success;
+}
+
+enum gptoss_status gptoss_metal_command_buffer_wait_completion(
+    const struct gptoss_metal_command_buffer* command_buffer,
+    double* elapsed_seconds)
+{
+    if (command_buffer->object == NULL) {
+        return gptoss_status_invalid_state;
+    }
+
+    id<MTLCommandBuffer> command_buffer_obj = (id<MTLCommandBuffer>) command_buffer->object;
+    [command_buffer_obj waitUntilCompleted];
+    if (elapsed_seconds != NULL) {
+        const CFTimeInterval start_time = [command_buffer_obj GPUStartTime];
+        const CFTimeInterval end_time = [command_buffer_obj GPUEndTime];
+        *elapsed_seconds = (double) end_time - (double) start_time;
+    }
+    return gptoss_status_success;
+}
+
+enum gptoss_status gptoss_metal_command_buffer_release(
+    struct gptoss_metal_command_buffer* command_buffer)
+{
+    if (command_buffer->object != NULL) {
+        id<MTLCommandBuffer> command_buffer_obj = (id<MTLCommandBuffer>) command_buffer->object;
+        [command_buffer_obj release];
+    }
+    memset(command_buffer, 0, sizeof(struct gptoss_metal_command_buffer));
+    return gptoss_status_success;
+}

gptoss_kernels/source/model.c

@@ -0,0 +1,581 @@
+#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/moematmul.metal

@@ -0,0 +1,702 @@
+#include <internal/kernel-args.h>
+#include <metal_common>
+#include <metal_compute>
+#include <metal_math>
+#include <metal_simdgroup>
+#include <metal_stdlib>
+
+#pragma METAL fp math_mode(safe)
+#pragma METAL fp contract(off)
+#define ceil_div(a, b) (((a) + (b) - 1) / (b))
+
+// Each simdgroup reduces all channels of the input and computes a single channel of the output
+// + Efficient synchronization
+// + Sequential memory access within a warp
+// Each threadgroup computes (simdgroups_per_threadgroup) consecutive output channels
+// + Reuse input vector from threadgroup memory
+// + Avoid synchronization across warps when doing reduction
+
+kernel void gptoss_f32_mf4w_moe_matmul_swiglu(
+    constant gptoss_moe_matmul_swiglu_args& args [[ buffer(0) ]],
+    const device float4* input [[ buffer(1) ]],
+    const device gptoss_expert_prediction* expert [[ buffer(2) ]],
+    const device uint4* weight_blocks [[ buffer(3) ]],
+    const device uchar* weight_scales [[ buffer(4) ]],
+    const device bfloat* bias [[ buffer(5) ]],
+    device float* output [[ buffer(6) ]],
+    const device gptoss_control* control [[ buffer(7) ]],
+    uint3 gid [[threadgroup_position_in_grid]],
+    uint tid [[thread_index_in_threadgroup]],
+    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;
+    threadgroup float threadgroup_buffer[32];
+    if (control->abort != 0) {
+        return;
+    }
+
+    const uint num_column_vecs = args.num_column_vecs;
+    const uint row = gid.x * num_simdgroups + simdgroup_idx;
+    const uint expert_id = expert[gid.y * args.num_active_experts + gid.z].expert_id;
+
+    input += 8 * (gid.y * num_column_vecs + simdgroup_tid);
+    weight_blocks = (const device uint4*) ((uintptr_t) (weight_blocks + num_column_vecs * row + simdgroup_tid) + expert_id * args.weight_expert_stride);
+    weight_scales = (const device uchar*) ((uintptr_t) (weight_scales + num_column_vecs * row + simdgroup_tid) + expert_id * args.weight_expert_stride);
+    bias = (const device bfloat*) ((uintptr_t) (bias + row) + expert_id * args.weight_expert_stride);
+    output += gid.y * args.num_rows + gid.x * (num_simdgroups / 2) + gid.z * args.output_expert_stride;
+
+    uint num_iter = (num_column_vecs - simdgroup_tid + (simdgroup_size - 1)) / simdgroup_size;
+
+    float4 sum4 = 0.0f;
+    do {
+        const uint4 wblock = *weight_blocks;
+        const float wscale = as_type<float>(static_cast<uint>(*weight_scales) << 23);
+        uint4 wblock02468ACEGIKMOQSU = wblock + wblock;
+        uint4 wblock13579BDFHJLNPRTV = wblock >> 3;
+        wblock02468ACEGIKMOQSU &= 0x1E1E1E1Eu;
+        wblock13579BDFHJLNPRTV &= 0x1E1E1E1Eu;
+        wblock02468ACEGIKMOQSU += 0x70707070u;
+        wblock13579BDFHJLNPRTV += 0x70707070u;
+        wblock02468ACEGIKMOQSU &= 0x8E8E8E8Eu;
+        wblock13579BDFHJLNPRTV &= 0x8E8E8E8Eu;
+        const uint4 wblock26AEIMQU = wblock02468ACEGIKMOQSU & 0xFF00FF00u;
+        const uint4 wblock048CGKOS = (wblock02468ACEGIKMOQSU << 8) & 0xFF00FF00u;
+        const uint4 wblock37BFJNRV = wblock13579BDFHJLNPRTV & 0xFF00FF00u;
+        const uint4 wblock159DHLPT = (wblock13579BDFHJLNPRTV << 8) & 0xFF00FF00u;
+        const float4 w048C = static_cast<float4>(as_type<half4>(wblock048CGKOS.xy));
+        const float4 wGKOS = static_cast<float4>(as_type<half4>(wblock048CGKOS.zw));
+        const float4 w26AE = static_cast<float4>(as_type<half4>(wblock26AEIMQU.xy));
+        const float4 wIMQU = static_cast<float4>(as_type<half4>(wblock26AEIMQU.zw));
+        const float4 w159D = static_cast<float4>(as_type<half4>(wblock159DHLPT.xy));
+        const float4 wHLPT = static_cast<float4>(as_type<half4>(wblock159DHLPT.zw));
+        const float4 w37BF = static_cast<float4>(as_type<half4>(wblock37BFJNRV.xy));
+        const float4 wJNRV = static_cast<float4>(as_type<half4>(wblock37BFJNRV.zw));
+
+        const float4 w0123 = (float4) { w048C.x, w159D.x, w26AE.x, w37BF.x };
+        const float4 w4567 = (float4) { w048C.y, w159D.y, w26AE.y, w37BF.y };
+        const float4 w89AB = (float4) { w048C.z, w159D.z, w26AE.z, w37BF.z };
+        const float4 wCDEF = (float4) { w048C.w, w159D.w, w26AE.w, w37BF.w };
+        const float4 wGHIJ = (float4) { wGKOS.x, wHLPT.x, wIMQU.x, wJNRV.x };
+        const float4 wKLMN = (float4) { wGKOS.y, wHLPT.y, wIMQU.y, wJNRV.y };
+        const float4 wOPQR = (float4) { wGKOS.z, wHLPT.z, wIMQU.z, wJNRV.z };
+        const float4 wSTUV = (float4) { wGKOS.w, wHLPT.w, wIMQU.w, wJNRV.w };
+
+        const float4 i0123 = input[0];
+        const float4 i4567 = input[1];
+        const float4 i89AB = input[2];
+        const float4 iCDEF = input[3];
+        const float4 iGHIJ = input[4];
+        const float4 iKLMN = input[5];
+        const float4 iOPQR = input[6];
+        const float4 iSTUV = input[7];
+
+        float4 psum0 = i0123 * w0123;
+        float4 psum1 = i4567 * w4567;
+        psum0 = metal::fma(i89AB, w89AB, psum0);
+        psum1 = metal::fma(iCDEF, wCDEF, psum1);
+        psum0 = metal::fma(iGHIJ, wGHIJ, psum0);
+        psum1 = metal::fma(iKLMN, wKLMN, psum1);
+        psum0 = metal::fma(iOPQR, wOPQR, psum0);
+        psum1 = metal::fma(iSTUV, wSTUV, psum1);
+        sum4 = metal::fma(psum0, wscale, sum4);
+        sum4 = metal::fma(psum1, wscale, sum4);
+
+        weight_blocks += simdgroup_size;
+        weight_scales += simdgroup_size;
+        input += 8 * 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);
+        threadgroup_buffer[simdgroup_idx] = sum;
+    }
+    metal::threadgroup_barrier(metal::mem_flags::mem_threadgroup);
+    if (tid * 2 < num_simdgroups) {
+        const float2 x = reinterpret_cast<const threadgroup float2*>(threadgroup_buffer)[tid];
+        const float swish_x = metal::min(x.x, args.swiglu_max);
+        const float linear_x = metal::clamp(x.y, args.swiglu_min, args.swiglu_max);
+        const float alpha = 1.702f;
+        const float swish_y = swish_x / (1.0f + metal::precise::exp(-alpha * swish_x));
+        const float swiglu_y = metal::fma(swish_y, linear_x, swish_y);
+        output[tid] = swiglu_y;
+    }
+}
+
+kernel void gptoss_f32_mf4w_moe_matmul(
+    constant gptoss_moe_matmul_args& args [[ buffer(0) ]],
+    const device float4* input [[ buffer(1) ]],
+    const device gptoss_expert_prediction* expert [[ buffer(2) ]],
+    const device uint4* weight_blocks [[ buffer(3) ]],
+    const device uchar* weight_scales [[ buffer(4) ]],
+    const device bfloat* bias [[ buffer(5) ]],
+    device float* output [[ buffer(6) ]],
+    const device gptoss_control* control [[ buffer(7) ]],
+    uint3 gid [[threadgroup_position_in_grid]],
+    uint tid [[thread_index_in_threadgroup]],
+    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;
+    const uint expert_id = expert[gid.y * args.num_active_experts + gid.z].expert_id;
+
+    input += 8 * (gid.y * num_column_vecs + simdgroup_tid + gid.z * args.input_expert_stride);
+    weight_blocks = (const device uint4*) ((uintptr_t) (weight_blocks + num_column_vecs * row + simdgroup_tid) + expert_id * args.weight_expert_stride);
+    weight_scales = (const device uchar*) ((uintptr_t) (weight_scales + num_column_vecs * row + simdgroup_tid) + expert_id * args.weight_expert_stride);
+    bias = (const device bfloat*) ((uintptr_t) (bias + row) + expert_id * args.weight_expert_stride);
+    output += gid.y * args.num_rows + row + gid.z * args.output_expert_stride;
+
+    uint num_iter = (num_column_vecs - simdgroup_tid + (simdgroup_size - 1)) / simdgroup_size;
+
+    float4 sum4 = 0.0f;
+    do {
+        const uint4 wblock = *weight_blocks;
+        const float wscale = as_type<float>(static_cast<uint>(*weight_scales) << 23);
+        uint4 wblock02468ACEGIKMOQSU = wblock + wblock;
+        uint4 wblock13579BDFHJLNPRTV = wblock >> 3;
+        wblock02468ACEGIKMOQSU &= 0x1E1E1E1Eu;
+        wblock13579BDFHJLNPRTV &= 0x1E1E1E1Eu;
+        wblock02468ACEGIKMOQSU += 0x70707070u;
+        wblock13579BDFHJLNPRTV += 0x70707070u;
+        wblock02468ACEGIKMOQSU &= 0x8E8E8E8Eu;
+        wblock13579BDFHJLNPRTV &= 0x8E8E8E8Eu;
+        const uint4 wblock26AEIMQU = wblock02468ACEGIKMOQSU & 0xFF00FF00u;
+        const uint4 wblock048CGKOS = (wblock02468ACEGIKMOQSU << 8) & 0xFF00FF00u;
+        const uint4 wblock37BFJNRV = wblock13579BDFHJLNPRTV & 0xFF00FF00u;
+        const uint4 wblock159DHLPT = (wblock13579BDFHJLNPRTV << 8) & 0xFF00FF00u;
+        const float4 w048C = static_cast<float4>(as_type<half4>(wblock048CGKOS.xy));
+        const float4 wGKOS = static_cast<float4>(as_type<half4>(wblock048CGKOS.zw));
+        const float4 w26AE = static_cast<float4>(as_type<half4>(wblock26AEIMQU.xy));
+        const float4 wIMQU = static_cast<float4>(as_type<half4>(wblock26AEIMQU.zw));
+        const float4 w159D = static_cast<float4>(as_type<half4>(wblock159DHLPT.xy));
+        const float4 wHLPT = static_cast<float4>(as_type<half4>(wblock159DHLPT.zw));
+        const float4 w37BF = static_cast<float4>(as_type<half4>(wblock37BFJNRV.xy));
+        const float4 wJNRV = static_cast<float4>(as_type<half4>(wblock37BFJNRV.zw));
+
+        const float4 w0123 = (float4) { w048C.x, w159D.x, w26AE.x, w37BF.x };
+        const float4 w4567 = (float4) { w048C.y, w159D.y, w26AE.y, w37BF.y };
+        const float4 w89AB = (float4) { w048C.z, w159D.z, w26AE.z, w37BF.z };
+        const float4 wCDEF = (float4) { w048C.w, w159D.w, w26AE.w, w37BF.w };
+        const float4 wGHIJ = (float4) { wGKOS.x, wHLPT.x, wIMQU.x, wJNRV.x };
+        const float4 wKLMN = (float4) { wGKOS.y, wHLPT.y, wIMQU.y, wJNRV.y };
+        const float4 wOPQR = (float4) { wGKOS.z, wHLPT.z, wIMQU.z, wJNRV.z };
+        const float4 wSTUV = (float4) { wGKOS.w, wHLPT.w, wIMQU.w, wJNRV.w };
+
+        const float4 i0123 = input[0];
+        const float4 i4567 = input[1];
+        const float4 i89AB = input[2];
+        const float4 iCDEF = input[3];
+        const float4 iGHIJ = input[4];
+        const float4 iKLMN = input[5];
+        const float4 iOPQR = input[6];
+        const float4 iSTUV = input[7];
+
+        float4 psum0 = i0123 * w0123;
+        float4 psum1 = i4567 * w4567;
+        psum0 = metal::fma(i89AB, w89AB, psum0);
+        psum1 = metal::fma(iCDEF, wCDEF, psum1);
+        psum0 = metal::fma(iGHIJ, wGHIJ, psum0);
+        psum1 = metal::fma(iKLMN, wKLMN, psum1);
+        psum0 = metal::fma(iOPQR, wOPQR, psum0);
+        psum1 = metal::fma(iSTUV, wSTUV, psum1);
+        sum4 = metal::fma(psum0, wscale, sum4);
+        sum4 = metal::fma(psum1, wscale, sum4);
+
+        weight_blocks += simdgroup_size;
+        weight_scales += simdgroup_size;
+        input += 8 * 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);
+        *output = sum;
+    }
+}
+
+kernel void gptoss_f32_mf4w_moe_dense_matmul_swiglu(
+    constant gptoss_moe_dense_matmul_swiglu_args& params [[ buffer(0) ]],
+    const device uint* __restrict__ expert_offsets [[ buffer(1) ]],
+    const device float* lhs [[ buffer(2) ]],
+    const device uint* weight_blocks [[ buffer(3) ]],
+    const device uchar* weight_scales [[ buffer(4) ]],
+    const device bfloat* __restrict__ bias [[ buffer(5) ]],
+    device float* out [[ buffer(6) ]],
+    uint sg_id [[simdgroup_index_in_threadgroup]],
+    uint3 threads_per_tg [[threads_per_threadgroup]],
+    uint sg_count_per_tg [[dispatch_simdgroups_per_threadgroup]],
+    uint3 gid [[thread_position_in_grid]],
+    uint3 tg_id [[threadgroup_position_in_grid]],
+    uint3 local_tid [[thread_position_in_threadgroup]]) 
+{
+    constexpr uint Bm = MOE_DENSE_MATMUL_SWIGLU_Bm;
+    constexpr uint Bn = MOE_DENSE_MATMUL_SWIGLU_Bn;
+    constexpr uint Bk = MOE_DENSE_MATMUL_SWIGLU_Bk;
+    constexpr uint Sg_Bm = MOE_DENSE_MATMUL_SWIGLU_Sg_Bm;
+    constexpr uint Sg_Bn = MOE_DENSE_MATMUL_SWIGLU_Sg_Bn;
+
+    // Assumptions about shapes.
+    assert(Bm % 8 == 0);
+    assert(Bn % 8 == 0);
+    assert(Bk % 8 == 0);
+    assert(Sg_Bm % 8 == 0);
+    assert(Sg_Bn % 8 == 0);
+    assert(Bm % Sg_Bm == 0);
+    assert(Bn % Sg_Bn == 0);
+
+    const uint K = params.k;
+    const uint N = params.n;
+    const uint M = expert_offsets[tg_id.z + 1] - expert_offsets[tg_id.z];
+    assert((K % 32) == 0);
+    assert((K % 8) == 0);
+    assert(N % Bn == 0);
+    assert(K % Bk == 0);
+    // Get row and col tg.
+    const uint row_tg = tg_id.y;
+    const uint col_tg = tg_id.x;
+    // Get row and col local tid.
+    const uint row_tg_offset = row_tg * Bm;
+    const uint col_tg_offset = col_tg * Bn;
+    if (row_tg_offset >= M || col_tg_offset >= N) {
+        return;
+    }
+    // Move lhs and output according to the passed offset.
+    const uint expert_offset = expert_offsets[tg_id.z];
+    lhs += expert_offset * K;
+    const uint N_output = N / 2;
+    out += expert_offset * N_output;
+
+    const uint S = params.weight_blocks_expert_stride_bytes;
+    const uint S_scales = params.weight_scales_expert_stride_bytes;
+    const uint S_bias = params.bias_expert_stride_bytes;
+
+    const device char* wb0 = reinterpret_cast<const device char*>(weight_blocks);
+    const device char* sc0 = reinterpret_cast<const device char*>(weight_scales);
+    const device char* bi0 = reinterpret_cast<const device char*>(bias);
+
+    weight_blocks = reinterpret_cast<const device uint*>(wb0 + tg_id.z * S);
+    weight_scales = reinterpret_cast<const device uchar*>(sc0 + tg_id.z * S_scales);
+    bias = reinterpret_cast<const device bfloat*>(bi0 + tg_id.z * S_bias);
+
+    const uint sg_col_count = Bn / Sg_Bn;
+    const uint row_sg = sg_id / sg_col_count;
+    const uint col_sg = sg_id % sg_col_count;
+
+    const uint row_sg_offset = row_sg * Sg_Bm;
+    const uint col_sg_offset = col_sg * Sg_Bn;
+    // Declare threadgroup blocks.
+    threadgroup float lhs_block[Bm * Bk];
+    // rhs_block will hold the scaled fp32 weights.
+    threadgroup float rhs_block[Bn * Bk];
+
+    constexpr uint temp_result_size = (Sg_Bm / 8) * (Sg_Bn / 8);
+    // Create an array of simdgroup_float8x8 to hold temp results.
+    metal::simdgroup_float8x8 OutTiles[temp_result_size];
+    for (uint i = 0; i < temp_result_size; i++) {
+        OutTiles[i] = metal::make_filled_simdgroup_matrix<float, 8, 8>(0.0);
+    }
+    // Linear thread id within TG (we launch 1-D TGs)
+    const uint lin_tid = local_tid.x;
+    const uint thread_count_per_tg = threads_per_tg.x * threads_per_tg.y * threads_per_tg.z;
+
+    // Iterate over all Bk blocks.
+    for (uint k_offset = 0; k_offset < K; k_offset += Bk) {
+        constexpr uint lhs_row_stride = Bk;
+        constexpr uint lhs_vec_cols = Bk / 4;
+        constexpr uint lhs_vec_total = Bm * lhs_vec_cols;
+
+        const uint LHS_ITERS = ceil_div(lhs_vec_total, thread_count_per_tg);
+
+        // #pragma clang loop unroll(full)
+        for (uint t = 0; t < LHS_ITERS; ++t) {
+            const uint i = t * thread_count_per_tg + lin_tid;
+            if (i < lhs_vec_total) {
+                const uint r = i / lhs_vec_cols;
+                const uint c4 = i % lhs_vec_cols;
+
+                const uint gr = row_tg_offset + r;
+                const uint gc4 = (k_offset / 4) + c4;
+
+                threadgroup float4* dst4 =
+                    reinterpret_cast<threadgroup float4*>(lhs_block + r * lhs_row_stride + (c4 << 2));
+                if (gr < M) {
+                    const device float4* src4 =
+                        reinterpret_cast<const device float4*>(lhs + gr * K + (gc4 << 2));
+
+                    *dst4 = *src4;
+                } else {
+                    *dst4 = float4(0.0);
+                }
+            }
+        }
+
+        // Load weights with vector loads.
+        constexpr uint rhs_row_stride = Bk;
+        constexpr uint weights_per_elem = 8;
+        constexpr uint rhs_loads_per_col = Bk / weights_per_elem;
+        constexpr uint rhs_loads_total = Bn * rhs_loads_per_col;
+        const uint RHS_ITERS = ceil_div(rhs_loads_total, thread_count_per_tg);
+        // #pragma clang loop unroll(full)
+        for (uint t = 0; t < RHS_ITERS; ++t) {
+            const uint i = t * thread_count_per_tg + lin_tid;
+            if (i < rhs_loads_total) {
+                const uint r = i / rhs_loads_per_col;
+                const uint c = i % rhs_loads_per_col;
+
+                const uint gr = col_tg_offset + r;
+                const uint gc = (k_offset / weights_per_elem) + c;
+                const uint gc_scale = (k_offset / 32) + (c >> 2);
+
+                const uint wblock = weight_blocks[gr * (K / weights_per_elem) + gc];
+                const float scale =
+                    as_type<float>(static_cast<uint>(weight_scales[gr * (K / 32) + gc_scale]) << 23);
+                uint wblock0246 = (wblock + wblock);
+                uint wblock1357 = (wblock >> 3);
+                wblock0246 &= 0x1E1E1E1Eu;
+                wblock1357 &= 0x1E1E1E1Eu;
+
+                wblock0246 += 0x70707070u;
+                wblock1357 += 0x70707070u;
+                wblock0246 &= 0x8E8E8E8Eu;
+                wblock1357 &= 0x8E8E8E8Eu;
+
+                uint wblock26 = (wblock0246) & 0xFF00FF00u;
+                uint wblock04 = ((wblock0246 << 8)) & 0xFF00FF00u;
+                uint wblock37 = (wblock1357) & 0xFF00FF00u;
+                uint wblock15 = ((wblock1357 << 8)) & 0xFF00FF00u;
+
+                half4 wblock0426 = as_type<half4>(uint2(wblock04, wblock26));
+                half4 wblock1537 = as_type<half4>(uint2(wblock15, wblock37));
+
+                // Convert to float scalars and apply scale
+                const float w0 = float(wblock0426.x) * scale;
+                const float w1 = float(wblock1537.x) * scale;
+                const float w2 = float(wblock0426.z) * scale;
+                const float w3 = float(wblock1537.z) * scale;
+                const float w4 = float(wblock0426.y) * scale;
+                const float w5 = float(wblock1537.y) * scale;
+                const float w6 = float(wblock0426.w) * scale;
+                const float w7 = float(wblock1537.w) * scale;
+                const uint rhs_offset = r * rhs_row_stride + c * 8;
+                rhs_block[rhs_offset] = w0;
+                rhs_block[rhs_offset + 1] = w1;
+                rhs_block[rhs_offset + 2] = w2;
+                rhs_block[rhs_offset + 3] = w3;
+                rhs_block[rhs_offset + 4] = w4;
+                rhs_block[rhs_offset + 5] = w5;
+                rhs_block[rhs_offset + 6] = w6;
+                rhs_block[rhs_offset + 7] = w7;
+            }
+        }
+        threadgroup_barrier(metal::mem_flags::mem_threadgroup);
+#pragma clang loop unroll(full)
+        for (uint k = 0; k < Bk; k += 8) {
+#pragma clang loop unroll(full)
+            for (uint m_subtile_ = 0; m_subtile_ < Sg_Bm; m_subtile_ += 8) {
+                const uint row_index_in_out_tile = m_subtile_ / 8;
+                metal::simdgroup_float8x8 lhs_frag;
+
+                simdgroup_load(lhs_frag, lhs_block, Bk, ulong2(k, m_subtile_ + row_sg_offset));
+#pragma clang loop unroll(full)
+                for (uint n_subtile_ = 0; n_subtile_ < Sg_Bn; n_subtile_ += 8) {
+                    const uint col_index_in_out_tile = n_subtile_ / 8;
+                    const uint current_index_out_tile =
+                        row_index_in_out_tile * (Sg_Bn / 8) + col_index_in_out_tile;
+                    metal::simdgroup_float8x8 rhs_frag;
+                    simdgroup_load(rhs_frag, rhs_block, Bk, ulong2(k, n_subtile_ + col_sg_offset), true);
+
+                    simdgroup_multiply_accumulate(OutTiles[current_index_out_tile], lhs_frag, rhs_frag,
+                        OutTiles[current_index_out_tile]);
+                }
+            }
+        }
+        threadgroup_barrier(metal::mem_flags::mem_threadgroup);
+    }
+
+    // Epilogue.
+    threadgroup float scratch[Bm * Bn];
+#pragma clang loop unroll(full)
+    for (uint n_subtile_ = 0; n_subtile_ < Sg_Bn; n_subtile_ += 8) {
+        const uint col_index_in_out_tile = n_subtile_ / 8;
+        const uint local_col_offset = col_sg_offset + n_subtile_;
+#pragma clang loop unroll(full)
+        for (uint m_subtile_ = 0; m_subtile_ < Sg_Bm; m_subtile_ += 8) {
+            const uint row_index_in_out_tile = m_subtile_ / 8;
+            const uint local_row_offset = row_sg_offset + m_subtile_;
+            const uint current_index_out_tile =
+                row_index_in_out_tile * (Sg_Bn / 8) + col_index_in_out_tile;
+            simdgroup_store(OutTiles[current_index_out_tile], scratch, Bn,
+                ulong2(local_col_offset, local_row_offset));
+        }
+    }
+    threadgroup float bias_tile[Bn];
+    // TODO(ibahmed): vectorize these loads an maybe unroll the loop.
+    for (uint c_local = local_tid.x; c_local < Bn; c_local += thread_count_per_tg) {
+        const uint c_global = col_tg_offset + c_local;
+        bias_tile[c_local] = (c_global < N) ? static_cast<float>(bias[c_global]) : 0.0f;
+    }
+
+    threadgroup_barrier(metal::mem_flags::mem_threadgroup);
+    const float alpha = 1.702f;
+    // TODO(ibahmed): vectorize these stores and maybe unroll the loop.
+    for (uint idx = local_tid.x; idx < Bm * Bn / 2; idx += thread_count_per_tg) {
+        const uint idx_swish = idx * 2;
+        const uint r = idx_swish / Bn;
+        const uint c_swish = idx_swish % Bn;
+
+        const uint out_row = row_tg_offset + r;
+        const uint out_col = (col_tg_offset / 2) + (c_swish / 2);
+
+        if (out_row < M && out_col < N_output) {
+            float acc_swish = scratch[idx_swish] + bias_tile[c_swish];
+            float acc_linear = scratch[idx_swish + 1] + bias_tile[c_swish + 1];
+            const float swish = metal::min(acc_swish, params.swiglu_max);
+            const float linear = metal::clamp(acc_linear, params.swiglu_min, params.swiglu_max);
+            const float swish_y = swish / (1.0f + metal::precise::exp(-alpha * swish));
+            const float swiglu_y = metal::fma(swish_y, linear, swish_y);
+            out[out_row * N_output + out_col] = swiglu_y;
+        }
+    }
+}
+
+kernel void gptoss_f32_mf4w_moe_dense_matmul(
+    constant gptoss_moe_dense_matmul_args& params [[ buffer(0) ]],
+    const device uint* __restrict__ expert_offsets [[ buffer(1) ]],
+    const device float* lhs [[ buffer(2) ]],
+    const device uint* weight_blocks [[ buffer(3) ]],
+    const device uchar* weight_scales [[ buffer(4) ]],
+    const device bfloat* __restrict__ bias [[ buffer(5) ]],
+    device float* out [[ buffer(6) ]],
+    uint sg_id [[simdgroup_index_in_threadgroup]],
+    uint3 threads_per_tg [[threads_per_threadgroup]],
+    uint sg_count_per_tg [[dispatch_simdgroups_per_threadgroup]],
+    uint3 gid [[thread_position_in_grid]],
+    uint3 tg_id [[threadgroup_position_in_grid]],
+    uint3 local_tid [[thread_position_in_threadgroup]]) 
+{
+    const uint Bm = MOE_DENSE_MATMUL_Bm;
+    const uint Bn = MOE_DENSE_MATMUL_Bn;
+    const uint Bk = MOE_DENSE_MATMUL_Bk;
+    const uint Sg_Bm = MOE_DENSE_MATMUL_Sg_Bm;
+    const uint Sg_Bn = MOE_DENSE_MATMUL_Sg_Bn;
+    assert(Bm % 8 == 0);
+    assert(Bn % 8 == 0);
+    assert(Bk % 8 == 0);
+    assert(Sg_Bm % 8 == 0);
+    assert(Sg_Bn % 8 == 0);
+    assert(Bm % Sg_Bm == 0);
+    assert(Bn % Sg_Bn == 0);
+
+    const uint K = params.k;
+    const uint N = params.n;
+    const uint M = expert_offsets[tg_id.z + 1] - expert_offsets[tg_id.z];
+    assert((K % 32) == 0);
+    assert((K % 8) == 0);
+    assert(N % Bn == 0);
+    assert(K % Bk == 0);
+    // Get row and col tg.
+    const uint row_tg = tg_id.y;
+    const uint col_tg = tg_id.x;
+    // Get row and col local tid.
+    const uint row_tg_offset = row_tg * Bm;
+    const uint col_tg_offset = col_tg * Bn;
+    if (row_tg_offset >= M || col_tg_offset >= N) {
+        return;
+    }
+    // Move lhs and output according to the passed offset.
+    const uint expert_offset = expert_offsets[tg_id.z];
+    lhs += expert_offset * K;
+    out += expert_offset * N;
+
+    const uint S = params.weight_blocks_expert_stride_bytes;
+    const uint S_scales = params.weight_scales_expert_stride_bytes;
+    const uint S_bias = params.bias_expert_stride_bytes;
+
+    const device char* wb0 = reinterpret_cast<const device char*>(weight_blocks);
+    const device char* sc0 = reinterpret_cast<const device char*>(weight_scales);
+    const device char* bi0 = reinterpret_cast<const device char*>(bias);
+
+    weight_blocks = reinterpret_cast<const device uint*>(wb0 + tg_id.z * S);
+    weight_scales = reinterpret_cast<const device uchar*>(sc0 + tg_id.z * S_scales);
+    bias = reinterpret_cast<const device bfloat*>(bi0 + tg_id.z * S_bias);
+
+    const uint sg_col_count = Bn / Sg_Bn;
+    const uint row_sg = sg_id / sg_col_count;
+    const uint col_sg = sg_id % sg_col_count;
+
+    const uint row_sg_offset = row_sg * Sg_Bm;
+    const uint col_sg_offset = col_sg * Sg_Bn;
+    // Declare threadgroup blocks.
+    threadgroup float lhs_block[Bm * Bk];
+    // rhs_block will hold the scaled fp32 weights.
+    threadgroup float rhs_block[Bn * Bk];
+
+    constexpr uint temp_result_size = (Sg_Bm / 8) * (Sg_Bn / 8);
+    // Create an array of simdgroup_float8x8 to hold temp results.
+    metal::simdgroup_float8x8 OutTiles[temp_result_size];
+    for (uint i = 0; i < temp_result_size; i++) {
+        OutTiles[i] = metal::make_filled_simdgroup_matrix<float, 8, 8>(0.0);
+    }
+    // Linear thread id within TG (we launch 1-D TGs)
+    const uint lin_tid = local_tid.x;
+
+    const uint thread_count_per_tg = threads_per_tg.x * threads_per_tg.y * threads_per_tg.z;
+    // Iterate over all Bk blocks.
+    for (uint k_offset = 0; k_offset < K; k_offset += Bk) {
+        constexpr uint lhs_row_stride = Bk;
+        constexpr uint lhs_vec_cols = Bk / 4;
+        constexpr uint lhs_vec_total = Bm * lhs_vec_cols;
+
+        const uint LHS_ITERS = ceil_div(lhs_vec_total, thread_count_per_tg);
+
+        for (uint t = 0; t < LHS_ITERS; ++t) {
+            const uint i = t * thread_count_per_tg + lin_tid;
+            if (i < lhs_vec_total) {
+                const uint r = i / lhs_vec_cols;
+                const uint c4 = i % lhs_vec_cols;
+
+                const uint gr = row_tg_offset + r;
+                const uint gc4 = (k_offset / 4) + c4;
+
+                threadgroup float4* dst4 =
+                    reinterpret_cast<threadgroup float4*>(lhs_block + r * lhs_row_stride + (c4 << 2));
+                if (gr < M) {
+                    const device float4* src4 =
+                        reinterpret_cast<const device float4*>(lhs + gr * K + (gc4 << 2));
+
+                    *dst4 = *src4;
+                } else {
+                    *dst4 = float4(0.0);
+                }
+            }
+        }
+
+        // Load weights with vector loads.
+        constexpr uint rhs_row_stride = Bk;
+        constexpr uint weights_per_elem = 8;
+        constexpr uint rhs_loads_per_col = Bk / weights_per_elem;
+        constexpr uint rhs_loads_total = Bn * rhs_loads_per_col;
+        const uint RHS_ITERS = ceil_div(rhs_loads_total, thread_count_per_tg);
+        // #pragma clang loop unroll(full)
+        for (uint t = 0; t < RHS_ITERS; ++t) {
+            const uint i = t * thread_count_per_tg + lin_tid;
+            if (i < rhs_loads_total) {
+                const uint r = i / rhs_loads_per_col;
+                const uint c = i % rhs_loads_per_col;
+
+                const uint gr = col_tg_offset + r;
+                const uint gc = (k_offset / weights_per_elem) + c;
+                const uint gc_scale = (k_offset / 32) + (c >> 2);
+
+                const uint wblock = weight_blocks[gr * (K / weights_per_elem) + gc];
+                const float scale =
+                    as_type<float>(static_cast<uint>(weight_scales[gr * (K / 32) + gc_scale]) << 23);
+
+                uint wblock0246 = (wblock + wblock);
+                uint wblock1357 = (wblock >> 3);
+                wblock0246 &= 0x1E1E1E1Eu;
+                wblock1357 &= 0x1E1E1E1Eu;
+
+                wblock0246 += 0x70707070u;
+                wblock1357 += 0x70707070u;
+                wblock0246 &= 0x8E8E8E8Eu;
+                wblock1357 &= 0x8E8E8E8Eu;
+
+                uint wblock26 = (wblock0246) & 0xFF00FF00u;
+                uint wblock04 = ((wblock0246 << 8)) & 0xFF00FF00u;
+                uint wblock37 = (wblock1357) & 0xFF00FF00u;
+                uint wblock15 = ((wblock1357 << 8)) & 0xFF00FF00u;
+
+                half4 wblock0426 = as_type<half4>(uint2(wblock04, wblock26));
+                half4 wblock1537 = as_type<half4>(uint2(wblock15, wblock37));
+
+                const float w0 = float(wblock0426.x) * scale;
+                const float w1 = float(wblock1537.x) * scale;
+                const float w2 = float(wblock0426.z) * scale;
+                const float w3 = float(wblock1537.z) * scale;
+                const float w4 = float(wblock0426.y) * scale;
+                const float w5 = float(wblock1537.y) * scale;
+                const float w6 = float(wblock0426.w) * scale;
+                const float w7 = float(wblock1537.w) * scale;
+                const uint rhs_offset = r * rhs_row_stride + c * 8;
+                rhs_block[rhs_offset] = w0;
+                rhs_block[rhs_offset + 1] = w1;
+                rhs_block[rhs_offset + 2] = w2;
+                rhs_block[rhs_offset + 3] = w3;
+                rhs_block[rhs_offset + 4] = w4;
+                rhs_block[rhs_offset + 5] = w5;
+                rhs_block[rhs_offset + 6] = w6;
+                rhs_block[rhs_offset + 7] = w7;
+            }
+        }
+        threadgroup_barrier(metal::mem_flags::mem_threadgroup);
+#pragma clang loop unroll(full)
+        for (uint k = 0; k < Bk; k += 8) {
+#pragma clang loop unroll(full)
+            for (uint m_subtile_ = 0; m_subtile_ < Sg_Bm; m_subtile_ += 8) {
+                const uint row_index_in_out_tile = m_subtile_ / 8;
+                metal::simdgroup_float8x8 lhs_frag;
+
+                simdgroup_load(lhs_frag, lhs_block, Bk, ulong2(k, m_subtile_ + row_sg_offset));
+#pragma clang loop unroll(full)
+                for (uint n_subtile_ = 0; n_subtile_ < Sg_Bn; n_subtile_ += 8) {
+                    const uint col_index_in_out_tile = n_subtile_ / 8;
+                    const uint current_index_out_tile =
+                        row_index_in_out_tile * (Sg_Bn / 8) + col_index_in_out_tile;
+                    metal::simdgroup_float8x8 rhs_frag;
+                    simdgroup_load(rhs_frag, rhs_block, Bk, ulong2(k, n_subtile_ + col_sg_offset), true);
+                    simdgroup_multiply_accumulate(OutTiles[current_index_out_tile], lhs_frag, rhs_frag,
+                        OutTiles[current_index_out_tile]);
+                }
+            }
+        }
+        threadgroup_barrier(metal::mem_flags::mem_threadgroup);
+    }
+
+    // Epilogue.
+    threadgroup float scratch[Bm * Bn];
+#pragma clang loop unroll(full)
+    for (uint n_subtile_ = 0; n_subtile_ < Sg_Bn; n_subtile_ += 8) {
+        const uint col_index_in_out_tile = n_subtile_ / 8;
+        const uint local_col_offset = col_sg_offset + n_subtile_;
+#pragma clang loop unroll(full)
+        for (uint m_subtile_ = 0; m_subtile_ < Sg_Bm; m_subtile_ += 8) {
+            const uint row_index_in_out_tile = m_subtile_ / 8;
+            const uint local_row_offset = row_sg_offset + m_subtile_;
+            const uint current_index_out_tile =
+                row_index_in_out_tile * (Sg_Bn / 8) + col_index_in_out_tile;
+            simdgroup_store(OutTiles[current_index_out_tile], scratch, Bn,
+                ulong2(local_col_offset, local_row_offset));
+        }
+    }
+    threadgroup float bias_tile[Bn];
+    for (uint c_local = local_tid.x; c_local < Bn; c_local += thread_count_per_tg) {
+        const uint c_global = col_tg_offset + c_local;
+        bias_tile[c_local] = (c_global < N) ? static_cast<float>(bias[c_global]) : 0.0f;
+    }
+
+    threadgroup_barrier(metal::mem_flags::mem_threadgroup);
+    for (uint idx = local_tid.x; idx < Bm * Bn; idx += thread_count_per_tg) {
+        const uint r = idx / Bn;
+        const uint c = idx % Bn;
+
+        const uint out_row = row_tg_offset + r;
+        const uint out_col = col_tg_offset + c;
+
+        if (out_row < M && out_col < N) {
+            float acc = scratch[idx] + bias_tile[c];
+            out[out_row * N + out_col] = acc;
+        }
+    }
+}

gptoss_kernels/source/random.metal

@@ -0,0 +1,97 @@
+#include <metal_integer>
+#include <metal_math>
+
+#include <internal/kernel-args.h>
+
+#pragma METAL fp math_mode(safe)
+#pragma METAL fp contract(off)
+
+
+inline static uint rng_squares32(ulong offset, ulong seed) {
+    const ulong y = offset * seed;
+    const ulong z = y + seed;
+
+    /* Round 1 */
+    ulong x = y * y + y;
+    x = metal::rotate(x, 32ul);
+
+    /* Round 2 */
+    x = x * x + z;
+    x = metal::rotate(x, 32ul);
+
+    /* Round 3 */
+    x = x * x + y;
+    x = metal::rotate(x, 32ul);
+
+    /* Round 4 */
+    x = x * x + z;
+    return as_type<uint2>(x).y;
+}
+
+kernel void gptoss_u32_fill_random(
+    constant gptoss_u32_fill_random_args& args [[ buffer(0) ]],
+    device uint* output [[ buffer(1) ]],
+    uint gid [[threadgroup_position_in_grid]],
+    uint tid [[thread_position_in_threadgroup]],
+    uint threadgroup_size [[ threads_per_threadgroup ]])
+{
+    const ulong num_vecs_per_threadgroup = args.num_vecs_per_threadgroup;
+    const ulong threadgroup_start = gid * num_vecs_per_threadgroup;
+    const ulong threadgroup_end = metal::min(threadgroup_start + num_vecs_per_threadgroup, args.num_vecs);
+    const ulong thread_start = threadgroup_start + tid;
+    uint num_iter = static_cast<uint>((threadgroup_end - thread_start + (threadgroup_size - 1)) / threadgroup_size);
+
+    output += thread_start;
+    ulong offset = args.offset + thread_start;
+    for (; num_iter != 0; num_iter--) {
+        *output = rng_squares32(offset, args.seed);
+        output += threadgroup_size;
+        offset += threadgroup_size;
+    }
+}
+
+kernel void gptoss_f32_fill_random(
+    constant gptoss_f32_fill_random_args& args [[ buffer(0) ]],
+    device float* output [[ buffer(1) ]],
+    uint gid [[threadgroup_position_in_grid]],
+    uint tid [[thread_position_in_threadgroup]],
+    uint threadgroup_size [[ threads_per_threadgroup ]])
+{
+    const ulong num_vecs_per_threadgroup = args.num_vecs_per_threadgroup;
+    const ulong threadgroup_start = gid * num_vecs_per_threadgroup;
+    const ulong threadgroup_end = metal::min(threadgroup_start + num_vecs_per_threadgroup, args.num_vecs);
+    const ulong thread_start = threadgroup_start + tid;
+    uint num_iter = static_cast<uint>((threadgroup_end - thread_start + (threadgroup_size - 1)) / threadgroup_size);
+
+    output += thread_start;
+    ulong offset = args.offset + thread_start;
+    for (; num_iter != 0; num_iter--) {
+        const uint word = rng_squares32(offset, args.seed);
+        *output = metal::fma(static_cast<float>(as_type<int>(word)), args.scale, args.bias);
+        output += threadgroup_size;
+        offset += threadgroup_size;
+    }
+}
+
+kernel void gptoss_bf16_fill_random(
+    constant gptoss_f32_fill_random_args& args [[ buffer(0) ]],
+    device bfloat* output [[ buffer(1) ]],
+    uint gid [[threadgroup_position_in_grid]],
+    uint tid [[thread_position_in_threadgroup]],
+    uint threadgroup_size [[ threads_per_threadgroup ]])
+{
+    const ulong num_vecs_per_threadgroup = args.num_vecs_per_threadgroup;
+    const ulong threadgroup_start = gid * num_vecs_per_threadgroup;
+    const ulong threadgroup_end = metal::min(threadgroup_start + num_vecs_per_threadgroup, args.num_vecs);
+    const ulong thread_start = threadgroup_start + tid;
+    uint num_iter = static_cast<uint>((threadgroup_end - thread_start + (threadgroup_size - 1)) / threadgroup_size);
+
+    output += thread_start;
+    ulong offset = args.offset + thread_start;
+    for (; num_iter != 0; num_iter--) {
+        const uint word = rng_squares32(offset, args.seed);
+        *output = static_cast<bfloat>(metal::fma(static_cast<float>(as_type<int>(word)), args.scale, args.bias));
+        output += threadgroup_size;
+        offset += threadgroup_size;
+    }
+}

gptoss_kernels/source/rmsnorm.metal

@@ -0,0 +1,58 @@
+#include <metal_compute>
+#include <metal_math>
+#include <metal_simdgroup>
+
+#include <internal/kernel-args.h>
+
+#pragma METAL fp math_mode(safe)
+#pragma METAL fp contract(off)
+
+
+[[max_total_threads_per_threadgroup(1024)]]
+kernel void gptoss_f32_bf16w_rmsnorm(
+    constant gptoss_rmsnorm_args& args [[ buffer(0) ]],
+    const device float4* input [[ buffer(1) ]],
+    const device bfloat4* weights [[ buffer(2) ]],
+    device float4* output [[ buffer(3) ]],
+    const device gptoss_control* control [[ buffer(4) ]],
+    uint gid [[threadgroup_position_in_grid]],
+    uint tid [[thread_position_in_threadgroup]],
+    uint threadgroup_size [[ threads_per_threadgroup ]])
+{
+    const uint simdgroup_size = 32;
+    threadgroup float threadgroup_buffer[32];
+    if (control->abort != 0) {
+        return;
+    }
+
+    input += gid * args.num_vecs;
+    output += gid * args.num_vecs;
+
+    float4 sumsq4 = 0.0f;
+    for (uint i = tid; i < args.num_vecs; i += threadgroup_size) {
+        const float4 val = input[i];
+        sumsq4 = metal::fma(val, val, sumsq4);
+    }
+
+    // Tree-reduce sumsq within thread, then all-reduce within threadgroup.
+    const float2 sumsq2 = sumsq4.xy + sumsq4.zw;
+    float sumsq = sumsq2.x + sumsq2.y;
+    // Warning: this all-reduce works only for simdgroup of 32 threads and threadgroup of 32*32=1024 threads.
+    sumsq = metal::simd_sum(sumsq);
+    if (metal::simd_is_first()) {
+        const uint simdgroup_idx = tid / simdgroup_size;
+        threadgroup_buffer[simdgroup_idx] = sumsq;
+    }
+    metal::threadgroup_barrier(metal::mem_flags::mem_threadgroup);
+    const uint simdgroup_tid = tid % simdgroup_size;
+    sumsq = threadgroup_buffer[simdgroup_tid];
+    sumsq = metal::simd_sum(sumsq);
+
+    const float avgsq = sumsq / args.num_channels;
+    const float scale = metal::precise::rsqrt(avgsq + args.epsilon);
+    for (uint i = tid; i < args.num_vecs; i += threadgroup_size) {
+        const float4 val = input[i] * scale;
+        const float4 weight_val = static_cast<float4>(weights[i]);
+        output[i] = val * weight_val;
+    }
+}

gptoss_kernels/source/rope.metal

@@ -0,0 +1,43 @@
+#include <metal_common>
+#include <metal_math>
+
+#include <internal/kernel-args.h>
+
+#pragma METAL fp math_mode(safe)
+#pragma METAL fp contract(off)
+
+
+// Each thread handles 2 head elements.
+// Each simdgroup handles one head (64 head elements).
+
+kernel void gptoss_f32_rope(
+    constant gptoss_rope_args& args [[ buffer(0) ]],
+    device float2* activations [[ buffer(1) ]],
+    const device gptoss_control* control [[ buffer(2) ]],
+    uint2 gid [[thread_position_in_grid]])
+{
+    const uint num_head_dims = 64;
+    if (control->abort != 0) {
+        return;
+    }
+
+    const float dim_idx = static_cast<float>(gid.x % (num_head_dims / 2));
+    const uint token_idx = args.token_offset + gid.y;
+    activations += gid.y * args.token_stride + gid.x;
+
+    const float2 input_vals = *activations;
+    const float inv_extrapolation_freq = metal::precise::exp(dim_idx * args.freq_scale);
+    const float inv_interpolation_freq = inv_extrapolation_freq * args.interpolation_scale;
+    const float alpha = metal::saturate(metal::fma(dim_idx, args.yarn_scale, args.yarn_offset));
+    const float inv_freq = metal::mix(inv_extrapolation_freq, inv_interpolation_freq, alpha);
+
+    const float phi = static_cast<float>(token_idx) * inv_freq;
+    const float yarn_multiplier = args.yarn_multiplier;
+    float cosphi;
+    const float sinphi = metal::precise::sincos(phi, cosphi) * yarn_multiplier;
+    cosphi *= yarn_multiplier;
+
+    const float output_re = input_vals.x * cosphi - input_vals.y * sinphi;
+    const float output_im = input_vals.x * sinphi + input_vals.y * cosphi;
+    *activations = (float2) { output_re, output_im };
+}

gptoss_kernels/source/sample.metal

@@ -0,0 +1,209 @@
+#include <metal_compute>
+#include <metal_integer>
+#include <metal_math>
+#include <metal_simdgroup>
+
+#include <internal/kernel-args.h>
+
+#pragma METAL fp math_mode(safe)
+#pragma METAL fp contract(off)
+
+
+inline static uint rng_squares32(ulong offset, ulong seed) {
+    const ulong y = offset * seed;
+    const ulong z = y + seed;
+
+    /* Round 1 */
+    ulong x = y * y + y;
+    x = metal::rotate(x, 32ul);
+
+    /* Round 2 */
+    x = x * x + z;
+    x = metal::rotate(x, 32ul);
+
+    /* Round 3 */
+    x = x * x + y;
+    x = metal::rotate(x, 32ul);
+
+    /* Round 4 */
+    x = x * x + z;
+    return as_type<uint2>(x).y;
+}
+
+kernel void gptoss_f32_softmax(
+    constant gptoss_softmax_args& args [[ buffer(0) ]],
+    const device float* score [[ buffer(1) ]],
+    const device uint2* argmax [[ buffer(2) ]],
+    device float* prob [[ buffer(3) ]],
+    device float* sum [[ buffer(4) ]],
+    const device gptoss_control* control [[ buffer(5) ]],
+    uint tidx [[thread_index_in_threadgroup]],
+    uint2 gid [[threadgroup_position_in_grid]],
+    uint2 threadgroup_size [[threads_per_threadgroup]],
+    uint simdgroup_tid [[thread_index_in_simdgroup]],
+    uint simdgroup_idx [[simdgroup_index_in_threadgroup]],
+    uint num_simdgroups [[simdgroups_per_threadgroup]])
+{
+    threadgroup float threadgroup_sumexp[32];
+    if (control->abort != 0) {
+        return;
+    }
+
+    score += gid.y * args.num_vecs + gid.x * args.num_vecs_per_threadgroup;
+    prob += gid.y * args.num_vecs + gid.x * args.num_vecs_per_threadgroup;
+    sum += gid.y * args.max_threadgroups;
+
+    uint max_bits = argmax[gid.y].y;
+    if (static_cast<int>(max_bits) >= 0) {
+        max_bits ^= 0x7FFFFFFFu;
+    }
+    const float max_val = as_type<float>(max_bits);
+    float sum_exp = 0.0f;
+    const uint num_vecs_per_threadgroup = metal::min(args.num_vecs - gid.x * args.num_vecs_per_threadgroup, args.num_vecs_per_threadgroup);
+    for (uint i = tidx; i < num_vecs_per_threadgroup; i += threadgroup_size.x) {
+        const float score_val = score[i];
+        const float prob_val = metal::precise::exp((score_val - max_val) * args.temperature);
+        prob[i] = prob_val;
+        sum_exp += prob_val;
+    }
+    sum_exp = metal::simd_sum(sum_exp);
+    if (metal::simd_is_first()) {
+        threadgroup_sumexp[simdgroup_idx] = sum_exp;
+    }
+    metal::threadgroup_barrier(metal::mem_flags::mem_threadgroup);
+    if (simdgroup_idx == 0) {
+        // Sum-Reduce threadgroup_sumexp
+        sum_exp = 0.0f;
+        if (simdgroup_tid < num_simdgroups) {
+            sum_exp = threadgroup_sumexp[simdgroup_tid];
+        }
+        sum_exp = metal::simd_sum(sum_exp);
+        if (metal::simd_is_first()) {
+            sum[gid.x] = sum_exp;
+        }
+    }
+}
+
+[[max_total_threads_per_threadgroup(1024)]]
+kernel void gptoss_f32_sample(
+    constant gptoss_sample_args& args [[ buffer(0) ]],
+    device const float* prob [[ buffer(1) ]],
+    device const float* sum [[ buffer(2) ]],
+    device uint* prediction [[ buffer(3) ]],
+    device gptoss_control* control [[ buffer(4) ]],
+    uint tid [[thread_position_in_threadgroup]],
+    uint threadgroup_size [[threads_per_threadgroup]],
+    uint simdgroup_tid [[thread_index_in_simdgroup]],
+    uint simdgroup_idx [[simdgroup_index_in_threadgroup]],
+    uint num_simdgroups [[simdgroups_per_threadgroup]])
+{
+    threadgroup float threadgroup_sum_buffer[32];
+    threadgroup uint threadgroup_idx_buffer[32];
+    threadgroup float threadgroup_cumsum_buffer[32];
+    if (control->abort != 0) {
+        return;
+    }
+
+    const uint sample_word = rng_squares32(args.rng_offset, args.rng_seed);
+    float sample_cdf = static_cast<float>(sample_word & 0x00FFFFFFu) * 0x1.0p-24f;
+
+    float cumsum = 0.0f;
+    if (tid < args.num_blocks) {
+        cumsum = sum[tid];
+    }
+    cumsum = metal::simd_prefix_inclusive_sum(cumsum);
+    if (simdgroup_tid == 31) {
+        threadgroup_sum_buffer[simdgroup_idx] = cumsum;
+    }
+    metal::threadgroup_barrier(metal::mem_flags::mem_threadgroup);
+    float threadgroup_cumsum = 0.0f, threadgroup_sum = 0.0f;
+    if (simdgroup_tid < num_simdgroups) {
+        threadgroup_sum = threadgroup_sum_buffer[simdgroup_tid];
+        if (simdgroup_tid < simdgroup_idx) {
+            threadgroup_cumsum = threadgroup_sum;
+        }
+    }
+    threadgroup_sum = metal::simd_sum(threadgroup_sum);
+    cumsum += metal::simd_sum(threadgroup_cumsum);
+
+    sample_cdf *= threadgroup_sum;
+    sample_cdf = metal::max(sample_cdf, 0x1.0p-149f);
+
+    // Find the block: the smallest tid where sample_cdf >= s
+    uint block_idx = args.num_blocks;
+    float block_sum = cumsum;
+    if (tid >= args.num_blocks - 1) {
+        block_idx = args.num_blocks - 1;
+        block_sum = 0.0f;
+    } else if (cumsum >= sample_cdf) {
+        block_idx = tid;
+        block_sum = 0.0f;
+    }
+    block_idx = metal::simd_min(block_idx);
+    block_sum = metal::simd_max(block_sum);
+    if (simdgroup_tid == 0) {
+        threadgroup_idx_buffer[simdgroup_idx] = block_idx;
+        threadgroup_cumsum_buffer[simdgroup_idx] = block_sum;
+    }
+    metal::threadgroup_barrier(metal::mem_flags::mem_threadgroup);
+    if (simdgroup_tid < num_simdgroups) {
+        block_idx = threadgroup_idx_buffer[simdgroup_tid];
+        block_sum = threadgroup_cumsum_buffer[simdgroup_tid];
+    }
+    block_idx = metal::simd_min(block_idx);
+    block_sum = metal::simd_max(block_sum);
+
+    const uint block_start = args.num_dims_per_block * block_idx;
+    const uint block_end = metal::min(block_start + args.num_dims_per_block, args.num_dims);
+    uint offset = block_start + tid;
+    float accumulated_sum = block_sum;
+    uint sample_idx;
+
+    // This loop must be threadgroup-uniform.
+    do {
+        // Find the token: the smallest tid where sample_cdf >= s
+        float cumsum = 0.0f;
+        if (offset < block_end) {
+            cumsum = prob[offset];
+        }
+        cumsum = metal::simd_prefix_inclusive_sum(cumsum);
+        if (simdgroup_tid == 31) {
+            threadgroup_sum_buffer[simdgroup_idx] = cumsum;
+        }
+        metal::threadgroup_barrier(metal::mem_flags::mem_threadgroup);
+        float threadgroup_cumsum = 0.0f, threadgroup_sum = 0.0f;
+        if (simdgroup_tid < num_simdgroups) {
+            threadgroup_sum = threadgroup_sum_buffer[simdgroup_tid];
+            if (simdgroup_tid < simdgroup_idx) {
+                threadgroup_cumsum = threadgroup_sum;
+            }
+        }
+        threadgroup_sum = metal::simd_sum(threadgroup_sum);
+        cumsum += metal::simd_sum(threadgroup_cumsum);
+        cumsum += accumulated_sum;
+
+        sample_idx = block_end;
+        if (offset >= block_end) {
+            // Trigger loop exit, with the last token in the block being sampled if no other candidate was found.
+            sample_idx = block_end - 1;
+        } else if (cumsum >= sample_cdf) {
+            sample_idx = offset;
+        }
+        sample_idx = metal::simd_min(sample_idx);
+        if (simdgroup_tid == 0) {
+            threadgroup_idx_buffer[simdgroup_idx] = sample_idx;
+        }
+        metal::threadgroup_barrier(metal::mem_flags::mem_threadgroup);
+        if (simdgroup_tid < num_simdgroups) {
+            sample_idx = threadgroup_idx_buffer[simdgroup_tid];
+        }
+        sample_idx = metal::simd_min(sample_idx);
+
+        offset += threadgroup_size;
+        accumulated_sum += threadgroup_sum;
+    } while (sample_idx == block_end);
+
+    if (tid == 0) {
+        *prediction = sample_idx;
+    }
+}

gptoss_kernels/source/scatter.metal

@@ -0,0 +1,65 @@
+#include <internal/kernel-args.h>
+#include <metal_integer>
+#include <metal_math>
+#include <metal_stdlib>
+
+// TODO(ibrahim): This is not optimal as each thread only scatters a single float4. To amortize the
+// cost of reading the expert id and offset for a token, we should let each thread scatter several
+// float4s.
+kernel void gptoss_f32_scatter_e4(
+    constant gptoss_scatter_args& args [[ buffer(0) ]],
+    const device float* in [[ buffer(1) ]],
+    const device gptoss_expert_prediction* __restrict__ expert_predictions [[ buffer(2) ]],
+    const device uint* __restrict__ expert_offsets [[ buffer(3) ]],
+    const device uint* __restrict__ intra_expert_offsets [[ buffer(4) ]],
+    device float* out [[ buffer(5) ]],
+    uint3 gid [[thread_position_in_grid]]) 
+{
+    const uint total_tokens = args.tokens;
+    const uint active_experts_per_token = args.active_experts_per_token;
+    const uint embedding_dim = args.token_stride;
+    assert(embedding_dim % 4 == 0);
+    // Hard coded to top4 for now.
+    assert(active_experts_per_token == 4);
+    const uint row_in = gid.y;
+    if (row_in >= total_tokens) {
+        return;
+    }
+    // Consecutive threads in a tg read consecutive columns of the input.
+    const uint col_in_vec4 = gid.x;
+    const uint col_in = col_in_vec4 * 4;
+    if (col_in >= embedding_dim) {
+        return;
+    }
+    // Pointer to the piece of the input that we will copy to the top4 experts.
+    const device float4* src4 =
+        reinterpret_cast<const device float4*>(in + row_in * embedding_dim + col_in);
+
+    // Get the 4 destinations -- 4 experts.
+    const uint base = row_in * active_experts_per_token;
+    const uint expert0_id = expert_predictions[base].expert_id;
+    const uint expert1_id = expert_predictions[base + 1].expert_id;
+    const uint expert2_id = expert_predictions[base + 2].expert_id;
+    const uint expert3_id = expert_predictions[base + 3].expert_id;
+    const uint expert0_offset = expert_offsets[expert0_id];
+    const uint expert1_offset = expert_offsets[expert1_id];
+    const uint expert2_offset = expert_offsets[expert2_id];
+    const uint expert3_offset = expert_offsets[expert3_id];
+    const uint expert0_intra_expert_offset = intra_expert_offsets[base];
+    const uint expert1_intra_expert_offset = intra_expert_offsets[base + 1];
+    const uint expert2_intra_expert_offset = intra_expert_offsets[base + 2];
+    const uint expert3_intra_expert_offset = intra_expert_offsets[base + 3];
+    device float4* dst4_0 = reinterpret_cast<device float4*>(
+        out + (expert0_offset + expert0_intra_expert_offset) * embedding_dim + col_in);
+    device float4* dst4_1 = reinterpret_cast<device float4*>(
+        out + (expert1_offset + expert1_intra_expert_offset) * embedding_dim + col_in);
+    device float4* dst4_2 = reinterpret_cast<device float4*>(
+        out + (expert2_offset + expert2_intra_expert_offset) * embedding_dim + col_in);
+    device float4* dst4_3 = reinterpret_cast<device float4*>(
+        out + (expert3_offset + expert3_intra_expert_offset) * embedding_dim + col_in);
+    const float4 data = *src4;
+    *dst4_0 = data;
+    *dst4_1 = data;
+    *dst4_2 = data;
+    *dst4_3 = data;
+}

gptoss_kernels/source/sdpa.metal

@@ -0,0 +1,293 @@
+#include <metal_geometric>
+#include <metal_integer>
+#include <metal_math>
+#include <metal_compute>
+#include <metal_simdgroup>
+
+#include <internal/kernel-args.h>
+
+#pragma METAL fp math_mode(safe)
+#pragma METAL fp contract(off)
+
+// Each threadgroup handles 8 Q heads / 1 KV head for 1 token
+
+kernel void gptoss_f32_sdpa_q8_d64(
+    constant gptoss_sdpa_args& args [[ buffer(0) ]],
+    const device float* q [[ buffer(1) ]],
+    const device float* kv [[ buffer(2) ]],
+    const device bfloat* s [[ buffer(3) ]],
+    device float* output [[ buffer(4) ]],
+    const device gptoss_control* control [[ buffer(6) ]],
+    threadgroup void* threadgroup_buffer [[ threadgroup(0) ]],
+    uint2 gid [[threadgroup_position_in_grid]],
+    uint2 tid [[thread_position_in_threadgroup]],
+    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_q_heads = 64;
+    const uint head_dim = 64;
+    const uint qmul = 8;
+
+    const uint token_stride = 2 * head_dim;
+
+    const uint qt = gid.x;  // Q token index
+    const uint h = gid.y;   // KV head index
+
+    q += qt * args.qkv_dim + h * (qmul * head_dim);
+    kv += h * args.kv_stride;
+    output += qt * (num_q_heads * head_dim) + h * (qmul * head_dim);
+
+    float m0 = static_cast<float>(s[h * qmul + 0]);
+    float m1 = static_cast<float>(s[h * qmul + 1]);
+    float m2 = static_cast<float>(s[h * qmul + 2]);
+    float m3 = static_cast<float>(s[h * qmul + 3]);
+    float m4 = static_cast<float>(s[h * qmul + 4]);
+    float m5 = static_cast<float>(s[h * qmul + 5]);
+    float m6 = static_cast<float>(s[h * qmul + 6]);
+    float m7 = static_cast<float>(s[h * qmul + 7]);
+
+    float l0 = simdgroup_idx == 0 ? 1.0f : 0.0f;
+    float l1 = simdgroup_idx == 0 ? 1.0f : 0.0f;
+    float l2 = simdgroup_idx == 0 ? 1.0f : 0.0f;
+    float l3 = simdgroup_idx == 0 ? 1.0f : 0.0f;
+    float l4 = simdgroup_idx == 0 ? 1.0f : 0.0f;
+    float l5 = simdgroup_idx == 0 ? 1.0f : 0.0f;
+    float l6 = simdgroup_idx == 0 ? 1.0f : 0.0f;
+    float l7 = simdgroup_idx == 0 ? 1.0f : 0.0f;
+
+    float2 out0 = 0.0f;
+    float2 out1 = 0.0f;
+    float2 out2 = 0.0f;
+    float2 out3 = 0.0f;
+    float2 out4 = 0.0f;
+    float2 out5 = 0.0f;
+    float2 out6 = 0.0f;
+    float2 out7 = 0.0f;
+
+    float2 q0 = reinterpret_cast<const device float2*>(q + 0 * head_dim)[simdgroup_tid];
+    float2 q1 = reinterpret_cast<const device float2*>(q + 1 * head_dim)[simdgroup_tid];
+    float2 q2 = reinterpret_cast<const device float2*>(q + 2 * head_dim)[simdgroup_tid];
+    float2 q3 = reinterpret_cast<const device float2*>(q + 3 * head_dim)[simdgroup_tid];
+    float2 q4 = reinterpret_cast<const device float2*>(q + 4 * head_dim)[simdgroup_tid];
+    float2 q5 = reinterpret_cast<const device float2*>(q + 5 * head_dim)[simdgroup_tid];
+    float2 q6 = reinterpret_cast<const device float2*>(q + 6 * head_dim)[simdgroup_tid];
+    float2 q7 = reinterpret_cast<const device float2*>(q + 7 * head_dim)[simdgroup_tid];
+
+    const uint kt_end = qt + args.num_kv_tokens + 1;
+    const uint kt_start = metal::subsat(kt_end, args.window) + simdgroup_idx;
+    kv += token_stride * kt_start;
+    for (uint kt = kt_start; kt < kt_end; kt += num_simdgroups) {
+        const float2 kval = reinterpret_cast<const device float2*>(kv)[simdgroup_tid];
+
+        float qk0 = metal::dot(q0, kval);
+        float qk1 = metal::dot(q1, kval);
+        float qk2 = metal::dot(q2, kval);
+        float qk3 = metal::dot(q3, kval);
+        float qk4 = metal::dot(q4, kval);
+        float qk5 = metal::dot(q5, kval);
+        float qk6 = metal::dot(q6, kval);
+        float qk7 = metal::dot(q7, kval);
+
+        qk0 = metal::simd_sum(qk0);
+        qk1 = metal::simd_sum(qk1);
+        qk2 = metal::simd_sum(qk2);
+        qk3 = metal::simd_sum(qk3);
+        qk4 = metal::simd_sum(qk4);
+        qk5 = metal::simd_sum(qk5);
+        qk6 = metal::simd_sum(qk6);
+        qk7 = metal::simd_sum(qk7);
+
+        const float new_m0 = metal::max(m0, qk0);
+        const float new_m1 = metal::max(m1, qk1);
+        const float new_m2 = metal::max(m2, qk2);
+        const float new_m3 = metal::max(m3, qk3);
+        const float new_m4 = metal::max(m4, qk4);
+        const float new_m5 = metal::max(m5, qk5);
+        const float new_m6 = metal::max(m6, qk6);
+        const float new_m7 = metal::max(m7, qk7);
+
+        const float alpha0 = metal::fast::exp(m0 - new_m0);
+        const float alpha1 = metal::fast::exp(m1 - new_m1);
+        const float alpha2 = metal::fast::exp(m2 - new_m2);
+        const float alpha3 = metal::fast::exp(m3 - new_m3);
+        const float alpha4 = metal::fast::exp(m4 - new_m4);
+        const float alpha5 = metal::fast::exp(m5 - new_m5);
+        const float alpha6 = metal::fast::exp(m6 - new_m6);
+        const float alpha7 = metal::fast::exp(m7 - new_m7);
+
+        qk0 = metal::fast::exp(qk0 - new_m0);
+        qk1 = metal::fast::exp(qk1 - new_m1);
+        qk2 = metal::fast::exp(qk2 - new_m2);
+        qk3 = metal::fast::exp(qk3 - new_m3);
+        qk4 = metal::fast::exp(qk4 - new_m4);
+        qk5 = metal::fast::exp(qk5 - new_m5);
+        qk6 = metal::fast::exp(qk6 - new_m6);
+        qk7 = metal::fast::exp(qk7 - new_m7);
+
+        l0 = metal::fma(l0, alpha0, qk0);
+        l1 = metal::fma(l1, alpha1, qk1);
+        l2 = metal::fma(l2, alpha2, qk2);
+        l3 = metal::fma(l3, alpha3, qk3);
+        l4 = metal::fma(l4, alpha4, qk4);
+        l5 = metal::fma(l5, alpha5, qk5);
+        l6 = metal::fma(l6, alpha6, qk6);
+        l7 = metal::fma(l7, alpha7, qk7);
+
+        m0 = new_m0;
+        m1 = new_m1;
+        m2 = new_m2;
+        m3 = new_m3;
+        m4 = new_m4;
+        m5 = new_m5;
+        m6 = new_m6;
+        m7 = new_m7;
+
+        const float2 vval = reinterpret_cast<const device float2*>(kv + head_dim)[simdgroup_tid];
+        kv += token_stride * num_simdgroups;
+        out0 = metal::fma(vval, qk0, out0 * alpha0);
+        out1 = metal::fma(vval, qk1, out1 * alpha1);
+        out2 = metal::fma(vval, qk2, out2 * alpha2);
+        out3 = metal::fma(vval, qk3, out3 * alpha3);
+        out4 = metal::fma(vval, qk4, out4 * alpha4);
+        out5 = metal::fma(vval, qk5, out5 * alpha5);
+        out6 = metal::fma(vval, qk6, out6 * alpha6);
+        out7 = metal::fma(vval, qk7, out7 * alpha7);
+    }
+    if (num_simdgroups > 1) {
+        if (metal::simd_is_first()) {
+            static_cast<threadgroup float*>(threadgroup_buffer)[0 * num_simdgroups + simdgroup_idx] = m0;
+            static_cast<threadgroup float*>(threadgroup_buffer)[1 * num_simdgroups + simdgroup_idx] = m1;
+            static_cast<threadgroup float*>(threadgroup_buffer)[2 * num_simdgroups + simdgroup_idx] = m2;
+            static_cast<threadgroup float*>(threadgroup_buffer)[3 * num_simdgroups + simdgroup_idx] = m3;
+            static_cast<threadgroup float*>(threadgroup_buffer)[4 * num_simdgroups + simdgroup_idx] = m4;
+            static_cast<threadgroup float*>(threadgroup_buffer)[5 * num_simdgroups + simdgroup_idx] = m5;
+            static_cast<threadgroup float*>(threadgroup_buffer)[6 * num_simdgroups + simdgroup_idx] = m6;
+            static_cast<threadgroup float*>(threadgroup_buffer)[7 * num_simdgroups + simdgroup_idx] = m7;
+
+            static_cast<threadgroup float*>(threadgroup_buffer)[ 8 * num_simdgroups + simdgroup_idx] = l0;
+            static_cast<threadgroup float*>(threadgroup_buffer)[ 9 * num_simdgroups + simdgroup_idx] = l1;
+            static_cast<threadgroup float*>(threadgroup_buffer)[10 * num_simdgroups + simdgroup_idx] = l2;
+            static_cast<threadgroup float*>(threadgroup_buffer)[11 * num_simdgroups + simdgroup_idx] = l3;
+            static_cast<threadgroup float*>(threadgroup_buffer)[12 * num_simdgroups + simdgroup_idx] = l4;
+            static_cast<threadgroup float*>(threadgroup_buffer)[13 * num_simdgroups + simdgroup_idx] = l5;
+            static_cast<threadgroup float*>(threadgroup_buffer)[14 * num_simdgroups + simdgroup_idx] = l6;
+            static_cast<threadgroup float*>(threadgroup_buffer)[15 * num_simdgroups + simdgroup_idx] = l7;
+        }
+        metal::threadgroup_barrier(metal::mem_flags::mem_threadgroup);
+        // Note: simdgroup refers not to the thread's current simdgroup, but to one with simdgroup_idx == thread's simdgroup_tid.
+        float simdgroup_m0 = m0;
+        float simdgroup_m1 = m1;
+        float simdgroup_m2 = m2;
+        float simdgroup_m3 = m3;
+        float simdgroup_m4 = m4;
+        float simdgroup_m5 = m5;
+        float simdgroup_m6 = m6;
+        float simdgroup_m7 = m7;
+        if (simdgroup_tid < num_simdgroups) {
+            simdgroup_m0 = static_cast<const threadgroup float*>(threadgroup_buffer)[0 * num_simdgroups + simdgroup_tid];
+            simdgroup_m1 = static_cast<const threadgroup float*>(threadgroup_buffer)[1 * num_simdgroups + simdgroup_tid];
+            simdgroup_m2 = static_cast<const threadgroup float*>(threadgroup_buffer)[2 * num_simdgroups + simdgroup_tid];
+            simdgroup_m3 = static_cast<const threadgroup float*>(threadgroup_buffer)[3 * num_simdgroups + simdgroup_tid];
+            simdgroup_m4 = static_cast<const threadgroup float*>(threadgroup_buffer)[4 * num_simdgroups + simdgroup_tid];
+            simdgroup_m5 = static_cast<const threadgroup float*>(threadgroup_buffer)[5 * num_simdgroups + simdgroup_tid];
+            simdgroup_m6 = static_cast<const threadgroup float*>(threadgroup_buffer)[6 * num_simdgroups + simdgroup_tid];
+            simdgroup_m7 = static_cast<const threadgroup float*>(threadgroup_buffer)[7 * num_simdgroups + simdgroup_tid];
+        }
+
+        const float threadgroup_m0 = metal::simd_max(simdgroup_m0);
+        const float threadgroup_m1 = metal::simd_max(simdgroup_m1);
+        const float threadgroup_m2 = metal::simd_max(simdgroup_m2);
+        const float threadgroup_m3 = metal::simd_max(simdgroup_m3);
+        const float threadgroup_m4 = metal::simd_max(simdgroup_m4);
+        const float threadgroup_m5 = metal::simd_max(simdgroup_m5);
+        const float threadgroup_m6 = metal::simd_max(simdgroup_m6);
+        const float threadgroup_m7 = metal::simd_max(simdgroup_m7);
+
+        out0 *= metal::fast::exp(m0 - threadgroup_m0);
+        out1 *= metal::fast::exp(m1 - threadgroup_m1);
+        out2 *= metal::fast::exp(m2 - threadgroup_m2);
+        out3 *= metal::fast::exp(m3 - threadgroup_m3);
+        out4 *= metal::fast::exp(m4 - threadgroup_m4);
+        out5 *= metal::fast::exp(m5 - threadgroup_m5);
+        out6 *= metal::fast::exp(m6 - threadgroup_m6);
+        out7 *= metal::fast::exp(m7 - threadgroup_m7);
+
+        if (simdgroup_idx == 0) {
+            l0 = 0.0f;
+            l1 = 0.0f;
+            l2 = 0.0f;
+            l3 = 0.0f;
+            l4 = 0.0f;
+            l5 = 0.0f;
+            l6 = 0.0f;
+            l7 = 0.0f;
+            if (simdgroup_tid < num_simdgroups) {
+                l0 = static_cast<const threadgroup float*>(threadgroup_buffer)[ 8 * num_simdgroups + simdgroup_tid];
+                l1 = static_cast<const threadgroup float*>(threadgroup_buffer)[ 9 * num_simdgroups + simdgroup_tid];
+                l2 = static_cast<const threadgroup float*>(threadgroup_buffer)[10 * num_simdgroups + simdgroup_tid];
+                l3 = static_cast<const threadgroup float*>(threadgroup_buffer)[11 * num_simdgroups + simdgroup_tid];
+                l4 = static_cast<const threadgroup float*>(threadgroup_buffer)[12 * num_simdgroups + simdgroup_tid];
+                l5 = static_cast<const threadgroup float*>(threadgroup_buffer)[13 * num_simdgroups + simdgroup_tid];
+                l6 = static_cast<const threadgroup float*>(threadgroup_buffer)[14 * num_simdgroups + simdgroup_tid];
+                l7 = static_cast<const threadgroup float*>(threadgroup_buffer)[15 * num_simdgroups + simdgroup_tid];
+            }
+
+            l0 = metal::simd_sum(l0 * metal::fast::exp(simdgroup_m0 - threadgroup_m0));
+            l1 = metal::simd_sum(l1 * metal::fast::exp(simdgroup_m1 - threadgroup_m1));
+            l2 = metal::simd_sum(l2 * metal::fast::exp(simdgroup_m2 - threadgroup_m2));
+            l3 = metal::simd_sum(l3 * metal::fast::exp(simdgroup_m3 - threadgroup_m3));
+            l4 = metal::simd_sum(l4 * metal::fast::exp(simdgroup_m4 - threadgroup_m4));
+            l5 = metal::simd_sum(l5 * metal::fast::exp(simdgroup_m5 - threadgroup_m5));
+            l6 = metal::simd_sum(l6 * metal::fast::exp(simdgroup_m6 - threadgroup_m6));
+            l7 = metal::simd_sum(l7 * metal::fast::exp(simdgroup_m7 - threadgroup_m7));
+        }
+
+        uint num_threads = num_simdgroups * simdgroup_size;
+        do {
+            const uint num_smem_threads = (num_threads / 2) & -simdgroup_size;
+            const uint num_half_threads = num_threads - num_smem_threads;
+
+            metal::threadgroup_barrier(metal::mem_flags::mem_threadgroup);
+            const uint smem_tid = tid.x - num_half_threads;
+            if (smem_tid < num_smem_threads) {
+                static_cast<threadgroup float2*>(threadgroup_buffer)[num_smem_threads * 0 + smem_tid] = out0;
+                static_cast<threadgroup float2*>(threadgroup_buffer)[num_smem_threads * 1 + smem_tid] = out1;
+                static_cast<threadgroup float2*>(threadgroup_buffer)[num_smem_threads * 2 + smem_tid] = out2;
+                static_cast<threadgroup float2*>(threadgroup_buffer)[num_smem_threads * 3 + smem_tid] = out3;
+                static_cast<threadgroup float2*>(threadgroup_buffer)[num_smem_threads * 4 + smem_tid] = out4;
+                static_cast<threadgroup float2*>(threadgroup_buffer)[num_smem_threads * 5 + smem_tid] = out5;
+                static_cast<threadgroup float2*>(threadgroup_buffer)[num_smem_threads * 6 + smem_tid] = out6;
+                static_cast<threadgroup float2*>(threadgroup_buffer)[num_smem_threads * 7 + smem_tid] = out7;
+            }
+            metal::threadgroup_barrier(metal::mem_flags::mem_threadgroup);
+            if (tid.x < num_smem_threads) {
+                out0 += static_cast<const threadgroup float2*>(threadgroup_buffer)[num_smem_threads * 0 + tid.x];
+                out1 += static_cast<const threadgroup float2*>(threadgroup_buffer)[num_smem_threads * 1 + tid.x];
+                out2 += static_cast<const threadgroup float2*>(threadgroup_buffer)[num_smem_threads * 2 + tid.x];
+                out3 += static_cast<const threadgroup float2*>(threadgroup_buffer)[num_smem_threads * 3 + tid.x];
+                out4 += static_cast<const threadgroup float2*>(threadgroup_buffer)[num_smem_threads * 4 + tid.x];
+                out5 += static_cast<const threadgroup float2*>(threadgroup_buffer)[num_smem_threads * 5 + tid.x];
+                out6 += static_cast<const threadgroup float2*>(threadgroup_buffer)[num_smem_threads * 6 + tid.x];
+                out7 += static_cast<const threadgroup float2*>(threadgroup_buffer)[num_smem_threads * 7 + tid.x];
+            }
+
+            num_threads = num_half_threads;
+        } while (num_threads > simdgroup_size);
+    }
+    if (simdgroup_idx == 0) {
+        reinterpret_cast<device float2*>(output + 0 * head_dim)[simdgroup_tid] = out0 / l0;
+        reinterpret_cast<device float2*>(output + 1 * head_dim)[simdgroup_tid] = out1 / l1;
+        reinterpret_cast<device float2*>(output + 2 * head_dim)[simdgroup_tid] = out2 / l2;
+        reinterpret_cast<device float2*>(output + 3 * head_dim)[simdgroup_tid] = out3 / l3;
+        reinterpret_cast<device float2*>(output + 4 * head_dim)[simdgroup_tid] = out4 / l4;
+        reinterpret_cast<device float2*>(output + 5 * head_dim)[simdgroup_tid] = out5 / l5;
+        reinterpret_cast<device float2*>(output + 6 * head_dim)[simdgroup_tid] = out6 / l6;
+        reinterpret_cast<device float2*>(output + 7 * head_dim)[simdgroup_tid] = out7 / l7;
+    }
+}

gptoss_kernels/source/tokenizer.c

@@ -0,0 +1,106 @@
+#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;
+}

gptoss_kernels/source/topk.metal

@@ -0,0 +1,205 @@
+#include <metal_compute>
+#include <metal_integer>
+#include <metal_math>
+#include <metal_simdgroup>
+
+#include <internal/kernel-args.h>
+
+#pragma METAL fp math_mode(safe)
+#pragma METAL fp contract(off)
+
+
+[[max_total_threads_per_threadgroup(32)]]
+kernel void gptoss_f32_topk_softmax_e128_k4(
+    constant gptoss_topk_args& args [[ buffer(0) ]],
+    const device float4* input [[ buffer(1) ]],
+    device gptoss_expert_prediction* output [[ buffer(2) ]],
+    const device gptoss_control* control [[ buffer(3) ]],
+    uint gid [[threadgroup_position_in_grid]],
+    uint tid [[thread_position_in_threadgroup]])
+{
+    const uint num_experts = 128;
+    const uint num_active_experts = 4;
+    if (control->abort != 0) {
+        return;
+    }
+
+    input += gid * (num_experts / 4);
+    output += gid * num_active_experts;
+
+    uint4 idx = tid * 4 + (uint4) {0, 1, 2, 3};
+    float4 val = input[tid];
+
+    const float topval0 = metal::simd_max(metal::max3(metal::max(val.x, val.y), val.z, val.w));
+    uint idx0 = 0xFFFFFFFFu;
+    if (val.w == topval0) {
+        idx0 = idx.w;
+    }
+    if (val.z == topval0) {
+        idx0 = idx.z;
+    }
+    if (val.y == topval0) {
+        idx0 = idx.y;
+    }
+    if (val.x == topval0) {
+        idx0 = idx.x;
+    }
+    const uint topidx0 = metal::simd_min(idx0);
+    const bool4 is_topidx0 = idx == topidx0;
+    val = metal::select(val, -INFINITY, is_topidx0);
+    idx = metal::select(idx, 0xFFFFFFFFu, is_topidx0);
+
+    const float topval1 = metal::simd_max(metal::max3(metal::max(val.x, val.y), val.z, val.w));
+    uint idx1 = 0xFFFFFFFFu;
+    if (val.w == topval1) {
+        idx1 = idx.w;
+    }
+    if (val.z == topval1) {
+        idx1 = idx.z;
+    }
+    if (val.y == topval1) {
+        idx1 = idx.y;
+    }
+    if (val.x == topval1) {
+        idx1 = idx.x;
+    }
+    const uint topidx1 = metal::simd_min(idx1);
+    const bool4 is_topidx1 = idx == topidx1;
+    val = metal::select(val, -INFINITY, is_topidx1);
+    idx = metal::select(idx, 0xFFFFFFFFu, is_topidx1);
+
+    const float topval2 = metal::simd_max(metal::max3(metal::max(val.x, val.y), val.z, val.w));
+    uint idx2 = 0xFFFFFFFFu;
+    if (val.w == topval2) {
+        idx2 = idx.w;
+    }
+    if (val.z == topval2) {
+        idx2 = idx.z;
+    }
+    if (val.y == topval2) {
+        idx2 = idx.y;
+    }
+    if (val.x == topval2) {
+        idx2 = idx.x;
+    }
+    const uint topidx2 = metal::simd_min(idx2);
+    const bool4 is_topidx2 = idx == topidx2;
+    val = metal::select(val, -INFINITY, is_topidx2);
+    idx = metal::select(idx, 0xFFFFFFFFu, is_topidx2);
+
+    const float topval3 = metal::simd_max(metal::max3(metal::max(val.x, val.y), val.z, val.w));
+    uint idx3 = 0xFFFFFFFFu;
+    if (val.w == topval3) {
+        idx3 = idx.w;
+    }
+    if (val.z == topval3) {
+        idx3 = idx.z;
+    }
+    if (val.y == topval3) {
+        idx3 = idx.y;
+    }
+    if (val.x == topval3) {
+        idx3 = idx.x;
+    }
+    const uint topidx3 = metal::simd_min(idx3);
+
+    if (metal::simd_is_first()) {
+        const float topexp0 = 1.0f;
+        const float topexp1 = metal::precise::exp(topval1 - topval0);
+        const float topexp2 = metal::precise::exp(topval2 - topval0);
+        const float topexp3 = metal::precise::exp(topval3 - topval0);
+
+        const float sum = (topexp0 + topexp1) + (topexp2 + topexp3);
+        const float scale = 1.0 / sum;
+
+        output[0] = (gptoss_expert_prediction) {
+            .expert_id = topidx0,
+            .score = topexp0 * scale,
+        };
+        output[1] = (gptoss_expert_prediction) {
+            .expert_id = topidx1,
+            .score = topexp1 * scale,
+        };
+        output[2] = (gptoss_expert_prediction) {
+            .expert_id = topidx2,
+            .score = topexp2 * scale,
+        };
+        output[3] = (gptoss_expert_prediction) {
+            .expert_id = topidx3,
+            .score = topexp3 * scale,
+        };
+    }
+}
+
+[[max_total_threads_per_threadgroup(32)]]
+kernel void gptoss_f32_topk_softmax_e32_k4(
+    constant gptoss_topk_args& args [[ buffer(0) ]],
+    const device float* input [[ buffer(1) ]],
+    device gptoss_expert_prediction* output [[ buffer(2) ]],
+    const device gptoss_control* control [[ buffer(3) ]],
+    uint gid [[threadgroup_position_in_grid]],
+    uint tid [[thread_position_in_threadgroup]])
+{
+    const uint num_experts = 32;
+    const uint num_active_experts = 4;
+    if (control->abort != 0) {
+        return;
+    }
+
+    input += gid * num_experts;
+    output += gid * num_active_experts;
+
+    float val = input[tid];
+    uint idx = tid;
+
+    const float topval0 = metal::simd_max(val);
+    const uint topidx0 = metal::simd_min(val == topval0 ? idx : 0xFFFFFFFFu);
+    if (idx == topidx0) {
+        val = -INFINITY;
+        idx = 0xFFFFFFFFu;
+    }
+
+    const float topval1 = metal::simd_max(val);
+    const uint topidx1 = metal::simd_min(val == topval1 ? idx : 0xFFFFFFFFu);
+    if (idx == topidx1) {
+        val = -INFINITY;
+        idx = 0xFFFFFFFFu;
+    }
+
+    const float topval2 = metal::simd_max(val);
+    const uint topidx2 = metal::simd_min(val == topval2 ? idx : 0xFFFFFFFFu);
+    if (idx == topidx2) {
+        val = -INFINITY;
+        idx = 0xFFFFFFFFu;
+    }
+
+    const float topval3 = metal::simd_max(val);
+    const uint topidx3 = metal::simd_min(val == topval3 ? idx : 0xFFFFFFFFu);
+
+    if (metal::simd_is_first()) {
+        const float topexp0 = 1.0f;
+        const float topexp1 = metal::precise::exp(topval1 - topval0);
+        const float topexp2 = metal::precise::exp(topval2 - topval0);
+        const float topexp3 = metal::precise::exp(topval3 - topval0);
+
+        const float sum = (topexp0 + topexp1) + (topexp2 + topexp3);
+        const float scale = 1.0 / sum;
+
+        output[0] = (gptoss_expert_prediction) {
+            .expert_id = topidx0,
+            .score = topexp0 * scale,
+        };
+        output[1] = (gptoss_expert_prediction) {
+            .expert_id = topidx1,
+            .score = topexp1 * scale,
+        };
+        output[2] = (gptoss_expert_prediction) {
+            .expert_id = topidx2,
+            .score = topexp2 * scale,
+        };
+        output[3] = (gptoss_expert_prediction) {
+            .expert_id = topidx3,
+            .score = topexp3 * scale,
+        };
+    }
+}

gptoss_kernels/test/bf16-f32-embeddings.cc

@@ -0,0 +1,33 @@
+#include <gtest/gtest.h>
+
+#include <cstddef>
+
+#include "embeddings-kernel-tester.hpp"
+
+
+using gptoss::EmbeddingsKernelTester;
+
+constexpr std::size_t kThreadgroupSize = 64;
+
+
+TEST(BF16_F32_EMBEDDINGS, single_token_single_tile) {
+    EmbeddingsKernelTester()
+        .num_channels(kThreadgroupSize)
+        .threadgroup_size(kThreadgroupSize)
+        .TestBF16_F32();
+}
+
+TEST(BF16_F32_EMBEDDINGS, single_token_multi_tile) {
+    EmbeddingsKernelTester()
+        .num_channels(kThreadgroupSize * 4 + 16)
+        .threadgroup_size(kThreadgroupSize)
+        .TestBF16_F32();
+}
+
+TEST(BF16_F32_EMBEDDINGS, multiple_tokens) {
+    EmbeddingsKernelTester()
+        .num_channels(kThreadgroupSize * 4 + 16)
+        .num_tokens(3)
+        .threadgroup_size(kThreadgroupSize)
+        .TestBF16_F32();
+}

gptoss_kernels/test/embeddings-kernel-tester.hpp

@@ -0,0 +1,123 @@
+#pragma once
+
+#include <gtest/gtest.h>
+
+#include <cstddef>
+#include <cstdint>
+
+#include <internal/datatype.hpp>
+#include <internal/metal.hpp>
+#include <internal/metal-kernels.h>
+
+
+namespace gptoss {
+
+class EmbeddingsKernelTester {
+public:
+    EmbeddingsKernelTester() { }
+
+    EmbeddingsKernelTester(const EmbeddingsKernelTester&) = delete;
+    EmbeddingsKernelTester(EmbeddingsKernelTester&&) = delete;
+    EmbeddingsKernelTester& operator=(const EmbeddingsKernelTester&) = delete;
+    EmbeddingsKernelTester& operator=(EmbeddingsKernelTester&&) = delete;
+
+    [[nodiscard]]
+    EmbeddingsKernelTester& num_channels(std::uint32_t num_channels) {
+        num_channels_ = num_channels;
+        return *this;
+    }
+
+    std::uint32_t num_channels() const {
+        return num_channels_;
+    }
+
+    [[nodiscard]]
+    EmbeddingsKernelTester& num_tokens(std::uint32_t num_tokens) {
+        num_tokens_ = num_tokens;
+        return *this;
+    }
+
+    std::uint32_t num_tokens() const {
+        return num_tokens_;
+    }
+
+    std::uint32_t vocabulary_size() const {
+        return num_tokens() + 1;
+    }
+
+    [[nodiscard]]
+    EmbeddingsKernelTester& threadgroup_size(std::size_t threadgroup_size) {
+        threadgroup_size_ = threadgroup_size;
+        return *this;
+    }
+
+    std::size_t threadgroup_size() const {
+        return threadgroup_size_;
+    }
+
+    void Validate() const {
+        ASSERT_NE(num_channels(), 0);
+        ASSERT_NE(num_tokens(), 0);
+        ASSERT_NE(threadgroup_size(), 0);
+        ASSERT_EQ(threadgroup_size() % 32, 0);
+    }
+
+    void TestBF16_F32() const {
+        Validate();
+
+        metal::CommandBuffer command_buffer{command_queue_};
+        metal::Buffer token_buffer{device_, sizeof(std::uint32_t)};
+        metal::Buffer weight_buffer{device_, vocabulary_size() * num_channels() * sizeof(gptoss_bfloat16)};
+        metal::Buffer output_buffer{device_, num_channels() * sizeof(float)};
+        metal::Buffer control_buffer{device_, sizeof(gptoss_control)};
+        std::memset(control_buffer.ptr(), 0, sizeof(gptoss_control));
+
+        std::uint32_t* token_ptr = static_cast<std::uint32_t*>(token_buffer.ptr());
+        for (std::uint32_t t = 0; t < num_tokens(); t++) {
+            token_ptr[t] = t + 1;
+        }
+
+        Check(gptoss_metal_command_buffer_encode_launch_bf16_f32_embeddings(
+                command_buffer.handle(),
+                bf16_f32_embeddings_fn.handle(),
+                threadgroup_size(),
+                token_buffer.handle(),
+                /*token_offset=*/0,
+                weight_buffer.handle(),
+                /*weight_offset=*/0,
+                output_buffer.handle(),
+                /*output_offset=*/0,
+                control_buffer.handle(),
+                /*control_offset=*/0,
+                num_tokens(),
+                num_channels()),
+            "gptoss_metal_command_buffer_encode_launch_bf16_f32_embeddings");
+
+        command_buffer.commit();
+        command_buffer.wait_completion();
+
+        const gptoss_bfloat16* weight_ptr = static_cast<const gptoss_bfloat16*>(weight_buffer.ptr());
+        const float* output_ptr = static_cast<const float*>(output_buffer.ptr());
+        for (std::uint32_t t = 0; t < num_tokens(); t++) {
+            const std::uint32_t token = token_ptr[t];
+            for (std::uint32_t i = 0; i < num_channels(); i++) {
+                const gptoss_bfloat16 input_val = weight_ptr[token * num_channels() + i];
+                const float ref_output = upcast<float>(input_val);
+                const float output = output_ptr[t * num_channels() + i];
+                ASSERT_EQ(output, ref_output)
+                    << "at token " << t << ", position " << i << " / " << num_channels() << ", input " << std::uint32_t(input_val.bits);
+            }
+        }
+    }
+
+private:
+    metal::Device device_{};
+    metal::CommandQueue command_queue_{device_};
+    metal::Library library_{device_};
+    metal::Function bf16_f32_embeddings_fn{library_, "gptoss_bf16_f32_embeddings"};
+    std::uint32_t num_tokens_{1};
+    std::uint32_t num_channels_{1};
+    std::size_t threadgroup_size_{32};
+};
+
+}  // namespace gptoss

gptoss_kernels/test/f32-bf16w-matmul.cc

@@ -0,0 +1,87 @@
+#include <gtest/gtest.h>
+
+#include <cstddef>
+#include <cstdint>
+
+#include "matmul-kernel-tester.hpp"
+
+
+using gptoss::MatMulKernelTester;
+
+constexpr size_t kSimdgroupSize = 32;  // fixed in the kernel
+
+TEST(F32_BF16W_MATMUL, single_simdgroup_single_iteration) {
+    MatMulKernelTester()
+        .num_rows(1)
+        .num_cols(kSimdgroupSize * 4)
+        .threadgroup_size(kSimdgroupSize)
+        .TestF32_BF16W();
+}
+
+TEST(F32_BF16W_MATMUL, single_simdgroup_multiple_iteration) {
+    MatMulKernelTester()
+        .num_rows(1)
+        .num_cols((2 * kSimdgroupSize + 1) * 4)
+        .threadgroup_size(kSimdgroupSize)
+        .TestF32_BF16W();
+}
+
+TEST(F32_BF16W_MATMUL, single_threadgroup) {
+    constexpr std::size_t threadgroup_size = 2 * kSimdgroupSize;
+
+    MatMulKernelTester()
+        .num_rows(threadgroup_size / kSimdgroupSize)
+        .num_cols((2 * kSimdgroupSize + 1) * 4)
+        .threadgroup_size(threadgroup_size)
+        .TestF32_BF16W();
+}
+
+TEST(F32_BF16W_MATMUL, multiple_threadgroups) {
+    constexpr std::size_t threadgroup_size = 2 * kSimdgroupSize;
+    constexpr std::uint32_t num_threadgroups = 3;
+
+    MatMulKernelTester()
+        .num_rows(num_threadgroups * threadgroup_size / kSimdgroupSize)
+        .num_cols((2 * kSimdgroupSize + 1) * 4)
+        .threadgroup_size(threadgroup_size)
+        .TestF32_BF16W();
+}
+
+TEST(F32_BF16W_MATMUL, multiple_tokens) {
+    constexpr std::size_t threadgroup_size = 2 * kSimdgroupSize;
+    constexpr std::uint32_t num_threadgroups = 3;
+
+    MatMulKernelTester()
+        .num_rows(num_threadgroups * threadgroup_size / kSimdgroupSize)
+        .num_cols((2 * kSimdgroupSize + 1) * 4)
+        .num_tokens(2)
+        .threadgroup_size(threadgroup_size)
+        .TestF32_BF16W();
+}
+
+TEST(F32_BF16W_DENSE_MATMUL_QKV, seq_len_1024) {
+    MatMulKernelTester()
+        .num_tokens(1024)
+        .num_rows(5120)
+        .num_cols(2880)
+        .TestF32_BF16W(
+            MatMulKernelTester::MatMulKernelType::PREFILL_QKV_OPTIMIZED);
+}
+
+TEST(F32_BF16W_DENSE_MATMUL_ATTN_OUTPUT, seq_len_1024) {
+    MatMulKernelTester()
+        .num_tokens(1024)
+        .num_rows(2880)
+        .num_cols(4096)
+        .TestF32_BF16W(
+            MatMulKernelTester::MatMulKernelType::PREFILL_ATTN_OUTPUT_OPTIMIZED);
+}
+
+TEST(F32_BF16W_DENSE_MATMUL_MLP_GATE, seq_len_1024) {
+    MatMulKernelTester()
+        .num_tokens(1024)
+        .num_rows(128)
+        .num_cols(2880)
+        .TestF32_BF16W(
+            MatMulKernelTester::MatMulKernelType::PREFILL_MLP_GATE_OPTIMIZED);
+}

gptoss_kernels/test/f32-bf16w-rmsnorm.cc

@@ -0,0 +1,36 @@
+#include <gtest/gtest.h>
+
+#include <cstdint>
+
+#include "rmsnorm-kernel-tester.hpp"
+
+
+using gptoss::RMSNormKernelTester;
+
+constexpr std::uint32_t kThreadgroupSize = 1024;  // fixed in the kernel
+constexpr std::uint32_t kVectorSize = 4;  // fixed in the kernel
+
+TEST(F32_BF16W_RMSNORM, single_iteration) {
+    RMSNormKernelTester()
+        .num_channels(kThreadgroupSize)
+        .TestF32_BF16W();
+}
+
+TEST(F32_BF16W_RMSNORM, multiple_iterations) {
+    RMSNormKernelTester()
+        .num_channels(kThreadgroupSize * 2)
+        .TestF32_BF16W();
+}
+
+TEST(F32_BF16W_RMSNORM, partial_iteration) {
+    RMSNormKernelTester()
+        .num_channels(kThreadgroupSize * 2 + kVectorSize)
+        .TestF32_BF16W();
+}
+
+TEST(F32_BF16W_RMSNORM, multiple_tokens) {
+    RMSNormKernelTester()
+        .num_tokens(3)
+        .num_channels(kThreadgroupSize * 2 + kVectorSize)
+        .TestF32_BF16W();
+}