feat: first cut at actual token generation and model loading
This commit is contained in:
Georg Bauer
275
metal/argsort.metal
Normal file
275
metal/argsort.metal
Normal file
@@ -0,0 +1,275 @@
|
||||
struct ds4_metal_args_argsort {
|
||||
int32_t ne00;
|
||||
int32_t ne01;
|
||||
int32_t ne02;
|
||||
int32_t ne03;
|
||||
uint64_t nb00;
|
||||
uint64_t nb01;
|
||||
uint64_t nb02;
|
||||
uint64_t nb03;
|
||||
int32_t ne0;
|
||||
int32_t ne1;
|
||||
int32_t ne2;
|
||||
int32_t ne3;
|
||||
int32_t top_k;
|
||||
};
|
||||
|
||||
struct ds4_metal_args_argsort_merge {
|
||||
int64_t ne00;
|
||||
int64_t ne01;
|
||||
int64_t ne02;
|
||||
int64_t ne03;
|
||||
uint64_t nb00;
|
||||
uint64_t nb01;
|
||||
uint64_t nb02;
|
||||
uint64_t nb03;
|
||||
int32_t ne0;
|
||||
int32_t ne1;
|
||||
int32_t ne2;
|
||||
int32_t ne3;
|
||||
int32_t top_k;
|
||||
int32_t len;
|
||||
};
|
||||
|
||||
typedef void (argsort_t)(
|
||||
constant ds4_metal_args_argsort & args,
|
||||
device const char * src0,
|
||||
device int32_t * dst,
|
||||
threadgroup int32_t * shmem_i32 [[threadgroup(0)]],
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
ushort3 tpitg[[thread_position_in_threadgroup]],
|
||||
ushort3 ntg[[threads_per_threadgroup]]);
|
||||
|
||||
// Sort one float row into an index row. DS4 only exports the descending
|
||||
// instance because router and indexer selection both need top-k order.
|
||||
template<ds4_sort_order order>
|
||||
kernel void kernel_argsort_f32_i32(
|
||||
constant ds4_metal_args_argsort & args,
|
||||
device const char * src0,
|
||||
device int32_t * dst,
|
||||
threadgroup int32_t * shmem_i32 [[threadgroup(0)]],
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
ushort3 tpitg[[thread_position_in_threadgroup]],
|
||||
ushort3 ntg[[threads_per_threadgroup]]) {
|
||||
// bitonic sort
|
||||
const int col = tpitg[0];
|
||||
const int ib = tgpig[0] / args.ne01;
|
||||
|
||||
const int i00 = ib*ntg.x;
|
||||
const int i01 = tgpig[0] % args.ne01;
|
||||
const int i02 = tgpig[1];
|
||||
const int i03 = tgpig[2];
|
||||
|
||||
device const float * src0_row = (device const float *) (src0 + args.nb01*i01 + args.nb02*i02 + args.nb03*i03);
|
||||
|
||||
// initialize indices
|
||||
shmem_i32[col] = i00 + col;
|
||||
|
||||
// Stage this block's score slice in threadgroup memory (indices stay in
|
||||
// [i00, i00+ntg.x), so shmem_f32[idx - i00] replaces the device gather).
|
||||
// The host allocates ntg.x extra floats after the index array. Values and
|
||||
// the comparison network are unchanged, so the permutation is identical.
|
||||
threadgroup float * shmem_f32 = (threadgroup float *) (shmem_i32 + ntg.x);
|
||||
if (i00 + col < args.ne00) {
|
||||
shmem_f32[col] = src0_row[i00 + col];
|
||||
}
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
for (int k = 2; k <= ntg.x; k *= 2) {
|
||||
for (int j = k / 2; j > 0; j /= 2) {
|
||||
int ixj = col ^ j;
|
||||
if (ixj > col) {
|
||||
if ((col & k) == 0) {
|
||||
if (shmem_i32[col] >= args.ne00 ||
|
||||
(shmem_i32[ixj] < args.ne00 && (order == DS4_SORT_ORDER_ASC ?
|
||||
shmem_f32[shmem_i32[col] - i00] > shmem_f32[shmem_i32[ixj] - i00] :
|
||||
shmem_f32[shmem_i32[col] - i00] < shmem_f32[shmem_i32[ixj] - i00]))
|
||||
) {
|
||||
SWAP(shmem_i32[col], shmem_i32[ixj]);
|
||||
}
|
||||
} else {
|
||||
if (shmem_i32[ixj] >= args.ne00 ||
|
||||
(shmem_i32[col] < args.ne00 && (order == DS4_SORT_ORDER_ASC ?
|
||||
shmem_f32[shmem_i32[col] - i00] < shmem_f32[shmem_i32[ixj] - i00] :
|
||||
shmem_f32[shmem_i32[col] - i00] > shmem_f32[shmem_i32[ixj] - i00]))
|
||||
) {
|
||||
SWAP(shmem_i32[col], shmem_i32[ixj]);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
}
|
||||
}
|
||||
|
||||
const int64_t i0 = ib*args.top_k;
|
||||
|
||||
// copy the result to dst without the padding
|
||||
if (i0 + col < args.ne0 && col < args.top_k) {
|
||||
dst += i0 + args.ne0*i01 + args.ne0*args.ne1*i02 + args.ne0*args.ne1*args.ne2*i03;
|
||||
|
||||
dst[col] = shmem_i32[col];
|
||||
}
|
||||
}
|
||||
|
||||
// Host-visible sort variant used by DS4 top-k selection.
|
||||
template [[host_name("kernel_argsort_f32_i32_desc")]] kernel argsort_t kernel_argsort_f32_i32<DS4_SORT_ORDER_DESC>;
|
||||
|
||||
typedef void (argsort_merge_t)(
|
||||
constant ds4_metal_args_argsort_merge & args,
|
||||
device const char * src0,
|
||||
device const int32_t * tmp,
|
||||
device int32_t * dst,
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
ushort3 tpitg[[thread_position_in_threadgroup]],
|
||||
ushort3 ntg[[threads_per_threadgroup]]);
|
||||
|
||||
// Merges sorted index runs produced by kernel_argsort_f32_i32. In the DS4 graph
|
||||
// this finishes top-k over router or compressed-attention score rows.
|
||||
template<ds4_sort_order order>
|
||||
kernel void kernel_argsort_merge_f32_i32(
|
||||
constant ds4_metal_args_argsort_merge & args,
|
||||
device const char * src0,
|
||||
device const int32_t * tmp,
|
||||
device int32_t * dst,
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
ushort3 tpitg[[thread_position_in_threadgroup]],
|
||||
ushort3 ntg[[threads_per_threadgroup]]) {
|
||||
|
||||
const int im = tgpig[0] / args.ne01;
|
||||
const int i01 = tgpig[0] % args.ne01;
|
||||
const int i02 = tgpig[1];
|
||||
const int i03 = tgpig[2];
|
||||
|
||||
const int start = im * (2 * args.len);
|
||||
|
||||
const int len0 = MIN(args.len, MAX(0, args.ne0 - (int)(start)));
|
||||
const int len1 = MIN(args.len, MAX(0, args.ne0 - (int)(start + args.len)));
|
||||
|
||||
const int total = len0 + len1;
|
||||
|
||||
device const int32_t * tmp0 = tmp + start
|
||||
+ i01*args.ne0
|
||||
+ i02*args.ne0*args.ne01
|
||||
+ i03*args.ne0*args.ne01*args.ne02;
|
||||
|
||||
device const int32_t * tmp1 = tmp0 + args.len;
|
||||
|
||||
dst += start
|
||||
+ i01*args.top_k
|
||||
+ i02*args.top_k*args.ne01
|
||||
+ i03*args.top_k*args.ne01*args.ne02;
|
||||
|
||||
device const float * src0_row = (device const float *)(src0
|
||||
+ args.nb01*i01
|
||||
+ args.nb02*i02
|
||||
+ args.nb03*i03);
|
||||
|
||||
if (total == 0) {
|
||||
return;
|
||||
}
|
||||
|
||||
const int chunk = (total + ntg.x - 1) / ntg.x;
|
||||
|
||||
const int k0 = tpitg.x * chunk;
|
||||
const int k1 = MIN(MIN(k0 + chunk, total), args.top_k);
|
||||
|
||||
if (k0 >= args.top_k) {
|
||||
return;
|
||||
}
|
||||
|
||||
if (k0 >= total) {
|
||||
return;
|
||||
}
|
||||
|
||||
int low = k0 > len1 ? k0 - len1 : 0;
|
||||
int high = MIN(k0, len0);
|
||||
|
||||
// binary-search partition (i, j) such that i + j = k
|
||||
while (low < high) {
|
||||
const int mid = (low + high) >> 1;
|
||||
|
||||
const int32_t idx0 = tmp0[mid];
|
||||
const int32_t idx1 = tmp1[k0 - mid - 1];
|
||||
|
||||
const float val0 = src0_row[idx0];
|
||||
const float val1 = src0_row[idx1];
|
||||
|
||||
bool take_left;
|
||||
if (order == DS4_SORT_ORDER_ASC) {
|
||||
take_left = (val0 <= val1);
|
||||
} else {
|
||||
take_left = (val0 >= val1);
|
||||
}
|
||||
|
||||
if (take_left) {
|
||||
low = mid + 1;
|
||||
} else {
|
||||
high = mid;
|
||||
}
|
||||
}
|
||||
|
||||
int i = low;
|
||||
int j = k0 - i;
|
||||
|
||||
// keep the merge fronts into registers
|
||||
int32_t idx0 = 0;
|
||||
float val0 = 0.0f;
|
||||
if (i < len0) {
|
||||
idx0 = tmp0[i];
|
||||
val0 = src0_row[idx0];
|
||||
}
|
||||
|
||||
int32_t idx1 = 0;
|
||||
float val1 = 0.0f;
|
||||
if (j < len1) {
|
||||
idx1 = tmp1[j];
|
||||
val1 = src0_row[idx1];
|
||||
}
|
||||
|
||||
for (int k = k0; k < k1; ++k) {
|
||||
int32_t out_idx;
|
||||
|
||||
if (i >= len0) {
|
||||
while (k < k1) {
|
||||
dst[k++] = tmp1[j++];
|
||||
}
|
||||
break;
|
||||
} else if (j >= len1) {
|
||||
while (k < k1) {
|
||||
dst[k++] = tmp0[i++];
|
||||
}
|
||||
break;
|
||||
} else {
|
||||
bool take_left;
|
||||
|
||||
if (order == DS4_SORT_ORDER_ASC) {
|
||||
take_left = (val0 <= val1);
|
||||
} else {
|
||||
take_left = (val0 >= val1);
|
||||
}
|
||||
|
||||
if (take_left) {
|
||||
out_idx = idx0;
|
||||
++i;
|
||||
if (i < len0) {
|
||||
idx0 = tmp0[i];
|
||||
val0 = src0_row[idx0];
|
||||
}
|
||||
} else {
|
||||
out_idx = idx1;
|
||||
++j;
|
||||
if (j < len1) {
|
||||
idx1 = tmp1[j];
|
||||
val1 = src0_row[idx1];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
dst[k] = out_idx;
|
||||
}
|
||||
}
|
||||
|
||||
// Host-visible merge variant used by DS4 top-k selection.
|
||||
template [[host_name("kernel_argsort_merge_f32_i32_desc")]] kernel argsort_merge_t kernel_argsort_merge_f32_i32<DS4_SORT_ORDER_DESC>;
|
||||
217
metal/bin.metal
Normal file
217
metal/bin.metal
Normal file
@@ -0,0 +1,217 @@
|
||||
struct ds4_metal_args_bin {
|
||||
int32_t ne00;
|
||||
int32_t ne01;
|
||||
int32_t ne02;
|
||||
int32_t ne03;
|
||||
uint64_t nb00;
|
||||
uint64_t nb01;
|
||||
uint64_t nb02;
|
||||
uint64_t nb03;
|
||||
int32_t ne10;
|
||||
int32_t ne11;
|
||||
int32_t ne12;
|
||||
int32_t ne13;
|
||||
uint64_t nb10;
|
||||
uint64_t nb11;
|
||||
uint64_t nb12;
|
||||
uint64_t nb13;
|
||||
int32_t ne0;
|
||||
int32_t ne1;
|
||||
int32_t ne2;
|
||||
int32_t ne3;
|
||||
uint64_t nb0;
|
||||
uint64_t nb1;
|
||||
uint64_t nb2;
|
||||
uint64_t nb3;
|
||||
uint64_t offs;
|
||||
uint64_t o1[8];
|
||||
};
|
||||
|
||||
struct ds4_metal_args_add3 {
|
||||
uint32_t n;
|
||||
};
|
||||
|
||||
constant short FC_bin_op [[function_constant(FC_BIN + 0)]];
|
||||
constant short FC_bin_f [[function_constant(FC_BIN + 1)]];
|
||||
constant bool FC_bin_rb [[function_constant(FC_BIN + 2)]];
|
||||
constant bool FC_bin_cb [[function_constant(FC_BIN + 3)]];
|
||||
|
||||
// Generic binary elementwise op with compile-time operation and broadcast
|
||||
// modes. DS4 currently instantiates this as add, multiply, scalar multiply, and
|
||||
// row division in the static graph.
|
||||
template <typename T0, typename T1, typename T>
|
||||
kernel void kernel_bin_fuse_impl(
|
||||
constant ds4_metal_args_bin & args,
|
||||
device const char * src0,
|
||||
device const char * src1,
|
||||
device char * dst,
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
ushort3 tpitg[[thread_position_in_threadgroup]],
|
||||
ushort3 ntg[[threads_per_threadgroup]]) {
|
||||
#define FC_OP FC_bin_op
|
||||
#define FC_F FC_bin_f
|
||||
#define FC_RB FC_bin_rb
|
||||
#define FC_CB FC_bin_cb
|
||||
|
||||
if (FC_RB) {
|
||||
const uint i0 = tgpig.y*args.ne00 + tgpig.x;
|
||||
const uint i1 = FC_CB ? tgpig.x%args.ne10 : tgpig.x;
|
||||
|
||||
device const T0 * src0_row = (device const T0 *) (src0);
|
||||
device T * dst_row = (device T *) (dst);
|
||||
|
||||
if (FC_F == 1) {
|
||||
device const T1 * src1_row = (device const T1 *) (src1 + args.o1[0]);
|
||||
|
||||
if (FC_OP == 0) {
|
||||
dst_row[i0] = src0_row[i0] + src1_row[i1];
|
||||
}
|
||||
|
||||
if (FC_OP == 1) {
|
||||
dst_row[i0] = src0_row[i0] - src1_row[i1];
|
||||
}
|
||||
|
||||
if (FC_OP == 2) {
|
||||
dst_row[i0] = src0_row[i0] * src1_row[i1];
|
||||
}
|
||||
|
||||
if (FC_OP == 3) {
|
||||
dst_row[i0] = src0_row[i0] / src1_row[i1];
|
||||
}
|
||||
} else {
|
||||
T0 res = src0_row[i0];
|
||||
|
||||
if (FC_OP == 0) {
|
||||
FOR_UNROLL (short j = 0; j < FC_F; ++j) {
|
||||
res += ((device const T1 *) (src1 + args.o1[j]))[i1];
|
||||
}
|
||||
}
|
||||
|
||||
if (FC_OP == 1) {
|
||||
FOR_UNROLL (short j = 0; j < FC_F; ++j) {
|
||||
res -= ((device const T1 *) (src1 + args.o1[j]))[i1];
|
||||
}
|
||||
}
|
||||
|
||||
if (FC_OP == 2) {
|
||||
FOR_UNROLL (short j = 0; j < FC_F; ++j) {
|
||||
res *= ((device const T1 *) (src1 + args.o1[j]))[i1];
|
||||
}
|
||||
}
|
||||
|
||||
if (FC_OP == 3) {
|
||||
FOR_UNROLL (short j = 0; j < FC_F; ++j) {
|
||||
res /= ((device const T1 *) (src1 + args.o1[j]))[i1];
|
||||
}
|
||||
}
|
||||
|
||||
dst_row[i0] = res;
|
||||
}
|
||||
} else {
|
||||
const int i03 = tgpig.z;
|
||||
const int i02 = tgpig.y;
|
||||
const int i01 = tgpig.x;
|
||||
|
||||
if (i01 >= args.ne01) {
|
||||
return;
|
||||
}
|
||||
|
||||
const int i13 = i03%args.ne13;
|
||||
const int i12 = i02%args.ne12;
|
||||
const int i11 = i01%args.ne11;
|
||||
|
||||
device const T0 * src0_ptr = (device const T0 *) (src0 + i03*args.nb03 + i02*args.nb02 + i01*args.nb01 + args.offs);
|
||||
device T * dst_ptr = (device T *) (dst + i03*args.nb3 + i02*args.nb2 + i01*args.nb1 + args.offs);
|
||||
|
||||
if (FC_F == 1) {
|
||||
device const T1 * src1_ptr = (device const T1 *) (src1 + args.o1[0] + i13*args.nb13 + i12*args.nb12 + i11*args.nb11);
|
||||
|
||||
for (int i0 = tpitg.x; i0 < args.ne0; i0 += ntg.x) {
|
||||
const int i10 = FC_CB ? i0%args.ne10 : i0;
|
||||
|
||||
if (FC_OP == 0) {
|
||||
dst_ptr[i0] = src0_ptr[i0] + src1_ptr[i10];
|
||||
}
|
||||
|
||||
if (FC_OP == 1) {
|
||||
dst_ptr[i0] = src0_ptr[i0] - src1_ptr[i10];
|
||||
}
|
||||
|
||||
if (FC_OP == 2) {
|
||||
dst_ptr[i0] = src0_ptr[i0] * src1_ptr[i10];
|
||||
}
|
||||
|
||||
if (FC_OP == 3) {
|
||||
dst_ptr[i0] = src0_ptr[i0] / src1_ptr[i10];
|
||||
}
|
||||
}
|
||||
} else {
|
||||
device const T1 * src1_ptr[8];
|
||||
FOR_UNROLL (short j = 0; j < FC_F; ++j) {
|
||||
src1_ptr[j] = (device const T1 *) (src1 + args.o1[j] + i13*args.nb13 + i12*args.nb12 + i11*args.nb11);
|
||||
}
|
||||
|
||||
for (int i0 = tpitg.x; i0 < args.ne0; i0 += ntg.x) {
|
||||
const int i10 = FC_CB ? i0%args.ne10 : i0;
|
||||
|
||||
T res = src0_ptr[i0];
|
||||
|
||||
if (FC_OP == 0) {
|
||||
FOR_UNROLL (short j = 0; j < FC_F; ++j) {
|
||||
res += src1_ptr[j][i10];
|
||||
}
|
||||
}
|
||||
|
||||
if (FC_OP == 1) {
|
||||
FOR_UNROLL (short j = 0; j < FC_F; ++j) {
|
||||
res -= src1_ptr[j][i10];
|
||||
}
|
||||
}
|
||||
|
||||
if (FC_OP == 2) {
|
||||
FOR_UNROLL (short j = 0; j < FC_F; ++j) {
|
||||
res *= src1_ptr[j][i10];
|
||||
}
|
||||
}
|
||||
|
||||
if (FC_OP == 3) {
|
||||
FOR_UNROLL (short j = 0; j < FC_F; ++j) {
|
||||
res /= src1_ptr[j][i10];
|
||||
}
|
||||
}
|
||||
|
||||
dst_ptr[i0] = res;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#undef FC_OP
|
||||
#undef FC_F
|
||||
#undef FC_RB
|
||||
#undef FC_CB
|
||||
}
|
||||
|
||||
typedef decltype(kernel_bin_fuse_impl<float, float, float>) kernel_bin_fuse_t;
|
||||
// Host-visible F32 binary op; function constants specialize it per use site.
|
||||
template [[host_name("kernel_bin_fuse_f32_f32_f32")]] kernel kernel_bin_fuse_t kernel_bin_fuse_impl<float, float, float>;
|
||||
|
||||
kernel void kernel_add2_f32(
|
||||
constant ds4_metal_args_add3 &args,
|
||||
device const float *a,
|
||||
device const float *b,
|
||||
device float *out,
|
||||
uint i [[thread_position_in_grid]]) {
|
||||
if (i >= args.n) return;
|
||||
out[i] = a[i] + b[i];
|
||||
}
|
||||
|
||||
kernel void kernel_add3_f32(
|
||||
constant ds4_metal_args_add3 &args,
|
||||
device const float *a,
|
||||
device const float *b,
|
||||
device const float *c,
|
||||
device float *out,
|
||||
uint i [[thread_position_in_grid]]) {
|
||||
if (i >= args.n) return;
|
||||
out[i] = a[i] + b[i] + c[i];
|
||||
}
|
||||
62
metal/concat.metal
Normal file
62
metal/concat.metal
Normal file
@@ -0,0 +1,62 @@
|
||||
// DS4 Metal concat kernel used by the graph.
|
||||
|
||||
struct ds4_metal_args_concat {
|
||||
int32_t ne00;
|
||||
int32_t ne01;
|
||||
int32_t ne02;
|
||||
int32_t ne03;
|
||||
uint64_t nb00;
|
||||
uint64_t nb01;
|
||||
uint64_t nb02;
|
||||
uint64_t nb03;
|
||||
int32_t ne10;
|
||||
int32_t ne11;
|
||||
int32_t ne12;
|
||||
int32_t ne13;
|
||||
uint64_t nb10;
|
||||
uint64_t nb11;
|
||||
uint64_t nb12;
|
||||
uint64_t nb13;
|
||||
int32_t ne0;
|
||||
int32_t ne1;
|
||||
int32_t ne2;
|
||||
int32_t ne3;
|
||||
uint64_t nb0;
|
||||
uint64_t nb1;
|
||||
uint64_t nb2;
|
||||
uint64_t nb3;
|
||||
int32_t dim;
|
||||
};
|
||||
|
||||
// Concatenates two float tensors along one dimension. In DS4 this is a graph
|
||||
// utility for assembling attention inputs with exactly the same tensor layout
|
||||
// expected by the downstream kernels.
|
||||
kernel void kernel_concat(
|
||||
constant ds4_metal_args_concat & args,
|
||||
device const char * src0,
|
||||
device const char * src1,
|
||||
device char * dst,
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
ushort3 tpitg[[thread_position_in_threadgroup]],
|
||||
ushort3 ntg[[threads_per_threadgroup]]) {
|
||||
const int i3 = tgpig.z;
|
||||
const int i2 = tgpig.y;
|
||||
const int i1 = tgpig.x;
|
||||
|
||||
int o[4] = {0, 0, 0, 0};
|
||||
o[args.dim] = args.dim == 0 ? args.ne00 : (args.dim == 1 ? args.ne01 : (args.dim == 2 ? args.ne02 : args.ne03));
|
||||
|
||||
device const float * x;
|
||||
|
||||
for (int i0 = tpitg.x; i0 < args.ne0; i0 += ntg.x) {
|
||||
if (i0 < args.ne00 && i1 < args.ne01 && i2 < args.ne02 && i3 < args.ne03) {
|
||||
x = (device const float *)(src0 + (i3 )*args.nb03 + (i2 )*args.nb02 + (i1 )*args.nb01 + (i0 )*args.nb00);
|
||||
} else {
|
||||
x = (device const float *)(src1 + (i3 - o[3])*args.nb13 + (i2 - o[2])*args.nb12 + (i1 - o[1])*args.nb11 + (i0 - o[0])*args.nb10);
|
||||
}
|
||||
|
||||
device float * y = (device float *)(dst + i3*args.nb3 + i2*args.nb2 + i1*args.nb1 + i0*args.nb0);
|
||||
|
||||
*y = *x;
|
||||
}
|
||||
}
|
||||
247
metal/cpy.metal
Normal file
247
metal/cpy.metal
Normal file
@@ -0,0 +1,247 @@
|
||||
struct ds4_metal_args_cpy {
|
||||
int64_t nk0;
|
||||
int64_t ne00;
|
||||
int64_t ne01;
|
||||
int64_t ne02;
|
||||
int64_t ne03;
|
||||
uint64_t nb00;
|
||||
uint64_t nb01;
|
||||
uint64_t nb02;
|
||||
uint64_t nb03;
|
||||
int64_t ne0;
|
||||
int64_t ne1;
|
||||
int64_t ne2;
|
||||
int64_t ne3;
|
||||
uint64_t nb0;
|
||||
uint64_t nb1;
|
||||
uint64_t nb2;
|
||||
uint64_t nb3;
|
||||
};
|
||||
|
||||
// Typed copy/conversion between graph tensors. DS4 uses this for layout
|
||||
// materialization and F32/F16 conversions at graph boundaries such as KV/cache
|
||||
// packing and compressor pooling.
|
||||
template<typename T0, typename T1>
|
||||
kernel void kernel_cpy_t_t(
|
||||
constant ds4_metal_args_cpy & args,
|
||||
device const char * src0,
|
||||
device char * dst,
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
ushort tiitg[[thread_index_in_threadgroup]],
|
||||
ushort3 ntg[[threads_per_threadgroup]]) {
|
||||
const int i03 = tgpig[2];
|
||||
const int i02 = tgpig[1];
|
||||
const int i01 = ntg[1] == 1 ? tgpig[0]%args.ne01 : tgpig[0]*ntg[1] + tiitg/ntg[0];
|
||||
const int iw0 = ntg[1] == 1 ? tgpig[0]/args.ne01 : 0;
|
||||
|
||||
const int64_t n = i03*args.ne02*args.ne01*args.ne00 + i02*args.ne01*args.ne00 + i01*args.ne00;
|
||||
|
||||
const int64_t i3 = n/(args.ne2*args.ne1*args.ne0);
|
||||
const int64_t i2 = (n - i3*args.ne2*args.ne1*args.ne0)/(args.ne1*args.ne0);
|
||||
const int64_t i1 = (n - i3*args.ne2*args.ne1*args.ne0 - i2*args.ne1*args.ne0)/args.ne0;
|
||||
const int64_t i0 = (n - i3*args.ne2*args.ne1*args.ne0 - i2*args.ne1*args.ne0 - i1*args.ne0);
|
||||
|
||||
device T1 * dst_data = (device T1 *) (dst + i3*args.nb3 + i2*args.nb2 + i1*args.nb1 + i0*args.nb0);
|
||||
|
||||
for (int64_t i00 = iw0*ntg[0] + tiitg%ntg[0]; i00 < args.ne00; ) {
|
||||
device const T0 * src = (device T0 *)(src0 + i03*args.nb03 + i02*args.nb02 + i01*args.nb01 + i00*args.nb00);
|
||||
dst_data[i00] = (T1) src[0];
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
typedef decltype(kernel_cpy_t_t<float, float>) kernel_cpy_t;
|
||||
// Host-visible copy/conversion variants used by the DS4 graph.
|
||||
template [[host_name("kernel_cpy_f32_f32")]] kernel kernel_cpy_t kernel_cpy_t_t<float, float>;
|
||||
template [[host_name("kernel_cpy_f32_f16")]] kernel kernel_cpy_t kernel_cpy_t_t<float, half>;
|
||||
template [[host_name("kernel_cpy_f16_f32")]] kernel kernel_cpy_t kernel_cpy_t_t<half, float>;
|
||||
template [[host_name("kernel_cpy_f16_f16")]] kernel kernel_cpy_t kernel_cpy_t_t<half, half>;
|
||||
|
||||
// Contiguous 1D conversions avoid the generic tensor-index reconstruction
|
||||
// above. Packed vector types retain scalar alignment, so tensor views whose
|
||||
// offsets are float/half aligned do not need additional 16/8-byte alignment.
|
||||
// The final vector is converted element-by-element when n is not divisible by
|
||||
// four, preserving the generic kernel's bounds and conversion semantics.
|
||||
kernel void kernel_cpy_contig_f32_f16_4(
|
||||
constant uint & n,
|
||||
device const packed_float4 * src,
|
||||
device packed_half4 * dst,
|
||||
uint gid [[thread_position_in_grid]]) {
|
||||
const uint i = gid * 4u;
|
||||
if (i >= n) {
|
||||
return;
|
||||
}
|
||||
|
||||
const uint remaining = n - i;
|
||||
if (remaining >= 4u) {
|
||||
const float4 value = float4(src[gid]);
|
||||
dst[gid] = packed_half4(half4(value));
|
||||
return;
|
||||
}
|
||||
|
||||
device const float * src_scalar = (device const float *)src;
|
||||
device half * dst_scalar = (device half *)dst;
|
||||
for (uint lane = 0; lane < remaining; ++lane) {
|
||||
dst_scalar[i + lane] = half(src_scalar[i + lane]);
|
||||
}
|
||||
}
|
||||
|
||||
kernel void kernel_cpy_contig_f16_f32_4(
|
||||
constant uint & n,
|
||||
device const packed_half4 * src,
|
||||
device packed_float4 * dst,
|
||||
uint gid [[thread_position_in_grid]]) {
|
||||
const uint i = gid * 4u;
|
||||
if (i >= n) {
|
||||
return;
|
||||
}
|
||||
|
||||
const uint remaining = n - i;
|
||||
if (remaining >= 4u) {
|
||||
const half4 value = half4(src[gid]);
|
||||
dst[gid] = packed_float4(float4(value));
|
||||
return;
|
||||
}
|
||||
|
||||
device const half * src_scalar = (device const half *)src;
|
||||
device float * dst_scalar = (device float *)dst;
|
||||
for (uint lane = 0; lane < remaining; ++lane) {
|
||||
dst_scalar[i + lane] = float(src_scalar[i + lane]);
|
||||
}
|
||||
}
|
||||
|
||||
// Bitwise F16 transport for cache staging. Use ushort rather than half so NaN
|
||||
// payloads and every other binary16 encoding pass through unchanged.
|
||||
kernel void kernel_cpy_contig_f16_f16_bits_4(
|
||||
constant uint & n,
|
||||
device const packed_ushort4 * src,
|
||||
device packed_ushort4 * dst,
|
||||
uint gid [[thread_position_in_grid]]) {
|
||||
const uint i = gid * 4u;
|
||||
if (i >= n) {
|
||||
return;
|
||||
}
|
||||
|
||||
const uint remaining = n - i;
|
||||
if (remaining >= 4u) {
|
||||
dst[gid] = src[gid];
|
||||
return;
|
||||
}
|
||||
|
||||
device const ushort * src_scalar = (device const ushort *)src;
|
||||
device ushort * dst_scalar = (device ushort *)dst;
|
||||
for (uint lane = 0; lane < remaining; ++lane) {
|
||||
dst_scalar[i + lane] = src_scalar[i + lane];
|
||||
}
|
||||
}
|
||||
|
||||
struct ds4_metal_args_flash_kv_stage_f16 {
|
||||
uint raw_cap;
|
||||
uint raw_start;
|
||||
uint n_raw;
|
||||
uint n_comp;
|
||||
uint pad_rows;
|
||||
uint shared_pad;
|
||||
};
|
||||
|
||||
// Decode-time gathered attention consumes a logical raw-cache ring followed
|
||||
// by an already-F16 compressed cache. Pack both regions into the contiguous
|
||||
// F16 FlashAttention scratch in one dispatch. The raw conversion expression
|
||||
// and compressed ushort4 transport exactly match the standalone copy kernels.
|
||||
kernel void kernel_dsv4_flash_kv_stage_f16(
|
||||
constant ds4_metal_args_flash_kv_stage_f16 & args,
|
||||
device const char * raw_src,
|
||||
device const char * comp_src,
|
||||
device char * dst,
|
||||
device const char * mask_src,
|
||||
device char * pad_dst,
|
||||
uint gid [[thread_position_in_grid]]) {
|
||||
constexpr uint row_vecs = 128;
|
||||
const uint raw_vecs = args.n_raw * row_vecs;
|
||||
const uint n_keys = args.n_raw + args.n_comp;
|
||||
const uint total_vecs = n_keys * row_vecs;
|
||||
|
||||
if (gid < raw_vecs) {
|
||||
const uint logical_row = gid >> 7;
|
||||
const uint col = gid & 127u;
|
||||
uint physical_row = args.raw_start + logical_row;
|
||||
if (physical_row >= args.raw_cap) {
|
||||
physical_row -= args.raw_cap;
|
||||
}
|
||||
device const packed_float4 *raw =
|
||||
(device const packed_float4 *)raw_src;
|
||||
device packed_half4 *dst_half = (device packed_half4 *)dst;
|
||||
const float4 value =
|
||||
float4(raw[physical_row * row_vecs + col]);
|
||||
dst_half[gid] = packed_half4(half4(value));
|
||||
return;
|
||||
}
|
||||
|
||||
if (gid < total_vecs) {
|
||||
device const packed_ushort4 *comp =
|
||||
(device const packed_ushort4 *)comp_src;
|
||||
device packed_ushort4 *dst_bits = (device packed_ushort4 *)dst;
|
||||
dst_bits[gid] = comp[gid - raw_vecs];
|
||||
return;
|
||||
}
|
||||
|
||||
// The vector FlashAttention kernel redirects its final partial block to
|
||||
// a compact K/V/mask buffer. When requested, append those writes to this
|
||||
// dispatch so gathered decode does not need a standalone pad dispatch.
|
||||
const uint pad_rows = args.pad_rows;
|
||||
const uint pad_vecs = pad_rows * row_vecs;
|
||||
uint pad_gid = gid - total_vecs;
|
||||
if (pad_gid < pad_vecs) {
|
||||
const uint row = pad_gid >> 7;
|
||||
const uint col = pad_gid & 127u;
|
||||
const uint valid_rows = n_keys % pad_rows;
|
||||
device packed_half4 *pad_half = (device packed_half4 *)pad_dst;
|
||||
device packed_ushort4 *pad_bits = (device packed_ushort4 *)pad_dst;
|
||||
if (row >= valid_rows) {
|
||||
const packed_half4 zero = packed_half4(half4(0.0h));
|
||||
pad_half[pad_gid] = zero;
|
||||
if (!args.shared_pad) {
|
||||
pad_half[pad_vecs + pad_gid] = zero;
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
const uint logical_row = n_keys - valid_rows + row;
|
||||
if (logical_row < args.n_raw) {
|
||||
uint physical_row = args.raw_start + logical_row;
|
||||
if (physical_row >= args.raw_cap) {
|
||||
physical_row -= args.raw_cap;
|
||||
}
|
||||
device const packed_float4 *raw =
|
||||
(device const packed_float4 *)raw_src;
|
||||
const float4 value =
|
||||
float4(raw[physical_row * row_vecs + col]);
|
||||
const packed_half4 value_half = packed_half4(half4(value));
|
||||
pad_half[pad_gid] = value_half;
|
||||
if (!args.shared_pad) {
|
||||
pad_half[pad_vecs + pad_gid] = value_half;
|
||||
}
|
||||
} else {
|
||||
device const packed_ushort4 *comp =
|
||||
(device const packed_ushort4 *)comp_src;
|
||||
const packed_ushort4 value_bits =
|
||||
comp[(logical_row - args.n_raw) * row_vecs + col];
|
||||
pad_bits[pad_gid] = value_bits;
|
||||
if (!args.shared_pad) {
|
||||
pad_bits[pad_vecs + pad_gid] = value_bits;
|
||||
}
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
pad_gid -= pad_vecs;
|
||||
if (pad_gid < pad_rows) {
|
||||
const uint valid_rows = n_keys % pad_rows;
|
||||
device const ushort *mask_bits = (device const ushort *)mask_src;
|
||||
device ushort *pad_mask_bits =
|
||||
(device ushort *)pad_dst + 2u * pad_vecs * 4u;
|
||||
pad_mask_bits[pad_gid] = pad_gid < valid_rows
|
||||
? mask_bits[n_keys - valid_rows + pad_gid]
|
||||
: 0xfbffu;
|
||||
}
|
||||
}
|
||||
2389
metal/dense.metal
Normal file
2389
metal/dense.metal
Normal file
File diff suppressed because it is too large
Load Diff
1017
metal/dsv4_hc.metal
Normal file
1017
metal/dsv4_hc.metal
Normal file
File diff suppressed because it is too large
Load Diff
369
metal/dsv4_kv.metal
Normal file
369
metal/dsv4_kv.metal
Normal file
@@ -0,0 +1,369 @@
|
||||
constant float dsv4_e4m3fn_exp_scale[16] = {
|
||||
0.0f, 0.015625f, 0.03125f, 0.0625f,
|
||||
0.125f, 0.25f, 0.5f, 1.0f,
|
||||
2.0f, 4.0f, 8.0f, 16.0f,
|
||||
32.0f, 64.0f, 128.0f, 256.0f,
|
||||
};
|
||||
|
||||
constant float dsv4_e2m1fn_values[8] = {
|
||||
0.0f, 0.5f, 1.0f, 1.5f, 2.0f, 3.0f, 4.0f, 6.0f,
|
||||
};
|
||||
|
||||
struct ds4_metal_args_dsv4_fp8_kv_quantize {
|
||||
int64_t ne00;
|
||||
int64_t ne01;
|
||||
int64_t ne02;
|
||||
int64_t ne03;
|
||||
ulong nb00;
|
||||
ulong nb01;
|
||||
ulong nb02;
|
||||
ulong nb03;
|
||||
ulong nb0;
|
||||
ulong nb1;
|
||||
ulong nb2;
|
||||
ulong nb3;
|
||||
int n_rot;
|
||||
};
|
||||
|
||||
struct ds4_metal_args_dsv4_kv_fp8_store {
|
||||
int32_t head_dim;
|
||||
int32_t n_rot;
|
||||
int32_t raw_row;
|
||||
};
|
||||
|
||||
struct ds4_metal_args_dsv4_indexer_qat {
|
||||
uint32_t n_rows;
|
||||
uint32_t head_dim;
|
||||
uint64_t row_stride;
|
||||
};
|
||||
|
||||
struct ds4_metal_args_dsv4_ratio4_shift {
|
||||
uint32_t width;
|
||||
};
|
||||
|
||||
struct ds4_metal_args_dsv4_compressor_pack_ratio4 {
|
||||
uint32_t head_dim;
|
||||
uint32_t n_comp;
|
||||
uint32_t replay;
|
||||
uint32_t n_threads;
|
||||
};
|
||||
|
||||
struct ds4_metal_args_dsv4_compressor_store_one {
|
||||
uint32_t width;
|
||||
uint32_t ratio;
|
||||
uint32_t pos;
|
||||
uint32_t ape_type;
|
||||
};
|
||||
|
||||
static inline float dsv4_e4m3fn_value(int i) {
|
||||
const int exp = (i >> 3) & 0x0f;
|
||||
const int mant = i & 0x07;
|
||||
return exp == 0
|
||||
? float(mant) * 0.001953125f
|
||||
: (1.0f + float(mant) * 0.125f) * dsv4_e4m3fn_exp_scale[exp];
|
||||
}
|
||||
|
||||
static inline float dsv4_e4m3fn_dequant(float x) {
|
||||
const float sign = x < 0.0f ? -1.0f : 1.0f;
|
||||
const float ax = min(abs(x), 448.0f);
|
||||
|
||||
int lo = 0;
|
||||
int hi = 126;
|
||||
while (lo < hi) {
|
||||
const int mid = (lo + hi + 1) >> 1;
|
||||
if (dsv4_e4m3fn_value(mid) <= ax) {
|
||||
lo = mid;
|
||||
} else {
|
||||
hi = mid - 1;
|
||||
}
|
||||
}
|
||||
|
||||
int best = lo;
|
||||
if (best < 126) {
|
||||
const float best_diff = abs(ax - dsv4_e4m3fn_value(best));
|
||||
const float next_diff = abs(ax - dsv4_e4m3fn_value(best + 1));
|
||||
if (next_diff < best_diff || (next_diff == best_diff && ((best + 1) & 1) == 0 && (best & 1) != 0)) {
|
||||
best = best + 1;
|
||||
}
|
||||
}
|
||||
|
||||
return sign * dsv4_e4m3fn_value(best);
|
||||
}
|
||||
|
||||
static inline float dsv4_e2m1fn_dequant(float x) {
|
||||
const float sign = x < 0.0f ? -1.0f : 1.0f;
|
||||
const float ax = min(abs(x), 6.0f);
|
||||
int best = 0;
|
||||
float best_diff = abs(ax - dsv4_e2m1fn_values[0]);
|
||||
for (int i = 1; i < 8; i++) {
|
||||
const float diff = abs(ax - dsv4_e2m1fn_values[i]);
|
||||
if (diff < best_diff || (diff == best_diff && ((i & 1) == 0) && ((best & 1) != 0))) {
|
||||
best = i;
|
||||
best_diff = diff;
|
||||
}
|
||||
}
|
||||
return sign * dsv4_e2m1fn_values[best];
|
||||
}
|
||||
|
||||
// Quantizes the non-RoPE part of a KV row through E4M3FN and writes the
|
||||
// dequantized value back as float. DS4 uses this to match the FP8 KV-cache
|
||||
// semantics while keeping the Metal graph's cache buffers float-addressable.
|
||||
kernel void kernel_dsv4_fp8_kv_quantize_f32(
|
||||
constant ds4_metal_args_dsv4_fp8_kv_quantize & args,
|
||||
device const char * src0,
|
||||
device char * dst,
|
||||
threadgroup float * scratch [[threadgroup(0)]],
|
||||
uint row [[threadgroup_position_in_grid]],
|
||||
uint tid [[thread_position_in_threadgroup]]) {
|
||||
const int64_t n_rows = args.ne01 * args.ne02 * args.ne03;
|
||||
if ((int64_t) row >= n_rows) {
|
||||
return;
|
||||
}
|
||||
|
||||
const int64_t i1 = row % args.ne01;
|
||||
const int64_t i2 = (row / args.ne01) % args.ne02;
|
||||
const int64_t i3 = row / (args.ne01 * args.ne02);
|
||||
|
||||
device const char * src_base = src0 + i1*args.nb01 + i2*args.nb02 + i3*args.nb03;
|
||||
device char * dst_base = dst + i1*args.nb1 + i2*args.nb2 + i3*args.nb3;
|
||||
|
||||
const int64_t n_nope = args.ne00 - args.n_rot;
|
||||
|
||||
for (int64_t off = 0; off < n_nope; off += 64) {
|
||||
float v = 0.0f;
|
||||
if (tid < 64) {
|
||||
v = *((device const float *) (src_base + (off + tid)*args.nb00));
|
||||
scratch[tid] = abs(v);
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
for (uint stride = 32; stride > 0; stride >>= 1) {
|
||||
if (tid < stride) {
|
||||
scratch[tid] = max(scratch[tid], scratch[tid + stride]);
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
}
|
||||
|
||||
const float amax = max(scratch[0], 1.0e-4f);
|
||||
const float scale = exp2(ceil(log2(amax / 448.0f)));
|
||||
if (tid < 64) {
|
||||
const float q = dsv4_e4m3fn_dequant(clamp(v / scale, -448.0f, 448.0f)) * scale;
|
||||
*((device float *) (dst_base + (off + tid)*args.nb0)) = q;
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
}
|
||||
|
||||
for (int64_t i = n_nope + tid; i < args.ne00; i += 64) {
|
||||
*((device float *) (dst_base + i*args.nb0)) = *((device const float *) (src_base + i*args.nb00));
|
||||
}
|
||||
}
|
||||
|
||||
// The official DS4 indexer applies a 128-wide Hadamard rotation and then an
|
||||
// inplace FP4 activation-simulation pass to both indexer Q and indexer KV.
|
||||
kernel void kernel_dsv4_indexer_hadamard_fp4_f32(
|
||||
constant ds4_metal_args_dsv4_indexer_qat & args,
|
||||
device char * x,
|
||||
threadgroup float * scratch [[threadgroup(0)]],
|
||||
uint row [[threadgroup_position_in_grid]],
|
||||
uint tid [[thread_position_in_threadgroup]]) {
|
||||
if (row >= args.n_rows || args.head_dim != 128u || tid >= 128u) {
|
||||
return;
|
||||
}
|
||||
|
||||
threadgroup float *vals = scratch;
|
||||
threadgroup float *absbuf = scratch + 128;
|
||||
device float *xr = (device float *)(x + (uint64_t)row * args.row_stride);
|
||||
|
||||
vals[tid] = xr[tid];
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
for (uint stride = 1u; stride < 128u; stride <<= 1u) {
|
||||
if ((tid & stride) == 0u) {
|
||||
const uint base = (tid & ~(2u * stride - 1u)) + (tid & (stride - 1u));
|
||||
const float a = vals[base];
|
||||
const float b = vals[base + stride];
|
||||
vals[base] = a + b;
|
||||
vals[base + stride] = a - b;
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
}
|
||||
|
||||
float v = vals[tid] * 0.08838834764831845f;
|
||||
const uint block = tid >> 5u;
|
||||
const uint lane = tid & 31u;
|
||||
const uint block_base = block * 32u;
|
||||
absbuf[tid] = abs(v);
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
for (uint stride = 16u; stride > 0u; stride >>= 1u) {
|
||||
if (lane < stride) {
|
||||
absbuf[block_base + lane] = max(absbuf[block_base + lane],
|
||||
absbuf[block_base + lane + stride]);
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
}
|
||||
|
||||
const float amax = max(absbuf[block_base], 7.052966104933725e-38f);
|
||||
const float scale = exp2(ceil(log2(amax / 6.0f)));
|
||||
xr[tid] = dsv4_e2m1fn_dequant(clamp(v / scale, -6.0f, 6.0f)) * scale;
|
||||
}
|
||||
|
||||
// Decode-side KV finalizer after RoPE. The normal RoPE kernel intentionally
|
||||
// remains separate because tiny trigonometric codegen changes can flip later
|
||||
// sampled tokens. This kernel only fuses the FP8 round-trip for the non-RoPE
|
||||
// prefix with the F16-rounded raw-cache row used by FlashAttention.
|
||||
kernel void kernel_dsv4_kv_fp8_store_f32(
|
||||
constant ds4_metal_args_dsv4_kv_fp8_store & args,
|
||||
device float * kv,
|
||||
device float * raw_cache,
|
||||
threadgroup float * scratch [[threadgroup(0)]],
|
||||
uint tid [[thread_position_in_threadgroup]]) {
|
||||
const int head_dim = args.head_dim;
|
||||
const int n_rot = args.n_rot;
|
||||
const int n_nope = head_dim - n_rot;
|
||||
if (head_dim <= 0 || n_rot < 0 || n_nope < 0 || tid >= 64) {
|
||||
return;
|
||||
}
|
||||
|
||||
device float * raw = raw_cache + (int64_t)args.raw_row * head_dim;
|
||||
|
||||
for (int off = 0; off < n_nope; off += 64) {
|
||||
float v = 0.0f;
|
||||
if (off + (int)tid < n_nope) {
|
||||
v = kv[off + tid];
|
||||
scratch[tid] = abs(v);
|
||||
} else {
|
||||
scratch[tid] = 0.0f;
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
for (uint stride = 32; stride > 0; stride >>= 1) {
|
||||
if (tid < stride) {
|
||||
scratch[tid] = max(scratch[tid], scratch[tid + stride]);
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
}
|
||||
|
||||
const float amax = max(scratch[0], 1.0e-4f);
|
||||
const float fp8_scale = exp2(ceil(log2(amax / 448.0f)));
|
||||
if (off + (int)tid < n_nope) {
|
||||
const float q = dsv4_e4m3fn_dequant(clamp(v / fp8_scale, -448.0f, 448.0f)) * fp8_scale;
|
||||
kv[off + tid] = q;
|
||||
// Diagnostic only: skip the FP16 round-trip that normally matches the
|
||||
// half-typed FlashAttention KV buffer's precision. With this enabled the
|
||||
// indexer will see higher-precision raw values than FlashAttention does,
|
||||
// which is informative but not a production-ready setting.
|
||||
#ifdef DS4_METAL_KV_RAW_F32
|
||||
raw[off + tid] = q;
|
||||
#else
|
||||
raw[off + tid] = (float)((half)q);
|
||||
#endif
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
}
|
||||
|
||||
for (int i = n_nope + tid; i < head_dim; i += 64) {
|
||||
#ifdef DS4_METAL_KV_RAW_F32
|
||||
raw[i] = kv[i];
|
||||
#else
|
||||
raw[i] = (float)((half)kv[i]);
|
||||
#endif
|
||||
}
|
||||
}
|
||||
|
||||
// Builds the two ratio-4 softmax-pool packs together. Each output plane holds
|
||||
// four previous-half rows followed by four current-half rows. Normal prefill
|
||||
// seeds plane zero with 0/-inf; replay seeds it from the previous compressor
|
||||
// state. Later planes take their previous half from the prior input group.
|
||||
kernel void kernel_dsv4_compressor_pack_ratio4(
|
||||
constant ds4_metal_args_dsv4_compressor_pack_ratio4 & args,
|
||||
device const uint * kv,
|
||||
device const uint * score,
|
||||
device const uint * state_kv,
|
||||
device const uint * state_score,
|
||||
device uint * packed_kv,
|
||||
device uint * packed_score,
|
||||
uint2 group [[threadgroup_position_in_grid]],
|
||||
uint tid [[thread_index_in_threadgroup]]) {
|
||||
if (group.x >= args.n_comp || group.y >= 8u ||
|
||||
args.head_dim == 0u || args.n_threads == 0u) {
|
||||
return;
|
||||
}
|
||||
|
||||
const uint plane = group.x;
|
||||
const uint row = group.y;
|
||||
const uint64_t input_row_stride = 2ull * args.head_dim;
|
||||
const uint64_t dst_row = ((uint64_t)plane * 8u + row) * args.head_dim;
|
||||
|
||||
for (uint col = tid; col < args.head_dim; col += args.n_threads) {
|
||||
const uint64_t dst = dst_row + col;
|
||||
if (row >= 4u) {
|
||||
const uint token = plane * 4u + (row - 4u);
|
||||
const uint64_t src = (uint64_t)token * input_row_stride +
|
||||
args.head_dim + col;
|
||||
packed_kv[dst] = kv[src];
|
||||
packed_score[dst] = score[src];
|
||||
} else if (plane != 0u) {
|
||||
const uint token = (plane - 1u) * 4u + row;
|
||||
const uint64_t src = (uint64_t)token * input_row_stride + col;
|
||||
packed_kv[dst] = kv[src];
|
||||
packed_score[dst] = score[src];
|
||||
} else if (args.replay != 0u) {
|
||||
const uint64_t src = (uint64_t)row * input_row_stride + col;
|
||||
packed_kv[dst] = state_kv[src];
|
||||
packed_score[dst] = state_score[src];
|
||||
} else {
|
||||
packed_kv[dst] = 0u;
|
||||
packed_score[dst] = 0xff800000u;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Ratio-4 compression keeps two 4-row halves of recurrent state. After an
|
||||
// emitted compressed row, the second half becomes the next window's previous
|
||||
// half. The old encoder expressed this as four generic copies; this DS4-specific
|
||||
// kernel performs the KV and score copies together.
|
||||
kernel void kernel_dsv4_ratio4_shift_f32(
|
||||
constant ds4_metal_args_dsv4_ratio4_shift & args,
|
||||
device float * state_kv,
|
||||
device float * state_score,
|
||||
uint gid [[thread_position_in_grid]]) {
|
||||
const uint n = 4u * args.width;
|
||||
if (gid >= n) return;
|
||||
|
||||
state_kv[gid] = state_kv[n + gid];
|
||||
state_score[gid] = state_score[n + gid];
|
||||
}
|
||||
|
||||
// One-token compressor frontier update. Decode appends exactly one projected KV
|
||||
// row and one score row into a small recurrent state. The generic batch helper
|
||||
// expresses this as APE copy, score add, and two set_rows operations; this
|
||||
// kernel writes both state tensors directly while preserving the same
|
||||
// score + APE arithmetic.
|
||||
kernel void kernel_dsv4_compressor_store_one(
|
||||
constant ds4_metal_args_dsv4_compressor_store_one & args,
|
||||
device const float * kv,
|
||||
device const float * score,
|
||||
device const char * ape,
|
||||
device float * state_kv,
|
||||
device float * state_score,
|
||||
uint gid [[thread_position_in_grid]]) {
|
||||
if (gid >= args.width || args.width == 0 || args.ratio == 0) {
|
||||
return;
|
||||
}
|
||||
|
||||
const uint pos_mod = args.pos % args.ratio;
|
||||
const uint dst_row = args.ratio == 4u ? args.ratio + pos_mod : pos_mod;
|
||||
const uint dst = dst_row * args.width + gid;
|
||||
const uint ape_i = pos_mod * args.width + gid;
|
||||
|
||||
float ape_v;
|
||||
if (args.ape_type == 1u) {
|
||||
ape_v = (float)(((device const half *)ape)[ape_i]);
|
||||
} else {
|
||||
ape_v = ((device const float *)ape)[ape_i];
|
||||
}
|
||||
|
||||
state_kv[dst] = kv[gid];
|
||||
state_score[dst] = score[gid] + ape_v;
|
||||
}
|
||||
6157
metal/dsv4_misc.metal
Normal file
6157
metal/dsv4_misc.metal
Normal file
File diff suppressed because it is too large
Load Diff
385
metal/dsv4_rope.metal
Normal file
385
metal/dsv4_rope.metal
Normal file
@@ -0,0 +1,385 @@
|
||||
struct ds4_metal_args_dsv4_rope_tail {
|
||||
int64_t ne00;
|
||||
int64_t ne01;
|
||||
int64_t ne02;
|
||||
int64_t ne03;
|
||||
uint64_t nb00;
|
||||
uint64_t nb01;
|
||||
uint64_t nb02;
|
||||
uint64_t nb03;
|
||||
uint64_t nb0;
|
||||
uint64_t nb1;
|
||||
uint64_t nb2;
|
||||
uint64_t nb3;
|
||||
int32_t n_dims;
|
||||
int32_t mode;
|
||||
int32_t n_ctx_orig;
|
||||
int32_t inverse;
|
||||
float freq_base;
|
||||
float freq_scale;
|
||||
float ext_factor;
|
||||
float attn_factor;
|
||||
float beta_fast;
|
||||
float beta_slow;
|
||||
bool src2;
|
||||
};
|
||||
|
||||
struct ds4_metal_args_dsv4_rope_affine_pair {
|
||||
uint64_t row_bytes;
|
||||
uint64_t token_bytes;
|
||||
int32_t head_dim;
|
||||
int32_t n_dims;
|
||||
int32_t n_ctx_orig;
|
||||
int32_t inverse;
|
||||
uint32_t pos0;
|
||||
uint32_t pos_step;
|
||||
float freq_base;
|
||||
float freq_scale;
|
||||
float ext_factor;
|
||||
float attn_factor;
|
||||
float beta_fast;
|
||||
float beta_slow;
|
||||
};
|
||||
|
||||
static float rope_yarn_ramp(const float low, const float high, const int i0) {
|
||||
const float y = (i0 / 2 - low) / max(0.001f, high - low);
|
||||
return 1.0f - min(1.0f, max(0.0f, y));
|
||||
}
|
||||
|
||||
// YaRN algorithm based on LlamaYaRNScaledRotaryEmbedding.py from https://github.com/jquesnelle/yarn
|
||||
// MIT licensed. Copyright (c) 2023 Jeffrey Quesnelle and Bowen Peng.
|
||||
static void rope_yarn(
|
||||
float theta_extrap, float freq_scale, float corr_dims[2], int i0, float ext_factor, float mscale,
|
||||
thread float * cos_theta, thread float * sin_theta) {
|
||||
// Get n-d rotational scaling corrected for extrapolation
|
||||
float theta_interp = freq_scale * theta_extrap;
|
||||
float theta = theta_interp;
|
||||
if (ext_factor != 0.0f) {
|
||||
float ramp_mix = rope_yarn_ramp(corr_dims[0], corr_dims[1], i0) * ext_factor;
|
||||
theta = theta_interp * (1 - ramp_mix) + theta_extrap * ramp_mix;
|
||||
|
||||
// Get n-d magnitude scaling corrected for interpolation
|
||||
mscale *= 1.0f + 0.1f * log(1.0f / freq_scale);
|
||||
}
|
||||
*cos_theta = cos(theta) * mscale;
|
||||
*sin_theta = sin(theta) * mscale;
|
||||
}
|
||||
|
||||
// Apparently solving `n_rot = 2pi * x * base^((2 * max_pos_emb) / n_dims)` for x, we get
|
||||
// `corr_fac(n_rot) = n_dims * log(max_pos_emb / (n_rot * 2pi)) / (2 * log(base))`
|
||||
static float rope_yarn_corr_factor(int n_dims, int n_ctx_orig, float n_rot, float base) {
|
||||
return n_dims * log(n_ctx_orig / (n_rot * 2 * M_PI_F)) / (2 * log(base));
|
||||
}
|
||||
|
||||
static void rope_yarn_corr_dims(
|
||||
int n_dims, int n_ctx_orig, float freq_base, float beta_fast, float beta_slow, float dims[2]
|
||||
) {
|
||||
// start and end correction dims
|
||||
dims[0] = max(0.0f, floor(rope_yarn_corr_factor(n_dims, n_ctx_orig, beta_fast, freq_base)));
|
||||
dims[1] = min(n_dims - 1.0f, ceil(rope_yarn_corr_factor(n_dims, n_ctx_orig, beta_slow, freq_base)));
|
||||
}
|
||||
|
||||
// Applies DeepSeek V4's partial RoPE: the no-position prefix is copied and only
|
||||
// the rotated tail is transformed. This is used for Q/K after their projections
|
||||
// and before writing/reading the attention KV state.
|
||||
kernel void kernel_dsv4_rope_tail_f32(
|
||||
constant ds4_metal_args_dsv4_rope_tail & args,
|
||||
device const char * src0,
|
||||
device const char * src1,
|
||||
device const char * src2,
|
||||
device char * dst,
|
||||
uint tid [[thread_index_in_threadgroup]],
|
||||
ushort3 ntg [[threads_per_threadgroup]],
|
||||
uint3 tgpig [[threadgroup_position_in_grid]]) {
|
||||
const int i1 = tgpig[0];
|
||||
const int i2 = tgpig[1];
|
||||
const int i3 = tgpig[2];
|
||||
|
||||
const int n_nope = args.ne00 - args.n_dims;
|
||||
if (n_nope < 0) {
|
||||
return;
|
||||
}
|
||||
|
||||
device const int32_t * pos = (device const int32_t *) src1;
|
||||
|
||||
float corr_dims[2];
|
||||
rope_yarn_corr_dims(args.n_dims, args.n_ctx_orig, args.freq_base, args.beta_fast, args.beta_slow, corr_dims);
|
||||
|
||||
const float theta_base = (float) pos[i2];
|
||||
const float inv_ndims = -1.f/args.n_dims;
|
||||
const bool is_neox = args.mode == 2;
|
||||
|
||||
for (int i0 = tid; i0 < args.ne00; i0 += ntg.x) {
|
||||
device const char * src_base = src0 + i3*args.nb03 + i2*args.nb02 + i1*args.nb01;
|
||||
device char * dst_base = dst + i3*args.nb3 + i2*args.nb2 + i1*args.nb1;
|
||||
|
||||
if (i0 < n_nope) {
|
||||
*((device float *) (dst_base + i0*args.nb0)) = *((device const float *) (src_base + i0*args.nb00));
|
||||
continue;
|
||||
}
|
||||
|
||||
const int r = i0 - n_nope;
|
||||
if (is_neox) {
|
||||
const int n_half = args.n_dims/2;
|
||||
if (r >= n_half) {
|
||||
continue;
|
||||
}
|
||||
|
||||
const int ic = r;
|
||||
const int rel_i0 = 2*ic;
|
||||
#ifdef DS4_METAL_ROPE_EXP2_LOG2
|
||||
// Equivalent to pow(freq_base, k) but expressed through IEEE-754
|
||||
// primitives that have tighter precision guarantees than Metal's pow().
|
||||
const float theta = theta_base * exp2(inv_ndims * (float)rel_i0 * log2(args.freq_base));
|
||||
#else
|
||||
const float theta = theta_base * pow(args.freq_base, inv_ndims*rel_i0);
|
||||
#endif
|
||||
const float freq_factor = args.src2 ? ((device const float *) src2)[ic] : 1.0f;
|
||||
|
||||
float cos_theta;
|
||||
float sin_theta;
|
||||
rope_yarn(theta/freq_factor, args.freq_scale, corr_dims, rel_i0, args.ext_factor, args.attn_factor, &cos_theta, &sin_theta);
|
||||
if (args.inverse) {
|
||||
sin_theta = -sin_theta;
|
||||
}
|
||||
|
||||
const int j0 = n_nope + ic;
|
||||
const int j1 = n_nope + ic + n_half;
|
||||
const float x0 = *((device const float *) (src_base + j0*args.nb00));
|
||||
const float x1 = *((device const float *) (src_base + j1*args.nb00));
|
||||
|
||||
*((device float *) (dst_base + j0*args.nb0)) = x0*cos_theta - x1*sin_theta;
|
||||
*((device float *) (dst_base + j1*args.nb0)) = x0*sin_theta + x1*cos_theta;
|
||||
} else {
|
||||
if ((r & 1) != 0) {
|
||||
continue;
|
||||
}
|
||||
|
||||
const int ic = r/2;
|
||||
#ifdef DS4_METAL_ROPE_EXP2_LOG2
|
||||
const float theta = theta_base * exp2(inv_ndims * (float)r * log2(args.freq_base));
|
||||
#else
|
||||
const float theta = theta_base * pow(args.freq_base, inv_ndims*r);
|
||||
#endif
|
||||
const float freq_factor = args.src2 ? ((device const float *) src2)[ic] : 1.0f;
|
||||
|
||||
float cos_theta;
|
||||
float sin_theta;
|
||||
rope_yarn(theta/freq_factor, args.freq_scale, corr_dims, r, args.ext_factor, args.attn_factor, &cos_theta, &sin_theta);
|
||||
if (args.inverse) {
|
||||
sin_theta = -sin_theta;
|
||||
}
|
||||
|
||||
const int j0 = n_nope + r;
|
||||
const int j1 = j0 + 1;
|
||||
const float x0 = *((device const float *) (src_base + j0*args.nb00));
|
||||
const float x1 = *((device const float *) (src_base + j1*args.nb00));
|
||||
|
||||
*((device float *) (dst_base + j0*args.nb0)) = x0*cos_theta - x1*sin_theta;
|
||||
*((device float *) (dst_base + j1*args.nb0)) = x0*sin_theta + x1*cos_theta;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// DS4 only calls the Metal RoPE helper in-place and uses the adjacent-pair
|
||||
// layout (mode 0). The generic kernel above still copies the no-position
|
||||
// prefix, even though source and destination alias, and dispatches enough
|
||||
// lanes for the full head. This specialization maps lanes directly to the
|
||||
// rotated pairs and deliberately never reads or writes the unchanged prefix.
|
||||
// Keep the arithmetic below in the same order as the mode-0 branch above so
|
||||
// the optimized and reference paths remain bit-identical.
|
||||
kernel void kernel_dsv4_rope_tail_f32_inplace_pair(
|
||||
constant ds4_metal_args_dsv4_rope_tail & args,
|
||||
device const char * src0,
|
||||
device const char * src1,
|
||||
device const char * src2,
|
||||
device char * dst,
|
||||
uint tid [[thread_index_in_threadgroup]],
|
||||
ushort3 ntg [[threads_per_threadgroup]],
|
||||
uint3 tgpig [[threadgroup_position_in_grid]]) {
|
||||
if (args.mode != 0) {
|
||||
return;
|
||||
}
|
||||
|
||||
const int i1 = tgpig[0];
|
||||
const int i2 = tgpig[1];
|
||||
const int i3 = tgpig[2];
|
||||
const int n_nope = args.ne00 - args.n_dims;
|
||||
if (n_nope < 0) {
|
||||
return;
|
||||
}
|
||||
|
||||
device const int32_t * pos = (device const int32_t *) src1;
|
||||
|
||||
float corr_dims[2];
|
||||
rope_yarn_corr_dims(args.n_dims, args.n_ctx_orig, args.freq_base, args.beta_fast, args.beta_slow, corr_dims);
|
||||
|
||||
const float theta_base = (float) pos[i2];
|
||||
const float inv_ndims = -1.f/args.n_dims;
|
||||
device const char * src_base = src0 + i3*args.nb03 + i2*args.nb02 + i1*args.nb01;
|
||||
device char * dst_base = dst + i3*args.nb3 + i2*args.nb2 + i1*args.nb1;
|
||||
|
||||
// Keep each pair on the same SIMD lane as the generic in-place dispatch.
|
||||
// n_nope is 32-aligned for the supported DS4 shapes, so logical tail index
|
||||
// r ran on lane r%32 in the reference kernel. Compacting pairs would move
|
||||
// them to lane (r/2)%32 and can perturb fast-math code generation.
|
||||
for (int r = tid; r < args.n_dims; r += ntg.x) {
|
||||
if ((r & 1) != 0) {
|
||||
continue;
|
||||
}
|
||||
const int ic = r/2;
|
||||
#ifdef DS4_METAL_ROPE_EXP2_LOG2
|
||||
const float theta = theta_base * exp2(inv_ndims * (float)r * log2(args.freq_base));
|
||||
#else
|
||||
const float theta = theta_base * pow(args.freq_base, inv_ndims*r);
|
||||
#endif
|
||||
const float freq_factor = args.src2 ? ((device const float *) src2)[ic] : 1.0f;
|
||||
|
||||
float cos_theta;
|
||||
float sin_theta;
|
||||
rope_yarn(theta/freq_factor, args.freq_scale, corr_dims, r, args.ext_factor, args.attn_factor, &cos_theta, &sin_theta);
|
||||
if (args.inverse) {
|
||||
sin_theta = -sin_theta;
|
||||
}
|
||||
|
||||
const int j0 = n_nope + r;
|
||||
const int j1 = j0 + 1;
|
||||
const float x0 = *((device const float *) (src_base + j0*args.nb00));
|
||||
const float x1 = *((device const float *) (src_base + j1*args.nb00));
|
||||
|
||||
*((device float *) (dst_base + j0*args.nb0)) = x0*cos_theta - x1*sin_theta;
|
||||
*((device float *) (dst_base + j1*args.nb0)) = x0*sin_theta + x1*cos_theta;
|
||||
}
|
||||
}
|
||||
|
||||
// The Q/K RoPE calls use the same position and scaling parameters for every
|
||||
// head. Group four heads into one threadgroup so one 64-thread cohort computes
|
||||
// the 32 adjacent-pair coefficients and the other cohorts reuse them. Keeping
|
||||
// r on the same r%32 SIMD lane as the per-head specialization preserves the
|
||||
// fast-math instruction mapping; only the redundant coefficient work changes.
|
||||
kernel void kernel_dsv4_rope_tail_f32_inplace_pair_shared4(
|
||||
constant ds4_metal_args_dsv4_rope_tail & args,
|
||||
device const char * src0,
|
||||
device const char * src1,
|
||||
device const char * src2,
|
||||
device char * dst,
|
||||
uint tid [[thread_index_in_threadgroup]],
|
||||
uint3 tgpig [[threadgroup_position_in_grid]]) {
|
||||
if (args.mode != 0 || args.n_dims != 64) {
|
||||
return;
|
||||
}
|
||||
|
||||
const int n_nope = args.ne00 - args.n_dims;
|
||||
if (n_nope < 0) {
|
||||
return;
|
||||
}
|
||||
|
||||
const uint cohort = tid >> 6;
|
||||
const int r = (int)(tid & 63u);
|
||||
threadgroup float cos_shared[32];
|
||||
threadgroup float sin_shared[32];
|
||||
|
||||
device const int32_t * pos = (device const int32_t *) src1;
|
||||
float corr_dims[2];
|
||||
rope_yarn_corr_dims(args.n_dims, args.n_ctx_orig, args.freq_base, args.beta_fast, args.beta_slow, corr_dims);
|
||||
|
||||
const int i2 = tgpig[1];
|
||||
const float theta_base = (float) pos[i2];
|
||||
const float inv_ndims = -1.f/args.n_dims;
|
||||
|
||||
if (cohort == 0 && (r & 1) == 0) {
|
||||
const int ic = r/2;
|
||||
#ifdef DS4_METAL_ROPE_EXP2_LOG2
|
||||
const float theta = theta_base * exp2(inv_ndims * (float)r * log2(args.freq_base));
|
||||
#else
|
||||
const float theta = theta_base * pow(args.freq_base, inv_ndims*r);
|
||||
#endif
|
||||
const float freq_factor = args.src2 ? ((device const float *) src2)[ic] : 1.0f;
|
||||
|
||||
float cos_theta;
|
||||
float sin_theta;
|
||||
rope_yarn(theta/freq_factor, args.freq_scale, corr_dims, r, args.ext_factor, args.attn_factor, &cos_theta, &sin_theta);
|
||||
if (args.inverse) {
|
||||
sin_theta = -sin_theta;
|
||||
}
|
||||
cos_shared[ic] = cos_theta;
|
||||
sin_shared[ic] = sin_theta;
|
||||
}
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
const int i1 = (int)tgpig[0]*4 + (int)cohort;
|
||||
if (i1 >= args.ne01 || (r & 1) != 0) {
|
||||
return;
|
||||
}
|
||||
|
||||
const int i3 = tgpig[2];
|
||||
device const char * src_base = src0 + i3*args.nb03 + i2*args.nb02 + i1*args.nb01;
|
||||
device char * dst_base = dst + i3*args.nb3 + i2*args.nb2 + i1*args.nb1;
|
||||
const int j0 = n_nope + r;
|
||||
const int j1 = j0 + 1;
|
||||
const float x0 = *((device const float *) (src_base + j0*args.nb00));
|
||||
const float x1 = *((device const float *) (src_base + j1*args.nb00));
|
||||
const float cos_theta = cos_shared[r/2];
|
||||
const float sin_theta = sin_shared[r/2];
|
||||
|
||||
*((device float *) (dst_base + j0*args.nb0)) = x0*cos_theta - x1*sin_theta;
|
||||
*((device float *) (dst_base + j1*args.nb0)) = x0*sin_theta + x1*cos_theta;
|
||||
}
|
||||
|
||||
// DS4 positions are always affine within one RoPE dispatch. This variant
|
||||
// reconstructs the same wrapped int32 position in-kernel, avoiding the host
|
||||
// position array and its buffer binding while preserving the pair lane mapping
|
||||
// and all floating-point operations of the specialization above.
|
||||
kernel void kernel_dsv4_rope_tail_f32_inplace_pair_affine(
|
||||
constant ds4_metal_args_dsv4_rope_affine_pair & args [[buffer(0)]],
|
||||
device const char * src0 [[buffer(1)]],
|
||||
device char * dst [[buffer(4)]],
|
||||
uint tid [[thread_index_in_threadgroup]],
|
||||
ushort3 ntg [[threads_per_threadgroup]],
|
||||
uint3 tgpig [[threadgroup_position_in_grid]]) {
|
||||
const int i1 = tgpig[0];
|
||||
const int i2 = tgpig[1];
|
||||
const int n_nope = args.head_dim - args.n_dims;
|
||||
if (n_nope < 0) {
|
||||
return;
|
||||
}
|
||||
|
||||
float corr_dims[2];
|
||||
rope_yarn_corr_dims(args.n_dims, args.n_ctx_orig, args.freq_base, args.beta_fast, args.beta_slow, corr_dims);
|
||||
|
||||
const uint raw_pos = args.pos0 + (uint)i2 * args.pos_step;
|
||||
const float theta_base = (float)as_type<int>(raw_pos);
|
||||
const float inv_ndims = -1.f/args.n_dims;
|
||||
device const char * src_base =
|
||||
src0 + (uint64_t)i2*args.token_bytes + (uint64_t)i1*args.row_bytes;
|
||||
device char * dst_base =
|
||||
dst + (uint64_t)i2*args.token_bytes + (uint64_t)i1*args.row_bytes;
|
||||
|
||||
for (int r = tid; r < args.n_dims; r += ntg.x) {
|
||||
if ((r & 1) != 0) {
|
||||
continue;
|
||||
}
|
||||
#ifdef DS4_METAL_ROPE_EXP2_LOG2
|
||||
const float theta = theta_base * exp2(inv_ndims * (float)r * log2(args.freq_base));
|
||||
#else
|
||||
const float theta = theta_base * pow(args.freq_base, inv_ndims*r);
|
||||
#endif
|
||||
const float freq_factor = 1.0f;
|
||||
|
||||
float cos_theta;
|
||||
float sin_theta;
|
||||
rope_yarn(theta/freq_factor, args.freq_scale, corr_dims, r, args.ext_factor, args.attn_factor, &cos_theta, &sin_theta);
|
||||
if (args.inverse) {
|
||||
sin_theta = -sin_theta;
|
||||
}
|
||||
|
||||
const int j0 = n_nope + r;
|
||||
const int j1 = j0 + 1;
|
||||
const float x0 = *((device const float *) (src_base + j0*sizeof(float)));
|
||||
const float x1 = *((device const float *) (src_base + j1*sizeof(float)));
|
||||
|
||||
*((device float *) (dst_base + j0*sizeof(float))) = x0*cos_theta - x1*sin_theta;
|
||||
*((device float *) (dst_base + j1*sizeof(float))) = x0*sin_theta + x1*cos_theta;
|
||||
}
|
||||
}
|
||||
1441
metal/flash_attn.metal
Normal file
1441
metal/flash_attn.metal
Normal file
File diff suppressed because it is too large
Load Diff
186
metal/get_rows.metal
Normal file
186
metal/get_rows.metal
Normal file
@@ -0,0 +1,186 @@
|
||||
// DS4 Metal get-rows kernel.
|
||||
|
||||
struct ds4_metal_args_get_rows {
|
||||
int32_t ne00t;
|
||||
int32_t ne00;
|
||||
uint64_t nb01;
|
||||
uint64_t nb02;
|
||||
uint64_t nb03;
|
||||
int32_t ne10;
|
||||
uint64_t nb10;
|
||||
uint64_t nb11;
|
||||
uint64_t nb12;
|
||||
uint64_t nb1;
|
||||
uint64_t nb2;
|
||||
uint64_t nb3;
|
||||
};
|
||||
|
||||
struct ds4_metal_args_get_rows_q8_0 {
|
||||
int32_t n_embd;
|
||||
int32_t n_vocab;
|
||||
int32_t n_tokens;
|
||||
uint64_t src_row_bytes;
|
||||
uint64_t dst_row_bytes;
|
||||
uint64_t token_stride;
|
||||
};
|
||||
|
||||
struct ds4_get_rows_block_q4_0 {
|
||||
half d;
|
||||
uchar qs[16];
|
||||
};
|
||||
|
||||
struct ds4_get_rows_block_q4_K {
|
||||
half d;
|
||||
half dmin;
|
||||
uchar scales[12];
|
||||
uchar qs[128];
|
||||
};
|
||||
|
||||
static inline uchar2 ds4_get_rows_q4_K_scale_min(int j, int k, device const uchar *q) {
|
||||
return j < 4 ? uchar2{uchar(q[j + 0 + k] & 63), uchar(q[j + 4 + k] & 63)}
|
||||
: uchar2{uchar((q[j + 4 + k] & 0x0f) | ((q[j - 4 + k] & 0xc0) >> 2)),
|
||||
uchar((q[j + 4 + k] >> 4) | ((q[j - 0 + k] & 0xc0) >> 2))};
|
||||
}
|
||||
|
||||
// Gathers embedding/table rows by integer ids. DS4 uses this for token
|
||||
// embeddings and small indexed tables such as router/hash lookup outputs.
|
||||
template<typename T0, typename T>
|
||||
kernel void kernel_get_rows_f(
|
||||
constant ds4_metal_args_get_rows & args,
|
||||
device const char * src0,
|
||||
device const char * src1,
|
||||
device char * dst,
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
ushort tiitg[[thread_index_in_threadgroup]],
|
||||
ushort3 ntg [[threads_per_threadgroup]]) {
|
||||
const int32_t iw0 = tgpig.x/args.ne10;
|
||||
const int32_t i10 = tgpig.x%args.ne10;
|
||||
const int32_t i11 = tgpig.y;
|
||||
const int32_t i12 = tgpig.z;
|
||||
|
||||
const int32_t r = ((const device int32_t *) (src1 + i12*args.nb12 + i11*args.nb11 + i10*args.nb10))[0];
|
||||
|
||||
const int32_t i02 = i11;
|
||||
const int32_t i03 = i12;
|
||||
|
||||
auto psrc = (const device T0 *) (src0 + i03*args.nb03 + i02*args.nb02 + r*args.nb01);
|
||||
auto pdst = ( device T *) (dst + i12*args.nb3 + i11*args.nb2 + i10*args.nb1);
|
||||
|
||||
for (int ind = iw0*ntg.x + tiitg; ind < args.ne00t;) {
|
||||
pdst[ind] = psrc[ind];
|
||||
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
typedef decltype(kernel_get_rows_f<float, float>) get_rows_f_t;
|
||||
|
||||
// Host-visible gather variants for F32, F16, and I32 tables.
|
||||
template [[host_name("kernel_get_rows_f32")]] kernel get_rows_f_t kernel_get_rows_f<float, float>;
|
||||
template [[host_name("kernel_get_rows_f16")]] kernel get_rows_f_t kernel_get_rows_f<half, float>;
|
||||
template [[host_name("kernel_get_rows_i32")]] kernel get_rows_f_t kernel_get_rows_f<int32_t, int32_t>;
|
||||
|
||||
kernel void kernel_get_rows_q8_0_f32(
|
||||
constant ds4_metal_args_get_rows_q8_0 & args,
|
||||
device const char * src0,
|
||||
device const char * src1,
|
||||
device char * dst,
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
ushort tiitg[[thread_index_in_threadgroup]],
|
||||
ushort3 ntg [[threads_per_threadgroup]]) {
|
||||
const int32_t block = (int32_t)tgpig.x;
|
||||
const int32_t tok_i = (int32_t)tgpig.y;
|
||||
if (tok_i >= args.n_tokens) return;
|
||||
|
||||
const int32_t token =
|
||||
((const device int32_t *)(src1 + (uint64_t)tok_i*args.token_stride))[0];
|
||||
if (token < 0 || token >= args.n_vocab) return;
|
||||
|
||||
const device block_q8_0 *row =
|
||||
(const device block_q8_0 *)(src0 + (uint64_t)token * args.src_row_bytes);
|
||||
const device block_q8_0 *qb = row + block;
|
||||
device float *out =
|
||||
(device float *)(dst + (uint64_t)tok_i * args.dst_row_bytes);
|
||||
|
||||
const int32_t i0 = block * QK8_0;
|
||||
const float d = (float)qb->d;
|
||||
for (int32_t i = (int32_t)tiitg; i < QK8_0; i += (int32_t)ntg.x) {
|
||||
const int32_t idx = i0 + i;
|
||||
if (idx < args.n_embd) {
|
||||
out[idx] = d * (float)qb->qs[i];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
kernel void kernel_get_rows_q4_0_f32(
|
||||
constant ds4_metal_args_get_rows_q8_0 & args,
|
||||
device const char * src0,
|
||||
device const char * src1,
|
||||
device char * dst,
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
ushort tiitg[[thread_index_in_threadgroup]],
|
||||
ushort3 ntg [[threads_per_threadgroup]]) {
|
||||
const int32_t block = (int32_t)tgpig.x;
|
||||
const int32_t tok_i = (int32_t)tgpig.y;
|
||||
if (tok_i >= args.n_tokens) return;
|
||||
|
||||
const int32_t token =
|
||||
((const device int32_t *)(src1 + (uint64_t)tok_i * args.token_stride))[0];
|
||||
if (token < 0 || token >= args.n_vocab) return;
|
||||
|
||||
const device ds4_get_rows_block_q4_0 *row =
|
||||
(const device ds4_get_rows_block_q4_0 *)(src0 + (uint64_t)token * args.src_row_bytes);
|
||||
const device ds4_get_rows_block_q4_0 *qb = row + block;
|
||||
device float *out =
|
||||
(device float *)(dst + (uint64_t)tok_i * args.dst_row_bytes);
|
||||
|
||||
const int32_t i0 = block * 32;
|
||||
const float d = (float)qb->d;
|
||||
for (int32_t i = (int32_t)tiitg; i < 32; i += (int32_t)ntg.x) {
|
||||
const int32_t idx = i0 + i;
|
||||
if (idx < args.n_embd) {
|
||||
/* ggml Q4_0: elems 0..15 = low nibbles of qs[0..15], 16..31 = high. */
|
||||
const uint8_t packed = qb->qs[(uint)(i & 15)];
|
||||
const uint8_t q = (i < 16) ? (packed & 0x0f) : (packed >> 4);
|
||||
out[idx] = d * ((float)q - 8.0f);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
kernel void kernel_get_rows_q4_K_f32(
|
||||
constant ds4_metal_args_get_rows_q8_0 & args,
|
||||
device const char * src0,
|
||||
device const char * src1,
|
||||
device char * dst,
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
ushort tiitg[[thread_index_in_threadgroup]],
|
||||
ushort3 ntg [[threads_per_threadgroup]]) {
|
||||
const int32_t block = (int32_t)tgpig.x;
|
||||
const int32_t tok_i = (int32_t)tgpig.y;
|
||||
if (tok_i >= args.n_tokens) return;
|
||||
|
||||
const int32_t token =
|
||||
((const device int32_t *)(src1 + (uint64_t)tok_i * args.token_stride))[0];
|
||||
if (token < 0 || token >= args.n_vocab) return;
|
||||
|
||||
const device ds4_get_rows_block_q4_K *row =
|
||||
(const device ds4_get_rows_block_q4_K *)(src0 + (uint64_t)token * args.src_row_bytes);
|
||||
const device ds4_get_rows_block_q4_K *qb = row + block;
|
||||
device float *out =
|
||||
(device float *)(dst + (uint64_t)tok_i * args.dst_row_bytes);
|
||||
|
||||
const int32_t i0 = block * 256;
|
||||
for (int32_t i = (int32_t)tiitg; i < 256; i += (int32_t)ntg.x) {
|
||||
const int32_t idx = i0 + i;
|
||||
if (idx < args.n_embd) {
|
||||
const uint group = (uint)i >> 5u;
|
||||
const uint l = (uint)i & 31u;
|
||||
const uchar2 sm = ds4_get_rows_q4_K_scale_min((int)group, 0, qb->scales);
|
||||
const uint byte_off = (group >> 1u) * 32u + l;
|
||||
const uint shift = (group & 1u) * 4u;
|
||||
const uint q = ((uint)qb->qs[byte_off] >> shift) & 0x0fu;
|
||||
out[idx] = (float)qb->d * (float)sm.x * (float)q -
|
||||
(float)qb->dmin * (float)sm.y;
|
||||
}
|
||||
}
|
||||
}
|
||||
63
metal/glu.metal
Normal file
63
metal/glu.metal
Normal file
@@ -0,0 +1,63 @@
|
||||
struct ds4_metal_args_glu {
|
||||
int32_t ne00;
|
||||
uint64_t nb01;
|
||||
int32_t ne10;
|
||||
uint64_t nb11;
|
||||
int32_t ne0;
|
||||
uint64_t nb1;
|
||||
int32_t i00;
|
||||
int32_t i10;
|
||||
float alpha;
|
||||
float limit;
|
||||
};
|
||||
|
||||
// SwiGLU activation for the FFN inner state. DS4 clamps the shared expert with
|
||||
// the same swiglu_limit used by routed experts.
|
||||
kernel void kernel_swiglu_f32(
|
||||
constant ds4_metal_args_glu & args,
|
||||
device const char * src0,
|
||||
device const char * src1,
|
||||
device char * dst,
|
||||
uint tgpig[[threadgroup_position_in_grid]],
|
||||
uint tpitg[[thread_position_in_threadgroup]],
|
||||
uint ntg[[threads_per_threadgroup]]) {
|
||||
device const float * src0_row = (device const float *) ((device const char *) src0 + tgpig*args.nb01) + args.i00;
|
||||
device const float * src1_row = (device const float *) ((device const char *) src1 + tgpig*args.nb11) + args.i10;
|
||||
device float * dst_row = (device float *) ((device char *) dst + tgpig*args.nb1);
|
||||
|
||||
for (int i0 = tpitg; i0 < args.ne0; i0 += ntg) {
|
||||
float x0 = src0_row[i0];
|
||||
float x1 = src1_row[i0];
|
||||
if (args.limit > 1.0e-6f) {
|
||||
x0 = min(x0, args.limit);
|
||||
x1 = clamp(x1, -args.limit, args.limit);
|
||||
}
|
||||
|
||||
const float silu = x0 / (1.0f + exp(-x0));
|
||||
|
||||
dst_row[i0] = silu*x1*args.alpha;
|
||||
}
|
||||
}
|
||||
|
||||
kernel void kernel_swiglu_flat_f32(
|
||||
constant ds4_metal_args_glu & args,
|
||||
device const char * src0,
|
||||
device const char * src1,
|
||||
device char * dst,
|
||||
uint i [[thread_position_in_grid]]) {
|
||||
if (i >= (uint)args.ne0) return;
|
||||
|
||||
device const float * src0_f32 = (device const float *) src0 + args.i00;
|
||||
device const float * src1_f32 = (device const float *) src1 + args.i10;
|
||||
device float * dst_f32 = (device float *) dst;
|
||||
|
||||
float x0 = src0_f32[i];
|
||||
float x1 = src1_f32[i];
|
||||
if (args.limit > 1.0e-6f) {
|
||||
x0 = min(x0, args.limit);
|
||||
x1 = clamp(x1, -args.limit, args.limit);
|
||||
}
|
||||
|
||||
const float silu = x0 / (1.0f + exp(-x0));
|
||||
dst_f32[i] = silu*x1*args.alpha;
|
||||
}
|
||||
7858
metal/moe.metal
Normal file
7858
metal/moe.metal
Normal file
File diff suppressed because it is too large
Load Diff
243
metal/norm.metal
Normal file
243
metal/norm.metal
Normal file
@@ -0,0 +1,243 @@
|
||||
struct ds4_metal_args_norm {
|
||||
int32_t ne00;
|
||||
int32_t ne00_t;
|
||||
uint64_t nb1;
|
||||
uint64_t nb2;
|
||||
uint64_t nb3;
|
||||
float eps;
|
||||
int32_t nef1[3];
|
||||
int32_t nef2[3];
|
||||
int32_t nef3[3];
|
||||
uint64_t nbf1[3];
|
||||
uint64_t nbf2[3];
|
||||
uint64_t nbf3[3];
|
||||
};
|
||||
|
||||
// RMSNorm over one activation row, optionally fusing the learned weight
|
||||
// multiply. DS4 calls this before attention, before the FFN, and for plain
|
||||
// diagnostics that need normalized but unweighted rows.
|
||||
template <typename T, short F>
|
||||
kernel void kernel_rms_norm_fuse_impl(
|
||||
constant ds4_metal_args_norm & args,
|
||||
device const char * src0,
|
||||
device const char * src1_0,
|
||||
device const char * src1_1,
|
||||
device char * dst,
|
||||
threadgroup float * shmem_f32 [[threadgroup(0)]],
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
ushort3 tpitg[[thread_position_in_threadgroup]],
|
||||
ushort sgitg[[simdgroup_index_in_threadgroup]],
|
||||
ushort tiisg[[thread_index_in_simdgroup]],
|
||||
ushort3 ntg[[threads_per_threadgroup]]) {
|
||||
if (sgitg == 0) {
|
||||
shmem_f32[tiisg] = 0.0f;
|
||||
}
|
||||
|
||||
const int i01 = tgpig.x;
|
||||
const int i02 = tgpig.y;
|
||||
const int i03 = tgpig.z;
|
||||
|
||||
device const T * x = (device const T *) (src0 + i03*args.nbf3[0] + i02*args.nbf2[0] + i01*args.nbf1[0]);
|
||||
|
||||
device const T * f0 = (device const T *) (src1_0 + (i03%args.nef3[1])*args.nbf3[1] + (i02%args.nef2[1])*args.nbf2[1] + (i01%args.nef1[1])*args.nbf1[1]);
|
||||
device const T * f1 = (device const T *) (src1_1 + (i03%args.nef3[2])*args.nbf3[2] + (i02%args.nef2[2])*args.nbf2[2] + (i01%args.nef1[2])*args.nbf1[2]);
|
||||
|
||||
float sumf = 0.0f;
|
||||
|
||||
// parallel sum
|
||||
for (int i00 = tpitg.x; i00 < args.ne00_t; i00 += ntg.x) {
|
||||
sumf += dot(x[i00], x[i00]);
|
||||
}
|
||||
sumf = simd_sum(sumf);
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
if (tiisg == 0) {
|
||||
shmem_f32[sgitg] = sumf;
|
||||
}
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
sumf = shmem_f32[tiisg];
|
||||
sumf = simd_sum(sumf);
|
||||
|
||||
const float mean = sumf/args.ne00;
|
||||
const float scale = 1.0f/sqrt(mean + args.eps);
|
||||
|
||||
device T * y = (device T *) (dst + i03*args.nb3 + i02*args.nb2 + i01*args.nb1);
|
||||
for (int i00 = tpitg.x; i00 < args.ne00_t; i00 += ntg.x) {
|
||||
if (F == 1) {
|
||||
y[i00] = (x[i00]*scale);
|
||||
}
|
||||
if (F == 2) {
|
||||
y[i00] = (x[i00]*scale)*f0[i00];
|
||||
}
|
||||
if (F == 3) {
|
||||
y[i00] = (x[i00]*scale)*f0[i00] + f1[i00];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
typedef decltype(kernel_rms_norm_fuse_impl<float4, 1>) kernel_rms_norm_fuse_t;
|
||||
|
||||
// Host-visible RMSNorm variants: plain norm and norm multiplied by weight.
|
||||
template [[host_name("kernel_rms_norm_f32_4")]] kernel kernel_rms_norm_fuse_t kernel_rms_norm_fuse_impl<float4, 1>;
|
||||
template [[host_name("kernel_rms_norm_mul_f32_4")]] kernel kernel_rms_norm_fuse_t kernel_rms_norm_fuse_impl<float4, 2>;
|
||||
|
||||
kernel void kernel_add_rms_norm_mul_f32_4(
|
||||
constant ds4_metal_args_norm & args,
|
||||
device const char * src0,
|
||||
device const char * src1,
|
||||
device const char * weight,
|
||||
device char * sum_dst,
|
||||
device char * norm_dst,
|
||||
threadgroup float * shmem_f32 [[threadgroup(0)]],
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
ushort3 tpitg[[thread_position_in_threadgroup]],
|
||||
ushort sgitg[[simdgroup_index_in_threadgroup]],
|
||||
ushort tiisg[[thread_index_in_simdgroup]],
|
||||
ushort3 ntg[[threads_per_threadgroup]]) {
|
||||
if (sgitg == 0) shmem_f32[tiisg] = 0.0f;
|
||||
|
||||
const int i01 = tgpig.x;
|
||||
const int i02 = tgpig.y;
|
||||
const int i03 = tgpig.z;
|
||||
device const float4 *a = (device const float4 *)
|
||||
(src0 + i03*args.nbf3[0] + i02*args.nbf2[0] + i01*args.nbf1[0]);
|
||||
device const float4 *b = (device const float4 *)
|
||||
(src1 + i03*args.nbf3[0] + i02*args.nbf2[0] + i01*args.nbf1[0]);
|
||||
device const float4 *w = (device const float4 *)
|
||||
(weight + (i03%args.nef3[1])*args.nbf3[1] + (i02%args.nef2[1])*args.nbf2[1] + (i01%args.nef1[1])*args.nbf1[1]);
|
||||
device float4 *sum = (device float4 *)
|
||||
(sum_dst + i03*args.nb3 + i02*args.nb2 + i01*args.nb1);
|
||||
device float4 *norm = (device float4 *)
|
||||
(norm_dst + i03*args.nb3 + i02*args.nb2 + i01*args.nb1);
|
||||
|
||||
float sumf = 0.0f;
|
||||
for (int i00 = tpitg.x; i00 < args.ne00_t; i00 += ntg.x) {
|
||||
const float4 v = a[i00] + b[i00];
|
||||
sum[i00] = v;
|
||||
sumf += dot(v, v);
|
||||
}
|
||||
sumf = simd_sum(sumf);
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
if (tiisg == 0) shmem_f32[sgitg] = sumf;
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
sumf = simd_sum(shmem_f32[tiisg]);
|
||||
|
||||
const float mean = sumf / args.ne00;
|
||||
const float scale = 1.0f / sqrt(mean + args.eps);
|
||||
for (int i00 = tpitg.x; i00 < args.ne00_t; i00 += ntg.x) {
|
||||
norm[i00] = (sum[i00] * scale) * w[i00];
|
||||
}
|
||||
}
|
||||
|
||||
// RMSNorm reduction used when the following F16 matmul applies the row scale
|
||||
// while staging its RHS tile.
|
||||
kernel void kernel_rms_norm_scale_f32_4(
|
||||
constant ds4_metal_args_norm & args,
|
||||
device const char * src0,
|
||||
device float * dst_scale,
|
||||
threadgroup float * shmem_f32 [[threadgroup(0)]],
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
ushort3 tpitg[[thread_position_in_threadgroup]],
|
||||
ushort sgitg[[simdgroup_index_in_threadgroup]],
|
||||
ushort tiisg[[thread_index_in_simdgroup]],
|
||||
ushort3 ntg[[threads_per_threadgroup]]) {
|
||||
if (sgitg == 0) shmem_f32[tiisg] = 0.0f;
|
||||
|
||||
const int i01 = tgpig.x;
|
||||
const int i02 = tgpig.y;
|
||||
const int i03 = tgpig.z;
|
||||
device const float4 *x = (device const float4 *)
|
||||
(src0 + i03*args.nbf3[0] + i02*args.nbf2[0] + i01*args.nbf1[0]);
|
||||
|
||||
float sumf = 0.0f;
|
||||
for (int i00 = tpitg.x; i00 < args.ne00_t; i00 += ntg.x) {
|
||||
sumf += dot(x[i00], x[i00]);
|
||||
}
|
||||
sumf = simd_sum(sumf);
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
if (tiisg == 0) shmem_f32[sgitg] = sumf;
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
sumf = simd_sum(shmem_f32[tiisg]);
|
||||
|
||||
const float mean = sumf / args.ne00;
|
||||
const float scale = 1.0f / sqrt(mean + args.eps);
|
||||
if (tpitg.x == 0) dst_scale[i01] = scale;
|
||||
}
|
||||
|
||||
struct ds4_metal_args_qkv_rms_norm {
|
||||
int32_t q_n;
|
||||
int32_t q_n4;
|
||||
int32_t kv_n;
|
||||
int32_t kv_n4;
|
||||
uint64_t q_row_stride;
|
||||
uint64_t kv_row_stride;
|
||||
float eps;
|
||||
};
|
||||
|
||||
// Normalizes DS4's q-lora row and KV row in one dispatch. The two reductions
|
||||
// deliberately mirror kernel_rms_norm_mul_f32_4: Q uses the full 256-thread
|
||||
// row shape for 1024 floats, while KV only has work in the first 128 lanes for
|
||||
// its 512 floats. This keeps the q/kv normalization math aligned with the
|
||||
// standalone kernels while removing one tiny launch from the attention setup.
|
||||
kernel void kernel_dsv4_qkv_rms_norm_f32_4(
|
||||
constant ds4_metal_args_qkv_rms_norm & args,
|
||||
device const float4 * q_src,
|
||||
device const float4 * q_weight,
|
||||
device float4 * q_dst,
|
||||
device const float4 * kv_src,
|
||||
device const float4 * kv_weight,
|
||||
device float4 * kv_dst,
|
||||
threadgroup float * shmem_f32 [[threadgroup(0)]],
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
ushort3 tpitg[[thread_position_in_threadgroup]],
|
||||
ushort sgitg[[simdgroup_index_in_threadgroup]],
|
||||
ushort tiisg[[thread_index_in_simdgroup]],
|
||||
ushort3 ntg[[threads_per_threadgroup]]) {
|
||||
if (sgitg == 0) {
|
||||
shmem_f32[tiisg] = 0.0f;
|
||||
}
|
||||
|
||||
const uint row = tgpig.x;
|
||||
const bool kv_task = tgpig.y != 0;
|
||||
const int n = kv_task ? args.kv_n : args.q_n;
|
||||
const int n4 = kv_task ? args.kv_n4 : args.q_n4;
|
||||
const uint64_t row_stride4 = (kv_task ? args.kv_row_stride : args.q_row_stride) / sizeof(float4);
|
||||
|
||||
device const float4 * x = kv_task ? kv_src + row * row_stride4 : q_src + row * row_stride4;
|
||||
device const float4 * w = kv_task ? kv_weight : q_weight;
|
||||
device float4 * y = kv_task ? kv_dst + row * row_stride4 : q_dst + row * row_stride4;
|
||||
|
||||
float sumf = 0.0f;
|
||||
for (int i = tpitg.x; i < n4; i += ntg.x) {
|
||||
const float4 v = x[i];
|
||||
sumf += dot(v, v);
|
||||
}
|
||||
sumf = simd_sum(sumf);
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
if (tiisg == 0) {
|
||||
shmem_f32[sgitg] = sumf;
|
||||
}
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
sumf = shmem_f32[tiisg];
|
||||
sumf = simd_sum(sumf);
|
||||
|
||||
#ifdef DS4_METAL_NORM_RSQRT_DISABLE
|
||||
// Match the formula used by kernel_rms_norm_fuse_impl above so both RMSNorm
|
||||
// entry points produce bit-identical scales. Hardware rsqrt() and 1.0f/sqrt()
|
||||
// can differ by ~1 ULP and that difference compounds across 43 layers.
|
||||
const float scale = 1.0f / sqrt(sumf / float(n) + args.eps);
|
||||
#else
|
||||
const float scale = rsqrt(sumf / float(n) + args.eps);
|
||||
#endif
|
||||
|
||||
for (int i = tpitg.x; i < n4; i += ntg.x) {
|
||||
y[i] = (x[i] * scale) * w[i];
|
||||
}
|
||||
}
|
||||
52
metal/repeat.metal
Normal file
52
metal/repeat.metal
Normal file
@@ -0,0 +1,52 @@
|
||||
// DS4 Metal repeat kernel used for HC embedding expansion.
|
||||
|
||||
struct ds4_metal_args_repeat {
|
||||
int32_t ne00;
|
||||
int32_t ne01;
|
||||
int32_t ne02;
|
||||
int32_t ne03;
|
||||
uint64_t nb00;
|
||||
uint64_t nb01;
|
||||
uint64_t nb02;
|
||||
uint64_t nb03;
|
||||
int32_t ne0;
|
||||
int32_t ne1;
|
||||
int32_t ne2;
|
||||
int32_t ne3;
|
||||
uint64_t nb0;
|
||||
uint64_t nb1;
|
||||
uint64_t nb2;
|
||||
uint64_t nb3;
|
||||
};
|
||||
|
||||
// Repeats a source row into the HC channel dimension. DS4 uses this when the
|
||||
// token embedding has to become an HC activation block before layer 0.
|
||||
template<typename T>
|
||||
kernel void kernel_repeat(
|
||||
constant ds4_metal_args_repeat & args,
|
||||
device const char * src0,
|
||||
device char * dst,
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
ushort3 tpitg[[thread_position_in_threadgroup]],
|
||||
ushort3 ntg[[threads_per_threadgroup]]) {
|
||||
const int i3 = tgpig.z;
|
||||
const int i2 = tgpig.y;
|
||||
const int i1 = tgpig.x;
|
||||
|
||||
const int i03 = i3%args.ne03;
|
||||
const int i02 = i2%args.ne02;
|
||||
const int i01 = i1%args.ne01;
|
||||
|
||||
device const char * src0_ptr = src0 + i03*args.nb03 + i02*args.nb02 + i01*args.nb01;
|
||||
device char * dst_ptr = dst + i3*args.nb3 + i2*args.nb2 + i1*args.nb1;
|
||||
|
||||
for (int i0 = tpitg.x; i0 < args.ne0; i0 += ntg.x) {
|
||||
const int i00 = i0%args.ne00;
|
||||
*((device T *)(dst_ptr + i0*args.nb0)) = *((device T *)(src0_ptr + i00*args.nb00));
|
||||
}
|
||||
}
|
||||
|
||||
typedef decltype(kernel_repeat<float>) kernel_repeat_t;
|
||||
|
||||
// Host-visible F32 repeat used for HC expansion of embeddings.
|
||||
template [[host_name("kernel_repeat_f32")]] kernel kernel_repeat_t kernel_repeat<float>;
|
||||
55
metal/set_rows.metal
Normal file
55
metal/set_rows.metal
Normal file
@@ -0,0 +1,55 @@
|
||||
// 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>;
|
||||
241
metal/softmax.metal
Normal file
241
metal/softmax.metal
Normal file
@@ -0,0 +1,241 @@
|
||||
// DS4 Metal softmax kernel used by the compressor pooling compatibility path.
|
||||
// The single-compressed-row path is intentionally left as soft_max -> mul ->
|
||||
// sum_rows instead of using the fused dsv4_softmax_pool kernel.
|
||||
|
||||
struct ds4_metal_args_soft_max {
|
||||
int32_t ne00;
|
||||
int32_t ne01;
|
||||
int32_t ne02;
|
||||
uint64_t nb01;
|
||||
uint64_t nb02;
|
||||
uint64_t nb03;
|
||||
int32_t ne11;
|
||||
int32_t ne12;
|
||||
int32_t ne13;
|
||||
uint64_t nb11;
|
||||
uint64_t nb12;
|
||||
uint64_t nb13;
|
||||
uint64_t nb1;
|
||||
uint64_t nb2;
|
||||
uint64_t nb3;
|
||||
float scale;
|
||||
float max_bias;
|
||||
float m0;
|
||||
float m1;
|
||||
int32_t n_head_log2;
|
||||
};
|
||||
|
||||
// Row softmax for score matrices. DS4 uses it in the literal one-compressor-row
|
||||
// path where preserving the original graph operation boundary avoids drift.
|
||||
template<typename T>
|
||||
kernel void kernel_soft_max(
|
||||
constant ds4_metal_args_soft_max & args,
|
||||
device const char * src0,
|
||||
device const char * src1,
|
||||
device const char * src2,
|
||||
device char * dst,
|
||||
threadgroup float * buf [[threadgroup(0)]],
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
uint3 tpitg[[thread_position_in_threadgroup]],
|
||||
uint sgitg[[simdgroup_index_in_threadgroup]],
|
||||
uint tiisg[[thread_index_in_simdgroup]],
|
||||
uint3 tptg[[threads_per_threadgroup]]) {
|
||||
const int32_t i03 = tgpig.z;
|
||||
const int32_t i02 = tgpig.y;
|
||||
const int32_t i01 = tgpig.x;
|
||||
|
||||
const int32_t i13 = i03%args.ne13;
|
||||
const int32_t i12 = i02%args.ne12;
|
||||
const int32_t i11 = i01;
|
||||
|
||||
device const float * psrc0 = (device const float *) (src0 + i01*args.nb01 + i02*args.nb02 + i03*args.nb03);
|
||||
device const T * pmask = src1 != src0 ? (device const T * ) (src1 + i11*args.nb11 + i12*args.nb12 + i13*args.nb13) : nullptr;
|
||||
device const float * psrc2 = src2 != src0 ? (device const float *) (src2) : nullptr;
|
||||
device float * pdst = (device float *) (dst + i01*args.nb1 + i02*args.nb2 + i03*args.nb3);
|
||||
|
||||
float slope = 1.0f;
|
||||
|
||||
if (args.max_bias > 0.0f) {
|
||||
const int32_t h = i02;
|
||||
|
||||
const float base = h < args.n_head_log2 ? args.m0 : args.m1;
|
||||
const int exp = h < args.n_head_log2 ? h + 1 : 2*(h - args.n_head_log2) + 1;
|
||||
|
||||
slope = pow(base, exp);
|
||||
}
|
||||
|
||||
float lmax = psrc2 ? psrc2[i02] : -INFINITY;
|
||||
|
||||
for (int i00 = tpitg.x; i00 < args.ne00; i00 += tptg.x) {
|
||||
lmax = MAX(lmax, psrc0[i00]*args.scale + (pmask ? slope*pmask[i00] : 0.0f));
|
||||
}
|
||||
|
||||
float max_val = simd_max(lmax);
|
||||
if (tptg.x > N_SIMDWIDTH) {
|
||||
if (sgitg == 0) {
|
||||
buf[tiisg] = -INFINITY;
|
||||
}
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
if (tiisg == 0) {
|
||||
buf[sgitg] = max_val;
|
||||
}
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
max_val = buf[tiisg];
|
||||
max_val = simd_max(max_val);
|
||||
}
|
||||
|
||||
float lsum = 0.0f;
|
||||
for (int i00 = tpitg.x; i00 < args.ne00; i00 += tptg.x) {
|
||||
const float exp_psrc0 = exp((psrc0[i00]*args.scale + (pmask ? slope*pmask[i00] : 0.0f)) - max_val);
|
||||
lsum += exp_psrc0;
|
||||
pdst[i00] = exp_psrc0;
|
||||
}
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_none);
|
||||
|
||||
float sum = simd_sum(lsum);
|
||||
|
||||
if (tptg.x > N_SIMDWIDTH) {
|
||||
if (sgitg == 0) {
|
||||
buf[tiisg] = 0.0f;
|
||||
}
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
if (tiisg == 0) {
|
||||
buf[sgitg] = sum;
|
||||
}
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
sum = buf[tiisg];
|
||||
sum = simd_sum(sum);
|
||||
}
|
||||
|
||||
if (psrc2) {
|
||||
sum += exp(psrc2[i02] - max_val);
|
||||
}
|
||||
|
||||
const float inv_sum = 1.0f/sum;
|
||||
|
||||
for (int i00 = tpitg.x; i00 < args.ne00; i00 += tptg.x) {
|
||||
pdst[i00] *= inv_sum;
|
||||
}
|
||||
}
|
||||
|
||||
// Vectorized float4 row softmax for contiguous score rows whose length is a
|
||||
// multiple of four; used by the same DS4 compressor/indexer graph path.
|
||||
template<typename T>
|
||||
kernel void kernel_soft_max_4(
|
||||
constant ds4_metal_args_soft_max & args,
|
||||
device const char * src0,
|
||||
device const char * src1,
|
||||
device const char * src2,
|
||||
device char * dst,
|
||||
threadgroup float * buf [[threadgroup(0)]],
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
uint3 tpitg[[thread_position_in_threadgroup]],
|
||||
uint sgitg[[simdgroup_index_in_threadgroup]],
|
||||
uint tiisg[[thread_index_in_simdgroup]],
|
||||
uint3 tptg[[threads_per_threadgroup]]) {
|
||||
const int32_t i03 = tgpig.z;
|
||||
const int32_t i02 = tgpig.y;
|
||||
const int32_t i01 = tgpig.x;
|
||||
|
||||
const int32_t i13 = i03%args.ne13;
|
||||
const int32_t i12 = i02%args.ne12;
|
||||
const int32_t i11 = i01;
|
||||
|
||||
device const float4 * psrc4 = (device const float4 *) (src0 + i01*args.nb01 + i02*args.nb02 + i03*args.nb03);
|
||||
device const T * pmask = src1 != src0 ? (device const T * ) (src1 + i11*args.nb11 + i12*args.nb12 + i13*args.nb13) : nullptr;
|
||||
device const float * psrc2 = src2 != src0 ? (device const float * ) (src2) : nullptr;
|
||||
device float4 * pdst4 = (device float4 *) (dst + i01*args.nb1 + i02*args.nb2 + i03*args.nb3);
|
||||
|
||||
float slope = 1.0f;
|
||||
|
||||
if (args.max_bias > 0.0f) {
|
||||
const int32_t h = i02;
|
||||
|
||||
const float base = h < args.n_head_log2 ? args.m0 : args.m1;
|
||||
const int exp = h < args.n_head_log2 ? h + 1 : 2*(h - args.n_head_log2) + 1;
|
||||
|
||||
slope = pow(base, exp);
|
||||
}
|
||||
|
||||
float4 lmax4 = psrc2 ? psrc2[i02] : -INFINITY;
|
||||
|
||||
for (int i00 = tpitg.x; i00 < args.ne00/4; i00 += tptg.x) {
|
||||
lmax4 = fmax(lmax4, psrc4[i00]*args.scale + (float4)((pmask ? slope*pmask[i00] : 0.0f)));
|
||||
}
|
||||
|
||||
const float lmax = MAX(MAX(lmax4[0], lmax4[1]), MAX(lmax4[2], lmax4[3]));
|
||||
|
||||
float max_val = simd_max(lmax);
|
||||
if (tptg.x > N_SIMDWIDTH) {
|
||||
if (sgitg == 0) {
|
||||
buf[tiisg] = -INFINITY;
|
||||
}
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
if (tiisg == 0) {
|
||||
buf[sgitg] = max_val;
|
||||
}
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
max_val = buf[tiisg];
|
||||
max_val = simd_max(max_val);
|
||||
}
|
||||
|
||||
float4 lsum4 = 0.0f;
|
||||
for (int i00 = tpitg.x; i00 < args.ne00/4; i00 += tptg.x) {
|
||||
const float4 exp_psrc4 = exp((psrc4[i00]*args.scale + (float4)((pmask ? slope*pmask[i00] : 0.0f))) - max_val);
|
||||
lsum4 += exp_psrc4;
|
||||
pdst4[i00] = exp_psrc4;
|
||||
}
|
||||
|
||||
const float lsum = lsum4[0] + lsum4[1] + lsum4[2] + lsum4[3];
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_none);
|
||||
|
||||
float sum = simd_sum(lsum);
|
||||
|
||||
if (tptg.x > N_SIMDWIDTH) {
|
||||
if (sgitg == 0) {
|
||||
buf[tiisg] = 0.0f;
|
||||
}
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
if (tiisg == 0) {
|
||||
buf[sgitg] = sum;
|
||||
}
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
sum = buf[tiisg];
|
||||
sum = simd_sum(sum);
|
||||
}
|
||||
|
||||
if (psrc2) {
|
||||
sum += exp(psrc2[i02] - max_val);
|
||||
}
|
||||
|
||||
const float inv_sum = 1.0f/sum;
|
||||
|
||||
for (int i00 = tpitg.x; i00 < args.ne00/4; i00 += tptg.x) {
|
||||
pdst4[i00] *= inv_sum;
|
||||
}
|
||||
}
|
||||
|
||||
typedef decltype(kernel_soft_max<float>) kernel_soft_max_t;
|
||||
typedef decltype(kernel_soft_max_4<float4>) kernel_soft_max_4_t;
|
||||
|
||||
// Host-visible F32 softmax variants used by compressor pooling.
|
||||
template [[host_name("kernel_soft_max_f32")]] kernel kernel_soft_max_t kernel_soft_max<float>;
|
||||
template [[host_name("kernel_soft_max_f32_4")]] kernel kernel_soft_max_4_t kernel_soft_max_4<float4>;
|
||||
102
metal/sum_rows.metal
Normal file
102
metal/sum_rows.metal
Normal file
@@ -0,0 +1,102 @@
|
||||
// DS4 Metal row-sum kernel.
|
||||
|
||||
#define FC_SUM_ROWS 1400
|
||||
|
||||
#define OP_SUM_ROWS_NUM_SUM_ROWS 10
|
||||
#define OP_SUM_ROWS_NUM_MEAN 11
|
||||
|
||||
struct ds4_metal_args_sum_rows {
|
||||
int64_t ne00;
|
||||
int64_t ne01;
|
||||
int64_t ne02;
|
||||
int64_t ne03;
|
||||
uint64_t nb00;
|
||||
uint64_t nb01;
|
||||
uint64_t nb02;
|
||||
uint64_t nb03;
|
||||
int64_t ne0;
|
||||
int64_t ne1;
|
||||
int64_t ne2;
|
||||
int64_t ne3;
|
||||
uint64_t nb0;
|
||||
uint64_t nb1;
|
||||
uint64_t nb2;
|
||||
uint64_t nb3;
|
||||
};
|
||||
|
||||
static inline float sum(float x) {
|
||||
return x;
|
||||
}
|
||||
|
||||
static inline float sum(float4 x) {
|
||||
return x[0] + x[1] + x[2] + x[3];
|
||||
}
|
||||
|
||||
constant short FC_sum_rows_op [[function_constant(FC_SUM_ROWS + 0)]];
|
||||
|
||||
// Reduces each row to a sum or mean. DS4 mainly uses the sum form to preserve
|
||||
// the compressor-pooling graph boundary in the single-compressor-row case.
|
||||
template <typename T0, typename T>
|
||||
kernel void kernel_sum_rows_impl(
|
||||
constant ds4_metal_args_sum_rows & args,
|
||||
device const char * src0,
|
||||
device char * dst,
|
||||
threadgroup char * shmem [[threadgroup(0)]],
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
ushort3 tpitg[[thread_position_in_threadgroup]],
|
||||
ushort sgitg[[simdgroup_index_in_threadgroup]],
|
||||
ushort tiisg[[thread_index_in_simdgroup]],
|
||||
ushort3 ntg[[threads_per_threadgroup]]) {
|
||||
#define FC_OP FC_sum_rows_op
|
||||
|
||||
const int i3 = tgpig.z;
|
||||
const int i2 = tgpig.y;
|
||||
const int i1 = tgpig.x;
|
||||
|
||||
threadgroup T0 * shmem_t = (threadgroup T0 *) shmem;
|
||||
|
||||
if (sgitg == 0) {
|
||||
shmem_t[tiisg] = 0.0f;
|
||||
}
|
||||
|
||||
device const T0 * src_row = (device const T0 *) (src0 + i1*args.nb01 + i2*args.nb02 + i3*args.nb03);
|
||||
device T * dst_row = (device T *) (dst + i1*args.nb1 + i2*args.nb2 + i3*args.nb3);
|
||||
|
||||
T0 sumf = T0(0.0f);
|
||||
|
||||
for (int64_t i0 = tpitg.x; i0 < args.ne00; i0 += ntg.x) {
|
||||
sumf += src_row[i0];
|
||||
}
|
||||
|
||||
sumf = simd_sum(sumf);
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
if (tiisg == 0) {
|
||||
shmem_t[sgitg] = sumf;
|
||||
}
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
sumf = shmem_t[tiisg];
|
||||
sumf = simd_sum(sumf);
|
||||
|
||||
if (tpitg.x == 0) {
|
||||
if (FC_OP == OP_SUM_ROWS_NUM_MEAN) {
|
||||
if (is_same<float4, T0>::value) {
|
||||
dst_row[0] = sum(sumf) / (4*args.ne00);
|
||||
} else {
|
||||
dst_row[0] = sum(sumf) / args.ne00;
|
||||
}
|
||||
} else {
|
||||
dst_row[0] = sum(sumf);
|
||||
}
|
||||
}
|
||||
|
||||
#undef FC_OP
|
||||
}
|
||||
|
||||
typedef decltype(kernel_sum_rows_impl<float, float>) kernel_sum_rows_t;
|
||||
|
||||
// Host-visible F32 row reduction used by compressor pooling.
|
||||
template [[host_name("kernel_sum_rows_f32_f32")]] kernel kernel_sum_rows_t kernel_sum_rows_impl<float, float>;
|
||||
312
metal/unary.metal
Normal file
312
metal/unary.metal
Normal file
@@ -0,0 +1,312 @@
|
||||
#define FC_UNARY 1200
|
||||
|
||||
#define OP_UNARY_NUM_SCALE 10
|
||||
#define OP_UNARY_NUM_FILL 11
|
||||
#define OP_UNARY_NUM_CLAMP 12
|
||||
#define OP_UNARY_NUM_SQR 13
|
||||
#define OP_UNARY_NUM_SQRT 14
|
||||
#define OP_UNARY_NUM_SIN 15
|
||||
#define OP_UNARY_NUM_COS 16
|
||||
#define OP_UNARY_NUM_LOG 17
|
||||
#define OP_UNARY_NUM_LEAKY_RELU 18
|
||||
|
||||
#define OP_UNARY_NUM_TANH 100
|
||||
#define OP_UNARY_NUM_RELU 101
|
||||
#define OP_UNARY_NUM_SIGMOID 102
|
||||
#define OP_UNARY_NUM_GELU 103
|
||||
#define OP_UNARY_NUM_GELU_ERF 104
|
||||
#define OP_UNARY_NUM_GELU_QUICK 105
|
||||
#define OP_UNARY_NUM_SILU 106
|
||||
#define OP_UNARY_NUM_ELU 107
|
||||
#define OP_UNARY_NUM_NEG 108
|
||||
#define OP_UNARY_NUM_ABS 109
|
||||
#define OP_UNARY_NUM_SGN 110
|
||||
#define OP_UNARY_NUM_STEP 111
|
||||
#define OP_UNARY_NUM_HARDSWISH 112
|
||||
#define OP_UNARY_NUM_HARDSIGMOID 113
|
||||
#define OP_UNARY_NUM_EXP 114
|
||||
#define OP_UNARY_NUM_SOFTPLUS 115
|
||||
#define OP_UNARY_NUM_EXPM1 116
|
||||
#define OP_UNARY_NUM_FLOOR 117
|
||||
#define OP_UNARY_NUM_CEIL 118
|
||||
#define OP_UNARY_NUM_ROUND 119
|
||||
#define OP_UNARY_NUM_TRUNC 120
|
||||
#define OP_UNARY_NUM_XIELU 121
|
||||
|
||||
struct ds4_metal_args_unary {
|
||||
int32_t ne00;
|
||||
int32_t ne01;
|
||||
int32_t ne02;
|
||||
int32_t ne03;
|
||||
uint64_t nb00;
|
||||
uint64_t nb01;
|
||||
uint64_t nb02;
|
||||
uint64_t nb03;
|
||||
int32_t ne0;
|
||||
int32_t ne1;
|
||||
int32_t ne2;
|
||||
int32_t ne3;
|
||||
uint64_t nb0;
|
||||
uint64_t nb1;
|
||||
uint64_t nb2;
|
||||
uint64_t nb3;
|
||||
float slope;
|
||||
float scale;
|
||||
float bias;
|
||||
float val;
|
||||
float min;
|
||||
float max;
|
||||
};
|
||||
|
||||
constant float GELU_COEF_A = 0.044715f;
|
||||
constant float GELU_QUICK_COEF = -1.702f;
|
||||
constant float SQRT_2_OVER_PI = 0.79788456080286535587989211986876f;
|
||||
constant float SQRT_2_INV = 0.70710678118654752440084436210484f;
|
||||
|
||||
// based on Abramowitz and Stegun formula 7.1.26 or similar Hastings' approximation
|
||||
// ref: https://www.johndcook.com/blog/python_erf/
|
||||
constant float p_erf = 0.3275911f;
|
||||
constant float a1_erf = 0.254829592f;
|
||||
constant float a2_erf = -0.284496736f;
|
||||
constant float a3_erf = 1.421413741f;
|
||||
constant float a4_erf = -1.453152027f;
|
||||
constant float a5_erf = 1.061405429f;
|
||||
|
||||
template<typename T>
|
||||
inline T erf_approx(T x) {
|
||||
T sign_x = sign(x);
|
||||
x = fabs(x);
|
||||
T t = 1.0f / (1.0f + p_erf * x);
|
||||
T y = 1.0f - (((((a5_erf * t + a4_erf) * t) + a3_erf) * t + a2_erf) * t + a1_erf) * t * exp(-x * x);
|
||||
return sign_x * y;
|
||||
}
|
||||
|
||||
template<typename T> T elu_approx(T x);
|
||||
|
||||
template<> inline float elu_approx<float>(float x) {
|
||||
return (x > 0.f) ? x : (exp(x) - 1);
|
||||
}
|
||||
|
||||
template<> inline float4 elu_approx<float4>(float4 x) {
|
||||
float4 res;
|
||||
|
||||
res[0] = (x[0] > 0.0f) ? x[0] : (exp(x[0]) - 1.0f);
|
||||
res[1] = (x[1] > 0.0f) ? x[1] : (exp(x[1]) - 1.0f);
|
||||
res[2] = (x[2] > 0.0f) ? x[2] : (exp(x[2]) - 1.0f);
|
||||
res[3] = (x[3] > 0.0f) ? x[3] : (exp(x[3]) - 1.0f);
|
||||
|
||||
return res;
|
||||
}
|
||||
|
||||
constant short FC_unary_op [[function_constant(FC_UNARY + 0)]];
|
||||
constant bool FC_unary_cnt[[function_constant(FC_UNARY + 1)]];
|
||||
|
||||
// Generic unary elementwise op selected by function constant. DS4 only uses a
|
||||
// small subset in inference, mainly sigmoid, SiLU, softplus, sqrt, clamp,
|
||||
// scale, and fill.
|
||||
template <typename T0, typename T, typename TC>
|
||||
kernel void kernel_unary_impl(
|
||||
constant ds4_metal_args_unary & args,
|
||||
device const char * src0,
|
||||
device char * dst,
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
ushort3 tpitg[[thread_position_in_threadgroup]],
|
||||
ushort3 ntg[[threads_per_threadgroup]]) {
|
||||
#define FC_OP FC_unary_op
|
||||
#define FC_CNT FC_unary_cnt
|
||||
|
||||
device const T0 * src0_ptr;
|
||||
device T * dst_ptr;
|
||||
|
||||
int i0;
|
||||
|
||||
if (FC_CNT) {
|
||||
i0 = tgpig.x;
|
||||
|
||||
src0_ptr = (device const T0 *) (src0);
|
||||
dst_ptr = (device T *) (dst);
|
||||
} else {
|
||||
const int i03 = tgpig.z;
|
||||
const int i02 = tgpig.y;
|
||||
const int k0 = tgpig.x/args.ne01;
|
||||
const int i01 = tgpig.x - k0*args.ne01;
|
||||
|
||||
i0 = k0*ntg.x + tpitg.x;
|
||||
|
||||
src0_ptr = (device const T0 *) (src0 + i03*args.nb03 + i02*args.nb02 + i01*args.nb01);
|
||||
dst_ptr = (device T *) (dst + i03*args.nb3 + i02*args.nb2 + i01*args.nb1 );
|
||||
}
|
||||
|
||||
{
|
||||
if (!FC_CNT) {
|
||||
if (i0 >= args.ne0) {
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
||||
const TC x = (TC) src0_ptr[i0];
|
||||
|
||||
if (FC_OP == OP_UNARY_NUM_SCALE) {
|
||||
dst_ptr[i0] = (T) (args.scale * x + args.bias);
|
||||
}
|
||||
|
||||
if (FC_OP == OP_UNARY_NUM_FILL) {
|
||||
dst_ptr[i0] = (T) args.val;
|
||||
}
|
||||
|
||||
if (FC_OP == OP_UNARY_NUM_CLAMP) {
|
||||
dst_ptr[i0] = (T) clamp(x, args.min, args.max);
|
||||
}
|
||||
|
||||
if (FC_OP == OP_UNARY_NUM_SQR) {
|
||||
dst_ptr[i0] = (T) (x * x);
|
||||
}
|
||||
|
||||
if (FC_OP == OP_UNARY_NUM_SQRT) {
|
||||
dst_ptr[i0] = (T) sqrt(x);
|
||||
}
|
||||
|
||||
if (FC_OP == OP_UNARY_NUM_SIN) {
|
||||
dst_ptr[i0] = (T) sin(x);
|
||||
}
|
||||
|
||||
if (FC_OP == OP_UNARY_NUM_COS) {
|
||||
dst_ptr[i0] = (T) cos(x);
|
||||
}
|
||||
|
||||
if (FC_OP == OP_UNARY_NUM_LOG) {
|
||||
dst_ptr[i0] = (T) log(x);
|
||||
}
|
||||
|
||||
if (FC_OP == OP_UNARY_NUM_LEAKY_RELU) {
|
||||
dst_ptr[i0] = (T) (TC(x > 0)*x + TC(x <= 0)*(x * args.slope));
|
||||
}
|
||||
|
||||
if (FC_OP == OP_UNARY_NUM_TANH) {
|
||||
dst_ptr[i0] = (T) precise::tanh(x);
|
||||
}
|
||||
|
||||
if (FC_OP == OP_UNARY_NUM_RELU) {
|
||||
dst_ptr[i0] = (T) fmax(0, x);
|
||||
}
|
||||
|
||||
if (FC_OP == OP_UNARY_NUM_SIGMOID) {
|
||||
dst_ptr[i0] = (T) (1 / (1 + exp(-x)));
|
||||
}
|
||||
|
||||
if (FC_OP == OP_UNARY_NUM_GELU) {
|
||||
dst_ptr[i0] = (T) (0.5*x*(1 + precise::tanh(SQRT_2_OVER_PI*x*(1 + GELU_COEF_A*x*x))));
|
||||
}
|
||||
|
||||
if (FC_OP == OP_UNARY_NUM_GELU_ERF) {
|
||||
dst_ptr[i0] = (T) (0.5*x*(1 + erf_approx(SQRT_2_INV*x)));
|
||||
}
|
||||
|
||||
if (FC_OP == OP_UNARY_NUM_GELU_QUICK) {
|
||||
dst_ptr[i0] = (T) (x * (1/(1 + exp(GELU_QUICK_COEF*x))));
|
||||
}
|
||||
|
||||
if (FC_OP == OP_UNARY_NUM_SILU) {
|
||||
dst_ptr[i0] = (T) (x / (1 + exp(-x)));
|
||||
}
|
||||
|
||||
if (FC_OP == OP_UNARY_NUM_ELU) {
|
||||
dst_ptr[i0] = (T) elu_approx(x);
|
||||
}
|
||||
|
||||
if (FC_OP == OP_UNARY_NUM_NEG) {
|
||||
dst_ptr[i0] = (T) -x;
|
||||
}
|
||||
|
||||
if (FC_OP == OP_UNARY_NUM_ABS) {
|
||||
dst_ptr[i0] = (T) fabs(x);
|
||||
}
|
||||
|
||||
if (FC_OP == OP_UNARY_NUM_SGN) {
|
||||
dst_ptr[i0] = T(x > 0) - T(x < 0);
|
||||
}
|
||||
|
||||
if (FC_OP == OP_UNARY_NUM_STEP) {
|
||||
dst_ptr[i0] = T(x > 0);
|
||||
}
|
||||
|
||||
if (FC_OP == OP_UNARY_NUM_HARDSWISH) {
|
||||
dst_ptr[i0] = (T) (x * fmax(0, fmin(1, x/6 + 0.5)));
|
||||
}
|
||||
|
||||
if (FC_OP == OP_UNARY_NUM_HARDSIGMOID) {
|
||||
dst_ptr[i0] = (T) fmax(0, fmin(1, x/6 + 0.5));
|
||||
}
|
||||
|
||||
if (FC_OP == OP_UNARY_NUM_EXP) {
|
||||
dst_ptr[i0] = (T) exp(x);
|
||||
}
|
||||
|
||||
if (FC_OP == OP_UNARY_NUM_SOFTPLUS) {
|
||||
dst_ptr[i0] = (T) select(log(1 + exp(x)), x, x > 20);
|
||||
}
|
||||
|
||||
if (FC_OP == OP_UNARY_NUM_EXPM1) {
|
||||
// Metal target profiles used here do not all expose expm1(); this
|
||||
// generic unary branch is not used by the DS4 inference graph.
|
||||
dst_ptr[i0] = (T) (exp(x) - 1);
|
||||
}
|
||||
|
||||
if (FC_OP == OP_UNARY_NUM_FLOOR) {
|
||||
dst_ptr[i0] = (T) floor(x);
|
||||
}
|
||||
|
||||
if (FC_OP == OP_UNARY_NUM_CEIL) {
|
||||
dst_ptr[i0] = (T) ceil(x);
|
||||
}
|
||||
|
||||
if (FC_OP == OP_UNARY_NUM_ROUND) {
|
||||
dst_ptr[i0] = (T) round(x);
|
||||
}
|
||||
|
||||
if (FC_OP == OP_UNARY_NUM_TRUNC) {
|
||||
dst_ptr[i0] = (T) trunc(x);
|
||||
}
|
||||
|
||||
if (FC_OP == OP_UNARY_NUM_XIELU) {
|
||||
const TC xi = x;
|
||||
const TC gate = TC(xi > TC(0.0f));
|
||||
const TC clamped = fmin(xi, TC(args.val));
|
||||
const TC y_pos = TC(args.scale) * xi * xi + TC(args.bias) * xi;
|
||||
const TC y_neg = (exp(clamped) - TC(1.0f) - xi) * TC(args.slope) + TC(args.bias) * xi;
|
||||
dst_ptr[i0] = (T) (gate * y_pos + (TC(1.0f) - gate) * y_neg);
|
||||
}
|
||||
}
|
||||
|
||||
#undef FC_OP
|
||||
#undef FC_CNT
|
||||
}
|
||||
|
||||
typedef decltype(kernel_unary_impl<float, float, float>) kernel_unary_t;
|
||||
|
||||
// Decode router probability transform. The generic path applies softplus and
|
||||
// sqrt as two elementwise kernels; DS4 decode always transforms one 256-wide
|
||||
// expert-logit row, so this vectorized kernel does both in one pass.
|
||||
kernel void kernel_dsv4_softplus_sqrt_f32_4(
|
||||
constant ds4_metal_args_unary & args,
|
||||
device const char *src,
|
||||
device char *dst,
|
||||
uint3 tgpig [[threadgroup_position_in_grid]],
|
||||
ushort3 tpitg [[thread_position_in_threadgroup]],
|
||||
ushort3 ntg [[threads_per_threadgroup]]) {
|
||||
const int k0 = tgpig.x/args.ne01;
|
||||
const int i01 = tgpig.x - k0*args.ne01;
|
||||
const int i0 = k0*ntg.x + tpitg.x;
|
||||
if (i0 >= args.ne0) return;
|
||||
|
||||
device const float4 *s = (device const float4 *)(src + i01*args.nb01);
|
||||
device float4 *d = (device float4 *)(dst + i01*args.nb1);
|
||||
const float4 x = s[i0];
|
||||
const float4 sp = select(log(1.0f + exp(x)), x, x > 20.0f);
|
||||
d[i0] = sqrt(sp);
|
||||
}
|
||||
|
||||
// Host-visible unary variants. Function constants select the actual DS4 op.
|
||||
template [[host_name("kernel_unary_f32_f32")]] kernel kernel_unary_t kernel_unary_impl<float, float, float>;
|
||||
template [[host_name("kernel_unary_f32_f32_4")]] kernel kernel_unary_t kernel_unary_impl<float4, float4, float4>;
|
||||
template [[host_name("kernel_unary_f16_f16")]] kernel kernel_unary_t kernel_unary_impl<half, half, float>;
|
||||
Reference in New Issue
Block a user