3918 lines
150 KiB
Metal
3918 lines
150 KiB
Metal
// DS4 Metal matvec kernels used by generation.
|
|
|
|
constant short FC_mul_mv_nsg [[function_constant(FC_MUL_MV + 0)]];
|
|
constant short FC_mul_mv_nxpsg [[function_constant(FC_MUL_MV + 1)]];
|
|
|
|
// Round to the nearest bfloat16 value in f32 storage (DeepSeek V4.1 keeps its
|
|
// activations at bf16 precision between operators).
|
|
static inline float ds4_bf16_round(float v) {
|
|
uint bits = as_type<uint>(v);
|
|
if ((bits & 0x7f800000u) != 0x7f800000u) bits += 0x7fffu + ((bits >> 16u) & 1u);
|
|
return as_type<float>(bits & 0xffff0000u);
|
|
}
|
|
static inline float4 ds4_bf16_round(float4 v) {
|
|
uint4 bits = as_type<uint4>(v);
|
|
const bool4 finite = (bits & 0x7f800000u) != 0x7f800000u;
|
|
bits += select(uint4(0), uint4(0x7fffu) + ((bits >> 16u) & 1u), finite);
|
|
return as_type<float4>(bits & 0xffff0000u);
|
|
}
|
|
|
|
struct ds4_metal_args_mul_mv {
|
|
int ne00;
|
|
int ne01;
|
|
int ne02;
|
|
ulong nb00;
|
|
ulong nb01;
|
|
ulong nb02;
|
|
ulong nb03;
|
|
int ne10;
|
|
int ne11;
|
|
int ne12;
|
|
ulong nb10;
|
|
ulong nb11;
|
|
ulong nb12;
|
|
ulong nb13;
|
|
int ne0;
|
|
int ne1;
|
|
int nr0;
|
|
short r2;
|
|
short r3;
|
|
};
|
|
|
|
struct ds4_metal_args_compressor_pair_store {
|
|
uint32_t width;
|
|
uint32_t ratio;
|
|
uint32_t pos;
|
|
uint32_t ape_type;
|
|
};
|
|
|
|
struct ds4_metal_args_mul_mm {
|
|
int32_t ne00;
|
|
int32_t ne02;
|
|
uint64_t nb01;
|
|
uint64_t nb02;
|
|
uint64_t nb03;
|
|
int32_t ne12;
|
|
uint64_t nb10;
|
|
uint64_t nb11;
|
|
uint64_t nb12;
|
|
uint64_t nb13;
|
|
int32_t ne0;
|
|
int32_t ne1;
|
|
int16_t r2;
|
|
int16_t r3;
|
|
};
|
|
|
|
struct ds4_metal_args_mul_mv_ext {
|
|
int32_t ne00;
|
|
int32_t ne01;
|
|
int32_t ne02;
|
|
uint64_t nb00;
|
|
uint64_t nb01;
|
|
uint64_t nb02;
|
|
uint64_t nb03;
|
|
int32_t ne10;
|
|
int32_t ne11;
|
|
int32_t ne12;
|
|
uint64_t nb10;
|
|
uint64_t nb11;
|
|
uint64_t nb12;
|
|
uint64_t nb13;
|
|
int32_t ne0;
|
|
int32_t ne1;
|
|
int16_t r2;
|
|
int16_t r3;
|
|
};
|
|
|
|
template<short NR0, bool ROUND = false>
|
|
static inline void helper_mv_reduce_and_write(
|
|
device float * dst_f32,
|
|
float sumf[NR0],
|
|
const int r0,
|
|
const int ne01,
|
|
ushort tiisg,
|
|
ushort sgitg,
|
|
threadgroup char * shmem) {
|
|
constexpr short NW = N_SIMDWIDTH;
|
|
|
|
threadgroup float * shmem_f32[NR0];
|
|
|
|
for (short row = 0; row < NR0; ++row) {
|
|
shmem_f32[row] = (threadgroup float *) shmem + NW*row;
|
|
|
|
if (sgitg == 0) {
|
|
shmem_f32[row][tiisg] = 0.0f;
|
|
}
|
|
|
|
sumf[row] = simd_sum(sumf[row]);
|
|
}
|
|
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
for (short row = 0; row < NR0; ++row) {
|
|
if (tiisg == 0) {
|
|
shmem_f32[row][sgitg] = sumf[row];
|
|
}
|
|
}
|
|
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
for (short row = 0; row < NR0 && r0 + row < ne01; ++row) {
|
|
float tot = simd_sum(shmem_f32[row][tiisg]);
|
|
|
|
if (tiisg == 0 && sgitg == 0) {
|
|
dst_f32[r0 + row] = ROUND ? ds4_bf16_round(tot) : tot;
|
|
}
|
|
}
|
|
}
|
|
|
|
// Rope folded into a matvec: the last 64 columns of every `width`-wide head
|
|
// rotate at position `start` (ROPE_IN on the activation as it loads, after
|
|
// rounding it to bf16; ROPE_OUT on the bf16 outputs before they store).
|
|
struct ds4_metal_args_mv_rope {
|
|
uint width;
|
|
uint start;
|
|
uint inverse;
|
|
float frequencies[32];
|
|
};
|
|
|
|
template<short NR0>
|
|
static inline void helper_mv_reduce_and_write_rope(
|
|
device float * dst_f32,
|
|
float sumf[NR0],
|
|
const int r0,
|
|
const int ne01,
|
|
ushort tiisg,
|
|
ushort sgitg,
|
|
threadgroup char * shmem,
|
|
constant ds4_metal_args_mv_rope & rope) {
|
|
static_assert(NR0 % 2 == 0, "rope pairs are adjacent output rows");
|
|
constexpr short NW = N_SIMDWIDTH;
|
|
|
|
threadgroup float * shmem_f32[NR0];
|
|
|
|
for (short row = 0; row < NR0; ++row) {
|
|
shmem_f32[row] = (threadgroup float *) shmem + NW*row;
|
|
|
|
if (sgitg == 0) {
|
|
shmem_f32[row][tiisg] = 0.0f;
|
|
}
|
|
|
|
sumf[row] = simd_sum(sumf[row]);
|
|
}
|
|
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
for (short row = 0; row < NR0; ++row) {
|
|
if (tiisg == 0) {
|
|
shmem_f32[row][sgitg] = sumf[row];
|
|
}
|
|
}
|
|
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
float tot[NR0];
|
|
for (short row = 0; row < NR0; ++row) {
|
|
tot[row] = ds4_bf16_round(simd_sum(shmem_f32[row][tiisg]));
|
|
}
|
|
|
|
if (tiisg == 0 && sgitg == 0) {
|
|
for (short p = 0; p < NR0 && r0 + p + 1 < ne01; p += 2) {
|
|
const uint col = (uint)(r0 + p) % rope.width;
|
|
if (col >= rope.width - 64u) {
|
|
const float theta = float(rope.start) * rope.frequencies[(col - (rope.width - 64u)) >> 1u];
|
|
const float c = precise::cos(theta);
|
|
const float s = rope.inverse ? -precise::sin(theta) : precise::sin(theta);
|
|
const float re = tot[p], im = tot[p + 1];
|
|
dst_f32[r0 + p] = ds4_bf16_round(fma(re, c, -(im * s)));
|
|
dst_f32[r0 + p + 1] = ds4_bf16_round(fma(im, c, re * s));
|
|
} else {
|
|
dst_f32[r0 + p] = tot[p];
|
|
dst_f32[r0 + p + 1] = tot[p + 1];
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// The HC expand (kernel_dsv4_hc_expand4, one token, four streams) applied to
|
|
// each bf16 output as it is written: the row joins `add` when has_add and
|
|
// spreads into the streams with the same per-stream accumulation order.
|
|
struct ds4_metal_args_mv_hc_expand4 {
|
|
uint n_embd;
|
|
uint has_add;
|
|
uint split_stride; /* floats between the rows' post/comb weights */
|
|
};
|
|
|
|
template<short NR0>
|
|
static inline void helper_mv_reduce_and_write_hc_expand4(
|
|
device float * dst_f32,
|
|
float sumf[NR0],
|
|
const int r0,
|
|
const int ne01,
|
|
ushort tiisg,
|
|
ushort sgitg,
|
|
threadgroup char * shmem,
|
|
constant ds4_metal_args_mv_hc_expand4 & hc,
|
|
device const float * add,
|
|
device const float * residual,
|
|
device const float * post,
|
|
device const float * comb) {
|
|
constexpr short NW = N_SIMDWIDTH;
|
|
|
|
threadgroup float * shmem_f32[NR0];
|
|
|
|
for (short row = 0; row < NR0; ++row) {
|
|
shmem_f32[row] = (threadgroup float *) shmem + NW*row;
|
|
|
|
if (sgitg == 0) {
|
|
shmem_f32[row][tiisg] = 0.0f;
|
|
}
|
|
|
|
sumf[row] = simd_sum(sumf[row]);
|
|
}
|
|
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
for (short row = 0; row < NR0; ++row) {
|
|
if (tiisg == 0) {
|
|
shmem_f32[row][sgitg] = sumf[row];
|
|
}
|
|
}
|
|
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
for (short row = 0; row < NR0 && r0 + row < ne01; ++row) {
|
|
float block_v = ds4_bf16_round(simd_sum(shmem_f32[row][tiisg]));
|
|
|
|
if (tiisg == 0 && sgitg == 0) {
|
|
const int d = r0 + row;
|
|
if (hc.has_add) {
|
|
block_v += add[d];
|
|
block_v = ds4_bf16_round(block_v);
|
|
}
|
|
const float r0v = residual[d];
|
|
const float r1v = residual[d + hc.n_embd];
|
|
const float r2v = residual[d + 2*hc.n_embd];
|
|
const float r3v = residual[d + 3*hc.n_embd];
|
|
|
|
for (int dst_hc = 0; dst_hc < 4; ++dst_hc) {
|
|
float acc = block_v * post[dst_hc];
|
|
|
|
acc += comb[dst_hc + 0*4] * r0v;
|
|
acc += comb[dst_hc + 1*4] * r1v;
|
|
acc += comb[dst_hc + 2*4] * r2v;
|
|
acc += comb[dst_hc + 3*4] * r3v;
|
|
|
|
dst_f32[d + dst_hc*hc.n_embd] = ds4_bf16_round(acc);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
template<short NR0, typename args_t, bool ROUND = false, bool ROUND_IN = false,
|
|
bool ROPE_IN = false, bool ROPE_OUT = false, bool HC_OUT = false>
|
|
void kernel_mul_mv_q8_0_f32_rope_impl(
|
|
args_t args,
|
|
device const char * src0,
|
|
device const char * src1,
|
|
device char * dst,
|
|
threadgroup char * shmem,
|
|
uint3 tgpig,
|
|
ushort tiisg,
|
|
ushort sgitg,
|
|
constant ds4_metal_args_mv_rope * rope,
|
|
constant ds4_metal_args_mv_hc_expand4 * hc = nullptr,
|
|
device const float * hc_add = nullptr,
|
|
device const float * hc_residual = nullptr,
|
|
device const float * hc_post = nullptr,
|
|
device const float * hc_comb = nullptr) {
|
|
const short NSG = FC_mul_mv_nsg;
|
|
|
|
constexpr short NW = N_SIMDWIDTH;
|
|
constexpr short NQ = 8;
|
|
|
|
const int nb = args.ne00/QK8_0;
|
|
|
|
const int r0 = tgpig.x*NR0;
|
|
const int r1 = tgpig.y;
|
|
const int im = tgpig.z;
|
|
|
|
const uint i12 = im%args.ne12;
|
|
const uint i13 = im/args.ne12;
|
|
|
|
const uint64_t offset1 = r1*args.nb11 + (i12)*args.nb12 + (i13)*args.nb13;
|
|
|
|
device const float * y = (device const float *) (src1 + offset1);
|
|
|
|
device const block_q8_0 * ax[NR0];
|
|
FOR_UNROLL (short row = 0; row < NR0; ++row) {
|
|
const uint64_t offset0 = (r0 + row)*args.nb01 + (i12/args.r2)*args.nb02 + (i13/args.r3)*args.nb03;
|
|
|
|
ax[row] = (device const block_q8_0 *) ((device char *) src0 + offset0);
|
|
}
|
|
|
|
float sumf[NR0] = { 0.f };
|
|
|
|
const short ix = tiisg/(NW/NQ);
|
|
const short il = tiisg%(NW/NQ);
|
|
|
|
const int ib0 = sgitg*NQ + ix;
|
|
|
|
float yl[NQ];
|
|
|
|
device const float * yb = y + ib0*QK8_0 + il*NQ;
|
|
|
|
for (int ib = ib0; ib < nb; ib += NSG*NQ) {
|
|
for (short i = 0; i < NQ; ++i) {
|
|
yl[i] = ROUND_IN ? ds4_bf16_round(yb[i]) : yb[i];
|
|
}
|
|
|
|
if (ROPE_IN) {
|
|
const uint col = (uint)(ib*QK8_0 + il*NQ) % rope->width;
|
|
if (col >= rope->width - 64u) {
|
|
const uint pair0 = (col - (rope->width - 64u)) >> 1u;
|
|
FOR_UNROLL (short p = 0; p < NQ/2; ++p) {
|
|
const float theta = float(rope->start) * rope->frequencies[pair0 + p];
|
|
float c;
|
|
const float sn = precise::sincos(theta, c);
|
|
const float s = rope->inverse ? -sn : sn;
|
|
const float re = yl[2*p], im = yl[2*p + 1];
|
|
yl[2*p] = ds4_bf16_round(fma(re, c, -(im * s)));
|
|
yl[2*p + 1] = ds4_bf16_round(fma(im, c, re * s));
|
|
}
|
|
}
|
|
}
|
|
|
|
for (short row = 0; row < NR0; row++) {
|
|
device const int8_t * qs = ax[row][ib].qs + il*NQ;
|
|
|
|
float sumq = 0.f;
|
|
FOR_UNROLL (short i = 0; i < NQ; ++i) {
|
|
sumq += qs[i] * yl[i];
|
|
}
|
|
|
|
sumf[row] += sumq*ax[row][ib].d;
|
|
}
|
|
|
|
yb += NSG*NQ*QK8_0;
|
|
}
|
|
|
|
device float * dst_f32 = (device float *) dst + (uint64_t)im*args.ne0*args.ne1 + (uint64_t)r1*args.ne0;
|
|
|
|
if (HC_OUT) {
|
|
helper_mv_reduce_and_write_hc_expand4<NR0>(dst_f32, sumf, r0, args.ne01, tiisg, sgitg, shmem,
|
|
*hc, hc_add, hc_residual, hc_post, hc_comb);
|
|
} else if (ROPE_OUT) {
|
|
helper_mv_reduce_and_write_rope<NR0>(dst_f32, sumf, r0, args.ne01, tiisg, sgitg, shmem, *rope);
|
|
} else {
|
|
helper_mv_reduce_and_write<NR0, ROUND>(dst_f32, sumf, r0, args.ne01, tiisg, sgitg, shmem);
|
|
}
|
|
}
|
|
|
|
template<short NR0, typename args_t, bool ROUND = false, bool ROUND_IN = false>
|
|
void kernel_mul_mv_q8_0_f32_impl(
|
|
args_t args,
|
|
device const char * src0,
|
|
device const char * src1,
|
|
device char * dst,
|
|
threadgroup char * shmem,
|
|
uint3 tgpig,
|
|
ushort tiisg,
|
|
ushort sgitg) {
|
|
kernel_mul_mv_q8_0_f32_rope_impl<NR0, args_t, ROUND, ROUND_IN>(
|
|
args, src0, src1, dst, shmem, tgpig, tiisg, sgitg, (constant ds4_metal_args_mv_rope *)nullptr);
|
|
}
|
|
|
|
// One simdgroup group per activation row, each running the single-row
|
|
// kernel's walk over the same weight blocks: every row's result matches the
|
|
// one-row kernel, and the rows share the weight loads through the cache.
|
|
template<short NR0, bool ROUND = false>
|
|
static inline void helper_mv_reduce_and_write_rows(
|
|
device float * dst_f32,
|
|
float sumf[NR0],
|
|
const int r0,
|
|
const int ne01,
|
|
ushort tiisg,
|
|
ushort sg,
|
|
short r1,
|
|
threadgroup char * shmem) {
|
|
constexpr short NW = N_SIMDWIDTH;
|
|
|
|
threadgroup float * shmem_f32[NR0];
|
|
|
|
for (short row = 0; row < NR0; ++row) {
|
|
shmem_f32[row] = (threadgroup float *) shmem + NW*(r1*NR0 + row);
|
|
|
|
if (sg == 0) {
|
|
shmem_f32[row][tiisg] = 0.0f;
|
|
}
|
|
|
|
sumf[row] = simd_sum(sumf[row]);
|
|
}
|
|
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
for (short row = 0; row < NR0; ++row) {
|
|
if (tiisg == 0) {
|
|
shmem_f32[row][sg] = sumf[row];
|
|
}
|
|
}
|
|
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
for (short row = 0; row < NR0 && r0 + row < ne01; ++row) {
|
|
float tot = simd_sum(shmem_f32[row][tiisg]);
|
|
|
|
if (tiisg == 0 && sg == 0) {
|
|
dst_f32[r0 + row] = ROUND ? ds4_bf16_round(tot) : tot;
|
|
}
|
|
}
|
|
}
|
|
|
|
// A few rows side by side: one simdgroup group per row in the threadgroup,
|
|
// each with the single-row K walk and reduction.
|
|
template<short NR0, typename args_t, bool ROUND = false, bool ROUND_IN = false>
|
|
void kernel_mul_mv_q8_0_f32_rows_impl(
|
|
args_t args,
|
|
device const char * src0,
|
|
device const char * src1,
|
|
device char * dst,
|
|
threadgroup char * shmem,
|
|
uint3 tgpig,
|
|
ushort tiisg,
|
|
ushort sgitg) {
|
|
const short NSG = FC_mul_mv_nsg;
|
|
|
|
constexpr short NW = N_SIMDWIDTH;
|
|
constexpr short NQ = 8;
|
|
|
|
const int nb = args.ne00/QK8_0;
|
|
|
|
const int r0 = tgpig.x*NR0;
|
|
const int r1 = sgitg/NSG;
|
|
const ushort sg = sgitg%NSG;
|
|
const int im = tgpig.z;
|
|
|
|
const uint i12 = im%args.ne12;
|
|
const uint i13 = im/args.ne12;
|
|
|
|
const uint64_t offset1 = r1*args.nb11 + (i12)*args.nb12 + (i13)*args.nb13;
|
|
|
|
device const float * y = (device const float *) (src1 + offset1);
|
|
|
|
device const block_q8_0 * ax[NR0];
|
|
FOR_UNROLL (short row = 0; row < NR0; ++row) {
|
|
const uint64_t offset0 = (r0 + row)*args.nb01 + (i12/args.r2)*args.nb02 + (i13/args.r3)*args.nb03;
|
|
|
|
ax[row] = (device const block_q8_0 *) ((device char *) src0 + offset0);
|
|
}
|
|
|
|
float sumf[NR0] = { 0.f };
|
|
|
|
const short ix = tiisg/(NW/NQ);
|
|
const short il = tiisg%(NW/NQ);
|
|
|
|
const int ib0 = sg*NQ + ix;
|
|
|
|
float yl[NQ];
|
|
|
|
device const float * yb = y + ib0*QK8_0 + il*NQ;
|
|
|
|
for (int ib = ib0; ib < nb; ib += NSG*NQ) {
|
|
for (short i = 0; i < NQ; ++i) {
|
|
yl[i] = ROUND_IN ? ds4_bf16_round(yb[i]) : yb[i];
|
|
}
|
|
|
|
for (short row = 0; row < NR0; row++) {
|
|
device const int8_t * qs = ax[row][ib].qs + il*NQ;
|
|
|
|
float sumq = 0.f;
|
|
FOR_UNROLL (short i = 0; i < NQ; ++i) {
|
|
sumq += qs[i] * yl[i];
|
|
}
|
|
|
|
sumf[row] += sumq*ax[row][ib].d;
|
|
}
|
|
|
|
yb += NSG*NQ*QK8_0;
|
|
}
|
|
|
|
device float * dst_f32 = (device float *) dst + (uint64_t)im*args.ne0*args.ne1 + (uint64_t)r1*args.ne0;
|
|
|
|
helper_mv_reduce_and_write_rows<NR0, ROUND>(dst_f32, sumf, r0, args.ne01, tiisg, sg, (short)r1, shmem);
|
|
}
|
|
|
|
// The rows one after another on the single-row simdgroup layout: the
|
|
// threadgroup keeps the weight lines hot across rows and runs at the
|
|
// single-row occupancy; every row's result matches the one-row kernel.
|
|
template<short NR0, typename args_t, bool ROUND = false, bool ROUND_IN = false, bool HC_OUT = false>
|
|
void kernel_mul_mv_q8_0_f32_rows_seq_impl(
|
|
args_t args,
|
|
device const char * src0,
|
|
device const char * src1,
|
|
device char * dst,
|
|
threadgroup char * shmem,
|
|
uint3 tgpig,
|
|
ushort tiisg,
|
|
ushort sgitg,
|
|
constant ds4_metal_args_mv_hc_expand4 * hc = nullptr,
|
|
device const float * hc_add = nullptr,
|
|
device const float * hc_residual = nullptr,
|
|
device const float * hc_post = nullptr,
|
|
device const float * hc_comb = nullptr) {
|
|
const short NSG = FC_mul_mv_nsg;
|
|
constexpr short NW = N_SIMDWIDTH;
|
|
constexpr short NQ = 8;
|
|
const int nb = args.ne00/QK8_0;
|
|
const int r0 = tgpig.x*NR0;
|
|
const int im = tgpig.z;
|
|
const uint i12 = im%args.ne12;
|
|
const uint i13 = im/args.ne12;
|
|
device const block_q8_0 * ax[NR0];
|
|
FOR_UNROLL (short row = 0; row < NR0; ++row) {
|
|
const uint64_t offset0 = (r0 + row)*args.nb01 + (i12/args.r2)*args.nb02 + (i13/args.r3)*args.nb03;
|
|
ax[row] = (device const block_q8_0 *) ((device char *) src0 + offset0);
|
|
}
|
|
const short ix = tiisg/(NW/NQ);
|
|
const short il = tiisg%(NW/NQ);
|
|
const int ib0 = sgitg*NQ + ix;
|
|
for (int r1 = 0; r1 < args.ne11; ++r1) {
|
|
const uint64_t offset1 = r1*args.nb11 + (i12)*args.nb12 + (i13)*args.nb13;
|
|
device const float * y = (device const float *) (src1 + offset1);
|
|
float sumf[NR0] = { 0.f };
|
|
float yl[NQ];
|
|
device const float * yb = y + ib0*QK8_0 + il*NQ;
|
|
for (int ib = ib0; ib < nb; ib += NSG*NQ) {
|
|
for (short i = 0; i < NQ; ++i) {
|
|
yl[i] = ROUND_IN ? ds4_bf16_round(yb[i]) : yb[i];
|
|
}
|
|
for (short row = 0; row < NR0; row++) {
|
|
device const int8_t * qs = ax[row][ib].qs + il*NQ;
|
|
float sumq = 0.f;
|
|
FOR_UNROLL (short i = 0; i < NQ; ++i) {
|
|
sumq += qs[i] * yl[i];
|
|
}
|
|
sumf[row] += sumq*ax[row][ib].d;
|
|
}
|
|
yb += NSG*NQ*QK8_0;
|
|
}
|
|
if (HC_OUT) {
|
|
const uint64_t hc_dim = 4ull * hc->n_embd, stride = hc->split_stride;
|
|
helper_mv_reduce_and_write_hc_expand4<NR0>((device float *) dst + r1*hc_dim, sumf, r0, args.ne01,
|
|
tiisg, sgitg, shmem, *hc, hc_add + r1*hc->n_embd, hc_residual + r1*hc_dim,
|
|
hc_post + r1*stride, hc_comb + r1*stride);
|
|
} else {
|
|
device float * dst_f32 = (device float *) dst + (uint64_t)im*args.ne0*args.ne1 + (uint64_t)r1*args.ne0;
|
|
helper_mv_reduce_and_write<NR0, ROUND>(dst_f32, sumf, r0, args.ne01, tiisg, sgitg, shmem);
|
|
}
|
|
threadgroup_barrier(mem_flags::mem_threadgroup); /* the reduce scratch is reused */
|
|
}
|
|
}
|
|
|
|
#define DS4_MUL_MV_Q8_0_ROWS_SEQ(NAME, NR0, ...) \
|
|
[[host_name(#NAME)]] \
|
|
kernel void NAME( \
|
|
constant ds4_metal_args_mul_mv & args, \
|
|
device const char * src0, \
|
|
device const char * src1, \
|
|
device char * dst, \
|
|
threadgroup char * shmem [[threadgroup(0)]], \
|
|
uint3 tgpig[[threadgroup_position_in_grid]], \
|
|
ushort tiisg[[thread_index_in_simdgroup]], \
|
|
ushort sgitg[[simdgroup_index_in_threadgroup]]) { \
|
|
kernel_mul_mv_q8_0_f32_rows_seq_impl<NR0, constant ds4_metal_args_mul_mv &, ##__VA_ARGS__>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg); \
|
|
}
|
|
DS4_MUL_MV_Q8_0_ROWS_SEQ(kernel_mul_mv_q8_0_f32_rows_seq4, 4)
|
|
DS4_MUL_MV_Q8_0_ROWS_SEQ(kernel_mul_mv_q8_0_f32_bf16_rows_seq4, 4, true)
|
|
DS4_MUL_MV_Q8_0_ROWS_SEQ(kernel_mul_mv_q8_0_f32_bf16io_rows_seq4, 4, true, true)
|
|
#undef DS4_MUL_MV_Q8_0_ROWS_SEQ
|
|
|
|
#define DS4_MUL_MV_Q8_0_HC_ROWS_SEQ(NAME, NR0) \
|
|
[[host_name(#NAME)]] \
|
|
kernel void NAME( \
|
|
constant ds4_metal_args_mul_mv & args, \
|
|
device const char * src0, \
|
|
device const char * src1, \
|
|
device char * dst, \
|
|
constant ds4_metal_args_mv_hc_expand4 & hc, \
|
|
device const float * add, \
|
|
device const float * residual, \
|
|
device const float * post, \
|
|
device const float * comb, \
|
|
threadgroup char * shmem [[threadgroup(0)]], \
|
|
uint3 tgpig[[threadgroup_position_in_grid]], \
|
|
ushort tiisg[[thread_index_in_simdgroup]], \
|
|
ushort sgitg[[simdgroup_index_in_threadgroup]]) { \
|
|
kernel_mul_mv_q8_0_f32_rows_seq_impl<NR0, constant ds4_metal_args_mul_mv &, true, true, true>( \
|
|
args, src0, src1, dst, shmem, tgpig, tiisg, sgitg, &hc, add, residual, post, comb); \
|
|
}
|
|
DS4_MUL_MV_Q8_0_HC_ROWS_SEQ(kernel_mul_mv_q8_0_f32_bf16io_hc_expand4_rows_seq4, 4)
|
|
#undef DS4_MUL_MV_Q8_0_HC_ROWS_SEQ
|
|
|
|
// Two to eight activation rows through simdgroup matrices: the simdgroups
|
|
// of a threadgroup each take a slice of K for the same eight weight rows,
|
|
// dequantize a 32-wide block once and multiply it against the eight
|
|
// (zero-padded) activation rows, then the partial tiles add in a fixed
|
|
// order. The weights are read once whatever the row count.
|
|
template<bool ROUND, bool ROUND_IN, bool HC_OUT = false, short NSG = 4,
|
|
typename args_t = constant ds4_metal_args_mul_mv &>
|
|
void kernel_mul_mv_q8_0_f32_mma_rows_impl(
|
|
args_t args,
|
|
device const char * src0,
|
|
device const char * src1,
|
|
device char * dst,
|
|
threadgroup char * shmem,
|
|
uint3 tgpig,
|
|
ushort tiisg,
|
|
ushort sgitg,
|
|
constant ds4_metal_args_mv_hc_expand4 * hc = nullptr,
|
|
device const float * hc_add = nullptr,
|
|
device const float * hc_residual = nullptr,
|
|
device const float * hc_post = nullptr,
|
|
device const float * hc_comb = nullptr) {
|
|
constexpr short NR = 8, KC = QK8_0;
|
|
// per simdgroup: 4 tiles of 8x8 for the activations, 4 for the weights
|
|
threadgroup float * A = (threadgroup float *) shmem + sgitg * (2 * NR * KC);
|
|
threadgroup float * W = A + NR * KC;
|
|
const int r0 = (int)tgpig.x * NR;
|
|
const int kq = args.ne00 / NSG; /* this simdgroup's K slice */
|
|
const int k_begin = sgitg * kq, k_end = k_begin + kq;
|
|
const short wr = tiisg / 4, wq = tiisg % 4; /* row, 8-wide quarter of the block */
|
|
device const block_q8_0 * wrow = (device const block_q8_0 *)(src0 + (uint64_t)(r0 + wr) * args.nb01);
|
|
const bool have_row = wr < args.ne11;
|
|
device const float * y = (device const float *)(src1 + (uint64_t)wr * args.nb11) + wq * 8;
|
|
threadgroup float * a = A + wq * 64 + wr * 8;
|
|
threadgroup float * w = W + wq * 64 + wr * 8;
|
|
simdgroup_float8x8 acc = make_filled_simdgroup_matrix<float, 8, 8>(0.0f);
|
|
// the next block's loads are issued before this block's matrix products
|
|
float av[8];
|
|
float4 q0, q1;
|
|
float d;
|
|
{
|
|
device const block_q8_0 * b = wrow + k_begin / QK8_0;
|
|
device const packed_char4 * q4 = (device const packed_char4 *)(b->qs + wq * 8);
|
|
d = (float)b->d; q0 = float4(q4[0]); q1 = float4(q4[1]);
|
|
for (short i = 0; i < 8; ++i) av[i] = have_row ? y[k_begin + i] : 0.0f;
|
|
}
|
|
for (int k0 = k_begin; k0 < k_end; k0 += KC) {
|
|
for (short i = 0; i < 8; ++i) a[i] = ROUND_IN ? ds4_bf16_round(av[i]) : av[i];
|
|
const float4 v0 = d * q0, v1 = d * q1;
|
|
w[0] = v0.x; w[1] = v0.y; w[2] = v0.z; w[3] = v0.w;
|
|
w[4] = v1.x; w[5] = v1.y; w[6] = v1.z; w[7] = v1.w;
|
|
if (k0 + KC < k_end) {
|
|
device const block_q8_0 * b = wrow + (k0 + KC) / QK8_0;
|
|
device const packed_char4 * q4 = (device const packed_char4 *)(b->qs + wq * 8);
|
|
d = (float)b->d; q0 = float4(q4[0]); q1 = float4(q4[1]);
|
|
for (short i = 0; i < 8; ++i) av[i] = have_row ? y[k0 + KC + i] : 0.0f;
|
|
}
|
|
simdgroup_barrier(mem_flags::mem_threadgroup);
|
|
for (short kt = 0; kt < KC / 8; ++kt) {
|
|
simdgroup_float8x8 wt, at;
|
|
simdgroup_load(wt, W + kt * 64, 8, 0, false);
|
|
simdgroup_load(at, A + kt * 64, 8, 0, true);
|
|
simdgroup_multiply_accumulate(acc, wt, at, acc);
|
|
}
|
|
simdgroup_barrier(mem_flags::mem_threadgroup);
|
|
}
|
|
// the partial tiles, summed in simdgroup order
|
|
threadgroup float * P = (threadgroup float *) shmem + NSG * 2 * NR * KC + sgitg * 64;
|
|
simdgroup_store(acc, P, 8, 0, false);
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
if (sgitg == 0) {
|
|
threadgroup const float * P0 = (threadgroup const float *) shmem + NSG * 2 * NR * KC;
|
|
const short j = tiisg / 4;
|
|
if (j < args.ne11) {
|
|
device float * o = (device float *) dst + (uint64_t)j * args.ne0 + r0;
|
|
for (short i = (tiisg % 4) * 2; i < (tiisg % 4) * 2 + 2; ++i) {
|
|
float v = P0[i * 8 + j];
|
|
for (short g = 1; g < NSG; ++g) v += P0[g * 64 + i * 8 + j];
|
|
if (HC_OUT) {
|
|
// the rows kernel's expansion, one (row, output) pair per lane step
|
|
const uint n_embd = hc->n_embd;
|
|
const uint64_t hc_dim = 4ull * n_embd, stride = hc->split_stride;
|
|
const int d = r0 + i;
|
|
float block_v = ds4_bf16_round(v);
|
|
if (hc->has_add) block_v = ds4_bf16_round(block_v + hc_add[(uint64_t)j * n_embd + d]);
|
|
device const float * res = hc_residual + (uint64_t)j * hc_dim;
|
|
device const float * pj = hc_post + (uint64_t)j * stride;
|
|
device const float * cj = hc_comb + (uint64_t)j * stride;
|
|
const float r0v = res[d], r1v = res[d + n_embd], r2v = res[d + 2 * n_embd], r3v = res[d + 3 * n_embd];
|
|
device float * dj = (device float *) dst + (uint64_t)j * hc_dim;
|
|
for (int dst_hc = 0; dst_hc < 4; ++dst_hc) {
|
|
float acc = block_v * pj[dst_hc];
|
|
acc += cj[dst_hc + 0 * 4] * r0v;
|
|
acc += cj[dst_hc + 1 * 4] * r1v;
|
|
acc += cj[dst_hc + 2 * 4] * r2v;
|
|
acc += cj[dst_hc + 3 * 4] * r3v;
|
|
dj[d + dst_hc * n_embd] = ds4_bf16_round(acc);
|
|
}
|
|
} else {
|
|
o[i] = ROUND ? ds4_bf16_round(v) : v;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
#define DS4_MUL_MV_Q8_0_MMA_ROWS(NAME, ...) \
|
|
[[host_name(#NAME)]] \
|
|
kernel void NAME( \
|
|
constant ds4_metal_args_mul_mv & args, \
|
|
device const char * src0, \
|
|
device const char * src1, \
|
|
device char * dst, \
|
|
threadgroup char * shmem [[threadgroup(0)]], \
|
|
uint3 tgpig[[threadgroup_position_in_grid]], \
|
|
ushort tiisg[[thread_index_in_simdgroup]], \
|
|
ushort sgitg[[simdgroup_index_in_threadgroup]]) { \
|
|
kernel_mul_mv_q8_0_f32_mma_rows_impl<__VA_ARGS__>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg); \
|
|
}
|
|
DS4_MUL_MV_Q8_0_MMA_ROWS(kernel_mul_mv_q8_0_f32_mma_rows, false, false)
|
|
DS4_MUL_MV_Q8_0_MMA_ROWS(kernel_mul_mv_q8_0_f32_bf16_mma_rows, true, false)
|
|
DS4_MUL_MV_Q8_0_MMA_ROWS(kernel_mul_mv_q8_0_f32_bf16io_mma_rows, true, true)
|
|
DS4_MUL_MV_Q8_0_MMA_ROWS(kernel_mul_mv_q8_0_f32_mma_rows8, false, false, false, 8)
|
|
DS4_MUL_MV_Q8_0_MMA_ROWS(kernel_mul_mv_q8_0_f32_bf16_mma_rows8, true, false, false, 8)
|
|
DS4_MUL_MV_Q8_0_MMA_ROWS(kernel_mul_mv_q8_0_f32_bf16io_mma_rows8, true, true, false, 8)
|
|
#undef DS4_MUL_MV_Q8_0_MMA_ROWS
|
|
[[host_name("kernel_mul_mv_q8_0_f32_bf16io_hc_expand4_mma_rows")]]
|
|
kernel void kernel_mul_mv_q8_0_f32_bf16io_hc_expand4_mma_rows(
|
|
constant ds4_metal_args_mul_mv & args,
|
|
device const char * src0,
|
|
device const char * src1,
|
|
device char * dst,
|
|
constant ds4_metal_args_mv_hc_expand4 & hc,
|
|
device const float * add,
|
|
device const float * residual,
|
|
device const float * post,
|
|
device const float * comb,
|
|
threadgroup char * shmem [[threadgroup(0)]],
|
|
uint3 tgpig[[threadgroup_position_in_grid]],
|
|
ushort tiisg[[thread_index_in_simdgroup]],
|
|
ushort sgitg[[simdgroup_index_in_threadgroup]]) {
|
|
kernel_mul_mv_q8_0_f32_mma_rows_impl<true, true, true>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg,
|
|
&hc, add, residual, post, comb);
|
|
}
|
|
[[host_name("kernel_mul_mv_q8_0_f32_bf16io_hc_expand4_mma_rows8")]]
|
|
kernel void kernel_mul_mv_q8_0_f32_bf16io_hc_expand4_mma_rows8(
|
|
constant ds4_metal_args_mul_mv & args,
|
|
device const char * src0,
|
|
device const char * src1,
|
|
device char * dst,
|
|
constant ds4_metal_args_mv_hc_expand4 & hc,
|
|
device const float * add,
|
|
device const float * residual,
|
|
device const float * post,
|
|
device const float * comb,
|
|
threadgroup char * shmem [[threadgroup(0)]],
|
|
uint3 tgpig[[threadgroup_position_in_grid]],
|
|
ushort tiisg[[thread_index_in_simdgroup]],
|
|
ushort sgitg[[simdgroup_index_in_threadgroup]]) {
|
|
kernel_mul_mv_q8_0_f32_mma_rows_impl<true, true, true, 8>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg,
|
|
&hc, add, residual, post, comb);
|
|
}
|
|
|
|
// Decode-time Q8_0 matrix-vector multiply. DS4 uses this for Q8_0 dense
|
|
// projections such as shared experts and output-side small matvecs.
|
|
[[host_name("kernel_mul_mv_q8_0_f32")]]
|
|
kernel void kernel_mul_mv_q8_0_f32(
|
|
constant ds4_metal_args_mul_mv & args,
|
|
device const char * src0,
|
|
device const char * src1,
|
|
device char * dst,
|
|
threadgroup char * shmem [[threadgroup(0)]],
|
|
uint3 tgpig[[threadgroup_position_in_grid]],
|
|
ushort tiisg[[thread_index_in_simdgroup]],
|
|
ushort sgitg[[simdgroup_index_in_threadgroup]]) {
|
|
kernel_mul_mv_q8_0_f32_impl<N_R0_Q8_0, constant ds4_metal_args_mul_mv &>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
|
|
}
|
|
|
|
[[host_name("kernel_mul_mv_q8_0_f32_bf16")]]
|
|
kernel void kernel_mul_mv_q8_0_f32_bf16(
|
|
constant ds4_metal_args_mul_mv & args,
|
|
device const char * src0,
|
|
device const char * src1,
|
|
device char * dst,
|
|
threadgroup char * shmem [[threadgroup(0)]],
|
|
uint3 tgpig[[threadgroup_position_in_grid]],
|
|
ushort tiisg[[thread_index_in_simdgroup]],
|
|
ushort sgitg[[simdgroup_index_in_threadgroup]]) {
|
|
kernel_mul_mv_q8_0_f32_impl<N_R0_Q8_0, constant ds4_metal_args_mul_mv &, true>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
|
|
}
|
|
|
|
[[host_name("kernel_mul_mv_q8_0_f32_rows8")]]
|
|
kernel void kernel_mul_mv_q8_0_f32_rows8(
|
|
constant ds4_metal_args_mul_mv & args,
|
|
device const char * src0,
|
|
device const char * src1,
|
|
device char * dst,
|
|
threadgroup char * shmem [[threadgroup(0)]],
|
|
uint3 tgpig[[threadgroup_position_in_grid]],
|
|
ushort tiisg[[thread_index_in_simdgroup]],
|
|
ushort sgitg[[simdgroup_index_in_threadgroup]]) {
|
|
kernel_mul_mv_q8_0_f32_impl<8, constant ds4_metal_args_mul_mv &, false>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
|
|
}
|
|
|
|
[[host_name("kernel_mul_mv_q8_0_f32_bf16io")]]
|
|
kernel void kernel_mul_mv_q8_0_f32_bf16io(
|
|
constant ds4_metal_args_mul_mv & args,
|
|
device const char * src0,
|
|
device const char * src1,
|
|
device char * dst,
|
|
threadgroup char * shmem [[threadgroup(0)]],
|
|
uint3 tgpig[[threadgroup_position_in_grid]],
|
|
ushort tiisg[[thread_index_in_simdgroup]],
|
|
ushort sgitg[[simdgroup_index_in_threadgroup]]) {
|
|
kernel_mul_mv_q8_0_f32_impl<N_R0_Q8_0, constant ds4_metal_args_mul_mv &, true, true>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
|
|
}
|
|
|
|
// The q projection with its rope applied to the bf16 outputs.
|
|
[[host_name("kernel_mul_mv_q8_0_f32_bf16_rope")]]
|
|
kernel void kernel_mul_mv_q8_0_f32_bf16_rope(
|
|
constant ds4_metal_args_mul_mv & args,
|
|
device const char * src0,
|
|
device const char * src1,
|
|
device char * dst,
|
|
constant ds4_metal_args_mv_rope & rope,
|
|
threadgroup char * shmem [[threadgroup(0)]],
|
|
uint3 tgpig[[threadgroup_position_in_grid]],
|
|
ushort tiisg[[thread_index_in_simdgroup]],
|
|
ushort sgitg[[simdgroup_index_in_threadgroup]]) {
|
|
kernel_mul_mv_q8_0_f32_rope_impl<N_R0_Q8_0, constant ds4_metal_args_mul_mv &, true, false, false, true>(
|
|
args, src0, src1, dst, shmem, tgpig, tiisg, sgitg, &rope);
|
|
}
|
|
|
|
// A bf16 in/out matvec whose outputs expand straight into the HC streams.
|
|
[[host_name("kernel_mul_mv_q8_0_f32_bf16io_hc_expand4")]]
|
|
kernel void kernel_mul_mv_q8_0_f32_bf16io_hc_expand4(
|
|
constant ds4_metal_args_mul_mv & args,
|
|
device const char * src0,
|
|
device const char * src1,
|
|
device char * dst,
|
|
constant ds4_metal_args_mv_hc_expand4 & hc,
|
|
device const float * add,
|
|
device const float * residual,
|
|
device const float * post,
|
|
device const float * comb,
|
|
threadgroup char * shmem [[threadgroup(0)]],
|
|
uint3 tgpig[[threadgroup_position_in_grid]],
|
|
ushort tiisg[[thread_index_in_simdgroup]],
|
|
ushort sgitg[[simdgroup_index_in_threadgroup]]) {
|
|
kernel_mul_mv_q8_0_f32_rope_impl<N_R0_Q8_0, constant ds4_metal_args_mul_mv &, true, true, false, false, true>(
|
|
args, src0, src1, dst, shmem, tgpig, tiisg, sgitg, nullptr, &hc, add, residual, post, comb);
|
|
}
|
|
|
|
// Two Q8_0 matrices over one activation row in one dispatch: the leading
|
|
// threadgroups take the first matrix's row groups, the rest the second's.
|
|
[[host_name("kernel_mul_mv_q8_0_f32_bf16_pair")]]
|
|
kernel void kernel_mul_mv_q8_0_f32_bf16_pair(
|
|
constant ds4_metal_args_mul_mv & args_a,
|
|
constant ds4_metal_args_mul_mv & args_b,
|
|
device const char * src0_a,
|
|
device const char * src0_b,
|
|
device const char * src1,
|
|
device char * dst_a,
|
|
device char * dst_b,
|
|
threadgroup char * shmem [[threadgroup(0)]],
|
|
uint3 tgpig[[threadgroup_position_in_grid]],
|
|
ushort tiisg[[thread_index_in_simdgroup]],
|
|
ushort sgitg[[simdgroup_index_in_threadgroup]]) {
|
|
const uint groups_a = ((uint)args_a.ne01 + N_R0_Q8_0 - 1u) / N_R0_Q8_0;
|
|
if (tgpig.x < groups_a) {
|
|
kernel_mul_mv_q8_0_f32_impl<N_R0_Q8_0, constant ds4_metal_args_mul_mv &, true>(args_a, src0_a, src1, dst_a, shmem, tgpig, tiisg, sgitg);
|
|
} else {
|
|
tgpig.x -= groups_a;
|
|
kernel_mul_mv_q8_0_f32_impl<N_R0_Q8_0, constant ds4_metal_args_mul_mv &, true>(args_b, src0_b, src1, dst_b, shmem, tgpig, tiisg, sgitg);
|
|
}
|
|
}
|
|
|
|
// Q8_0 matvec whose output is this rank's TP partial in its slab slot: same
|
|
// K walk and reduction tree as kernel_mul_mv_q8_0_f32_impl, plus the checked
|
|
// poll-gate flag published by the last-arriving threadgroup (see
|
|
// kernel_dsv4_add2_f32_tp_flag_checked). Writes exactly the values the plain
|
|
// kernel writes; the checksum is the integer sum of the stored words.
|
|
template<short NR0>
|
|
void kernel_dsv4_mul_mv_q8_0_f32_tp_flag_impl(
|
|
constant ds4_metal_args_mul_mv & args,
|
|
device atomic_uint & flag,
|
|
device atomic_uint & check,
|
|
constant uint & value,
|
|
device atomic_uint * ctl,
|
|
constant uint & ntg,
|
|
device const char * src0,
|
|
device const char * src1,
|
|
device char * dst,
|
|
threadgroup char * shmem,
|
|
threadgroup uint * ctl_shmem,
|
|
uint3 tgpig,
|
|
ushort tiisg,
|
|
ushort sgitg) {
|
|
const short NSG = FC_mul_mv_nsg;
|
|
|
|
constexpr short NW = N_SIMDWIDTH;
|
|
constexpr short NQ = 8;
|
|
|
|
const int nb = args.ne00/QK8_0;
|
|
|
|
const int r0 = tgpig.x*NR0;
|
|
const int r1 = tgpig.y;
|
|
const int im = tgpig.z;
|
|
|
|
const uint i12 = im%args.ne12;
|
|
const uint i13 = im/args.ne12;
|
|
|
|
const uint64_t offset1 = r1*args.nb11 + (i12)*args.nb12 + (i13)*args.nb13;
|
|
|
|
device const float * y = (device const float *) (src1 + offset1);
|
|
|
|
device const block_q8_0 * ax[NR0];
|
|
FOR_UNROLL (short row = 0; row < NR0; ++row) {
|
|
const uint64_t offset0 = (r0 + row)*args.nb01 + (i12/args.r2)*args.nb02 + (i13/args.r3)*args.nb03;
|
|
|
|
ax[row] = (device const block_q8_0 *) ((device char *) src0 + offset0);
|
|
}
|
|
|
|
float sumf[NR0] = { 0.f };
|
|
|
|
const short ix = tiisg/(NW/NQ);
|
|
const short il = tiisg%(NW/NQ);
|
|
|
|
const int ib0 = sgitg*NQ + ix;
|
|
|
|
float yl[NQ];
|
|
|
|
device const float * yb = y + ib0*QK8_0 + il*NQ;
|
|
|
|
for (int ib = ib0; ib < nb; ib += NSG*NQ) {
|
|
for (short i = 0; i < NQ; ++i) {
|
|
yl[i] = yb[i];
|
|
}
|
|
|
|
for (short row = 0; row < NR0; row++) {
|
|
device const int8_t * qs = ax[row][ib].qs + il*NQ;
|
|
|
|
float sumq = 0.f;
|
|
FOR_UNROLL (short i = 0; i < NQ; ++i) {
|
|
sumq += qs[i] * yl[i];
|
|
}
|
|
|
|
sumf[row] += sumq*ax[row][ib].d;
|
|
}
|
|
|
|
yb += NSG*NQ*QK8_0;
|
|
}
|
|
|
|
device float * dst_f32 = (device float *) dst + (uint64_t)im*args.ne0*args.ne1 + (uint64_t)r1*args.ne0;
|
|
|
|
// Reduce and write exactly as the plain kernel does, keeping the stored
|
|
// words for the checksum.
|
|
{
|
|
constexpr short NWR = N_SIMDWIDTH;
|
|
threadgroup float * shmem_f32[NR0];
|
|
for (short row = 0; row < NR0; ++row) {
|
|
shmem_f32[row] = (threadgroup float *) shmem + NWR*row;
|
|
if (sgitg == 0) {
|
|
shmem_f32[row][tiisg] = 0.0f;
|
|
}
|
|
sumf[row] = simd_sum(sumf[row]);
|
|
}
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
for (short row = 0; row < NR0; ++row) {
|
|
if (tiisg == 0) {
|
|
shmem_f32[row][sgitg] = sumf[row];
|
|
}
|
|
}
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
uint words = 0u;
|
|
for (short row = 0; row < NR0 && r0 + row < args.ne01; ++row) {
|
|
float tot = simd_sum(shmem_f32[row][tiisg]);
|
|
if (tiisg == 0 && sgitg == 0) {
|
|
dst_f32[r0 + row] = tot;
|
|
words += as_type<uint>(tot);
|
|
}
|
|
}
|
|
// Publish: accumulate this threadgroup's words, then arrive; the
|
|
// last arriver writes the checksum and the flag.
|
|
if (tiisg == 0 && sgitg == 0) ctl_shmem[0] = words;
|
|
threadgroup_barrier(mem_flags::mem_device | mem_flags::mem_threadgroup);
|
|
if (tiisg == 0 && sgitg == 0) {
|
|
atomic_fetch_add_explicit(&ctl[1], ctl_shmem[0], memory_order_relaxed);
|
|
}
|
|
threadgroup_barrier(mem_flags::mem_device);
|
|
if (tiisg == 0 && sgitg == 0) {
|
|
const uint old = atomic_fetch_add_explicit(&ctl[0], 1u, memory_order_relaxed);
|
|
if (old + 1u == ntg) {
|
|
const uint total = atomic_exchange_explicit(&ctl[1], 0u, memory_order_relaxed);
|
|
atomic_store_explicit(&ctl[0], 0u, memory_order_relaxed);
|
|
atomic_store_explicit(&check, total ^ (value * 0x9E3779B9u), memory_order_relaxed);
|
|
atomic_store_explicit(&flag, value, memory_order_relaxed);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
[[host_name("kernel_dsv4_mul_mv_q8_0_f32_tp_flag_checked")]]
|
|
kernel void kernel_dsv4_mul_mv_q8_0_f32_tp_flag_checked(
|
|
constant ds4_metal_args_mul_mv & args,
|
|
device const char * src0,
|
|
device const char * src1,
|
|
device char * dst,
|
|
device atomic_uint & flag,
|
|
device atomic_uint & check,
|
|
constant uint & value,
|
|
device atomic_uint * ctl,
|
|
constant uint & ntg,
|
|
threadgroup char * shmem [[threadgroup(0)]],
|
|
threadgroup uint * ctl_shmem [[threadgroup(1)]],
|
|
uint3 tgpig[[threadgroup_position_in_grid]],
|
|
ushort tiisg[[thread_index_in_simdgroup]],
|
|
ushort sgitg[[simdgroup_index_in_threadgroup]]) {
|
|
kernel_dsv4_mul_mv_q8_0_f32_tp_flag_impl<N_R0_Q8_0>(
|
|
args, flag, check, value, ctl, ntg, src0, src1, dst, shmem,
|
|
ctl_shmem, tgpig, tiisg, sgitg);
|
|
}
|
|
|
|
// Decode Q-A/KV pair. Both projections consume the same activation row but
|
|
// have independent weight ranges and output extents. Keep the standalone Q8_0
|
|
// lane/block traversal and two-stage reduction verbatim for each bank; only
|
|
// the activation load and threadgroup scheduling are shared.
|
|
[[host_name("kernel_mul_mv_q8_0_f32_pair")]]
|
|
kernel void kernel_mul_mv_q8_0_f32_pair(
|
|
constant ds4_metal_args_mul_mv & args0,
|
|
constant ds4_metal_args_mul_mv & args1,
|
|
device const char * src0_a,
|
|
device const char * src0_b,
|
|
device const char * src1,
|
|
device char * dst_a,
|
|
device char * dst_b,
|
|
threadgroup char * shmem [[threadgroup(0)]],
|
|
uint3 tgpig[[threadgroup_position_in_grid]],
|
|
ushort tiisg[[thread_index_in_simdgroup]],
|
|
ushort sgitg[[simdgroup_index_in_threadgroup]]) {
|
|
const short NSG = FC_mul_mv_nsg;
|
|
constexpr short NW = N_SIMDWIDTH;
|
|
constexpr short NQ = 8;
|
|
constexpr short NR0 = N_R0_Q8_0;
|
|
|
|
const int nb = args0.ne00 / QK8_0;
|
|
const int r0 = tgpig.x * NR0;
|
|
const bool active_a = r0 < args0.ne01;
|
|
const bool active_b = r0 < args1.ne01;
|
|
|
|
device const float *y = (device const float *)src1;
|
|
device const block_q8_0 *ax_a[NR0];
|
|
device const block_q8_0 *ax_b[NR0];
|
|
FOR_UNROLL (short row = 0; row < NR0; ++row) {
|
|
const int out_row = r0 + row;
|
|
ax_a[row] = active_a && out_row < args0.ne01
|
|
? (device const block_q8_0 *)(src0_a + (uint64_t)out_row * args0.nb01)
|
|
: (device const block_q8_0 *)src0_a;
|
|
ax_b[row] = active_b && out_row < args1.ne01
|
|
? (device const block_q8_0 *)(src0_b + (uint64_t)out_row * args1.nb01)
|
|
: (device const block_q8_0 *)src0_b;
|
|
}
|
|
|
|
float suma[NR0] = { 0.f };
|
|
float sumb[NR0] = { 0.f };
|
|
|
|
const short ix = tiisg / (NW / NQ);
|
|
const short il = tiisg % (NW / NQ);
|
|
const int ib0 = sgitg * NQ + ix;
|
|
float yl[NQ];
|
|
device const float *yb = y + ib0 * QK8_0 + il * NQ;
|
|
|
|
for (int ib = ib0; ib < nb; ib += NSG * NQ) {
|
|
FOR_UNROLL (short i = 0; i < NQ; ++i) {
|
|
yl[i] = yb[i];
|
|
}
|
|
|
|
FOR_UNROLL (short row = 0; row < NR0; ++row) {
|
|
const int out_row = r0 + row;
|
|
if (active_a && out_row < args0.ne01) {
|
|
device const int8_t *qs = ax_a[row][ib].qs + il * NQ;
|
|
float sumq = 0.f;
|
|
FOR_UNROLL (short i = 0; i < NQ; ++i) {
|
|
sumq += qs[i] * yl[i];
|
|
}
|
|
suma[row] += sumq * ax_a[row][ib].d;
|
|
}
|
|
if (active_b && out_row < args1.ne01) {
|
|
device const int8_t *qs = ax_b[row][ib].qs + il * NQ;
|
|
float sumq = 0.f;
|
|
FOR_UNROLL (short i = 0; i < NQ; ++i) {
|
|
sumq += qs[i] * yl[i];
|
|
}
|
|
sumb[row] += sumq * ax_b[row][ib].d;
|
|
}
|
|
}
|
|
|
|
yb += NSG * NQ * QK8_0;
|
|
}
|
|
|
|
threadgroup float *shared = (threadgroup float *)shmem;
|
|
threadgroup float *sha[NR0];
|
|
threadgroup float *shb[NR0];
|
|
FOR_UNROLL (short row = 0; row < NR0; ++row) {
|
|
sha[row] = shared + NW * row;
|
|
shb[row] = shared + NW * (NR0 + row);
|
|
if (sgitg == 0) {
|
|
sha[row][tiisg] = 0.0f;
|
|
if (active_b) shb[row][tiisg] = 0.0f;
|
|
}
|
|
suma[row] = simd_sum(suma[row]);
|
|
if (active_b) sumb[row] = simd_sum(sumb[row]);
|
|
}
|
|
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
FOR_UNROLL (short row = 0; row < NR0; ++row) {
|
|
if (tiisg == 0) {
|
|
sha[row][sgitg] = suma[row];
|
|
if (active_b) shb[row][sgitg] = sumb[row];
|
|
}
|
|
}
|
|
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
device float *out_a = (device float *)dst_a;
|
|
device float *out_b = (device float *)dst_b;
|
|
FOR_UNROLL (short row = 0; row < NR0; ++row) {
|
|
const float total_a = simd_sum(sha[row][tiisg]);
|
|
if (tiisg == 0 && sgitg == 0) {
|
|
const int out_row = r0 + row;
|
|
if (active_a && out_row < args0.ne01) out_a[out_row] = total_a;
|
|
}
|
|
if (active_b) {
|
|
const float total_b = simd_sum(shb[row][tiisg]);
|
|
if (tiisg == 0 && sgitg == 0) {
|
|
const int out_row = r0 + row;
|
|
if (out_row < args1.ne01) out_b[out_row] = total_b;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// Decode shared-expert gate/up projections followed by SwiGLU:
|
|
//
|
|
// mid = silu(min(gate, limit)) * clamp(up, -limit, limit)
|
|
//
|
|
// DS4's shared expert uses two Q8_0 matrices with the same input row. This
|
|
// kernel preserves the exact Q8_0 dot-product reduction shape for both
|
|
// projections, still writes gate/up for diagnostics, and derives `mid` in the
|
|
// same lane that owns the reduced output row. The point is not to fuse two
|
|
// independent weight streams into one matmul; it is to remove the separate
|
|
// activation pass and its reread of the two 2048-wide rows.
|
|
template<short NR0, bool STORE_GATE_UP, bool ROUND_GATE_UP = false>
|
|
void kernel_dsv4_shared_gate_up_swiglu_q8_0_impl(
|
|
constant ds4_metal_args_mul_mv & args,
|
|
device const char * src0_gate,
|
|
device const char * src0_up,
|
|
device const char * src1,
|
|
device char * dst_gate,
|
|
device char * dst_up,
|
|
device char * dst_mid,
|
|
constant float &clamp_value,
|
|
threadgroup char * shmem,
|
|
uint3 tgpig[[threadgroup_position_in_grid]],
|
|
ushort tiisg[[thread_index_in_simdgroup]],
|
|
ushort sgitg[[simdgroup_index_in_threadgroup]]) {
|
|
const short NSG = FC_mul_mv_nsg;
|
|
constexpr short NW = N_SIMDWIDTH;
|
|
constexpr short NQ = 8;
|
|
|
|
const int nb = args.ne00 / QK8_0;
|
|
const int r0 = tgpig.x * NR0;
|
|
const int r1 = tgpig.y;
|
|
const int im = tgpig.z;
|
|
|
|
const uint i12 = im % args.ne12;
|
|
const uint i13 = im / args.ne12;
|
|
const uint64_t offset1 = r1 * args.nb11 + i12 * args.nb12 + i13 * args.nb13;
|
|
device const float *y = (device const float *)(src1 + offset1);
|
|
|
|
device const block_q8_0 *ag[NR0];
|
|
device const block_q8_0 *au[NR0];
|
|
FOR_UNROLL (short row = 0; row < NR0; ++row) {
|
|
const uint64_t offset0 = (r0 + row) * args.nb01 +
|
|
(i12 / args.r2) * args.nb02 +
|
|
(i13 / args.r3) * args.nb03;
|
|
ag[row] = (device const block_q8_0 *)((device const char *)src0_gate + offset0);
|
|
au[row] = (device const block_q8_0 *)((device const char *)src0_up + offset0);
|
|
}
|
|
|
|
float sumg[NR0] = { 0.f };
|
|
float sumu[NR0] = { 0.f };
|
|
|
|
const short ix = tiisg / (NW / NQ);
|
|
const short il = tiisg % (NW / NQ);
|
|
const int ib0 = sgitg * NQ + ix;
|
|
float yl[NQ];
|
|
device const float *yb = y + ib0 * QK8_0 + il * NQ;
|
|
|
|
for (int ib = ib0; ib < nb; ib += NSG * NQ) {
|
|
FOR_UNROLL (short i = 0; i < NQ; ++i) {
|
|
yl[i] = yb[i];
|
|
}
|
|
|
|
FOR_UNROLL (short row = 0; row < NR0; ++row) {
|
|
device const int8_t *qg = ag[row][ib].qs + il * NQ;
|
|
device const int8_t *qu = au[row][ib].qs + il * NQ;
|
|
|
|
float sg = 0.f;
|
|
float su = 0.f;
|
|
FOR_UNROLL (short i = 0; i < NQ; ++i) {
|
|
sg += qg[i] * yl[i];
|
|
su += qu[i] * yl[i];
|
|
}
|
|
|
|
sumg[row] += sg * ag[row][ib].d;
|
|
sumu[row] += su * au[row][ib].d;
|
|
}
|
|
|
|
yb += NSG * NQ * QK8_0;
|
|
}
|
|
|
|
threadgroup float *shmem_f32 = (threadgroup float *)shmem;
|
|
threadgroup float *sh_gate[NR0];
|
|
threadgroup float *sh_up[NR0];
|
|
FOR_UNROLL (short row = 0; row < NR0; ++row) {
|
|
sh_gate[row] = shmem_f32 + NW * row;
|
|
sh_up[row] = shmem_f32 + NW * (NR0 + row);
|
|
if (sgitg == 0) {
|
|
sh_gate[row][tiisg] = 0.0f;
|
|
sh_up[row][tiisg] = 0.0f;
|
|
}
|
|
sumg[row] = simd_sum(sumg[row]);
|
|
sumu[row] = simd_sum(sumu[row]);
|
|
}
|
|
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
FOR_UNROLL (short row = 0; row < NR0; ++row) {
|
|
if (tiisg == 0) {
|
|
sh_gate[row][sgitg] = sumg[row];
|
|
sh_up[row][sgitg] = sumu[row];
|
|
}
|
|
}
|
|
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
device float *gate_f32 = (device float *)dst_gate +
|
|
(uint64_t)im * args.ne0 * args.ne1 + (uint64_t)r1 * args.ne0;
|
|
device float *up_f32 = (device float *)dst_up +
|
|
(uint64_t)im * args.ne0 * args.ne1 + (uint64_t)r1 * args.ne0;
|
|
device float *mid_f32 = (device float *)dst_mid +
|
|
(uint64_t)im * args.ne0 * args.ne1 + (uint64_t)r1 * args.ne0;
|
|
|
|
FOR_UNROLL (short row = 0; row < NR0 && r0 + row < args.ne01; ++row) {
|
|
const float gate = simd_sum(sh_gate[row][tiisg]);
|
|
const float up = simd_sum(sh_up[row][tiisg]);
|
|
if (tiisg == 0 && sgitg == 0) {
|
|
const uint out_row = r0 + row;
|
|
if (STORE_GATE_UP) {
|
|
gate_f32[out_row] = gate;
|
|
up_f32[out_row] = up;
|
|
}
|
|
float g = ROUND_GATE_UP ? ds4_bf16_round(gate) : gate;
|
|
float u = ROUND_GATE_UP ? ds4_bf16_round(up) : up;
|
|
if (clamp_value > 1.0e-6f) {
|
|
g = min(g, clamp_value);
|
|
u = clamp(u, -clamp_value, clamp_value);
|
|
}
|
|
const float silu = g / (1.0f + exp(-g));
|
|
mid_f32[out_row] = silu * u;
|
|
}
|
|
}
|
|
}
|
|
|
|
[[host_name("kernel_dsv4_shared_gate_up_swiglu_q8_0")]]
|
|
kernel void kernel_dsv4_shared_gate_up_swiglu_q8_0(
|
|
constant ds4_metal_args_mul_mv & args,
|
|
device const char * src0_gate,
|
|
device const char * src0_up,
|
|
device const char * src1,
|
|
device char * dst_gate,
|
|
device char * dst_up,
|
|
device char * dst_mid,
|
|
constant float &clamp_value,
|
|
threadgroup char * shmem [[threadgroup(0)]],
|
|
uint3 tgpig[[threadgroup_position_in_grid]],
|
|
ushort tiisg[[thread_index_in_simdgroup]],
|
|
ushort sgitg[[simdgroup_index_in_threadgroup]]) {
|
|
kernel_dsv4_shared_gate_up_swiglu_q8_0_impl<N_R0_Q8_0, true>(
|
|
args, src0_gate, src0_up, src1, dst_gate, dst_up, dst_mid,
|
|
clamp_value, shmem, tgpig, tiisg, sgitg);
|
|
}
|
|
|
|
// Same body as kernel_dsv4_shared_gate_up_swiglu_q8_0_impl with the row
|
|
// bound taken from a thread value instead of args.ne01 (GPU-decided lane
|
|
// count). Arguments stay in constant memory.
|
|
template<short NR0, bool STORE_GATE_UP>
|
|
void kernel_dsv4_shared_gate_up_swiglu_q8_0_rows_impl(
|
|
const int row_count,
|
|
constant ds4_metal_args_mul_mv & args,
|
|
device const char * src0_gate,
|
|
device const char * src0_up,
|
|
device const char * src1,
|
|
device char * dst_gate,
|
|
device char * dst_up,
|
|
device char * dst_mid,
|
|
constant float &clamp_value,
|
|
threadgroup char * shmem,
|
|
uint3 tgpig[[threadgroup_position_in_grid]],
|
|
ushort tiisg[[thread_index_in_simdgroup]],
|
|
ushort sgitg[[simdgroup_index_in_threadgroup]]) {
|
|
const short NSG = FC_mul_mv_nsg;
|
|
constexpr short NW = N_SIMDWIDTH;
|
|
constexpr short NQ = 8;
|
|
|
|
const int nb = args.ne00 / QK8_0;
|
|
const int r0 = tgpig.x * NR0;
|
|
const int r1 = tgpig.y;
|
|
const int im = tgpig.z;
|
|
|
|
const uint i12 = im % args.ne12;
|
|
const uint i13 = im / args.ne12;
|
|
const uint64_t offset1 = r1 * args.nb11 + i12 * args.nb12 + i13 * args.nb13;
|
|
device const float *y = (device const float *)(src1 + offset1);
|
|
|
|
device const block_q8_0 *ag[NR0];
|
|
device const block_q8_0 *au[NR0];
|
|
FOR_UNROLL (short row = 0; row < NR0; ++row) {
|
|
const uint64_t offset0 = (r0 + row) * args.nb01 +
|
|
(i12 / args.r2) * args.nb02 +
|
|
(i13 / args.r3) * args.nb03;
|
|
ag[row] = (device const block_q8_0 *)((device const char *)src0_gate + offset0);
|
|
au[row] = (device const block_q8_0 *)((device const char *)src0_up + offset0);
|
|
}
|
|
|
|
float sumg[NR0] = { 0.f };
|
|
float sumu[NR0] = { 0.f };
|
|
|
|
const short ix = tiisg / (NW / NQ);
|
|
const short il = tiisg % (NW / NQ);
|
|
const int ib0 = sgitg * NQ + ix;
|
|
float yl[NQ];
|
|
device const float *yb = y + ib0 * QK8_0 + il * NQ;
|
|
|
|
for (int ib = ib0; ib < nb; ib += NSG * NQ) {
|
|
FOR_UNROLL (short i = 0; i < NQ; ++i) {
|
|
yl[i] = yb[i];
|
|
}
|
|
|
|
FOR_UNROLL (short row = 0; row < NR0; ++row) {
|
|
device const int8_t *qg = ag[row][ib].qs + il * NQ;
|
|
device const int8_t *qu = au[row][ib].qs + il * NQ;
|
|
|
|
float sg = 0.f;
|
|
float su = 0.f;
|
|
FOR_UNROLL (short i = 0; i < NQ; ++i) {
|
|
sg += qg[i] * yl[i];
|
|
su += qu[i] * yl[i];
|
|
}
|
|
|
|
sumg[row] += sg * ag[row][ib].d;
|
|
sumu[row] += su * au[row][ib].d;
|
|
}
|
|
|
|
yb += NSG * NQ * QK8_0;
|
|
}
|
|
|
|
threadgroup float *shmem_f32 = (threadgroup float *)shmem;
|
|
threadgroup float *sh_gate[NR0];
|
|
threadgroup float *sh_up[NR0];
|
|
FOR_UNROLL (short row = 0; row < NR0; ++row) {
|
|
sh_gate[row] = shmem_f32 + NW * row;
|
|
sh_up[row] = shmem_f32 + NW * (NR0 + row);
|
|
if (sgitg == 0) {
|
|
sh_gate[row][tiisg] = 0.0f;
|
|
sh_up[row][tiisg] = 0.0f;
|
|
}
|
|
sumg[row] = simd_sum(sumg[row]);
|
|
sumu[row] = simd_sum(sumu[row]);
|
|
}
|
|
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
FOR_UNROLL (short row = 0; row < NR0; ++row) {
|
|
if (tiisg == 0) {
|
|
sh_gate[row][sgitg] = sumg[row];
|
|
sh_up[row][sgitg] = sumu[row];
|
|
}
|
|
}
|
|
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
device float *gate_f32 = (device float *)dst_gate +
|
|
(uint64_t)im * args.ne0 * args.ne1 + (uint64_t)r1 * args.ne0;
|
|
device float *up_f32 = (device float *)dst_up +
|
|
(uint64_t)im * args.ne0 * args.ne1 + (uint64_t)r1 * args.ne0;
|
|
device float *mid_f32 = (device float *)dst_mid +
|
|
(uint64_t)im * args.ne0 * args.ne1 + (uint64_t)r1 * args.ne0;
|
|
|
|
FOR_UNROLL (short row = 0; row < NR0 && r0 + row < row_count; ++row) {
|
|
const float gate = simd_sum(sh_gate[row][tiisg]);
|
|
const float up = simd_sum(sh_up[row][tiisg]);
|
|
if (tiisg == 0 && sgitg == 0) {
|
|
const uint out_row = r0 + row;
|
|
if (STORE_GATE_UP) {
|
|
gate_f32[out_row] = gate;
|
|
up_f32[out_row] = up;
|
|
}
|
|
float g = gate;
|
|
float u = up;
|
|
if (clamp_value > 1.0e-6f) {
|
|
g = min(g, clamp_value);
|
|
u = clamp(u, -clamp_value, clamp_value);
|
|
}
|
|
const float silu = g / (1.0f + exp(-g));
|
|
mid_f32[out_row] = silu * u;
|
|
}
|
|
}
|
|
}
|
|
|
|
// Same K walk and reduction tree as kernel_mul_mv_q8_0_f32_impl with the
|
|
// K length taken from a thread value (GPU-decided lane count); the caller
|
|
// bumps src0 to the lane base inside every row.
|
|
template<short NR0>
|
|
void kernel_dsv4_shared_down_q8_0_k_impl(
|
|
const int k_count,
|
|
constant ds4_metal_args_mul_mv & args,
|
|
device const char * src0,
|
|
device const char * src1,
|
|
device char * dst,
|
|
threadgroup char * shmem,
|
|
uint3 tgpig,
|
|
ushort tiisg,
|
|
ushort sgitg) {
|
|
const short NSG = FC_mul_mv_nsg;
|
|
|
|
constexpr short NW = N_SIMDWIDTH;
|
|
constexpr short NQ = 8;
|
|
|
|
const int nb = k_count/QK8_0;
|
|
|
|
const int r0 = tgpig.x*NR0;
|
|
const int r1 = tgpig.y;
|
|
const int im = tgpig.z;
|
|
|
|
const uint i12 = im%args.ne12;
|
|
const uint i13 = im/args.ne12;
|
|
|
|
const uint64_t offset1 = r1*args.nb11 + (i12)*args.nb12 + (i13)*args.nb13;
|
|
|
|
device const float * y = (device const float *) (src1 + offset1);
|
|
|
|
device const block_q8_0 * ax[NR0];
|
|
FOR_UNROLL (short row = 0; row < NR0; ++row) {
|
|
const uint64_t offset0 = (r0 + row)*args.nb01 + (i12/args.r2)*args.nb02 + (i13/args.r3)*args.nb03;
|
|
|
|
ax[row] = (device const block_q8_0 *) ((device char *) src0 + offset0);
|
|
}
|
|
|
|
float sumf[NR0] = { 0.f };
|
|
|
|
const short ix = tiisg/(NW/NQ);
|
|
const short il = tiisg%(NW/NQ);
|
|
|
|
const int ib0 = sgitg*NQ + ix;
|
|
|
|
float yl[NQ];
|
|
|
|
device const float * yb = y + ib0*QK8_0 + il*NQ;
|
|
|
|
for (int ib = ib0; ib < nb; ib += NSG*NQ) {
|
|
for (short i = 0; i < NQ; ++i) {
|
|
yl[i] = yb[i];
|
|
}
|
|
|
|
for (short row = 0; row < NR0; row++) {
|
|
device const int8_t * qs = ax[row][ib].qs + il*NQ;
|
|
|
|
float sumq = 0.f;
|
|
FOR_UNROLL (short i = 0; i < NQ; ++i) {
|
|
sumq += qs[i] * yl[i];
|
|
}
|
|
|
|
sumf[row] += sumq*ax[row][ib].d;
|
|
}
|
|
|
|
yb += NSG*NQ*QK8_0;
|
|
}
|
|
|
|
device float * dst_f32 = (device float *) dst + (uint64_t)im*args.ne0*args.ne1 + (uint64_t)r1*args.ne0;
|
|
|
|
helper_mv_reduce_and_write<NR0>(dst_f32, sumf, r0, args.ne01, tiisg, sgitg, shmem);
|
|
}
|
|
|
|
// GPU-decided shared-expert lane split for two-rank TP decode.
|
|
//
|
|
// Routed experts are owned statically by expert id (ds4_tp_owns_expert), so
|
|
// the number of selected experts each rank streams changes per token (0..6
|
|
// of 6) while the shared expert used to be split in fixed halves. The rank
|
|
// owning more selected experts is the critical path of the layer and the
|
|
// other rank idles at the gate. Both ranks read the same selected ids, so
|
|
// every threadgroup recomputes the same partition without any host round
|
|
// trip: rank 0 takes lanes [0, lanes0), rank 1 takes [lanes0, shared_dim),
|
|
// lanes0 = shared_dim/2 + (n1 - n0) * shift * shared_dim, where n0/n1 count
|
|
// the selected experts owned by rank 0/1 and shift = routed expert bytes /
|
|
// (2 * shared expert bytes). With the MXFP4 routed / Q8 shared mix of the
|
|
// V4 Flash file, one extra routed expert moves ~half of the shared expert to
|
|
// the other rank; one owning 4 of 6 hands over the whole shared expert and
|
|
// both ranks stream the same byte count. shift_q16 == 0 reproduces the
|
|
// static halves bit-exactly. The union of both ranges is always the full
|
|
// shared width, and a rank with zero lanes writes a zero partial.
|
|
typedef struct {
|
|
int32_t tp_rank;
|
|
int32_t tp_world;
|
|
int32_t n_expert;
|
|
int32_t n_expert_used;
|
|
int32_t shared_dim;
|
|
int32_t lane_granule;
|
|
int32_t shift_q16;
|
|
} ds4_metal_shared_split_args;
|
|
|
|
static inline void ds4_shared_split_range(
|
|
constant ds4_metal_shared_split_args &sp,
|
|
device const int32_t *ids,
|
|
thread int &base,
|
|
thread int &count) {
|
|
const int per_rank = sp.n_expert / sp.tp_world;
|
|
int n0 = 0;
|
|
for (int i = 0; i < sp.n_expert_used; ++i) {
|
|
n0 += ids[i] < per_rank ? 1 : 0;
|
|
}
|
|
const int n1 = sp.n_expert_used - n0;
|
|
// |n1 - n0| <= 64, shift_q16 <= 65536, shared_dim <= 32768: fits int64.
|
|
const int64_t delta_q16 =
|
|
(int64_t)(n1 - n0) * (int64_t)sp.shift_q16 * (int64_t)sp.shared_dim;
|
|
const int delta = (int)(delta_q16 >= 0 ? (delta_q16 >> 16)
|
|
: -((-delta_q16) >> 16));
|
|
int lanes0 = sp.shared_dim / 2 + delta;
|
|
lanes0 = clamp(lanes0, 0, sp.shared_dim);
|
|
const int granule = sp.lane_granule;
|
|
lanes0 = ((lanes0 + granule / 2) / granule) * granule;
|
|
lanes0 = min(lanes0, sp.shared_dim);
|
|
if (sp.tp_rank == 0) {
|
|
base = 0;
|
|
count = (int)lanes0;
|
|
} else {
|
|
base = (int)lanes0;
|
|
count = sp.shared_dim - (int)lanes0;
|
|
}
|
|
}
|
|
|
|
[[host_name("kernel_dsv4_shared_gate_up_swiglu_q8_0_split")]]
|
|
kernel void kernel_dsv4_shared_gate_up_swiglu_q8_0_split(
|
|
constant ds4_metal_args_mul_mv & args,
|
|
constant ds4_metal_shared_split_args & sp,
|
|
device const char * src0_gate,
|
|
device const char * src0_up,
|
|
device const char * src1,
|
|
device char * dst_gate,
|
|
device char * dst_up,
|
|
device char * dst_mid,
|
|
constant float &clamp_value,
|
|
device const char * ids,
|
|
threadgroup char * shmem [[threadgroup(0)]],
|
|
uint3 tgpig[[threadgroup_position_in_grid]],
|
|
ushort tiisg[[thread_index_in_simdgroup]],
|
|
ushort sgitg[[simdgroup_index_in_threadgroup]]) {
|
|
int base = 0;
|
|
int count = 0;
|
|
ds4_shared_split_range(sp, (device const int32_t *)ids, base, count);
|
|
// Uniform per threadgroup: rows past this rank's lane count exit before
|
|
// any barrier. The grid always covers the full shared width.
|
|
if ((int)tgpig.x * N_R0_Q8_0 >= count) return;
|
|
const uint64_t lane_bytes = (uint64_t)base * (uint64_t)args.nb01;
|
|
kernel_dsv4_shared_gate_up_swiglu_q8_0_rows_impl<N_R0_Q8_0, true>(
|
|
count, args, src0_gate + lane_bytes, src0_up + lane_bytes, src1,
|
|
dst_gate, dst_up, dst_mid, clamp_value, shmem,
|
|
tgpig, tiisg, sgitg);
|
|
}
|
|
|
|
// Shared-expert down projection over this rank's lane range: the K slice
|
|
// starts at the lane base inside every Q8_0 row and runs for the lane
|
|
// count; the SwiGLU intermediate is compact at the buffer base. A rank
|
|
// with no lanes runs an empty K loop and writes zeros for every row.
|
|
[[host_name("kernel_dsv4_shared_down_q8_0_split")]]
|
|
kernel void kernel_dsv4_shared_down_q8_0_split(
|
|
constant ds4_metal_args_mul_mv & args,
|
|
constant ds4_metal_shared_split_args & sp,
|
|
device const char * src0,
|
|
device const char * src1,
|
|
device char * dst,
|
|
device const char * ids,
|
|
threadgroup char * shmem [[threadgroup(0)]],
|
|
uint3 tgpig[[threadgroup_position_in_grid]],
|
|
ushort tiisg[[thread_index_in_simdgroup]],
|
|
ushort sgitg[[simdgroup_index_in_threadgroup]]) {
|
|
int base = 0;
|
|
int count = 0;
|
|
ds4_shared_split_range(sp, (device const int32_t *)ids, base, count);
|
|
const uint64_t k_bytes = ((uint64_t)base / QK8_0) * (uint64_t)args.nb00;
|
|
kernel_dsv4_shared_down_q8_0_k_impl<N_R0_Q8_0>(
|
|
count, args, src0 + k_bytes, src1, dst, shmem, tgpig, tiisg, sgitg);
|
|
}
|
|
|
|
// Decode-only fusion of the router logits matvec (F16, embd -> n_expert)
|
|
// with the shared-expert gate/up SwiGLU (Q8_0, embd -> shared). Both read
|
|
// the same normalized FFN input back to back; one dispatch removes one
|
|
// launch per decode layer. Router threadgroups replicate
|
|
// kernel_mul_mv_f16_f32_4 (nsg=8, nr0=2); shared threadgroups host two
|
|
// virtual 4-simdgroup cohorts replicating
|
|
// kernel_dsv4_shared_gate_up_swiglu_q8_0 (nsg=4, nr0=2), including its
|
|
// per-row simd/shmem reduction trees. Bit-exact by construction.
|
|
kernel void kernel_dsv4_router_shared_gate_up_q8_0(
|
|
constant ds4_metal_args_mul_mv & args,
|
|
constant ds4_metal_args_mul_mv & sargs,
|
|
device const char * src0_router,
|
|
device const char * src0_gate,
|
|
device const char * src0_up,
|
|
device const char * src1,
|
|
device char * dst_router,
|
|
device char * dst_gate,
|
|
device char * dst_up,
|
|
device char * dst_mid,
|
|
constant float &clamp_value,
|
|
threadgroup char * shmem [[threadgroup(0)]],
|
|
uint3 tgpig [[threadgroup_position_in_grid]],
|
|
ushort tiisg [[thread_index_in_simdgroup]],
|
|
ushort sgitg [[simdgroup_index_in_threadgroup]]) {
|
|
constexpr short NW = N_SIMDWIDTH;
|
|
const uint router_tgs = ((uint)args.ne01 + 1u) / 2u;
|
|
|
|
if (tgpig.x < router_tgs) {
|
|
// Exact replica of kernel_mul_mv_f16_f32_4 with NSG=8, NR0=2.
|
|
constexpr short NSG = 8;
|
|
constexpr short NR0 = 2;
|
|
constexpr short NB = 32;
|
|
constexpr short NF = 16;
|
|
constexpr short NF4 = NF/4;
|
|
|
|
const int nb = args.ne00/NB;
|
|
const int r0 = tgpig.x*NR0;
|
|
|
|
device const float4 * y4 = (device const float4 *) src1;
|
|
|
|
device const half4 * ax4[NR0];
|
|
FOR_UNROLL (short row = 0; row < NR0; ++row) {
|
|
ax4[row] = (device const half4 *)
|
|
(src0_router + (uint64_t)(r0 + row)*args.nb01);
|
|
}
|
|
|
|
float sumf[NR0] = { 0.f };
|
|
|
|
const short ix = tiisg/(NW/NF);
|
|
const short il = tiisg%(NW/NF);
|
|
const int ib0 = sgitg*NF + ix;
|
|
|
|
device const float4 * yb4 = y4 + (ib0*NB + il*NF)/4;
|
|
|
|
for (int ib = ib0; ib < nb; ib += NSG*NF) {
|
|
float4 yl4[NF4];
|
|
FOR_UNROLL (short i = 0; i < NF4; ++i) {
|
|
yl4[i] = yb4[i];
|
|
}
|
|
|
|
FOR_UNROLL (short row = 0; row < NR0; row++) {
|
|
device const half4 * xb4 = ax4[row] + (ib*NB + il*NF)/4;
|
|
|
|
float sumq = 0.f;
|
|
FOR_UNROLL (short i = 0; i < NF4; ++i) {
|
|
sumq += dot(float4(xb4[i]), yl4[i]);
|
|
}
|
|
|
|
sumf[row] += sumq;
|
|
}
|
|
|
|
yb4 += NSG*NF*NW/4;
|
|
}
|
|
|
|
device float * dst_f32 = (device float *) dst_router;
|
|
helper_mv_reduce_and_write<NR0>(dst_f32, sumf, r0, args.ne01,
|
|
tiisg, sgitg, shmem);
|
|
return;
|
|
}
|
|
|
|
// Shared-expert part: two virtual nsg=4 cohorts per threadgroup, each an
|
|
// exact replica of kernel_dsv4_shared_gate_up_swiglu_q8_0 (NR0=2).
|
|
constexpr short NSG = 4;
|
|
constexpr short NR0 = 2;
|
|
constexpr short NQ = 8;
|
|
|
|
const uint cohort = sgitg >> 2;
|
|
const ushort vsg = sgitg & 3u;
|
|
const uint vt = (tgpig.x - router_tgs) * 2u + cohort;
|
|
|
|
const int nb = sargs.ne00 / QK8_0;
|
|
const int r0 = vt * NR0;
|
|
|
|
device const float *y = (device const float *) src1;
|
|
|
|
device const block_q8_0 *ag[NR0];
|
|
device const block_q8_0 *au[NR0];
|
|
FOR_UNROLL (short row = 0; row < NR0; ++row) {
|
|
const uint64_t offset0 = (uint64_t)(r0 + row) * sargs.nb01;
|
|
ag[row] = (device const block_q8_0 *)(src0_gate + offset0);
|
|
au[row] = (device const block_q8_0 *)(src0_up + offset0);
|
|
}
|
|
|
|
float sumg[NR0] = { 0.f };
|
|
float sumu[NR0] = { 0.f };
|
|
|
|
const short ix = tiisg / (NW / NQ);
|
|
const short il = tiisg % (NW / NQ);
|
|
const int ib0 = vsg * NQ + ix;
|
|
float yl[NQ];
|
|
device const float *yb = y + ib0 * QK8_0 + il * NQ;
|
|
|
|
for (int ib = ib0; ib < nb; ib += NSG * NQ) {
|
|
FOR_UNROLL (short i = 0; i < NQ; ++i) {
|
|
yl[i] = yb[i];
|
|
}
|
|
|
|
FOR_UNROLL (short row = 0; row < NR0; ++row) {
|
|
device const int8_t *qg = ag[row][ib].qs + il * NQ;
|
|
device const int8_t *qu = au[row][ib].qs + il * NQ;
|
|
|
|
float sg = 0.f;
|
|
float su = 0.f;
|
|
FOR_UNROLL (short i = 0; i < NQ; ++i) {
|
|
sg += qg[i] * yl[i];
|
|
su += qu[i] * yl[i];
|
|
}
|
|
|
|
sumg[row] += sg * ag[row][ib].d;
|
|
sumu[row] += su * au[row][ib].d;
|
|
}
|
|
|
|
yb += NSG * NQ * QK8_0;
|
|
}
|
|
|
|
threadgroup float *shmem_f32 = (threadgroup float *)shmem + cohort * (2*NR0*NW);
|
|
threadgroup float *sh_gate[NR0];
|
|
threadgroup float *sh_up[NR0];
|
|
FOR_UNROLL (short row = 0; row < NR0; ++row) {
|
|
sh_gate[row] = shmem_f32 + NW * row;
|
|
sh_up[row] = shmem_f32 + NW * (NR0 + row);
|
|
if (vsg == 0) {
|
|
sh_gate[row][tiisg] = 0.0f;
|
|
sh_up[row][tiisg] = 0.0f;
|
|
}
|
|
sumg[row] = simd_sum(sumg[row]);
|
|
sumu[row] = simd_sum(sumu[row]);
|
|
}
|
|
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
FOR_UNROLL (short row = 0; row < NR0; ++row) {
|
|
if (tiisg == 0) {
|
|
sh_gate[row][vsg] = sumg[row];
|
|
sh_up[row][vsg] = sumu[row];
|
|
}
|
|
}
|
|
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
device float *gate_f32 = (device float *)dst_gate;
|
|
device float *up_f32 = (device float *)dst_up;
|
|
device float *mid_f32 = (device float *)dst_mid;
|
|
|
|
FOR_UNROLL (short row = 0; row < NR0 && r0 + row < sargs.ne01; ++row) {
|
|
const float gate = simd_sum(sh_gate[row][tiisg]);
|
|
const float up = simd_sum(sh_up[row][tiisg]);
|
|
if (tiisg == 0 && vsg == 0) {
|
|
const uint out_row = r0 + row;
|
|
gate_f32[out_row] = gate;
|
|
up_f32[out_row] = up;
|
|
float g = gate;
|
|
float u = up;
|
|
if (clamp_value > 1.0e-6f) {
|
|
g = min(g, clamp_value);
|
|
u = clamp(u, -clamp_value, clamp_value);
|
|
}
|
|
const float silu = g / (1.0f + exp(-g));
|
|
mid_f32[out_row] = silu * u;
|
|
}
|
|
}
|
|
}
|
|
|
|
|
|
[[host_name("kernel_dsv4_shared_mid_swiglu_q8_0")]]
|
|
kernel void kernel_dsv4_shared_mid_swiglu_q8_0(
|
|
constant ds4_metal_args_mul_mv & args,
|
|
device const char * src0_gate,
|
|
device const char * src0_up,
|
|
device const char * src1,
|
|
device char * dst_gate,
|
|
device char * dst_up,
|
|
device char * dst_mid,
|
|
constant float &clamp_value,
|
|
threadgroup char * shmem [[threadgroup(0)]],
|
|
uint3 tgpig[[threadgroup_position_in_grid]],
|
|
ushort tiisg[[thread_index_in_simdgroup]],
|
|
ushort sgitg[[simdgroup_index_in_threadgroup]]) {
|
|
kernel_dsv4_shared_gate_up_swiglu_q8_0_impl<N_R0_Q8_0, false>(
|
|
args, src0_gate, src0_up, src1, dst_gate, dst_up, dst_mid,
|
|
clamp_value, shmem, tgpig, tiisg, sgitg);
|
|
}
|
|
|
|
// The gate and up sums round to bf16 before the SwiGLU, as the separate
|
|
// bf16 matvecs would have stored them.
|
|
[[host_name("kernel_dsv4_shared_mid_swiglu_q8_0_bf16")]]
|
|
kernel void kernel_dsv4_shared_mid_swiglu_q8_0_bf16(
|
|
constant ds4_metal_args_mul_mv & args,
|
|
device const char * src0_gate,
|
|
device const char * src0_up,
|
|
device const char * src1,
|
|
device char * dst_gate,
|
|
device char * dst_up,
|
|
device char * dst_mid,
|
|
constant float &clamp_value,
|
|
threadgroup char * shmem [[threadgroup(0)]],
|
|
uint3 tgpig[[threadgroup_position_in_grid]],
|
|
ushort tiisg[[thread_index_in_simdgroup]],
|
|
ushort sgitg[[simdgroup_index_in_threadgroup]]) {
|
|
kernel_dsv4_shared_gate_up_swiglu_q8_0_impl<N_R0_Q8_0, false, true>(
|
|
args, src0_gate, src0_up, src1, dst_gate, dst_up, dst_mid,
|
|
clamp_value, shmem, tgpig, tiisg, sgitg);
|
|
}
|
|
|
|
// A few rows of the bf16 shared expert: grid x is the row so that the rows'
|
|
// threadgroups for the same weight rows run back to back.
|
|
[[host_name("kernel_dsv4_shared_mid_swiglu_q8_0_bf16_rows")]]
|
|
kernel void kernel_dsv4_shared_mid_swiglu_q8_0_bf16_rows(
|
|
constant ds4_metal_args_mul_mv & args,
|
|
device const char * src0_gate,
|
|
device const char * src0_up,
|
|
device const char * src1,
|
|
device char * dst_gate,
|
|
device char * dst_up,
|
|
device char * dst_mid,
|
|
constant float &clamp_value,
|
|
threadgroup char * shmem [[threadgroup(0)]],
|
|
uint3 tgpig[[threadgroup_position_in_grid]],
|
|
ushort tiisg[[thread_index_in_simdgroup]],
|
|
ushort sgitg[[simdgroup_index_in_threadgroup]]) {
|
|
kernel_dsv4_shared_gate_up_swiglu_q8_0_impl<N_R0_Q8_0, false, true>(
|
|
args, src0_gate, src0_up, src1, dst_gate, dst_up, dst_mid,
|
|
clamp_value, shmem, uint3(tgpig.y, tgpig.x, tgpig.z), tiisg, sgitg);
|
|
}
|
|
|
|
// The matrix rows form of the bf16 shared expert: the gate and up rows
|
|
// ride the same activation tiles, the SwiGLU follows the reduce.
|
|
template<short NSG>
|
|
void kernel_dsv4_shared_mid_swiglu_q8_0_bf16_mma_impl(
|
|
constant ds4_metal_args_mul_mv & args,
|
|
device const char * src0_gate,
|
|
device const char * src0_up,
|
|
device const char * src1,
|
|
device char * dst_mid,
|
|
constant float &clamp_value,
|
|
threadgroup char * shmem,
|
|
uint3 tgpig,
|
|
ushort tiisg,
|
|
ushort sgitg) {
|
|
constexpr short NR = 8, KC = QK8_0;
|
|
threadgroup float * A = (threadgroup float *) shmem + sgitg * (3 * NR * KC);
|
|
threadgroup float * Wg = A + NR * KC;
|
|
threadgroup float * Wu = Wg + NR * KC;
|
|
const int r0 = (int)tgpig.x * NR;
|
|
const int kq = args.ne00 / NSG;
|
|
const int k_begin = sgitg * kq, k_end = k_begin + kq;
|
|
const short wr = tiisg / 4, wq = tiisg % 4;
|
|
device const block_q8_0 * grow = (device const block_q8_0 *)(src0_gate + (uint64_t)(r0 + wr) * args.nb01);
|
|
device const block_q8_0 * urow = (device const block_q8_0 *)(src0_up + (uint64_t)(r0 + wr) * args.nb01);
|
|
const bool have_row = wr < args.ne11;
|
|
device const float * y = (device const float *)(src1 + (uint64_t)wr * args.nb11) + wq * 8;
|
|
threadgroup float * a = A + wq * 64 + wr * 8;
|
|
threadgroup float * wg = Wg + wq * 64 + wr * 8;
|
|
threadgroup float * wu = Wu + wq * 64 + wr * 8;
|
|
simdgroup_float8x8 accg = make_filled_simdgroup_matrix<float, 8, 8>(0.0f);
|
|
simdgroup_float8x8 accu = make_filled_simdgroup_matrix<float, 8, 8>(0.0f);
|
|
float av[8];
|
|
float4 g0, g1, u0, u1;
|
|
float gd, ud;
|
|
{
|
|
device const block_q8_0 * gb = grow + k_begin / QK8_0, * ub = urow + k_begin / QK8_0;
|
|
device const packed_char4 * gq = (device const packed_char4 *)(gb->qs + wq * 8);
|
|
device const packed_char4 * uq = (device const packed_char4 *)(ub->qs + wq * 8);
|
|
gd = (float)gb->d; g0 = float4(gq[0]); g1 = float4(gq[1]);
|
|
ud = (float)ub->d; u0 = float4(uq[0]); u1 = float4(uq[1]);
|
|
for (short i = 0; i < 8; ++i) av[i] = have_row ? y[k_begin + i] : 0.0f;
|
|
}
|
|
for (int k0 = k_begin; k0 < k_end; k0 += KC) {
|
|
for (short i = 0; i < 8; ++i) a[i] = av[i];
|
|
const float4 vg0 = gd * g0, vg1 = gd * g1, vu0 = ud * u0, vu1 = ud * u1;
|
|
wg[0] = vg0.x; wg[1] = vg0.y; wg[2] = vg0.z; wg[3] = vg0.w;
|
|
wg[4] = vg1.x; wg[5] = vg1.y; wg[6] = vg1.z; wg[7] = vg1.w;
|
|
wu[0] = vu0.x; wu[1] = vu0.y; wu[2] = vu0.z; wu[3] = vu0.w;
|
|
wu[4] = vu1.x; wu[5] = vu1.y; wu[6] = vu1.z; wu[7] = vu1.w;
|
|
if (k0 + KC < k_end) {
|
|
device const block_q8_0 * gb = grow + (k0 + KC) / QK8_0, * ub = urow + (k0 + KC) / QK8_0;
|
|
device const packed_char4 * gq = (device const packed_char4 *)(gb->qs + wq * 8);
|
|
device const packed_char4 * uq = (device const packed_char4 *)(ub->qs + wq * 8);
|
|
gd = (float)gb->d; g0 = float4(gq[0]); g1 = float4(gq[1]);
|
|
ud = (float)ub->d; u0 = float4(uq[0]); u1 = float4(uq[1]);
|
|
for (short i = 0; i < 8; ++i) av[i] = have_row ? y[k0 + KC + i] : 0.0f;
|
|
}
|
|
simdgroup_barrier(mem_flags::mem_threadgroup);
|
|
for (short kt = 0; kt < KC / 8; ++kt) {
|
|
simdgroup_float8x8 at, wt;
|
|
simdgroup_load(at, A + kt * 64, 8, 0, true);
|
|
simdgroup_load(wt, Wg + kt * 64, 8, 0, false);
|
|
simdgroup_multiply_accumulate(accg, wt, at, accg);
|
|
simdgroup_load(wt, Wu + kt * 64, 8, 0, false);
|
|
simdgroup_multiply_accumulate(accu, wt, at, accu);
|
|
}
|
|
simdgroup_barrier(mem_flags::mem_threadgroup);
|
|
}
|
|
threadgroup float * P = (threadgroup float *) shmem + NSG * 3 * NR * KC;
|
|
simdgroup_store(accg, P + sgitg * 64, 8, 0, false);
|
|
simdgroup_store(accu, P + (NSG + sgitg) * 64, 8, 0, false);
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
if (sgitg == 0) {
|
|
const short j = tiisg / 4;
|
|
if (j < args.ne11) {
|
|
device float * mid_f32 = (device float *) dst_mid + (uint64_t)j * args.ne0;
|
|
for (short i = (tiisg % 4) * 2; i < (tiisg % 4) * 2 + 2; ++i) {
|
|
float gate = P[i * 8 + j], up = P[NSG * 64 + i * 8 + j];
|
|
for (short g = 1; g < NSG; ++g) {
|
|
gate += P[g * 64 + i * 8 + j];
|
|
up += P[(NSG + g) * 64 + i * 8 + j];
|
|
}
|
|
float gv = ds4_bf16_round(gate), uv = ds4_bf16_round(up);
|
|
if (clamp_value > 1.0e-6f) {
|
|
gv = min(gv, clamp_value);
|
|
uv = clamp(uv, -clamp_value, clamp_value);
|
|
}
|
|
const float silu = gv / (1.0f + exp(-gv));
|
|
mid_f32[r0 + i] = silu * uv;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
[[host_name("kernel_dsv4_shared_mid_swiglu_q8_0_bf16_mma_rows8")]]
|
|
kernel void kernel_dsv4_shared_mid_swiglu_q8_0_bf16_mma_rows8(
|
|
constant ds4_metal_args_mul_mv & args,
|
|
device const char * src0_gate,
|
|
device const char * src0_up,
|
|
device const char * src1,
|
|
device char * dst_gate,
|
|
device char * dst_up,
|
|
device char * dst_mid,
|
|
constant float &clamp_value,
|
|
threadgroup char * shmem [[threadgroup(0)]],
|
|
uint3 tgpig[[threadgroup_position_in_grid]],
|
|
ushort tiisg[[thread_index_in_simdgroup]],
|
|
ushort sgitg[[simdgroup_index_in_threadgroup]]) {
|
|
(void)dst_gate; (void)dst_up;
|
|
kernel_dsv4_shared_mid_swiglu_q8_0_bf16_mma_impl<8>(args, src0_gate, src0_up, src1, dst_mid,
|
|
clamp_value, shmem, tgpig, tiisg, sgitg);
|
|
}
|
|
|
|
template<typename T0, typename T1, short NR0, typename args_t>
|
|
void kernel_mul_mv_t_t_impl(
|
|
args_t args,
|
|
device const char * src0,
|
|
device const char * src1,
|
|
device char * dst,
|
|
threadgroup char * shmem,
|
|
uint3 tgpig,
|
|
ushort tiisg,
|
|
ushort sgitg) {
|
|
const short NSG = FC_mul_mv_nsg;
|
|
|
|
constexpr short NW = N_SIMDWIDTH;
|
|
constexpr short NB = 32;
|
|
constexpr short NF = 8;
|
|
|
|
const int nb = args.ne00/NB;
|
|
|
|
const int r0 = tgpig.x*NR0;
|
|
const int r1 = tgpig.y;
|
|
const int im = tgpig.z;
|
|
|
|
const uint i12 = im%args.ne12;
|
|
const uint i13 = im/args.ne12;
|
|
|
|
const uint64_t offset1 = r1*args.nb11 + (i12)*args.nb12 + (i13)*args.nb13;
|
|
|
|
device const T1 * y = (device const T1 *) (src1 + offset1);
|
|
|
|
device const T0 * ax[NR0];
|
|
FOR_UNROLL (short row = 0; row < NR0; ++row) {
|
|
const uint64_t offset0 = (r0 + row)*args.nb01 + (i12/args.r2)*args.nb02 + (i13/args.r3)*args.nb03;
|
|
|
|
ax[row] = (device const T0 *) ((device char *) src0 + offset0);
|
|
}
|
|
|
|
float sumf[NR0] = { 0.f };
|
|
|
|
const short ix = tiisg/(NW/NF);
|
|
const short il = tiisg%(NW/NF);
|
|
|
|
const int ib0 = sgitg*NF + ix;
|
|
|
|
T1 yl[NF];
|
|
|
|
device const T1 * yb = y + (ib0*NB + il*NF);
|
|
|
|
for (int ib = ib0; ib < nb; ib += NSG*NF) {
|
|
for (short i = 0; i < NF; ++i) {
|
|
yl[i] = yb[i];
|
|
}
|
|
|
|
for (short row = 0; row < NR0; row++) {
|
|
device const T0 * xb = ax[row] + (ib*NB + il*NF);
|
|
|
|
float sumq = 0.f;
|
|
FOR_UNROLL (short i = 0; i < NF; ++i) {
|
|
sumq += xb[i] * yl[i];
|
|
}
|
|
|
|
sumf[row] += sumq;
|
|
}
|
|
|
|
yb += NSG*NF*NW;
|
|
}
|
|
|
|
for (int i = nb*NB + sgitg*NW + tiisg; i < args.ne00; i += NW*NSG) {
|
|
for (short row = 0; row < NR0; row++) {
|
|
sumf[row] += ax[row][i] * y[i];
|
|
}
|
|
}
|
|
|
|
device float * dst_f32 = (device float *) dst + (uint64_t)im*args.ne0*args.ne1 + (uint64_t)r1*args.ne0;
|
|
|
|
helper_mv_reduce_and_write<NR0>(dst_f32, sumf, r0, args.ne01, tiisg, sgitg, shmem);
|
|
}
|
|
|
|
template<typename T0, typename T1, typename args_t>
|
|
void kernel_mul_mv_t_t_disp(
|
|
args_t args,
|
|
device const char * src0,
|
|
device const char * src1,
|
|
device char * dst,
|
|
threadgroup char * shmem,
|
|
uint3 tgpig,
|
|
ushort tiisg,
|
|
ushort sgitg) {
|
|
switch (args.nr0) {
|
|
case 2: kernel_mul_mv_t_t_impl<T0, T1, 2, args_t>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg); break;
|
|
case 4: kernel_mul_mv_t_t_impl<T0, T1, 4, args_t>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg); break;
|
|
}
|
|
}
|
|
|
|
// Decode-time dense F32/F16 matrix-vector multiply. The instantiated kernels
|
|
// handle unquantized DS4 weights and activations that are already float rows.
|
|
template<typename T0, typename T1>
|
|
kernel void kernel_mul_mv_t_t(
|
|
constant ds4_metal_args_mul_mv & args,
|
|
device const char * src0,
|
|
device const char * src1,
|
|
device char * dst,
|
|
threadgroup char * shmem [[threadgroup(0)]],
|
|
uint3 tgpig[[threadgroup_position_in_grid]],
|
|
ushort tiisg[[thread_index_in_simdgroup]],
|
|
ushort sgitg[[simdgroup_index_in_threadgroup]]) {
|
|
kernel_mul_mv_t_t_disp<T0, T1, constant ds4_metal_args_mul_mv &>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
|
|
}
|
|
|
|
typedef decltype(kernel_mul_mv_t_t<half, half>) mul_mv_t_t;
|
|
|
|
// Host-visible dense matvec variants used by the graph for F32 and F16 weights.
|
|
template [[host_name("kernel_mul_mv_f32_f32")]] kernel mul_mv_t_t kernel_mul_mv_t_t<float, float>;
|
|
template [[host_name("kernel_mul_mv_f16_f32")]] kernel mul_mv_t_t kernel_mul_mv_t_t<half, float>;
|
|
|
|
template<typename T0, typename T04, typename T1, typename T14, short NR0, typename args_t, bool ROUND = false>
|
|
void kernel_mul_mv_t_t_4_impl(
|
|
args_t args,
|
|
device const char * src0,
|
|
device const char * src1,
|
|
device char * dst,
|
|
threadgroup char * shmem,
|
|
uint3 tgpig,
|
|
ushort tiisg,
|
|
ushort sgitg) {
|
|
const short NSG = FC_mul_mv_nsg;
|
|
|
|
constexpr short NW = N_SIMDWIDTH;
|
|
constexpr short NB = 32;
|
|
constexpr short NF = 16;
|
|
constexpr short NF4 = NF/4;
|
|
|
|
const int nb = args.ne00/NB;
|
|
|
|
const int r0 = tgpig.x*NR0;
|
|
const int r1 = tgpig.y;
|
|
const int im = tgpig.z;
|
|
|
|
const uint i12 = im%args.ne12;
|
|
const uint i13 = im/args.ne12;
|
|
|
|
const uint64_t offset1 = r1*args.nb11 + (i12)*args.nb12 + (i13)*args.nb13;
|
|
|
|
device const T1 * y = (device const T1 *) (src1 + offset1);
|
|
device const T14 * y4 = (device const T14 *) (src1 + offset1);
|
|
|
|
device const T0 * ax [NR0];
|
|
device const T04 * ax4[NR0];
|
|
FOR_UNROLL (short row = 0; row < NR0; ++row) {
|
|
const uint64_t offset0 = (r0 + row)*args.nb01 + (i12/args.r2)*args.nb02 + (i13/args.r3)*args.nb03;
|
|
|
|
ax [row] = (device const T0 *) ((device char *) src0 + offset0);
|
|
ax4[row] = (device const T04 *) ((device char *) src0 + offset0);
|
|
}
|
|
|
|
float sumf[NR0] = { 0.f };
|
|
|
|
const short ix = tiisg/(NW/NF);
|
|
const short il = tiisg%(NW/NF);
|
|
|
|
const int ib0 = sgitg*NF + ix;
|
|
|
|
T14 yl4[NF4];
|
|
|
|
device const T14 * yb4 = y4 + (ib0*NB + il*NF)/4;
|
|
|
|
for (int ib = ib0; ib < nb; ib += NSG*NF) {
|
|
for (short i = 0; i < NF4; ++i) {
|
|
yl4[i] = yb4[i];
|
|
}
|
|
|
|
for (short row = 0; row < NR0; row++) {
|
|
device const T04 * xb4 = ax4[row] + (ib*NB + il*NF)/4;
|
|
|
|
float sumq = 0.f;
|
|
FOR_UNROLL (short i = 0; i < NF4; ++i) {
|
|
sumq += dot(float4(xb4[i]), float4(yl4[i]));
|
|
}
|
|
|
|
sumf[row] += sumq;
|
|
}
|
|
|
|
yb4 += NSG*NF*NW/4;
|
|
}
|
|
|
|
for (int i = nb*NB + sgitg*NW + tiisg; i < args.ne00; i += NW*NSG) {
|
|
for (short row = 0; row < NR0; row++) {
|
|
sumf[row] += ax[row][i] * y[i];
|
|
}
|
|
}
|
|
|
|
device float * dst_f32 = (device float *) dst + (uint64_t)im*args.ne0*args.ne1 + (uint64_t)r1*args.ne0;
|
|
|
|
helper_mv_reduce_and_write<NR0, ROUND>(dst_f32, sumf, r0, args.ne01, tiisg, sgitg, shmem);
|
|
}
|
|
|
|
template<typename T0, typename T04, typename T1, typename T14, typename args_t, bool ROUND = false>
|
|
void kernel_mul_mv_t_t_4_disp(
|
|
args_t args,
|
|
device const char * src0,
|
|
device const char * src1,
|
|
device char * dst,
|
|
threadgroup char * shmem,
|
|
uint3 tgpig,
|
|
ushort tiisg,
|
|
ushort sgitg) {
|
|
switch (args.nr0) {
|
|
case 1: kernel_mul_mv_t_t_4_impl<T0, T04, T1, T14, 1, args_t, ROUND>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg); break;
|
|
case 2: kernel_mul_mv_t_t_4_impl<T0, T04, T1, T14, 2, args_t, ROUND>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg); break;
|
|
case 4: kernel_mul_mv_t_t_4_impl<T0, T04, T1, T14, 4, args_t, ROUND>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg); break;
|
|
};
|
|
}
|
|
|
|
// Vectorized dense matvec using float4/half4 loads. DS4 uses this where the
|
|
// inner dimension and alignment make vector loads cheaper than scalar lanes.
|
|
template<typename T0, typename T04, typename T1, typename T14>
|
|
kernel void kernel_mul_mv_t_t_4(
|
|
constant ds4_metal_args_mul_mv & args,
|
|
device const char * src0,
|
|
device const char * src1,
|
|
device char * dst,
|
|
threadgroup char * shmem [[threadgroup(0)]],
|
|
uint3 tgpig[[threadgroup_position_in_grid]],
|
|
ushort tiisg[[thread_index_in_simdgroup]],
|
|
ushort sgitg[[simdgroup_index_in_threadgroup]]) {
|
|
kernel_mul_mv_t_t_4_disp<T0, T04, T1, T14, constant ds4_metal_args_mul_mv &>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
|
|
}
|
|
|
|
typedef decltype(kernel_mul_mv_t_t_4<half, half4, half, half4>) mul_mv_t_t_4;
|
|
|
|
// Host-visible vectorized dense matvec variants for F32 and F16 weights.
|
|
template [[host_name("kernel_mul_mv_f32_f32_4")]] kernel mul_mv_t_t_4 kernel_mul_mv_t_t_4<float, float4, float, float4>;
|
|
template [[host_name("kernel_mul_mv_f16_f32_4")]] kernel mul_mv_t_t_4 kernel_mul_mv_t_t_4<half, half4, float, float4>;
|
|
|
|
// The same matvec with its output rounded to bf16.
|
|
template<typename T0, typename T04, typename T1, typename T14>
|
|
kernel void kernel_mul_mv_t_t_4_bf16(
|
|
constant ds4_metal_args_mul_mv & args,
|
|
device const char * src0,
|
|
device const char * src1,
|
|
device char * dst,
|
|
threadgroup char * shmem [[threadgroup(0)]],
|
|
uint3 tgpig[[threadgroup_position_in_grid]],
|
|
ushort tiisg[[thread_index_in_simdgroup]],
|
|
ushort sgitg[[simdgroup_index_in_threadgroup]]) {
|
|
kernel_mul_mv_t_t_4_disp<T0, T04, T1, T14, constant ds4_metal_args_mul_mv &, true>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
|
|
}
|
|
|
|
typedef decltype(kernel_mul_mv_t_t_4_bf16<half, half4, half, half4>) mul_mv_t_t_4_bf16;
|
|
|
|
template [[host_name("kernel_mul_mv_f16_f32_4_bf16")]] kernel mul_mv_t_t_4_bf16 kernel_mul_mv_t_t_4_bf16<half, half4, float, float4>;
|
|
|
|
// The rows form of the vectorized dense matvec: one simdgroup group per
|
|
// activation row, each with the single-row kernel's K walk and reduction.
|
|
template<typename T0, typename T04, typename T1, typename T14, short NR0, typename args_t, bool ROUND = false>
|
|
void kernel_mul_mv_t_t_4_rows_impl(
|
|
args_t args,
|
|
device const char * src0,
|
|
device const char * src1,
|
|
device char * dst,
|
|
threadgroup char * shmem,
|
|
uint3 tgpig,
|
|
ushort tiisg,
|
|
ushort sgitg) {
|
|
const short NSG = FC_mul_mv_nsg;
|
|
|
|
constexpr short NW = N_SIMDWIDTH;
|
|
constexpr short NB = 32;
|
|
constexpr short NF = 16;
|
|
constexpr short NF4 = NF/4;
|
|
|
|
const int nb = args.ne00/NB;
|
|
|
|
const int r0 = tgpig.x*NR0;
|
|
const int r1 = sgitg/NSG;
|
|
const ushort sg = sgitg%NSG;
|
|
const int im = tgpig.z;
|
|
|
|
const uint i12 = im%args.ne12;
|
|
const uint i13 = im/args.ne12;
|
|
|
|
const uint64_t offset1 = r1*args.nb11 + (i12)*args.nb12 + (i13)*args.nb13;
|
|
|
|
device const T1 * y = (device const T1 *) (src1 + offset1);
|
|
device const T14 * y4 = (device const T14 *) (src1 + offset1);
|
|
|
|
device const T0 * ax [NR0];
|
|
device const T04 * ax4[NR0];
|
|
FOR_UNROLL (short row = 0; row < NR0; ++row) {
|
|
const uint64_t offset0 = (r0 + row)*args.nb01 + (i12/args.r2)*args.nb02 + (i13/args.r3)*args.nb03;
|
|
|
|
ax [row] = (device const T0 *) ((device char *) src0 + offset0);
|
|
ax4[row] = (device const T04 *) ((device char *) src0 + offset0);
|
|
}
|
|
|
|
float sumf[NR0] = { 0.f };
|
|
|
|
const short ix = tiisg/(NW/NF);
|
|
const short il = tiisg%(NW/NF);
|
|
|
|
const int ib0 = sg*NF + ix;
|
|
|
|
T14 yl4[NF4];
|
|
|
|
device const T14 * yb4 = y4 + (ib0*NB + il*NF)/4;
|
|
|
|
for (int ib = ib0; ib < nb; ib += NSG*NF) {
|
|
for (short i = 0; i < NF4; ++i) {
|
|
yl4[i] = yb4[i];
|
|
}
|
|
|
|
for (short row = 0; row < NR0; row++) {
|
|
device const T04 * xb4 = ax4[row] + (ib*NB + il*NF)/4;
|
|
|
|
float sumq = 0.f;
|
|
FOR_UNROLL (short i = 0; i < NF4; ++i) {
|
|
sumq += dot(float4(xb4[i]), float4(yl4[i]));
|
|
}
|
|
|
|
sumf[row] += sumq;
|
|
}
|
|
|
|
yb4 += NSG*NF*NW/4;
|
|
}
|
|
|
|
for (int i = nb*NB + sg*NW + tiisg; i < args.ne00; i += NW*NSG) {
|
|
for (short row = 0; row < NR0; row++) {
|
|
sumf[row] += ax[row][i] * y[i];
|
|
}
|
|
}
|
|
|
|
device float * dst_f32 = (device float *) dst + (uint64_t)im*args.ne0*args.ne1 + (uint64_t)r1*args.ne0;
|
|
|
|
helper_mv_reduce_and_write_rows<NR0, ROUND>(dst_f32, sumf, r0, args.ne01, tiisg, sg, (short)r1, shmem);
|
|
}
|
|
|
|
template<typename T0, typename T04, typename T1, typename T14>
|
|
kernel void kernel_mul_mv_t_t_4_rows(
|
|
constant ds4_metal_args_mul_mv & args,
|
|
device const char * src0,
|
|
device const char * src1,
|
|
device char * dst,
|
|
threadgroup char * shmem [[threadgroup(0)]],
|
|
uint3 tgpig[[threadgroup_position_in_grid]],
|
|
ushort tiisg[[thread_index_in_simdgroup]],
|
|
ushort sgitg[[simdgroup_index_in_threadgroup]]) {
|
|
switch (args.nr0) {
|
|
case 1: kernel_mul_mv_t_t_4_rows_impl<T0, T04, T1, T14, 1, constant ds4_metal_args_mul_mv &>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg); break;
|
|
case 2: kernel_mul_mv_t_t_4_rows_impl<T0, T04, T1, T14, 2, constant ds4_metal_args_mul_mv &>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg); break;
|
|
case 4: kernel_mul_mv_t_t_4_rows_impl<T0, T04, T1, T14, 4, constant ds4_metal_args_mul_mv &>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg); break;
|
|
};
|
|
}
|
|
|
|
typedef decltype(kernel_mul_mv_t_t_4_rows<half, half4, half, half4>) mul_mv_t_t_4_rows;
|
|
|
|
template [[host_name("kernel_mul_mv_f16_f32_4_rows")]] kernel mul_mv_t_t_4_rows kernel_mul_mv_t_t_4_rows<half, half4, float, float4>;
|
|
|
|
template<typename T0, typename T04, typename T1, typename T14>
|
|
kernel void kernel_mul_mv_t_t_4_rows_bf16(
|
|
constant ds4_metal_args_mul_mv & args,
|
|
device const char * src0,
|
|
device const char * src1,
|
|
device char * dst,
|
|
threadgroup char * shmem [[threadgroup(0)]],
|
|
uint3 tgpig[[threadgroup_position_in_grid]],
|
|
ushort tiisg[[thread_index_in_simdgroup]],
|
|
ushort sgitg[[simdgroup_index_in_threadgroup]]) {
|
|
switch (args.nr0) {
|
|
case 1: kernel_mul_mv_t_t_4_rows_impl<T0, T04, T1, T14, 1, constant ds4_metal_args_mul_mv &, true>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg); break;
|
|
case 2: kernel_mul_mv_t_t_4_rows_impl<T0, T04, T1, T14, 2, constant ds4_metal_args_mul_mv &, true>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg); break;
|
|
case 4: kernel_mul_mv_t_t_4_rows_impl<T0, T04, T1, T14, 4, constant ds4_metal_args_mul_mv &, true>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg); break;
|
|
};
|
|
}
|
|
|
|
typedef decltype(kernel_mul_mv_t_t_4_rows_bf16<half, half4, half, half4>) mul_mv_t_t_4_rows_bf16;
|
|
|
|
template [[host_name("kernel_mul_mv_f16_f32_4_rows_bf16")]] kernel mul_mv_t_t_4_rows_bf16 kernel_mul_mv_t_t_4_rows_bf16<half, half4, float, float4>;
|
|
|
|
// DS4 compressor projections always compute two same-shaped F16 matvecs from
|
|
// the same normalized activation: one for projected KV and one for pooling
|
|
// scores. This paired variant keeps the exact dense F16 row-reduction shape
|
|
// for each matrix, but shares one dispatch and one activation stream.
|
|
template<short NR0, typename args_t>
|
|
void kernel_mul_mv_f16_f32_pair_4_impl(
|
|
args_t args,
|
|
device const char * src0_a,
|
|
device const char * src0_b,
|
|
device const char * src1,
|
|
device char * dst_a,
|
|
device char * dst_b,
|
|
threadgroup char * shmem,
|
|
uint3 tgpig,
|
|
ushort tiisg,
|
|
ushort sgitg) {
|
|
const short NSG = FC_mul_mv_nsg;
|
|
|
|
constexpr short NW = N_SIMDWIDTH;
|
|
constexpr short NB = 32;
|
|
constexpr short NF = 16;
|
|
constexpr short NF4 = NF/4;
|
|
|
|
const int nb = args.ne00/NB;
|
|
|
|
const int r0 = tgpig.x*NR0;
|
|
const int r1 = tgpig.y;
|
|
const int im = tgpig.z;
|
|
|
|
const uint i12 = im%args.ne12;
|
|
const uint i13 = im/args.ne12;
|
|
|
|
const uint64_t offset1 = r1*args.nb11 + (i12)*args.nb12 + (i13)*args.nb13;
|
|
|
|
device const float * y = (device const float *) (src1 + offset1);
|
|
device const float4 * y4 = (device const float4 *) (src1 + offset1);
|
|
|
|
device const half * ax_a [NR0];
|
|
device const half4 * ax4_a[NR0];
|
|
device const half * ax_b [NR0];
|
|
device const half4 * ax4_b[NR0];
|
|
FOR_UNROLL (short row = 0; row < NR0; ++row) {
|
|
const uint64_t offset0 = (r0 + row)*args.nb01 + (i12/args.r2)*args.nb02 + (i13/args.r3)*args.nb03;
|
|
|
|
ax_a [row] = (device const half *) ((device char *) src0_a + offset0);
|
|
ax4_a[row] = (device const half4 *) ((device char *) src0_a + offset0);
|
|
ax_b [row] = (device const half *) ((device char *) src0_b + offset0);
|
|
ax4_b[row] = (device const half4 *) ((device char *) src0_b + offset0);
|
|
}
|
|
|
|
float sum_a[NR0] = { 0.f };
|
|
float sum_b[NR0] = { 0.f };
|
|
|
|
const short ix = tiisg/(NW/NF);
|
|
const short il = tiisg%(NW/NF);
|
|
|
|
const int ib0 = sgitg*NF + ix;
|
|
|
|
float4 yl4[NF4];
|
|
|
|
device const float4 * yb4 = y4 + (ib0*NB + il*NF)/4;
|
|
|
|
for (int ib = ib0; ib < nb; ib += NSG*NF) {
|
|
for (short i = 0; i < NF4; ++i) {
|
|
yl4[i] = yb4[i];
|
|
}
|
|
|
|
for (short row = 0; row < NR0; row++) {
|
|
device const half4 * xb4_a = ax4_a[row] + (ib*NB + il*NF)/4;
|
|
device const half4 * xb4_b = ax4_b[row] + (ib*NB + il*NF)/4;
|
|
|
|
float suma = 0.f;
|
|
float sumb = 0.f;
|
|
FOR_UNROLL (short i = 0; i < NF4; ++i) {
|
|
const float4 yv = float4(yl4[i]);
|
|
suma += dot(float4(xb4_a[i]), yv);
|
|
sumb += dot(float4(xb4_b[i]), yv);
|
|
}
|
|
|
|
sum_a[row] += suma;
|
|
sum_b[row] += sumb;
|
|
}
|
|
|
|
yb4 += NSG*NF*NW/4;
|
|
}
|
|
|
|
for (int i = nb*NB + sgitg*NW + tiisg; i < args.ne00; i += NW*NSG) {
|
|
for (short row = 0; row < NR0; row++) {
|
|
const float yi = y[i];
|
|
sum_a[row] += ax_a[row][i] * yi;
|
|
sum_b[row] += ax_b[row][i] * yi;
|
|
}
|
|
}
|
|
|
|
device float * dst_a_f32 = (device float *) dst_a + (uint64_t)im*args.ne0*args.ne1 + (uint64_t)r1*args.ne0;
|
|
device float * dst_b_f32 = (device float *) dst_b + (uint64_t)im*args.ne0*args.ne1 + (uint64_t)r1*args.ne0;
|
|
|
|
helper_mv_reduce_and_write<NR0>(dst_a_f32, sum_a, r0, args.ne01, tiisg, sgitg, shmem);
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
helper_mv_reduce_and_write<NR0>(dst_b_f32, sum_b, r0, args.ne01, tiisg, sgitg, shmem);
|
|
}
|
|
|
|
template<typename args_t>
|
|
void kernel_mul_mv_f16_f32_pair_4_disp(
|
|
args_t args,
|
|
device const char * src0_a,
|
|
device const char * src0_b,
|
|
device const char * src1,
|
|
device char * dst_a,
|
|
device char * dst_b,
|
|
threadgroup char * shmem,
|
|
uint3 tgpig,
|
|
ushort tiisg,
|
|
ushort sgitg) {
|
|
switch (args.nr0) {
|
|
case 2: kernel_mul_mv_f16_f32_pair_4_impl<2>(args, src0_a, src0_b, src1, dst_a, dst_b, shmem, tgpig, tiisg, sgitg); break;
|
|
case 4: kernel_mul_mv_f16_f32_pair_4_impl<4>(args, src0_a, src0_b, src1, dst_a, dst_b, shmem, tgpig, tiisg, sgitg); break;
|
|
}
|
|
}
|
|
|
|
kernel void kernel_mul_mv_f16_f32_pair_4(
|
|
constant ds4_metal_args_mul_mv & args,
|
|
device const char * src0_a,
|
|
device const char * src0_b,
|
|
device const char * src1,
|
|
device char * dst_a,
|
|
device char * dst_b,
|
|
threadgroup char * shmem [[threadgroup(0)]],
|
|
uint3 tgpig[[threadgroup_position_in_grid]],
|
|
ushort tiisg[[thread_index_in_simdgroup]],
|
|
ushort sgitg[[simdgroup_index_in_threadgroup]]) {
|
|
kernel_mul_mv_f16_f32_pair_4_disp<constant ds4_metal_args_mul_mv &>(
|
|
args, src0_a, src0_b, src1, dst_a, dst_b, shmem, tgpig, tiisg, sgitg);
|
|
}
|
|
|
|
// Decode compressor projection plus recurrent-state append. The paired
|
|
// matvec remains unchanged and still materializes both F32 outputs. After a
|
|
// device-memory barrier, the first NR0 threads reload those exact stored bits
|
|
// and perform the same state write and score+APE addition as
|
|
// kernel_dsv4_compressor_store_one.
|
|
kernel void kernel_mul_mv_f16_f32_pair_compressor_store_4(
|
|
constant ds4_metal_args_mul_mv & args,
|
|
constant ds4_metal_args_compressor_pair_store & store,
|
|
device const char * src0_a,
|
|
device const char * src0_b,
|
|
device const char * src1,
|
|
device char * dst_a,
|
|
device char * dst_b,
|
|
device const char * ape,
|
|
device float * state_kv,
|
|
device float * state_score,
|
|
threadgroup char * shmem [[threadgroup(0)]],
|
|
uint3 tgpig [[threadgroup_position_in_grid]],
|
|
ushort tiitg [[thread_index_in_threadgroup]],
|
|
ushort tiisg [[thread_index_in_simdgroup]],
|
|
ushort sgitg [[simdgroup_index_in_threadgroup]]) {
|
|
kernel_mul_mv_f16_f32_pair_4_disp<constant ds4_metal_args_mul_mv &>(
|
|
args, src0_a, src0_b, src1, dst_a, dst_b,
|
|
shmem, tgpig, tiisg, sgitg);
|
|
|
|
threadgroup_barrier(mem_flags::mem_device);
|
|
|
|
if (tiitg >= args.nr0 || store.width == 0u || store.ratio == 0u) {
|
|
return;
|
|
}
|
|
const uint col = tgpig.x * (uint)args.nr0 + tiitg;
|
|
if (col >= store.width) return;
|
|
|
|
const uint pos_mod = store.pos % store.ratio;
|
|
const uint dst_row = store.ratio == 4u ? store.ratio + pos_mod : pos_mod;
|
|
const uint dst = dst_row * store.width + col;
|
|
const uint ape_i = pos_mod * store.width + col;
|
|
|
|
device volatile const float * projected_kv =
|
|
(device volatile const float *)dst_a;
|
|
device volatile const float * projected_score =
|
|
(device volatile const float *)dst_b;
|
|
float ape_v;
|
|
if (store.ape_type == 1u) {
|
|
ape_v = (float)(((device const half *)ape)[ape_i]);
|
|
} else {
|
|
ape_v = ((device const float *)ape)[ape_i];
|
|
}
|
|
|
|
state_kv[dst] = projected_kv[col];
|
|
state_score[dst] = projected_score[col] + ape_v;
|
|
}
|
|
|
|
// Decode compressor + indexer-compressor projection in one dispatch. Both
|
|
// pairs read the same normalized activation with the same F16 matvec shape,
|
|
// so one launch covers all four matrices: threadgroups below the first
|
|
// range boundary run the exact paired matvec + state store of
|
|
// kernel_mul_mv_f16_f32_pair_compressor_store_4 for the attention
|
|
// compressor, the rest for the indexer compressor. Per-row reduction trees
|
|
// and the per-threadgroup state stores are unchanged, keeping the fused
|
|
// result bit-identical to the two separate dispatches while removing one
|
|
// dispatch per decode layer.
|
|
kernel void kernel_mul_mv_f16_f32_quad_compressor_store_4(
|
|
constant ds4_metal_args_mul_mv & args,
|
|
constant ds4_metal_args_compressor_pair_store & store0,
|
|
constant ds4_metal_args_compressor_pair_store & store1,
|
|
device const char * src0_a0,
|
|
device const char * src0_b0,
|
|
device const char * src0_a1,
|
|
device const char * src0_b1,
|
|
device const char * src1,
|
|
device char * dst_a0,
|
|
device char * dst_b0,
|
|
device char * dst_a1,
|
|
device char * dst_b1,
|
|
device const char * ape0,
|
|
device const char * ape1,
|
|
device float * state0_kv,
|
|
device float * state0_score,
|
|
device float * state1_kv,
|
|
device float * state1_score,
|
|
threadgroup char * shmem [[threadgroup(0)]],
|
|
uint3 tgpig [[threadgroup_position_in_grid]],
|
|
ushort tiitg [[thread_index_in_threadgroup]],
|
|
ushort tiisg [[thread_index_in_simdgroup]],
|
|
ushort sgitg [[simdgroup_index_in_threadgroup]]) {
|
|
constexpr short NR0 = 2;
|
|
const uint tgs0 = ((uint)store0.width + NR0 - 1u) / NR0;
|
|
const bool second = tgpig.x >= tgs0;
|
|
|
|
uint3 local_tgpig = tgpig;
|
|
if (second) local_tgpig.x = tgpig.x - tgs0;
|
|
|
|
ds4_metal_args_mul_mv largs = args;
|
|
largs.nr0 = NR0;
|
|
largs.ne01 = second ? (int32_t)store1.width : (int32_t)store0.width;
|
|
|
|
if (!second) {
|
|
kernel_mul_mv_f16_f32_pair_4_impl<NR0>(
|
|
largs, src0_a0, src0_b0, src1, dst_a0, dst_b0,
|
|
shmem, local_tgpig, tiisg, sgitg);
|
|
} else {
|
|
kernel_mul_mv_f16_f32_pair_4_impl<NR0>(
|
|
largs, src0_a1, src0_b1, src1, dst_a1, dst_b1,
|
|
shmem, local_tgpig, tiisg, sgitg);
|
|
}
|
|
|
|
threadgroup_barrier(mem_flags::mem_device);
|
|
|
|
// State append: identical to the paired store kernel, scoped to the
|
|
// range this threadgroup just projected (its own outputs only).
|
|
constant ds4_metal_args_compressor_pair_store & store = second ? store1 : store0;
|
|
if (tiitg >= NR0 || store.width == 0u || store.ratio == 0u) {
|
|
return;
|
|
}
|
|
const uint col = local_tgpig.x * (uint)NR0 + tiitg;
|
|
if (col >= store.width) return;
|
|
|
|
const uint pos_mod = store.pos % store.ratio;
|
|
const uint dst_row = store.ratio == 4u ? store.ratio + pos_mod : pos_mod;
|
|
const uint dst = dst_row * store.width + col;
|
|
const uint ape_i = pos_mod * store.width + col;
|
|
|
|
device volatile const float * projected_kv = second
|
|
? (device volatile const float *)dst_a1
|
|
: (device volatile const float *)dst_a0;
|
|
device volatile const float * projected_score = second
|
|
? (device volatile const float *)dst_b1
|
|
: (device volatile const float *)dst_b0;
|
|
device const char * ape = second ? ape1 : ape0;
|
|
device float * state_kv = second ? state1_kv : state0_kv;
|
|
device float * state_score = second ? state1_score : state0_score;
|
|
|
|
float ape_v;
|
|
if (store.ape_type == 1u) {
|
|
ape_v = (float)(((device const half *)ape)[ape_i]);
|
|
} else {
|
|
ape_v = ((device const float *)ape)[ape_i];
|
|
}
|
|
|
|
state_kv[dst] = projected_kv[col];
|
|
state_score[dst] = projected_score[col] + ape_v;
|
|
}
|
|
|
|
/* Decode-only fusion: one dispatch covers the q_a/kv Q8 pair projection and
|
|
* the four F16 compressor projections (attention + indexer) with their
|
|
* state-store epilogue. Both stages read the same normalized attention
|
|
* input and write disjoint outputs. The q_a/kv range hosts two virtual
|
|
* NSG=4 cohorts per threadgroup, each an exact replica of
|
|
* kernel_mul_mv_q8_0_f32_pair (same per-lane K walk and reduction tree, cf.
|
|
* kernel_dsv4_router_shared_gate_up_q8_0); the compressor ranges run
|
|
* kernel_mul_mv_f16_f32_pair_4_impl<2> and the paired store epilogue
|
|
* verbatim, so every output bit matches the two separate dispatches. */
|
|
kernel void kernel_dsv4_qkv_pair_quad_compressor_store_q8_0(
|
|
constant ds4_metal_args_mul_mv & args0,
|
|
constant ds4_metal_args_mul_mv & args1,
|
|
constant ds4_metal_args_mul_mv & cargs,
|
|
constant ds4_metal_args_compressor_pair_store & store0,
|
|
constant ds4_metal_args_compressor_pair_store & store1,
|
|
constant uint & pair_vtgs,
|
|
device const char * qw0,
|
|
device const char * qw1,
|
|
device const char * cw0a,
|
|
device const char * cw0b,
|
|
device const char * cw1a,
|
|
device const char * cw1b,
|
|
device const char * src1,
|
|
device char * dst0,
|
|
device char * dst1,
|
|
device char * cdst_a0,
|
|
device char * cdst_b0,
|
|
device char * cdst_a1,
|
|
device char * cdst_b1,
|
|
device const char * ape0,
|
|
device const char * ape1,
|
|
device float * state0_kv,
|
|
device float * state0_score,
|
|
device float * state1_kv,
|
|
device float * state1_score,
|
|
threadgroup char * shmem [[threadgroup(0)]],
|
|
uint3 tgpig [[threadgroup_position_in_grid]],
|
|
ushort tiitg [[thread_index_in_threadgroup]],
|
|
ushort tiisg [[thread_index_in_simdgroup]],
|
|
ushort sgitg [[simdgroup_index_in_threadgroup]]) {
|
|
constexpr short NW = N_SIMDWIDTH;
|
|
const uint pair_ctgs = (pair_vtgs + 1u) / 2u;
|
|
|
|
if (tgpig.x < pair_ctgs) {
|
|
/* Q8 pair range: cohort c of threadgroup t runs virtual pair
|
|
* threadgroup 2t+c with the original NSG=4 mapping. */
|
|
constexpr short NSG = 4;
|
|
constexpr short NQ = 8;
|
|
constexpr short NR0 = 2;
|
|
const uint cohort = sgitg >> 2;
|
|
const ushort vsg = sgitg & 3u;
|
|
const uint vt = tgpig.x * 2u + cohort;
|
|
const bool valid = vt < pair_vtgs;
|
|
|
|
const int r0 = vt * NR0;
|
|
const bool active_a = valid && r0 < args0.ne01;
|
|
const bool active_b = valid && r0 < args1.ne01;
|
|
const int nb = args0.ne00 / QK8_0;
|
|
|
|
device const float *y = (device const float *)src1;
|
|
device const block_q8_0 *ax_a[NR0];
|
|
device const block_q8_0 *ax_b[NR0];
|
|
FOR_UNROLL (short row = 0; row < NR0; ++row) {
|
|
const int out_row = r0 + row;
|
|
ax_a[row] = active_a && out_row < args0.ne01
|
|
? (device const block_q8_0 *)(qw0 + (uint64_t)out_row * args0.nb01)
|
|
: (device const block_q8_0 *)qw0;
|
|
ax_b[row] = active_b && out_row < args1.ne01
|
|
? (device const block_q8_0 *)(qw1 + (uint64_t)out_row * args1.nb01)
|
|
: (device const block_q8_0 *)qw1;
|
|
}
|
|
|
|
float suma[NR0] = { 0.f };
|
|
float sumb[NR0] = { 0.f };
|
|
const short ix = tiisg / (NW / NQ);
|
|
const short il = tiisg % (NW / NQ);
|
|
const int ib0 = vsg * NQ + ix;
|
|
float yl[NQ];
|
|
device const float *yb = y + ib0 * QK8_0 + il * NQ;
|
|
|
|
if (valid) {
|
|
for (int ib = ib0; ib < nb; ib += NSG * NQ) {
|
|
FOR_UNROLL (short i = 0; i < NQ; ++i) {
|
|
yl[i] = yb[i];
|
|
}
|
|
FOR_UNROLL (short row = 0; row < NR0; ++row) {
|
|
const int out_row = r0 + row;
|
|
if (active_a && out_row < args0.ne01) {
|
|
device const int8_t *qs = ax_a[row][ib].qs + il * NQ;
|
|
float sumq = 0.f;
|
|
FOR_UNROLL (short i = 0; i < NQ; ++i) {
|
|
sumq += qs[i] * yl[i];
|
|
}
|
|
suma[row] += sumq * ax_a[row][ib].d;
|
|
}
|
|
if (active_b && out_row < args1.ne01) {
|
|
device const int8_t *qs = ax_b[row][ib].qs + il * NQ;
|
|
float sumq = 0.f;
|
|
FOR_UNROLL (short i = 0; i < NQ; ++i) {
|
|
sumq += qs[i] * yl[i];
|
|
}
|
|
sumb[row] += sumq * ax_b[row][ib].d;
|
|
}
|
|
}
|
|
yb += NSG * NQ * QK8_0;
|
|
}
|
|
}
|
|
|
|
threadgroup float *shared =
|
|
(threadgroup float *)shmem + cohort * (2 * NR0 * NW);
|
|
threadgroup float *sha[NR0];
|
|
threadgroup float *shb[NR0];
|
|
FOR_UNROLL (short row = 0; row < NR0; ++row) {
|
|
sha[row] = shared + NW * row;
|
|
shb[row] = shared + NW * (NR0 + row);
|
|
if (vsg == 0) {
|
|
sha[row][tiisg] = 0.0f;
|
|
if (active_b) shb[row][tiisg] = 0.0f;
|
|
}
|
|
suma[row] = simd_sum(suma[row]);
|
|
if (active_b) sumb[row] = simd_sum(sumb[row]);
|
|
}
|
|
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
FOR_UNROLL (short row = 0; row < NR0; ++row) {
|
|
if (tiisg == 0) {
|
|
sha[row][vsg] = suma[row];
|
|
if (active_b) shb[row][vsg] = sumb[row];
|
|
}
|
|
}
|
|
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
device float *out_a = (device float *)dst0;
|
|
device float *out_b = (device float *)dst1;
|
|
FOR_UNROLL (short row = 0; row < NR0; ++row) {
|
|
const float total_a = simd_sum(sha[row][tiisg]);
|
|
if (tiisg == 0 && vsg == 0) {
|
|
const int out_row = r0 + row;
|
|
if (active_a && out_row < args0.ne01) out_a[out_row] = total_a;
|
|
}
|
|
if (active_b) {
|
|
const float total_b = simd_sum(shb[row][tiisg]);
|
|
if (tiisg == 0 && vsg == 0) {
|
|
const int out_row = r0 + row;
|
|
if (out_row < args1.ne01) out_b[out_row] = total_b;
|
|
}
|
|
}
|
|
}
|
|
return;
|
|
}
|
|
|
|
/* Compressor quad range: verbatim body of
|
|
* kernel_mul_mv_f16_f32_quad_compressor_store_4 on the shifted grid. */
|
|
constexpr short NR0 = 2;
|
|
const uint lx = tgpig.x - pair_ctgs;
|
|
const uint tgs0 = ((uint)store0.width + NR0 - 1u) / NR0;
|
|
const bool second = lx >= tgs0;
|
|
|
|
uint3 local_tgpig = tgpig;
|
|
local_tgpig.x = second ? lx - tgs0 : lx;
|
|
|
|
ds4_metal_args_mul_mv largs = cargs;
|
|
largs.nr0 = NR0;
|
|
largs.ne01 = second ? (int32_t)store1.width : (int32_t)store0.width;
|
|
|
|
if (!second) {
|
|
kernel_mul_mv_f16_f32_pair_4_impl<NR0>(
|
|
largs, cw0a, cw0b, src1, cdst_a0, cdst_b0,
|
|
shmem, local_tgpig, tiisg, sgitg);
|
|
} else {
|
|
kernel_mul_mv_f16_f32_pair_4_impl<NR0>(
|
|
largs, cw1a, cw1b, src1, cdst_a1, cdst_b1,
|
|
shmem, local_tgpig, tiisg, sgitg);
|
|
}
|
|
|
|
threadgroup_barrier(mem_flags::mem_device);
|
|
|
|
// State append: identical to the paired store kernel, scoped to the
|
|
// range this threadgroup just projected (its own outputs only).
|
|
constant ds4_metal_args_compressor_pair_store & store = second ? store1 : store0;
|
|
if (tiitg >= NR0 || store.width == 0u || store.ratio == 0u) {
|
|
return;
|
|
}
|
|
const uint col = local_tgpig.x * (uint)NR0 + tiitg;
|
|
if (col >= store.width) return;
|
|
|
|
const uint pos_mod = store.pos % store.ratio;
|
|
const uint dst_row = store.ratio == 4u ? store.ratio + pos_mod : pos_mod;
|
|
const uint dst = dst_row * store.width + col;
|
|
const uint ape_i = pos_mod * store.width + col;
|
|
|
|
device volatile const float * projected_kv = second
|
|
? (device volatile const float *)cdst_a1
|
|
: (device volatile const float *)cdst_a0;
|
|
device volatile const float * projected_score = second
|
|
? (device volatile const float *)cdst_b1
|
|
: (device volatile const float *)cdst_b0;
|
|
device const char * ape = second ? ape1 : ape0;
|
|
device float * state_kv = second ? state1_kv : state0_kv;
|
|
device float * state_score = second ? state1_score : state0_score;
|
|
|
|
float ape_v;
|
|
if (store.ape_type == 1u) {
|
|
ape_v = (float)(((device const half *)ape)[ape_i]);
|
|
} else {
|
|
ape_v = ((device const float *)ape)[ape_i];
|
|
}
|
|
|
|
state_kv[dst] = projected_kv[col];
|
|
state_score[dst] = projected_score[col] + ape_v;
|
|
}
|
|
|
|
template<typename T0, typename T1, typename args_t>
|
|
void kernel_mul_mv_t_t_short_impl(
|
|
args_t args,
|
|
device const char * src0,
|
|
device const char * src1,
|
|
device char * dst,
|
|
uint3 tgpig,
|
|
ushort tiisg) {
|
|
const int r0 = tgpig.x*32 + tiisg;
|
|
const int r1 = tgpig.y;
|
|
const int im = tgpig.z;
|
|
|
|
if (r0 >= args.ne01) {
|
|
return;
|
|
}
|
|
|
|
const uint i12 = im%args.ne12;
|
|
const uint i13 = im/args.ne12;
|
|
|
|
const uint64_t offset0 = r0*args.nb01 + (i12/args.r2)*args.nb02 + (i13/args.r3)*args.nb03;
|
|
|
|
device const T0 * x = (device const T0 *) (src0 + offset0);
|
|
|
|
device float * dst_f32 = (device float *) dst + (uint64_t)im*args.ne0*args.ne1;
|
|
|
|
const uint64_t offset1 = r1*args.nb11 + (i12)*args.nb12 + (i13)*args.nb13;
|
|
|
|
device const T1 * y = (device const T1 *) (src1 + offset1);
|
|
|
|
float res = 0.0f;
|
|
|
|
for (int i = 0; i < args.ne00; ++i) {
|
|
res += (float) x[i] * (float) y[i];
|
|
}
|
|
|
|
dst_f32[(uint64_t)r1*args.ne0 + r0] = res;
|
|
}
|
|
|
|
// Scalar fallback for short rows. It trades parallelism for lower dispatch and
|
|
// reduction overhead when DS4 asks for tiny dense matvecs.
|
|
template<typename T0, typename T1>
|
|
kernel void kernel_mul_mv_t_t_short(
|
|
constant ds4_metal_args_mul_mv & args,
|
|
device const char * src0,
|
|
device const char * src1,
|
|
device char * dst,
|
|
uint3 tgpig[[threadgroup_position_in_grid]],
|
|
ushort tiisg[[thread_index_in_simdgroup]]) {
|
|
kernel_mul_mv_t_t_short_impl<T0, T1, constant ds4_metal_args_mul_mv &>(
|
|
args,
|
|
src0,
|
|
src1,
|
|
dst,
|
|
tgpig,
|
|
tiisg);
|
|
}
|
|
|
|
typedef decltype(kernel_mul_mv_t_t_short<half, half>) mul_mv_t_t_short_t;
|
|
|
|
// Host-visible short-row dense matvec variants.
|
|
template [[host_name("kernel_mul_mv_f32_f32_short")]] kernel mul_mv_t_t_short_t kernel_mul_mv_t_t_short<float, float>;
|
|
template [[host_name("kernel_mul_mv_f16_f32_short")]] kernel mul_mv_t_t_short_t kernel_mul_mv_t_t_short<half, float>;
|
|
|
|
template <typename type4x4>
|
|
void dequantize_f32(device const float4x4 * src, short il, thread type4x4 & reg) {
|
|
reg = (type4x4)(*src);
|
|
}
|
|
|
|
template <typename type4x4>
|
|
void dequantize_f16(device const half4x4 * src, short il, thread type4x4 & reg) {
|
|
reg = (type4x4)(*src);
|
|
}
|
|
|
|
template <typename type4x4>
|
|
void dequantize_q8_0(device const block_q8_0 *xb, short il, thread type4x4 & reg) {
|
|
device const int8_t * qs = ((device const int8_t *)xb->qs);
|
|
const float d = xb->d;
|
|
|
|
float4x4 reg_f;
|
|
|
|
for (int i = 0; i < 16; i++) {
|
|
reg_f[i/4][i%4] = (qs[i + 16*il] * d);
|
|
}
|
|
|
|
reg = (type4x4) reg_f;
|
|
}
|
|
|
|
struct ds4_dense_block_q4_0 {
|
|
half d;
|
|
uchar qs[16];
|
|
};
|
|
|
|
struct ds4_dense_block_q4_K {
|
|
half d;
|
|
half dmin;
|
|
uchar scales[12];
|
|
uchar qs[128];
|
|
};
|
|
|
|
static inline uchar2 ds4_dense_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))};
|
|
}
|
|
|
|
template <typename type4x4>
|
|
void dequantize_dense_q4_0(device const ds4_dense_block_q4_0 *xb, short il, thread type4x4 ®) {
|
|
float4x4 reg_f;
|
|
const float d = (float)xb->d;
|
|
const int base = 16 * (int)il;
|
|
for (int i = 0; i < 16; i++) {
|
|
const int k = base + i;
|
|
/* ggml Q4_0: elems 0..15 = low nibbles of qs[0..15], 16..31 = high. */
|
|
const uchar packed = xb->qs[(uint)(k & 15)];
|
|
const uchar q = (k < 16) ? (packed & 0x0f) : (packed >> 4);
|
|
reg_f[i / 4][i % 4] = d * ((float)q - 8.0f);
|
|
}
|
|
reg = (type4x4)reg_f;
|
|
}
|
|
|
|
template <typename type4x4>
|
|
void dequantize_dense_q4_K(device const ds4_dense_block_q4_K *xb, short il, thread type4x4 ®) {
|
|
device const uchar *q = xb->qs;
|
|
short is = (il / 4) * 2;
|
|
q = q + (il / 4) * 32 + 16 * (il & 1);
|
|
il = il & 3;
|
|
const uchar2 sc = ds4_dense_q4_K_scale_min(is, il / 2, xb->scales);
|
|
const float d = il < 2 ? (float)xb->d : (float)xb->d * (1.0f / 16.0f);
|
|
const float min = (float)xb->dmin;
|
|
const float dl = d * sc[0];
|
|
const float ml = min * sc[1];
|
|
const ushort mask = il < 2 ? 0x0F : 0xF0;
|
|
|
|
float4x4 reg_f;
|
|
for (int i = 0; i < 16; ++i) {
|
|
reg_f[i / 4][i % 4] = dl * (q[i] & mask) - ml;
|
|
}
|
|
reg = (type4x4)reg_f;
|
|
}
|
|
|
|
/*
|
|
* Bit-identical twin of dequantize_q8_0 for the MPP staging loop: same
|
|
* half(float(qs[i]) * d) per element, but the 16 consecutive int8 lanes are
|
|
* fetched as eight aligned 16-bit loads instead of sixteen byte loads.
|
|
* xb->qs is always 2-byte aligned (block_q8_0.d is a half).
|
|
*/
|
|
void dequantize_q8_0_pairs(device const block_q8_0 *xb, short il, thread half4x4 & reg) {
|
|
device const ushort *qs16 = (device const ushort *)(xb->qs + 16*il);
|
|
const float d = xb->d;
|
|
|
|
float4x4 reg_f;
|
|
|
|
FOR_UNROLL (short i = 0; i < 8; i++) {
|
|
const ushort u = qs16[i];
|
|
reg_f[i/2][(i%2)*2 + 0] = ((float)(int8_t)(u & 0xFF)) * d;
|
|
reg_f[i/2][(i%2)*2 + 1] = ((float)(int8_t)(u >> 8)) * d;
|
|
}
|
|
|
|
reg = (half4x4) reg_f;
|
|
}
|
|
|
|
/* Materialize the same half-rounded Q8 weights used by tiled matmul once
|
|
* per projection. Eight threads cover each block, four adjacent values each. */
|
|
kernel void kernel_q8_prefill_unpack(
|
|
constant uint &n_blocks,
|
|
device const block_q8_0 *src,
|
|
device half4 *dst,
|
|
uint gid [[thread_position_in_grid]]) {
|
|
const uint block = gid / 8u;
|
|
if (block >= n_blocks) return;
|
|
const uint i = (gid % 8u) * 4u;
|
|
const float d = (float)src[block].d;
|
|
float4 values;
|
|
for (uint j = 0; j < 4; j++) values[j] = (float)src[block].qs[i+j] * d;
|
|
dst[gid] = half4(values);
|
|
}
|
|
|
|
template <typename type4>
|
|
void dequantize_q8_0_t4(device const block_q8_0 *xb, short il, thread type4 & reg) {
|
|
device const int8_t * qs = ((device const int8_t *)xb->qs);
|
|
const float d = xb->d;
|
|
|
|
for (int i = 0; i < 4; i++) {
|
|
reg[i] = (qs[4*(il%4) + i + 16*(il/4)] * d);
|
|
}
|
|
}
|
|
|
|
template <typename type4>
|
|
void dequantize_dense_q4_0_t4(device const ds4_dense_block_q4_0 *xb, short il, thread type4 ®) {
|
|
const float d = (float)xb->d;
|
|
const int base = 4 * (int)il;
|
|
for (int i = 0; i < 4; i++) {
|
|
const int k = base + i;
|
|
/* ggml Q4_0: elems 0..15 = low nibbles of qs[0..15], 16..31 = high. */
|
|
const uchar packed = xb->qs[(uint)(k & 15)];
|
|
const uchar q = (k < 16) ? (packed & 0x0f) : (packed >> 4);
|
|
reg[i] = d * ((float)q - 8.0f);
|
|
}
|
|
}
|
|
|
|
template <typename type4>
|
|
void dequantize_dense_q4_K_t4(device const ds4_dense_block_q4_K *xb, short il, thread type4 ®) {
|
|
float4x4 tmp;
|
|
dequantize_dense_q4_K(xb, il / 4, tmp);
|
|
const short row = il & 3;
|
|
for (int i = 0; i < 4; i++) {
|
|
reg[i] = tmp[row][i];
|
|
}
|
|
}
|
|
|
|
// DS4 small-batch mat-vec kernel used for 2..8 prompt tokens.
|
|
template<short r1ptg, typename q_t, short chpb, void (*deq_t4)(device const q_t *, short, thread float4 &) >
|
|
void kernel_mul_mv_ext_q4_f32_impl(
|
|
constant ds4_metal_args_mul_mv_ext & args,
|
|
device const char * src0,
|
|
device const char * src1,
|
|
device char * dst,
|
|
uint3 tgpig[[threadgroup_position_in_grid]],
|
|
ushort tiisg[[thread_index_in_simdgroup]],
|
|
ushort sgitg[[simdgroup_index_in_threadgroup]]) {
|
|
const short NSG = FC_mul_mv_nsg;
|
|
const short nxpsg = FC_mul_mv_nxpsg;
|
|
|
|
const short chpt = 4; // chunks per thread
|
|
|
|
const short nypsg = (32/nxpsg);
|
|
|
|
const short tx = tiisg%nxpsg;
|
|
const short ty = tiisg/nxpsg;
|
|
|
|
const int i01 = tgpig.x*(nypsg*NSG) + nypsg*sgitg + ty;
|
|
const int i11 = tgpig.y*r1ptg;
|
|
const int i1m = tgpig.z;
|
|
|
|
const int i12 = i1m%args.ne12;
|
|
const int i13 = i1m/args.ne12;
|
|
|
|
const uint64_t offset0 = i01*args.nb01 + (i12/args.r2)*args.nb02 + (i13/args.r3)*args.nb03;
|
|
const uint64_t offset1 = i11*args.nb11 + (i12 )*args.nb12 + (i13 )*args.nb13;
|
|
|
|
device const q_t * xq = (i01 < args.ne01) ? (device const q_t *) (src0 + offset0) + tx/chpb : (device const q_t *) src0;
|
|
|
|
device const float4 * y4[r1ptg];
|
|
|
|
for (int ir1 = 0; ir1 < r1ptg; ++ir1) {
|
|
y4[ir1] = (i11 + ir1 < args.ne11) ? (device const float4 *) (src1 + offset1 + ir1*args.nb11) + tx : (device const float4 *) src1;
|
|
}
|
|
|
|
float sumf[r1ptg] = { [ 0 ... r1ptg - 1 ] = 0.0f };
|
|
|
|
short cch = tx%chpb; // current chunk index
|
|
|
|
for (int ich = tx; 4*ich < args.ne00; ich += chpt*nxpsg) {
|
|
float4 lx[chpt];
|
|
|
|
#pragma unroll(chpt)
|
|
for (short ch = 0; ch < chpt; ++ch) {
|
|
deq_t4(xq, cch, lx[ch]);
|
|
|
|
cch += nxpsg;
|
|
if (cch >= chpb) {
|
|
xq += cch/chpb;
|
|
cch %= chpb;
|
|
}
|
|
}
|
|
|
|
#pragma unroll(chpt)
|
|
for (short ch = 0; ch < chpt; ++ch) {
|
|
#pragma unroll(r1ptg)
|
|
for (short ir1 = 0; ir1 < r1ptg; ++ir1) {
|
|
sumf[ir1] += dot(lx[ch], y4[ir1][ch*nxpsg]);
|
|
}
|
|
}
|
|
|
|
#pragma unroll(r1ptg)
|
|
for (short ir1 = 0; ir1 < r1ptg; ++ir1) {
|
|
y4[ir1] += chpt*nxpsg;
|
|
}
|
|
}
|
|
|
|
// reduce only the threads in each row
|
|
for (short ir1 = 0; ir1 < r1ptg; ++ir1) {
|
|
if (nxpsg >= 32) {
|
|
sumf[ir1] += simd_shuffle_down(sumf[ir1], 16);
|
|
}
|
|
if (nxpsg >= 16) {
|
|
sumf[ir1] += simd_shuffle_down(sumf[ir1], 8);
|
|
}
|
|
if (nxpsg >= 8) {
|
|
sumf[ir1] += simd_shuffle_down(sumf[ir1], 4);
|
|
}
|
|
if (nxpsg >= 4) {
|
|
sumf[ir1] += simd_shuffle_down(sumf[ir1], 2);
|
|
}
|
|
if (nxpsg >= 2) {
|
|
sumf[ir1] += simd_shuffle_down(sumf[ir1], 1);
|
|
}
|
|
}
|
|
|
|
if (tx == 0) {
|
|
for (short ir1 = 0; ir1 < r1ptg && i11 + ir1 < args.ne11; ++ir1) {
|
|
device float * dst_f32 = (device float *) dst + (uint64_t)i1m*args.ne0*args.ne1 + (uint64_t)(i11 + ir1)*args.ne0;
|
|
|
|
if (i01 < args.ne01) {
|
|
dst_f32[i01] = sumf[ir1];
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// Small-batch prompt matvec for 2..5 tokens. It bridges decode-style matvec and
|
|
// full matmul when DS4 prefill chunks are too small to amortize matrix tiles.
|
|
template<short r1ptg, typename q_t, short epb, void (*deq_t4)(device const q_t *, short, thread float4 &)>
|
|
kernel void kernel_mul_mv_ext_q4_f32_disp(
|
|
constant ds4_metal_args_mul_mv_ext & args,
|
|
device const char * src0,
|
|
device const char * src1,
|
|
device char * dst,
|
|
uint3 tgpig[[threadgroup_position_in_grid]],
|
|
ushort tiisg[[thread_index_in_simdgroup]],
|
|
ushort sgitg[[simdgroup_index_in_threadgroup]]) {
|
|
kernel_mul_mv_ext_q4_f32_impl<r1ptg, q_t, epb/4, deq_t4>(args, src0, src1, dst, tgpig, tiisg, sgitg);
|
|
}
|
|
|
|
typedef decltype(kernel_mul_mv_ext_q4_f32_disp<2, block_q8_0, 32, dequantize_q8_0_t4>) mul_mv_ext_q4_f32_t;
|
|
|
|
// Host-visible small-batch variants for r1=2..5 during tiny prompt/support
|
|
// paths.
|
|
template [[host_name("kernel_mul_mv_ext_f32_f32_r1_2")]] kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<2, float4, 4, dequantize_f32_t4>;
|
|
template [[host_name("kernel_mul_mv_ext_f32_f32_r1_3")]] kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<3, float4, 4, dequantize_f32_t4>;
|
|
template [[host_name("kernel_mul_mv_ext_f32_f32_r1_4")]] kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<4, float4, 4, dequantize_f32_t4>;
|
|
template [[host_name("kernel_mul_mv_ext_f32_f32_r1_5")]] kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<5, float4, 4, dequantize_f32_t4>;
|
|
|
|
template<short r1ptg>
|
|
void kernel_mul_mv_ext_q8_0_pair_swiglu_f32_impl(
|
|
constant ds4_metal_args_mul_mv_ext & args,
|
|
device const char * src0_gate,
|
|
device const char * src0_up,
|
|
device const char * src1,
|
|
device char * dst_gate,
|
|
device char * dst_up,
|
|
device char * dst_mid,
|
|
constant float &clamp_value,
|
|
uint3 tgpig[[threadgroup_position_in_grid]],
|
|
ushort tiisg[[thread_index_in_simdgroup]],
|
|
ushort sgitg[[simdgroup_index_in_threadgroup]]) {
|
|
const short NSG = FC_mul_mv_nsg;
|
|
const short nxpsg = FC_mul_mv_nxpsg;
|
|
|
|
const short chpt = 4;
|
|
const short chpb = 8;
|
|
const short nypsg = (32/nxpsg);
|
|
|
|
const short tx = tiisg%nxpsg;
|
|
const short ty = tiisg/nxpsg;
|
|
|
|
const int i01 = tgpig.x*(nypsg*NSG) + nypsg*sgitg + ty;
|
|
const int i11 = tgpig.y*r1ptg;
|
|
const int i1m = tgpig.z;
|
|
|
|
const int i12 = i1m%args.ne12;
|
|
const int i13 = i1m/args.ne12;
|
|
|
|
const uint64_t offset0 = i01*args.nb01 + (i12/args.r2)*args.nb02 + (i13/args.r3)*args.nb03;
|
|
const uint64_t offset1 = i11*args.nb11 + (i12 )*args.nb12 + (i13 )*args.nb13;
|
|
|
|
device const block_q8_0 * xq_gate =
|
|
(i01 < args.ne01) ? (device const block_q8_0 *)(src0_gate + offset0) + tx/chpb
|
|
: (device const block_q8_0 *)src0_gate;
|
|
device const block_q8_0 * xq_up =
|
|
(i01 < args.ne01) ? (device const block_q8_0 *)(src0_up + offset0) + tx/chpb
|
|
: (device const block_q8_0 *)src0_up;
|
|
|
|
device const float4 * y4[r1ptg];
|
|
for (int ir1 = 0; ir1 < r1ptg; ++ir1) {
|
|
y4[ir1] = (i11 + ir1 < args.ne11)
|
|
? (device const float4 *)(src1 + offset1 + ir1*args.nb11) + tx
|
|
: (device const float4 *)src1;
|
|
}
|
|
|
|
float sum_gate[r1ptg] = { [ 0 ... r1ptg - 1 ] = 0.0f };
|
|
float sum_up[r1ptg] = { [ 0 ... r1ptg - 1 ] = 0.0f };
|
|
|
|
short cch = tx%chpb;
|
|
for (int ich = tx; 4*ich < args.ne00; ich += chpt*nxpsg) {
|
|
float4 lg[chpt];
|
|
float4 lu[chpt];
|
|
|
|
#pragma unroll(chpt)
|
|
for (short ch = 0; ch < chpt; ++ch) {
|
|
dequantize_q8_0_t4(xq_gate, cch, lg[ch]);
|
|
dequantize_q8_0_t4(xq_up, cch, lu[ch]);
|
|
|
|
cch += nxpsg;
|
|
if (cch >= chpb) {
|
|
xq_gate += cch/chpb;
|
|
xq_up += cch/chpb;
|
|
cch %= chpb;
|
|
}
|
|
}
|
|
|
|
#pragma unroll(chpt)
|
|
for (short ch = 0; ch < chpt; ++ch) {
|
|
#pragma unroll(r1ptg)
|
|
for (short ir1 = 0; ir1 < r1ptg; ++ir1) {
|
|
const float4 y = y4[ir1][ch*nxpsg];
|
|
sum_gate[ir1] += dot(lg[ch], y);
|
|
sum_up[ir1] += dot(lu[ch], y);
|
|
}
|
|
}
|
|
|
|
#pragma unroll(r1ptg)
|
|
for (short ir1 = 0; ir1 < r1ptg; ++ir1) {
|
|
y4[ir1] += chpt*nxpsg;
|
|
}
|
|
}
|
|
|
|
for (short ir1 = 0; ir1 < r1ptg; ++ir1) {
|
|
if (nxpsg >= 32) {
|
|
sum_gate[ir1] += simd_shuffle_down(sum_gate[ir1], 16);
|
|
sum_up[ir1] += simd_shuffle_down(sum_up[ir1], 16);
|
|
}
|
|
if (nxpsg >= 16) {
|
|
sum_gate[ir1] += simd_shuffle_down(sum_gate[ir1], 8);
|
|
sum_up[ir1] += simd_shuffle_down(sum_up[ir1], 8);
|
|
}
|
|
if (nxpsg >= 8) {
|
|
sum_gate[ir1] += simd_shuffle_down(sum_gate[ir1], 4);
|
|
sum_up[ir1] += simd_shuffle_down(sum_up[ir1], 4);
|
|
}
|
|
if (nxpsg >= 4) {
|
|
sum_gate[ir1] += simd_shuffle_down(sum_gate[ir1], 2);
|
|
sum_up[ir1] += simd_shuffle_down(sum_up[ir1], 2);
|
|
}
|
|
if (nxpsg >= 2) {
|
|
sum_gate[ir1] += simd_shuffle_down(sum_gate[ir1], 1);
|
|
sum_up[ir1] += simd_shuffle_down(sum_up[ir1], 1);
|
|
}
|
|
}
|
|
|
|
if (tx == 0 && i01 < args.ne01) {
|
|
for (short ir1 = 0; ir1 < r1ptg && i11 + ir1 < args.ne11; ++ir1) {
|
|
const uint64_t dst_base =
|
|
(uint64_t)i1m*args.ne0*args.ne1 + (uint64_t)(i11 + ir1)*args.ne0;
|
|
device float * gate_f32 = (device float *)dst_gate + dst_base;
|
|
device float * up_f32 = (device float *)dst_up + dst_base;
|
|
device float * mid_f32 = (device float *)dst_mid + dst_base;
|
|
|
|
const float gate = sum_gate[ir1];
|
|
const float up = sum_up[ir1];
|
|
gate_f32[i01] = gate;
|
|
up_f32[i01] = up;
|
|
|
|
float g = gate;
|
|
float u = up;
|
|
if (clamp_value > 1.0e-6f) {
|
|
g = min(g, clamp_value);
|
|
u = clamp(u, -clamp_value, clamp_value);
|
|
}
|
|
const float silu = g / (1.0f + exp(-g));
|
|
mid_f32[i01] = silu * u;
|
|
}
|
|
}
|
|
}
|
|
|
|
template<short r1ptg>
|
|
kernel void kernel_mul_mv_ext_q8_0_pair_swiglu_f32_disp(
|
|
constant ds4_metal_args_mul_mv_ext & args,
|
|
device const char * src0_gate,
|
|
device const char * src0_up,
|
|
device const char * src1,
|
|
device char * dst_gate,
|
|
device char * dst_up,
|
|
device char * dst_mid,
|
|
constant float &clamp_value,
|
|
uint3 tgpig[[threadgroup_position_in_grid]],
|
|
ushort tiisg[[thread_index_in_simdgroup]],
|
|
ushort sgitg[[simdgroup_index_in_threadgroup]]) {
|
|
kernel_mul_mv_ext_q8_0_pair_swiglu_f32_impl<r1ptg>(
|
|
args, src0_gate, src0_up, src1, dst_gate, dst_up, dst_mid, clamp_value,
|
|
tgpig, tiisg, sgitg);
|
|
}
|
|
|
|
typedef decltype(kernel_mul_mv_ext_q8_0_pair_swiglu_f32_disp<2>) mul_mv_ext_q8_0_pair_swiglu_f32_t;
|
|
|
|
// Host-visible small-batch variants. DS4 currently needs F16 and Q8_0 weights
|
|
// for r1=2..5 during the prompt path.
|
|
template [[host_name("kernel_mul_mv_ext_f16_f32_r1_2")]] kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<2, half4, 4, dequantize_f16_t4>;
|
|
template [[host_name("kernel_mul_mv_ext_f16_f32_r1_3")]] kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<3, half4, 4, dequantize_f16_t4>;
|
|
template [[host_name("kernel_mul_mv_ext_f16_f32_r1_4")]] kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<4, half4, 4, dequantize_f16_t4>;
|
|
template [[host_name("kernel_mul_mv_ext_f16_f32_r1_5")]] kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<5, half4, 4, dequantize_f16_t4>;
|
|
|
|
template [[host_name("kernel_mul_mv_ext_q8_0_f32_r1_2")]] kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<2, block_q8_0, 32, dequantize_q8_0_t4>;
|
|
template [[host_name("kernel_mul_mv_ext_q8_0_f32_r1_3")]] kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<3, block_q8_0, 32, dequantize_q8_0_t4>;
|
|
template [[host_name("kernel_mul_mv_ext_q8_0_f32_r1_4")]] kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<4, block_q8_0, 32, dequantize_q8_0_t4>;
|
|
template [[host_name("kernel_mul_mv_ext_q8_0_f32_r1_5")]] kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<5, block_q8_0, 32, dequantize_q8_0_t4>;
|
|
|
|
template [[host_name("kernel_mul_mv_ext_q4_0_f32_r1_1")]] kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<1, ds4_dense_block_q4_0, 32, dequantize_dense_q4_0_t4>;
|
|
template [[host_name("kernel_mul_mv_ext_q4_0_f32_r1_2")]] kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<2, ds4_dense_block_q4_0, 32, dequantize_dense_q4_0_t4>;
|
|
template [[host_name("kernel_mul_mv_ext_q4_0_f32_r1_3")]] kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<3, ds4_dense_block_q4_0, 32, dequantize_dense_q4_0_t4>;
|
|
template [[host_name("kernel_mul_mv_ext_q4_0_f32_r1_4")]] kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<4, ds4_dense_block_q4_0, 32, dequantize_dense_q4_0_t4>;
|
|
template [[host_name("kernel_mul_mv_ext_q4_0_f32_r1_5")]] kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<5, ds4_dense_block_q4_0, 32, dequantize_dense_q4_0_t4>;
|
|
|
|
template [[host_name("kernel_mul_mv_ext_q4_K_f32_r1_1")]] kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<1, ds4_dense_block_q4_K, 256, dequantize_dense_q4_K_t4>;
|
|
template [[host_name("kernel_mul_mv_ext_q4_K_f32_r1_2")]] kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<2, ds4_dense_block_q4_K, 256, dequantize_dense_q4_K_t4>;
|
|
template [[host_name("kernel_mul_mv_ext_q4_K_f32_r1_3")]] kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<3, ds4_dense_block_q4_K, 256, dequantize_dense_q4_K_t4>;
|
|
template [[host_name("kernel_mul_mv_ext_q4_K_f32_r1_4")]] kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<4, ds4_dense_block_q4_K, 256, dequantize_dense_q4_K_t4>;
|
|
template [[host_name("kernel_mul_mv_ext_q4_K_f32_r1_5")]] kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<5, ds4_dense_block_q4_K, 256, dequantize_dense_q4_K_t4>;
|
|
|
|
template [[host_name("kernel_mul_mv_ext_q8_0_pair_swiglu_f32_r1_2")]] kernel mul_mv_ext_q8_0_pair_swiglu_f32_t kernel_mul_mv_ext_q8_0_pair_swiglu_f32_disp<2>;
|
|
template [[host_name("kernel_mul_mv_ext_q8_0_pair_swiglu_f32_r1_3")]] kernel mul_mv_ext_q8_0_pair_swiglu_f32_t kernel_mul_mv_ext_q8_0_pair_swiglu_f32_disp<3>;
|
|
template [[host_name("kernel_mul_mv_ext_q8_0_pair_swiglu_f32_r1_4")]] kernel mul_mv_ext_q8_0_pair_swiglu_f32_t kernel_mul_mv_ext_q8_0_pair_swiglu_f32_disp<4>;
|
|
template [[host_name("kernel_mul_mv_ext_q8_0_pair_swiglu_f32_r1_5")]] kernel mul_mv_ext_q8_0_pair_swiglu_f32_t kernel_mul_mv_ext_q8_0_pair_swiglu_f32_disp<5>;
|
|
|
|
constant bool FC_mul_mm_bc_inp [[function_constant(FC_MUL_MM + 0)]];
|
|
constant bool FC_mul_mm_bc_out [[function_constant(FC_MUL_MM + 1)]];
|
|
|
|
#ifdef DS4_METAL_HAS_TENSOR
|
|
// Retained Metal4/TensorOps dense prefill kernel. The legacy MPP prototype
|
|
// staged both operands in threadgroup memory; this version stages only the
|
|
// model weight tile and lets MPP read the dense RHS activation matrix directly
|
|
// from device memory. That direct-RHS shape was the clear win for DS4's large
|
|
// aligned F16/Q8_0 prompt matmuls. The host selects the widest token tile that
|
|
// evenly divides the batch, with 128-token tiles retained after the 64-token
|
|
// retest was neutral or slower.
|
|
//
|
|
// The host dispatch guarantees M % NR0 == 0 and K % NK == 0, so the dequant
|
|
// stage does no bounds work. The weight tile is double-buffered: the next
|
|
// k-step's dequant overlaps the current cooperative matmul, so the k-loop
|
|
// needs one threadgroup barrier per step instead of two.
|
|
template<
|
|
short NR1,
|
|
typename SA, typename SA_4x4, typename block_q, short nl,
|
|
void (*dequantize_func)(device const block_q *, short, thread SA_4x4 &),
|
|
typename T0, typename T0_4x4, typename T1>
|
|
kernel void kernel_mul_mm_mpp_direct_rhs(
|
|
constant ds4_metal_args_mul_mm & args,
|
|
device const char * srcA,
|
|
device const char * srcB,
|
|
device char * dst,
|
|
threadgroup char * shmem [[threadgroup(0)]],
|
|
uint3 tgpig [[threadgroup_position_in_grid]],
|
|
ushort tiitg [[thread_index_in_threadgroup]],
|
|
ushort sgitg [[simdgroup_index_in_threadgroup]]) {
|
|
(void) sgitg;
|
|
|
|
constexpr int NR0 = 64;
|
|
constexpr int NK = 32;
|
|
constexpr int NL = NK/16;
|
|
constexpr int NUM_THREADS = 128;
|
|
|
|
const int K = args.ne00;
|
|
const int M = args.ne0;
|
|
const int N = args.ne1;
|
|
const int im = tgpig.z;
|
|
const int i12 = im%args.ne12;
|
|
const int i13 = im/args.ne12;
|
|
const int r0 = tgpig.y*NR0;
|
|
const int r1 = tgpig.x*NR1;
|
|
|
|
const uint64_t offset0 = (i12/args.r2)*args.nb02 + (i13/args.r3)*args.nb03;
|
|
|
|
threadgroup SA *sa = (threadgroup SA *)shmem;
|
|
auto tA0 = tensor(sa, dextents<int32_t, 2>(NK, NR0));
|
|
auto tA1 = tensor(sa + NR0*NK, dextents<int32_t, 2>(NK, NR0));
|
|
|
|
device T1 *ptrB = (device T1 *)(srcB + args.nb12*i12 + args.nb13*i13);
|
|
const int strideB = args.nb11/sizeof(T1);
|
|
auto tB = tensor(ptrB, dextents<int32_t, 2>(K, N), array<int, 2>({1, strideB}));
|
|
|
|
matmul2d<
|
|
matmul2d_descriptor(NR1, NR0, NK, false, true, true,
|
|
matmul2d_descriptor::mode::multiply_accumulate),
|
|
execution_simdgroups<4>> mm;
|
|
|
|
auto cT = mm.template get_destination_cooperative_tensor<decltype(tB), decltype(tA0), float>();
|
|
|
|
#pragma unroll
|
|
for (uint16_t i = 0; i < cT.get_capacity(); ++i) {
|
|
if (cT.is_valid_element(i)) {
|
|
cT[i] = 0.0f;
|
|
}
|
|
}
|
|
|
|
// NR0*NL/NUM_THREADS 16-value weight chunks per thread (1 at NR0=64).
|
|
auto stage_tile = [&](const int loop_k, threadgroup SA *buf) {
|
|
FOR_UNROLL (int work = tiitg; work < NR0*NL; work += NUM_THREADS) {
|
|
const int row = work / NL;
|
|
const int k_chunk = work % NL;
|
|
const int k_pos = loop_k + k_chunk*16;
|
|
const short k_base = k_chunk*16;
|
|
if (is_same<T0_4x4, block_q>::value && FC_mul_mm_bc_inp) {
|
|
device const T0 *row_ptr_f =
|
|
(device const T0 *)(srcA + args.nb01*(r0 + row) + offset0);
|
|
FOR_UNROLL (short i = 0; i < 16; i++) {
|
|
buf[row*NK + k_base + i] = (SA)row_ptr_f[k_pos + i];
|
|
}
|
|
} else {
|
|
device const block_q *row_ptr =
|
|
(device const block_q *)(srcA + args.nb01*(r0 + row) + offset0);
|
|
SA_4x4 temp_a;
|
|
dequantize_func(row_ptr + k_pos/(16*nl), (k_pos/16)%nl, temp_a);
|
|
typedef vec<SA, 4> SA4;
|
|
threadgroup SA4 *dst4 = (threadgroup SA4 *)(buf + row*NK + k_base);
|
|
dst4[0] = temp_a[0];
|
|
dst4[1] = temp_a[1];
|
|
dst4[2] = temp_a[2];
|
|
dst4[3] = temp_a[3];
|
|
}
|
|
}
|
|
};
|
|
|
|
stage_tile(0, sa);
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
uint buf_sel = 0;
|
|
for (int loop_k = 0; loop_k < K; loop_k += NK) {
|
|
auto mA = (buf_sel ? tA1 : tA0).slice(0, 0);
|
|
auto mB = tB.slice(loop_k, r1);
|
|
mm.run(mB, mA, cT);
|
|
|
|
const int next_k = loop_k + NK;
|
|
if (next_k < K) {
|
|
buf_sel ^= 1u;
|
|
stage_tile(next_k, buf_sel ? sa + NR0*NK : sa);
|
|
}
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
}
|
|
|
|
device float *dst_batch = (device float *)dst + im*N*M;
|
|
auto tD = tensor(dst_batch, dextents<int32_t, 2>(M, N), array<int, 2>({1, M}));
|
|
auto mD = tD.slice(r0, r1);
|
|
cT.store(mD);
|
|
}
|
|
|
|
typedef decltype(kernel_mul_mm_mpp_direct_rhs<32, half, half4x4, float4x4, 1, dequantize_f32, float, float4x4, float>) mul_mm_mpp_direct_rhs_t;
|
|
|
|
template [[host_name("kernel_mul_mm_f16_f32_mpp_direct_rhs")]] kernel mul_mm_mpp_direct_rhs_t kernel_mul_mm_mpp_direct_rhs<32, half, half4x4, half4x4, 1, dequantize_f16, half, half4x4, float>;
|
|
template [[host_name("kernel_mul_mm_f16_f32_mpp_direct_rhs_n64")]] kernel mul_mm_mpp_direct_rhs_t kernel_mul_mm_mpp_direct_rhs<64, half, half4x4, half4x4, 1, dequantize_f16, half, half4x4, float>;
|
|
template [[host_name("kernel_mul_mm_f16_f32_mpp_direct_rhs_n128")]] kernel mul_mm_mpp_direct_rhs_t kernel_mul_mm_mpp_direct_rhs<128, half, half4x4, half4x4, 1, dequantize_f16, half, half4x4, float>;
|
|
template [[host_name("kernel_mul_mm_q4_0_f32_nax_direct_rhs")]] kernel mul_mm_mpp_direct_rhs_t kernel_mul_mm_mpp_direct_rhs<32, half, half4x4, ds4_dense_block_q4_0, 2, dequantize_dense_q4_0, float, float4x4, float>;
|
|
template [[host_name("kernel_mul_mm_q4_0_f32_nax_direct_rhs_n64")]] kernel mul_mm_mpp_direct_rhs_t kernel_mul_mm_mpp_direct_rhs<64, half, half4x4, ds4_dense_block_q4_0, 2, dequantize_dense_q4_0, float, float4x4, float>;
|
|
template [[host_name("kernel_mul_mm_q4_0_f32_nax_direct_rhs_n128")]] kernel mul_mm_mpp_direct_rhs_t kernel_mul_mm_mpp_direct_rhs<128, half, half4x4, ds4_dense_block_q4_0, 2, dequantize_dense_q4_0, float, float4x4, float>;
|
|
template [[host_name("kernel_mul_mm_q4_K_f32_nax_direct_rhs")]] kernel mul_mm_mpp_direct_rhs_t kernel_mul_mm_mpp_direct_rhs<32, half, half4x4, ds4_dense_block_q4_K, 16, dequantize_dense_q4_K, float, float4x4, float>;
|
|
template [[host_name("kernel_mul_mm_q4_K_f32_nax_direct_rhs_n64")]] kernel mul_mm_mpp_direct_rhs_t kernel_mul_mm_mpp_direct_rhs<64, half, half4x4, ds4_dense_block_q4_K, 16, dequantize_dense_q4_K, float, float4x4, float>;
|
|
template [[host_name("kernel_mul_mm_q4_K_f32_nax_direct_rhs_n128")]] kernel mul_mm_mpp_direct_rhs_t kernel_mul_mm_mpp_direct_rhs<128, half, half4x4, ds4_dense_block_q4_K, 16, dequantize_dense_q4_K, float, float4x4, float>;
|
|
|
|
template [[host_name("kernel_mul_mm_q8_0_f32_nax_direct_rhs")]] kernel mul_mm_mpp_direct_rhs_t kernel_mul_mm_mpp_direct_rhs<32, half, half4x4, block_q8_0, 2, dequantize_q8_0_pairs, float, float4x4, float>;
|
|
template [[host_name("kernel_mul_mm_q8_0_f32_nax_direct_rhs_n64")]] kernel mul_mm_mpp_direct_rhs_t kernel_mul_mm_mpp_direct_rhs<64, half, half4x4, block_q8_0, 2, dequantize_q8_0_pairs, float, float4x4, float>;
|
|
template [[host_name("kernel_mul_mm_q8_0_f32_nax_direct_rhs_n128")]] kernel mul_mm_mpp_direct_rhs_t kernel_mul_mm_mpp_direct_rhs<128, half, half4x4, block_q8_0, 2, dequantize_q8_0_pairs, float, float4x4, float>;
|
|
#endif
|
|
|
|
// Tiled matrix-matrix kernel used for prompt batches larger than 8. DS4 uses
|
|
// this to turn prefill into large simdgroup matrix operations; each block_q
|
|
// contains 16*nl weights.
|
|
template<typename S0, typename S0_4x4, typename S0_8x8, typename S1, typename S1_2x4, typename S1_8x8, typename block_q, short nl, void (*dequantize_func)(device const block_q *, short, thread S0_4x4 &), typename T0, typename T0_4x4, typename T1, typename T1_2x4>
|
|
kernel void kernel_mul_mm(
|
|
constant ds4_metal_args_mul_mm & args,
|
|
device const char * src0,
|
|
device const char * src1,
|
|
device char * dst,
|
|
threadgroup char * shmem [[threadgroup(0)]],
|
|
uint3 tgpig[[threadgroup_position_in_grid]],
|
|
ushort tiitg[[thread_index_in_threadgroup]],
|
|
ushort sgitg[[simdgroup_index_in_threadgroup]]) {
|
|
|
|
threadgroup S0 * sa = (threadgroup S0 *)(shmem);
|
|
threadgroup S1 * sb = (threadgroup S1 *)(shmem + 64u*32u*sizeof(S0));
|
|
|
|
constexpr int NR0 = 64;
|
|
constexpr int NR1 = 32;
|
|
|
|
constexpr int NK = 32;
|
|
constexpr int NL0 = NK/16;
|
|
constexpr int NL1 = NK/8;
|
|
|
|
const int im = tgpig.z;
|
|
const int r0 = tgpig.y*NR0;
|
|
const int r1 = tgpig.x*NR1;
|
|
|
|
// if this block is of 64x32 shape or smaller
|
|
const short nr0 = (args.ne0 - r0 < NR0) ? (args.ne0 - r0) : NR0;
|
|
const short nr1 = (args.ne1 - r1 < NR1) ? (args.ne1 - r1) : NR1;
|
|
|
|
// a thread shouldn't load data outside of the matrix
|
|
const short lr0 = ((short)tiitg/NL0) < nr0 ? ((short)tiitg/NL0) : nr0 - 1; // 0 .. 63
|
|
const short lr1 = ((short)tiitg/NL1) < nr1 ? ((short)tiitg/NL1) : nr1 - 1; // 0 .. 31
|
|
|
|
const short il0 = (tiitg % NL0);
|
|
|
|
short il = il0;
|
|
|
|
const int i12 = im%args.ne12;
|
|
const int i13 = im/args.ne12;
|
|
|
|
const uint64_t offset0 = (i12/args.r2)*args.nb02 + (i13/args.r3)*args.nb03;
|
|
const short offset1 = il0/nl;
|
|
|
|
device const block_q * x = (device const block_q *)(src0 + args.nb01*(r0 + lr0) + offset0) + offset1;
|
|
|
|
const short iy = 8*(tiitg % NL1);
|
|
|
|
device const T1 * y = (device const T1 *)(src1
|
|
+ args.nb13*i13
|
|
+ args.nb12*i12
|
|
+ args.nb11*(r1 + lr1)
|
|
+ args.nb10*iy);
|
|
|
|
S0_8x8 ma[4];
|
|
S1_8x8 mb[2];
|
|
|
|
simdgroup_float8x8 mc[8];
|
|
|
|
for (short i = 0; i < 8; i++){
|
|
mc[i] = make_filled_simdgroup_matrix<float, 8>(0.f);
|
|
}
|
|
|
|
for (int loop_k = 0; loop_k < args.ne00; loop_k += NK) {
|
|
// load data and store to threadgroup memory
|
|
if (is_same<T0_4x4, block_q>::value && FC_mul_mm_bc_inp) {
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
// no need for dequantization
|
|
for (short i = 0; i < 16; i++) {
|
|
const short sx = 2*il0 + i/8;
|
|
const short sy = (tiitg/NL0)/8;
|
|
|
|
const short lx = (tiitg/NL0)%8;
|
|
const short ly = i%8;
|
|
|
|
const short ib = 8*sx + sy;
|
|
|
|
*(sa + 64*ib + 8*ly + lx) = loop_k + 16*il + i < args.ne00 ? *((device T0 *) x + i) : 0;
|
|
}
|
|
} else {
|
|
S0_4x4 temp_a;
|
|
dequantize_func(x, il, temp_a);
|
|
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
FOR_UNROLL (short i = 0; i < 16; i++) {
|
|
const short sx = 2*il0 + i/8;
|
|
const short sy = (tiitg/NL0)/8;
|
|
|
|
const short lx = (tiitg/NL0)%8;
|
|
const short ly = i%8;
|
|
|
|
const short ib = 8*sx + sy;
|
|
|
|
// Pointer-form store avoids a slower address-lowering path in
|
|
// current Apple Metal compilers for this dequantized tile write.
|
|
*(sa + 64*ib + 8*ly + lx) = temp_a[i/4][i%4];
|
|
}
|
|
}
|
|
|
|
if (FC_mul_mm_bc_inp) {
|
|
for (short i = 0; i < 8; ++i) {
|
|
const short sx = (tiitg%NL1);
|
|
const short sy = (tiitg/NL1)/8;
|
|
|
|
const short lx = i;
|
|
const short ly = (tiitg/NL1)%8;
|
|
|
|
const short ib = 4*sx + sy;
|
|
|
|
*(sb + 64*ib + 8*ly + lx) = loop_k + iy + i < args.ne00 ? (S1) *((device T1 *) y + i) : 0;
|
|
}
|
|
} else {
|
|
const short sx = (tiitg%NL1);
|
|
const short sy = (tiitg/NL1)/8;
|
|
|
|
const short ly = (tiitg/NL1)%8;
|
|
|
|
const short ib = 4*sx + sy;
|
|
|
|
*(threadgroup S1_2x4 *)(sb + 64*ib + 8*ly) = (S1_2x4)(*((device T1_2x4 *) y));
|
|
}
|
|
|
|
il = (il + 2 < nl) ? il + 2 : il % 2;
|
|
x = (il < 2) ? x + (2 + nl - 1)/nl : x;
|
|
|
|
y += NK;
|
|
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
// load matrices from threadgroup memory and conduct outer products
|
|
threadgroup const S0 * lsma = (sa + 4*64*(sgitg%2));
|
|
threadgroup const S1 * lsmb = (sb + 2*64*(sgitg/2));
|
|
|
|
FOR_UNROLL (short ik = 0; ik < NK/8; ik++) {
|
|
simdgroup_barrier(mem_flags::mem_none);
|
|
|
|
FOR_UNROLL (short i = 0; i < 4; i++) {
|
|
simdgroup_load(ma[i], lsma + 64*i, 8, 0, false);
|
|
}
|
|
|
|
simdgroup_barrier(mem_flags::mem_none);
|
|
|
|
FOR_UNROLL (short i = 0; i < 2; i++) {
|
|
simdgroup_load(mb[i], lsmb + 64*i, 8, 0, false);
|
|
}
|
|
|
|
simdgroup_barrier(mem_flags::mem_none);
|
|
|
|
FOR_UNROLL (short i = 0; i < 8; i++){
|
|
simdgroup_multiply_accumulate(mc[i], mb[i/4], ma[i%4], mc[i]);
|
|
}
|
|
|
|
lsma += 8*64;
|
|
lsmb += 4*64;
|
|
}
|
|
}
|
|
|
|
if (!FC_mul_mm_bc_out || (r0 + NR0 <= args.ne0 && r1 + NR1 <= args.ne1)) {
|
|
// if no bounds checks on the output are needed, we can directly write to device memory
|
|
device float * C = (device float *) dst +
|
|
(r0 + 32*(sgitg & 1)) + \
|
|
(r1 + 16*(sgitg >> 1)) * args.ne0 + im*args.ne1*args.ne0;
|
|
|
|
for (short i = 0; i < 8; i++) {
|
|
simdgroup_store(mc[i], C + 8*(i%4) + 8*args.ne0*(i/4), args.ne0, 0, false);
|
|
}
|
|
} else {
|
|
// block is smaller than 64x32, we should avoid writing data outside of the matrix
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
threadgroup float * temp_str = ((threadgroup float *) shmem) + 32*(sgitg&1) + (16*(sgitg >> 1))*NR0;
|
|
|
|
for (short i = 0; i < 8; i++) {
|
|
simdgroup_store(mc[i], temp_str + 8*(i%4) + 8*NR0*(i/4), NR0, 0, false);
|
|
}
|
|
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
if (sgitg == 0) {
|
|
for (int j = tiitg; j < nr1; j += NR1) {
|
|
device float * D = (device float *) dst + r0 + (r1 + j)*args.ne0 + im*args.ne1*args.ne0;
|
|
device float4 * D4 = (device float4 *) D;
|
|
|
|
threadgroup float * C = temp_str + (j*NR0);
|
|
threadgroup float4 * C4 = (threadgroup float4 *) C;
|
|
|
|
int i = 0;
|
|
for (; i < nr0/4; i++) {
|
|
*(D4 + i) = *(C4 + i);
|
|
}
|
|
|
|
i *= 4;
|
|
for (; i < nr0; i++) {
|
|
*(D + i) = *(C + i);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// Legacy F16-weight/F32-RHS prefill matmul with a per-row RMS scale applied at
|
|
// the existing F32-to-F16 RHS staging boundary. The tile layout, half inputs,
|
|
// float accumulators, and output path intentionally mirror kernel_mul_mm so the
|
|
// only arithmetic change versus materializing RMSNorm first is where x*scale is
|
|
// rounded from F32 to F16.
|
|
kernel void kernel_mul_mm_f16_f32_scaled(
|
|
constant ds4_metal_args_mul_mm & args,
|
|
device const char * src0,
|
|
device const char * src1,
|
|
device char * dst,
|
|
device const float * scales,
|
|
threadgroup char * shmem [[threadgroup(0)]],
|
|
uint3 tgpig[[threadgroup_position_in_grid]],
|
|
ushort tiitg[[thread_index_in_threadgroup]],
|
|
ushort sgitg[[simdgroup_index_in_threadgroup]]) {
|
|
|
|
threadgroup half * sa = (threadgroup half *)(shmem);
|
|
threadgroup half * sb = (threadgroup half *)(shmem + 4096);
|
|
|
|
constexpr int NR0 = 64;
|
|
constexpr int NR1 = 32;
|
|
|
|
constexpr int NK = 32;
|
|
constexpr int NL0 = NK/16;
|
|
constexpr int NL1 = NK/8;
|
|
|
|
const int im = tgpig.z;
|
|
const int r0 = tgpig.y*NR0;
|
|
const int r1 = tgpig.x*NR1;
|
|
|
|
// if this block is of 64x32 shape or smaller
|
|
const short nr0 = (args.ne0 - r0 < NR0) ? (args.ne0 - r0) : NR0;
|
|
const short nr1 = (args.ne1 - r1 < NR1) ? (args.ne1 - r1) : NR1;
|
|
|
|
// a thread shouldn't load data outside of the matrix
|
|
const short lr0 = ((short)tiitg/NL0) < nr0 ? ((short)tiitg/NL0) : nr0 - 1; // 0 .. 63
|
|
const short lr1 = ((short)tiitg/NL1) < nr1 ? ((short)tiitg/NL1) : nr1 - 1; // 0 .. 31
|
|
|
|
const short il0 = (tiitg % NL0);
|
|
|
|
short il = il0;
|
|
|
|
const int i12 = im%args.ne12;
|
|
const int i13 = im/args.ne12;
|
|
|
|
const uint64_t offset0 = (i12/args.r2)*args.nb02 + (i13/args.r3)*args.nb03;
|
|
const short offset1 = il0;
|
|
|
|
device const half4x4 * x = (device const half4x4 *)(src0 + args.nb01*(r0 + lr0) + offset0) + offset1;
|
|
|
|
const short iy = 8*(tiitg % NL1);
|
|
|
|
device const float * y = (device const float *)(src1
|
|
+ args.nb13*i13
|
|
+ args.nb12*i12
|
|
+ args.nb11*(r1 + lr1)
|
|
+ args.nb10*iy);
|
|
const float row_scale = scales[r1 + lr1];
|
|
|
|
simdgroup_half8x8 ma[4];
|
|
simdgroup_half8x8 mb[2];
|
|
|
|
simdgroup_float8x8 mc[8];
|
|
|
|
for (short i = 0; i < 8; i++){
|
|
mc[i] = make_filled_simdgroup_matrix<float, 8>(0.f);
|
|
}
|
|
|
|
for (int loop_k = 0; loop_k < args.ne00; loop_k += NK) {
|
|
// load data and store to threadgroup memory
|
|
if (FC_mul_mm_bc_inp) {
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
// no need for dequantization
|
|
for (short i = 0; i < 16; i++) {
|
|
const short sx = 2*il0 + i/8;
|
|
const short sy = (tiitg/NL0)/8;
|
|
|
|
const short lx = (tiitg/NL0)%8;
|
|
const short ly = i%8;
|
|
|
|
const short ib = 8*sx + sy;
|
|
|
|
*(sa + 64*ib + 8*ly + lx) = loop_k + 16*il + i < args.ne00 ? *((device half *) x + i) : 0;
|
|
}
|
|
} else {
|
|
half4x4 temp_a;
|
|
dequantize_f16(x, il, temp_a);
|
|
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
FOR_UNROLL (short i = 0; i < 16; i++) {
|
|
const short sx = 2*il0 + i/8;
|
|
const short sy = (tiitg/NL0)/8;
|
|
|
|
const short lx = (tiitg/NL0)%8;
|
|
const short ly = i%8;
|
|
|
|
const short ib = 8*sx + sy;
|
|
|
|
// Pointer-form store matches the legacy F16 tile layout.
|
|
*(sa + 64*ib + 8*ly + lx) = temp_a[i/4][i%4];
|
|
}
|
|
}
|
|
|
|
if (FC_mul_mm_bc_inp) {
|
|
for (short i = 0; i < 8; ++i) {
|
|
const short sx = (tiitg%NL1);
|
|
const short sy = (tiitg/NL1)/8;
|
|
|
|
const short lx = i;
|
|
const short ly = (tiitg/NL1)%8;
|
|
|
|
const short ib = 4*sx + sy;
|
|
const float scaled = loop_k + iy + i < args.ne00 ? y[i] * row_scale : 0.0f;
|
|
|
|
*(sb + 64*ib + 8*ly + lx) = (half)scaled;
|
|
}
|
|
} else {
|
|
const short sx = (tiitg%NL1);
|
|
const short sy = (tiitg/NL1)/8;
|
|
|
|
const short ly = (tiitg/NL1)%8;
|
|
|
|
const short ib = 4*sx + sy;
|
|
|
|
const float2x4 raw = *((device const float2x4 *) y);
|
|
float2x4 scaled;
|
|
scaled[0] = raw[0] * row_scale;
|
|
scaled[1] = raw[1] * row_scale;
|
|
*(threadgroup half2x4 *)(sb + 64*ib + 8*ly) = (half2x4)scaled;
|
|
}
|
|
|
|
il = (il + 2 < 1) ? il + 2 : il % 2;
|
|
x = (il < 2) ? x + 2 : x;
|
|
|
|
y += NK;
|
|
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
// load matrices from threadgroup memory and conduct outer products
|
|
threadgroup const half * lsma = (sa + 4*64*(sgitg%2));
|
|
threadgroup const half * lsmb = (sb + 2*64*(sgitg/2));
|
|
|
|
FOR_UNROLL (short ik = 0; ik < NK/8; ik++) {
|
|
simdgroup_barrier(mem_flags::mem_none);
|
|
|
|
FOR_UNROLL (short i = 0; i < 4; i++) {
|
|
simdgroup_load(ma[i], lsma + 64*i, 8, 0, false);
|
|
}
|
|
|
|
simdgroup_barrier(mem_flags::mem_none);
|
|
|
|
FOR_UNROLL (short i = 0; i < 2; i++) {
|
|
simdgroup_load(mb[i], lsmb + 64*i, 8, 0, false);
|
|
}
|
|
|
|
simdgroup_barrier(mem_flags::mem_none);
|
|
|
|
FOR_UNROLL (short i = 0; i < 8; i++){
|
|
simdgroup_multiply_accumulate(mc[i], mb[i/4], ma[i%4], mc[i]);
|
|
}
|
|
|
|
lsma += 8*64;
|
|
lsmb += 4*64;
|
|
}
|
|
}
|
|
|
|
if (!FC_mul_mm_bc_out || (r0 + NR0 <= args.ne0 && r1 + NR1 <= args.ne1)) {
|
|
// if no bounds checks on the output are needed, we can directly write to device memory
|
|
device float * C = (device float *) dst +
|
|
(r0 + 32*(sgitg & 1)) + \
|
|
(r1 + 16*(sgitg >> 1)) * args.ne0 + im*args.ne1*args.ne0;
|
|
|
|
for (short i = 0; i < 8; i++) {
|
|
simdgroup_store(mc[i], C + 8*(i%4) + 8*args.ne0*(i/4), args.ne0, 0, false);
|
|
}
|
|
} else {
|
|
// block is smaller than 64x32, we should avoid writing data outside of the matrix
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
threadgroup float * temp_str = ((threadgroup float *) shmem) + 32*(sgitg&1) + (16*(sgitg >> 1))*NR0;
|
|
|
|
for (short i = 0; i < 8; i++) {
|
|
simdgroup_store(mc[i], temp_str + 8*(i%4) + 8*NR0*(i/4), NR0, 0, false);
|
|
}
|
|
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
if (sgitg == 0) {
|
|
for (int j = tiitg; j < nr1; j += NR1) {
|
|
device float * D = (device float *) dst + r0 + (r1 + j)*args.ne0 + im*args.ne1*args.ne0;
|
|
device float4 * D4 = (device float4 *) D;
|
|
|
|
threadgroup float * C = temp_str + (j*NR0);
|
|
threadgroup float4 * C4 = (threadgroup float4 *) C;
|
|
|
|
int i = 0;
|
|
for (; i < nr0/4; i++) {
|
|
*(D4 + i) = *(C4 + i);
|
|
}
|
|
|
|
i *= 4;
|
|
for (; i < nr0; i++) {
|
|
*(D + i) = *(C + i);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
typedef decltype(kernel_mul_mm<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, float4x4, 1, dequantize_f32, float, float4x4, float, float2x4>) mul_mm_t;
|
|
|
|
// Dense prefill variants, including a full-F32 router without half conversion.
|
|
template [[host_name("kernel_mul_mm_f32_f32_full")]] kernel mul_mm_t kernel_mul_mm<float, float4x4, simdgroup_float8x8, float, float2x4, simdgroup_float8x8, float4x4, 1, dequantize_f32, float, float4x4, float, float2x4>;
|
|
template [[host_name("kernel_mul_mm_f16_f32")]] kernel mul_mm_t kernel_mul_mm<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, half4x4, 1, dequantize_f16, half, half4x4, float, float2x4>;
|
|
template [[host_name("kernel_mul_mm_q8_0_f32")]] kernel mul_mm_t kernel_mul_mm<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q8_0, 2, dequantize_q8_0, float, float4x4, float, float2x4>;
|
|
template [[host_name("kernel_mul_mm_q4_0_f32")]] kernel mul_mm_t kernel_mul_mm<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, ds4_dense_block_q4_0, 2, dequantize_dense_q4_0, float, float4x4, float, float2x4>;
|
|
template [[host_name("kernel_mul_mm_q4_K_f32")]] kernel mul_mm_t kernel_mul_mm<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, ds4_dense_block_q4_K, 16, dequantize_dense_q4_K, float, float4x4, float, float2x4>;
|