wehub-resource-sync
56 lines
1.9 KiB
Metal
56 lines
1.9 KiB
Metal
// DS4 Metal set-rows kernel used for KV writes.
|
|
|
|
struct ds4_metal_args_set_rows {
|
|
int32_t nk0;
|
|
int32_t ne01;
|
|
uint64_t nb01;
|
|
uint64_t nb02;
|
|
uint64_t nb03;
|
|
int32_t ne11;
|
|
int32_t ne12;
|
|
uint64_t nb10;
|
|
uint64_t nb11;
|
|
uint64_t nb12;
|
|
uint64_t nb1;
|
|
uint64_t nb2;
|
|
uint64_t nb3;
|
|
};
|
|
|
|
// Scatters rows into the KV cache by token position. DS4 uses this after Q/K/V
|
|
// preparation so decode and later prefill chunks can attend to previous tokens.
|
|
template<typename T, typename TI>
|
|
kernel void kernel_set_rows_f(
|
|
constant ds4_metal_args_set_rows & args,
|
|
device const char * src0,
|
|
device const char * src1,
|
|
device float * dst,
|
|
uint3 tgpig[[threadgroup_position_in_grid]],
|
|
uint tiitg[[thread_index_in_threadgroup]],
|
|
uint3 tptg [[threads_per_threadgroup]]) {
|
|
const int32_t i03 = tgpig.z;
|
|
const int32_t i02 = tgpig.y;
|
|
|
|
const int32_t i12 = i03%args.ne12;
|
|
const int32_t i11 = i02%args.ne11;
|
|
|
|
const int32_t i01 = tgpig.x*tptg.y + tiitg/tptg.x;
|
|
if (i01 >= args.ne01) {
|
|
return;
|
|
}
|
|
|
|
const int32_t i10 = i01;
|
|
const TI i1 = ((const device TI *) (src1 + i10*args.nb10 + i11*args.nb11 + i12*args.nb12))[0];
|
|
|
|
device T * dst_row = ( device T *) ((device char *) dst + i1*args.nb1 + i02*args.nb2 + i03*args.nb3);
|
|
const device float * src_row = (const device float *) ( src0 + i01*args.nb01 + i02*args.nb02 + i03*args.nb03);
|
|
|
|
for (int ind = tiitg%tptg.x; ind < args.nk0; ind += tptg.x) {
|
|
dst_row[ind] = (T) src_row[ind];
|
|
}
|
|
}
|
|
|
|
typedef decltype(kernel_set_rows_f<float, int64_t>) set_rows_f_t;
|
|
|
|
// Host-visible F32/I32 scatter variant used by KV-cache writes.
|
|
template [[host_name("kernel_set_rows_f32_i32")]] kernel set_rows_f_t kernel_set_rows_f<float, int32_t>;
|