Skip to content

Commit 3300b1e

Browse files
MollySophiaMollySophia
committed
WKV7 Metal
Signed-off-by: Molly Sophia <[email protected]>
1 parent ebc9ed3 commit 3300b1e

File tree

2 files changed

+141
-4
lines changed

2 files changed

+141
-4
lines changed

‎ggml/src/ggml-metal/ggml-metal.m

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -181,6 +181,7 @@ static void ggml_backend_metal_device_rel(struct ggml_backend_metal_device_conte
181181
GGML_METAL_KERNEL_TYPE_SSM_CONV_F32,
182182
GGML_METAL_KERNEL_TYPE_SSM_SCAN_F32,
183183
GGML_METAL_KERNEL_TYPE_RWKV_WKV6_F32,
184+
GGML_METAL_KERNEL_TYPE_RWKV_WKV7_F32,
184185
GGML_METAL_KERNEL_TYPE_MUL_MV_F32_F32,
185186
GGML_METAL_KERNEL_TYPE_MUL_MV_F16_F32,
186187
GGML_METAL_KERNEL_TYPE_MUL_MV_F16_F32_1ROW,
@@ -790,6 +791,7 @@ @implementation GGMLMetalClass
790791
GGML_METAL_ADD_KERNEL(GGML_METAL_KERNEL_TYPE_SSM_CONV_F32, ssm_conv_f32, true);
791792
GGML_METAL_ADD_KERNEL(GGML_METAL_KERNEL_TYPE_SSM_SCAN_F32, ssm_scan_f32, true);
792793
GGML_METAL_ADD_KERNEL(GGML_METAL_KERNEL_TYPE_RWKV_WKV6_F32, rwkv_wkv6_f32, true);
794+
GGML_METAL_ADD_KERNEL(GGML_METAL_KERNEL_TYPE_RWKV_WKV7_F32, rwkv_wkv7_f32, true);
793795
GGML_METAL_ADD_KERNEL(GGML_METAL_KERNEL_TYPE_MUL_MV_F32_F32, mul_mv_f32_f32, has_simdgroup_reduction);
794796
GGML_METAL_ADD_KERNEL(GGML_METAL_KERNEL_TYPE_MUL_MV_BF16_F32, mul_mv_bf16_f32, has_simdgroup_reduction && use_bfloat);
795797
GGML_METAL_ADD_KERNEL(GGML_METAL_KERNEL_TYPE_MUL_MV_BF16_F32_1ROW, mul_mv_bf16_f32_1row, has_simdgroup_reduction && use_bfloat);
@@ -1254,6 +1256,7 @@ static bool ggml_metal_supports_op(const struct ggml_backend_metal_device_contex
12541256
case GGML_OP_SSM_CONV:
12551257
case GGML_OP_SSM_SCAN:
12561258
case GGML_OP_RWKV_WKV6:
1259+
case GGML_OP_RWKV_WKV7:
12571260
return true;
12581261
case GGML_OP_MUL_MAT:
12591262
case GGML_OP_MUL_MAT_ID:
@@ -2216,6 +2219,46 @@ static void ggml_metal_encode_node(
22162219
[encoder setBytes:&C length:sizeof(C) atIndex:9];
22172220
[encoder setBytes:&H length:sizeof(H) atIndex:10];
22182221

2222+
[encoder dispatchThreadgroups:MTLSizeMake(B * H, 1, 1) threadsPerThreadgroup:MTLSizeMake(C/ H, 1, 1)];
2223+
} break;
2224+
case GGML_OP_RWKV_WKV7:
2225+
{
2226+
const int64_t B = dst->src[6]->ne[1];
2227+
const int64_t T = dst->src[0]->ne[2];
2228+
const int64_t C = dst->ne[0];
2229+
const int64_t H = dst->src[0]->ne[1];
2230+
2231+
GGML_ASSERT(dst->src[6]->type == GGML_TYPE_F32);
2232+
GGML_ASSERT(C % H == 0);
2233+
GGML_ASSERT(C / H == 64);
2234+
2235+
size_t offs_src3 = 0;
2236+
size_t offs_src4 = 0;
2237+
size_t offs_src5 = 0;
2238+
size_t offs_src6 = 0;
2239+
2240+
id<MTLBuffer> id_src3 = dst->src[3] ? ggml_metal_get_buffer(dst->src[3], &offs_src3) : nil;
2241+
id<MTLBuffer> id_src4 = dst->src[4] ? ggml_metal_get_buffer(dst->src[4], &offs_src4) : nil;
2242+
id<MTLBuffer> id_src5 = dst->src[5] ? ggml_metal_get_buffer(dst->src[5], &offs_src5) : nil;
2243+
id<MTLBuffer> id_src6 = dst->src[6] ? ggml_metal_get_buffer(dst->src[6], &offs_src6) : nil;
2244+
2245+
id<MTLComputePipelineState> pipeline = ctx->kernels[GGML_METAL_KERNEL_TYPE_RWKV_WKV7_F32].pipeline;
2246+
2247+
[encoder setComputePipelineState:pipeline];
2248+
[encoder setBuffer:id_src0 offset:offs_src0 atIndex:0];
2249+
[encoder setBuffer:id_src1 offset:offs_src1 atIndex:1];
2250+
[encoder setBuffer:id_src2 offset:offs_src2 atIndex:2];
2251+
[encoder setBuffer:id_src3 offset:offs_src3 atIndex:3];
2252+
[encoder setBuffer:id_src4 offset:offs_src4 atIndex:4];
2253+
[encoder setBuffer:id_src5 offset:offs_src5 atIndex:5];
2254+
[encoder setBuffer:id_src6 offset:offs_src6 atIndex:6];
2255+
[encoder setBuffer:id_dst offset:offs_dst atIndex:7];
2256+
2257+
[encoder setBytes:&B length:sizeof(B) atIndex:8];
2258+
[encoder setBytes:&T length:sizeof(T) atIndex:9];
2259+
[encoder setBytes:&C length:sizeof(C) atIndex:10];
2260+
[encoder setBytes:&H length:sizeof(H) atIndex:11];
2261+
22192262
[encoder dispatchThreadgroups:MTLSizeMake(B * H, 1, 1) threadsPerThreadgroup:MTLSizeMake(C/ H, 1, 1)];
22202263
} break;
22212264
case GGML_OP_MUL_MAT:

‎ggml/src/ggml-metal/ggml-metal.metal

Lines changed: 98 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1453,10 +1453,10 @@ kernel void kernel_rwkv_wkv6_f32(
14531453
y += dot(r_vec, temp);
14541454

14551455
s_vec = s_vec * td_vec + kv;
1456-
state[j] = s_vec.x;
1457-
state[j+1] = s_vec.y;
1458-
state[j+2] = s_vec.z;
1459-
state[j+3] = s_vec.w;
1456+
state[j] = s_vec[0];
1457+
state[j+1] = s_vec[1];
1458+
state[j+2] = s_vec[2];
1459+
state[j+3] = s_vec[3];
14601460
}
14611461

14621462
dst[t] = y;
@@ -1468,6 +1468,100 @@ kernel void kernel_rwkv_wkv6_f32(
14681468
}
14691469
}
14701470

1471+
kernel void kernel_rwkv_wkv7_f32(
1472+
device const float * r,
1473+
device const float * w,
1474+
device const float * k,
1475+
device const float * v,
1476+
device const float * a,
1477+
device const float * b,
1478+
device const float * state_in,
1479+
device float * dst,
1480+
constant uint & B,
1481+
constant uint & T,
1482+
constant uint & C,
1483+
constant uint & H,
1484+
uint3 tgpig[[threadgroup_position_in_grid]],
1485+
uint3 tpitg[[thread_position_in_threadgroup]],
1486+
uint3 ntg[[threads_per_threadgroup]]) {
1487+
1488+
const uint head_size = 64;
1489+
const uint batch_id = tgpig.x / H;
1490+
const uint head_id = tgpig.x % H;
1491+
const uint tid = tpitg.x;
1492+
1493+
if (batch_id >= B || head_id >= H) {
1494+
return;
1495+
}
1496+
1497+
const uint state_size = C * head_size;
1498+
const uint n_seq_tokens = T / B;
1499+
1500+
threadgroup float _r[head_size];
1501+
threadgroup float _w[head_size];
1502+
threadgroup float _k[head_size];
1503+
threadgroup float _a[head_size];
1504+
threadgroup float _b[head_size];
1505+
1506+
float state[head_size];
1507+
#pragma unroll(64)
1508+
for (uint i = 0; i < head_size; i++) {
1509+
state[i] = state_in[batch_id * state_size + head_id * head_size * head_size
1510+
+ tid * head_size + i];
1511+
}
1512+
1513+
const uint start_t = batch_id * n_seq_tokens * C + head_id * head_size + tid;
1514+
const uint end_t = (batch_id + 1) * n_seq_tokens * C + head_id * head_size + tid;
1515+
1516+
for (uint t = start_t; t < end_t; t += C) {
1517+
threadgroup_barrier(mem_flags::mem_threadgroup);
1518+
_r[tid] = r[t];
1519+
_w[tid] = w[t];
1520+
_k[tid] = k[t];
1521+
_a[tid] = a[t];
1522+
_b[tid] = b[t];
1523+
threadgroup_barrier(mem_flags::mem_threadgroup);
1524+
1525+
const float v_val = v[t];
1526+
float y = 0.0, sa = 0.0;
1527+
1528+
float4 sa_vec(0.0);
1529+
#pragma unroll(64)
1530+
for (int j = 0; j < head_size; j += 4) {
1531+
float4 a_vec = float4(_a[j], _a[j+1], _a[j+2], _a[j+3]);
1532+
float4 s_vec = float4(state[j], state[j+1], state[j+2], state[j+3]);
1533+
sa_vec += a_vec * s_vec;
1534+
}
1535+
sa = sa_vec[0] + sa_vec[1] + sa_vec[2] + sa_vec[3];
1536+
1537+
#pragma unroll(64)
1538+
for (uint j = 0; j < head_size; j += 4) {
1539+
float4 r_vec = float4(_r[j], _r[j+1], _r[j+2], _r[j+3]);
1540+
float4 w_vec = float4(_w[j], _w[j+1], _w[j+2], _w[j+3]);
1541+
float4 k_vec = float4(_k[j], _k[j+1], _k[j+2], _k[j+3]);
1542+
float4 b_vec = float4(_b[j], _b[j+1], _b[j+2], _b[j+3]);
1543+
float4 s_vec = float4(state[j], state[j+1], state[j+2], state[j+3]);
1544+
1545+
float4 kv = k_vec * v_val;
1546+
1547+
s_vec = s_vec * w_vec + kv + sa * b_vec;
1548+
y += dot(s_vec, r_vec);
1549+
1550+
state[j] = s_vec[0];
1551+
state[j+1] = s_vec[1];
1552+
state[j+2] = s_vec[2];
1553+
state[j+3] = s_vec[3];
1554+
}
1555+
1556+
dst[t] = y;
1557+
}
1558+
#pragma unroll(64)
1559+
for (uint i = 0; i < head_size; i++) {
1560+
dst[T * C + batch_id * state_size + head_id * head_size * head_size
1561+
+ tid * head_size + i] = state[i];
1562+
}
1563+
}
1564+
14711565
kernel void kernel_argmax(
14721566
device const void * x,
14731567
device int32_t * dst,

0 commit comments

Comments
 (0)