new builds

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

@@ -1,8 +1,174 @@
 from ._ops import ops
 import torch
 
-def f32_bf16w_matmul(input: torch.Tensor, weight_bf16: torch.Tensor, bias_bf16: torch.Tensor, output: torch.Tensor, num_tokens: int, num_cols: int, num_rows: int, threadgroup_size: int) -> None:
-    ops.f32_bf16w_matmul(input, weight_bf16, bias_bf16, output, num_tokens, num_cols, num_rows, threadgroup_size)
+def f32_bf16w_matmul(input: torch.Tensor,
+                     weight_bf16: torch.Tensor,
+                     bias_bf16: torch.Tensor,
+                     output: torch.Tensor,
+                     num_tokens: int,
+                     num_cols: int,
+                     num_rows: int,
+                     threadgroup_size: int) -> torch.Tensor:
+    ops.f32_bf16w_matmul(input, weight_bf16, bias_bf16, output,
+                         num_tokens, num_cols, num_rows, threadgroup_size)
     return output
 
-__all__ = ["f32_bf16w_matmul"]
+def bf16_f32_embeddings(token_ids: torch.Tensor,
+                        weight_bf16: torch.Tensor,
+                        output: torch.Tensor,
+                        threadgroup_size: int) -> torch.Tensor:
+    ops.bf16_f32_embeddings(token_ids, weight_bf16, output, threadgroup_size)
+    return output
+
+def f32_bf16w_rmsnorm(input: torch.Tensor,
+                      weight_bf16: torch.Tensor,
+                      output: torch.Tensor,
+                      epsilon: float) -> torch.Tensor:
+    ops.f32_bf16w_rmsnorm(input, weight_bf16, output, epsilon)
+    return output
+
+def f32_bf16w_dense_matmul_qkv(input: torch.Tensor,
+                               weight_bf16: torch.Tensor,
+                               bias_bf16: torch.Tensor,
+                               output: torch.Tensor) -> torch.Tensor:
+    ops.f32_bf16w_dense_matmul_qkv(input, weight_bf16, bias_bf16, output)
+    return output
+
+def f32_bf16w_dense_matmul_attn_output(input: torch.Tensor,
+                                       weight_bf16: torch.Tensor,
+                                       bias_bf16: torch.Tensor,
+                                       output: torch.Tensor) -> torch.Tensor:
+    ops.f32_bf16w_dense_matmul_attn_output(input, weight_bf16, bias_bf16, output)
+    return output
+
+def f32_bf16w_dense_matmul_mlp_gate(input: torch.Tensor,
+                                    weight_bf16: torch.Tensor,
+                                    bias_bf16: torch.Tensor,
+                                    output: torch.Tensor) -> torch.Tensor:
+    ops.f32_bf16w_dense_matmul_mlp_gate(input, weight_bf16, bias_bf16, output)
+    return output
+
+def f32_rope(activations: torch.Tensor,
+             rope_base: float,
+             interpolation_scale: float,
+             yarn_offset: float,
+             yarn_scale: float,
+             yarn_multiplier: float,
+             num_tokens: int,
+             num_q_heads: int,
+             num_kv_heads: int,
+             attn_head_dim: int,
+             token_offset: int,
+             threadgroup_size: int) -> torch.Tensor:
+    ops.f32_rope(activations, rope_base, interpolation_scale, yarn_offset,
+                 yarn_scale, yarn_multiplier, num_tokens, num_q_heads,
+                 num_kv_heads, attn_head_dim, token_offset, threadgroup_size)
+    return activations
+
+def f32_bf16w_matmul_qkv(input: torch.Tensor,
+                         weight_bf16: torch.Tensor,
+                         bias_bf16: torch.Tensor,
+                         output: torch.Tensor,
+                         kv_cache: torch.Tensor,
+                         kv_cache_offset_bytes: int,
+                         num_tokens: int,
+                         num_cols: int,
+                         num_q_heads: int,
+                         num_kv_heads: int,
+                         attn_head_dim: int,
+                         token_offset: int,
+                         max_tokens: int,
+                         rope_base: float,
+                         interpolation_scale: float,
+                         yarn_offset: float,
+                         yarn_scale: float,
+                         yarn_multiplier: float,
+                         threadgroup_size: int) -> torch.Tensor:
+    ops.f32_bf16w_matmul_qkv(input, weight_bf16, bias_bf16, output, kv_cache,
+                             kv_cache_offset_bytes, num_tokens, num_cols,
+                             num_q_heads, num_kv_heads, attn_head_dim,
+                             token_offset, max_tokens, rope_base,
+                             interpolation_scale, yarn_offset, yarn_scale,
+                             yarn_multiplier, threadgroup_size)
+    return output
+
+def f32_sdpa(q: torch.Tensor,
+             q_offset_bytes: int,
+             kv: torch.Tensor,
+             kv_offset_bytes: int,
+             s_bf16: torch.Tensor,
+             s_offset_bytes: int,
+             output: torch.Tensor,
+             output_offset_bytes: int,
+             window: int,
+             kv_stride: int,
+             num_q_tokens: int,
+             num_kv_tokens: int,
+             num_q_heads: int,
+             num_kv_heads: int,
+             head_dim: int) -> torch.Tensor:
+    ops.f32_sdpa(q, q_offset_bytes, kv, kv_offset_bytes, s_bf16, s_offset_bytes,
+                 output, output_offset_bytes, window, kv_stride,
+                 num_q_tokens, num_kv_tokens, num_q_heads, num_kv_heads, head_dim)
+    return output
+
+def f32_topk(scores: torch.Tensor,
+             expert_ids: torch.Tensor,
+             expert_scores: torch.Tensor,
+             num_tokens: int,
+             num_experts: int,
+             num_active_experts: int) -> None:
+    ops.f32_topk(scores, expert_ids, expert_scores,
+                 num_tokens, num_experts, num_active_experts)
+
+def expert_routing_metadata(expert_ids: torch.Tensor,
+                            expert_scores: torch.Tensor,
+                            expert_offsets: torch.Tensor,
+                            intra_expert_offsets: torch.Tensor,
+                            num_tokens: int,
+                            num_experts: int) -> None:
+    ops.expert_routing_metadata(expert_ids, expert_scores,
+                                expert_offsets, intra_expert_offsets,
+                                num_tokens, num_experts)
+
+def f32_scatter(input: torch.Tensor,
+                expert_ids: torch.Tensor,
+                expert_scores: torch.Tensor,
+                expert_offsets: torch.Tensor,
+                intra_expert_offsets: torch.Tensor,
+                output: torch.Tensor,
+                num_channels: int,
+                num_tokens: int,
+                num_active_experts: int) -> torch.Tensor:
+    ops.f32_scatter(input, expert_ids, expert_scores,
+                    expert_offsets, intra_expert_offsets,
+                    output, num_channels, num_tokens, num_active_experts)
+    return output
+
+def f32_bf16w_matmul_add(input: torch.Tensor,
+                         weight_bf16: torch.Tensor,
+                         bias_bf16: torch.Tensor,
+                         output: torch.Tensor,
+                         num_tokens: int,
+                         num_cols: int,
+                         num_rows: int,
+                         threadgroup_size: int) -> torch.Tensor:
+    ops.f32_bf16w_matmul_add(input, weight_bf16, bias_bf16, output,
+                             num_tokens, num_cols, num_rows, threadgroup_size)
+    return output
+
+__all__ = [
+    "f32_bf16w_matmul",
+    "bf16_f32_embeddings",
+    "f32_bf16w_rmsnorm",
+    "f32_bf16w_dense_matmul_qkv",
+    "f32_bf16w_dense_matmul_attn_output",
+    "f32_bf16w_dense_matmul_mlp_gate",
+    "f32_rope",
+    "f32_bf16w_matmul_qkv",
+    "f32_sdpa",
+    "f32_topk",
+    "expert_routing_metadata",
+    "f32_scatter",
+    "f32_bf16w_matmul_add",
+]

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

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

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

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

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

@@ -1,8 +1,174 @@
 from ._ops import ops
 import torch
 
-def f32_bf16w_matmul(input: torch.Tensor, weight_bf16: torch.Tensor, bias_bf16: torch.Tensor, output: torch.Tensor, num_tokens: int, num_cols: int, num_rows: int, threadgroup_size: int) -> None:
-    ops.f32_bf16w_matmul(input, weight_bf16, bias_bf16, output, num_tokens, num_cols, num_rows, threadgroup_size)
+def f32_bf16w_matmul(input: torch.Tensor,
+                     weight_bf16: torch.Tensor,
+                     bias_bf16: torch.Tensor,
+                     output: torch.Tensor,
+                     num_tokens: int,
+                     num_cols: int,
+                     num_rows: int,
+                     threadgroup_size: int) -> torch.Tensor:
+    ops.f32_bf16w_matmul(input, weight_bf16, bias_bf16, output,
+                         num_tokens, num_cols, num_rows, threadgroup_size)
     return output
 
