add builds

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

@@ -0,0 +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) -> torch.Tensor:
+    ops.f32_bf16w_matmul(input, weight_bf16, bias_bf16, output,
+                         num_tokens, num_cols, num_rows, threadgroup_size)
+    return output
+
+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/gpt_oss_metal_kernels/_gpt_oss_metal_kernels_4db85a7.abi3.so

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

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

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

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

@@ -0,0 +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) -> torch.Tensor:
+    ops.f32_bf16w_matmul(input, weight_bf16, bias_bf16, output,
+                         num_tokens, num_cols, num_rows, threadgroup_size)
+    return output
+
+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/gpt_oss_metal_kernels/_gpt_oss_metal_kernels_4db85a7.abi3.so

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:51bd6eb3c95b6602dd1ddff1137b8fab4db2cb9781c08f96cf12147f7f1c9fa0
+size 410328

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

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