) for a message.
+ """
+ return jinja_env.from_string(_message_template).render(
+ role=message["role"],
+ content=message["content"],
+ variant=message.get("variant", None),
+ )
+
+
+jinja_env.globals["message_to_html"] = message_to_html
+
+
+_report_template = """
+
+
+
+
+
+
+ {% if metrics %}
+
Metrics
+
+
+ | Metric |
+ Value |
+
+
+ | Score |
+ {{ score | float | round(3) }} |
+
+ {% for name, value in metrics.items() %}
+
+ | {{ name }} |
+ {{ value }} |
+
+ {% endfor %}
+
+ {% endif %}
+
Examples
+ {% for html in htmls %}
+ {{ html | safe }}
+
+ {% endfor %}
+
+
+"""
+
+
+def make_report(eval_result: EvalResult) -> str:
+ """
+ Create a standalone HTML report from an EvalResult.
+ """
+ return jinja_env.from_string(_report_template).render(
+ score=eval_result.score,
+ metrics=eval_result.metrics,
+ htmls=eval_result.htmls,
+ )
diff --git a/gpt_oss/evals/responses_sampler.py b/gpt_oss/evals/responses_sampler.py
new file mode 100644
index 0000000000000000000000000000000000000000..134303f5db715f7b6a0261e45cfcf8b5c902f762
--- /dev/null
+++ b/gpt_oss/evals/responses_sampler.py
@@ -0,0 +1,85 @@
+import time
+from typing import Any
+
+import openai
+from openai import OpenAI
+
+from .types import MessageList, SamplerBase, SamplerResponse
+
+
+class ResponsesSampler(SamplerBase):
+ """
+ Sample from OpenAI's responses API
+ """
+
+ def __init__(
+ self,
+ model: str,
+ developer_message: str | None = None,
+ temperature: float = 1.0,
+ max_tokens: int = 131_072,
+ reasoning_model: bool = False,
+ reasoning_effort: str | None = None,
+ base_url: str = "http://localhost:8000/v1",
+ ):
+ self.client = OpenAI(base_url=base_url, timeout=24*60*60)
+ self.model = model
+ self.developer_message = developer_message
+ self.temperature = temperature
+ self.max_tokens = max_tokens
+ self.image_format = "url"
+ self.reasoning_model = reasoning_model
+ self.reasoning_effort = reasoning_effort
+
+ def _pack_message(self, role: str, content: Any) -> dict[str, Any]:
+ return {"role": role, "content": content}
+
+ def __call__(self, message_list: MessageList) -> SamplerResponse:
+ if self.developer_message:
+ message_list = [
+ self._pack_message("developer", self.developer_message)
+ ] + message_list
+ trial = 0
+ while True:
+ try:
+ request_kwargs = {
+ "model": self.model,
+ "input": message_list,
+ "temperature": self.temperature,
+ "max_output_tokens": self.max_tokens,
+ }
+ if self.reasoning_model:
+ request_kwargs["reasoning"] = (
+ {"effort": self.reasoning_effort} if self.reasoning_effort else None
+ )
+ response = self.client.responses.create(**request_kwargs)
+
+ for output in response.output:
+ if hasattr(output, "text"):
+ message_list.append(self._pack_message(getattr(output, "role", "assistant"), output.text))
+ elif hasattr(output, "content"):
+ for c in output.content:
+ # c.text handled below
+ pass
+
+ return SamplerResponse(
+ response_text=response.output_text,
+ response_metadata={"usage": response.usage},
+ actual_queried_message_list=message_list,
+ )
+ except openai.BadRequestError as e:
+ print("Bad Request Error", e)
+ return SamplerResponse(
+ response_text="",
+ response_metadata={"usage": None},
+ actual_queried_message_list=message_list,
+ )
+ except Exception as e:
+ exception_backoff = 2**trial # expontial back off
+ print(
+ f"Rate limit exception so wait and retry {trial} after {exception_backoff} sec",
+ e,
+ )
+ time.sleep(exception_backoff)
+ trial += 1
+ # unknown error shall throw exception
diff --git a/gpt_oss/evals/types.py b/gpt_oss/evals/types.py
new file mode 100644
index 0000000000000000000000000000000000000000..8f2b42d75705d4717d3204254bb365ccb326013d
--- /dev/null
+++ b/gpt_oss/evals/types.py
@@ -0,0 +1,66 @@
+from dataclasses import dataclass, field
+from typing import Any, Literal, overload
+
+Message = dict[str, Any] # keys role, content
+MessageList = list[Message]
+
+
+
+@dataclass
+class SamplerResponse:
+ """
+ Response from a sampler.
+ """
+ response_text: str
+ actual_queried_message_list: MessageList
+ response_metadata: dict[str, Any]
+
+class SamplerBase:
+ """
+ Base class for defining a sampling model, which can be evaluated,
+ or used as part of the grading process.
+ """
+
+ def __call__(
+ self,
+ message_list: MessageList,
+ ) -> SamplerResponse:
+ raise NotImplementedError
+
+
+@dataclass
+class EvalResult:
+ """
+ Result of running an evaluation (usually consisting of many samples)
+ """
+
+ score: float | None # top-line metric
+ metrics: dict[str, float] | None # other metrics
+ htmls: list[str] # strings of valid HTML
+ convos: list[MessageList] # sampled conversations
+ metadata: dict[str, Any] | None # Extra data such as rubric scores or sollen
+
+
+@dataclass
+class SingleEvalResult:
+ """
+ Result of evaluating a single sample
+ """
+
+ score: float | None
+ metrics: dict[str, float] = field(default_factory=dict)
+ html: str | None = None
+ convo: MessageList | None = None # sampled conversation
+ example_level_metadata: dict[str, Any] | None = (
+ None # Extra data such as rubric scores or sollen
+ )
+
+
+class Eval:
+ """
+ Base class for defining an evaluation.
+ """
+
+ def __call__(self, sampler: SamplerBase) -> EvalResult:
+ raise NotImplementedError
+
diff --git a/gpt_oss/generate.py b/gpt_oss/generate.py
new file mode 100644
index 0000000000000000000000000000000000000000..c0755805eb11063c3ce32158b596cb026fd145e5
--- /dev/null
+++ b/gpt_oss/generate.py
@@ -0,0 +1,95 @@
+# Model parallel inference
+# Note: This script is for demonstration purposes only. It is not designed for production use.
+# See gpt_oss.chat for a more complete example with the Harmony parser.
+# torchrun --nproc-per-node=4 -m gpt_oss.generate -p "why did the chicken cross the road?" model/
+
+import argparse
+
+from gpt_oss.tokenizer import get_tokenizer
+
+
+def main(args):
+ match args.backend:
+ case "torch":
+ from gpt_oss.torch.utils import init_distributed
+ from gpt_oss.torch.model import TokenGenerator as TorchGenerator
+ device = init_distributed()
+ generator = TorchGenerator(args.checkpoint, device=device)
+ case "triton":
+ from gpt_oss.torch.utils import init_distributed
+ from gpt_oss.triton.model import TokenGenerator as TritonGenerator
+ device = init_distributed()
+ generator = TritonGenerator(args.checkpoint, context=args.context_length, device=device)
+ case "vllm":
+ from gpt_oss.vllm.token_generator import TokenGenerator as VLLMGenerator
+ generator = VLLMGenerator(args.checkpoint, tensor_parallel_size=args.tensor_parallel_size)
+ case _:
+ raise ValueError(f"Invalid backend: {args.backend}")
+
+ tokenizer = get_tokenizer()
+ tokens = tokenizer.encode(args.prompt)
+ max_tokens = None if args.limit == 0 else args.limit
+ for token, logprob in generator.generate(tokens, stop_tokens=[tokenizer.eot_token], temperature=args.temperature, max_tokens=max_tokens, return_logprobs=True):
+ tokens.append(token)
+ token_text = tokenizer.decode([token])
+ print(
+ f"Generated token: {repr(token_text)}, logprob: {logprob}"
+ )
+
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser(description="Text generation example")
+ parser.add_argument(
+ "checkpoint",
+ metavar="FILE",
+ type=str,
+ help="Path to the SafeTensors checkpoint",
+ )
+ parser.add_argument(
+ "-p",
+ "--prompt",
+ metavar="PROMPT",
+ type=str,
+ default="How are you?",
+ help="LLM prompt",
+ )
+ parser.add_argument(
+ "-t",
+ "--temperature",
+ metavar="TEMP",
+ type=float,
+ default=0.0,
+ help="Sampling temperature",
+ )
+ parser.add_argument(
+ "-l",
+ "--limit",
+ metavar="LIMIT",
+ type=int,
+ default=0,
+ help="Limit on the number of tokens (0 to disable)",
+ )
+ parser.add_argument(
+ "-b",
+ "--backend",
+ metavar="BACKEND",
+ type=str,
+ default="torch",
+ choices=["triton", "torch", "vllm"],
+ help="Inference backend",
+ )
+ parser.add_argument(
+ "--tensor-parallel-size",
+ type=int,
+ default=2,
+ help="Tensor parallel size for vLLM backend",
+ )
+ parser.add_argument(
+ "--context-length",
+ type=int,
+ default=4096,
+ help="Context length for Triton backend",
+ )
+ args = parser.parse_args()
+
+ main(args)
diff --git a/gpt_oss/metal/CMakeLists.txt b/gpt_oss/metal/CMakeLists.txt
new file mode 100644
index 0000000000000000000000000000000000000000..264696e1dce9f30c8c959fcd955a6e283e0e2c45
--- /dev/null
+++ b/gpt_oss/metal/CMakeLists.txt
@@ -0,0 +1,189 @@
+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/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/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/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_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}
+ $)
+
+# 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)
+# ------------------------------------------------------------------------
\ No newline at end of file
diff --git a/gpt_oss/metal/__init__.py b/gpt_oss/metal/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..fbddbe918a81bdc37793c013d0f268b6c273a9de
--- /dev/null
+++ b/gpt_oss/metal/__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
diff --git a/gpt_oss/metal/benchmark/end-to-end-threadgroup.cc b/gpt_oss/metal/benchmark/end-to-end-threadgroup.cc
new file mode 100644
index 0000000000000000000000000000000000000000..82048cd1dece65d4ce72102f318854dd7a25114b
--- /dev/null
+++ b/gpt_oss/metal/benchmark/end-to-end-threadgroup.cc
@@ -0,0 +1,590 @@
+#include
+#include
+
+#include
+#include
+#include
+#include
+#include
+#include
+#include
+#include
+
+#include
+
+
+constexpr std::uint32_t kNumGeneratedTokens = 100;
+
+
+static void attn_qkv_tgsize(benchmark::State& state, const char* env_var_name) {
+ const char* model_path = getenv(env_var_name);
+ if (model_path == NULL) {
+ state.SkipWithError(std::format("environment variable {} is not set", env_var_name));
+ return;
+ }
+
+ gptoss_model_t model_ptr = nullptr;
+ gptoss_status status = gptoss_model_create_from_file(model_path, &model_ptr);
+ if (status != gptoss_status_success) {
+ state.SkipWithError(std::format("failed to load model from file {}", model_path));
+ return;
+ }
+ std::unique_ptr, decltype(&gptoss_model_release)> model(model_ptr, gptoss_model_release);
+ model->attn_qkv_threadgroup_size = static_cast(state.range(0));
+
+ gptoss_context_t context_ptr = nullptr;
+ status = gptoss_context_create(model.get(), /*context_length=*/0, /*max_batch_tokens=*/0, &context_ptr);
+ if (status != gptoss_status_success) {
+ state.SkipWithError("failed to create Context object");
+ return;
+ }
+ std::unique_ptr, decltype(&gptoss_context_release)> context(context_ptr, gptoss_context_release);
+
+ const char* prompt = "why did the chicken cross the road?";
+ std::size_t num_prompt_tokens = 0;
+ status = gptoss_context_append_chars(context.get(), prompt, strlen(prompt), &num_prompt_tokens);
+ if (status != gptoss_status_success) {
+ state.SkipWithError(std::format("failed to tokenize prompt \"{}\"", prompt));
+ return;
+ }
+
+ // Prefill
+ status = gptoss_context_process(context.get());
+ if (status != gptoss_status_success) {
+ state.SkipWithError("failed to prefill Context object");
+ return;
+ }
+ const std::size_t num_kvcache_tokens = context->num_kv_tokens;
+
+ std::uint64_t rng_seed = 0;
+ for (auto _ : state) {
+ const std::uint64_t current_rng_seed = rng_seed++;
+ context->num_kv_tokens = num_prompt_tokens;
+ context->num_tokens = num_prompt_tokens;
+
+ std::array tokens;
+ std::size_t num_generated_tokens = 0;
+ do {
+ std::size_t num_current_generated_tokens = 0;
+ status = gptoss_context_sample(context.get(), /*temperature=*/1.0f, /*rng_state=*/current_rng_seed,
+ /*max_tokens=*/kNumGeneratedTokens - num_generated_tokens, tokens.data(), &num_current_generated_tokens);
+ if (status != gptoss_status_success) {
+ state.SkipWithError("failed to sample from the Context object");
+ return;
+ }
+ num_generated_tokens += num_current_generated_tokens;
+ } while (num_generated_tokens < kNumGeneratedTokens);
+ }
+
+ state.counters["generations"] =
+ benchmark::Counter(state.iterations(), benchmark::Counter::kIsRate);
+ state.counters["tokens"] =
+ benchmark::Counter(state.iterations() * kNumGeneratedTokens, benchmark::Counter::kIsRate);
+}
+
+static void AttnQKVThreadgroupSizeArguments(benchmark::internal::Benchmark* b) {
+ b->ArgNames({"tgsize"});
+ for (auto attn_qkv_threadgroup_size = 32; attn_qkv_threadgroup_size <= 1024; attn_qkv_threadgroup_size += 32) {
+ const auto num_simdgroups = attn_qkv_threadgroup_size / 32;
+ if (5120 % num_simdgroups != 0) {
+ // Skip incompatible threadgroup sizes
+ continue;
+ }
+ b->Args({attn_qkv_threadgroup_size});
+ }
+}
+
+BENCHMARK_CAPTURE(attn_qkv_tgsize, gpt_oss_20b, "GPT_OSS_20B_PATH")
+ ->UseRealTime()->Unit(benchmark::kMillisecond)->Apply(AttnQKVThreadgroupSizeArguments);
+BENCHMARK_CAPTURE(attn_qkv_tgsize, gpt_oss_120b, "GPT_OSS_120B_PATH")
+ ->UseRealTime()->Unit(benchmark::kMillisecond)->Apply(AttnQKVThreadgroupSizeArguments);
+
+static void attn_out_tgsize(benchmark::State& state, const char* env_var_name) {
+ const char* model_path = getenv(env_var_name);
+ if (model_path == NULL) {
+ state.SkipWithError(std::format("environment variable {} is not set", env_var_name));
+ return;
+ }
+
+ gptoss_model_t model_ptr = nullptr;
+ gptoss_status status = gptoss_model_create_from_file(model_path, &model_ptr);
+ if (status != gptoss_status_success) {
+ state.SkipWithError(std::format("failed to load model from file {}", model_path));
+ return;
+ }
+ std::unique_ptr, decltype(&gptoss_model_release)> model(model_ptr, gptoss_model_release);
+ model->attn_out_threadgroup_size = static_cast(state.range(0));
+
+ gptoss_context_t context_ptr = nullptr;
+ status = gptoss_context_create(model.get(), /*context_length=*/0, /*max_batch_tokens=*/0, &context_ptr);
+ if (status != gptoss_status_success) {
+ state.SkipWithError("failed to create Context object");
+ return;
+ }
+ std::unique_ptr, decltype(&gptoss_context_release)> context(context_ptr, gptoss_context_release);
+
+ const char* prompt = "why did the chicken cross the road?";
+ std::size_t num_prompt_tokens = 0;
+ status = gptoss_context_append_chars(context.get(), prompt, strlen(prompt), &num_prompt_tokens);
+ if (status != gptoss_status_success) {
+ state.SkipWithError(std::format("failed to tokenize prompt \"{}\"", prompt));
+ return;
+ }
+
+ // Prefill
+ status = gptoss_context_process(context.get());
+ if (status != gptoss_status_success) {
+ state.SkipWithError("failed to prefill Context object");
+ return;
+ }
+ const std::size_t num_kvcache_tokens = context->num_kv_tokens;
+
+ std::uint64_t rng_seed = 0;
+ for (auto _ : state) {
+ const std::uint64_t current_rng_seed = rng_seed++;
+ context->num_kv_tokens = num_prompt_tokens;
+ context->num_tokens = num_prompt_tokens;
+
+ std::array tokens;
+ std::size_t num_generated_tokens = 0;
+ do {
+ std::size_t num_current_generated_tokens = 0;
+ status = gptoss_context_sample(context.get(), /*temperature=*/1.0f, /*rng_state=*/current_rng_seed,
+ /*max_tokens=*/kNumGeneratedTokens - num_generated_tokens, tokens.data(), &num_current_generated_tokens);
+ if (status != gptoss_status_success) {
+ state.SkipWithError("failed to sample from the Context object");
+ return;
+ }
+ num_generated_tokens += num_current_generated_tokens;
+ } while (num_generated_tokens < kNumGeneratedTokens);
+ }
+
+ state.counters["generations"] =
+ benchmark::Counter(state.iterations(), benchmark::Counter::kIsRate);
+ state.counters["tokens"] =
+ benchmark::Counter(state.iterations() * kNumGeneratedTokens, benchmark::Counter::kIsRate);
+}
+
+static void AttnOutThreadgroupSizeArguments(benchmark::internal::Benchmark* b) {
+ b->ArgNames({"tgsize"});
+ for (auto attn_out_threadgroup_size = 32; attn_out_threadgroup_size <= 1024; attn_out_threadgroup_size += 32) {
+ const auto num_simdgroups = attn_out_threadgroup_size / 32;
+ if (2880 % num_simdgroups != 0) {
+ // Skip incompatible threadgroup sizes
+ continue;
+ }
+ b->Args({attn_out_threadgroup_size});
+ }
+}
+
+BENCHMARK_CAPTURE(attn_out_tgsize, gpt_oss_20b, "GPT_OSS_20B_PATH")
+ ->UseRealTime()->Unit(benchmark::kMillisecond)->Apply(AttnOutThreadgroupSizeArguments);
+BENCHMARK_CAPTURE(attn_out_tgsize, gpt_oss_120b, "GPT_OSS_120B_PATH")
+ ->UseRealTime()->Unit(benchmark::kMillisecond)->Apply(AttnOutThreadgroupSizeArguments);
+
+static void mlp_gate_tgsize(benchmark::State& state, const char* env_var_name) {
+ const char* model_path = getenv(env_var_name);
+ if (model_path == NULL) {
+ state.SkipWithError(std::format("environment variable {} is not set", env_var_name));
+ return;
+ }
+
+ gptoss_model_t model_ptr = nullptr;
+ gptoss_status status = gptoss_model_create_from_file(model_path, &model_ptr);
+ if (status != gptoss_status_success) {
+ state.SkipWithError(std::format("failed to load model from file {}", model_path));
+ return;
+ }
+ std::unique_ptr, decltype(&gptoss_model_release)> model(model_ptr, gptoss_model_release);
+ model->mlp_gate_threadgroup_size = static_cast(state.range(0));
+
+ gptoss_context_t context_ptr = nullptr;
+ status = gptoss_context_create(model.get(), /*context_length=*/0, /*max_batch_tokens=*/0, &context_ptr);
+ if (status != gptoss_status_success) {
+ state.SkipWithError("failed to create Context object");
+ return;
+ }
+ std::unique_ptr, decltype(&gptoss_context_release)> context(context_ptr, gptoss_context_release);
+
+ const char* prompt = "why did the chicken cross the road?";
+ std::size_t num_prompt_tokens = 0;
+ status = gptoss_context_append_chars(context.get(), prompt, strlen(prompt), &num_prompt_tokens);
+ if (status != gptoss_status_success) {
+ state.SkipWithError(std::format("failed to tokenize prompt \"{}\"", prompt));
+ return;
+ }
+
+ // Prefill
+ status = gptoss_context_process(context.get());
+ if (status != gptoss_status_success) {
+ state.SkipWithError("failed to prefill Context object");
+ return;
+ }
+ const std::size_t num_kvcache_tokens = context->num_kv_tokens;
+
+ std::uint64_t rng_seed = 0;
+ for (auto _ : state) {
+ const std::uint64_t current_rng_seed = rng_seed++;
+ context->num_kv_tokens = num_prompt_tokens;
+ context->num_tokens = num_prompt_tokens;
+
+ std::array tokens;
+ std::size_t num_generated_tokens = 0;
+ do {
+ std::size_t num_current_generated_tokens = 0;
+ status = gptoss_context_sample(context.get(), /*temperature=*/1.0f, /*rng_state=*/current_rng_seed,
+ /*max_tokens=*/kNumGeneratedTokens - num_generated_tokens, tokens.data(), &num_current_generated_tokens);
+ if (status != gptoss_status_success) {
+ state.SkipWithError("failed to sample from the Context object");
+ return;
+ }
+ num_generated_tokens += num_current_generated_tokens;
+ } while (num_generated_tokens < kNumGeneratedTokens);
+ }
+
+ state.counters["generations"] =
+ benchmark::Counter(state.iterations(), benchmark::Counter::kIsRate);
+ state.counters["tokens"] =
+ benchmark::Counter(state.iterations() * kNumGeneratedTokens, benchmark::Counter::kIsRate);
+}
+
+static void MlpGateThreadgroupSizeArguments(benchmark::internal::Benchmark* b) {
+ b->ArgNames({"tgsize"});
+ for (auto mlp_gate_threadgroup_size = 32; mlp_gate_threadgroup_size <= 1024; mlp_gate_threadgroup_size += 32) {
+ const auto num_simdgroups = mlp_gate_threadgroup_size / 32;
+ if (128 % num_simdgroups != 0) {
+ // Skip incompatible threadgroup sizes
+ continue;
+ }
+ b->Args({mlp_gate_threadgroup_size});
+ }
+}
+
+BENCHMARK_CAPTURE(mlp_gate_tgsize, gpt_oss_20b, "GPT_OSS_20B_PATH")
+ ->UseRealTime()->Unit(benchmark::kMillisecond)->Apply(MlpGateThreadgroupSizeArguments);
+BENCHMARK_CAPTURE(mlp_gate_tgsize, gpt_oss_120b, "GPT_OSS_120B_PATH")
+ ->UseRealTime()->Unit(benchmark::kMillisecond)->Apply(MlpGateThreadgroupSizeArguments);
+
+static void mlp_swiglu_tgsize(benchmark::State& state, const char* env_var_name) {
+ const char* model_path = getenv(env_var_name);
+ if (model_path == NULL) {
+ state.SkipWithError(std::format("environment variable {} is not set", env_var_name));
+ return;
+ }
+
+ gptoss_model_t model_ptr = nullptr;
+ gptoss_status status = gptoss_model_create_from_file(model_path, &model_ptr);
+ if (status != gptoss_status_success) {
+ state.SkipWithError(std::format("failed to load model from file {}", model_path));
+ return;
+ }
+ std::unique_ptr, decltype(&gptoss_model_release)> model(model_ptr, gptoss_model_release);
+ model->mlp_swiglu_threadgroup_size = static_cast(state.range(0));
+
+ gptoss_context_t context_ptr = nullptr;
+ status = gptoss_context_create(model.get(), /*context_length=*/0, /*max_batch_tokens=*/0, &context_ptr);
+ if (status != gptoss_status_success) {
+ state.SkipWithError("failed to create Context object");
+ return;
+ }
+ std::unique_ptr, decltype(&gptoss_context_release)> context(context_ptr, gptoss_context_release);
+
+ const char* prompt = "why did the chicken cross the road?";
+ std::size_t num_prompt_tokens = 0;
+ status = gptoss_context_append_chars(context.get(), prompt, strlen(prompt), &num_prompt_tokens);
+ if (status != gptoss_status_success) {
+ state.SkipWithError(std::format("failed to tokenize prompt \"{}\"", prompt));
+ return;
+ }
+
+ // Prefill
+ status = gptoss_context_process(context.get());
+ if (status != gptoss_status_success) {
+ state.SkipWithError("failed to prefill Context object");
+ return;
+ }
+ const std::size_t num_kvcache_tokens = context->num_kv_tokens;
+
+ std::uint64_t rng_seed = 0;
+ for (auto _ : state) {
+ const std::uint64_t current_rng_seed = rng_seed++;
+ context->num_kv_tokens = num_prompt_tokens;
+ context->num_tokens = num_prompt_tokens;
+
+ std::array tokens;
+ std::size_t num_generated_tokens = 0;
+ do {
+ std::size_t num_current_generated_tokens = 0;
+ status = gptoss_context_sample(context.get(), /*temperature=*/1.0f, /*rng_state=*/current_rng_seed,
+ /*max_tokens=*/kNumGeneratedTokens - num_generated_tokens, tokens.data(), &num_current_generated_tokens);
+ if (status != gptoss_status_success) {
+ state.SkipWithError("failed to sample from the Context object");
+ return;
+ }
+ num_generated_tokens += num_current_generated_tokens;
+ } while (num_generated_tokens < kNumGeneratedTokens);
+ }
+
+ state.counters["generations"] =
+ benchmark::Counter(state.iterations(), benchmark::Counter::kIsRate);
+ state.counters["tokens"] =
+ benchmark::Counter(state.iterations() * kNumGeneratedTokens, benchmark::Counter::kIsRate);
+}
+
+static void MlpSwigluThreadgroupSizeArguments(benchmark::internal::Benchmark* b) {
+ b->ArgNames({"tgsize"});
+ for (auto threadgroup_size = 64; threadgroup_size <= 1024; threadgroup_size += 64) {
+ const auto num_simdgroups = threadgroup_size / 32;
+ if (5760 % num_simdgroups != 0) {
+ // Skip incompatible threadgroup sizes
+ continue;
+ }
+ b->Args({threadgroup_size});
+ }
+}
+
+BENCHMARK_CAPTURE(mlp_swiglu_tgsize, gpt_oss_20b, "GPT_OSS_20B_PATH")
+ ->UseRealTime()->Unit(benchmark::kMillisecond)->Apply(MlpSwigluThreadgroupSizeArguments);
+BENCHMARK_CAPTURE(mlp_swiglu_tgsize, gpt_oss_120b, "GPT_OSS_120B_PATH")
+ ->UseRealTime()->Unit(benchmark::kMillisecond)->Apply(MlpSwigluThreadgroupSizeArguments);
+
+static void mlp_out_tgsize(benchmark::State& state, const char* env_var_name) {
+ const char* model_path = getenv(env_var_name);
+ if (model_path == NULL) {
+ state.SkipWithError(std::format("environment variable {} is not set", env_var_name));
+ return;
+ }
+
+ gptoss_model_t model_ptr = nullptr;
+ gptoss_status status = gptoss_model_create_from_file(model_path, &model_ptr);
+ if (status != gptoss_status_success) {
+ state.SkipWithError(std::format("failed to load model from file {}", model_path));
+ return;
+ }
+ std::unique_ptr, decltype(&gptoss_model_release)> model(model_ptr, gptoss_model_release);
+ model->mlp_out_threadgroup_size = static_cast(state.range(0));
+
+ gptoss_context_t context_ptr = nullptr;
+ status = gptoss_context_create(model.get(), /*context_length=*/0, /*max_batch_tokens=*/0, &context_ptr);
+ if (status != gptoss_status_success) {
+ state.SkipWithError("failed to create Context object");
+ return;
+ }
+ std::unique_ptr, decltype(&gptoss_context_release)> context(context_ptr, gptoss_context_release);
+
+ const char* prompt = "why did the chicken cross the road?";
+ std::size_t num_prompt_tokens = 0;
+ status = gptoss_context_append_chars(context.get(), prompt, strlen(prompt), &num_prompt_tokens);
+ if (status != gptoss_status_success) {
+ state.SkipWithError(std::format("failed to tokenize prompt \"{}\"", prompt));
+ return;
+ }
+
+ // Prefill
+ status = gptoss_context_process(context.get());
+ if (status != gptoss_status_success) {
+ state.SkipWithError("failed to prefill Context object");
+ return;
+ }
+ const std::size_t num_kvcache_tokens = context->num_kv_tokens;
+
+ std::uint64_t rng_seed = 0;
+ for (auto _ : state) {
+ const std::uint64_t current_rng_seed = rng_seed++;
+ context->num_kv_tokens = num_prompt_tokens;
+ context->num_tokens = num_prompt_tokens;
+
+ std::array tokens;
+ std::size_t num_generated_tokens = 0;
+ do {
+ std::size_t num_current_generated_tokens = 0;
+ status = gptoss_context_sample(context.get(), /*temperature=*/1.0f, /*rng_state=*/current_rng_seed,
+ /*max_tokens=*/kNumGeneratedTokens - num_generated_tokens, tokens.data(), &num_current_generated_tokens);
+ if (status != gptoss_status_success) {
+ state.SkipWithError("failed to sample from the Context object");
+ return;
+ }
+ num_generated_tokens += num_current_generated_tokens;
+ } while (num_generated_tokens < kNumGeneratedTokens);
+ }
+
+ state.counters["generations"] =
+ benchmark::Counter(state.iterations(), benchmark::Counter::kIsRate);
+ state.counters["tokens"] =
+ benchmark::Counter(state.iterations() * kNumGeneratedTokens, benchmark::Counter::kIsRate);
+}
+
+static void MlpOutThreadgroupSizeArguments(benchmark::internal::Benchmark* b) {
+ b->ArgNames({"tgsize"});
+ for (auto threadgroup_size = 64; threadgroup_size <= 1024; threadgroup_size += 64) {
+ const auto num_simdgroups = threadgroup_size / 32;
+ if (5760 % num_simdgroups != 0) {
+ // Skip incompatible threadgroup sizes
+ continue;
+ }
+ b->Args({threadgroup_size});
+ }
+}
+
+BENCHMARK_CAPTURE(mlp_out_tgsize, gpt_oss_20b, "GPT_OSS_20B_PATH")
+ ->UseRealTime()->Unit(benchmark::kMillisecond)->Apply(MlpOutThreadgroupSizeArguments);
+BENCHMARK_CAPTURE(mlp_out_tgsize, gpt_oss_120b, "GPT_OSS_120B_PATH")
+ ->UseRealTime()->Unit(benchmark::kMillisecond)->Apply(MlpOutThreadgroupSizeArguments);
+
+static void mlp_acc_tgsize(benchmark::State& state, const char* env_var_name) {
+ const char* model_path = getenv(env_var_name);
+ if (model_path == NULL) {
+ state.SkipWithError(std::format("environment variable {} is not set", env_var_name));
+ return;
+ }
+
+ gptoss_model_t model_ptr = nullptr;
+ gptoss_status status = gptoss_model_create_from_file(model_path, &model_ptr);
+ if (status != gptoss_status_success) {
+ state.SkipWithError(std::format("failed to load model from file {}", model_path));
+ return;
+ }
+ std::unique_ptr, decltype(&gptoss_model_release)> model(model_ptr, gptoss_model_release);
+ model->mlp_acc_threadgroup_size = static_cast(state.range(0));
+
+ gptoss_context_t context_ptr = nullptr;
+ status = gptoss_context_create(model.get(), /*context_length=*/0, /*max_batch_tokens=*/0, &context_ptr);
+ if (status != gptoss_status_success) {
+ state.SkipWithError("failed to create Context object");
+ return;
+ }
+ std::unique_ptr, decltype(&gptoss_context_release)> context(context_ptr, gptoss_context_release);
+
+ const char* prompt = "why did the chicken cross the road?";
+ std::size_t num_prompt_tokens = 0;
+ status = gptoss_context_append_chars(context.get(), prompt, strlen(prompt), &num_prompt_tokens);
+ if (status != gptoss_status_success) {
+ state.SkipWithError(std::format("failed to tokenize prompt \"{}\"", prompt));
+ return;
+ }
+
+ // Prefill
+ status = gptoss_context_process(context.get());
+ if (status != gptoss_status_success) {
+ state.SkipWithError("failed to prefill Context object");
+ return;
+ }
+ const std::size_t num_kvcache_tokens = context->num_kv_tokens;
+
+ std::uint64_t rng_seed = 0;
+ for (auto _ : state) {
+ const std::uint64_t current_rng_seed = rng_seed++;
+ context->num_kv_tokens = num_prompt_tokens;
+ context->num_tokens = num_prompt_tokens;
+
+ std::array tokens;
+ std::size_t num_generated_tokens = 0;
+ do {
+ std::size_t num_current_generated_tokens = 0;
+ status = gptoss_context_sample(context.get(), /*temperature=*/1.0f, /*rng_state=*/current_rng_seed,
+ /*max_tokens=*/kNumGeneratedTokens - num_generated_tokens, tokens.data(), &num_current_generated_tokens);
+ if (status != gptoss_status_success) {
+ state.SkipWithError("failed to sample from the Context object");
+ return;
+ }
+ num_generated_tokens += num_current_generated_tokens;
+ } while (num_generated_tokens < kNumGeneratedTokens);
+ }
+
+ state.counters["generations"] =
+ benchmark::Counter(state.iterations(), benchmark::Counter::kIsRate);
+ state.counters["tokens"] =
+ benchmark::Counter(state.iterations() * kNumGeneratedTokens, benchmark::Counter::kIsRate);
+}
+
+static void MlpAccThreadgroupSizeArguments(benchmark::internal::Benchmark* b) {
+ b->ArgNames({"tgsize"});
+ for (auto threadgroup_size = 32; threadgroup_size <= 1024; threadgroup_size += 32) {
+ b->Args({threadgroup_size});
+ }
+}
+
+BENCHMARK_CAPTURE(mlp_acc_tgsize, gpt_oss_20b, "GPT_OSS_20B_PATH")
+ ->UseRealTime()->Unit(benchmark::kMillisecond)->Apply(MlpAccThreadgroupSizeArguments);
+BENCHMARK_CAPTURE(mlp_acc_tgsize, gpt_oss_120b, "GPT_OSS_120B_PATH")
+ ->UseRealTime()->Unit(benchmark::kMillisecond)->Apply(MlpAccThreadgroupSizeArguments);
+
+static void unembedding_tgsize(benchmark::State& state, const char* env_var_name) {
+ const char* model_path = getenv(env_var_name);
+ if (model_path == NULL) {
+ state.SkipWithError(std::format("environment variable {} is not set", env_var_name));
+ return;
+ }
+
+ gptoss_model_t model_ptr = nullptr;
+ gptoss_status status = gptoss_model_create_from_file(model_path, &model_ptr);
+ if (status != gptoss_status_success) {
+ state.SkipWithError(std::format("failed to load model from file {}", model_path));
+ return;
+ }
+ std::unique_ptr, decltype(&gptoss_model_release)> model(model_ptr, gptoss_model_release);
+ model->unembedding_threadgroup_size = static_cast(state.range(0));
+
+ gptoss_context_t context_ptr = nullptr;
+ status = gptoss_context_create(model.get(), /*context_length=*/0, /*max_batch_tokens=*/0, &context_ptr);
+ if (status != gptoss_status_success) {
+ state.SkipWithError("failed to create Context object");
+ return;
+ }
+ std::unique_ptr, decltype(&gptoss_context_release)> context(context_ptr, gptoss_context_release);
+
+ const char* prompt = "why did the chicken cross the road?";
+ std::size_t num_prompt_tokens = 0;
+ status = gptoss_context_append_chars(context.get(), prompt, strlen(prompt), &num_prompt_tokens);
+ if (status != gptoss_status_success) {
+ state.SkipWithError(std::format("failed to tokenize prompt \"{}\"", prompt));
+ return;
+ }
+
+ // Prefill
+ status = gptoss_context_process(context.get());
+ if (status != gptoss_status_success) {
+ state.SkipWithError("failed to prefill Context object");
+ return;
+ }
+ const std::size_t num_kvcache_tokens = context->num_kv_tokens;
+
+ std::uint64_t rng_seed = 0;
+ for (auto _ : state) {
+ const std::uint64_t current_rng_seed = rng_seed++;
+ context->num_kv_tokens = num_prompt_tokens;
+ context->num_tokens = num_prompt_tokens;
+
+ std::array tokens;
+ std::size_t num_generated_tokens = 0;
+ do {
+ std::size_t num_current_generated_tokens = 0;
+ status = gptoss_context_sample(context.get(), /*temperature=*/1.0f, /*rng_state=*/current_rng_seed,
+ /*max_tokens=*/kNumGeneratedTokens - num_generated_tokens, tokens.data(), &num_current_generated_tokens);
+ if (status != gptoss_status_success) {
+ state.SkipWithError("failed to sample from the Context object");
+ return;
+ }
+ num_generated_tokens += num_current_generated_tokens;
+ } while (num_generated_tokens < kNumGeneratedTokens);
+ }
+
+ state.counters["generations"] =
+ benchmark::Counter(state.iterations(), benchmark::Counter::kIsRate);
+ state.counters["tokens"] =
+ benchmark::Counter(state.iterations() * kNumGeneratedTokens, benchmark::Counter::kIsRate);
+}
+
+static void UnembeddingThreadgroupSizeArguments(benchmark::internal::Benchmark* b) {
+ b->ArgNames({"tgsize"});
+ for (auto threadgroup_size = 32; threadgroup_size <= 1024; threadgroup_size += 32) {
+ b->Args({threadgroup_size});
+ }
+}
+
+BENCHMARK_CAPTURE(unembedding_tgsize, gpt_oss_20b, "GPT_OSS_20B_PATH")
+ ->UseRealTime()->Unit(benchmark::kMillisecond)->Apply(UnembeddingThreadgroupSizeArguments);
+BENCHMARK_CAPTURE(unembedding_tgsize, gpt_oss_120b, "GPT_OSS_120B_PATH")
+ ->UseRealTime()->Unit(benchmark::kMillisecond)->Apply(UnembeddingThreadgroupSizeArguments);
+
+BENCHMARK_MAIN();
diff --git a/gpt_oss/metal/benchmark/end-to-end.cc b/gpt_oss/metal/benchmark/end-to-end.cc
new file mode 100644
index 0000000000000000000000000000000000000000..ba87e3af7928de00860af7a3dc4f9dec2c65d734
--- /dev/null
+++ b/gpt_oss/metal/benchmark/end-to-end.cc
@@ -0,0 +1,227 @@
+#include
+#include
+
+#include
+#include
+#include
+#include
+#include
+#include
+#include
+#include
+#include
+
+#include
+
+constexpr std::uint32_t kNumGeneratedTokens = 100;
+
+static void end2end_decode(benchmark::State& state, const char* env_var_name) {
+ const char* model_path = getenv(env_var_name);
+ if (model_path == NULL) {
+ state.SkipWithError(std::format("environment variable {} is not set", env_var_name));
+ return;
+ }
+
+ gptoss_model_t model_ptr = nullptr;
+ gptoss_status status = gptoss_model_create_from_file(model_path, &model_ptr);
+ if (status != gptoss_status_success) {
+ state.SkipWithError(std::format("failed to load model from file {}", model_path));
+ return;
+ }
+ std::unique_ptr, decltype(&gptoss_model_release)> model(model_ptr, gptoss_model_release);
+
+ gptoss_context_t context_ptr = nullptr;
+ status = gptoss_context_create(model.get(), /*context_length=*/0, /*max_batch_tokens=*/0, &context_ptr);
+ if (status != gptoss_status_success) {
+ state.SkipWithError("failed to create Context object");
+ return;
+ }
+ std::unique_ptr, decltype(&gptoss_context_release)> context(context_ptr, gptoss_context_release);
+
+ const char* prompt = "why did the chicken cross the road?";
+ std::size_t num_prompt_tokens = 0;
+ status = gptoss_context_append_chars(context.get(), prompt, strlen(prompt), &num_prompt_tokens);
+ if (status != gptoss_status_success) {
+ state.SkipWithError(std::format("failed to tokenize prompt \"{}\"", prompt));
+ return;
+ }
+
+ // Prefill
+ status = gptoss_context_process(context.get());
+ if (status != gptoss_status_success) {
+ state.SkipWithError("failed to prefill Context object");
+ return;
+ }
+ std::uint64_t rng_seed = 0;
+
+ for (auto _ : state) {
+ const std::uint64_t current_rng_seed = rng_seed++;
+ context->num_kv_tokens = num_prompt_tokens;
+ context->num_tokens = num_prompt_tokens;
+
+ std::array tokens;
+ std::size_t num_generated_tokens = 0;
+ do {
+ std::size_t num_current_generated_tokens = 0;
+ status = gptoss_context_sample(context.get(), /*temperature=*/1.0f, /*rng_state=*/current_rng_seed,
+ /*max_tokens=*/kNumGeneratedTokens - num_generated_tokens, tokens.data(), &num_current_generated_tokens);
+ if (status != gptoss_status_success) {
+ state.SkipWithError("failed to sample from the Context object");
+ return;
+ }
+ num_generated_tokens += num_current_generated_tokens;
+ } while (num_generated_tokens < kNumGeneratedTokens);
+ }
+
+ state.counters["generations"] =
+ benchmark::Counter(state.iterations(), benchmark::Counter::kIsRate);
+ state.counters["tokens"] =
+ benchmark::Counter(state.iterations() * kNumGeneratedTokens, benchmark::Counter::kIsRate);
+}
+
+static void end2end_prefill(benchmark::State& state,
+ const char* model_path_env_var_name,
+ const char* prompt_env_var_name,
+ size_t context_length = 0) {
+ const char* model_path = getenv(model_path_env_var_name);
+ if (model_path == NULL) {
+ state.SkipWithError(std::format("environment variable {} is not set",
+ model_path_env_var_name));
+ return;
+ }
+
+ const char* prompt_file_path = getenv(prompt_env_var_name);
+ if (prompt_file_path == NULL) {
+ state.SkipWithError(std::format("environment variable {} is not set",
+ prompt_env_var_name));
+ return;
+ }
+
+ // Read prompt contents from file into a std::string
+ std::ifstream prompt_file(prompt_file_path,
+ std::ios::in | std::ios::binary);
+ if (!prompt_file) {
+ state.SkipWithError(
+ std::format("failed to open prompt file {}", prompt_file_path));
+ return;
+ }
+ std::string prompt_str;
+ prompt_file.seekg(0, std::ios::end);
+ std::streampos file_size = prompt_file.tellg();
+ if (file_size < 0) {
+ state.SkipWithError(std::format("failed to read prompt file size {}",
+ prompt_file_path));
+ return;
+ }
+ prompt_str.resize(static_cast(file_size));
+ prompt_file.seekg(0, std::ios::beg);
+ if (file_size > 0) {
+ prompt_file.read(prompt_str.data(), file_size);
+ }
+ if (!prompt_file) {
+ state.SkipWithError(
+ std::format("failed to read prompt file {}", prompt_file_path));
+ return;
+ }
+
+ gptoss_model_t model_ptr = nullptr;
+ gptoss_status status = gptoss_model_create_from_file(model_path, &model_ptr);
+ if (status != gptoss_status_success) {
+ state.SkipWithError(
+ std::format("failed to load model from file {}", model_path));
+ return;
+ }
+ std::unique_ptr,
+ decltype(&gptoss_model_release)>
+ model(model_ptr, gptoss_model_release);
+
+ gptoss_tokenizer_t tokenizer_ptr = nullptr;
+ status = gptoss_model_get_tokenizer(model.get(), &tokenizer_ptr);
+ if (status != gptoss_status_success) {
+ state.SkipWithError("failed to retrieve Tokenizer");
+ return;
+ }
+ std::unique_ptr,
+ decltype(&gptoss_tokenizer_release)>
+ tokenizer(tokenizer_ptr, gptoss_tokenizer_release);
+
+ gptoss_context_t context_ptr = nullptr;
+ status = gptoss_context_create(model.get(),
+ /*context_lenght=*/0,
+ /*max_batch_tokens=*/1024,
+ &context_ptr);
+ if (status != gptoss_status_success) {
+ state.SkipWithError("failed to create Context object");
+ return;
+ }
+ std::unique_ptr,
+ decltype(&gptoss_context_release)>
+ context(context_ptr, gptoss_context_release);
+
+ const char* prompt = prompt_str.c_str();
+ status = gptoss_context_append_chars(context.get(), prompt,
+ prompt_str.size(), nullptr);
+ if (status != gptoss_status_success) {
+ state.SkipWithError(std::format(
+ "failed to tokenize prompt from file {}", prompt_file_path));
+ return;
+ }
+
+ size_t num_tokens;
+ status = gptoss_context_get_num_tokens(context.get(), &num_tokens);
+ if (status != gptoss_status_success) {
+ state.SkipWithError("failed to get number of tokens");
+ return;
+ }
+ if (context_length != 0) {
+ assert(context_length <= num_tokens);
+ context->num_tokens = context_length;
+ }
+ status = gptoss_context_get_num_tokens(context.get(), &num_tokens);
+ if (status != gptoss_status_success) {
+ state.SkipWithError("failed to get number of tokens");
+ return;
+ }
+ // Prefill
+ for (auto _ : state) {
+ status = gptoss_context_process(context.get());
+ if (status != gptoss_status_success) {
+ state.SkipWithError("failed to prefill Context object");
+ return;
+ }
+ context->num_kv_tokens = 0;
+ }
+
+ state.counters["tokens"] = num_tokens;
+ state.counters["tokens/s"] = benchmark::Counter(
+ state.iterations() * num_tokens, benchmark::Counter::kIsRate);
+}
+
+// Decode end-to-end benchmark
+BENCHMARK_CAPTURE(end2end_decode, gpt_oss_20b_decode, "GPT_OSS_20B_PATH")
+ ->UseRealTime()
+ ->Unit(benchmark::kMillisecond);
+BENCHMARK_CAPTURE(end2end_decode, gpt_oss_120b_decode, "GPT_OSS_120B_PATH")
+ ->UseRealTime()
+ ->Unit(benchmark::kMillisecond);
+
+// Prefill end-to-end benchmark
+BENCHMARK_CAPTURE(end2end_prefill, gpt_oss_120b_prefill_1024,
+ "GPT_OSS_120B_PATH", "GPT_OSS_PROMPT_FILE_PATH", 1024)
+ ->UseRealTime()
+ ->Unit(benchmark::kMillisecond);
+BENCHMARK_CAPTURE(end2end_prefill, gpt_oss_20b_prefill_1024, "GPT_OSS_20B_PATH",
+ "GPT_OSS_PROMPT_FILE_PATH", 1024)
+ ->UseRealTime()
+ ->Unit(benchmark::kMillisecond);
+
+BENCHMARK_CAPTURE(end2end_prefill, gpt_oss_120b_prefill_3072,
+ "GPT_OSS_120B_PATH", "GPT_OSS_PROMPT_FILE_PATH", 3072)
+ ->UseRealTime()
+ ->Unit(benchmark::kMillisecond);
+BENCHMARK_CAPTURE(end2end_prefill, gpt_oss_20b_prefill_3072, "GPT_OSS_20B_PATH",
+ "GPT_OSS_PROMPT_FILE_PATH", 3072)
+ ->UseRealTime()
+ ->Unit(benchmark::kMillisecond);
+
+BENCHMARK_MAIN();
diff --git a/gpt_oss/metal/benchmark/f32-bf16w-rmsnorm.cc b/gpt_oss/metal/benchmark/f32-bf16w-rmsnorm.cc
new file mode 100644
index 0000000000000000000000000000000000000000..ee7551c252a0200e79d34d11c165b8718b3b2221
--- /dev/null
+++ b/gpt_oss/metal/benchmark/f32-bf16w-rmsnorm.cc
@@ -0,0 +1,99 @@
+#include
+#include
+#include
+#include
+
+#include
+
+#include
+
+using gptoss::Check;
+using namespace gptoss::metal;
+
+constexpr float kEpsilon = 1.0e-5f;
+constexpr uint64_t kSeed = UINT64_C(1019827666124465388);
+
+static void f32_bf16w_rnsnorm(benchmark::State& state) {
+ const size_t num_tokens = 1;
+ const size_t num_channels = state.range(0);
+
+ Device device;
+ CommandQueue command_queue{device};
+ Library library{device};
+ Function f32_fill_random_fn{library, "gptoss_f32_fill_random"};
+ Function bf16_fill_random_fn{library, "gptoss_bf16_fill_random"};
+ Function f32_bf16w_rmsnorm_fn{library, "gptoss_f32_bf16w_rmsnorm"};
+ Buffer input_buffer{device, num_tokens * num_channels * sizeof(float)};
+ Buffer weight_buffer{device, num_channels * sizeof(gptoss_bfloat16)};
+ Buffer output_buffer{device, num_tokens * num_channels * sizeof(float)};
+ Buffer control_buffer{device, sizeof(gptoss_control)};
+ std::memset(control_buffer.ptr(), 0, sizeof(gptoss_control));
+
+ {
+ CommandBuffer command_buffer{command_queue};
+
+ size_t offset = 0;
+ Check(gptoss_metal_command_buffer_encode_launch_f32_fill_random(
+ command_buffer.handle(),
+ f32_fill_random_fn.handle(),
+ /*threadgroup_size=*/0,
+ /*max_threadgroups=*/10,
+ /*output_buffer=*/input_buffer.handle(),
+ /*output_offset=*/0,
+ num_channels, kSeed, offset, /*min=*/-1.0f, /*max=*/1.0),
+ "gptoss_metal_command_buffer_encode_launch_f32_fill_random");
+ offset += num_channels;
+
+ Check(gptoss_metal_command_buffer_encode_launch_bf16_fill_random(
+ command_buffer.handle(),
+ bf16_fill_random_fn.handle(),
+ /*threadgroup_size=*/0,
+ /*max_threadgroups=*/10,
+ /*output_buffer=*/weight_buffer.handle(),
+ /*output_offset=*/0,
+ num_channels, kSeed, offset, /*min=*/-1.0f, /*max=*/1.0),
+ "gptoss_metal_command_buffer_encode_launch_bf16_fill_random");
+ offset += num_channels;
+
+ command_buffer.commit();
+ command_buffer.wait_completion();
+ }
+
+ for (auto _ : state) {
+ CommandBuffer command_buffer{command_queue};
+
+ Check(gptoss_metal_command_buffer_encode_launch_f32_bf16w_rmsnorm(
+ command_buffer.handle(),
+ f32_bf16w_rmsnorm_fn.handle(),
+ input_buffer.handle(),
+ /*input_offset=*/0,
+ weight_buffer.handle(),
+ /*weight_offset=*/0,
+ output_buffer.handle(),
+ /*output_offset=*/0,
+ control_buffer.handle(),
+ /*control_offset=*/0,
+ num_tokens,
+ num_channels,
+ kEpsilon),
+ "gptoss_metal_command_buffer_encode_launch_f32_bf16w_rmsnorm");
+
+ command_buffer.commit();
+ const double elapsed_seconds = command_buffer.wait_completion();
+ state.SetIterationTime(elapsed_seconds);
+ }
+
+ const size_t num_elements = num_tokens * num_channels;
+ state.counters["elements"] =
+ benchmark::Counter(state.iterations() * num_elements,
+ benchmark::Counter::kIsRate);
+
+ const int64_t bytes_per_iteration = input_buffer.size() + weight_buffer.size() + output_buffer.size();
+ state.counters["bytes"] =
+ benchmark::Counter(state.iterations() * bytes_per_iteration,
+ benchmark::Counter::kIsRate);
+}
+
+BENCHMARK(f32_bf16w_rnsnorm)->Arg(2880)->UseManualTime()->Unit(benchmark::kMicrosecond);
+
+BENCHMARK_MAIN();
diff --git a/gpt_oss/metal/benchmark/f32-random.cc b/gpt_oss/metal/benchmark/f32-random.cc
new file mode 100644
index 0000000000000000000000000000000000000000..62455e3460f0ec32b51bf5c8d4771d3de6cff73f
--- /dev/null
+++ b/gpt_oss/metal/benchmark/f32-random.cc
@@ -0,0 +1,55 @@
+#include
+#include
+#include
+
+#include
+
+using gptoss::Check;
+using namespace gptoss::metal;
+
+static void f32_fill_random(benchmark::State& state) {
+ const size_t numel = state.range(0);
+
+ Device device;
+ CommandQueue command_queue{device};
+ Library library{device};
+ Function f32_fill_random_fn{library, "gptoss_f32_fill_random"};
+ Buffer buffer{device, numel * sizeof(float)};
+
+ constexpr uint64_t seed = UINT64_C(1019827666124465388);
+ constexpr uint64_t offset = UINT64_C(12345678901234567890);
+ const float min = -1.0f;
+ const float max = 7.0f;
+ for (auto _ : state) {
+ CommandBuffer command_buffer{command_queue};
+
+ Check(gptoss_metal_command_buffer_encode_launch_f32_fill_random(
+ command_buffer.handle(),
+ f32_fill_random_fn.handle(),
+ /*threadgroup_size=*/0,
+ /*max_threadgroups=*/120,
+ /*output_buffer=*/buffer.handle(),
+ /*output_offset=*/0,
+ numel, seed, offset, min, max),
+ "gptoss_metal_command_buffer_encode_launch_f32_fill_random");
+
+ command_buffer.commit();
+ const double elapsed_seconds = command_buffer.wait_completion();
+ state.SetIterationTime(elapsed_seconds);
+ }
+
+ const int64_t elements_per_iteration = numel;
+ state.counters["elements"] =
+ benchmark::Counter(state.iterations() * elements_per_iteration,
+ benchmark::Counter::kIsRate);
+
+ const int64_t bytes_per_iteration = numel * sizeof(float);
+ state.counters["bytes"] =
+ benchmark::Counter(state.iterations() * bytes_per_iteration,
+ benchmark::Counter::kIsRate);
+}
+
+constexpr int64_t giga = INT64_C(1073741824);
+BENCHMARK(f32_fill_random)->Arg(2 * giga)->UseManualTime()->Unit(benchmark::kMicrosecond);
+
+BENCHMARK_MAIN();
diff --git a/gpt_oss/metal/benchmark/mf4-f32-convert.cc b/gpt_oss/metal/benchmark/mf4-f32-convert.cc
new file mode 100644
index 0000000000000000000000000000000000000000..dff5871814abbb895773b570edb899ab84a611f8
--- /dev/null
+++ b/gpt_oss/metal/benchmark/mf4-f32-convert.cc
@@ -0,0 +1,65 @@
+#include
+#include
+#include
+#include
+
+#include
+
+#include
+
+using gptoss::Check;
+using namespace gptoss::metal;
+
+static void mf4_f32_convert(benchmark::State& state) {
+ const size_t num_blocks = state.range(0);
+ const size_t num_elements = num_blocks * 32;
+ const size_t num_bytes = num_elements / 2;
+
+ Device device;
+ CommandQueue command_queue{device};
+ Library library{device};
+ Function mf4_f32_convert_fn{library, "gptoss_mf4_f32_convert"};
+ Buffer block_buffer{device, num_bytes};
+ Buffer scale_buffer{device, num_blocks * sizeof(gptoss_float8ue8m0)};
+ Buffer output_buffer{device, num_elements * sizeof(float)};
+
+ std::memset(block_buffer.ptr(), 0x91, num_bytes); // force subnormals
+ std::memset(scale_buffer.ptr(), 128, num_blocks * sizeof(uint8_t)); // scale = 2.0
+
+ for (auto _ : state) {
+ CommandBuffer command_buffer{command_queue};
+
+ Check(gptoss_metal_command_buffer_encode_launch_mf4_f32_convert(
+ command_buffer.handle(),
+ mf4_f32_convert_fn.handle(),
+ /*threadgroup_size=*/0,
+ /*max_threadgroups=*/120,
+ block_buffer.handle(),
+ scale_buffer.handle(),
+ output_buffer.handle(),
+ num_elements),
+ "gptoss_metal_command_buffer_encode_launch_mf4_f32_convert");
+
+ command_buffer.commit();
+ const double elapsed_seconds = command_buffer.wait_completion();
+ state.SetIterationTime(elapsed_seconds);
+ }
+
+ state.counters["blocks"] =
+ benchmark::Counter(state.iterations() * num_blocks,
+ benchmark::Counter::kIsRate);
+
+ state.counters["elements"] =
+ benchmark::Counter(state.iterations() * num_elements,
+ benchmark::Counter::kIsRate);
+
+ const int64_t bytes_per_iteration = num_bytes + num_blocks + num_elements * sizeof(float);
+ state.counters["bytes"] =
+ benchmark::Counter(state.iterations() * bytes_per_iteration,
+ benchmark::Counter::kIsRate);
+}
+
+constexpr int64_t mega = INT64_C(1048576);
+BENCHMARK(mf4_f32_convert)->Arg(256 * mega)->UseManualTime()->Unit(benchmark::kMicrosecond);
+
+BENCHMARK_MAIN();
diff --git a/gpt_oss/metal/benchmark/u32-random.cc b/gpt_oss/metal/benchmark/u32-random.cc
new file mode 100644
index 0000000000000000000000000000000000000000..1d73e2210d44300330e6dac16c1ab706b9181384
--- /dev/null
+++ b/gpt_oss/metal/benchmark/u32-random.cc
@@ -0,0 +1,53 @@
+#include
+#include
+#include
+
+#include
+
+using gptoss::Check;
+using namespace gptoss::metal;
+
+static void u32_fill_random(benchmark::State& state) {
+ const size_t numel = state.range(0);
+
+ Device device;
+ CommandQueue command_queue{device};
+ Library library{device};
+ Function u32_fill_random_fn{library, "gptoss_u32_fill_random"};
+ Buffer buffer{device, numel * sizeof(float)};
+
+ constexpr uint64_t seed = UINT64_C(1019827666124465388);
+ constexpr uint64_t offset = UINT64_C(12345678901234567890);
+ for (auto _ : state) {
+ CommandBuffer command_buffer{command_queue};
+
+ Check(gptoss_metal_command_buffer_encode_launch_u32_fill_random(
+ command_buffer.handle(),
+ u32_fill_random_fn.handle(),
+ /*threadgroup_size=*/0,
+ /*max_threadgroups=*/120,
+ /*output_buffer=*/buffer.handle(),
+ /*output_offset=*/0,
+ numel, seed, offset),
+ "gptoss_metal_command_buffer_encode_launch_u32_fill_random");
+
+ command_buffer.commit();
+ const double elapsed_seconds = command_buffer.wait_completion();
+ state.SetIterationTime(elapsed_seconds);
+ }
+
+ const int64_t elements_per_iteration = numel;
+ state.counters["elements"] =
+ benchmark::Counter(state.iterations() * elements_per_iteration,
+ benchmark::Counter::kIsRate);
+
+ const int64_t bytes_per_iteration = numel * sizeof(float);
+ state.counters["bytes"] =
+ benchmark::Counter(state.iterations() * bytes_per_iteration,
+ benchmark::Counter::kIsRate);
+}
+
+constexpr int64_t giga = INT64_C(1073741824);
+BENCHMARK(u32_fill_random)->Arg(2 * giga)->UseManualTime()->Unit(benchmark::kMicrosecond);
+
+BENCHMARK_MAIN();
diff --git a/gpt_oss/metal/examples/chat.py b/gpt_oss/metal/examples/chat.py
new file mode 100644
index 0000000000000000000000000000000000000000..f29cec3b5989b533a6e8a1d884e6ab217026a568
--- /dev/null
+++ b/gpt_oss/metal/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:])
diff --git a/gpt_oss/metal/examples/generate.py b/gpt_oss/metal/examples/generate.py
new file mode 100644
index 0000000000000000000000000000000000000000..3b78199994e66d8322a8cec980bed18b8098308c
--- /dev/null
+++ b/gpt_oss/metal/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:])
diff --git a/gpt_oss/metal/include/gpt-oss.h b/gpt_oss/metal/include/gpt-oss.h
new file mode 100644
index 0000000000000000000000000000000000000000..f24a7b1b34336b09bd64695f111321b519986cb5
--- /dev/null
+++ b/gpt_oss/metal/include/gpt-oss.h
@@ -0,0 +1,5 @@
+#pragma once
+
+#include
+#include
+#include
diff --git a/gpt_oss/metal/include/gpt-oss/functions.h b/gpt_oss/metal/include/gpt-oss/functions.h
new file mode 100644
index 0000000000000000000000000000000000000000..758756c7a69dbc0bc22ad6c6f4a80d62be08fffb
--- /dev/null
+++ b/gpt_oss/metal/include/gpt-oss/functions.h
@@ -0,0 +1,401 @@
+#pragma once
+
+#include
+#include
+
+#include
+#include
+
+#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
diff --git a/gpt_oss/metal/include/gpt-oss/macros.h b/gpt_oss/metal/include/gpt-oss/macros.h
new file mode 100644
index 0000000000000000000000000000000000000000..e1c7f34e10202b9537225c87ddbc48ea7ed37d21
--- /dev/null
+++ b/gpt_oss/metal/include/gpt-oss/macros.h
@@ -0,0 +1,5 @@
+#pragma once
+
+#ifndef GPTOSS_ABI
+ #define GPTOSS_ABI
+#endif // GPTOSS_ABI
diff --git a/gpt_oss/metal/include/gpt-oss/types.h b/gpt_oss/metal/include/gpt-oss/types.h
new file mode 100644
index 0000000000000000000000000000000000000000..f60aed5b2170a2217b671783a5d5cf79ee247612
--- /dev/null
+++ b/gpt_oss/metal/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;
diff --git a/gpt_oss/metal/python/context.c b/gpt_oss/metal/python/context.c
new file mode 100644
index 0000000000000000000000000000000000000000..0bf016006e60c7e962f4dbc8875c12947bf62320
--- /dev/null
+++ b/gpt_oss/metal/python/context.c
@@ -0,0 +1,295 @@
+#include
+
+#include
+
+#include "module.h"
+
+
+static int PyGPTOSSContext_init(PyGPTOSSContext* self, PyObject* args, PyObject* kwargs) {
+ static char *kwlist[] = {"model", "context_length", "max_batch_tokens", NULL};
+ PyObject* model = NULL;
+ Py_ssize_t context_length = 0; // Default to 0 if None
+ Py_ssize_t max_batch_tokens = 0; // Default to 0 if None
+
+ PyObject* model = NULL;
+ if (!PyArg_ParseTupleAndKeywords(args, kwargs, "O!|$ii", kwlist,
+ &PyGPTOSSModel_Type, &model,
+ &context_length, &max_batch_tokens))
+ {
+ return -1;
+ }
+ if (context_length < 0) {
+ PyErr_SetString(PyExc_ValueError, "context_length must be a positive integer");
+ return -1;
+ }
+ if (max_batch_tokens < 0) {
+ PyErr_SetString(PyExc_ValueError, "max_batch_tokens must be a positive integer");
+ return -1;
+ }
+
+ enum gptoss_status status = gptoss_context_create(
+ ((const PyGPTOSSModel*) model)->handle,
+ (size_t) context_length,
+ (size_t) max_batch_tokens,
+ &self->handle);
+ if (status != gptoss_status_success) {
+ // TODO: set exception
+ goto error;
+ }
+
+ return 0;
+
+error:
+ gptoss_context_release(self->handle);
+ self->handle = NULL;
+ return -1;
+}
+
+static void PyGPTOSSContext_dealloc(PyGPTOSSContext* self) {
+ (void) gptoss_context_release(self->handle);
+ self->handle = NULL;
+ PyObject_Del((PyObject*) self);
+}
+
+static PyObject* PyGPTOSSContext_copy(PyGPTOSSContext *self) {
+ PyGPTOSSContext* copy = (PyGPTOSSContext*) PyObject_New(PyGPTOSSContext, Py_TYPE(self));
+ if (copy == NULL) {
+ return NULL;
+ }
+
+ (void) gptoss_context_retain(self->handle);
+ copy->handle = self->handle;
+ return (PyObject*) copy;
+}
+
+static PyObject* PyGPTOSSContext_append(PyGPTOSSContext* self, PyObject* arg) {
+ if (PyBytes_Check(arg)) {
+ char* string_ptr = NULL;
+ Py_ssize_t string_size = 0;
+ if (PyBytes_AsStringAndSize(arg, &string_ptr, &string_size) < 0) {
+ return NULL;
+ }
+
+ const enum gptoss_status status = gptoss_context_append_chars(
+ self->handle, string_ptr, string_size, /*num_tokens_out=*/NULL);
+ if (status != gptoss_status_success) {
+ // TODO: set exception
+ return NULL;
+ }
+
+ Py_RETURN_NONE;
+ } else if (PyUnicode_Check(arg)) {
+ Py_ssize_t string_size = 0;
+ const char* string_ptr = PyUnicode_AsUTF8AndSize(arg, &string_size);
+ if (string_ptr == NULL) {
+ return NULL;
+ }
+
+ const enum gptoss_status status = gptoss_context_append_chars(
+ self->handle, string_ptr, string_size, /*num_tokens_out=*/NULL);
+ if (status != gptoss_status_success) {
+ // TODO: set exception
+ return NULL;
+ }
+
+ Py_RETURN_NONE;
+ } else if (PyLong_Check(arg)) {
+ const unsigned long token_as_ulong = PyLong_AsUnsignedLong(arg);
+ if (token_as_ulong == (unsigned long) -1 && PyErr_Occurred()) {
+ return NULL;
+ }
+
+ const uint32_t token = (uint32_t) token_as_ulong;
+ const enum gptoss_status status = gptoss_context_append_tokens(
+ self->handle, /*num_tokens=*/1, &token);
+ if (status != gptoss_status_success) {
+ // TODO: set exception
+ return NULL;
+ }
+
+ Py_RETURN_NONE;
+ } else {
+ PyErr_SetString(PyExc_TypeError, "expected a bytes or integer argument");
+ return NULL;
+ }
+}
+
+static PyObject* PyGPTOSSContext_process(PyGPTOSSContext* self) {
+ const enum gptoss_status status = gptoss_context_process(self->handle);
+ if (status != gptoss_status_success) {
+ // TODO: set exception
+ return NULL;
+ }
+
+ Py_RETURN_NONE;
+}
+
+static PyObject* PyGPTOSSContext_sample(PyGPTOSSContext* self, PyObject* args, PyObject* kwargs) {
+ static char *kwlist[] = {"max_output_tokens", "temperature", "seed", NULL};
+ PyObject* token_list_obj = NULL;
+ uint32_t* token_ptr = NULL;
+
+ unsigned int max_output_tokens = 0;
+ unsigned long long seed = 0;
+ float temperature = 1.0f;
+ if (!PyArg_ParseTupleAndKeywords(args, kwargs, "I|$fK", kwlist,
+ &max_output_tokens, &temperature, &seed))
+ {
+ return NULL;
+ }
+
+ token_ptr = (uint32_t*) PyMem_Malloc(max_output_tokens * sizeof(uint32_t));
+ if (token_ptr == NULL) {
+ goto error;
+ }
+
+ size_t num_tokens = 0;
+ const enum gptoss_status status = gptoss_context_sample(
+ self->handle, temperature, (uint64_t) seed,
+ (size_t) max_output_tokens, token_ptr, &num_tokens);
+ if (status != gptoss_status_success) {
+ // TODO: set exception
+ goto error;
+ }
+
+ token_list_obj = PyList_New((Py_ssize_t) num_tokens);
+ if (token_list_obj == NULL) {
+ goto error;
+ }
+
+ for (size_t t = 0; t < num_tokens; t++) {
+ PyObject* token_obj = PyLong_FromUnsignedLong((unsigned long) token_ptr[t]);
+ if (token_obj == NULL) {
+ goto error;
+ }
+
+ PyList_SET_ITEM(token_list_obj, (Py_ssize_t) t, token_obj);
+ }
+
+ PyMem_Free(token_ptr);
+ return token_list_obj;
+
+error:
+ PyMem_Free(token_ptr);
+ Py_XDECREF(token_list_obj);
+ return NULL;
+}
+
+static PyObject* PyGPTOSSContext_reset(PyGPTOSSContext* self) {
+ const enum gptoss_status status = gptoss_context_reset(self->handle);
+ if (status != gptoss_status_success) {
+ // TODO: set exception
+ return NULL;
+ }
+
+ Py_RETURN_NONE;
+}
+
+static PyMethodDef PyGPTOSSContext_methods[] = {
+ {"__copy__", (PyCFunction) PyGPTOSSContext_copy, METH_NOARGS, "Create a copy of the Context"},
+ {"append", (PyCFunction) PyGPTOSSContext_append, METH_O, "Append bytes to the Context"},
+ {"process", (PyCFunction) PyGPTOSSContext_process, METH_NOARGS, "Process tokens in the Context"},
+ {"sample", (PyCFunction) PyGPTOSSContext_sample, METH_VARARGS | METH_KEYWORDS, "Sample token predictions from the Context"},
+ {"reset", (PyCFunction) PyGPTOSSContext_reset, METH_NOARGS, "Discard the content of the Context"},
+ {NULL},
+};
+
+static PyObject* PyGPTOSSContext_get_num_tokens(PyGPTOSSContext* self, void* closure) {
+ size_t num_tokens = 0;
+ const enum gptoss_status status = gptoss_context_get_num_tokens(self->handle, &num_tokens);
+ if (status != gptoss_status_success) {
+ // TODO: set exception
+ return NULL;
+ }
+
+ return PyLong_FromSize_t(num_tokens);
+}
+
+static PyObject* PyGPTOSSContext_get_max_tokens(PyGPTOSSContext* self, void* closure) {
+ size_t max_tokens = 0;
+ const enum gptoss_status status = gptoss_context_get_max_tokens(self->handle, &max_tokens);
+ if (status != gptoss_status_success) {
+ // TODO: set exception
+ return NULL;
+ }
+
+ return PyLong_FromSize_t(max_tokens);
+}
+
+static PyObject* PyGPTOSSContext_get_tokens(PyGPTOSSContext* self, void* closure) {
+ PyObject* token_list_obj = NULL;
+ uint32_t* token_ptr = NULL;
+
+ size_t num_tokens = 0;
+ gptoss_context_get_tokens(self->handle, /*tokens_out=*/NULL, /*max_tokens=*/0, &num_tokens);
+
+ if (num_tokens != 0) {
+ token_ptr = (uint32_t*) PyMem_Malloc(num_tokens * sizeof(uint32_t));
+ if (token_ptr == NULL) {
+ // TODO: set exception
+ goto error;
+ }
+
+ enum gptoss_status status = gptoss_context_get_tokens(self->handle, token_ptr, /*max_tokens=*/num_tokens, &num_tokens);
+ if (status != gptoss_status_success) {
+ // TODO: set exception
+ goto error;
+ }
+ }
+
+ token_list_obj = PyList_New((Py_ssize_t) num_tokens);
+ if (token_list_obj == NULL) {
+ goto error;
+ }
+
+ for (size_t t = 0; t < num_tokens; t++) {
+ PyObject* token_obj = PyLong_FromUnsignedLong((unsigned long) token_ptr[t]);
+ if (token_obj == NULL) {
+ goto error;
+ }
+
+ PyList_SET_ITEM(token_list_obj, (Py_ssize_t) t, token_obj);
+ }
+
+ PyMem_Free(token_ptr);
+ return token_list_obj;
+
+error:
+ PyMem_Free(token_ptr);
+ Py_XDECREF(token_list_obj);
+ return NULL;
+}
+
+static PyGetSetDef PyGPTOSSContext_getseters[] = {
+ (PyGetSetDef) {
+ .name = "num_tokens",
+ .get = (getter) PyGPTOSSContext_get_num_tokens,
+ .doc = "Current number of tokens in the context",
+ },
+ (PyGetSetDef) {
+ .name = "max_tokens",
+ .get = (getter) PyGPTOSSContext_get_max_tokens,
+ .doc = "Maximum number of tokens in the context",
+ },
+ (PyGetSetDef) {
+ .name = "tokens",
+ .get = (getter) PyGPTOSSContext_get_tokens,
+ .doc = "List of token IDs in the context",
+ },
+ {NULL} /* Sentinel */
+};
+
+PyTypeObject PyGPTOSSContext_Type = {
+ PyVarObject_HEAD_INIT(NULL, 0)
+ .tp_name = "gptoss.Context",
+ .tp_basicsize = sizeof(PyGPTOSSContext),
+ .tp_flags = 0
+ | Py_TPFLAGS_DEFAULT
+ | Py_TPFLAGS_BASETYPE,
+ .tp_doc = "Context object",
+ .tp_methods = PyGPTOSSContext_methods,
+ .tp_getset = PyGPTOSSContext_getseters,
+ .tp_new = PyType_GenericNew,
+ .tp_init = (initproc) PyGPTOSSContext_init,
+ .tp_dealloc = (destructor) PyGPTOSSContext_dealloc,
+};
diff --git a/gpt_oss/metal/python/model.c b/gpt_oss/metal/python/model.c
new file mode 100644
index 0000000000000000000000000000000000000000..49202a2c5cbdd8e57f0f60a538fb8c296aa391e0
--- /dev/null
+++ b/gpt_oss/metal/python/model.c
@@ -0,0 +1,94 @@
+#include
+
+#include
+
+#include "module.h"
+
+
+static int PyGPTOSSModel_init(PyGPTOSSModel* self, PyObject* args, PyObject* kwargs) {
+ enum gptoss_status status;
+ const char* filepath;
+
+ if (!PyArg_ParseTuple(args, "s", &filepath)) {
+ return -1;
+ }
+ status = gptoss_model_create_from_file(filepath, &self->handle);
+ if (status != gptoss_status_success) {
+ // TODO: set exception
+ return -1;
+ }
+ return 0;
+}
+
+static void PyGPTOSSModel_dealloc(PyGPTOSSModel* self) {
+ (void) gptoss_model_release(self->handle);
+ self->handle = NULL;
+ PyObject_Del((PyObject*) self);
+}
+
+static PyObject* PyGPTOSSModel_copy(PyGPTOSSModel* self) {
+ PyGPTOSSModel* copy = (PyGPTOSSModel*) PyObject_New(PyGPTOSSModel, Py_TYPE(self));
+ if (copy == NULL) {
+ return NULL;
+ }
+
+ (void) gptoss_model_retain(self->handle);
+ copy->handle = self->handle;
+ return (PyObject*) copy;
+}
+
+static PyMethodDef PyGPTOSSModel_methods[] = {
+ {"__copy__", (PyCFunction) PyGPTOSSModel_copy, METH_NOARGS, "Create a copy of the Model"},
+ {NULL},
+};
+
+static PyObject *PyGPTOSSModel_get_max_context_length(PyGPTOSSModel* self, void* closure) {
+ size_t max_context_length = 0;
+ const enum gptoss_status status = gptoss_model_get_max_context_length(self->handle, &max_context_length);
+ if (status != gptoss_status_success) {
+ // TODO: set exception
+ return NULL;
+ }
+
+ return PyLong_FromSize_t(max_context_length);
+}
+
+static PyObject *PyGPTOSSModel_get_tokenizer(PyGPTOSSModel* self, void* closure) {
+ PyObject* args = PyTuple_Pack(1, self);
+ if (args == NULL) {
+ return NULL;
+ }
+
+ PyObject* tokenizer = PyObject_CallObject((PyObject*) &PyGPTOSSTokenizer_Type, args);
+ Py_DECREF(args);
+ return tokenizer;
+}
+
+static PyGetSetDef PyGPTOSSModel_getseters[] = {
+ (PyGetSetDef) {
+ .name = "max_context_length",
+ .get = (getter) PyGPTOSSModel_get_max_context_length,
+ .doc = "Maximum context length supported by the model",
+ },
+ (PyGetSetDef) {
+ .name = "tokenizer",
+ .get = (getter) PyGPTOSSModel_get_tokenizer,
+ .doc = "Tokenizer object associated with the model",
+ },
+ {NULL} // Sentinel
+};
+
+PyTypeObject PyGPTOSSModel_Type = {
+ PyVarObject_HEAD_INIT(NULL, 0)
+ .tp_name = "gptoss.Model",
+ .tp_basicsize = sizeof(PyGPTOSSModel),
+ .tp_flags = 0
+ | Py_TPFLAGS_DEFAULT
+ | Py_TPFLAGS_BASETYPE,
+ .tp_doc = "Model object",
+ .tp_methods = PyGPTOSSModel_methods,
+ .tp_getset = PyGPTOSSModel_getseters,
+ .tp_new = PyType_GenericNew,
+ .tp_init = (initproc) PyGPTOSSModel_init,
+ .tp_dealloc = (destructor) PyGPTOSSModel_dealloc,
+};
diff --git a/gpt_oss/metal/python/module.c b/gpt_oss/metal/python/module.c
new file mode 100644
index 0000000000000000000000000000000000000000..2910c8f1d0280981b1d2a04a1d0438f6faaa8d7e
--- /dev/null
+++ b/gpt_oss/metal/python/module.c
@@ -0,0 +1,67 @@
+#include
+
+#include "module.h"
+
+
+static PyMethodDef module_methods[] = {
+ {NULL, NULL, 0, NULL}
+};
+
+static PyModuleDef metal_module = {
+ PyModuleDef_HEAD_INIT,
+ "_metal",
+ "Local GPT-OSS inference",
+ -1,
+ module_methods
+};
+
+PyMODINIT_FUNC PyInit__metal(void) {
+ PyObject* module = NULL;
+ PyObject* model_type = NULL;
+ PyObject* tokenizer_type = NULL;
+ PyObject* context_type = NULL;
+
+ if (PyType_Ready(&PyGPTOSSModel_Type) < 0) {
+ goto error;
+ }
+ model_type = (PyObject*) &PyGPTOSSModel_Type;
+ Py_INCREF(model_type);
+
+ if (PyType_Ready(&PyGPTOSSTokenizer_Type) < 0) {
+ goto error;
+ }
+ tokenizer_type = (PyObject*) &PyGPTOSSTokenizer_Type;
+ Py_INCREF(tokenizer_type);
+
+ if (PyType_Ready(&PyGPTOSSContext_Type) < 0) {
+ goto error;
+ }
+ context_type = (PyObject*) &PyGPTOSSContext_Type;
+ Py_INCREF(context_type);
+
+ module = PyModule_Create(&metal_module);
+ if (module == NULL) {
+ goto error;
+ }
+
+ if (PyModule_AddObject(module, "Model", model_type) < 0) {
+ goto error;
+ }
+
+ if (PyModule_AddObject(module, "Tokenizer", tokenizer_type) < 0) {
+ goto error;
+ }
+
+ if (PyModule_AddObject(module, "Context", context_type) < 0) {
+ goto error;
+ }
+
+ return module;
+
+error:
+ Py_XDECREF(context_type);
+ Py_XDECREF(tokenizer_type);
+ Py_XDECREF(model_type);
+ Py_XDECREF(module);
+ return NULL;
+}
diff --git a/gpt_oss/metal/python/module.h b/gpt_oss/metal/python/module.h
new file mode 100644
index 0000000000000000000000000000000000000000..a837eee07ddb2a34f65e6648015a3b29d1066fea
--- /dev/null
+++ b/gpt_oss/metal/python/module.h
@@ -0,0 +1,22 @@
+#include
+
+#include
+
+typedef struct {
+ PyObject_HEAD
+ gptoss_model_t handle;
+} PyGPTOSSModel;
+
+typedef struct {
+ PyObject_HEAD
+ gptoss_tokenizer_t handle;
+} PyGPTOSSTokenizer;
+
+typedef struct {
+ PyObject_HEAD
+ gptoss_context_t handle;
+} PyGPTOSSContext;
+
+extern PyTypeObject PyGPTOSSModel_Type;
+extern PyTypeObject PyGPTOSSTokenizer_Type;
+extern PyTypeObject PyGPTOSSContext_Type;
diff --git a/gpt_oss/metal/python/tokenizer.c b/gpt_oss/metal/python/tokenizer.c
new file mode 100644
index 0000000000000000000000000000000000000000..3107dc6db5882b5dfabe36678054db6c31aa7c75
--- /dev/null
+++ b/gpt_oss/metal/python/tokenizer.c
@@ -0,0 +1,185 @@
+#include
+
+#include
+
+#include "module.h"
+
+static PyObject* PyGPTOSSTokenizer_new(PyTypeObject* subtype, PyObject* args, PyObject* kwargs) {
+ static char *kwlist[] = {"model", NULL};
+ PyObject* model = NULL;
+ if (!PyArg_ParseTupleAndKeywords(args, kwargs, "O!", kwlist, &PyGPTOSSModel_Type, &model)) {
+ return NULL;
+ }
+
+ PyGPTOSSTokenizer* self = (PyGPTOSSTokenizer*) subtype->tp_alloc(subtype, 0);
+ if (self == NULL) {
+ return NULL;
+ }
+
+ const enum gptoss_status status = gptoss_model_get_tokenizer(
+ ((const PyGPTOSSModel*) model)->handle,
+ &self->handle);
+ if (status != gptoss_status_success) {
+ // TODO: set exception
+ return NULL;
+ }
+
+ return (PyObject*) self;
+}
+
+static void PyGPTOSSTokenizer_dealloc(PyGPTOSSTokenizer* self) {
+ (void) gptoss_tokenizer_release(self->handle);
+ self->handle = NULL;
+ PyObject_Del((PyObject*) self);
+}
+
+static PyObject* PyGPTOSSTokenizer_copy(PyGPTOSSTokenizer* self) {
+ PyGPTOSSTokenizer* copy = (PyGPTOSSTokenizer*) PyObject_New(PyGPTOSSTokenizer, Py_TYPE(self));
+ if (copy == NULL) {
+ return NULL;
+ }
+
+ (void) gptoss_tokenizer_retain(self->handle);
+ copy->handle = self->handle;
+ return (PyObject*) copy;
+}
+
+static PyObject* PyGPTOSSTokenizer_encode_special_token(PyGPTOSSTokenizer* self, PyObject* arg) {
+ if (PyUnicode_Check(arg)) {
+ const char* string_ptr = PyUnicode_AsUTF8(arg);
+ if (string_ptr == NULL) {
+ return NULL;
+ }
+
+ enum gptoss_special_token token_type = gptoss_special_token_invalid;
+ if (strcmp(string_ptr, "<|return|>") == 0) {
+ token_type = gptoss_special_token_return;
+ } else if (strcmp(string_ptr, "<|start|>") == 0) {
+ token_type = gptoss_special_token_start;
+ } else if (strcmp(string_ptr, "<|message|>") == 0) {
+ token_type = gptoss_special_token_message;
+ } else if (strcmp(string_ptr, "<|end|>") == 0) {
+ token_type = gptoss_special_token_end;
+ } else if (strcmp(string_ptr, "<|refusal|>") == 0) {
+ token_type = gptoss_special_token_refusal;
+ } else if (strcmp(string_ptr, "<|constrain|>") == 0) {
+ token_type = gptoss_special_token_constrain;
+ } else if (strcmp(string_ptr, "<|channel|>") == 0) {
+ token_type = gptoss_special_token_channel;
+ } else if (strcmp(string_ptr, "<|call|>") == 0) {
+ token_type = gptoss_special_token_call;
+ } else if (strcmp(string_ptr, "<|untrusted|>") == 0) {
+ token_type = gptoss_special_token_untrusted;
+ } else if (strcmp(string_ptr, "<|end_untrusted|>") == 0) {
+ token_type = gptoss_special_token_end_untrusted;
+ } else {
+ PyErr_Format(PyExc_ValueError, "unrecognized special token: %s", string_ptr);
+ return NULL;
+ }
+
+ uint32_t token_id = UINT32_MAX;
+ const enum gptoss_status status = gptoss_tokenizer_get_special_token_id(
+ self->handle, token_type, &token_id);
+ if (status != gptoss_status_success || token_id == UINT32_MAX) {
+ PyErr_Format(PyExc_ValueError, "tokenizer does not support the %s token", string_ptr);
+ return NULL;
+ }
+
+ return PyLong_FromUnsignedLong((unsigned long) token_id);
+ } else {
+ PyErr_SetString(PyExc_TypeError, "string argument expected");
+ return NULL;
+ }
+}
+
+static PyObject* PyGPTOSSTokenizer_decode(PyGPTOSSTokenizer* self, PyObject* args, PyObject* kwargs) {
+ static char *kwlist[] = {"token", NULL};
+ unsigned int token = 0; // Default to 0 if None
+
+ if (!PyArg_ParseTupleAndKeywords(args, kwargs, "I", kwlist, &token)) {
+ return NULL;
+ }
+
+ const void* token_ptr = NULL;
+ size_t token_size = 0;
+ const enum gptoss_status status = gptoss_tokenizer_decode(self->handle, (uint32_t) token, &token_ptr, &token_size);
+ if (status != gptoss_status_success) {
+ // TODO: set exception
+ return NULL;
+ }
+
+ return PyBytes_FromStringAndSize((const char*) token_ptr, (Py_ssize_t) token_size);
+}
+
+static PyMethodDef PyGPTOSSTokenizer_methods[] = {
+ {"__copy__", (PyCFunction) PyGPTOSSTokenizer_copy, METH_NOARGS, "Create a copy of the Tokenizer"},
+ {"encode_special_token", (PyCFunction) PyGPTOSSTokenizer_encode_special_token, METH_O, "Query ID of a special token"},
+ {"decode", (PyCFunction) PyGPTOSSTokenizer_decode, METH_VARARGS | METH_KEYWORDS, "Convert text token ID to bytes"},
+ {NULL},
+};
+
+static PyObject* PyGPTOSSTokenizer_get_num_text_tokens(PyGPTOSSTokenizer* self, void* closure) {
+ uint32_t num_text_tokens = 0;
+ const enum gptoss_status status = gptoss_tokenizer_get_num_text_tokens(self->handle, &num_text_tokens);
+ if (status != gptoss_status_success) {
+ // TODO: set exception
+ return NULL;
+ }
+
+ return PyLong_FromUnsignedLong((unsigned long) num_text_tokens);
+}
+
+static PyObject* PyGPTOSSTokenizer_get_num_special_tokens(PyGPTOSSTokenizer* self, void* closure) {
+ uint32_t num_special_tokens = 0;
+ const enum gptoss_status status = gptoss_tokenizer_get_num_special_tokens(self->handle, &num_special_tokens);
+ if (status != gptoss_status_success) {
+ // TODO: set exception
+ return NULL;
+ }
+
+ return PyLong_FromUnsignedLong((unsigned long) num_special_tokens);
+}
+
+static PyObject* PyGPTOSSTokenizer_get_num_tokens(PyGPTOSSTokenizer* self, void* closure) {
+ uint32_t num_tokens = 0;
+ const enum gptoss_status status = gptoss_tokenizer_get_num_tokens(self->handle, &num_tokens);
+ if (status != gptoss_status_success) {
+ // TODO: set exception
+ return NULL;
+ }
+
+ return PyLong_FromUnsignedLong((unsigned long) num_tokens);
+}
+
+static PyGetSetDef PyGPTOSSTokenizer_getseters[] = {
+ (PyGetSetDef) {
+ .name = "num_tokens",
+ .get = (getter) PyGPTOSSTokenizer_get_num_tokens,
+ .doc = "Total number of tokens in the tokenizer dictionary",
+ },
+ (PyGetSetDef) {
+ .name = "num_text_tokens",
+ .get = (getter) PyGPTOSSTokenizer_get_num_text_tokens,
+ .doc = "Number of text tokens in the tokenizer dictionary",
+ },
+ (PyGetSetDef) {
+ .name = "num_special_tokens",
+ .get = (getter) PyGPTOSSTokenizer_get_num_special_tokens,
+ .doc = "Number of special tokens in the tokenizer dictionary",
+ },
+ {NULL} /* Sentinel */
+};
+
+PyTypeObject PyGPTOSSTokenizer_Type = {
+ PyVarObject_HEAD_INIT(NULL, 0)
+ .tp_name = "gptoss.Tokenizer",
+ .tp_basicsize = sizeof(PyGPTOSSTokenizer),
+ .tp_flags = 0
+ | Py_TPFLAGS_DEFAULT
+ | Py_TPFLAGS_BASETYPE,
+ .tp_doc = "Tokenizer object",
+ .tp_methods = PyGPTOSSTokenizer_methods,
+ .tp_getset = PyGPTOSSTokenizer_getseters,
+ .tp_new = PyGPTOSSTokenizer_new,
+ .tp_dealloc = (destructor) PyGPTOSSTokenizer_dealloc,
+};
diff --git a/gpt_oss/metal/scripts/create-local-model.py b/gpt_oss/metal/scripts/create-local-model.py
new file mode 100644
index 0000000000000000000000000000000000000000..7f16ba9912fa91079c609301b7c1c0f818178f9d
--- /dev/null
+++ b/gpt_oss/metal/scripts/create-local-model.py
@@ -0,0 +1,358 @@
+import argparse
+import os
+import math
+import sys
+import json
+import itertools
+import struct
+from uuid import UUID
+
+import tiktoken
+
+import torch
+from safetensors import safe_open
+from tqdm import tqdm
+from openai_harmony import load_harmony_encoding, HarmonyEncodingName
+
+parser = argparse.ArgumentParser(prog='create-local-model.py', description='Convert a checkpoint directory to a local model file')
+parser.add_argument('-s', '--src', metavar='DIR', type=str, required=True, help='Path to the input checkpoint directory')
+parser.add_argument('-d', '--dst', metavar='FILE', type=str, required=True, help='Path to the output model file')
+
+
+o200k_base = tiktoken.get_encoding("o200k_base")
+harmony_encoding = load_harmony_encoding(HarmonyEncodingName.HARMONY_GPT_OSS)
+
+o200k_gptoss = tiktoken.Encoding(
+ name="o200k_gptoss",
+ pat_str=o200k_base._pat_str,
+ mergeable_ranks=o200k_base._mergeable_ranks,
+ special_tokens={
+ "<|reversed199998|>": 199998, # unused
+ "<|endoftext|>": 199999,
+ "<|untrusted|>": 200000,
+ "<|endofuntrusted|>": 200001,
+ "<|return|>": 200002,
+ "<|constrain|>": 200003,
+ "<|reversed200004|>": 200004, # unused
+ "<|channel|>": 200005,
+ "<|start|>": 200006,
+ "<|end|>": 200007,
+ "<|message|>": 200008,
+ "<|reversed200008|>": 200008, # unused
+ "<|reversed200009|>": 200009, # unused
+ "<|reversed200010|>": 200010, # unused
+ "<|reversed200011|>": 200011, # unused
+ "<|call|>": 200012,
+ "<|refusal|>": 200013,
+ }
+)
+
+FILE_MAGIC = struct.pack('ccccccccccccI', b'G', b'P', b'T', b'-', b'O', b'S', b'S', b' ', b'v', b'1', b'.', b'0', 0)
+SPECIAL_TOKEN_UUID = {
+ '<|start|>': UUID('55a77c2f-8a01-4c54-8ac2-313bfc7e208d').bytes,
+ '<|message|>': UUID('16e40431-f47f-4b22-b59b-8b278fc30a54').bytes,
+ '<|end|>': UUID('fcac2f6d-4705-4f6b-b228-642accac7238').bytes,
+ '<|return|>': UUID('f799ff69-1992-43c4-a3d8-d831f475dc75').bytes,
+ '<|refusal|>': UUID('e15ba702-28c4-4292-ab8f-ffa434709128').bytes,
+ '<|constrain|>': UUID('c0bb14c7-6022-49da-ad08-792d67e8b470').bytes,
+ '<|channel|>': UUID('fd3dda11-c8ab-4033-876e-d93deb172c93').bytes,
+ '<|call|>': UUID('1220f796-e388-4de5-b487-fe2eb5fe03c0').bytes,
+ '<|untrusted|>': UUID('07d7da55-b346-4cff-8b37-7cefacf8a3e8').bytes,
+ '<|end_untrusted|>': UUID('f265bd9c-c717-469e-a447-920687d65d90').bytes,
+}
+
+INCLUDE_SPECIAL_TOKENS = [
+ "<|start|>",
+ "<|message|>",
+ "<|end|>",
+ "<|return|>",
+ "<|refusal|>",
+ "<|constrain|>",
+ "<|channel|>",
+ "<|call|>",
+ "<|untrusted|>",
+ "<|end_untrusted|>",
+]
+
+GPTOSS_MODEL_UUID = UUID('df52dc86-1789-4ed0-a295-66f10508145b').bytes
+APPLE_GPU_LAYOUT_UUID = UUID('229177a8-5775-4268-bfd8-d588b351c56d').bytes
+TIKTOKEN_TOKENIZER_UUID = UUID('7401aded-2a95-40cb-b782-9ccebaafe72b').bytes
+
+UE8_OFFSET = 14 # bias to MXFP4 block scales
+
+def write_file_header(f):
+ f.write(FILE_MAGIC)
+
+def write_tokenizer_header(f,
+ num_special_tokens: int,
+ num_text_tokens: int,
+ regex_size: int,
+ tokens_size: int):
+ f.write(TIKTOKEN_TOKENIZER_UUID)
+ f.write(struct.pack(' 0
+ tokens_size += len(token_bytes) + 2 # uint16_t string length + string data
+ num_text_tokens += 1
+ # Then add all special tokens
+ num_included_tokens = 200013 + 1
+ print(f"Tokenizer: {num_included_tokens} tokens")
+
+ # Read from all files ending with .safetensors in the checkpoint directory
+ safetensor_files = [
+ os.path.join(options.src, fname)
+ for fname in os.listdir(options.src)
+ if fname.endswith(".safetensors")
+ ]
+ # Build a mapping from tensor name to filepath
+ tensor_name_to_file = {}
+ for safetensor_file in safetensor_files:
+ with safe_open(safetensor_file, framework="pt", device="cpu") as src:
+ for key in src.keys():
+ tensor_name_to_file[key] = safetensor_file
+
+ def get_tensor(name):
+ with safe_open(tensor_name_to_file[name], framework="pt", device="cpu") as src:
+ return src.get_tensor(name)
+
+ with open(options.dst, "wb") as dst:
+ write_file_header(dst)
+
+ yarn_low = (
+ head_dim / 2
+ * math.log(initial_context_length / (rope_ntk_beta * 2 * math.pi))
+ / math.log(rope_theta)
+ )
+ yarn_high = (
+ head_dim / 2
+ * math.log(initial_context_length / (rope_ntk_alpha * 2 * math.pi))
+ / math.log(rope_theta)
+ )
+
+ write_model_header(dst,
+ context_length=int(initial_context_length * rope_scaling_factor),
+ num_blocks=num_blocks,
+ num_experts=num_experts,
+ num_active_experts=num_active_experts,
+ embedding_dim=embedding_dim,
+ mlp_dim=mlp_dim,
+ swiglu_limit=swiglu_limit,
+ head_dim=head_dim,
+ num_heads=num_q_heads,
+ num_kv_heads=num_kv_heads,
+ attention_window=attention_window,
+ rope_theta=rope_theta,
+ interpolation_scale=1.0 / rope_scaling_factor,
+ yarn_offset=-yarn_low / (yarn_high - yarn_low),
+ yarn_scale=1.0 / (yarn_high - yarn_low),
+ yarn_multiplier=0.1 * math.log(rope_scaling_factor) + 1.0,
+ rmsnorm_epsilon=1.0e-5)
+
+ write_tokenizer_header(dst,
+ num_special_tokens=num_included_tokens - num_text_tokens,
+ num_text_tokens=num_text_tokens,
+ regex_size=len(o200k_gptoss._pat_str.encode("ascii")) + 1,
+ tokens_size=tokens_size)
+
+ ### Tokenizer
+ # Special tokens
+ for token_idx in range(num_text_tokens, num_included_tokens):
+ token = o200k_gptoss.decode_single_token_bytes(token_idx).decode('ascii')
+ if token in INCLUDE_SPECIAL_TOKENS:
+ dst.write(SPECIAL_TOKEN_UUID[token])
+ else:
+ dst.write(bytes(16))
+ # Regex
+ dst.write(o200k_gptoss._pat_str.encode("ascii"))
+ dst.write(struct.pack('B', 0))
+ # Text tokens
+ tokenizer_bytes_written = 0
+ for t in range(num_text_tokens):
+ token_bytes = o200k_gptoss.decode_single_token_bytes(t)
+ assert len(token_bytes) > 0
+ dst.write(struct.pack('
+#include
+
+#include
+
+#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((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;
+ }
+}
diff --git a/gpt_oss/metal/source/context.c b/gpt_oss/metal/source/context.c
new file mode 100644
index 0000000000000000000000000000000000000000..7e4b786db78af571f139aac0387748ca770b0a32
--- /dev/null
+++ b/gpt_oss/metal/source/context.c
@@ -0,0 +1,1154 @@
+#include
+#include
+#include
+#include
+#include
+#include
+#include
+
+#include
+
+#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;
+ }
+ status = gptoss_metal_buffer_create(&model->device, 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) {
+ // Commit and wait for the command buffer to complete.
+ // As we need topk output to compute routing metadata.
+ status = gptoss_metal_command_buffer_commit(command_buffer);
+ if (status != gptoss_status_success) {
+ return status;
+ }
+
+ status = gptoss_metal_command_buffer_wait_completion(command_buffer, NULL);
+ if (status != gptoss_status_success) {
+ return status;
+ }
+ gptoss_metal_command_buffer_release(command_buffer);
+
+ const size_t E = model->num_experts;
+ const size_t T = num_block_output_tokens * model->num_active_experts;
+
+ const struct gptoss_expert_prediction* preds =
+ (const struct gptoss_expert_prediction*) context->expert_activation_buffer.ptr;
+
+ uint32_t* token_to_expert_routing = (uint32_t*) context->token_to_expert_routing_buffer.ptr;
+ uint32_t* expert_offset = (uint32_t*) context->expert_offset_buffer.ptr;
+ // Zero out the expert offset buffer.
+ memset(expert_offset, 0, E * sizeof(uint32_t));
+
+ for (size_t i = 0; i < T; i++) {
+ const uint32_t expert_id = preds[i].expert_id;
+ token_to_expert_routing[i] = expert_offset[expert_id];
+ expert_offset[expert_id]++;
+ }
+
+ uint32_t total = 0;
+ // Prefix sum.
+ for (size_t i = 0; i < model->num_experts; i++) {
+ const uint32_t bin_size = expert_offset[i];
+ expert_offset[i] = total;
+ total += bin_size;
+ }
+
+ // Create a new command buffer.
+ status = gptoss_metal_command_buffer_create(&context->model->command_queue, command_buffer);
+ if (status != gptoss_status_success) {
+ 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 -- iterate over all experts.
+ const size_t total_tokens = num_block_output_tokens * model->num_active_experts;
+ for (size_t e = 0; e < model->num_experts; e++) {
+ bool last_expert = e == model->num_experts - 1;
+ uint32_t expert_tokens = last_expert ? total_tokens - expert_offset[e] : expert_offset[e + 1] - expert_offset[e];
+ if (expert_tokens == 0) {
+ continue;
+ }
+ status = gptoss_metal_command_buffer_encode_launch_f32_mf4w_moe_dense_matmul_swiglu(
+ command_buffer,
+ &model->f32_mf4w_moe_dense_matmul_swiglu_fn,
+ &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,
+ expert_tokens,
+ expert_offset[e],
+ e,
+ 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 -- again iterate over all experts.
+ for (size_t e = 0; e < model->num_experts; e++) {
+ bool last_expert = e == model->num_experts - 1;
+ uint32_t expert_tokens = last_expert ? total_tokens - expert_offset[e] : expert_offset[e + 1] - expert_offset[e];
+ if (expert_tokens == 0) {
+ continue;
+ }
+ status = gptoss_metal_command_buffer_encode_launch_f32_mf4w_moe_dense_matmul(
+ command_buffer,
+ &model->f32_mf4w_moe_dense_matmul_fn,
+ &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,
+ expert_tokens,
+ expert_offset[e],
+ e,
+ 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;
+ }
+ }
+
+ 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;
+}
diff --git a/gpt_oss/metal/source/convert.metal b/gpt_oss/metal/source/convert.metal
new file mode 100644
index 0000000000000000000000000000000000000000..1c44c39e0459438a92b27a480138743564c6db01
--- /dev/null
+++ b/gpt_oss/metal/source/convert.metal
@@ -0,0 +1,64 @@
+#include
+
+#include
+
+#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((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((static_cast(*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(as_type(block048CGKOS.xy)) * scale;
+ const float4 blockGKOS = static_cast(as_type(block048CGKOS.zw)) * scale;
+ const float4 block26AE = static_cast(as_type(block26AEIMQU.xy)) * scale;
+ const float4 blockIMQU = static_cast(as_type(block26AEIMQU.zw)) * scale;
+ const float4 block159D = static_cast(as_type(block159DHLPT.xy)) * scale;
+ const float4 blockHLPT = static_cast(as_type(block159DHLPT.zw)) * scale;
+ const float4 block37BF = static_cast(as_type(block37BFJNRV.xy)) * scale;
+ const float4 blockJNRV = static_cast(as_type(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;
+ }
+}
diff --git a/gpt_oss/metal/source/embeddings.metal b/gpt_oss/metal/source/embeddings.metal
new file mode 100644
index 0000000000000000000000000000000000000000..9cc7d1214fe100a897060081ddc0f8855b163b3e
--- /dev/null
+++ b/gpt_oss/metal/source/embeddings.metal
@@ -0,0 +1,29 @@
+#include
+
+#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(w);
+ }
+}
diff --git a/gpt_oss/metal/source/gather_and_accumulate.metal b/gpt_oss/metal/source/gather_and_accumulate.metal
new file mode 100644
index 0000000000000000000000000000000000000000..fd70a94fec8571b871dc765c67897fa333b9f316
--- /dev/null
+++ b/gpt_oss/metal/source/gather_and_accumulate.metal
@@ -0,0 +1,74 @@
+#include
+#include
+#include
+#include
+
+// 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(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(&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(in + r0 * D + col);
+ const device float4* src1 =
+ reinterpret_cast(in + r1 * D + col);
+ const device float4* src2 =
+ reinterpret_cast(in + r2 * D + col);
+ const device float4* src3 =
+ reinterpret_cast(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;
+}
\ No newline at end of file
diff --git a/gpt_oss/metal/source/generate.c b/gpt_oss/metal/source/generate.c
new file mode 100644
index 0000000000000000000000000000000000000000..ae6e680d30cb0c8e3de524a450a2429d348916ee
--- /dev/null
+++ b/gpt_oss/metal/source/generate.c
@@ -0,0 +1,317 @@
+#include
+#include
+#include
+#include
+#include
+#include
+#include
+#include
+#include
+#include
+
+#include
+
+#include
+
+#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 [-p ] [-n ]\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;
+}
diff --git a/gpt_oss/metal/source/include/internal/datatype.h b/gpt_oss/metal/source/include/internal/datatype.h
new file mode 100644
index 0000000000000000000000000000000000000000..27e6117c4907d93c661832ff7aaa864a060435b3
--- /dev/null
+++ b/gpt_oss/metal/source/include/internal/datatype.h
@@ -0,0 +1,41 @@
+#pragma once
+
+#include
+
+#include
+
+
+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");
diff --git a/gpt_oss/metal/source/include/internal/datatype.hpp b/gpt_oss/metal/source/include/internal/datatype.hpp
new file mode 100644
index 0000000000000000000000000000000000000000..4093927056748a7de160546ff89e9067f0264b08
--- /dev/null
+++ b/gpt_oss/metal/source/include/internal/datatype.hpp
@@ -0,0 +1,87 @@
+#pragma once
+
+#include
+
+#include
+
+
+namespace gptoss {
+
+template
+WideT upcast(NarrowT);
+
+template <>
+inline float upcast(gptoss_bfloat16 bf16_value) {
+ const uint32_t bits = static_cast(bf16_value.bits) << 16;
+ return std::bit_cast(bits);
+}
+
+template <>
+inline float upcast(gptoss_float16 fp16_value) {
+ return static_cast(std::bit_cast<_Float16>(fp16_value.bits));
+}
+
+template <>
+inline float upcast(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(bf16_value);
+}
+
+template <>
+inline double upcast(float fp32_value) {
+ return static_cast(fp32_value);
+}
+
+template <>
+inline double upcast(gptoss_bfloat16 bf16_value) {
+ const float fp32_value = upcast(bf16_value);
+ return upcast(fp32_value);
+}
+
+template <>
+inline double upcast(gptoss_float16 fp16_value) {
+ const float fp32_value = upcast(fp16_value);
+ return upcast(fp32_value);
+}
+
+template <>
+inline double upcast(gptoss_float8e4m3 fp8_value) {
+ const float fp32_value = upcast(fp8_value);
+ return upcast(fp32_value);
+}
+
+} // namespace gptoss
diff --git a/gpt_oss/metal/source/include/internal/kernel-args.h b/gpt_oss/metal/source/include/internal/kernel-args.h
new file mode 100644
index 0000000000000000000000000000000000000000..6d398b7a41e22b8911243afc12175b489dcfc2be
--- /dev/null
+++ b/gpt_oss/metal/source/include/internal/kernel-args.h
@@ -0,0 +1,202 @@
+#pragma once
+
+#if !defined(__METAL_VERSION__)
+#include
+#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 expert_token_count;
+ uint32_t k;
+ uint32_t n;
+ uint32_t expert_id;
+ uint32_t expert_token_offset;
+ 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 expert_token_count;
+ uint32_t k;
+ uint32_t n;
+ uint32_t expert_id;
+ uint32_t expert_token_offset;
+ uint32_t weight_blocks_expert_stride_bytes;
+ uint32_t weight_scales_expert_stride_bytes;
+ uint32_t bias_expert_stride_bytes;
+};
+
+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;
+};
diff --git a/gpt_oss/metal/source/include/internal/log.h b/gpt_oss/metal/source/include/internal/log.h
new file mode 100644
index 0000000000000000000000000000000000000000..0e94d501350e8541b8ebb3ecbf1c3bf321f31fe1
--- /dev/null
+++ b/gpt_oss/metal/source/include/internal/log.h
@@ -0,0 +1,20 @@
+#pragma once
+
+#include
+
+
+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__)
diff --git a/gpt_oss/metal/source/include/internal/macros.h b/gpt_oss/metal/source/include/internal/macros.h
new file mode 100644
index 0000000000000000000000000000000000000000..90361df3076a3ee71ee202500347145e2c8b5c6d
--- /dev/null
+++ b/gpt_oss/metal/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
diff --git a/gpt_oss/metal/source/include/internal/math.h b/gpt_oss/metal/source/include/internal/math.h
new file mode 100644
index 0000000000000000000000000000000000000000..06f2b1f1dfd049ed94e2c79d6de424319c654e6b
--- /dev/null
+++ b/gpt_oss/metal/source/include/internal/math.h
@@ -0,0 +1,40 @@
+#pragma once
+
+#include
+#include
+#include
+
+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;
+}
diff --git a/gpt_oss/metal/source/include/internal/metal-kernels.h b/gpt_oss/metal/source/include/internal/metal-kernels.h
new file mode 100644
index 0000000000000000000000000000000000000000..3addb3cc6803a32391237797bb724fd6ee2d9505
--- /dev/null
+++ b/gpt_oss/metal/source/include/internal/metal-kernels.h
@@ -0,0 +1,471 @@
+#pragma once
+
+#include
+#include
+
+#include
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+#include
+#include
+
+#include
+#include
+#include
+
+
+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_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* 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 expert_token_offset,
+ uint32_t expert_id,
+ 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* 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 expert_token_offset,
+ uint32_t expert_id,
+ 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
diff --git a/gpt_oss/metal/source/include/internal/metal.h b/gpt_oss/metal/source/include/internal/metal.h
new file mode 100644
index 0000000000000000000000000000000000000000..f38190f06e693de0b4817c2365dec5c6b9d7569c
--- /dev/null
+++ b/gpt_oss/metal/source/include/internal/metal.h
@@ -0,0 +1,138 @@
+#pragma once
+
+#include
+
+#include
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+struct gptoss_metal_device {
+ void* object; // id
+ 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
+};
+
+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
+ void* pipeline_state_object; // id
+ 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
+ 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
+};
+
+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
+};
+
+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
diff --git a/gpt_oss/metal/source/include/internal/metal.hpp b/gpt_oss/metal/source/include/internal/metal.hpp
new file mode 100644
index 0000000000000000000000000000000000000000..a143a11ab319cd5ce864b88affebd60d2eb29c6d
--- /dev/null
+++ b/gpt_oss/metal/source/include/internal/metal.hpp
@@ -0,0 +1,342 @@
+#pragma once
+
+#include
+#include
+#include
+#include
+#include
+
+#include
+#include
+#include
+
+
+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& threadgroup_size,
+ const std::array& num_threadgroups,
+ size_t params_size, const void* params,
+ std::initializer_list device_buffers = {},
+ size_t threadgroup_buffer_size = 0)
+ {
+ std::vector 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
diff --git a/gpt_oss/metal/source/include/internal/model.h b/gpt_oss/metal/source/include/internal/model.h
new file mode 100644
index 0000000000000000000000000000000000000000..d27743b2b4643517c0673e89247f20326782fd28
--- /dev/null
+++ b/gpt_oss/metal/source/include/internal/model.h
@@ -0,0 +1,177 @@
+#pragma once
+
+#ifndef __cplusplus
+ #include
+#endif
+#include
+#include
+#include
+
+#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_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;
+};
diff --git a/gpt_oss/metal/source/include/internal/rng.h b/gpt_oss/metal/source/include/internal/rng.h
new file mode 100644
index 0000000000000000000000000000000000000000..048b8b7d49764f86bb7ee6b4d36ef99eba6f2884
--- /dev/null
+++ b/gpt_oss/metal/source/include/internal/rng.h
@@ -0,0 +1,24 @@
+#pragma once
+
+#include
+
+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);
+}
diff --git a/gpt_oss/metal/source/include/internal/rng.hpp b/gpt_oss/metal/source/include/internal/rng.hpp
new file mode 100644
index 0000000000000000000000000000000000000000..e327a3398004c4bb9e37ff11a6a545712787a7c4
--- /dev/null
+++ b/gpt_oss/metal/source/include/internal/rng.hpp
@@ -0,0 +1,32 @@
+#pragma once
+
+#include
+
+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(x >> 32);
+}
+
+} // namespace rng
+
+} // namespace gptoss
diff --git a/gpt_oss/metal/source/include/internal/storage.h b/gpt_oss/metal/source/include/internal/storage.h
new file mode 100644
index 0000000000000000000000000000000000000000..cc7ed7b4fee10cb407b78224db9d4b1ea8445625
--- /dev/null
+++ b/gpt_oss/metal/source/include/internal/storage.h
@@ -0,0 +1,36 @@
+#pragma once
+
+#include
+#include
+
+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;
+};
diff --git a/gpt_oss/metal/source/include/internal/uuid.h b/gpt_oss/metal/source/include/internal/uuid.h
new file mode 100644
index 0000000000000000000000000000000000000000..0f52f8a425b32d2fc3d9b1edf1354e18213eb1b2
--- /dev/null
+++ b/gpt_oss/metal/source/include/internal/uuid.h
@@ -0,0 +1,114 @@
+#pragma once
+
+#include
+#include
+#include
+
+#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;
+ }
+}
diff --git a/gpt_oss/metal/source/log.c b/gpt_oss/metal/source/log.c
new file mode 100644
index 0000000000000000000000000000000000000000..ff7b2375eb97b4a0c1ca8bd28fdebf121c69c725
--- /dev/null
+++ b/gpt_oss/metal/source/log.c
@@ -0,0 +1,50 @@
+#include // assert
+#include // va_list, va_copy, va_end
+#include // vsnprintf
+#include // malloc, free
+
+#include 
+