-__all__ = ["f32_bf16w_matmul"]
+def bf16_f32_embeddings(token_ids: torch.Tensor,
+                        weight_bf16: torch.Tensor,
+                        output: torch.Tensor,
+                        threadgroup_size: int) -> torch.Tensor:
+    ops.bf16_f32_embeddings(token_ids, weight_bf16, output, threadgroup_size)
+    return output
+
+def f32_bf16w_rmsnorm(input: torch.Tensor,
+                      weight_bf16: torch.Tensor,
+                      output: torch.Tensor,
+                      epsilon: float) -> torch.Tensor:
+    ops.f32_bf16w_rmsnorm(input, weight_bf16, output, epsilon)
+    return output
+
+def f32_bf16w_dense_matmul_qkv(input: torch.Tensor,
+                               weight_bf16: torch.Tensor,
+                               bias_bf16: torch.Tensor,
+                               output: torch.Tensor) -> torch.Tensor:
+    ops.f32_bf16w_dense_matmul_qkv(input, weight_bf16, bias_bf16, output)
+    return output
+
+def f32_bf16w_dense_matmul_attn_output(input: torch.Tensor,
+                                       weight_bf16: torch.Tensor,
+                                       bias_bf16: torch.Tensor,
+                                       output: torch.Tensor) -> torch.Tensor:
+    ops.f32_bf16w_dense_matmul_attn_output(input, weight_bf16, bias_bf16, output)
+    return output
+
+def f32_bf16w_dense_matmul_mlp_gate(input: torch.Tensor,
+                                    weight_bf16: torch.Tensor,
+                                    bias_bf16: torch.Tensor,
+                                    output: torch.Tensor) -> torch.Tensor:
+    ops.f32_bf16w_dense_matmul_mlp_gate(input, weight_bf16, bias_bf16, output)
+    return output
+
+def f32_rope(activations: torch.Tensor,
+             rope_base: float,
+             interpolation_scale: float,
+             yarn_offset: float,
+             yarn_scale: float,
+             yarn_multiplier: float,
+             num_tokens: int,
+             num_q_heads: int,
+             num_kv_heads: int,
+             attn_head_dim: int,
+             token_offset: int,
+             threadgroup_size: int) -> torch.Tensor:
+    ops.f32_rope(activations, rope_base, interpolation_scale, yarn_offset,
+                 yarn_scale, yarn_multiplier, num_tokens, num_q_heads,
+                 num_kv_heads, attn_head_dim, token_offset, threadgroup_size)
+    return activations
+
+def f32_bf16w_matmul_qkv(input: torch.Tensor,
+                         weight_bf16: torch.Tensor,
+                         bias_bf16: torch.Tensor,
+                         output: torch.Tensor,
+                         kv_cache: torch.Tensor,
+                         kv_cache_offset_bytes: int,
+                         num_tokens: int,
+                         num_cols: int,
+                         num_q_heads: int,
+                         num_kv_heads: int,
+                         attn_head_dim: int,
+                         token_offset: int,
+                         max_tokens: int,
+                         rope_base: float,
+                         interpolation_scale: float,
+                         yarn_offset: float,
+                         yarn_scale: float,
+                         yarn_multiplier: float,
+                         threadgroup_size: int) -> torch.Tensor:
+    ops.f32_bf16w_matmul_qkv(input, weight_bf16, bias_bf16, output, kv_cache,
+                             kv_cache_offset_bytes, num_tokens, num_cols,
+                             num_q_heads, num_kv_heads, attn_head_dim,
+                             token_offset, max_tokens, rope_base,
+                             interpolation_scale, yarn_offset, yarn_scale,
+                             yarn_multiplier, threadgroup_size)
+    return output
+
+def f32_sdpa(q: torch.Tensor,
+             q_offset_bytes: int,
+             kv: torch.Tensor,
+             kv_offset_bytes: int,
+             s_bf16: torch.Tensor,
+             s_offset_bytes: int,
+             output: torch.Tensor,
+             output_offset_bytes: int,
+             window: int,
+             kv_stride: int,
+             num_q_tokens: int,
+             num_kv_tokens: int,
+             num_q_heads: int,
+             num_kv_heads: int,
+             head_dim: int) -> torch.Tensor:
+    ops.f32_sdpa(q, q_offset_bytes, kv, kv_offset_bytes, s_bf16, s_offset_bytes,
+                 output, output_offset_bytes, window, kv_stride,
+                 num_q_tokens, num_kv_tokens, num_q_heads, num_kv_heads, head_dim)
+    return output
+
+def f32_topk(scores: torch.Tensor,
+             expert_ids: torch.Tensor,
+             expert_scores: torch.Tensor,
+             num_tokens: int,
+             num_experts: int,
+             num_active_experts: int) -> None:
+    ops.f32_topk(scores, expert_ids, expert_scores,
+                 num_tokens, num_experts, num_active_experts)
+
+def expert_routing_metadata(expert_ids: torch.Tensor,
+                            expert_scores: torch.Tensor,
+                            expert_offsets: torch.Tensor,
+                            intra_expert_offsets: torch.Tensor,
+                            num_tokens: int,
+                            num_experts: int) -> None:
+    ops.expert_routing_metadata(expert_ids, expert_scores,
+                                expert_offsets, intra_expert_offsets,
+                                num_tokens, num_experts)
+
+def f32_scatter(input: torch.Tensor,
+                expert_ids: torch.Tensor,
+                expert_scores: torch.Tensor,
+                expert_offsets: torch.Tensor,
+                intra_expert_offsets: torch.Tensor,
+                output: torch.Tensor,
+                num_channels: int,
+                num_tokens: int,
+                num_active_experts: int) -> torch.Tensor:
+    ops.f32_scatter(input, expert_ids, expert_scores,
+                    expert_offsets, intra_expert_offsets,
+                    output, num_channels, num_tokens, num_active_experts)
+    return output
+
+def f32_bf16w_matmul_add(input: torch.Tensor,
+                         weight_bf16: torch.Tensor,
+                         bias_bf16: torch.Tensor,
+                         output: torch.Tensor,
+                         num_tokens: int,
+                         num_cols: int,
+                         num_rows: int,
+                         threadgroup_size: int) -> torch.Tensor:
+    ops.f32_bf16w_matmul_add(input, weight_bf16, bias_bf16, output,
+                             num_tokens, num_cols, num_rows, threadgroup_size)
+    return output
+
+__all__ = [
+    "f32_bf16w_matmul",
+    "bf16_f32_embeddings",
+    "f32_bf16w_rmsnorm",
+    "f32_bf16w_dense_matmul_qkv",
+    "f32_bf16w_dense_matmul_attn_output",
+    "f32_bf16w_dense_matmul_mlp_gate",
+    "f32_rope",
+    "f32_bf16w_matmul_qkv",
+    "f32_sdpa",
+    "f32_topk",
+    "expert_routing_metadata",
+    "f32_scatter",
+    "f32_bf16w_matmul_add",
+]

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

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

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

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

gptoss_kernels/source/tensor_wrappers.cpp

@@ -1,6 +1,227 @@
+#include <ATen/Functions.h>
+#include <ATen/Tensor.h>
+
 #include <internal/metal-kernels.h>
 #include <internal/metal.h>
