feat(metail): Parallelize SDPA across multiple simdgroups (#144)

Author: Maratyszcza

gpt_oss/metal/source/include/internal/math.h

@@ -1,5 +1,6 @@
 #pragma once
 
+#include <assert.h>
 #include <stddef.h>
 #include <stdint.h>
 
@@ -19,11 +20,21 @@ inline static size_t math_sub_sat(size_t a, size_t b) {
     return a > b ? a - b : 0;
 }
 
-static size_t math_round_up_po2(size_t bytes, size_t multiple) {
+static size_t math_round_down_po2(size_t number, size_t multiple) {
+    assert(multiple != 0);
+    assert((multiple & (multiple - 1)) == 0);
+
+    return number & -multiple;
+}
+
+static size_t math_round_up_po2(size_t number, size_t multiple) {
+    assert(multiple != 0);
+    assert((multiple & (multiple - 1)) == 0);
+
     const size_t multiple_mask = multiple - 1;
-    if ((bytes & multiple_mask) != 0) {
-        bytes |= multiple_mask;
-        bytes += 1;
+    if ((number & multiple_mask) != 0) {
+        number |= multiple_mask;
+        number += 1;
     }
-    return bytes;
+    return number;
 }

gpt_oss/metal/source/include/internal/metal.h

@@ -118,9 +118,10 @@ enum gptoss_status gptoss_metal_command_buffer_encode_launch_kernel(
     size_t num_threadgroups_z,
     size_t params_size,
     const void* params,
-    size_t num_buffers,
-    const struct gptoss_metal_buffer** buffers,
-    const size_t* buffer_offsets);
+    size_t num_device_buffers,
+    const struct gptoss_metal_buffer** device_buffers,
+    const size_t* device_buffer_offsets,
+    size_t threadgroup_buffer_size);
 
 enum gptoss_status gptoss_metal_command_buffer_commit(
     const struct gptoss_metal_command_buffer* command_buffer);

gpt_oss/metal/source/include/internal/metal.hpp

@@ -246,10 +246,11 @@ class CommandBuffer {
                                      const std::array<size_t, 3>& threadgroup_size,
                                      const std::array<size_t, 3>& num_threadgroups,
                                      size_t params_size, const void* params,
-                                     std::initializer_list<const Buffer*> buffers = {})
+                                     std::initializer_list<const Buffer*> device_buffers = {},
+                                     size_t threadgroup_buffer_size = 0)
     {
-        std::vector<const gptoss_metal_buffer*> buffer_handles(buffers.size());
-        std::transform(buffers.begin(), buffers.end(), buffer_handles.begin(),
+        std::vector<const gptoss_metal_buffer*> buffer_handles(device_buffers.size());
+        std::transform(device_buffers.begin(), device_buffers.end(), buffer_handles.begin(),
             [](const Buffer* buffer) -> const gptoss_metal_buffer* { return buffer->handle(); });
         Check(gptoss_metal_command_buffer_encode_launch_kernel(
                 &command_buffer_, function.handle(),
@@ -258,7 +259,8 @@ class CommandBuffer {
                 params_size, params,
                 buffer_handles.size(),
                 buffer_handles.data(),
-                /*buffer_offsets=*/nullptr),
+                /*buffer_offsets=*/nullptr,
+                threadgroup_buffer_size),
             "gptoss_metal_command_buffer_encode_launch_kernel");
     }
 

gpt_oss/metal/source/metal-kernels.c

@@ -46,7 +46,8 @@ enum gptoss_status gptoss_metal_command_buffer_encode_launch_u32_fill_random(
         threadgroup_size, 1, 1,
         num_threadgroups, 1, 1,
         sizeof(args), &args,
-        1, &output_buffer, &output_offset);
+        1, &output_buffer, &output_offset,
+        /*threadgroup_buffer_size=*/0);
 }
 
 enum gptoss_status gptoss_metal_command_buffer_encode_launch_f32_fill_random(
@@ -93,7 +94,8 @@ enum gptoss_status gptoss_metal_command_buffer_encode_launch_f32_fill_random(
         threadgroup_size, 1, 1,
         num_threadgroups, 1, 1,
         sizeof(args), &args,
-        1, &output_buffer, &output_offset);
+        1, &output_buffer, &output_offset,
+        /*threadgroup_buffer_size=*/0);
 }
 
 enum gptoss_status gptoss_metal_command_buffer_encode_launch_bf16_fill_random(
@@ -140,7 +142,8 @@ enum gptoss_status gptoss_metal_command_buffer_encode_launch_bf16_fill_random(
         threadgroup_size, 1, 1,
         num_threadgroups, 1, 1,
         sizeof(args), &args,
-        1, &output_buffer, &output_offset);
+        1, &output_buffer, &output_offset,
+        /*threadgroup_buffer_size=*/0);
 }
 
 enum gptoss_status gptoss_metal_command_buffer_encode_launch_mf4_f32_convert(
@@ -180,7 +183,8 @@ enum gptoss_status gptoss_metal_command_buffer_encode_launch_mf4_f32_convert(
         threadgroup_size, 1, 1,
         num_threadgroups, 1, 1,
         sizeof(args), &args,
-        3, (const struct gptoss_metal_buffer *[]) {block_buffer, scale_buffer, output_buffer}, NULL);
+        3, (const struct gptoss_metal_buffer *[]) {block_buffer, scale_buffer, output_buffer}, NULL,
+        /*threadgroup_buffer_size=*/0);
 }
 
 enum gptoss_status gptoss_metal_command_buffer_encode_launch_bf16_f32_embeddings(
@@ -222,7 +226,8 @@ enum gptoss_status gptoss_metal_command_buffer_encode_launch_bf16_f32_embeddings
         sizeof(args), &args,
         3,
         (const struct gptoss_metal_buffer *[]) {token_buffer, weight_buffer, output_buffer},
-        (const size_t[]) {token_offset, weight_offset, output_offset});
+        (const size_t[]) {token_offset, weight_offset, output_offset},
+        /*threadgroup_buffer_size=*/0);
 }
 
 enum gptoss_status gptoss_metal_command_buffer_encode_launch_f32_bf16w_rmsnorm(
@@ -268,7 +273,8 @@ enum gptoss_status gptoss_metal_command_buffer_encode_launch_f32_bf16w_rmsnorm(
         sizeof(args), &args,
         3,
         (const struct gptoss_metal_buffer *[]) {input_buffer, weight_buffer, output_buffer},
-        (const size_t[]) {input_offset, weight_offset, output_offset});
+        (const size_t[]) {input_offset, weight_offset, output_offset},
+        /*threadgroup_buffer_size=*/0);
 }
 
 enum gptoss_status gptoss_metal_command_buffer_encode_launch_f32_bf16w_matmul(
@@ -325,7 +331,8 @@ enum gptoss_status gptoss_metal_command_buffer_encode_launch_f32_bf16w_matmul(
         sizeof(args), &args,
         4,
         (const struct gptoss_metal_buffer *[]) {input_buffer, weight_buffer, bias_buffer, output_buffer},
-        (const size_t[]) {input_offset, weight_offset, bias_offset, output_offset});
+        (const size_t[]) {input_offset, weight_offset, bias_offset, output_offset},
+        /*threadgroup_buffer_size=*/0);
 }
 
 enum gptoss_status gptoss_metal_command_buffer_encode_launch_f32_bf16w_matmul_add(
@@ -382,7 +389,8 @@ enum gptoss_status gptoss_metal_command_buffer_encode_launch_f32_bf16w_matmul_ad
         sizeof(args), &args,
         4,
         (const struct gptoss_metal_buffer *[]) {input_buffer, weight_buffer, bias_buffer, output_buffer},
-        (const size_t[]) {input_offset, weight_offset, bias_offset, output_offset});
+        (const size_t[]) {input_offset, weight_offset, bias_offset, output_offset},
+        /*threadgroup_buffer_size=*/0);
 }
 
 enum gptoss_status gptoss_metal_command_buffer_encode_launch_f32_bf16w_unembedding(
@@ -437,7 +445,8 @@ enum gptoss_status gptoss_metal_command_buffer_encode_launch_f32_bf16w_unembeddi
         sizeof(args), &args,
         4,
         (const struct gptoss_metal_buffer *[]) {input_buffer, weight_buffer, output_buffer, argmax_buffer},
-        (const size_t[]) {input_offset, weight_offset, output_offset, argmax_offset});
+        (const size_t[]) {input_offset, weight_offset, output_offset, argmax_offset},
+        /*threadgroup_buffer_size=*/0);
 }
 
 enum gptoss_status gptoss_metal_command_buffer_encode_launch_f32_mf4w_moe_matmul_swiglu(
@@ -510,7 +519,8 @@ enum gptoss_status gptoss_metal_command_buffer_encode_launch_f32_mf4w_moe_matmul
         sizeof(args), &args,
         6,
         (const struct gptoss_metal_buffer *[]) {input_buffer, expert_buffer, weight_block_buffer, weight_scale_buffer, bias_buffer, output_buffer},
-        (const size_t[]) {input_offset, expert_offset, weight_block_offset, weight_scale_offset, bias_offset, output_offset});
+        (const size_t[]) {input_offset, expert_offset, weight_block_offset, weight_scale_offset, bias_offset, output_offset},
+        /*threadgroup_buffer_size=*/0);
 }
 
 enum gptoss_status gptoss_metal_command_buffer_encode_launch_f32_mf4w_moe_matmul(
@@ -581,7 +591,8 @@ enum gptoss_status gptoss_metal_command_buffer_encode_launch_f32_mf4w_moe_matmul
         sizeof(args), &args,
         6,
         (const struct gptoss_metal_buffer *[]) {input_buffer, expert_buffer, weight_block_buffer, weight_scale_buffer, bias_buffer, output_buffer},
-        (const size_t[]) {input_offset, expert_offset, weight_block_offset, weight_scale_offset, bias_offset, output_offset});
+        (const size_t[]) {input_offset, expert_offset, weight_block_offset, weight_scale_offset, bias_offset, output_offset},
+        /*threadgroup_buffer_size=*/0);
 }
 
 enum gptoss_status gptoss_metal_command_buffer_encode_launch_f32_rope(
@@ -631,7 +642,8 @@ enum gptoss_status gptoss_metal_command_buffer_encode_launch_f32_rope(
         threadgroup_size, 1, 1,
         num_qk_heads / num_simdgroups, num_tokens, 1,
         sizeof(args), &args,
-        1, (const struct gptoss_metal_buffer *[]) {activations_buffer}, NULL);
+        1, (const struct gptoss_metal_buffer *[]) {activations_buffer}, NULL,
+        /*threadgroup_buffer_size=*/0);
 }
 
 enum gptoss_status gptoss_metal_command_buffer_encode_launch_f32_accumulate(
@@ -680,7 +692,8 @@ enum gptoss_status gptoss_metal_command_buffer_encode_launch_f32_accumulate(
         sizeof(args), &args,
         3,
         (const struct gptoss_metal_buffer *[]) {input_buffer, expert_buffer, output_buffer},
-        (const size_t[]) {input_offset, expert_offset, output_offset});
+        (const size_t[]) {input_offset, expert_offset, output_offset},
+        /*threadgroup_buffer_size=*/0);
 }
 
 enum gptoss_status gptoss_metal_command_buffer_encode_launch_f32_topk(
@@ -715,7 +728,8 @@ enum gptoss_status gptoss_metal_command_buffer_encode_launch_f32_topk(
         sizeof(args), &args,
         2,
         (const struct gptoss_metal_buffer *[]) {input_buffer, output_buffer},
-        (const size_t[]) {input_offset, output_offset});
+        (const size_t[]) {input_offset, output_offset},
+        /*threadgroup_buffer_size=*/0);
 }
 
 enum gptoss_status gptoss_metal_command_buffer_encode_launch_f32_sdpa(
@@ -753,6 +767,11 @@ enum gptoss_status gptoss_metal_command_buffer_encode_launch_f32_sdpa(
         return gptoss_status_invalid_argument;
     }
 
+    const size_t max_context_tokens = math_min(num_q_tokens + num_kv_tokens + 1, window);
+    const size_t threadgroup_size = math_min(f32_sdpa_fn->max_threadgroup_threads,
+        max_context_tokens * f32_sdpa_fn->simdgroup_threads);
+    const size_t half_threadgroup_size = math_round_down_po2(threadgroup_size / 2, f32_sdpa_fn->simdgroup_threads);
+
     const struct gptoss_sdpa_args args = {
         .qkv_dim = head_dim * (num_q_heads + 2 * num_kv_heads),
         .num_kv_tokens = num_kv_tokens,
@@ -761,12 +780,13 @@ enum gptoss_status gptoss_metal_command_buffer_encode_launch_f32_sdpa(
 
     return gptoss_metal_command_buffer_encode_launch_kernel(
         command_buffer, f32_sdpa_fn,
-        /*threadgroup_size=*/32, 1, 1,
+        threadgroup_size, 1, 1,
         num_q_tokens, num_kv_heads, 1,
         sizeof(args), &args,
         5,
         (const struct gptoss_metal_buffer *[]) {q_buffer, k_buffer, v_buffer, s_buffer, output_buffer},
-        (const size_t[]) {q_offset, k_offset, v_offset, s_offset, output_offset});
+        (const size_t[]) {q_offset, k_offset, v_offset, s_offset, output_offset},
+        /*threadgroup_buffer_size=*/half_threadgroup_size * 8 * 4 * sizeof(float));
 }
 
 enum gptoss_status gptoss_metal_command_buffer_encode_launch_f32_softmax(
@@ -813,5 +833,6 @@ enum gptoss_status gptoss_metal_command_buffer_encode_launch_f32_softmax(
         sizeof(args), &args,
         4,
         (const struct gptoss_metal_buffer *[]) {score_buffer, argmax_buffer, prob_buffer, sum_buffer},
-        (const size_t[]) {score_offset, argmax_offset, prob_offset, sum_offset});
+        (const size_t[]) {score_offset, argmax_offset, prob_offset, sum_offset},
+        /*threadgroup_buffer_size=*/0);
 }

