42 lines
1.2 KiB
C
42 lines
1.2 KiB
C
/* SPDX-License-Identifier: AGPL-3.0-or-later */
|
|
/* Copyright (C) 2026 David (Corey) Herrera */
|
|
/**
|
|
* @file fm_kernel_wkv7.h
|
|
* @brief WKV-7 (Generalized Delta Rule) recurrence kernels — single timestep
|
|
* and whole-sequence. One threadgroup per head, head_size threads per group.
|
|
*
|
|
* @status None
|
|
* @issues None
|
|
* @todo None
|
|
*/
|
|
|
|
#ifndef FOUNDRY_METAL_KERNEL_WKV7_H
|
|
#define FOUNDRY_METAL_KERNEL_WKV7_H
|
|
|
|
#include "metal/fm_metal.h"
|
|
|
|
#ifdef __cplusplus
|
|
extern "C" {
|
|
#endif
|
|
|
|
/* WKV-7 single timestep (Generalized Delta Rule).
|
|
* All inputs/outputs are flat float32 arrays.
|
|
* state is [n_head, head_size, head_size], modified in place. */
|
|
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);
|
|
|
|
/* WKV-7 sequence of timesteps.
|
|
* Inputs are [T, n_head * head_size] except state. */
|
|
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);
|
|
|
|
#ifdef __cplusplus
|
|
}
|
|
#endif
|
|
|
|
#endif /* FOUNDRY_METAL_KERNEL_WKV7_H */
|