-#include <ATen/Tensor.h>
+
+#include <algorithm>
+#include <cstddef>
+#include <cstdint>
+#include <cstring>
+#include <utility>
+#include <vector>
+
+namespace {
+
+class MetalBuffer {
+public:
+    MetalBuffer() = default;
+    MetalBuffer(const MetalBuffer&) = delete;
+    MetalBuffer& operator=(const MetalBuffer&) = delete;
+
+    MetalBuffer(MetalBuffer&& other) noexcept
+        : buffer_(other.buffer_), has_value_(other.has_value_) {
+        other.buffer_ = {};
+        other.has_value_ = false;
+    }
+
+    MetalBuffer& operator=(MetalBuffer&& other) noexcept {
+        if (this != &other) {
+            reset();
+            buffer_ = other.buffer_;
+            has_value_ = other.has_value_;
+            other.buffer_ = {};
+            other.has_value_ = false;
+        }
+        return *this;
+    }
+
+    ~MetalBuffer() {
+        reset();
+    }
+
+    gptoss_metal_buffer* get() {
+        return &buffer_;
+    }
+
+    const gptoss_metal_buffer* get() const {
+        return &buffer_;
+    }
+
+    void* ptr() const {
+        return buffer_.ptr;
+    }
+
+    size_t size_bytes() const {
+        return buffer_.size;
+    }
+
+    bool valid() const {
+        return has_value_;
+    }
+
+    void wrap(const gptoss_metal_device* device, size_t size, const void* data) {
+        reset();
+        TORCH_CHECK(gptoss_metal_buffer_wrap(device, size, data, &buffer_) == gptoss_status_success,
+            "metal_buffer_wrap failed");
+        has_value_ = true;
+    }
+
+    void create(const gptoss_metal_device* device, size_t size, const void* data = nullptr) {
+        reset();
+        TORCH_CHECK(gptoss_metal_buffer_create(device, size, data, &buffer_) == gptoss_status_success,
+            "metal_buffer_create failed");
+        has_value_ = true;
+    }
+
+    void reset() {
+        if (has_value_) {
+            (void) gptoss_metal_buffer_release(&buffer_);
+            buffer_ = {};
+            has_value_ = false;
+        }
+    }
+
+private:
+    gptoss_metal_buffer buffer_{};
+    bool has_value_ = false;
+};
+
+template <typename EncodeFn>
+void run_metal_kernel(const char* kernel_symbol, EncodeFn&& encode_fn) {
+    gptoss_metal_device device{};
+    gptoss_metal_library library{};
+    gptoss_metal_function fn{};
+    gptoss_metal_command_queue cq{};
+    gptoss_metal_command_buffer cb{};
+
+    auto cleanup = [&]() {
+        (void) gptoss_metal_command_buffer_release(&cb);
+        (void) gptoss_metal_command_queue_release(&cq);
+        (void) gptoss_metal_function_release(&fn);
+        (void) gptoss_metal_library_release(&library);
+        (void) gptoss_metal_device_release(&device);
+    };
+
+    TORCH_CHECK(gptoss_metal_device_create_system_default(&device) == gptoss_status_success,
+        "device_create failed");
+    try {
+        TORCH_CHECK(gptoss_metal_library_create_default(&device, &library) == gptoss_status_success,
+            "library_create failed");
+        TORCH_CHECK(gptoss_metal_function_create(&library, kernel_symbol, &fn) == gptoss_status_success,
+            "function_create failed");
+        TORCH_CHECK(gptoss_metal_command_queue_create(&device, &cq) == gptoss_status_success,
+            "cq_create failed");
+        TORCH_CHECK(gptoss_metal_command_buffer_create(&cq, &cb) == gptoss_status_success,
+            "cb_create failed");
+
+        encode_fn(device, fn, cb);
+
+        TORCH_CHECK(gptoss_metal_command_buffer_commit(&cb) == gptoss_status_success,
+            "commit failed");
+        TORCH_CHECK(gptoss_metal_command_buffer_wait_completion(&cb, nullptr) == gptoss_status_success,
+            "wait failed");
+    } catch (...) {
+        cleanup();
+        throw;
+    }
+
+    cleanup();
+}
+
+at::Tensor to_cpu_contiguous(const at::Tensor& tensor) {
+    if (tensor.device().is_cpu() && tensor.is_contiguous()) {
+        return tensor;
+    }
+    return tensor.contiguous().to(at::kCPU);
+}
+
+at::Tensor empty_cpu_like(const at::Tensor& tensor) {
+    return at::empty_like(tensor, tensor.options().device(at::kCPU)).contiguous();
+}
+
+void copy_back(at::Tensor& dst, const at::Tensor& src_cpu) {
+    dst.copy_(src_cpu.to(dst.device(), /*non_blocking=*/false, /*copy=*/true));
+}
+
+void create_control_buffer(const gptoss_metal_device* device, MetalBuffer& buffer) {
+    struct gptoss_control ctrl {0};
+    buffer.create(device, sizeof(ctrl), &ctrl);
+}
+
+template <typename LaunchFn>
+void run_dense_matmul_bf16(const char* kernel_symbol,
+    LaunchFn&& launch_fn,
+    const at::Tensor& input,
+    const at::Tensor& weight_bf16,
+    const at::Tensor& bias_bf16,
+    at::Tensor& output)
+{
+    TORCH_CHECK(input.dtype() == at::kFloat, "input must be float32");
+    TORCH_CHECK(weight_bf16.dtype() == at::kBFloat16, "weight must be bfloat16");
+    TORCH_CHECK(bias_bf16.dtype() == at::kBFloat16, "bias must be bfloat16");
+    TORCH_CHECK(output.dtype() == at::kFloat, "output must be float32");
+
+    TORCH_CHECK(input.dim() == 2, "input must be 2D");
+    TORCH_CHECK(weight_bf16.dim() == 2, "weight must be 2D");
+    TORCH_CHECK(bias_bf16.dim() == 1, "bias must be 1D");
+    TORCH_CHECK(output.dim() == 2, "output must be 2D");
+
+    const int64_t num_tokens = input.size(0);
+    const int64_t num_cols = input.size(1);
+    const int64_t num_rows = output.size(1);
+
+    TORCH_CHECK(output.size(0) == num_tokens,
+        "output first dimension must match number of tokens");
+    TORCH_CHECK(weight_bf16.size(0) == num_cols && weight_bf16.size(1) == num_rows,
+        "weight shape must be [num_cols, num_rows]");
+    TORCH_CHECK(bias_bf16.size(0) == num_rows,
+        "bias length must equal number of rows");
+
+    auto input_cpu = to_cpu_contiguous(input);
+    auto weight_cpu = weight_bf16.transpose(0, 1).contiguous().to(at::kCPU);
+    auto bias_cpu = to_cpu_contiguous(bias_bf16);
+    auto out_cpu = empty_cpu_like(output);
+
+    const size_t in_bytes = static_cast<size_t>(input_cpu.numel()) * input_cpu.element_size();
+    const size_t weight_bytes = static_cast<size_t>(weight_cpu.numel()) * weight_cpu.element_size();
+    const size_t bias_bytes = static_cast<size_t>(bias_cpu.numel()) * bias_cpu.element_size();
+    const size_t out_bytes = static_cast<size_t>(out_cpu.numel()) * out_cpu.element_size();
+
+    MetalBuffer input_buf;
+    MetalBuffer weight_buf;
+    MetalBuffer bias_buf;
+    MetalBuffer out_buf;
+    MetalBuffer control_buf;
+
+    run_metal_kernel(kernel_symbol, [&](const gptoss_metal_device& device,
+                                        const gptoss_metal_function& fn,
+                                        gptoss_metal_command_buffer& cb) {
+        input_buf.wrap(&device, in_bytes, input_cpu.data_ptr());
+        weight_buf.wrap(&device, weight_bytes, weight_cpu.data_ptr());
+        bias_buf.wrap(&device, bias_bytes, bias_cpu.data_ptr());
+        out_buf.create(&device, out_bytes, nullptr);
+        create_control_buffer(&device, control_buf);
+
+        TORCH_CHECK(
+            launch_fn(
+                &cb, &fn,
+                input_buf.get(), 0,
+                weight_buf.get(), 0,
+                bias_buf.get(), 0,
+                out_buf.get(), 0,
+                control_buf.get(), 0,
+                static_cast<uint32_t>(num_tokens),
+                static_cast<uint32_t>(num_cols),
+                static_cast<uint32_t>(num_rows)) == gptoss_status_success,
+            "encode dense matmul failed");
+    });
+
+    std::memcpy(out_cpu.data_ptr(), out_buf.ptr(), out_bytes);
+    copy_back(output, out_cpu);
+}
+
+}  // namespace
 
 void f32_bf16w_matmul_torch(const at::Tensor &input,
     const at::Tensor &weight_bf16,
@@ -25,53 +246,719 @@ void f32_bf16w_matmul_torch(const at::Tensor &input,
     TORCH_CHECK(output.size(0) == num_tokens && output.size(1) == num_rows,
                 "output shape must be [num_tokens, num_rows]");
 
-    auto input_cpu  = input.contiguous().to(at::kCPU);
+    auto input_cpu  = to_cpu_contiguous(input);
     auto weight_cpu = weight_bf16.transpose(0, 1).contiguous().to(at::kCPU);
-    auto bias_cpu   = bias_bf16.contiguous().to(at::kCPU);
-    auto out_cpu = output.detach().to(at::kCPU).contiguous().clone();
+    auto bias_cpu   = to_cpu_contiguous(bias_bf16);
+    auto out_cpu    = empty_cpu_like(output);
 
-    gptoss_metal_device device{}; gptoss_metal_library library{};
-    gptoss_metal_function fn{};   gptoss_metal_command_queue cq{};
-    gptoss_metal_command_buffer cb{};
+    const size_t in_bytes  = static_cast<size_t>(num_tokens) * static_cast<size_t>(num_cols) * sizeof(float);
+    const size_t wt_bytes  = static_cast<size_t>(num_rows)   * static_cast<size_t>(num_cols) * sizeof(uint16_t);
+    const size_t bs_bytes  = static_cast<size_t>(num_rows) * sizeof(uint16_t);
+    const size_t out_bytes = static_cast<size_t>(num_tokens) * static_cast<size_t>(num_rows) * sizeof(float);
+
+    MetalBuffer in_buf;
+    MetalBuffer wt_buf;
+    MetalBuffer bs_buf;
+    MetalBuffer out_buf;
+    MetalBuffer ctrl_buf;
+
+    run_metal_kernel("gptoss_f32_bf16w_matmul", [&](const gptoss_metal_device& device,
+                                                    const gptoss_metal_function& fn,
+                                                    gptoss_metal_command_buffer& cb) {
+        in_buf.wrap(&device, in_bytes, input_cpu.data_ptr());
+        wt_buf.wrap(&device, wt_bytes, weight_cpu.data_ptr());
+        bs_buf.wrap(&device, bs_bytes, bias_cpu.data_ptr());
+        out_buf.create(&device, out_bytes, nullptr);
+        create_control_buffer(&device, ctrl_buf);
+
+        TORCH_CHECK(
+            gptoss_metal_command_buffer_encode_launch_f32_bf16w_matmul(
+                &cb, &fn, static_cast<size_t>(threadgroup_size),
+                in_buf.get(), 0,
+                wt_buf.get(), 0,
+                bs_buf.get(), 0,
+                out_buf.get(), 0,
+                ctrl_buf.get(), 0,
+                static_cast<uint32_t>(num_tokens),
+                static_cast<uint32_t>(num_cols),
+                static_cast<uint32_t>(num_rows)) == gptoss_status_success,
+            "encode failed");
+    });
+
+    std::memcpy(out_cpu.data_ptr(), out_buf.ptr(), out_bytes);
+    copy_back(output, out_cpu);
+}
+
+void bf16_f32_embeddings_torch(const at::Tensor& token_ids,
+    const at::Tensor& weight_bf16,
+    at::Tensor& output,
+    int64_t threadgroup_size)
+{
+    TORCH_CHECK(token_ids.dtype() == at::kInt || token_ids.dtype() == at::kLong,
+        "token_ids must be int32 or int64");
+    TORCH_CHECK(weight_bf16.dtype() == at::kBFloat16, "weight must be bfloat16");
+    TORCH_CHECK(output.dtype() == at::kFloat, "output must be float32");
+
+    TORCH_CHECK(token_ids.dim() == 1, "token_ids must be 1D");
+    TORCH_CHECK(weight_bf16.dim() == 2, "weight must be 2D");
+    TORCH_CHECK(output.dim() == 2, "output must be 2D");
+
+    const int64_t num_tokens = token_ids.size(0);
+    TORCH_CHECK(output.size(0) == num_tokens, "output first dimension must match num_tokens");
+    const int64_t num_channels = output.size(1);
+    TORCH_CHECK(num_channels % 4 == 0, "num_channels must be divisible by 4");
+    TORCH_CHECK(weight_bf16.size(1) == num_channels,
+        "weight second dimension must equal embedding dimension (num_channels)");
+
+    TORCH_CHECK(threadgroup_size >= 0, "threadgroup_size must be non-negative");
+
+    auto tokens_cpu = token_ids.dtype() == at::kInt
+        ? to_cpu_contiguous(token_ids)
+        : token_ids.to(at::kInt).contiguous().to(at::kCPU);
+    auto weight_cpu = to_cpu_contiguous(weight_bf16);
+    auto out_cpu = empty_cpu_like(output);
+
+    const size_t token_bytes = static_cast<size_t>(num_tokens) * sizeof(uint32_t);
+    const size_t weight_bytes = static_cast<size_t>(weight_cpu.numel()) * weight_cpu.element_size();
+    const size_t out_bytes = static_cast<size_t>(out_cpu.numel()) * out_cpu.element_size();
+
+    MetalBuffer tokens_buf;
+    MetalBuffer weight_buf;
+    MetalBuffer out_buf;
+    MetalBuffer control_buf;
+
+    run_metal_kernel("gptoss_bf16_f32_embeddings", [&](const gptoss_metal_device& device,
+                                                         const gptoss_metal_function& fn,
+                                                         gptoss_metal_command_buffer& cb) {
+        tokens_buf.wrap(&device, token_bytes, tokens_cpu.data_ptr());
+        weight_buf.wrap(&device, weight_bytes, weight_cpu.data_ptr());
+        out_buf.create(&device, out_bytes, nullptr);
+        create_control_buffer(&device, control_buf);
+
+        TORCH_CHECK(
+            gptoss_metal_command_buffer_encode_launch_bf16_f32_embeddings(
+                &cb, &fn, static_cast<size_t>(threadgroup_size),
+                tokens_buf.get(), 0,
+                weight_buf.get(), 0,
+                out_buf.get(), 0,
+                control_buf.get(), 0,
+                static_cast<uint32_t>(num_tokens),
+                static_cast<uint32_t>(num_channels)) == gptoss_status_success,
+            "encode embeddings failed");
+    });
+
+    std::memcpy(out_cpu.data_ptr(), out_buf.ptr(), out_bytes);
+    copy_back(output, out_cpu);
+}
+
+void f32_bf16w_rmsnorm_torch(const at::Tensor& input,
+    const at::Tensor& weight_bf16,
+    at::Tensor& output,
+    double epsilon)
+{
+    TORCH_CHECK(input.dtype() == at::kFloat, "input must be float32");
+    TORCH_CHECK(weight_bf16.dtype() == at::kBFloat16, "weight must be bfloat16");
+    TORCH_CHECK(output.dtype() == at::kFloat, "output must be float32");
+
+    TORCH_CHECK(input.dim() == 2, "input must be 2D");
+    TORCH_CHECK(weight_bf16.dim() == 1, "weight must be 1D");
+    TORCH_CHECK(output.dim() == 2, "output must be 2D");
+
+    const int64_t num_tokens = input.size(0);
+    const int64_t num_channels = input.size(1);
+    TORCH_CHECK(output.size(0) == num_tokens && output.size(1) == num_channels,
+        "output shape must match input shape");
+    TORCH_CHECK(weight_bf16.size(0) == num_channels,
+        "weight length must equal number of channels");
+    TORCH_CHECK(num_channels % 4 == 0, "num_channels must be divisible by 4");
+
+    auto input_cpu = to_cpu_contiguous(input);
+    auto weight_cpu = to_cpu_contiguous(weight_bf16);
+    auto out_cpu = empty_cpu_like(output);
+
+    const size_t in_bytes = static_cast<size_t>(input_cpu.numel()) * input_cpu.element_size();
+    const size_t weight_bytes = static_cast<size_t>(weight_cpu.numel()) * weight_cpu.element_size();
+    const size_t out_bytes = static_cast<size_t>(out_cpu.numel()) * out_cpu.element_size();
+
+    MetalBuffer input_buf;
+    MetalBuffer weight_buf;
+    MetalBuffer out_buf;
+    MetalBuffer control_buf;
+
+    run_metal_kernel("gptoss_f32_bf16w_rmsnorm", [&](const gptoss_metal_device& device,
+                                                       const gptoss_metal_function& fn,
+                                                       gptoss_metal_command_buffer& cb) {
+        input_buf.wrap(&device, in_bytes, input_cpu.data_ptr());
+        weight_buf.wrap(&device, weight_bytes, weight_cpu.data_ptr());
+        out_buf.create(&device, out_bytes, nullptr);
+        create_control_buffer(&device, control_buf);
+
+        TORCH_CHECK(
+            gptoss_metal_command_buffer_encode_launch_f32_bf16w_rmsnorm(
+                &cb, &fn,
+                input_buf.get(), 0,
+                weight_buf.get(), 0,
+                out_buf.get(), 0,
+                control_buf.get(), 0,
+                static_cast<uint32_t>(num_tokens),
+                static_cast<uint32_t>(num_channels),
+                static_cast<float>(epsilon)) == gptoss_status_success,
+            "encode rmsnorm failed");
+    });
+
+    std::memcpy(out_cpu.data_ptr(), out_buf.ptr(), out_bytes);
+    copy_back(output, out_cpu);
+}
+
+void f32_bf16w_dense_matmul_qkv_torch(const at::Tensor& input,
+    const at::Tensor& weight_bf16,
+    const at::Tensor& bias_bf16,
+    at::Tensor& output)
+{
+    run_dense_matmul_bf16(
+        "gptoss_f32_bf16w_dense_matmul_qkv",
+        gptoss_metal_command_buffer_encode_launch_f32_bf16w_dense_matmul_qkv,
+        input, weight_bf16, bias_bf16, output);
+}
+
+void f32_bf16w_dense_matmul_attn_output_torch(const at::Tensor& input,
+    const at::Tensor& weight_bf16,
+    const at::Tensor& bias_bf16,
+    at::Tensor& output)
+{
+    run_dense_matmul_bf16(
+        "gptoss_f32_bf16w_dense_matmul_attn_output",
+        gptoss_metal_command_buffer_encode_launch_f32_bf16w_dense_matmul_attn_output,
+        input, weight_bf16, bias_bf16, output);
+}
+
+void f32_bf16w_dense_matmul_mlp_gate_torch(const at::Tensor& input,
+    const at::Tensor& weight_bf16,
+    const at::Tensor& bias_bf16,
+    at::Tensor& output)
+{
+    run_dense_matmul_bf16(
+        "gptoss_f32_bf16w_dense_matmul_mlp_gate",
+        gptoss_metal_command_buffer_encode_launch_f32_bf16w_dense_matmul_mlp_gate,
+        input, weight_bf16, bias_bf16, output);
+}
+
+void f32_rope_torch(at::Tensor& activations,
+    double rope_base,
+    double interpolation_scale,
+    double yarn_offset,
+    double yarn_scale,
+    double yarn_multiplier,
+    int64_t num_tokens,
+    int64_t num_q_heads,
+    int64_t num_kv_heads,
+    int64_t attn_head_dim,
+    int64_t token_offset,
+    int64_t threadgroup_size)
+{
+    TORCH_CHECK(activations.dtype() == at::kFloat, "activations must be float32");
+    TORCH_CHECK(num_tokens >= 0 && num_q_heads >= 0 && num_kv_heads >= 0 && attn_head_dim >= 0,
+        "shape parameters must be non-negative");
+    TORCH_CHECK(threadgroup_size >= 0, "threadgroup_size must be non-negative");
+
+    auto activations_cpu = to_cpu_contiguous(activations);
+    MetalBuffer activations_buf;
+    MetalBuffer control_buf;
+
+    const size_t activations_bytes = static_cast<size_t>(activations_cpu.numel()) * activations_cpu.element_size();
+
+    run_metal_kernel("gptoss_f32_rope", [&](const gptoss_metal_device& device,
+                                             const gptoss_metal_function& fn,
+                                             gptoss_metal_command_buffer& cb) {
+        activations_buf.wrap(&device, activations_bytes, activations_cpu.data_ptr());
+        create_control_buffer(&device, control_buf);
+
+        TORCH_CHECK(
+            gptoss_metal_command_buffer_encode_launch_f32_rope(
+                &cb, &fn,
+                static_cast<size_t>(threadgroup_size),
+                activations_buf.get(), 0,
+                control_buf.get(), 0,
+                static_cast<float>(rope_base),
+                static_cast<float>(interpolation_scale),
+                static_cast<float>(yarn_offset),
+                static_cast<float>(yarn_scale),
+                static_cast<float>(yarn_multiplier),
+                static_cast<uint32_t>(num_tokens),
+                static_cast<uint32_t>(num_q_heads),
+                static_cast<uint32_t>(num_kv_heads),
+                static_cast<uint32_t>(attn_head_dim),
+                static_cast<uint32_t>(token_offset)) == gptoss_status_success,
+            "encode rope failed");
+    });
+
+    copy_back(activations, activations_cpu);
+}
+
+void f32_bf16w_matmul_qkv_torch(const at::Tensor& input,
+    const at::Tensor& weight_bf16,
+    const at::Tensor& bias_bf16,
+    at::Tensor& output,
+    at::Tensor& kv_cache,
+    int64_t kv_cache_offset_bytes,
+    int64_t num_tokens,
+    int64_t num_cols,
+    int64_t num_q_heads,
+    int64_t num_kv_heads,
+    int64_t attn_head_dim,
+    int64_t token_offset,
+    int64_t max_tokens,
+    double rope_base,
+    double interpolation_scale,
+    double yarn_offset,
+    double yarn_scale,
+    double yarn_multiplier,
+    int64_t threadgroup_size)
+{
+    TORCH_CHECK(input.dtype() == at::kFloat, "input must be float32");
+    TORCH_CHECK(weight_bf16.dtype() == at::kBFloat16, "weight must be bfloat16");
+    TORCH_CHECK(bias_bf16.dtype() == at::kBFloat16, "bias must be bfloat16");
+    TORCH_CHECK(output.dtype() == at::kFloat, "output must be float32");
+    TORCH_CHECK(kv_cache.dtype() == at::kFloat, "kv_cache must be float32");
+
+    TORCH_CHECK(input.dim() == 2, "input must be 2D");
+    TORCH_CHECK(weight_bf16.dim() == 2, "weight must be 2D");
+    TORCH_CHECK(bias_bf16.dim() == 1, "bias must be 1D");
+    TORCH_CHECK(output.dim() == 2, "output must be 2D");
+
+    TORCH_CHECK(num_tokens >= 0 && num_cols >= 0 && num_q_heads >= 0 && num_kv_heads >= 0 && attn_head_dim >= 0 && max_tokens >= 0,
+        "shape parameters must be non-negative");
+    TORCH_CHECK(threadgroup_size >= 0, "threadgroup_size must be non-negative");
+    TORCH_CHECK(kv_cache_offset_bytes >= 0, "kv_cache_offset_bytes must be non-negative");
+
+    TORCH_CHECK(input.size(0) == num_tokens && input.size(1) == num_cols,
+        "input shape must be [num_tokens, num_cols]");
+    const int64_t num_rows = (num_q_heads + 2 * num_kv_heads) * attn_head_dim;
+    TORCH_CHECK(output.size(0) == num_tokens && output.size(1) == num_rows,
+        "output shape must be [num_tokens, (num_q_heads + 2 * num_kv_heads) * attn_head_dim]");
+    TORCH_CHECK(weight_bf16.size(0) == num_cols && weight_bf16.size(1) == num_rows,
+        "weight shape must be [num_cols, (num_q_heads + 2 * num_kv_heads) * attn_head_dim]");
+    TORCH_CHECK(bias_bf16.size(0) == num_rows,
+        "bias length must equal output feature dimension");
+
+    auto input_cpu = to_cpu_contiguous(input);
+    auto weight_cpu = weight_bf16.transpose(0, 1).contiguous().to(at::kCPU);
+    auto bias_cpu = to_cpu_contiguous(bias_bf16);
+    auto out_cpu = empty_cpu_like(output);
+    auto kv_cpu = to_cpu_contiguous(kv_cache);
+
+    const size_t in_bytes = static_cast<size_t>(input_cpu.numel()) * input_cpu.element_size();
+    const size_t weight_bytes = static_cast<size_t>(weight_cpu.numel()) * weight_cpu.element_size();
+    const size_t bias_bytes = static_cast<size_t>(bias_cpu.numel()) * bias_cpu.element_size();
+    const size_t out_bytes = static_cast<size_t>(out_cpu.numel()) * out_cpu.element_size();
+    const size_t kv_bytes = static_cast<size_t>(kv_cpu.numel()) * kv_cpu.element_size();
+
+    MetalBuffer input_buf;
+    MetalBuffer weight_buf;
+    MetalBuffer bias_buf;
+    MetalBuffer out_buf;
+    MetalBuffer kv_buf;
+    MetalBuffer control_buf;
+
+    run_metal_kernel("gptoss_f32_bf16w_matmul_qkv", [&](const gptoss_metal_device& device,
+                                                          const gptoss_metal_function& fn,
+                                                          gptoss_metal_command_buffer& cb) {
+        input_buf.wrap(&device, in_bytes, input_cpu.data_ptr());
+        weight_buf.wrap(&device, weight_bytes, weight_cpu.data_ptr());
+        bias_buf.wrap(&device, bias_bytes, bias_cpu.data_ptr());
+        out_buf.create(&device, out_bytes, nullptr);
+        kv_buf.wrap(&device, kv_bytes, kv_cpu.data_ptr());
+        create_control_buffer(&device, control_buf);
+
+        TORCH_CHECK(
+            gptoss_metal_command_buffer_encode_launch_f32_bf16w_matmul_qkv(
+                &cb, &fn,
+                static_cast<size_t>(threadgroup_size),
+                input_buf.get(), 0,
+                weight_buf.get(), 0,
+                bias_buf.get(), 0,
+                out_buf.get(), 0,
+                kv_buf.get(), static_cast<size_t>(kv_cache_offset_bytes),
+                control_buf.get(), 0,
+                static_cast<uint32_t>(num_tokens),
+                static_cast<uint32_t>(num_cols),
+                static_cast<uint32_t>(num_q_heads),
+                static_cast<uint32_t>(num_kv_heads),
+                static_cast<uint32_t>(attn_head_dim),
+                static_cast<uint32_t>(token_offset),
+                static_cast<uint32_t>(max_tokens),
+                static_cast<float>(rope_base),
+                static_cast<float>(interpolation_scale),
+                static_cast<float>(yarn_offset),
+                static_cast<float>(yarn_scale),
+                static_cast<float>(yarn_multiplier)) == gptoss_status_success,
+            "encode matmul_qkv failed");
+    });
+
+    std::memcpy(out_cpu.data_ptr(), out_buf.ptr(), out_bytes);
+    copy_back(output, out_cpu);
+    copy_back(kv_cache, kv_cpu);
+}
+
+void f32_sdpa_torch(const at::Tensor& q,
+    int64_t q_offset_bytes,
+    const at::Tensor& kv,
+    int64_t kv_offset_bytes,
+    const at::Tensor& s_bf16,
+    int64_t s_offset_bytes,
+    at::Tensor& output,
+    int64_t output_offset_bytes,
+    int64_t window,
+    int64_t kv_stride,
+    int64_t num_q_tokens,
+    int64_t num_kv_tokens,
+    int64_t num_q_heads,
+    int64_t num_kv_heads,
+    int64_t head_dim)
+{
+    TORCH_CHECK(q.dtype() == at::kFloat, "q must be float32");
+    TORCH_CHECK(kv.dtype() == at::kFloat, "kv must be float32");
+    TORCH_CHECK(s_bf16.dtype() == at::kBFloat16, "s must be bfloat16");
+    TORCH_CHECK(output.dtype() == at::kFloat, "output must be float32");
+
+    TORCH_CHECK(q_offset_bytes >= 0 && kv_offset_bytes >= 0 && s_offset_bytes >= 0 && output_offset_bytes >= 0,
+        "offsets must be non-negative");
+    TORCH_CHECK(window >= 0 && kv_stride >= 0 && num_q_tokens >= 0 && num_kv_tokens >= 0 && num_q_heads >= 0 && num_kv_heads >= 0 && head_dim >= 0,
+        "shape parameters must be non-negative");
+
+    auto q_cpu = to_cpu_contiguous(q);
+    auto kv_cpu = to_cpu_contiguous(kv);
+    auto s_cpu = to_cpu_contiguous(s_bf16);
+    auto out_cpu = empty_cpu_like(output);
+
+    const size_t q_bytes = static_cast<size_t>(q_cpu.numel()) * q_cpu.element_size();
+    const size_t kv_bytes = static_cast<size_t>(kv_cpu.numel()) * kv_cpu.element_size();
+    const size_t s_bytes = static_cast<size_t>(s_cpu.numel()) * s_cpu.element_size();
+    const size_t out_bytes = static_cast<size_t>(out_cpu.numel()) * out_cpu.element_size();
+
+    MetalBuffer q_buf;
+    MetalBuffer kv_buf;
+    MetalBuffer s_buf;
+    MetalBuffer out_buf;
+    MetalBuffer control_buf;
+
+    run_metal_kernel("gptoss_f32_sdpa_q8_d64", [&](const gptoss_metal_device& device,
+                                                     const gptoss_metal_function& fn,
+                                                     gptoss_metal_command_buffer& cb) {
+        q_buf.wrap(&device, q_bytes, q_cpu.data_ptr());
+        kv_buf.wrap(&device, kv_bytes, kv_cpu.data_ptr());
+        s_buf.wrap(&device, s_bytes, s_cpu.data_ptr());
+        out_buf.create(&device, out_bytes, nullptr);
+        create_control_buffer(&device, control_buf);
+
+        TORCH_CHECK(
+            gptoss_metal_command_buffer_encode_launch_f32_sdpa(
+                &cb, &fn,
+                q_buf.get(), static_cast<size_t>(q_offset_bytes),
+                kv_buf.get(), static_cast<size_t>(kv_offset_bytes),
+                s_buf.get(), static_cast<size_t>(s_offset_bytes),
+                out_buf.get(), static_cast<size_t>(output_offset_bytes),
+                control_buf.get(), 0,
+                static_cast<uint32_t>(window),
+                static_cast<uint32_t>(kv_stride),
+                static_cast<uint32_t>(num_q_tokens),
+                static_cast<uint32_t>(num_kv_tokens),
+                static_cast<uint32_t>(num_q_heads),
+                static_cast<uint32_t>(num_kv_heads),
+                static_cast<uint32_t>(head_dim)) == gptoss_status_success,
+            "encode sdpa failed");
+    });
+
+    std::memcpy(out_cpu.data_ptr(), out_buf.ptr(), out_bytes);
+    copy_back(output, out_cpu);
+}
+
+void f32_topk_torch(const at::Tensor& scores,
+    at::Tensor& expert_ids,
+    at::Tensor& expert_scores,
+    int64_t num_tokens,
+    int64_t num_experts,
+    int64_t num_active_experts)
+{
+    TORCH_CHECK(scores.dtype() == at::kFloat, "scores must be float32");
+    TORCH_CHECK(expert_ids.dtype() == at::kInt, "expert_ids must be int32");
+    TORCH_CHECK(expert_scores.dtype() == at::kFloat, "expert_scores must be float32");
+
+    TORCH_CHECK(num_tokens >= 0 && num_experts >= 0 && num_active_experts >= 0,
+        "shape parameters must be non-negative");
+
+    TORCH_CHECK(scores.size(0) == num_tokens,
+        "scores first dimension must match num_tokens");
+    TORCH_CHECK(scores.numel() == num_tokens * num_experts,
+        "scores must have num_tokens * num_experts elements");
+    TORCH_CHECK(expert_ids.numel() == num_tokens * num_active_experts,
+        "expert_ids must have num_tokens * num_active_experts elements");
+    TORCH_CHECK(expert_scores.numel() == num_tokens * num_active_experts,
+        "expert_scores must have num_tokens * num_active_experts elements");
+
+    auto scores_cpu = to_cpu_contiguous(scores);
+    std::vector<gptoss_expert_prediction> predictions(static_cast<size_t>(num_tokens) * static_cast<size_t>(num_active_experts));
+
+    const size_t score_bytes = static_cast<size_t>(scores_cpu.numel()) * scores_cpu.element_size();
+    const size_t pred_bytes = predictions.size() * sizeof(gptoss_expert_prediction);
+
+    MetalBuffer score_buf;
+    MetalBuffer pred_buf;
+    MetalBuffer control_buf;
+
+    run_metal_kernel("gptoss_f32_topk_softmax_e128_k4", [&](const gptoss_metal_device& device,
+                                                             const gptoss_metal_function& fn,
+                                                             gptoss_metal_command_buffer& cb) {
+        score_buf.wrap(&device, score_bytes, scores_cpu.data_ptr());
+        pred_buf.wrap(&device, pred_bytes, predictions.data());
+        create_control_buffer(&device, control_buf);
+
+        TORCH_CHECK(
+            gptoss_metal_command_buffer_encode_launch_f32_topk(
+                &cb, &fn,
+                score_buf.get(), 0,
+                pred_buf.get(), 0,
+                control_buf.get(), 0,
+                static_cast<uint32_t>(num_tokens),
+                static_cast<uint32_t>(num_experts),
+                static_cast<uint32_t>(num_active_experts)) == gptoss_status_success,
+            "encode topk failed");
+    });
+
+    auto ids_cpu = expert_ids.to(at::kCPU).contiguous();
+    auto scores_out_cpu = expert_scores.to(at::kCPU).contiguous();
+    auto* ids_ptr = ids_cpu.data_ptr<int32_t>();
+    auto* scores_ptr = scores_out_cpu.data_ptr<float>();
+    const size_t total = predictions.size();
+    for (size_t i = 0; i < total; ++i) {
+        ids_ptr[i] = static_cast<int32_t>(predictions[i].expert_id);
+        scores_ptr[i] = predictions[i].score;
+    }
+    copy_back(expert_ids, ids_cpu);
+    copy_back(expert_scores, scores_out_cpu);
+}
+
+void expert_routing_metadata_torch(const at::Tensor& expert_ids,
+    const at::Tensor& expert_scores,
+    at::Tensor& expert_offsets,
+    at::Tensor& intra_expert_offsets,
+    int64_t num_tokens,
+    int64_t num_experts)
+{
+    TORCH_CHECK(expert_ids.dtype() == at::kInt, "expert_ids must be int32");
+    TORCH_CHECK(expert_scores.dtype() == at::kFloat, "expert_scores must be float32");
+    TORCH_CHECK(expert_offsets.dtype() == at::kInt, "expert_offsets must be int32");
+    TORCH_CHECK(intra_expert_offsets.dtype() == at::kInt, "intra_expert_offsets must be int32");
+
+    TORCH_CHECK(num_tokens >= 0 && num_experts >= 0, "shape parameters must be non-negative");
+    TORCH_CHECK(expert_ids.numel() == num_tokens,
+        "expert_ids must have num_tokens elements");
+    TORCH_CHECK(expert_scores.numel() == num_tokens,
+        "expert_scores must have num_tokens elements");
+    TORCH_CHECK(intra_expert_offsets.numel() == num_tokens,
+        "intra_expert_offsets must have num_tokens elements");
+    TORCH_CHECK(expert_offsets.numel() == num_experts + 1,
+        "expert_offsets must have num_experts + 1 elements");
+
+    auto ids_cpu = to_cpu_contiguous(expert_ids);
+    auto scores_cpu = to_cpu_contiguous(expert_scores);
+    auto offsets_cpu = to_cpu_contiguous(expert_offsets);
+    auto intra_offsets_cpu = to_cpu_contiguous(intra_expert_offsets);
+
+    std::vector<gptoss_expert_prediction> predictions(static_cast<size_t>(num_tokens));
+    const auto* ids_ptr = ids_cpu.data_ptr<int32_t>();
+    const auto* scores_ptr = scores_cpu.data_ptr<float>();
+    for (int64_t i = 0; i < num_tokens; ++i) {
+        predictions[static_cast<size_t>(i)] = gptoss_expert_prediction {
+            .expert_id = static_cast<uint32_t>(ids_ptr[i]),
+            .score = scores_ptr[i],
+        };
+    }
+
+    const size_t pred_bytes = predictions.size() * sizeof(gptoss_expert_prediction);
+    const size_t offsets_bytes = static_cast<size_t>(offsets_cpu.numel()) * offsets_cpu.element_size();
+    const size_t intra_bytes = static_cast<size_t>(intra_offsets_cpu.numel()) * intra_offsets_cpu.element_size();
+
+    MetalBuffer pred_buf;
+    MetalBuffer offsets_buf;
+    MetalBuffer intra_offsets_buf;
+
+    run_metal_kernel("gptoss_f32_expert_routing_metadata", [&](const gptoss_metal_device& device,
+                                                                  const gptoss_metal_function& fn,
+                                                                  gptoss_metal_command_buffer& cb) {
+        pred_buf.wrap(&device, pred_bytes, predictions.data());
+        offsets_buf.wrap(&device, offsets_bytes, offsets_cpu.data_ptr());
+        intra_offsets_buf.wrap(&device, intra_bytes, intra_offsets_cpu.data_ptr());
+
+        TORCH_CHECK(
+            gptoss_metal_command_buffer_encode_launch_expert_routing_metadata(
+                &cb, &fn,
+                pred_buf.get(), 0,
+                offsets_buf.get(), 0,
+                intra_offsets_buf.get(), 0,
+                static_cast<uint32_t>(num_tokens),
+                static_cast<uint32_t>(num_experts)) == gptoss_status_success,
+            "encode expert_routing_metadata failed");
+    });
+
+    copy_back(expert_offsets, offsets_cpu);
+    copy_back(intra_expert_offsets, intra_offsets_cpu);
+}
+
+void f32_scatter_torch(const at::Tensor& input,
+    const at::Tensor& expert_ids,
+    const at::Tensor& expert_scores,
+    const at::Tensor& expert_offsets,
+    const at::Tensor& intra_expert_offsets,
+    at::Tensor& output,
+    int64_t num_channels,
+    int64_t num_tokens,
+    int64_t num_active_experts)
+{
+    TORCH_CHECK(input.dtype() == at::kFloat, "input must be float32");
+    TORCH_CHECK(expert_ids.dtype() == at::kInt, "expert_ids must be int32");
+    TORCH_CHECK(expert_scores.dtype() == at::kFloat, "expert_scores must be float32");
+    TORCH_CHECK(expert_offsets.dtype() == at::kInt, "expert_offsets must be int32");
+    TORCH_CHECK(intra_expert_offsets.dtype() == at::kInt, "intra_expert_offsets must be int32");
+    TORCH_CHECK(output.dtype() == at::kFloat, "output must be float32");
+
+    TORCH_CHECK(num_channels >= 0 && num_tokens >= 0 && num_active_experts >= 0,
+        "shape parameters must be non-negative");
+
+    TORCH_CHECK(input.numel() == static_cast<int64_t>(num_tokens / num_active_experts) * num_channels,
+        "input size mismatch");
+    TORCH_CHECK(expert_ids.numel() == num_tokens,
+        "expert_ids must have num_tokens elements");
+    TORCH_CHECK(expert_scores.numel() == num_tokens,
+        "expert_scores must have num_tokens elements");
+    TORCH_CHECK(intra_expert_offsets.numel() == num_tokens,
+        "intra_expert_offsets must have num_tokens elements");
+    TORCH_CHECK(output.numel() == num_tokens * num_channels / num_active_experts,
+        "output size mismatch");
+
+    auto input_cpu = to_cpu_contiguous(input);
+    auto expert_offsets_cpu = to_cpu_contiguous(expert_offsets);
+    auto intra_offsets_cpu = to_cpu_contiguous(intra_expert_offsets);
+    auto output_cpu = empty_cpu_like(output);
+
+    std::vector<gptoss_expert_prediction> predictions(static_cast<size_t>(num_tokens));
+    const auto* ids_ptr = expert_ids.to(at::kCPU).contiguous().data_ptr<int32_t>();
+    const auto* scores_ptr = expert_scores.to(at::kCPU).contiguous().data_ptr<float>();
+    for (int64_t i = 0; i < num_tokens; ++i) {
+        predictions[static_cast<size_t>(i)] = gptoss_expert_prediction {
+            .expert_id = static_cast<uint32_t>(ids_ptr[i]),
+            .score = scores_ptr[i],
+        };
+    }
+
+    const size_t input_bytes = static_cast<size_t>(input_cpu.numel()) * input_cpu.element_size();
+    const size_t pred_bytes = predictions.size() * sizeof(gptoss_expert_prediction);
+    const size_t offsets_bytes = static_cast<size_t>(expert_offsets_cpu.numel()) * expert_offsets_cpu.element_size();
+    const size_t intra_bytes = static_cast<size_t>(intra_offsets_cpu.numel()) * intra_offsets_cpu.element_size();
+    const size_t output_bytes = static_cast<size_t>(output_cpu.numel()) * output_cpu.element_size();
+
+    MetalBuffer input_buf;
+    MetalBuffer pred_buf;
+    MetalBuffer offsets_buf;
+    MetalBuffer intra_offsets_buf;
+    MetalBuffer output_buf;
+
+    run_metal_kernel("gptoss_f32_scatter_e4", [&](const gptoss_metal_device& device,
+                                                    const gptoss_metal_function& fn,
+                                                    gptoss_metal_command_buffer& cb) {
+        input_buf.wrap(&device, input_bytes, input_cpu.data_ptr());
+        pred_buf.wrap(&device, pred_bytes, predictions.data());
+        offsets_buf.wrap(&device, offsets_bytes, expert_offsets_cpu.data_ptr());
+        intra_offsets_buf.wrap(&device, intra_bytes, intra_offsets_cpu.data_ptr());
+        output_buf.create(&device, output_bytes, nullptr);
+
+        TORCH_CHECK(
+            gptoss_metal_command_buffer_encode_launch_f32_scatter(
+                &cb, &fn,
+                input_buf.get(), 0,
+                pred_buf.get(), 0,
+                offsets_buf.get(), 0,
+                intra_offsets_buf.get(), 0,
+                output_buf.get(), 0,
+                static_cast<uint32_t>(num_channels),
+                static_cast<uint32_t>(num_tokens / num_active_experts),
+                static_cast<uint32_t>(num_active_experts)) == gptoss_status_success,
+            "encode scatter failed");
+    });
+
+    std::memcpy(output_cpu.data_ptr(), output_buf.ptr(), output_bytes);
+    copy_back(output, output_cpu);
+}
+
+void f32_bf16w_matmul_add_torch(const at::Tensor& input,
+    const at::Tensor& weight_bf16,
+    const at::Tensor& bias_bf16,
+    at::Tensor& output,
+    int64_t num_tokens,
+    int64_t num_cols,
+    int64_t num_rows,
+    int64_t threadgroup_size)
+{
+    TORCH_CHECK(input.dtype() == at::kFloat, "input must be float32");
+    TORCH_CHECK(weight_bf16.dtype() == at::kBFloat16, "weight must be bfloat16");
+    TORCH_CHECK(bias_bf16.dtype() == at::kBFloat16, "bias must be bfloat16");
+    TORCH_CHECK(output.dtype() == at::kFloat, "output must be float32");
+
+    TORCH_CHECK(input.dim() == 2, "input must be 2D");
+    TORCH_CHECK(weight_bf16.dim() == 2, "weight must be 2D");
+    TORCH_CHECK(bias_bf16.dim() == 1, "bias must be 1D");
+    TORCH_CHECK(output.dim() == 2, "output must be 2D");
+
+    TORCH_CHECK(input.size(0) == num_tokens && input.size(1) == num_cols,
+        "input shape must be [num_tokens, num_cols]");
+    TORCH_CHECK(weight_bf16.size(0) == num_cols && weight_bf16.size(1) == num_rows,
+        "weight shape must be [num_cols, num_rows]");
+    TORCH_CHECK(bias_bf16.size(0) == num_rows,
+        "bias length must equal num_rows");
+    TORCH_CHECK(output.size(0) == num_tokens && output.size(1) == num_rows,
+        "output shape must be [num_tokens, num_rows]");
+
+    auto input_cpu = to_cpu_contiguous(input);
+    auto weight_cpu = weight_bf16.transpose(0, 1).contiguous().to(at::kCPU);
+    auto bias_cpu = to_cpu_contiguous(bias_bf16);
+    auto out_cpu = to_cpu_contiguous(output);
+
+    const size_t in_bytes = static_cast<size_t>(input_cpu.numel()) * input_cpu.element_size();
+    const size_t weight_bytes = static_cast<size_t>(weight_cpu.numel()) * weight_cpu.element_size();
+    const size_t bias_bytes = static_cast<size_t>(bias_cpu.numel()) * bias_cpu.element_size();
+    const size_t out_bytes = static_cast<size_t>(out_cpu.numel()) * out_cpu.element_size();
+
+    MetalBuffer input_buf;
+    MetalBuffer weight_buf;
+    MetalBuffer bias_buf;
+    MetalBuffer out_buf;
+    MetalBuffer control_buf;
+
+    run_metal_kernel("gptoss_f32_bf16w_matmul", [&](const gptoss_metal_device& device,
+                                                     const gptoss_metal_function& fn,
+                                                     gptoss_metal_command_buffer& cb) {
+        input_buf.wrap(&device, in_bytes, input_cpu.data_ptr());
+        weight_buf.wrap(&device, weight_bytes, weight_cpu.data_ptr());
+        bias_buf.wrap(&device, bias_bytes, bias_cpu.data_ptr());
+        out_buf.create(&device, out_bytes, nullptr);
+        std::memcpy(out_buf.ptr(), out_cpu.data_ptr(), out_bytes);
+        create_control_buffer(&device, control_buf);
+
+        TORCH_CHECK(
+            gptoss_metal_command_buffer_encode_launch_f32_bf16w_matmul_add(
+                &cb, &fn,
+                static_cast<size_t>(threadgroup_size),
+                input_buf.get(), 0,
+                weight_buf.get(), 0,
+                bias_buf.get(), 0,
+                out_buf.get(), 0,
+                control_buf.get(), 0,
+                static_cast<uint32_t>(num_tokens),
+                static_cast<uint32_t>(num_cols),
+                static_cast<uint32_t>(num_rows)) == gptoss_status_success,
+            "encode matmul_add failed");
+    });
 
-    TORCH_CHECK(gptoss_metal_device_create_system_default(&device) == gptoss_status_success, "device_create failed");
-    TORCH_CHECK(gptoss_metal_library_create_default(&device, &library) == gptoss_status_success, "library_create failed");
-    TORCH_CHECK(gptoss_metal_function_create(&library, "gptoss_f32_bf16w_matmul", &fn) == gptoss_status_success, "function_create failed");
-    TORCH_CHECK(gptoss_metal_command_queue_create(&device, &cq) == gptoss_status_success, "cq_create failed");
-    TORCH_CHECK(gptoss_metal_command_buffer_create(&cq, &cb) == gptoss_status_success, "cb_create failed");
-
-    const size_t in_bytes  = (size_t)num_tokens * (size_t)num_cols * sizeof(float);
-    const size_t wt_bytes  = (size_t)num_rows   * (size_t)num_cols * sizeof(uint16_t);
-    const size_t bs_bytes  = (size_t)num_rows * sizeof(uint16_t);
-    const size_t out_bytes = (size_t)num_tokens * (size_t)num_rows * sizeof(float);
-
-    gptoss_metal_buffer in_buf{}, wt_buf{}, bs_buf{}, out_buf{}, ctrl_buf{};
-    TORCH_CHECK(gptoss_metal_buffer_wrap(&device, in_bytes,  input_cpu.data_ptr(),  &in_buf)  == gptoss_status_success, "wrap input failed");
-    TORCH_CHECK(gptoss_metal_buffer_wrap(&device, wt_bytes,  weight_cpu.data_ptr(), &wt_buf)  == gptoss_status_success, "wrap weight failed");
-    TORCH_CHECK(gptoss_metal_buffer_wrap(&device, bs_bytes,  bias_cpu.data_ptr(),   &bs_buf)  == gptoss_status_success, "wrap bias failed");
-    TORCH_CHECK(gptoss_metal_buffer_create(&device, out_bytes, nullptr, &out_buf)   == gptoss_status_success, "alloc out failed");
-    uint32_t ctrl_zero = 0;
-    TORCH_CHECK(gptoss_metal_buffer_create(&device, sizeof(uint32_t), &ctrl_zero, &ctrl_buf) == gptoss_status_success, "alloc ctrl failed");
-
-    TORCH_CHECK(gptoss_metal_command_buffer_encode_launch_f32_bf16w_matmul(
-    &cb, &fn, (size_t)threadgroup_size,
-    &in_buf, 0, &wt_buf, 0, &bs_buf, 0, &out_buf, 0, &ctrl_buf, 0,
-    (uint32_t)num_tokens, (uint32_t)num_cols, (uint32_t)num_rows) == gptoss_status_success, "encode failed");
-
-    TORCH_CHECK(gptoss_metal_command_buffer_commit(&cb) == gptoss_status_success, "commit failed");
-    TORCH_CHECK(gptoss_metal_command_buffer_wait_completion(&cb, nullptr) == gptoss_status_success, "wait failed");
-
-    std::memcpy(out_cpu.data_ptr(), out_buf.ptr, out_bytes);
-    output.copy_(out_cpu.to(output.device(), /*non_blocking=*/false, /*copy=*/true));
-
-    (void) gptoss_metal_command_buffer_release(&cb);
-    (void) gptoss_metal_command_queue_release(&cq);
-    (void) gptoss_metal_function_release(&fn);
-    (void) gptoss_metal_library_release(&library);
-    (void) gptoss_metal_device_release(&device);
-    (void) gptoss_metal_buffer_release(&ctrl_buf);
-    (void) gptoss_metal_buffer_release(&out_buf);
-    (void) gptoss_metal_buffer_release(&bs_buf);
-    (void) gptoss_metal_buffer_release(&wt_buf);
-    (void) gptoss_metal_buffer_release(&in_buf);
+    std::memcpy(out_cpu.data_ptr(), out_buf.ptr(), out_bytes);
+    copy_back(output, out_cpu);
 }

torch-ext/gptoss_kernels/__init__.py

@@ -1,8 +1,174 @@
 from ._ops import ops
 import torch
 
-def f32_bf16w_matmul(input: torch.Tensor, weight_bf16: torch.Tensor, bias_bf16: torch.Tensor, output: torch.Tensor, num_tokens: int, num_cols: int, num_rows: int, threadgroup_size: int) -> None:
-    ops.f32_bf16w_matmul(input, weight_bf16, bias_bf16, output, num_tokens, num_cols, num_rows, threadgroup_size)
+def f32_bf16w_matmul(input: torch.Tensor,
+                     weight_bf16: torch.Tensor,
+                     bias_bf16: torch.Tensor,
+                     output: torch.Tensor,
+                     num_tokens: int,
+                     num_cols: int,
+                     num_rows: int,
+                     threadgroup_size: int) -> torch.Tensor:
+    ops.f32_bf16w_matmul(input, weight_bf16, bias_bf16, output,
+                         num_tokens, num_cols, num_rows, threadgroup_size)
     return output
 
-__all__ = ["f32_bf16w_matmul"]
+def bf16_f32_embeddings(token_ids: torch.Tensor,
+                        weight_bf16: torch.Tensor,
+                        output: torch.Tensor,
+                        threadgroup_size: int) -> torch.Tensor:
+    ops.bf16_f32_embeddings(token_ids, weight_bf16, output, threadgroup_size)
+    return output
+
+def f32_bf16w_rmsnorm(input: torch.Tensor,
+                      weight_bf16: torch.Tensor,
+                      output: torch.Tensor,
+                      epsilon: float) -> torch.Tensor:
+    ops.f32_bf16w_rmsnorm(input, weight_bf16, output, epsilon)
+    return output
+
+def f32_bf16w_dense_matmul_qkv(input: torch.Tensor,
+                               weight_bf16: torch.Tensor,
+                               bias_bf16: torch.Tensor,
+                               output: torch.Tensor) -> torch.Tensor:
+    ops.f32_bf16w_dense_matmul_qkv(input, weight_bf16, bias_bf16, output)
+    return output
+
+def f32_bf16w_dense_matmul_attn_output(input: torch.Tensor,
+                                       weight_bf16: torch.Tensor,
+                                       bias_bf16: torch.Tensor,
+                                       output: torch.Tensor) -> torch.Tensor:
+    ops.f32_bf16w_dense_matmul_attn_output(input, weight_bf16, bias_bf16, output)
+    return output
+
+def f32_bf16w_dense_matmul_mlp_gate(input: torch.Tensor,
+                                    weight_bf16: torch.Tensor,
+                                    bias_bf16: torch.Tensor,
+                                    output: torch.Tensor) -> torch.Tensor:
+    ops.f32_bf16w_dense_matmul_mlp_gate(input, weight_bf16, bias_bf16, output)
+    return output
+
+def f32_rope(activations: torch.Tensor,
+             rope_base: float,
+             interpolation_scale: float,
+             yarn_offset: float,
+             yarn_scale: float,
+             yarn_multiplier: float,
+             num_tokens: int,
+             num_q_heads: int,
+             num_kv_heads: int,
+             attn_head_dim: int,
+             token_offset: int,
+             threadgroup_size: int) -> torch.Tensor:
+    ops.f32_rope(activations, rope_base, interpolation_scale, yarn_offset,
+                 yarn_scale, yarn_multiplier, num_tokens, num_q_heads,
+                 num_kv_heads, attn_head_dim, token_offset, threadgroup_size)
+    return activations
+
+def f32_bf16w_matmul_qkv(input: torch.Tensor,
+                         weight_bf16: torch.Tensor,
+                         bias_bf16: torch.Tensor,
+                         output: torch.Tensor,
+                         kv_cache: torch.Tensor,
+                         kv_cache_offset_bytes: int,
+                         num_tokens: int,
+                         num_cols: int,
+                         num_q_heads: int,
+                         num_kv_heads: int,
+                         attn_head_dim: int,
+                         token_offset: int,
+                         max_tokens: int,
+                         rope_base: float,
+                         interpolation_scale: float,
+                         yarn_offset: float,
+                         yarn_scale: float,
+                         yarn_multiplier: float,
+                         threadgroup_size: int) -> torch.Tensor:
+    ops.f32_bf16w_matmul_qkv(input, weight_bf16, bias_bf16, output, kv_cache,
+                             kv_cache_offset_bytes, num_tokens, num_cols,
+                             num_q_heads, num_kv_heads, attn_head_dim,
+                             token_offset, max_tokens, rope_base,
+                             interpolation_scale, yarn_offset, yarn_scale,
+                             yarn_multiplier, threadgroup_size)
+    return output
+
+def f32_sdpa(q: torch.Tensor,
+             q_offset_bytes: int,
+             kv: torch.Tensor,
+             kv_offset_bytes: int,
+             s_bf16: torch.Tensor,
+             s_offset_bytes: int,
+             output: torch.Tensor,
+             output_offset_bytes: int,
+             window: int,
+             kv_stride: int,
+             num_q_tokens: int,
+             num_kv_tokens: int,
+             num_q_heads: int,
+             num_kv_heads: int,
+             head_dim: int) -> torch.Tensor:
+    ops.f32_sdpa(q, q_offset_bytes, kv, kv_offset_bytes, s_bf16, s_offset_bytes,
+                 output, output_offset_bytes, window, kv_stride,
+                 num_q_tokens, num_kv_tokens, num_q_heads, num_kv_heads, head_dim)
+    return output
+
+def f32_topk(scores: torch.Tensor,
+             expert_ids: torch.Tensor,
+             expert_scores: torch.Tensor,
+             num_tokens: int,
+             num_experts: int,
+             num_active_experts: int) -> None:
+    ops.f32_topk(scores, expert_ids, expert_scores,
+                 num_tokens, num_experts, num_active_experts)
+
+def expert_routing_metadata(expert_ids: torch.Tensor,
+                            expert_scores: torch.Tensor,
+                            expert_offsets: torch.Tensor,
+                            intra_expert_offsets: torch.Tensor,
+                            num_tokens: int,
+                            num_experts: int) -> None:
+    ops.expert_routing_metadata(expert_ids, expert_scores,
+                                expert_offsets, intra_expert_offsets,
+                                num_tokens, num_experts)
+
+def f32_scatter(input: torch.Tensor,
+                expert_ids: torch.Tensor,
+                expert_scores: torch.Tensor,
+                expert_offsets: torch.Tensor,
+                intra_expert_offsets: torch.Tensor,
+                output: torch.Tensor,
+                num_channels: int,
+                num_tokens: int,
+                num_active_experts: int) -> torch.Tensor:
+    ops.f32_scatter(input, expert_ids, expert_scores,
+                    expert_offsets, intra_expert_offsets,
+                    output, num_channels, num_tokens, num_active_experts)
+    return output
+
+def f32_bf16w_matmul_add(input: torch.Tensor,
+                         weight_bf16: torch.Tensor,
+                         bias_bf16: torch.Tensor,
+                         output: torch.Tensor,
+                         num_tokens: int,
+                         num_cols: int,
+                         num_rows: int,
+                         threadgroup_size: int) -> torch.Tensor:
+    ops.f32_bf16w_matmul_add(input, weight_bf16, bias_bf16, output,
+                             num_tokens, num_cols, num_rows, threadgroup_size)
+    return output
+
+__all__ = [
+    "f32_bf16w_matmul",
+    "bf16_f32_embeddings",
+    "f32_bf16w_rmsnorm",
+    "f32_bf16w_dense_matmul_qkv",
+    "f32_bf16w_dense_matmul_attn_output",
+    "f32_bf16w_dense_matmul_mlp_gate",
+    "f32_rope",
+    "f32_bf16w_matmul_qkv",
+    "f32_sdpa",
+    "f32_topk",
+    "expert_routing_metadata",
+    "f32_scatter",
+    "f32_bf16w_matmul_add",
+]

torch-ext/gptoss_kernels/_ops.py

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

torch-ext/torch_binding.cpp

@@ -3,8 +3,58 @@
 #include "registration.h"
 
 TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
-  ops.def("f32_bf16w_matmul(Tensor input, Tensor weight_bf16, Tensor bias_bf16, Tensor output, int num_tokens, int num_cols, int num_rows, int threadgroup_size) -> ()");
+  ops.def("f32_bf16w_matmul(Tensor input, Tensor weight_bf16, Tensor bias_bf16, "
+          "Tensor! output, int num_tokens, int num_cols, int num_rows, int threadgroup_size) -> ()");
   ops.impl("f32_bf16w_matmul", torch::kMPS, &f32_bf16w_matmul_torch);
+
+  ops.def("bf16_f32_embeddings(Tensor token_ids, Tensor weight_bf16, Tensor! output, "
+          "int threadgroup_size) -> ()");
+  ops.impl("bf16_f32_embeddings", torch::kMPS, &bf16_f32_embeddings_torch);
+
+  ops.def("f32_bf16w_rmsnorm(Tensor input, Tensor weight_bf16, Tensor! output, float epsilon) -> ()");
+  ops.impl("f32_bf16w_rmsnorm", torch::kMPS, &f32_bf16w_rmsnorm_torch);
+
+  ops.def("f32_bf16w_dense_matmul_qkv(Tensor input, Tensor weight_bf16, Tensor bias_bf16, Tensor! output) -> ()");
+  ops.impl("f32_bf16w_dense_matmul_qkv", torch::kMPS, &f32_bf16w_dense_matmul_qkv_torch);
+
+  ops.def("f32_bf16w_dense_matmul_attn_output(Tensor input, Tensor weight_bf16, Tensor bias_bf16, Tensor! output) -> ()");
+  ops.impl("f32_bf16w_dense_matmul_attn_output", torch::kMPS, &f32_bf16w_dense_matmul_attn_output_torch);
+
+  ops.def("f32_bf16w_dense_matmul_mlp_gate(Tensor input, Tensor weight_bf16, Tensor bias_bf16, Tensor! output) -> ()");
+  ops.impl("f32_bf16w_dense_matmul_mlp_gate", torch::kMPS, &f32_bf16w_dense_matmul_mlp_gate_torch);
+
+  ops.def("f32_rope(Tensor! activations, float rope_base, float interpolation_scale, float yarn_offset, "
+          "float yarn_scale, float yarn_multiplier, int num_tokens, int num_q_heads, int num_kv_heads, "
+          "int attn_head_dim, int token_offset, int threadgroup_size) -> ()");
+  ops.impl("f32_rope", torch::kMPS, &f32_rope_torch);
+
+  ops.def("f32_bf16w_matmul_qkv(Tensor input, Tensor weight_bf16, Tensor bias_bf16, Tensor! output, Tensor kv_cache, "
+          "int kv_cache_offset_bytes, int num_tokens, int num_cols, int num_q_heads, int num_kv_heads, "
+          "int attn_head_dim, int token_offset, int max_tokens, float rope_base, float interpolation_scale, "
+          "float yarn_offset, float yarn_scale, float yarn_multiplier, int threadgroup_size) -> ()");
+  ops.impl("f32_bf16w_matmul_qkv", torch::kMPS, &f32_bf16w_matmul_qkv_torch);
+
+  ops.def("f32_sdpa(Tensor q, int q_offset_bytes, Tensor kv, int kv_offset_bytes, Tensor s_bf16, int s_offset_bytes, "
+          "Tensor! output, int output_offset_bytes, int window, int kv_stride, int num_q_tokens, int num_kv_tokens, "
+          "int num_q_heads, int num_kv_heads, int head_dim) -> ()");
+  ops.impl("f32_sdpa", torch::kMPS, &f32_sdpa_torch);
+
+  ops.def("f32_topk(Tensor scores, Tensor expert_ids, Tensor expert_scores, int num_tokens, int num_experts, "
+          "int num_active_experts) -> ()");
+  ops.impl("f32_topk", torch::kMPS, &f32_topk_torch);
+
+  ops.def("expert_routing_metadata(Tensor expert_ids, Tensor expert_scores, Tensor expert_offsets, "
+          "Tensor intra_expert_offsets, int num_tokens, int num_experts) -> ()");
+  ops.impl("expert_routing_metadata", torch::kMPS, &expert_routing_metadata_torch);
+
+  ops.def("f32_scatter(Tensor input, Tensor expert_ids, Tensor expert_scores, Tensor expert_offsets, "
+          "Tensor intra_expert_offsets, Tensor! output, int num_channels, int num_tokens, "
+          "int num_active_experts) -> ()");
+  ops.impl("f32_scatter", torch::kMPS, &f32_scatter_torch);
+
+  ops.def("f32_bf16w_matmul_add(Tensor input, Tensor weight_bf16, Tensor bias_bf16, Tensor! output, "
+          "int num_tokens, int num_cols, int num_rows, int threadgroup_size) -> ()");
+  ops.impl("f32_bf16w_matmul_add", torch::kMPS, &f32_bf16w_matmul_add_torch);
 }
 
 REGISTER_EXTENSION(TORCH_EXTENSION_NAME)

