Metal: fused QKV projection (matmul+RoPE+KV cache init) kernel (#184)

Author: Maratyszcza

gpt_oss/metal/source/context.c

@@ -253,10 +253,51 @@ static enum gptoss_status process_tokens(
                     GPTOSS_LOG_ERROR("failed to encode f32_bf16w_dense_matmul_qkv kernel launch");
                     return status;
                 }
+
+                status = gptoss_metal_command_buffer_encode_launch_f32_rope(
+                    command_buffer,
+                    &model->f32_rope_fn,
+                    /*threadgroup_size=*/32,
+                    &context->qkv_activation_buffer,
+                    /*input_offset=*/0,
+                    &context->control_buffer,
+                    /*control_offset=*/0,
+                    model->rope_theta,
+                    model->interpolation_scale,
+                    model->yarn_offset,
+                    model->yarn_scale,
+                    model->yarn_multiplier,
+                    input_batch_size,
+                    model->num_heads,
+                    model->num_kv_heads,
+                    model->head_dim,
+                    /*token_offset=*/input_batch_start);
+                if (status != gptoss_status_success) {
+                    GPTOSS_LOG_ERROR("failed to encode f32_rope kernel launch");
+                    return status;
+                }
+
+                for (uint32_t t = 0; t < input_batch_size; t++) {
+                    for (uint32_t kv = 0; kv < 2; kv++) {
+                        for (uint32_t h = 0; h < model->num_kv_heads; h++) {
+                            status = gptoss_metal_command_buffer_encode_copy_buffer(
+                                command_buffer,
+                                &context->qkv_activation_buffer,
+                                /*input_offset=*/(t * attn_qkv_dim + (model->num_heads + kv * model->num_kv_heads + h) * model->head_dim) * sizeof(float),
+                                &context->kvcache_buffer,
+                                /*output_offset=*/(((n * model->num_kv_heads + h) * context->max_tokens + input_batch_start + t) * 2 + kv) * model->head_dim * sizeof(float),
+                                /*size=*/model->head_dim * sizeof(float));
+                            if (status != gptoss_status_success) {
+                                GPTOSS_LOG_ERROR("failed to encode copy of token %" PRIu32 " to KV cache", t);
+                                return status;
+                            }
+                        }
+                    }
+                }
             } else {
-                status = gptoss_metal_command_buffer_encode_launch_f32_bf16w_matmul(
+                status = gptoss_metal_command_buffer_encode_launch_f32_bf16w_matmul_qkv(
                     command_buffer,
-                    &model->f32_bf16w_matmul_fn,
+                    &model->f32_bf16w_matmul_qkv_fn,
                     model->attn_qkv_threadgroup_size,
                     &context->rmsnorm_activation_buffer,
                     /*input_offset=*/0,
@@ -266,49 +307,24 @@ static enum gptoss_status process_tokens(
                     /*bias_offset=*/model->attn_qkv_bias_offset + model->per_block_shared_weights_size * n,
                     &context->qkv_activation_buffer,
                     /*output_offset=*/0,
+                    &context->kvcache_buffer,
+                    /*kv_offset=*/n * model->num_kv_heads * context->max_tokens * 2 * model->head_dim * sizeof(float),
                     &context->control_buffer,
                     /*control_offset=*/0,
                     /*num_tokens=*/input_batch_size,
                     /*num_cols=*/model->embedding_dim,
-                    /*num_rows=*/attn_qkv_dim);
-                if (status != gptoss_status_success) {
-                    GPTOSS_LOG_ERROR("failed to encode f32_bf16w_matmul kernel launch");
-                    return status;
-                }
-            }
-            status = gptoss_metal_command_buffer_encode_launch_f32_rope(
-                command_buffer,
-                &model->f32_rope_fn,
-                /*threadgroup_size=*/32,
-                &context->qkv_activation_buffer,
-                /*input_offset=*/0,
-                &context->control_buffer,
-                /*control_offset=*/0,
-                model->rope_theta,
-                model->interpolation_scale,
-                model->yarn_offset,
-                model->yarn_scale,
-                model->yarn_multiplier,
-                input_batch_size,
-                model->num_heads,
-                model->num_kv_heads,
-                model->head_dim,
-                /*token_offset=*/input_batch_start);
-            if (status != gptoss_status_success) {
-                GPTOSS_LOG_ERROR("failed to encode f32_rope kernel launch");
-                return status;
-            }
-
-            for (uint32_t t = 0; t < input_batch_size; t++) {
-                status = gptoss_metal_command_buffer_encode_copy_buffer(
-                    command_buffer,
-                    &context->qkv_activation_buffer,
-                    /*input_offset=*/(t * attn_qkv_dim + model->num_heads * model->head_dim) * sizeof(float),
-                    &context->kvcache_buffer,
-                    /*output_offset=*/(n * context->max_tokens + input_batch_start + t) * 2 * model->num_kv_heads * model->head_dim * sizeof(float),
-                    /*size=*/2 * model->num_kv_heads * model->head_dim * sizeof(float));
+                    /*num_q_heads=*/model->num_heads,
+                    /*num_kv_heads=*/model->num_kv_heads,
+                    /*attn_head_dim=*/model->head_dim,
+                    /*token_offset=*/input_batch_start,
+                    /*max_tokens=*/context->max_tokens,
+                    /*rope_base=*/model->rope_theta,
+                    /*interpolation_scale=*/model->interpolation_scale,
+                    /*yarn_offset=*/model->yarn_offset,
+                    /*yarn_scale=*/model->yarn_scale,
+                    /*yarn_multiplier=*/model->yarn_multiplier);
                 if (status != gptoss_status_success) {
-                    GPTOSS_LOG_ERROR("failed to encode copy of token %" PRIu32 " to KV cache", t);
+                    GPTOSS_LOG_ERROR("failed to encode f32_bf16w_matmul_qkv kernel launch");
                     return status;
                 }
             }
@@ -320,16 +336,15 @@ static enum gptoss_status process_tokens(
                     &context->qkv_activation_buffer,
                     /*q_offset=*/attn_qkv_dim * (input_batch_size - num_block_output_tokens) * sizeof(float),
                     &context->kvcache_buffer,
-                    /*k_offset=*/n * context->max_tokens * 2 * model->num_kv_heads * model->head_dim * sizeof(float),
-                    &context->kvcache_buffer,
-                    /*v_offset=*/(n * context->max_tokens * 2 + 1) * model->num_kv_heads * model->head_dim * sizeof(float),
+                    /*kv_offset=*/n * model->num_kv_heads * context->max_tokens * 2 * model->head_dim * sizeof(float),
                     &model->shared_weight_buffer,
                     /*s_offset=*/model->attn_sdpa_sink_offset + model->per_block_shared_weights_size * n,
                     &context->sdpa_activation_buffer,
                     /*output_offset=*/0,
                     &context->control_buffer,
                     /*control_offset=*/0,
                     /*window=*/n % 2 == 0 ? model->attention_window : UINT32_MAX,
+                    /*kv_stride=*/2 * context->max_tokens * model->head_dim,
                     num_block_output_tokens,
                     input_batch_start + input_batch_size - num_block_output_tokens,
                     model->num_heads, model->num_kv_heads, model->head_dim);

gpt_oss/metal/source/include/internal/kernel-args.h

@@ -39,6 +39,7 @@ struct gptoss_topk_args {
 struct gptoss_sdpa_args {
     uint32_t qkv_dim;
     uint32_t num_kv_tokens;
+    uint32_t kv_stride;
     uint32_t window;
 };
 
@@ -126,6 +127,18 @@ struct gptoss_rope_args {
     float yarn_multiplier;
 };
 
+struct gptoss_qkv_args {
+    uint32_t num_column_vecs;
+    uint32_t num_rows;
+    uint32_t token_offset;
+    float freq_scale;
+    float interpolation_scale;
+    float yarn_offset;
+    float yarn_scale;
+    float yarn_multiplier;
+    uint32_t max_tokens;
+};
+
 struct gptoss_softmax_args {
     uint32_t num_vecs;
     uint32_t num_vecs_per_threadgroup;

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

@@ -112,6 +112,35 @@ enum gptoss_status gptoss_metal_command_buffer_encode_launch_f32_bf16w_matmul(
     uint32_t num_cols,
     uint32_t num_rows);
 
+enum gptoss_status gptoss_metal_command_buffer_encode_launch_f32_bf16w_matmul_qkv(
+    const struct gptoss_metal_command_buffer* command_buffer,
+    const struct gptoss_metal_function* f32_bf16w_matmul_qkv_fn,
+    size_t threadgroup_size,
+    const struct gptoss_metal_buffer* input_buffer,
+    size_t input_offset,
+    const struct gptoss_metal_buffer* weight_buffer,
+    size_t weight_offset,
+    const struct gptoss_metal_buffer* bias_buffer,
+    size_t bias_offset,
+    const struct gptoss_metal_buffer* output_buffer,
+    size_t output_offset,
+    const struct gptoss_metal_buffer* kv_buffer,
+    size_t kv_offset,
+    const struct gptoss_metal_buffer* control_buffer,
+    size_t control_offset,
+    uint32_t num_tokens,
+    uint32_t num_cols,
+    uint32_t num_q_heads,
+    uint32_t num_kv_heads,
+    uint32_t attn_head_dim,
+    uint32_t token_offset,
+    uint32_t max_tokens,
+    float rope_base,
+    float interpolation_scale,
+    float yarn_offset,
+    float yarn_scale,
+    float yarn_multiplier);
+
 enum gptoss_status gptoss_metal_command_buffer_encode_launch_f32_bf16w_matmul_add(
     const struct gptoss_metal_command_buffer* command_buffer,
     const struct gptoss_metal_function* f32_bf16w_matmul_fn,
@@ -306,17 +335,16 @@ enum gptoss_status gptoss_metal_command_buffer_encode_launch_f32_sdpa(
     const struct gptoss_metal_function* f32_sdpa_fn,
     const struct gptoss_metal_buffer* q_buffer,
     size_t q_offset,
-    const struct gptoss_metal_buffer* k_buffer,
-    size_t k_offset,
-    const struct gptoss_metal_buffer* v_buffer,
-    size_t v_offset,
+    const struct gptoss_metal_buffer* kv_buffer,
+    size_t kv_offset,
     const struct gptoss_metal_buffer* s_buffer,
     size_t s_offset,
     const struct gptoss_metal_buffer* output_buffer,
     size_t output_offset,
     const struct gptoss_metal_buffer* control_buffer,
     size_t control_offset,
     uint32_t window,
+    uint32_t kv_stride,
     uint32_t num_q_tokens,
     uint32_t num_kv_tokens,
     uint32_t num_q_heads,

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

@@ -78,6 +78,7 @@ struct gptoss_model {
     struct gptoss_metal_function bf16_f32_embeddings_fn;
     struct gptoss_metal_function f32_bf16w_rmsnorm_fn;
     struct gptoss_metal_function f32_bf16w_matmul_fn;
+    struct gptoss_metal_function f32_bf16w_matmul_qkv_fn;
     struct gptoss_metal_function f32_bf16w_dense_matmul_qkv_fn;
     struct gptoss_metal_function f32_bf16w_dense_matmul_attn_output_fn;
     struct gptoss_metal_function f32_bf16w_dense_matmul_mlp_gate_fn;

gpt_oss/metal/source/matmul.metal

@@ -67,6 +67,94 @@ kernel void gptoss_f32_bf16w_matmul(
     }
 }
 
+kernel void gptoss_f32_bf16w_matmul_qkv(
+    constant gptoss_qkv_args& args [[ buffer(0) ]],
+    const device float4* input [[ buffer(1) ]],
+    const device bfloat4* weight [[ buffer(2) ]],
+    const device bfloat* bias [[ buffer(3) ]],
+    device float* q [[ buffer(4) ]],
+    device float* kv [[ buffer(5) ]],
+    const device gptoss_control* control [[ buffer(6) ]],
+    threadgroup void* scratch [[ threadgroup(0) ]],
+    uint2 gid [[threadgroup_position_in_grid]],
+    uint simdgroup_tid [[thread_index_in_simdgroup]],
+    uint simdgroup_idx [[simdgroup_index_in_threadgroup]],
+    uint num_simdgroups [[simdgroups_per_threadgroup]])
+{
+    const uint simdgroup_size = 32;
+    const uint head_dim = 64;
+    const uint num_q_heads = 64;
+    const uint num_kv_heads = 8;
+    if (control->abort != 0) {
+        return;
+    }
+
+    const uint num_column_vecs = args.num_column_vecs;
+    const uint row = gid.x * num_simdgroups + simdgroup_idx;
+
+    input += gid.y * num_column_vecs + simdgroup_tid;
+    weight += num_column_vecs * row + simdgroup_tid;
+    bias += row;
+    q += gid.y * args.num_rows;
+
+    uint num_iter = (num_column_vecs - simdgroup_tid + (simdgroup_size - 1)) / simdgroup_size;
+
+    float4 sum4 = 0.0f;
+    do {
+        const bfloat4 w = *weight;
+        const float4 i = *input;
+        sum4 = metal::fma(static_cast<float4>(w), i, sum4);
+
+        weight += simdgroup_size;
+        input += simdgroup_size;
+    } while (--num_iter != 0);
+    const float2 sum2 = sum4.xy + sum4.zw;
+    float sum = sum2.x + sum2.y;
+    sum = metal::simd_sum(sum);
+    if (metal::simd_is_first()) {
+        sum += static_cast<float>(*bias);
+        static_cast<threadgroup float*>(scratch)[simdgroup_idx] = sum;
+    }
+    metal::threadgroup_barrier(metal::mem_flags::mem_threadgroup);
+    if (simdgroup_idx == 0) {
+        const uint num_half_simdgroups = num_simdgroups / 2;
+        if (simdgroup_tid < num_half_simdgroups) {
+            float2 vals = static_cast<const threadgroup float2*>(scratch)[simdgroup_tid];
+            const uint idx = gid.x * num_half_simdgroups + simdgroup_tid;
+            const uint head_idx = idx / (head_dim / 2);
+            const uint token_idx = args.token_offset + gid.y;
+            const uint dim_idx = idx % (head_dim / 2);
+            if (head_idx < num_q_heads + num_kv_heads) {
+                const float dim_idx_val = static_cast<float>(dim_idx);
+                const float inv_extrapolation_freq = metal::precise::exp(dim_idx_val * args.freq_scale);
+                const float inv_interpolation_freq = inv_extrapolation_freq * args.interpolation_scale;
+                const float alpha = metal::saturate(metal::fma(dim_idx_val, args.yarn_scale, args.yarn_offset));
+                const float inv_freq = metal::mix(inv_extrapolation_freq, inv_interpolation_freq, alpha);
+
+                const float phi = static_cast<float>(token_idx) * inv_freq;
+                const float yarn_multiplier = args.yarn_multiplier;
+                float cosphi;
+                const float sinphi = metal::precise::sincos(phi, cosphi) * yarn_multiplier;
+                cosphi *= yarn_multiplier;
+
+                vals = (float2) {
+                    vals.x * cosphi - vals.y * sinphi,
+                    vals.x * sinphi + vals.y * cosphi,
+                };
+            }
+            if (head_idx < num_q_heads) {
+                reinterpret_cast<device float2*>(q)[idx] = vals;
+            } else if (head_idx < num_q_heads + num_kv_heads) {
+                const uint h = head_idx - num_q_heads;
+                reinterpret_cast<device float2*>(kv + (h * args.max_tokens + token_idx) * 2 * head_dim)[dim_idx] = vals;
+            } else {
+                const uint h = head_idx - num_q_heads - num_kv_heads;
+                reinterpret_cast<device float2*>(kv + (h * args.max_tokens + token_idx) * 2 * head_dim + head_dim)[dim_idx] = vals;
+            }
+        }
+    }
+}
+
 kernel void gptoss_f32_bf16w_unembedding(
     constant gptoss_unembedding_args& args [[ buffer(0) ]],
     const device float4* input [[ buffer(1) ]],