gpt_oss/metal/source/metal.m

@@ -380,9 +380,10 @@ enum gptoss_status gptoss_metal_command_buffer_encode_launch_kernel(
     size_t num_threadgroups_z,
     size_t params_size,
     const void* params,
-    size_t num_buffers,
-    const struct gptoss_metal_buffer** buffers,
-    const size_t* buffer_offsets)
+    size_t num_device_buffers,
+    const struct gptoss_metal_buffer** device_buffers,
+    const size_t* device_buffer_offsets,
+    size_t threadgroup_buffer_size)
 {
     if (command_buffer->object == NULL || function->pipeline_state_object == NULL) {
         return gptoss_status_invalid_state;
@@ -396,11 +397,14 @@ enum gptoss_status gptoss_metal_command_buffer_encode_launch_kernel(
     // Set kernel arguments
     [command_encoder_obj setComputePipelineState:pipeline_state_obj];
     [command_encoder_obj setBytes:params length:params_size atIndex:0];
-    for (size_t i = 0; i < num_buffers; ++i) {
-        id<MTLBuffer> buffer_obj = (id<MTLBuffer>) buffers[i]->object;
-        const NSUInteger offset = buffer_offsets == NULL ? 0 : (NSUInteger) buffer_offsets[i];
+    for (size_t i = 0; i < num_device_buffers; ++i) {
+        id<MTLBuffer> buffer_obj = (id<MTLBuffer>) device_buffers[i]->object;
+        const NSUInteger offset = device_buffer_offsets == NULL ? 0 : (NSUInteger) device_buffer_offsets[i];
         [command_encoder_obj setBuffer:buffer_obj offset:offset atIndex:i + 1];
     }
+    if (threadgroup_buffer_size != 0) {
+        [command_encoder_obj setThreadgroupMemoryLength:threadgroup_buffer_size atIndex:0];
+    }
 
     // Dispatch kernel
     const MTLSize threadgroup_size = MTLSizeMake(threadgroup_size_x, threadgroup_size_y, threadgroup_size_z);