torch-ext/torch_binding.h

@@ -2,4 +2,118 @@
 
 #include <torch/torch.h>
 
-void f32_bf16w_matmul_torch(const at::Tensor &input, const at::Tensor &weight_bf16, const at::Tensor &bias_bf16, at::Tensor &output, int64_t num_tokens, int64_t num_cols, int64_t num_rows, int64_t threadgroup_size);
+void f32_bf16w_matmul_torch(const at::Tensor& input,
+                            const at::Tensor& weight_bf16,
+                            const at::Tensor& bias_bf16,
+                            at::Tensor& output,
+                            int64_t num_tokens,
+                            int64_t num_cols,
+                            int64_t num_rows,
+                            int64_t threadgroup_size);
+
+void bf16_f32_embeddings_torch(const at::Tensor& token_ids,
+                               const at::Tensor& weight_bf16,
+                               at::Tensor& output,
+                               int64_t threadgroup_size);
+
+void f32_bf16w_rmsnorm_torch(const at::Tensor& input,
+                             const at::Tensor& weight_bf16,
+                             at::Tensor& output,
+                             double epsilon);
+
+void f32_bf16w_dense_matmul_qkv_torch(const at::Tensor& input,
+                                      const at::Tensor& weight_bf16,
+                                      const at::Tensor& bias_bf16,
+                                      at::Tensor& output);
+
+void f32_bf16w_dense_matmul_attn_output_torch(const at::Tensor& input,
+                                              const at::Tensor& weight_bf16,
+                                              const at::Tensor& bias_bf16,
+                                              at::Tensor& output);
+
+void f32_bf16w_dense_matmul_mlp_gate_torch(const at::Tensor& input,
+                                           const at::Tensor& weight_bf16,
+                                           const at::Tensor& bias_bf16,
+                                           at::Tensor& output);
+
+void f32_rope_torch(at::Tensor& activations,
+                    double rope_base,
+                    double interpolation_scale,
+                    double yarn_offset,
+                    double yarn_scale,
+                    double yarn_multiplier,
+                    int64_t num_tokens,
+                    int64_t num_q_heads,
+                    int64_t num_kv_heads,
+                    int64_t attn_head_dim,
+                    int64_t token_offset,
+                    int64_t threadgroup_size);
+
+void f32_bf16w_matmul_qkv_torch(const at::Tensor& input,
+                                const at::Tensor& weight_bf16,
+                                const at::Tensor& bias_bf16,
+                                at::Tensor& output,
+                                at::Tensor& kv_cache,
+                                int64_t kv_cache_offset_bytes,
+                                int64_t num_tokens,
+                                int64_t num_cols,
+                                int64_t num_q_heads,
+                                int64_t num_kv_heads,
+                                int64_t attn_head_dim,
+                                int64_t token_offset,
+                                int64_t max_tokens,
+                                double rope_base,
+                                double interpolation_scale,
+                                double yarn_offset,
+                                double yarn_scale,
+                                double yarn_multiplier,
+                                int64_t threadgroup_size);
+
+void f32_sdpa_torch(const at::Tensor& q,
+                    int64_t q_offset_bytes,
+                    const at::Tensor& kv,
+                    int64_t kv_offset_bytes,
+                    const at::Tensor& s_bf16,
+                    int64_t s_offset_bytes,
+                    at::Tensor& output,
+                    int64_t output_offset_bytes,
+                    int64_t window,
+                    int64_t kv_stride,
+                    int64_t num_q_tokens,
+                    int64_t num_kv_tokens,
+                    int64_t num_q_heads,
+                    int64_t num_kv_heads,
+                    int64_t head_dim);
+
+void f32_topk_torch(const at::Tensor& scores,
+                    at::Tensor& expert_ids,
+                    at::Tensor& expert_scores,
+                    int64_t num_tokens,
+                    int64_t num_experts,
+                    int64_t num_active_experts);
+
+void expert_routing_metadata_torch(const at::Tensor& expert_ids,
+                                   const at::Tensor& expert_scores,
+                                   at::Tensor& expert_offsets,
+                                   at::Tensor& intra_expert_offsets,
+                                   int64_t num_tokens,
+                                   int64_t num_experts);
+
+void f32_scatter_torch(const at::Tensor& input,
+                       const at::Tensor& expert_ids,
+                       const at::Tensor& expert_scores,
+                       const at::Tensor& expert_offsets,
+                       const at::Tensor& intra_expert_offsets,
+                       at::Tensor& output,
+                       int64_t num_channels,
+                       int64_t num_tokens,
+                       int64_t num_active_experts);
+
+void f32_bf16w_matmul_add_torch(const at::Tensor& input,
+                                const at::Tensor& weight_bf16,
+                                const at::Tensor& bias_bf16,
+                                at::Tensor& output,
+                                int64_t num_tokens,
+                                int64_t num_cols,
+                                int64_t num_rows,
+                                int64_t threadgroup_size);