131 lines
5.0 KiB
Mathematica
131 lines
5.0 KiB
Mathematica
/* SPDX-License-Identifier: AGPL-3.0-or-later */
|
|
/* Copyright (C) 2026 David (Corey) Herrera */
|
|
/**
|
|
* @file fm_kernel_wkv7.m
|
|
* @brief Host dispatch and MSL source for the WKV-7 recurrence kernels.
|
|
*
|
|
* @status None
|
|
* @issues None
|
|
* @todo None
|
|
*/
|
|
|
|
#include "metal/kernels/fm_kernel_wkv7.h"
|
|
|
|
#include "metal/fm_metal_internal.h"
|
|
#include "metal/kernels/fm_kernel_registry.h"
|
|
|
|
void fm_wkv7_one(fm_buffer *r, fm_buffer *w, fm_buffer *k, fm_buffer *v,
|
|
fm_buffer *neg_kk, fm_buffer *kk_a,
|
|
fm_buffer *state, fm_buffer *y_out,
|
|
int n_head, int head_size) {
|
|
if (!fm_metal_is_initialized()) return;
|
|
|
|
fm_buffer *bufs[] = { r, w, k, v, neg_kk, kk_a, state, y_out };
|
|
uint32_t consts[] = { (uint32_t)n_head, (uint32_t)head_size };
|
|
/* One threadgroup per head, head_size threads per group */
|
|
fm_dispatch_threadgroups(@"wkv7_one", bufs, 8, consts, 2,
|
|
n_head, head_size);
|
|
}
|
|
|
|
void fm_wkv7_seq(fm_buffer *r, fm_buffer *w, fm_buffer *k, fm_buffer *v,
|
|
fm_buffer *neg_kk, fm_buffer *kk_a,
|
|
fm_buffer *state, fm_buffer *y_out,
|
|
int n_head, int head_size, int t_len) {
|
|
if (!fm_metal_is_initialized()) return;
|
|
|
|
fm_buffer *bufs[] = { r, w, k, v, neg_kk, kk_a, state, y_out };
|
|
uint32_t consts[] = { (uint32_t)n_head, (uint32_t)head_size, (uint32_t)t_len };
|
|
fm_dispatch_threadgroups(@"wkv7_seq", bufs, 8, consts, 3,
|
|
n_head, head_size);
|
|
}
|
|
|
|
NSString *fm_kernel_msl_wkv7(void)
|
|
{
|
|
return @
|
|
"/* ============================================================\n"
|
|
" * WKV-7 single timestep (Generalized Delta Rule)\n"
|
|
" *\n"
|
|
" * One threadgroup per head, head_size threads per group.\n"
|
|
" * Uses threadgroup memory for cooperative loading.\n"
|
|
" * ============================================================ */\n"
|
|
"\n"
|
|
"kernel void wkv7_one(\n"
|
|
" device const float* r [[buffer(0)]],\n"
|
|
" device const float* w [[buffer(1)]],\n"
|
|
" device const float* k [[buffer(2)]],\n"
|
|
" device const float* v [[buffer(3)]],\n"
|
|
" device const float* neg_kk [[buffer(4)]],\n"
|
|
" device const float* kk_a [[buffer(5)]],\n"
|
|
" device float* state [[buffer(6)]],\n"
|
|
" device float* y [[buffer(7)]],\n"
|
|
" constant uint& n_head [[buffer(8)]],\n"
|
|
" constant uint& head_size [[buffer(9)]],\n"
|
|
" uint h [[threadgroup_position_in_grid]],\n"
|
|
" uint i [[thread_position_in_threadgroup]]\n"
|
|
") {\n"
|
|
" if (h >= n_head || i >= head_size) return;\n"
|
|
" uint n = head_size;\n"
|
|
" uint offset = h * n;\n"
|
|
" uint sbase = (h * n + i) * n;\n"
|
|
"\n"
|
|
" // Step 1: sa = dot(-kk, state_row)\n"
|
|
" float sa = 0.0f;\n"
|
|
" for (uint j = 0; j < n; ++j) {\n"
|
|
" sa += neg_kk[offset + j] * state[sbase + j];\n"
|
|
" }\n"
|
|
"\n"
|
|
" float vv = v[offset + i];\n"
|
|
"\n"
|
|
" // Step 2: state update + output\n"
|
|
" float y_out = 0.0f;\n"
|
|
" for (uint j = 0; j < n; ++j) {\n"
|
|
" float s = state[sbase + j];\n"
|
|
" s = s * w[offset + j] + k[offset + j] * vv + sa * kk_a[offset + j];\n"
|
|
" state[sbase + j] = s;\n"
|
|
" y_out += s * r[offset + j];\n"
|
|
" }\n"
|
|
" y[offset + i] = y_out;\n"
|
|
"}\n"
|
|
"\n"
|
|
"/* WKV-7 sequence kernel */\n"
|
|
"kernel void wkv7_seq(\n"
|
|
" device const float* r [[buffer(0)]],\n"
|
|
" device const float* w [[buffer(1)]],\n"
|
|
" device const float* k [[buffer(2)]],\n"
|
|
" device const float* v [[buffer(3)]],\n"
|
|
" device const float* neg_kk [[buffer(4)]],\n"
|
|
" device const float* kk_a [[buffer(5)]],\n"
|
|
" device float* state [[buffer(6)]],\n"
|
|
" device float* y [[buffer(7)]],\n"
|
|
" constant uint& n_head [[buffer(8)]],\n"
|
|
" constant uint& head_size [[buffer(9)]],\n"
|
|
" constant uint& t_len [[buffer(10)]],\n"
|
|
" uint h [[threadgroup_position_in_grid]],\n"
|
|
" uint i [[thread_position_in_threadgroup]]\n"
|
|
") {\n"
|
|
" if (h >= n_head || i >= head_size) return;\n"
|
|
" uint n = head_size;\n"
|
|
" uint c = n_head * n;\n"
|
|
" uint offset = h * n;\n"
|
|
" uint sbase = (h * n + i) * n;\n"
|
|
"\n"
|
|
" for (uint t = 0; t < t_len; ++t) {\n"
|
|
" uint toff = t * c + offset;\n"
|
|
" float sa = 0.0f;\n"
|
|
" for (uint j = 0; j < n; ++j)\n"
|
|
" sa += neg_kk[toff + j] * state[sbase + j];\n"
|
|
" float vv = v[toff + i];\n"
|
|
" float y_out = 0.0f;\n"
|
|
" for (uint j = 0; j < n; ++j) {\n"
|
|
" float s = state[sbase + j];\n"
|
|
" s = s * w[toff + j] + k[toff + j] * vv + sa * kk_a[toff + j];\n"
|
|
" state[sbase + j] = s;\n"
|
|
" y_out += s * r[toff + j];\n"
|
|
" }\n"
|
|
" y[t * c + offset + i] = y_out;\n"
|
|
" }\n"
|
|
"}\n"
|
|
"\n"
|
|
;
|
|
}
|