Support DeepSeek V4 Flash 0731
This commit is contained in:
Georg Bauer
@@ -197,33 +197,6 @@ kernel void kernel_mul_mv_q8_0_f32(
|
||||
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_r4")]]
|
||||
kernel void kernel_mul_mv_q8_0_f32_r4(
|
||||
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<4, constant ds4_metal_args_mul_mv &>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
|
||||
}
|
||||
|
||||
// Output projection alias used by the optimized host dispatch.
|
||||
[[host_name("kernel_mul_mv_q8_0_f32_nr4")]]
|
||||
kernel void kernel_mul_mv_q8_0_f32_nr4(
|
||||
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<4, constant ds4_metal_args_mul_mv &>(
|
||||
args, src0, src1, dst, 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
|
||||
@@ -497,25 +470,190 @@ kernel void kernel_dsv4_shared_gate_up_swiglu_q8_0(
|
||||
clamp_value, shmem, tgpig, tiisg, sgitg);
|
||||
}
|
||||
|
||||
[[host_name("kernel_dsv4_shared_gate_up_swiglu_q8_0_r4")]]
|
||||
kernel void kernel_dsv4_shared_gate_up_swiglu_q8_0_r4(
|
||||
// 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]]) {
|
||||
kernel_dsv4_shared_gate_up_swiglu_q8_0_impl<4, true>(
|
||||
args, src0_gate, src0_up, src1, dst_gate, dst_up, dst_mid,
|
||||
clamp_value, shmem, tgpig, tiisg, sgitg);
|
||||
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,
|
||||
@@ -535,24 +673,6 @@ kernel void kernel_dsv4_shared_mid_swiglu_q8_0(
|
||||
clamp_value, shmem, tgpig, tiisg, sgitg);
|
||||
}
|
||||
|
||||
[[host_name("kernel_dsv4_shared_mid_swiglu_q8_0_r4")]]
|
||||
kernel void kernel_dsv4_shared_mid_swiglu_q8_0_r4(
|
||||
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<4, false>(
|
||||
args, src0_gate, src0_up, src1, dst_gate, dst_up, 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(
|
||||
@@ -974,6 +1094,311 @@ kernel void kernel_mul_mv_f16_f32_pair_compressor_store_4(
|
||||
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,
|
||||
@@ -1476,125 +1901,6 @@ 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
|
||||
template<
|
||||
short NR0, 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(
|
||||
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 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;
|
||||
threadgroup SA *sb = sa + NR0*NK;
|
||||
auto tA = tensor(sa, dextents<int32_t, 2>(NK, NR0));
|
||||
auto tB = tensor(sb, dextents<int32_t, 2>(NK, NR1));
|
||||
|
||||
device const T1 *ptrB = (device const T1 *)(srcB + args.nb12*i12 + args.nb13*i13);
|
||||
const int strideB = args.nb11/sizeof(T1);
|
||||
|
||||
matmul2d<
|
||||
matmul2d_descriptor(NR1, NR0, NK, false, true, false,
|
||||
matmul2d_descriptor::mode::multiply_accumulate),
|
||||
execution_simdgroups<4>> mm;
|
||||
|
||||
auto cT = mm.template get_destination_cooperative_tensor<decltype(tB), decltype(tA), float>();
|
||||
|
||||
#pragma unroll
|
||||
for (uint16_t i = 0; i < cT.get_capacity(); ++i) {
|
||||
if (cT.is_valid_element(i)) {
|
||||
cT[i] = 0.0f;
|
||||
}
|
||||
}
|
||||
|
||||
for (int loop_k = 0; loop_k < K; loop_k += NK) {
|
||||
for (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 (!FC_mul_mm_bc_out || r0 + row < M) {
|
||||
if (is_same<T0_4x4, block_q>::value && FC_mul_mm_bc_inp) {
|
||||
device const T0 *row_ptr = (device const T0 *)(srcA + args.nb01*(r0 + row) + offset0);
|
||||
FOR_UNROLL (short i = 0; i < 16; i++) {
|
||||
sa[row*NK + k_base + i] = (k_pos + i < K) ? (SA)row_ptr[k_pos + i] : (SA)0;
|
||||
}
|
||||
} else {
|
||||
const int block_idx = k_pos/(16*nl);
|
||||
const short il = (k_pos/16)%nl;
|
||||
device const block_q *row_ptr = (device const block_q *)(srcA + args.nb01*(r0 + row) + offset0);
|
||||
|
||||
SA_4x4 temp_a;
|
||||
dequantize_func(row_ptr + block_idx, il, temp_a);
|
||||
FOR_UNROLL (short i = 0; i < 16; i++) {
|
||||
sa[row*NK + k_base + i] = (k_pos + i < K) ? temp_a[i/4][i%4] : (SA)0;
|
||||
}
|
||||
}
|
||||
} else {
|
||||
FOR_UNROLL (short i = 0; i < 16; i++) {
|
||||
sa[row*NK + k_base + i] = (SA)0;
|
||||
}
|
||||
}
|
||||
}
|
||||
for (int work = tiitg; work < NK*NR1; work += NUM_THREADS) {
|
||||
const int col = work/NK;
|
||||
const int k = work%NK;
|
||||
if ((!FC_mul_mm_bc_out && !FC_mul_mm_bc_inp) ||
|
||||
(r1 + col < N && loop_k + k < K)) {
|
||||
sb[col*NK + k] = (SA)ptrB[(uint64_t)(r1 + col)*strideB + loop_k + k];
|
||||
} else {
|
||||
sb[col*NK + k] = (SA)0;
|
||||
}
|
||||
}
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
auto mA = tA.slice(0, 0);
|
||||
auto mB = tB.slice(0, 0);
|
||||
mm.run(mB, mA, cT);
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
}
|
||||
|
||||
device float *dst_batch = (device float *)dst + im*N*M;
|
||||
if (!FC_mul_mm_bc_out) {
|
||||
device float *dst_tile = dst_batch + r0 + (uint64_t)r1*M;
|
||||
auto tD = tensor(dst_tile, dextents<int32_t, 2>(NR0, NR1), array<int, 2>({1, M}));
|
||||
cT.store(tD);
|
||||
} else {
|
||||
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<64, 32, half, half4x4, float4x4, 1, dequantize_f32, float, float4x4, float>) mul_mm_mpp_t;
|
||||
|
||||
template [[host_name("kernel_mul_mm_f16_f32_mpp")]] kernel mul_mm_mpp_t kernel_mul_mm_mpp<64, 32, half, half4x4, half4x4, 1, dequantize_f16, half, half4x4, float>;
|
||||
|
||||
// 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
|
||||
@@ -2144,242 +2450,6 @@ kernel void kernel_mul_mm_f16_f32_scaled(
|
||||
}
|
||||
}
|
||||
|
||||
kernel void kernel_mul_mm_f16_f32_pair(
|
||||
constant ds4_metal_args_mul_mm & 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 tiitg[[thread_index_in_threadgroup]],
|
||||
ushort sgitg[[simdgroup_index_in_threadgroup]]) {
|
||||
threadgroup half * sa_a = (threadgroup half *)(shmem);
|
||||
threadgroup half * sa_b = (threadgroup half *)(shmem + 4096);
|
||||
threadgroup half * sb = (threadgroup half *)(shmem + 8192);
|
||||
|
||||
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;
|
||||
|
||||
const short nr0 = (args.ne0 - r0 < NR0) ? (args.ne0 - r0) : NR0;
|
||||
const short nr1 = (args.ne1 - r1 < NR1) ? (args.ne1 - r1) : NR1;
|
||||
|
||||
const short lr0 = ((short)tiitg/NL0) < nr0 ? ((short)tiitg/NL0) : nr0 - 1;
|
||||
const short lr1 = ((short)tiitg/NL1) < nr1 ? ((short)tiitg/NL1) : nr1 - 1;
|
||||
|
||||
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 * xa = (device const half4x4 *)(src0_a + args.nb01*(r0 + lr0) + offset0) + offset1;
|
||||
device const half4x4 * xb = (device const half4x4 *)(src0_b + 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);
|
||||
|
||||
simdgroup_half8x8 ma[4];
|
||||
simdgroup_half8x8 mb[2];
|
||||
|
||||
simdgroup_float8x8 mc_a[8];
|
||||
simdgroup_float8x8 mc_b[8];
|
||||
|
||||
for (short i = 0; i < 8; i++) {
|
||||
mc_a[i] = make_filled_simdgroup_matrix<float, 8>(0.f);
|
||||
mc_b[i] = make_filled_simdgroup_matrix<float, 8>(0.f);
|
||||
}
|
||||
|
||||
for (int loop_k = 0; loop_k < args.ne00; loop_k += NK) {
|
||||
half4x4 temp_a;
|
||||
half4x4 temp_b;
|
||||
dequantize_f16(xa, il, temp_a);
|
||||
dequantize_f16(xb, il, temp_b);
|
||||
|
||||
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;
|
||||
|
||||
*(sa_a + 64*ib + 8*ly + lx) = temp_a[i/4][i%4];
|
||||
*(sa_b + 64*ib + 8*ly + lx) = temp_b[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 ? (half) *((device float *) 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 half2x4 *)(sb + 64*ib + 8*ly) = (half2x4)(*((device float2x4 *) y));
|
||||
}
|
||||
|
||||
il = (il + 2 < 1) ? il + 2 : il % 2;
|
||||
xa = (il < 2) ? xa + 2 : xa;
|
||||
xb = (il < 2) ? xb + 2 : xb;
|
||||
|
||||
y += NK;
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
threadgroup const half * lsma_a = (sa_a + 4*64*(sgitg%2));
|
||||
threadgroup const half * lsma_b = (sa_b + 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 < 2; i++) {
|
||||
simdgroup_load(mb[i], lsmb + 64*i, 8, 0, false);
|
||||
}
|
||||
|
||||
simdgroup_barrier(mem_flags::mem_none);
|
||||
|
||||
FOR_UNROLL (short i = 0; i < 4; i++) {
|
||||
simdgroup_load(ma[i], lsma_a + 64*i, 8, 0, false);
|
||||
}
|
||||
|
||||
simdgroup_barrier(mem_flags::mem_none);
|
||||
|
||||
FOR_UNROLL (short i = 0; i < 8; i++) {
|
||||
simdgroup_multiply_accumulate(mc_a[i], mb[i/4], ma[i%4], mc_a[i]);
|
||||
}
|
||||
|
||||
simdgroup_barrier(mem_flags::mem_none);
|
||||
|
||||
FOR_UNROLL (short i = 0; i < 4; i++) {
|
||||
simdgroup_load(ma[i], lsma_b + 64*i, 8, 0, false);
|
||||
}
|
||||
|
||||
simdgroup_barrier(mem_flags::mem_none);
|
||||
|
||||
FOR_UNROLL (short i = 0; i < 8; i++) {
|
||||
simdgroup_multiply_accumulate(mc_b[i], mb[i/4], ma[i%4], mc_b[i]);
|
||||
}
|
||||
|
||||
lsma_a += 8*64;
|
||||
lsma_b += 8*64;
|
||||
lsmb += 4*64;
|
||||
}
|
||||
}
|
||||
|
||||
if (!FC_mul_mm_bc_out || (r0 + NR0 <= args.ne0 && r1 + NR1 <= args.ne1)) {
|
||||
device float * C_a = (device float *) dst_a +
|
||||
(r0 + 32*(sgitg & 1)) +
|
||||
(r1 + 16*(sgitg >> 1)) * args.ne0 + im*args.ne1*args.ne0;
|
||||
device float * C_b = (device float *) dst_b +
|
||||
(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_a[i], C_a + 8*(i%4) + 8*args.ne0*(i/4), args.ne0, 0, false);
|
||||
simdgroup_store(mc_b[i], C_b + 8*(i%4) + 8*args.ne0*(i/4), args.ne0, 0, false);
|
||||
}
|
||||
} else {
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
threadgroup float * temp_str = (threadgroup float *) shmem;
|
||||
|
||||
for (short i = 0; i < 8; i++) {
|
||||
simdgroup_store(mc_a[i],
|
||||
temp_str + 32*(sgitg&1) + (16*(sgitg >> 1))*NR0 + 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_a + 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);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
for (short i = 0; i < 8; i++) {
|
||||
simdgroup_store(mc_b[i],
|
||||
temp_str + 32*(sgitg&1) + (16*(sgitg >> 1))*NR0 + 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_b + 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;
|
||||
|
||||
// Host-visible prefill matmul variants for F16 and Q8_0 weights.
|
||||
|
||||
@@ -21,20 +21,6 @@ struct ds4_metal_args_dsv4_hc_weighted_sum {
|
||||
uint64_t nb1;
|
||||
};
|
||||
|
||||
struct ds4_metal_args_dsv4_hc_weighted_sum_norm {
|
||||
int64_t n_embd;
|
||||
int64_t n_hc;
|
||||
int64_t n_tokens;
|
||||
uint64_t nb_x0;
|
||||
uint64_t nb_x1;
|
||||
uint64_t nb_x2;
|
||||
uint64_t nb_w0;
|
||||
uint64_t nb_w1;
|
||||
uint64_t nb0;
|
||||
uint64_t nb1;
|
||||
uint64_t nb_norm1;
|
||||
float norm_eps;
|
||||
};
|
||||
|
||||
struct ds4_metal_args_dsv4_output_hc_weights4 {
|
||||
float post_scale;
|
||||
@@ -411,6 +397,68 @@ kernel void kernel_dsv4_hc_split_weighted_sum(
|
||||
// kernel_dsv4_hc_split_weighted_sum, stores the HC-pre row for diagnostics, and
|
||||
// reuses the just-collapsed values from threadgroup memory for the RMSNorm
|
||||
// reduction.
|
||||
static __attribute__((always_inline)) inline void ds4_hc_comb_weights4_exact(
|
||||
constant ds4_metal_args_dsv4_hc_split_weighted_sum_norm & args,
|
||||
device volatile const float *mix,
|
||||
device const float *scale,
|
||||
device const float *base,
|
||||
device float *out) {
|
||||
const float epsv = args.eps;
|
||||
const float comb_scale = scale[2];
|
||||
|
||||
float4 r0 =
|
||||
*((device volatile const float4 *)(mix + 8)) * comb_scale +
|
||||
*((device const float4 *)(base + 8));
|
||||
float4 r1 =
|
||||
*((device volatile const float4 *)(mix + 12)) * comb_scale +
|
||||
*((device const float4 *)(base + 12));
|
||||
float4 r2 =
|
||||
*((device volatile const float4 *)(mix + 16)) * comb_scale +
|
||||
*((device const float4 *)(base + 16));
|
||||
float4 r3 =
|
||||
*((device volatile const float4 *)(mix + 20)) * comb_scale +
|
||||
*((device const float4 *)(base + 20));
|
||||
|
||||
const float m0 = max(max(r0.x, r0.y), max(r0.z, r0.w));
|
||||
const float m1 = max(max(r1.x, r1.y), max(r1.z, r1.w));
|
||||
const float m2 = max(max(r2.x, r2.y), max(r2.z, r2.w));
|
||||
const float m3 = max(max(r3.x, r3.y), max(r3.z, r3.w));
|
||||
|
||||
r0 = exp(r0 - m0);
|
||||
r1 = exp(r1 - m1);
|
||||
r2 = exp(r2 - m2);
|
||||
r3 = exp(r3 - m3);
|
||||
|
||||
r0 = r0 * (1.0f / (r0.x + r0.y + r0.z + r0.w)) + epsv;
|
||||
r1 = r1 * (1.0f / (r1.x + r1.y + r1.z + r1.w)) + epsv;
|
||||
r2 = r2 * (1.0f / (r2.x + r2.y + r2.z + r2.w)) + epsv;
|
||||
r3 = r3 * (1.0f / (r3.x + r3.y + r3.z + r3.w)) + epsv;
|
||||
|
||||
float4 col_inv = 1.0f / (r0 + r1 + r2 + r3 + epsv);
|
||||
r0 *= col_inv;
|
||||
r1 *= col_inv;
|
||||
r2 *= col_inv;
|
||||
r3 *= col_inv;
|
||||
|
||||
for (int iter = 1; iter < args.sinkhorn_iters; ++iter) {
|
||||
r0 *= 1.0f / (r0.x + r0.y + r0.z + r0.w + epsv);
|
||||
r1 *= 1.0f / (r1.x + r1.y + r1.z + r1.w + epsv);
|
||||
r2 *= 1.0f / (r2.x + r2.y + r2.z + r2.w + epsv);
|
||||
r3 *= 1.0f / (r3.x + r3.y + r3.z + r3.w + epsv);
|
||||
|
||||
col_inv = 1.0f / (r0 + r1 + r2 + r3 + epsv);
|
||||
r0 *= col_inv;
|
||||
r1 *= col_inv;
|
||||
r2 *= col_inv;
|
||||
r3 *= col_inv;
|
||||
}
|
||||
|
||||
*((device float4 *)(out + 8)) = r0;
|
||||
*((device float4 *)(out + 12)) = r1;
|
||||
*((device float4 *)(out + 16)) = r2;
|
||||
*((device float4 *)(out + 20)) = r3;
|
||||
}
|
||||
|
||||
kernel void kernel_dsv4_hc_split_weighted_sum_norm4(
|
||||
constant ds4_metal_args_dsv4_hc_split_weighted_sum_norm & args,
|
||||
device const char * mixes,
|
||||
@@ -519,7 +567,6 @@ kernel void kernel_dsv4_hc_split_weighted_sum_norm4(
|
||||
}
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
float sumf = 0.0f;
|
||||
for (uint i = tid; i < n4; i += ntg) {
|
||||
device const float4 *x0 = (device const float4 *)(x + 0 * args.nb_x1 + (uint64_t)row * args.nb_x2);
|
||||
@@ -884,6 +931,119 @@ kernel void kernel_dsv4_q8_hc_expand4_q8_0(
|
||||
}
|
||||
}
|
||||
|
||||
kernel void kernel_dsv4_q8_hc_expand4_q8_0_vec_hc(
|
||||
constant ds4_metal_args_mul_mv & mv,
|
||||
constant ds4_metal_args_dsv4_hc_expand & hc,
|
||||
device const char * weight,
|
||||
device const char * input,
|
||||
device char * block_out,
|
||||
device const char * residual,
|
||||
device const char * post,
|
||||
device const char * comb,
|
||||
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]]) {
|
||||
if (hc.n_hc != 4 || hc.n_tokens != 1) {
|
||||
return;
|
||||
}
|
||||
|
||||
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 = mv.ne00 / QK8_0;
|
||||
const int row0 = tgpig.x * NR0;
|
||||
|
||||
const short ix = tiisg / (NW / NQ);
|
||||
const short il = tiisg % (NW / NQ);
|
||||
const int ib0 = sgitg * NQ + ix;
|
||||
|
||||
device const float *y = (device const float *)(input);
|
||||
device const float *yb = y + ib0 * QK8_0 + il * NQ;
|
||||
|
||||
device const block_q8_0 *ax[NR0];
|
||||
FOR_UNROLL(short row = 0; row < NR0; ++row) {
|
||||
const uint64_t off0 = (uint64_t)(row0 + row) * mv.nb01;
|
||||
ax[row] = (device const block_q8_0 *)(weight + off0);
|
||||
}
|
||||
|
||||
float sumf[NR0] = { 0.0f };
|
||||
float yl[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 *qs = ax[row][ib].qs + il * NQ;
|
||||
|
||||
float sumq = 0.0f;
|
||||
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;
|
||||
}
|
||||
|
||||
threadgroup float *shmem_f32[NR0];
|
||||
FOR_UNROLL(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_UNROLL(short row = 0; row < NR0; ++row) {
|
||||
if (tiisg == 0) {
|
||||
shmem_f32[row][sgitg] = sumf[row];
|
||||
}
|
||||
}
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
FOR_UNROLL(short row = 0; row < NR0; ++row) {
|
||||
const int d = row0 + row;
|
||||
if (d >= mv.ne01) {
|
||||
continue;
|
||||
}
|
||||
|
||||
const float block_v = simd_sum(shmem_f32[row][tiisg]);
|
||||
if (tiisg == 0 && sgitg == 0) {
|
||||
*((device float *)(block_out + (uint64_t)d * sizeof(float))) = block_v;
|
||||
|
||||
const float r0 = *((device const float *)(residual + (uint64_t)d * hc.nb_res0 + 0 * hc.nb_res1));
|
||||
const float r1 = *((device const float *)(residual + (uint64_t)d * hc.nb_res0 + 1 * hc.nb_res1));
|
||||
const float r2 = *((device const float *)(residual + (uint64_t)d * hc.nb_res0 + 2 * hc.nb_res1));
|
||||
const float r3 = *((device const float *)(residual + (uint64_t)d * hc.nb_res0 + 3 * hc.nb_res1));
|
||||
|
||||
const float4 post4 = *((device const float4 *)post);
|
||||
const float4 comb0 = *((device const float4 *)(comb + 0 * hc.nb_comb1));
|
||||
const float4 comb1 = *((device const float4 *)(comb + 1 * hc.nb_comb1));
|
||||
const float4 comb2 = *((device const float4 *)(comb + 2 * hc.nb_comb1));
|
||||
const float4 comb3 = *((device const float4 *)(comb + 3 * hc.nb_comb1));
|
||||
float4 acc = block_v * post4;
|
||||
acc += comb0 * r0;
|
||||
acc += comb1 * r1;
|
||||
acc += comb2 * r2;
|
||||
acc += comb3 * r3;
|
||||
FOR_UNROLL (short dst_hc = 0; dst_hc < 4; ++dst_hc) {
|
||||
*((device float *)(dst + (uint64_t)d * hc.nb0 +
|
||||
(uint64_t)dst_hc * hc.nb1)) = acc[dst_hc];
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Reduces HC channels to a normal embedding row with the learned pre weights.
|
||||
// This is the input adapter before the attention block and before the FFN block.
|
||||
kernel void kernel_dsv4_hc_weighted_sum(
|
||||
@@ -910,76 +1070,6 @@ kernel void kernel_dsv4_hc_weighted_sum(
|
||||
*((device float *) (dst + d*args.nb0 + t*args.nb1)) = acc;
|
||||
}
|
||||
|
||||
// The one-row output head immediately applies a learned RMSNorm after reducing
|
||||
// its four HC streams. Preserve the standalone scalar HC accumulation, write
|
||||
// the collapsed row for diagnostics, then reload its F32 values from
|
||||
// threadgroup memory using the standalone RMSNorm's float4 reduction mapping.
|
||||
kernel void kernel_dsv4_hc_weighted_sum_norm4(
|
||||
constant ds4_metal_args_dsv4_hc_weighted_sum_norm & args,
|
||||
device const char * x,
|
||||
device const char * weights,
|
||||
device char * dst,
|
||||
device const char * norm_weight,
|
||||
device char * norm_dst,
|
||||
threadgroup float * shared [[threadgroup(0)]],
|
||||
ushort tid [[thread_position_in_threadgroup]],
|
||||
ushort sgitg [[simdgroup_index_in_threadgroup]],
|
||||
ushort tiisg [[thread_index_in_simdgroup]],
|
||||
ushort ntg [[threads_per_threadgroup]]) {
|
||||
if (args.n_tokens != 1 || args.n_hc != 4 ||
|
||||
args.n_embd <= 0 || (args.n_embd & 3) != 0) {
|
||||
return;
|
||||
}
|
||||
|
||||
const uint n_embd = uint(args.n_embd);
|
||||
const uint n4 = n_embd >> 2;
|
||||
threadgroup float *row_shmem = shared;
|
||||
threadgroup float *sum_shmem = shared + n_embd;
|
||||
|
||||
if (sgitg == 0) {
|
||||
sum_shmem[tiisg] = 0.0f;
|
||||
}
|
||||
|
||||
for (uint d = tid; d < n_embd; d += ntg) {
|
||||
float acc = 0.0f;
|
||||
for (int64_t h = 0; h < args.n_hc; ++h) {
|
||||
const float xv = *((device const float *)(
|
||||
x + (uint64_t)d*args.nb_x0 + (uint64_t)h*args.nb_x1));
|
||||
const float wv = *((device const float *)(
|
||||
weights + (uint64_t)h*args.nb_w0));
|
||||
acc += xv * wv;
|
||||
}
|
||||
row_shmem[d] = acc;
|
||||
*((device float *)(dst + (uint64_t)d*args.nb0)) = acc;
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
threadgroup const float4 *row4 =
|
||||
(threadgroup const float4 *)row_shmem;
|
||||
float sumf = 0.0f;
|
||||
for (uint i = tid; i < n4; i += ntg) {
|
||||
sumf += dot(row4[i], row4[i]);
|
||||
}
|
||||
sumf = simd_sum(sumf);
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
if (tiisg == 0) {
|
||||
sum_shmem[sgitg] = sumf;
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
sumf = sum_shmem[tiisg];
|
||||
sumf = simd_sum(sumf);
|
||||
|
||||
const float mean = sumf/args.n_embd;
|
||||
const float scale = 1.0f/sqrt(mean + args.norm_eps);
|
||||
device const float4 *w4 = (device const float4 *)norm_weight;
|
||||
device float4 *norm4 = (device float4 *)norm_dst;
|
||||
for (uint i = tid; i < n4; i += ntg) {
|
||||
norm4[i] = (row4[i]*scale)*w4[i];
|
||||
}
|
||||
}
|
||||
|
||||
// The one-row HC=4 output head historically materializes four device-F32
|
||||
// stages across separate launches. Collapse those launches into one tiny
|
||||
// two-thread group while preserving the scalar/vector lane mapping and every
|
||||
@@ -1015,3 +1105,440 @@ kernel void kernel_dsv4_output_hc_weights4(
|
||||
args.post_scale * x + args.eps;
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
struct ds4_metal_args_hc_norm_mix {
|
||||
int32_t n;
|
||||
int32_t out_dim;
|
||||
float eps;
|
||||
};
|
||||
|
||||
// Fused unweighted RMSNorm + F16 HC-mix projection for DS4 decode HC-pre.
|
||||
// The standalone decode path runs kernel_rms_norm_f32_4 over the flattened
|
||||
// 4*embd HC row (1024 threads, one threadgroup) and then
|
||||
// kernel_mul_mv_f16_f32_4 (nsg=8, nr0=2) over the normalized row. Both
|
||||
// stages are reproduced bit-exactly in one dispatch: every threadgroup
|
||||
// redundantly recomputes the norm partials with the original 1024-thread
|
||||
// mapping (each real lane covers one virtual thread of each 256-thread
|
||||
// slice, preserving every simd_sum tree), and the matvec keeps the original
|
||||
// per-row accumulation order with y = x*scale computed on the fly, which
|
||||
// rounds identically to the materialized normalized row. The host wrapper
|
||||
// gates this to n == 16384 && out_dim == 24, where the virtual-thread count
|
||||
// is exactly 1024 and the mv tail loop is empty.
|
||||
kernel void kernel_dsv4_hc_rms_norm_mix_f16(
|
||||
constant ds4_metal_args_hc_norm_mix & args,
|
||||
device const char * x,
|
||||
device const char * weight,
|
||||
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]]) {
|
||||
constexpr short NSG = 8; // ds4_gpu_make_plain_mv_dispatch(16384)
|
||||
constexpr short NW = N_SIMDWIDTH;
|
||||
constexpr short NR0 = 2; // plain mv nr0
|
||||
constexpr short NB = 32;
|
||||
constexpr short NF = 16;
|
||||
constexpr short NF4 = NF/4;
|
||||
constexpr uint VTHREADS = 1024u; // rms norm threads at n == 16384
|
||||
constexpr short VSLICES = VTHREADS/(NSG*NW); // virtual 256-thread slices
|
||||
|
||||
const uint n = (uint)args.n;
|
||||
const uint n4 = n >> 2;
|
||||
|
||||
device const float4 *x4 = (device const float4 *)x;
|
||||
|
||||
threadgroup float *norm_shmem = (threadgroup float *)shmem; // NW slots
|
||||
threadgroup float *mv_shmem = (threadgroup float *)shmem + NW; // NW*NR0 slots
|
||||
|
||||
// Phase A: exact replica of kernel_rms_norm_f32_4's reduction tree with
|
||||
// the 1024 virtual threads folded onto this threadgroup's 8 simdgroups.
|
||||
for (short v = 0; v < VSLICES; ++v) {
|
||||
const uint vt = (uint)(sgitg + NSG*v)*NW + tiisg;
|
||||
float sumf = 0.0f;
|
||||
for (uint i00 = vt; i00 < n4; i00 += VTHREADS) {
|
||||
sumf += dot(x4[i00], x4[i00]);
|
||||
}
|
||||
sumf = simd_sum(sumf);
|
||||
if (tiisg == 0) {
|
||||
norm_shmem[sgitg + NSG*v] = sumf;
|
||||
}
|
||||
}
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
float total = norm_shmem[tiisg];
|
||||
total = simd_sum(total);
|
||||
const float mean = total/(float)args.n;
|
||||
const float scale = 1.0f/sqrt(mean + args.eps);
|
||||
|
||||
// Phase B: exact replica of kernel_mul_mv_f16_f32_4 (nsg=8, nr0=2) with
|
||||
// the normalized operand recomputed as x*scale instead of reloaded.
|
||||
const int nb = args.n/NB;
|
||||
const int r0 = tgpig.x*NR0;
|
||||
|
||||
device const half4 * ax4[NR0];
|
||||
FOR_UNROLL (short row = 0; row < NR0; ++row) {
|
||||
ax4[row] = (device const half4 *)
|
||||
(weight + (uint64_t)(r0 + row)*(uint64_t)n*sizeof(half));
|
||||
}
|
||||
|
||||
float sumf_mv[NR0] = { 0.f };
|
||||
|
||||
const short ix = tiisg/(NW/NF);
|
||||
const short il = tiisg%(NW/NF);
|
||||
const int ib0 = sgitg*NF + ix;
|
||||
|
||||
for (int ib = ib0; ib < nb; ib += NSG*NF) {
|
||||
float4 yl4[NF4];
|
||||
FOR_UNROLL (short i = 0; i < NF4; ++i) {
|
||||
yl4[i] = x4[(ib*NB + il*NF)/4 + i]*scale;
|
||||
}
|
||||
|
||||
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_mv[row] += sumq;
|
||||
}
|
||||
}
|
||||
|
||||
// n == 16384 makes the scalar tail loop of the original empty.
|
||||
device float * dst_f32 = (device float *) dst;
|
||||
helper_mv_reduce_and_write<NR0>(dst_f32, sumf_mv, r0, args.out_dim,
|
||||
tiisg, sgitg, (threadgroup char *)mv_shmem);
|
||||
}
|
||||
|
||||
// M5 specialization: pack two exact NR0=2 HC-mix producer groups into one
|
||||
// 512-thread group. Two independent eight-simdgroup clusters retain the
|
||||
// matvec reductions while the exact RMS scale is redundantly formed six,
|
||||
// rather than twelve, times.
|
||||
kernel void kernel_dsv4_hc_rms_norm_mix_f16_cluster2(
|
||||
constant ds4_metal_args_hc_norm_mix & args,
|
||||
device const char * x,
|
||||
device const char * weight,
|
||||
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]]) {
|
||||
constexpr short NSG_CLUSTER = 8;
|
||||
constexpr short NCLUSTER = 2;
|
||||
constexpr short NSG_TOTAL = NSG_CLUSTER * NCLUSTER;
|
||||
constexpr short NW = N_SIMDWIDTH;
|
||||
constexpr short NR0 = 2;
|
||||
constexpr short NB = 32;
|
||||
constexpr short NF = 16;
|
||||
constexpr short NF4 = NF/4;
|
||||
constexpr uint VTHREADS = 1024u;
|
||||
constexpr short VSLICES = VTHREADS/(NSG_TOTAL*NW);
|
||||
|
||||
const uint n = (uint)args.n;
|
||||
const uint n4 = n >> 2;
|
||||
device const float4 *x4 = (device const float4 *)x;
|
||||
threadgroup float *norm_shmem = (threadgroup float *)shmem;
|
||||
threadgroup float *mv_shmem = norm_shmem + NW;
|
||||
|
||||
// Exact 1024-virtual-thread RMS reduction, now folded two ways over
|
||||
// the 16 physical simdgroups instead of four ways over eight.
|
||||
for (short v = 0; v < VSLICES; ++v) {
|
||||
const uint vt = (uint)(sgitg + NSG_TOTAL*v)*NW + tiisg;
|
||||
float sumf = 0.0f;
|
||||
for (uint i00 = vt; i00 < n4; i00 += VTHREADS) {
|
||||
sumf += dot(x4[i00], x4[i00]);
|
||||
}
|
||||
sumf = simd_sum(sumf);
|
||||
if (tiisg == 0) {
|
||||
norm_shmem[sgitg + NSG_TOTAL*v] = sumf;
|
||||
}
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
float total = norm_shmem[tiisg];
|
||||
total = simd_sum(total);
|
||||
const float mean = total/(float)args.n;
|
||||
const float scale = 1.0f/sqrt(mean + args.eps);
|
||||
|
||||
// Two independent eight-simdgroup clusters reproduce two original
|
||||
// NR0=2 matvec threadgroups inside this 512-thread threadgroup.
|
||||
const short cluster = sgitg / NSG_CLUSTER;
|
||||
const short local_sg = sgitg - cluster*NSG_CLUSTER;
|
||||
const int nb = args.n/NB;
|
||||
const int r0 = (int)tgpig.x*(NCLUSTER*NR0) + cluster*NR0;
|
||||
|
||||
device const half4 *ax4[NR0];
|
||||
FOR_UNROLL (short row = 0; row < NR0; ++row) {
|
||||
ax4[row] = (device const half4 *)
|
||||
(weight + (uint64_t)(r0 + row)*(uint64_t)n*sizeof(half));
|
||||
}
|
||||
|
||||
float sumf_mv[NR0] = { 0.f };
|
||||
const short ix = tiisg/(NW/NF);
|
||||
const short il = tiisg%(NW/NF);
|
||||
const int ib0 = local_sg*NF + ix;
|
||||
for (int ib = ib0; ib < nb; ib += NSG_CLUSTER*NF) {
|
||||
float4 yl4[NF4];
|
||||
FOR_UNROLL (short i = 0; i < NF4; ++i) {
|
||||
yl4[i] = x4[(ib*NB + il*NF)/4 + i]*scale;
|
||||
}
|
||||
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_mv[row] += sumq;
|
||||
}
|
||||
}
|
||||
|
||||
threadgroup float *cluster_shmem[NR0];
|
||||
FOR_UNROLL (short row = 0; row < NR0; ++row) {
|
||||
cluster_shmem[row] = mv_shmem +
|
||||
((uint)cluster*NR0 + row)*NW;
|
||||
if (local_sg == 0) {
|
||||
cluster_shmem[row][tiisg] = 0.0f;
|
||||
}
|
||||
sumf_mv[row] = simd_sum(sumf_mv[row]);
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
FOR_UNROLL (short row = 0; row < NR0; ++row) {
|
||||
if (tiisg == 0) {
|
||||
cluster_shmem[row][local_sg] = sumf_mv[row];
|
||||
}
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
device float *mixes_f32 = (device float *)dst;
|
||||
FOR_UNROLL (short row = 0; row < NR0; ++row) {
|
||||
const float tot = simd_sum(cluster_shmem[row][tiisg]);
|
||||
if (tiisg == 0 && local_sg == 0 && r0 + row < args.out_dim) {
|
||||
mixes_f32[r0 + row] = tot;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
kernel void kernel_dsv4_hc_rms_norm_mix_f16_cluster2_pre_norm(
|
||||
constant ds4_metal_args_hc_norm_mix & args,
|
||||
constant ds4_metal_args_dsv4_hc_split_weighted_sum_norm & split_args,
|
||||
device const char * x,
|
||||
device const char * weight,
|
||||
device char * dst,
|
||||
device const float * hc_scale,
|
||||
device const float * hc_base,
|
||||
device char * split,
|
||||
device char * collapse_dst,
|
||||
device const char * norm_weight,
|
||||
device char * norm_dst,
|
||||
device atomic_uint * completion,
|
||||
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 NSG_CLUSTER = 8;
|
||||
constexpr short NCLUSTER = 2;
|
||||
constexpr short NSG_TOTAL = NSG_CLUSTER * NCLUSTER;
|
||||
constexpr short NW = N_SIMDWIDTH;
|
||||
constexpr short NR0 = 2;
|
||||
constexpr short NB = 32;
|
||||
constexpr short NF = 16;
|
||||
constexpr short NF4 = NF/4;
|
||||
constexpr uint VTHREADS = 1024u;
|
||||
constexpr short VSLICES = VTHREADS/(NSG_TOTAL*NW);
|
||||
|
||||
const uint n = (uint)args.n;
|
||||
const uint n4 = n >> 2;
|
||||
device const float4 *x4 = (device const float4 *)x;
|
||||
threadgroup float *norm_shmem = (threadgroup float *)shmem;
|
||||
threadgroup float *mv_shmem = norm_shmem + NW;
|
||||
|
||||
// Exact 1024-virtual-thread RMS reduction, now folded two ways over
|
||||
// the 16 physical simdgroups instead of four ways over eight.
|
||||
for (short v = 0; v < VSLICES; ++v) {
|
||||
const uint vt = (uint)(sgitg + NSG_TOTAL*v)*NW + tiisg;
|
||||
float sumf = 0.0f;
|
||||
for (uint i00 = vt; i00 < n4; i00 += VTHREADS) {
|
||||
sumf += dot(x4[i00], x4[i00]);
|
||||
}
|
||||
sumf = simd_sum(sumf);
|
||||
if (tiisg == 0) {
|
||||
norm_shmem[sgitg + NSG_TOTAL*v] = sumf;
|
||||
}
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
float total = norm_shmem[tiisg];
|
||||
total = simd_sum(total);
|
||||
const float mean = total/(float)args.n;
|
||||
const float scale = 1.0f/sqrt(mean + args.eps);
|
||||
|
||||
// Two independent eight-simdgroup clusters reproduce two original
|
||||
// NR0=2 matvec threadgroups inside this 512-thread threadgroup.
|
||||
const short cluster = sgitg / NSG_CLUSTER;
|
||||
const short local_sg = sgitg - cluster*NSG_CLUSTER;
|
||||
const int nb = args.n/NB;
|
||||
const int r0 = (int)tgpig.x*(NCLUSTER*NR0) + cluster*NR0;
|
||||
|
||||
device const half4 *ax4[NR0];
|
||||
FOR_UNROLL (short row = 0; row < NR0; ++row) {
|
||||
ax4[row] = (device const half4 *)
|
||||
(weight + (uint64_t)(r0 + row)*(uint64_t)n*sizeof(half));
|
||||
}
|
||||
|
||||
float sumf_mv[NR0] = { 0.f };
|
||||
const short ix = tiisg/(NW/NF);
|
||||
const short il = tiisg%(NW/NF);
|
||||
const int ib0 = local_sg*NF + ix;
|
||||
for (int ib = ib0; ib < nb; ib += NSG_CLUSTER*NF) {
|
||||
float4 yl4[NF4];
|
||||
FOR_UNROLL (short i = 0; i < NF4; ++i) {
|
||||
yl4[i] = x4[(ib*NB + il*NF)/4 + i]*scale;
|
||||
}
|
||||
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_mv[row] += sumq;
|
||||
}
|
||||
}
|
||||
|
||||
threadgroup float *cluster_shmem[NR0];
|
||||
FOR_UNROLL (short row = 0; row < NR0; ++row) {
|
||||
cluster_shmem[row] = mv_shmem +
|
||||
((uint)cluster*NR0 + row)*NW;
|
||||
if (local_sg == 0) {
|
||||
cluster_shmem[row][tiisg] = 0.0f;
|
||||
}
|
||||
sumf_mv[row] = simd_sum(sumf_mv[row]);
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
FOR_UNROLL (short row = 0; row < NR0; ++row) {
|
||||
if (tiisg == 0) {
|
||||
cluster_shmem[row][local_sg] = sumf_mv[row];
|
||||
}
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
device volatile float *mixes_f32 =
|
||||
(device volatile float *)dst;
|
||||
if (local_sg == 0) {
|
||||
FOR_UNROLL (short row = 0; row < NR0; ++row) {
|
||||
const float tot = simd_sum(cluster_shmem[row][tiisg]);
|
||||
if (tiisg == 0 && r0 + row < args.out_dim) {
|
||||
mixes_f32[r0 + row] = tot;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// The first producer group owns mix[0:4]. After materializing and
|
||||
// reloading those values, fold the established 1024-thread HC collapse
|
||||
// and RMS reduction over this group's 512 physical threads as two
|
||||
// independent virtual slices. This retains the original 32-partial tree.
|
||||
threadgroup_barrier(mem_flags::mem_device_and_threadgroup);
|
||||
const uint tid = (uint)sgitg * (uint)NW + (uint)tiisg;
|
||||
threadgroup float *pre_shmem = norm_shmem + 32u + 4u*NW;
|
||||
threadgroup float *sum_shmem = pre_shmem + 4;
|
||||
|
||||
if (tgpig.x == 0) {
|
||||
device float *out = (device float *)split;
|
||||
if (tid == 0) {
|
||||
const float4 pre_z =
|
||||
*((device volatile const float4 *)mixes_f32) * hc_scale[0] +
|
||||
*((device const float4 *)hc_base);
|
||||
const float4 pre =
|
||||
1.0f / (1.0f + exp(-pre_z)) + split_args.eps;
|
||||
*((device float4 *)out) = pre;
|
||||
pre_shmem[0] = pre.x;
|
||||
pre_shmem[1] = pre.y;
|
||||
pre_shmem[2] = pre.z;
|
||||
pre_shmem[3] = pre.w;
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
const uint n4_collapse = uint(split_args.n_embd) >> 2;
|
||||
const uint i0 = tid;
|
||||
const uint i1 = tid + 512u;
|
||||
device const float4 *x0 = (device const float4 *)(
|
||||
x + 0 * split_args.nb_x1);
|
||||
device const float4 *x1 = (device const float4 *)(
|
||||
x + 1 * split_args.nb_x1);
|
||||
device const float4 *x2 = (device const float4 *)(
|
||||
x + 2 * split_args.nb_x1);
|
||||
device const float4 *x3 = (device const float4 *)(
|
||||
x + 3 * split_args.nb_x1);
|
||||
|
||||
float4 v0 = 0.0f;
|
||||
v0 += x0[i0] * pre_shmem[0];
|
||||
v0 += x1[i0] * pre_shmem[1];
|
||||
v0 += x2[i0] * pre_shmem[2];
|
||||
v0 += x3[i0] * pre_shmem[3];
|
||||
float sum0 = simd_sum(dot(v0, v0));
|
||||
|
||||
float4 v1 = 0.0f;
|
||||
if (i1 < n4_collapse) {
|
||||
v1 += x0[i1] * pre_shmem[0];
|
||||
v1 += x1[i1] * pre_shmem[1];
|
||||
v1 += x2[i1] * pre_shmem[2];
|
||||
v1 += x3[i1] * pre_shmem[3];
|
||||
}
|
||||
float sum1 = simd_sum(dot(v1, v1));
|
||||
if (tiisg == 0) {
|
||||
sum_shmem[sgitg] = sum0;
|
||||
sum_shmem[sgitg + 16] = sum1;
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
float sumf = sum_shmem[tiisg];
|
||||
sumf = simd_sum(sumf);
|
||||
const float norm_arg =
|
||||
sumf / float(split_args.n_embd) + split_args.norm_eps;
|
||||
const float norm_scale = rsqrt(norm_arg);
|
||||
device float4 *dst4 = (device float4 *)collapse_dst;
|
||||
device const float4 *w4 = (device const float4 *)norm_weight;
|
||||
device float4 *norm4 = (device float4 *)norm_dst;
|
||||
dst4[i0] = v0;
|
||||
norm4[i0] = (v0 * norm_scale) * w4[i0];
|
||||
if (i1 < n4_collapse) {
|
||||
dst4[i1] = v1;
|
||||
norm4[i1] = (v1 * norm_scale) * w4[i1];
|
||||
}
|
||||
} else if (tgpig.x == 1 && tid == 0) {
|
||||
device float *out = (device float *)split;
|
||||
const float4 post_z =
|
||||
*((device volatile const float4 *)(mixes_f32 + 4)) * hc_scale[1] +
|
||||
*((device const float4 *)(hc_base + 4));
|
||||
*((device float4 *)(out + 4)) = 2.0f / (1.0f + exp(-post_z));
|
||||
}
|
||||
|
||||
// Groups 2..5 own exactly the comb range consumed by the
|
||||
// continuation. Their four-way completion overlaps TG0's independent
|
||||
// pre-collapse/RMS epilogue. Every writer crosses the uniform publish
|
||||
// fence; only lane zero then participates in the completion protocol.
|
||||
atomic_thread_fence(mem_flags::mem_device,
|
||||
memory_order_seq_cst,
|
||||
thread_scope_device);
|
||||
if (tgpig.x < 2 || tid != 0) {
|
||||
return;
|
||||
}
|
||||
|
||||
const uint old = atomic_fetch_add_explicit(
|
||||
completion, 1u, memory_order_relaxed);
|
||||
if (old + 1u != 4u) {
|
||||
return;
|
||||
}
|
||||
atomic_thread_fence(mem_flags::mem_device,
|
||||
memory_order_seq_cst,
|
||||
thread_scope_device);
|
||||
ds4_hc_comb_weights4_exact(
|
||||
split_args, mixes_f32, hc_scale, hc_base,
|
||||
(device float *)split);
|
||||
atomic_thread_fence(mem_flags::mem_device,
|
||||
memory_order_seq_cst,
|
||||
thread_scope_device);
|
||||
atomic_store_explicit(completion, 0u, memory_order_relaxed);
|
||||
}
|
||||
|
||||
@@ -319,6 +319,228 @@ kernel void kernel_dsv4_compressor_pack_ratio4(
|
||||
}
|
||||
}
|
||||
|
||||
// Decode already holds the complete ratio-4 recurrent window in state layout:
|
||||
// eight rows of two head_dim planes. Pack the previous plane from rows 0..3
|
||||
// and the current plane from rows 4..7 directly into the transposed [head_dim,
|
||||
// 8] layout consumed by the exact GGML softmax/multiply/sum sequence. KV and
|
||||
// score move together; no arithmetic or reduction order changes.
|
||||
kernel void kernel_dsv4_compressor_pack_ratio4_decode_ggml(
|
||||
constant ds4_metal_args_dsv4_compressor_pack_ratio4 & args,
|
||||
device const uint * state_kv,
|
||||
device const uint * state_score,
|
||||
device uint * packed_kv,
|
||||
device uint * packed_score,
|
||||
uint row [[threadgroup_position_in_grid]],
|
||||
uint tid [[thread_index_in_threadgroup]]) {
|
||||
if (row >= 8u || args.head_dim == 0u || args.n_threads == 0u) {
|
||||
return;
|
||||
}
|
||||
|
||||
const uint64_t state_row_stride = 2ull * args.head_dim;
|
||||
const uint64_t src_plane = row >= 4u ? args.head_dim : 0u;
|
||||
for (uint col = tid; col < args.head_dim; col += args.n_threads) {
|
||||
const uint64_t src = (uint64_t)row * state_row_stride +
|
||||
src_plane + col;
|
||||
const uint64_t dst = (uint64_t)col * 8u + row;
|
||||
packed_kv[dst] = state_kv[src];
|
||||
packed_score[dst] = state_score[src];
|
||||
}
|
||||
}
|
||||
|
||||
// Exact decode specialization for the first two operations in GGML's
|
||||
// softmax -> multiply -> sum_rows compressor reduction. The normalized
|
||||
// softmax values are deliberately materialized in device memory and reloaded
|
||||
// after a device barrier before the in-place product, preserving the dispatch
|
||||
// boundary's float store/load semantics. The final sum remains the standalone
|
||||
// eight-thread sum_rows kernel: changing a 32-thread group into the original
|
||||
// eight-thread reduction inside this kernel would make its threadgroup
|
||||
// barriers non-uniform or alter simd_sum's active-lane topology.
|
||||
kernel void kernel_dsv4_compressor_exact_softmax_product_ratio4(
|
||||
constant ds4_metal_args_dsv4_compressor_pack_ratio4 & args,
|
||||
device const float * packed_kv,
|
||||
device const float * packed_score,
|
||||
device float * softmax,
|
||||
device float * product,
|
||||
threadgroup float * softmax_scratch [[threadgroup(0)]],
|
||||
uint row [[threadgroup_position_in_grid]],
|
||||
uint tid [[thread_position_in_threadgroup]]) {
|
||||
if (row >= args.head_dim || args.n_comp != 1u ||
|
||||
args.n_threads != 32u) {
|
||||
return;
|
||||
}
|
||||
|
||||
device const float4 * score4 =
|
||||
(device const float4 *)(packed_score + (uint64_t)row * 8u);
|
||||
device float4 * softmax4 =
|
||||
(device float4 *)(softmax + (uint64_t)row * 8u);
|
||||
const float scale = (float)args.replay;
|
||||
const float zero = (float)(args.n_comp - 1u);
|
||||
|
||||
// Match kernel_soft_max_f32_4(width=8, nth=32) literally. Only lanes zero
|
||||
// and one own float4s, while all 32 lanes participate in both reductions.
|
||||
float4 lmax4 = -INFINITY;
|
||||
for (int i00 = (int)tid; i00 < 2; i00 += 32) {
|
||||
lmax4 = fmax(lmax4, score4[i00] * scale + (float4)zero);
|
||||
}
|
||||
|
||||
const float lmax =
|
||||
MAX(MAX(lmax4[0], lmax4[1]), MAX(lmax4[2], lmax4[3]));
|
||||
const float max_val = simd_max(lmax);
|
||||
|
||||
float4 lsum4 = 0.0f;
|
||||
for (int i00 = (int)tid; i00 < 2; i00 += 32) {
|
||||
const float4 exp_score4 =
|
||||
exp((score4[i00] * scale + (float4)zero) - max_val);
|
||||
lsum4 += exp_score4;
|
||||
softmax4[i00] = exp_score4;
|
||||
}
|
||||
|
||||
const float lsum =
|
||||
lsum4[0] + lsum4[1] + lsum4[2] + lsum4[3];
|
||||
threadgroup_barrier(mem_flags::mem_none);
|
||||
const float sum = simd_sum(lsum);
|
||||
const float inv_sum = 1.0f / sum;
|
||||
|
||||
for (int i00 = (int)tid; i00 < 2; i00 += 32) {
|
||||
softmax4[i00] *= inv_sum;
|
||||
}
|
||||
|
||||
// Force the same normalized-softmax device store/reload boundary that the
|
||||
// separate multiply dispatch observes.
|
||||
threadgroup_barrier(mem_flags::mem_device);
|
||||
device volatile const float * reloaded_softmax =
|
||||
(device volatile const float *)(softmax + (uint64_t)row * 8u);
|
||||
device const float * kv_row = packed_kv + (uint64_t)row * 8u;
|
||||
device float * product_row = product + (uint64_t)row * 8u;
|
||||
|
||||
// Match kernel_bin_fuse_f32_f32_f32(width=8, nth=4): four lanes each
|
||||
// process their low element followed by the element four positions later.
|
||||
if (tid < 4u) {
|
||||
for (uint i0 = tid; i0 < 8u; i0 += 4u) {
|
||||
float value = kv_row[i0];
|
||||
value *= reloaded_softmax[i0];
|
||||
product_row[i0] = value;
|
||||
}
|
||||
}
|
||||
|
||||
// All 32 lanes reach the final device barrier. The following standalone
|
||||
// sum_rows dispatch performs the required global reload and exact TG8
|
||||
// two-stage simd_sum topology.
|
||||
threadgroup_barrier(mem_flags::mem_device);
|
||||
(void)softmax_scratch;
|
||||
}
|
||||
|
||||
// Exact one-dispatch ratio-4 decode pool. This specializes the three-dispatch
|
||||
// pack -> exact softmax/product -> sum_rows chain above without changing any
|
||||
// floating-point operation or reduction topology. The normalized softmax and
|
||||
// product are still materialized and volatile-reloaded through device memory.
|
||||
// The two simd_sum calls in the final reduction execute under an eight-lane
|
||||
// active mask, exactly matching kernel_sum_rows_f32_f32's original TG8.
|
||||
kernel void kernel_dsv4_compressor_exact_pool_ratio4_decode_ggml(
|
||||
constant ds4_metal_args_dsv4_compressor_pack_ratio4 & args,
|
||||
device const float * state_kv,
|
||||
device const float * state_score,
|
||||
device float * softmax,
|
||||
device float * product,
|
||||
device float * dst,
|
||||
threadgroup float * sum_scratch [[threadgroup(0)]],
|
||||
uint col [[threadgroup_position_in_grid]],
|
||||
uint tid [[thread_position_in_threadgroup]]) {
|
||||
if (col >= args.head_dim || args.n_comp != 1u ||
|
||||
args.n_threads != 32u) {
|
||||
return;
|
||||
}
|
||||
|
||||
const uint64_t state_row_stride = 2ull * args.head_dim;
|
||||
const float scale = (float)args.replay;
|
||||
const float zero = (float)(args.n_comp - 1u);
|
||||
|
||||
// Match the packed float4 ownership: lane 0 owns rows 0..3 and lane 1
|
||||
// rows 4..7. The gather itself is an integer-addressed bit-preserving load.
|
||||
float4 score_values = -INFINITY;
|
||||
if (tid < 2u) {
|
||||
const uint row0 = 4u * tid;
|
||||
for (uint j = 0u; j < 4u; ++j) {
|
||||
const uint row = row0 + j;
|
||||
const uint64_t src = (uint64_t)row * state_row_stride +
|
||||
(row >= 4u ? args.head_dim : 0u) + col;
|
||||
score_values[j] = state_score[src];
|
||||
}
|
||||
}
|
||||
|
||||
const uint64_t scratch_base = (uint64_t)col * 8u;
|
||||
device float4 * softmax4 =
|
||||
(device float4 *)(softmax + scratch_base);
|
||||
|
||||
// Verbatim kernel_soft_max_f32_4(width=8, nth=32) arithmetic.
|
||||
float4 lmax4 = -INFINITY;
|
||||
for (int i00 = (int)tid; i00 < 2; i00 += 32) {
|
||||
lmax4 = fmax(lmax4, score_values * scale + (float4)zero);
|
||||
}
|
||||
const float lmax =
|
||||
MAX(MAX(lmax4[0], lmax4[1]), MAX(lmax4[2], lmax4[3]));
|
||||
const float max_val = simd_max(lmax);
|
||||
|
||||
float4 lsum4 = 0.0f;
|
||||
for (int i00 = (int)tid; i00 < 2; i00 += 32) {
|
||||
const float4 exp_score4 =
|
||||
exp((score_values * scale + (float4)zero) - max_val);
|
||||
lsum4 += exp_score4;
|
||||
softmax4[i00] = exp_score4;
|
||||
}
|
||||
const float lsum =
|
||||
lsum4[0] + lsum4[1] + lsum4[2] + lsum4[3];
|
||||
threadgroup_barrier(mem_flags::mem_none);
|
||||
const float sum = simd_sum(lsum);
|
||||
const float inv_sum = 1.0f / sum;
|
||||
for (int i00 = (int)tid; i00 < 2; i00 += 32) {
|
||||
softmax4[i00] *= inv_sum;
|
||||
}
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_device);
|
||||
device volatile const float * reloaded_softmax =
|
||||
(device volatile const float *)(softmax + scratch_base);
|
||||
device float * product_row = product + scratch_base;
|
||||
|
||||
// Verbatim width=8, TG4 multiply ownership: low element, then +4.
|
||||
if (tid < 4u) {
|
||||
for (uint i0 = tid; i0 < 8u; i0 += 4u) {
|
||||
const uint64_t src = (uint64_t)i0 * state_row_stride +
|
||||
(i0 >= 4u ? args.head_dim : 0u) + col;
|
||||
float value = state_kv[src];
|
||||
value *= reloaded_softmax[i0];
|
||||
product_row[i0] = value;
|
||||
}
|
||||
}
|
||||
|
||||
// Preserve the product dispatch's device store/reload boundary.
|
||||
threadgroup_barrier(mem_flags::mem_device);
|
||||
device volatile const float * reloaded_product =
|
||||
(device volatile const float *)product_row;
|
||||
|
||||
// Reproduce kernel_sum_rows_f32_f32(width=8, TG8) literally. MSL defines
|
||||
// simdgroup collectives over active lanes, so the branch recreates the
|
||||
// original eight-lane partial SIMD group inside this 32-thread group.
|
||||
sum_scratch[tid] = 0.0f;
|
||||
float row_sum = 0.0f;
|
||||
if (tid < 8u) {
|
||||
row_sum += reloaded_product[tid];
|
||||
row_sum = simd_sum(row_sum);
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
if (tid == 0u) {
|
||||
sum_scratch[0] = row_sum;
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
if (tid < 8u) {
|
||||
row_sum = sum_scratch[tid];
|
||||
row_sum = simd_sum(row_sum);
|
||||
if (tid == 0u) {
|
||||
dst[col] = row_sum;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Ratio-4 compression keeps two 4-row halves of recurrent state. After an
|
||||
// emitted compressed row, the second half becomes the next window's previous
|
||||
// half. The old encoder expressed this as four generic copies; this DS4-specific
|
||||
|
||||
@@ -70,7 +70,7 @@ struct ds4_metal_args_dsv4_indexed_attention {
|
||||
uint32_t window;
|
||||
uint32_t ratio;
|
||||
uint32_t comp_kv_f16;
|
||||
uint32_t pad0;
|
||||
uint32_t n_splits;
|
||||
uint64_t q_token_stride;
|
||||
uint64_t q_head_stride;
|
||||
uint64_t raw_row_stride;
|
||||
@@ -4915,6 +4915,270 @@ kernel void kernel_dsv4_router_finalize_weights_one_simd(
|
||||
}
|
||||
}
|
||||
|
||||
// M3 decode specialization that materializes the probability
|
||||
// transform in device memory before running the exact SIMD selection and
|
||||
// weight normalization above. The volatile reload after the device barrier
|
||||
// pins the same float store/load boundary as the standalone transform dispatch.
|
||||
kernel void kernel_dsv4_router_transform_finalize_weights_one_simd(
|
||||
constant ds4_metal_args_dsv4_router_select_one & args,
|
||||
device const float *logits,
|
||||
device float *probs,
|
||||
device const float *bias,
|
||||
device const int32_t *hash,
|
||||
device const int32_t *tokens,
|
||||
device int32_t *selected,
|
||||
device float *weights,
|
||||
threadgroup float *scratch [[threadgroup(0)]],
|
||||
uint tid [[thread_position_in_threadgroup]]) {
|
||||
if (tid >= 256 || args.hash_mode) return;
|
||||
|
||||
if (tid < 64) {
|
||||
device const float4 *s = (device const float4 *)logits;
|
||||
device float4 *d = (device float4 *)probs;
|
||||
const float4 x = s[tid];
|
||||
const float4 sp = select(log(1.0f + exp(x)), x, x > 20.0f);
|
||||
d[tid] = sqrt(sp);
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_device);
|
||||
device volatile const float *reloaded_probs =
|
||||
(device volatile const float *)probs;
|
||||
|
||||
(void)hash;
|
||||
(void)tokens;
|
||||
threadgroup float *score0_tg = scratch;
|
||||
threadgroup int32_t *idx0_tg =
|
||||
(threadgroup int32_t *)(scratch + 256);
|
||||
threadgroup float *score1_tg = scratch + 512;
|
||||
threadgroup int32_t *idx1_tg =
|
||||
(threadgroup int32_t *)(scratch + 768);
|
||||
const float p = reloaded_probs[tid];
|
||||
float score = args.has_bias ? p + bias[tid] : p;
|
||||
int32_t idx = (int32_t)tid;
|
||||
uint cross_stage = 0;
|
||||
|
||||
for (uint k = 2; k <= 256; k <<= 1) {
|
||||
for (uint j = k >> 1; j > 0; j >>= 1) {
|
||||
float peer_score;
|
||||
int32_t peer_idx;
|
||||
bool take_peer;
|
||||
const bool lower = (tid & j) == 0;
|
||||
const bool descending = (tid & k) == 0;
|
||||
|
||||
if (j < 32) {
|
||||
peer_score = simd_shuffle_xor(score, (ushort)j);
|
||||
peer_idx = simd_shuffle_xor(idx, (ushort)j);
|
||||
take_peer = descending
|
||||
? (lower ? score < peer_score : score > peer_score)
|
||||
: (lower ? score > peer_score : score < peer_score);
|
||||
if (take_peer) {
|
||||
score = peer_score;
|
||||
idx = peer_idx;
|
||||
}
|
||||
} else {
|
||||
threadgroup float *score_tg =
|
||||
(cross_stage & 1u) != 0u ? score1_tg : score0_tg;
|
||||
threadgroup int32_t *idx_tg =
|
||||
(cross_stage & 1u) != 0u ? idx1_tg : idx0_tg;
|
||||
score_tg[tid] = score;
|
||||
idx_tg[tid] = idx;
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
const uint other = tid ^ j;
|
||||
peer_score = score_tg[other];
|
||||
peer_idx = idx_tg[other];
|
||||
take_peer = descending
|
||||
? (lower ? score < peer_score : score > peer_score)
|
||||
: (lower ? score > peer_score : score < peer_score);
|
||||
if (take_peer) {
|
||||
score = peer_score;
|
||||
idx = peer_idx;
|
||||
}
|
||||
cross_stage++;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if (tid < 6) {
|
||||
selected[tid] = idx;
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_device);
|
||||
|
||||
threadgroup volatile float *norm_scratch =
|
||||
(threadgroup volatile float *)scratch;
|
||||
if (tid == 0) {
|
||||
device const int32_t *s = selected;
|
||||
norm_scratch[0] = 0.0f;
|
||||
for (uint i = 0; i < 6; i++) {
|
||||
norm_scratch[0] =
|
||||
norm_scratch[0] + reloaded_probs[s[i]];
|
||||
}
|
||||
norm_scratch[0] = max(norm_scratch[0], 6.103515625e-5f);
|
||||
norm_scratch[1] = 1.5f / norm_scratch[0];
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
if (tid < 6) {
|
||||
device const int32_t *s = selected;
|
||||
weights[tid] = reloaded_probs[s[tid]] * norm_scratch[1];
|
||||
}
|
||||
}
|
||||
|
||||
kernel void kernel_dsv4_router_project_select_fused(
|
||||
constant ds4_metal_args_mul_mv & args,
|
||||
constant ds4_metal_args_dsv4_router_select_one & select_args,
|
||||
device const char * src0_router,
|
||||
device const char * src1,
|
||||
device float * logits,
|
||||
device float * probs,
|
||||
device const float * bias,
|
||||
device int32_t * selected,
|
||||
device float * weights,
|
||||
device atomic_uint * completion,
|
||||
threadgroup char * shmem_raw [[threadgroup(0)]],
|
||||
uint3 tgpig [[threadgroup_position_in_grid]],
|
||||
uint3 tpitg [[thread_position_in_threadgroup]],
|
||||
ushort tiisg [[thread_index_in_simdgroup]],
|
||||
ushort sgitg [[simdgroup_index_in_threadgroup]]) {
|
||||
constexpr short NSG = 8;
|
||||
constexpr short NR0 = 2;
|
||||
constexpr short NB = 32;
|
||||
constexpr short NF = 16;
|
||||
constexpr short NF4 = NF/4;
|
||||
constexpr short NW = N_SIMDWIDTH;
|
||||
const uint tid = tpitg.x;
|
||||
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;
|
||||
}
|
||||
helper_mv_reduce_and_write<NR0>(logits, sumf, r0, args.ne01,
|
||||
tiisg, sgitg, shmem_raw);
|
||||
|
||||
threadgroup float *scratch = (threadgroup float *)shmem_raw;
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
atomic_thread_fence(mem_flags::mem_device,
|
||||
memory_order_seq_cst,
|
||||
thread_scope_device);
|
||||
if (tid == 0) {
|
||||
const uint old = atomic_fetch_add_explicit(
|
||||
completion, 1u, memory_order_relaxed);
|
||||
scratch[0] = old == 127u ? 1.0f : 0.0f;
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
if (scratch[0] == 0.0f) return;
|
||||
atomic_thread_fence(mem_flags::mem_device,
|
||||
memory_order_seq_cst,
|
||||
thread_scope_device);
|
||||
|
||||
if (tid < 64) {
|
||||
device volatile const float4 *s =
|
||||
(device volatile const float4 *)logits;
|
||||
device float4 *d = (device float4 *)probs;
|
||||
const float4 xv = s[tid];
|
||||
const float4 sp = select(log(1.0f + exp(xv)), xv, xv > 20.0f);
|
||||
d[tid] = sqrt(sp);
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_device);
|
||||
device volatile const float *reloaded_probs =
|
||||
(device volatile const float *)probs;
|
||||
|
||||
threadgroup float *score0_tg = scratch;
|
||||
threadgroup int32_t *idx0_tg =
|
||||
(threadgroup int32_t *)(scratch + 256);
|
||||
threadgroup float *score1_tg = scratch + 512;
|
||||
threadgroup int32_t *idx1_tg =
|
||||
(threadgroup int32_t *)(scratch + 768);
|
||||
const float p = reloaded_probs[tid];
|
||||
float score = select_args.has_bias ? p + bias[tid] : p;
|
||||
int32_t idx = (int32_t)tid;
|
||||
uint cross_stage = 0;
|
||||
for (uint k = 2; k <= 256; k <<= 1) {
|
||||
for (uint j = k >> 1; j > 0; j >>= 1) {
|
||||
float peer_score;
|
||||
int32_t peer_idx;
|
||||
bool take_peer;
|
||||
const bool lower = (tid & j) == 0;
|
||||
const bool descending = (tid & k) == 0;
|
||||
if (j < 32) {
|
||||
peer_score = simd_shuffle_xor(score, (ushort)j);
|
||||
peer_idx = simd_shuffle_xor(idx, (ushort)j);
|
||||
take_peer = descending
|
||||
? (lower ? score < peer_score : score > peer_score)
|
||||
: (lower ? score > peer_score : score < peer_score);
|
||||
if (take_peer) {
|
||||
score = peer_score;
|
||||
idx = peer_idx;
|
||||
}
|
||||
} else {
|
||||
threadgroup float *score_tg =
|
||||
(cross_stage & 1u) != 0u ? score1_tg : score0_tg;
|
||||
threadgroup int32_t *idx_tg =
|
||||
(cross_stage & 1u) != 0u ? idx1_tg : idx0_tg;
|
||||
score_tg[tid] = score;
|
||||
idx_tg[tid] = idx;
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
const uint other = tid ^ j;
|
||||
peer_score = score_tg[other];
|
||||
peer_idx = idx_tg[other];
|
||||
take_peer = descending
|
||||
? (lower ? score < peer_score : score > peer_score)
|
||||
: (lower ? score > peer_score : score < peer_score);
|
||||
if (take_peer) {
|
||||
score = peer_score;
|
||||
idx = peer_idx;
|
||||
}
|
||||
cross_stage++;
|
||||
}
|
||||
}
|
||||
}
|
||||
if (tid < 6) selected[tid] = idx;
|
||||
threadgroup_barrier(mem_flags::mem_device);
|
||||
threadgroup volatile float *norm_scratch =
|
||||
(threadgroup volatile float *)scratch;
|
||||
if (tid == 0) {
|
||||
norm_scratch[0] = 0.0f;
|
||||
for (uint i = 0; i < 6; ++i) {
|
||||
norm_scratch[0] = norm_scratch[0] + reloaded_probs[selected[i]];
|
||||
}
|
||||
norm_scratch[0] = max(norm_scratch[0], 6.103515625e-5f);
|
||||
norm_scratch[1] = 1.5f / norm_scratch[0];
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
if (tid < 6) {
|
||||
weights[tid] = reloaded_probs[selected[tid]] * norm_scratch[1];
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
atomic_thread_fence(mem_flags::mem_device,
|
||||
memory_order_seq_cst,
|
||||
thread_scope_device);
|
||||
if (tid == 0) {
|
||||
atomic_store_explicit(completion, 0u, memory_order_relaxed);
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
// Fills the dense compressed-attention mask with -inf. The selected top-k rows
|
||||
// are enabled by kernel_dsv4_topk_mask_scatter in a second ordered dispatch.
|
||||
kernel void kernel_dsv4_topk_mask(
|
||||
@@ -5319,6 +5583,117 @@ kernel void kernel_dsv4_indexed_mixed_attention_heads8(
|
||||
dst4[lane + 96] = o3 * inv_s;
|
||||
}
|
||||
|
||||
// Each simdgroup owns two heads and updates both from one staged K/V row.
|
||||
// This doubles row reuse without increasing the 256-thread workgroup.
|
||||
kernel void kernel_dsv4_indexed_mixed_attention_heads16_dual(
|
||||
constant ds4_metal_args_dsv4_indexed_attention &args,
|
||||
device const char *q,
|
||||
device const char *raw_kv,
|
||||
device const char *comp_kv,
|
||||
device const char *topk,
|
||||
device const char *sinks,
|
||||
device char *dst,
|
||||
threadgroup half4 *kv_shared [[threadgroup(0)]],
|
||||
uint2 tgpig [[threadgroup_position_in_grid]],
|
||||
ushort tid [[thread_index_in_threadgroup]],
|
||||
ushort lane [[thread_index_in_simdgroup]],
|
||||
ushort sg [[simdgroup_index_in_threadgroup]]) {
|
||||
const uint token = tgpig.x;
|
||||
const uint head0 = tgpig.y*16u + (uint)sg;
|
||||
const uint head1 = head0 + 8u;
|
||||
if (token >= args.n_tokens || head0 >= args.n_head) return;
|
||||
|
||||
device const float4 *qa = (device const float4 *)(q +
|
||||
(uint64_t)token*args.q_token_stride +
|
||||
(uint64_t)head0*args.q_head_stride);
|
||||
half4 qa0 = (half4)qa[lane + 0];
|
||||
half4 qa1 = (half4)qa[lane + 32];
|
||||
half4 qa2 = (half4)qa[lane + 64];
|
||||
half4 qa3 = (half4)qa[lane + 96];
|
||||
half4 qb0 = half4(0.0h), qb1 = half4(0.0h);
|
||||
half4 qb2 = half4(0.0h), qb3 = half4(0.0h);
|
||||
if (head1 < args.n_head) {
|
||||
device const float4 *qb = (device const float4 *)(q +
|
||||
(uint64_t)token*args.q_token_stride +
|
||||
(uint64_t)head1*args.q_head_stride);
|
||||
qb0 = (half4)qb[lane + 0];
|
||||
qb1 = (half4)qb[lane + 32];
|
||||
qb2 = (half4)qb[lane + 64];
|
||||
qb3 = (half4)qb[lane + 96];
|
||||
}
|
||||
|
||||
float Ma = -FLT_MAX/2.0f, Sa = 0.0f;
|
||||
float Mb = -FLT_MAX/2.0f, Sb = 0.0f;
|
||||
float4 ao0 = 0.0f, ao1 = 0.0f, ao2 = 0.0f, ao3 = 0.0f;
|
||||
float4 bo0 = 0.0f, bo1 = 0.0f, bo2 = 0.0f, bo3 = 0.0f;
|
||||
|
||||
const uint qpos = args.pos0 + token;
|
||||
const uint last_pos = args.pos0 + args.n_tokens - 1u;
|
||||
const uint first_raw_pos = last_pos + 1u - args.n_raw;
|
||||
const uint raw_last_pos = first_raw_pos + args.n_raw - 1u;
|
||||
const uint window_first = (args.window != 0u && qpos + 1u > args.window) ?
|
||||
qpos + 1u - args.window : 0u;
|
||||
const uint first = max(first_raw_pos, window_first);
|
||||
const uint last = min(qpos, raw_last_pos);
|
||||
if (first <= last) {
|
||||
for (uint pos = first; pos <= last; pos++) {
|
||||
const uint logical = pos - first_raw_pos;
|
||||
const uint row = (args.raw_start + logical)%args.raw_cap;
|
||||
device const float4 *src = (device const float4 *)(raw_kv +
|
||||
(uint64_t)row*args.raw_row_stride);
|
||||
if (tid < 128) kv_shared[tid] = (half4)src[tid];
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
dsv4_attend_shared_h4_row(kv_shared, qa0, qa1, qa2, qa3,
|
||||
args.scale, lane, Ma, Sa, ao0, ao1, ao2, ao3);
|
||||
if (head1 < args.n_head) {
|
||||
dsv4_attend_shared_h4_row(kv_shared, qb0, qb1, qb2, qb3,
|
||||
args.scale, lane, Mb, Sb, bo0, bo1, bo2, bo3);
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
}
|
||||
}
|
||||
|
||||
const uint visible = min((qpos + 1u)/args.ratio, args.n_comp);
|
||||
device const int32_t *row_topk = (device const int32_t *)(topk +
|
||||
(uint64_t)token*args.topk_token_stride);
|
||||
for (uint i = 0; i < args.top_k; i++) {
|
||||
const int32_t idx = row_topk[i];
|
||||
if (idx < 0) continue;
|
||||
if ((uint)idx >= visible) break;
|
||||
if (tid < 128) {
|
||||
kv_shared[tid] = dsv4_load_cache_h4(comp_kv,
|
||||
args.comp_row_stride, (uint)idx, tid, args.comp_kv_f16 != 0u);
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
dsv4_attend_shared_h4_row(kv_shared, qa0, qa1, qa2, qa3,
|
||||
args.scale, lane, Ma, Sa, ao0, ao1, ao2, ao3);
|
||||
if (head1 < args.n_head) {
|
||||
dsv4_attend_shared_h4_row(kv_shared, qb0, qb1, qb2, qb3,
|
||||
args.scale, lane, Mb, Sb, bo0, bo1, bo2, bo3);
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
}
|
||||
|
||||
dsv4_attend_sink(((device const float *)sinks)[head0],
|
||||
Ma, Sa, ao0, ao1, ao2, ao3);
|
||||
const float ia = Sa == 0.0f ? 0.0f : 1.0f/Sa;
|
||||
device float4 *da = (device float4 *)(dst +
|
||||
(uint64_t)token*args.dst_token_stride +
|
||||
(uint64_t)head0*args.dst_head_stride);
|
||||
da[lane + 0] = ao0*ia; da[lane + 32] = ao1*ia;
|
||||
da[lane + 64] = ao2*ia; da[lane + 96] = ao3*ia;
|
||||
if (head1 < args.n_head) {
|
||||
dsv4_attend_sink(((device const float *)sinks)[head1],
|
||||
Mb, Sb, bo0, bo1, bo2, bo3);
|
||||
const float ib = Sb == 0.0f ? 0.0f : 1.0f/Sb;
|
||||
device float4 *db = (device float4 *)(dst +
|
||||
(uint64_t)token*args.dst_token_stride +
|
||||
(uint64_t)head1*args.dst_head_stride);
|
||||
db[lane + 0] = bo0*ib; db[lane + 32] = bo1*ib;
|
||||
db[lane + 64] = bo2*ib; db[lane + 96] = bo3*ib;
|
||||
}
|
||||
}
|
||||
|
||||
// Decode specialization of kernel_dsv4_indexed_mixed_attention_heads8.
|
||||
// Generation attends one token at a time, so the ratio-4 indexed path spends a
|
||||
// visible amount of time repeatedly staging the same K/V row for the eight
|
||||
@@ -5450,6 +5825,190 @@ kernel void kernel_dsv4_indexed_mixed_attention_heads8_rb16(
|
||||
dst4[lane + 96] = o3 * inv_s;
|
||||
}
|
||||
|
||||
// Long-context decode specialization of the indexed mixed-attention path.
|
||||
//
|
||||
// The ordinary heads8 kernel reuses each K/V row across eight heads, but only
|
||||
// launches one threadgroup per head group. Long-context decode therefore has
|
||||
// too little parallel work while each group scans its raw and selected rows.
|
||||
// This kernel retains the same eight-head reuse while splitting that row
|
||||
// sequence across args.n_splits workgroups. A second kernel merges the online
|
||||
// softmax partials and applies the attention sink.
|
||||
kernel void kernel_dsv4_indexed_mixed_attention_heads8_split(
|
||||
constant ds4_metal_args_dsv4_indexed_attention & args,
|
||||
device const char *q,
|
||||
device const char *raw_kv,
|
||||
device const char *comp_kv,
|
||||
device const char *topk,
|
||||
device char *tmp,
|
||||
threadgroup half4 *kv_shared [[threadgroup(0)]],
|
||||
uint3 tgpig [[threadgroup_position_in_grid]],
|
||||
ushort tid [[thread_index_in_threadgroup]],
|
||||
ushort lane [[thread_index_in_simdgroup]],
|
||||
ushort sg [[simdgroup_index_in_threadgroup]]) {
|
||||
constexpr uint rows_per_block = 16u;
|
||||
constexpr uint vecs_per_row = 128u;
|
||||
|
||||
const uint token = tgpig.x;
|
||||
const uint head = tgpig.y * 8u + (uint)sg;
|
||||
const uint split = tgpig.z;
|
||||
const uint n_splits = args.n_splits;
|
||||
if (token >= args.n_tokens || head >= args.n_head ||
|
||||
n_splits < 2u || n_splits > 31u || split >= n_splits) {
|
||||
return;
|
||||
}
|
||||
|
||||
device const float4 *q4 = (device const float4 *)(q +
|
||||
(uint64_t)token * args.q_token_stride +
|
||||
(uint64_t)head * args.q_head_stride);
|
||||
const half4 q0 = (half4)q4[lane + 0];
|
||||
const half4 q1 = (half4)q4[lane + 32];
|
||||
const half4 q2 = (half4)q4[lane + 64];
|
||||
const half4 q3 = (half4)q4[lane + 96];
|
||||
|
||||
float M = -FLT_MAX/2.0f;
|
||||
float S = 0.0f;
|
||||
float4 o0 = 0.0f;
|
||||
float4 o1 = 0.0f;
|
||||
float4 o2 = 0.0f;
|
||||
float4 o3 = 0.0f;
|
||||
|
||||
const uint qpos = args.pos0 + token;
|
||||
const uint last_pos = args.pos0 + args.n_tokens - 1u;
|
||||
const uint first_raw_pos = last_pos + 1u - args.n_raw;
|
||||
const uint raw_last_pos = first_raw_pos + args.n_raw - 1u;
|
||||
const uint window_first = (args.window != 0u && qpos + 1u > args.window) ?
|
||||
qpos + 1u - args.window : 0u;
|
||||
const uint raw_first = max(first_raw_pos, window_first);
|
||||
const uint raw_last = min(qpos, raw_last_pos);
|
||||
const uint raw_count = raw_first <= raw_last ?
|
||||
raw_last - raw_first + 1u : 0u;
|
||||
const uint total_rows = raw_count + args.top_k;
|
||||
const uint rows_per_split =
|
||||
(total_rows + n_splits - 1u) / n_splits;
|
||||
const uint split_first = min(split * rows_per_split, total_rows);
|
||||
const uint split_last = min(split_first + rows_per_split, total_rows);
|
||||
const uint visible = min((qpos + 1u) / args.ratio, args.n_comp);
|
||||
device const int32_t *row_topk = (device const int32_t *)(topk +
|
||||
(uint64_t)token * args.topk_token_stride);
|
||||
|
||||
for (uint seq0 = split_first; seq0 < split_last;
|
||||
seq0 += rows_per_block) {
|
||||
const uint n_rows = min(rows_per_block, split_last - seq0);
|
||||
for (uint off = (uint)tid;
|
||||
off < n_rows * vecs_per_row;
|
||||
off += 256u) {
|
||||
const uint r = off / vecs_per_row;
|
||||
const uint c = off - r * vecs_per_row;
|
||||
const uint seq = seq0 + r;
|
||||
half4 value = half4(0.0h);
|
||||
if (seq < raw_count) {
|
||||
const uint pos = raw_first + seq;
|
||||
const uint logical = pos - first_raw_pos;
|
||||
const uint row = (args.raw_start + logical) % args.raw_cap;
|
||||
device const float4 *src = (device const float4 *)(raw_kv +
|
||||
(uint64_t)row * args.raw_row_stride);
|
||||
value = (half4)src[c];
|
||||
} else {
|
||||
const int32_t idx = row_topk[seq - raw_count];
|
||||
if (idx >= 0 && (uint)idx < visible) {
|
||||
value = dsv4_load_cache_h4(comp_kv,
|
||||
args.comp_row_stride,
|
||||
(uint)idx,
|
||||
c,
|
||||
args.comp_kv_f16 != 0u);
|
||||
}
|
||||
}
|
||||
kv_shared[off] = value;
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
for (uint r = 0; r < n_rows; r++) {
|
||||
const uint seq = seq0 + r;
|
||||
bool valid = true;
|
||||
if (seq >= raw_count) {
|
||||
const int32_t idx = row_topk[seq - raw_count];
|
||||
valid = idx >= 0 && (uint)idx < visible;
|
||||
}
|
||||
if (valid) {
|
||||
dsv4_attend_shared_h4_row_at(kv_shared,
|
||||
r,
|
||||
q0, q1, q2, q3,
|
||||
args.scale,
|
||||
lane,
|
||||
M, S,
|
||||
o0, o1, o2, o3);
|
||||
}
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
}
|
||||
|
||||
const uint64_t n_rows = (uint64_t)args.n_tokens * args.n_head;
|
||||
const uint64_t row = (uint64_t)token * args.n_head + head;
|
||||
device float4 *partials = (device float4 *)tmp;
|
||||
partials[(row * vecs_per_row + lane + 0u) * n_splits + split] = o0;
|
||||
partials[(row * vecs_per_row + lane + 32u) * n_splits + split] = o1;
|
||||
partials[(row * vecs_per_row + lane + 64u) * n_splits + split] = o2;
|
||||
partials[(row * vecs_per_row + lane + 96u) * n_splits + split] = o3;
|
||||
|
||||
if (lane == 0u) {
|
||||
device float *stats = (device float *)(partials +
|
||||
n_rows * vecs_per_row * n_splits);
|
||||
const uint64_t stat = (row * n_splits + split) * 2u;
|
||||
stats[stat + 0u] = S;
|
||||
stats[stat + 1u] = M;
|
||||
}
|
||||
}
|
||||
|
||||
kernel void kernel_dsv4_indexed_mixed_attention_heads8_split_reduce(
|
||||
constant ds4_metal_args_dsv4_indexed_attention & args,
|
||||
device const char *tmp,
|
||||
device const char *sinks,
|
||||
device char *dst,
|
||||
uint tgpig [[threadgroup_position_in_grid]],
|
||||
ushort lane [[thread_index_in_simdgroup]],
|
||||
ushort sg [[simdgroup_index_in_threadgroup]]) {
|
||||
constexpr uint vecs_per_row = 128u;
|
||||
const uint n_splits = args.n_splits;
|
||||
const uint64_t n_rows = (uint64_t)args.n_tokens * args.n_head;
|
||||
const uint64_t row = tgpig;
|
||||
if (row >= n_rows || n_splits < 2u || n_splits > 31u) {
|
||||
return;
|
||||
}
|
||||
|
||||
device const float4 *partials = (device const float4 *)tmp;
|
||||
device const float *stats = (device const float *)(partials +
|
||||
n_rows * vecs_per_row * n_splits);
|
||||
float part_sum = 0.0f;
|
||||
float part_max = -FLT_MAX/2.0f;
|
||||
if ((uint)lane < n_splits) {
|
||||
const uint64_t stat = (row * n_splits + (uint)lane) * 2u;
|
||||
part_sum = stats[stat + 0u];
|
||||
part_max = stats[stat + 1u];
|
||||
} else if ((uint)lane == n_splits) {
|
||||
const uint head = (uint)(row % args.n_head);
|
||||
part_sum = 1.0f;
|
||||
part_max = ((device const float *)sinks)[head];
|
||||
}
|
||||
|
||||
const float global_max = simd_max(part_max);
|
||||
const float part_scale = part_sum > 0.0f ?
|
||||
exp(part_max - global_max) : 0.0f;
|
||||
const float total_sum = simd_sum(part_sum * part_scale);
|
||||
const float inv_sum = total_sum > 0.0f ? 1.0f / total_sum : 0.0f;
|
||||
|
||||
device float4 *out = (device float4 *)dst + row * vecs_per_row;
|
||||
for (uint i = (uint)sg; i < vecs_per_row; i += 4u) {
|
||||
float4 value = float4(0.0f);
|
||||
if ((uint)lane < n_splits) {
|
||||
value = partials[(row * vecs_per_row + i) * n_splits +
|
||||
(uint)lane] * part_scale;
|
||||
}
|
||||
value = simd_sum(value);
|
||||
if (lane == 0u) {
|
||||
out[i] = value * inv_sum;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
static inline float dsv4_indexer_dot128_shared_q(
|
||||
float4 c0,
|
||||
float4 c1,
|
||||
|
||||
@@ -41,6 +41,23 @@ struct ds4_metal_args_dsv4_rope_affine_pair {
|
||||
float beta_slow;
|
||||
};
|
||||
|
||||
struct ds4_metal_args_dsv4_head_norm_rope {
|
||||
int32_t n_head;
|
||||
int32_t head_dim;
|
||||
int32_t head_dim4;
|
||||
int32_t n_dims;
|
||||
int32_t n_ctx_orig;
|
||||
int32_t pos0;
|
||||
int32_t inverse;
|
||||
float eps;
|
||||
float freq_base;
|
||||
float freq_scale;
|
||||
float ext_factor;
|
||||
float attn_factor;
|
||||
float beta_fast;
|
||||
float beta_slow;
|
||||
};
|
||||
|
||||
static float rope_yarn_ramp(const float low, const float high, const int i0) {
|
||||
const float y = (i0 / 2 - low) / max(0.001f, high - low);
|
||||
return 1.0f - min(1.0f, max(0.0f, y));
|
||||
@@ -327,36 +344,109 @@ kernel void kernel_dsv4_rope_tail_f32_inplace_pair_shared4(
|
||||
*((device float *) (dst_base + j1*args.nb0)) = x0*sin_theta + x1*cos_theta;
|
||||
}
|
||||
|
||||
// DS4 positions are always affine within one RoPE dispatch. This variant
|
||||
// reconstructs the same wrapped int32 position in-kernel, avoiding the host
|
||||
// position array and its buffer binding while preserving the pair lane mapping
|
||||
// and all floating-point operations of the specialization above.
|
||||
kernel void kernel_dsv4_rope_tail_f32_inplace_pair_affine(
|
||||
constant ds4_metal_args_dsv4_rope_affine_pair & args [[buffer(0)]],
|
||||
device const char * src0 [[buffer(1)]],
|
||||
device char * dst [[buffer(4)]],
|
||||
uint tid [[thread_index_in_threadgroup]],
|
||||
ushort3 ntg [[threads_per_threadgroup]],
|
||||
uint3 tgpig [[threadgroup_position_in_grid]]) {
|
||||
const int i1 = tgpig[0];
|
||||
const int i2 = tgpig[1];
|
||||
// Fuses the per-head RMSNorm and partial Q RoPE while retaining the standalone
|
||||
// norm reduction tree and the mode-0 RoPE lane mapping.
|
||||
kernel void kernel_dsv4_head_rms_norm_rope_tail_f32(
|
||||
constant ds4_metal_args_dsv4_head_norm_rope & args,
|
||||
device char * xraw,
|
||||
threadgroup float * shmem_f32 [[threadgroup(0)]],
|
||||
uint3 tgpig [[threadgroup_position_in_grid]],
|
||||
ushort3 tpitg [[thread_position_in_threadgroup]],
|
||||
ushort sgitg [[simdgroup_index_in_threadgroup]],
|
||||
ushort tiisg [[thread_index_in_simdgroup]],
|
||||
ushort3 ntg [[threads_per_threadgroup]]) {
|
||||
if (sgitg == 0) {
|
||||
shmem_f32[tiisg] = 0.0f;
|
||||
}
|
||||
|
||||
const uint head = tgpig.x;
|
||||
const uint tok = tgpig.y;
|
||||
device float4 * x4 = (device float4 *)xraw +
|
||||
((uint64_t)tok * (uint64_t)args.n_head + head) *
|
||||
(uint64_t)args.head_dim4;
|
||||
|
||||
float sumf = 0.0f;
|
||||
for (int i00 = tpitg.x; i00 < args.head_dim4; i00 += ntg.x) {
|
||||
sumf += dot(x4[i00], x4[i00]);
|
||||
}
|
||||
sumf = simd_sum(sumf);
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
if (tiisg == 0) {
|
||||
shmem_f32[sgitg] = sumf;
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
sumf = simd_sum(shmem_f32[tiisg]);
|
||||
const float scale = 1.0f / sqrt(sumf / args.head_dim + args.eps);
|
||||
const int n_nope = args.head_dim - args.n_dims;
|
||||
if (n_nope < 0) {
|
||||
return;
|
||||
}
|
||||
|
||||
float corr_dims[2];
|
||||
rope_yarn_corr_dims(args.n_dims, args.n_ctx_orig, args.freq_base,
|
||||
args.beta_fast, args.beta_slow, corr_dims);
|
||||
const float theta_base = (float)(args.pos0 + (int)tok);
|
||||
const float inv_ndims = -1.0f / args.n_dims;
|
||||
device float * xs = (device float *)x4;
|
||||
|
||||
for (int i0 = tpitg.x; i0 < args.head_dim; i0 += ntg.x) {
|
||||
if (i0 < n_nope) {
|
||||
xs[i0] = xs[i0] * scale;
|
||||
continue;
|
||||
}
|
||||
const int r = i0 - n_nope;
|
||||
if ((r & 1) != 0) {
|
||||
continue;
|
||||
}
|
||||
#ifdef DS4_METAL_ROPE_EXP2_LOG2
|
||||
const float theta =
|
||||
theta_base * exp2(inv_ndims * (float)r * log2(args.freq_base));
|
||||
#else
|
||||
const float theta =
|
||||
theta_base * pow(args.freq_base, inv_ndims * r);
|
||||
#endif
|
||||
float cos_theta;
|
||||
float sin_theta;
|
||||
rope_yarn(theta, args.freq_scale, corr_dims, r,
|
||||
args.ext_factor, args.attn_factor,
|
||||
&cos_theta, &sin_theta);
|
||||
if (args.inverse) {
|
||||
sin_theta = -sin_theta;
|
||||
}
|
||||
|
||||
const float x0 = xs[i0] * scale;
|
||||
const float x1 = xs[i0 + 1] * scale;
|
||||
xs[i0] = x0 * cos_theta - x1 * sin_theta;
|
||||
xs[i0 + 1] = x0 * sin_theta + x1 * cos_theta;
|
||||
}
|
||||
}
|
||||
|
||||
// DS4 positions are always affine within one RoPE dispatch. This variant
|
||||
// reconstructs the same wrapped int32 position in-kernel, avoiding the host
|
||||
// position array and its buffer binding while preserving the pair lane mapping
|
||||
// and all floating-point operations of the specialization above.
|
||||
|
||||
/* Shared, deliberately noinline so that every caller gets bit-identical
|
||||
* trigonometric codegen. The header note about tiny trig codegen changes
|
||||
* flipping sampled tokens is exactly why this body must be compiled once and
|
||||
* shared rather than inlined separately into each kernel. */
|
||||
static __attribute__((noinline)) void ds4_rope_tail_pair_affine_row(
|
||||
constant ds4_metal_args_dsv4_rope_affine_pair & args,
|
||||
device const char * src_base,
|
||||
device char * dst_base,
|
||||
int n_nope,
|
||||
uint raw_pos,
|
||||
uint tid,
|
||||
uint nthreads) {
|
||||
float corr_dims[2];
|
||||
rope_yarn_corr_dims(args.n_dims, args.n_ctx_orig, args.freq_base, args.beta_fast, args.beta_slow, corr_dims);
|
||||
|
||||
const uint raw_pos = args.pos0 + (uint)i2 * args.pos_step;
|
||||
const float theta_base = (float)as_type<int>(raw_pos);
|
||||
const float theta_base = (float)as_type<int>(raw_pos);
|
||||
const float inv_ndims = -1.f/args.n_dims;
|
||||
device const char * src_base =
|
||||
src0 + (uint64_t)i2*args.token_bytes + (uint64_t)i1*args.row_bytes;
|
||||
device char * dst_base =
|
||||
dst + (uint64_t)i2*args.token_bytes + (uint64_t)i1*args.row_bytes;
|
||||
|
||||
for (int r = tid; r < args.n_dims; r += ntg.x) {
|
||||
for (int r = tid; r < args.n_dims; r += nthreads) {
|
||||
if ((r & 1) != 0) {
|
||||
continue;
|
||||
}
|
||||
@@ -381,5 +471,413 @@ kernel void kernel_dsv4_rope_tail_f32_inplace_pair_affine(
|
||||
|
||||
*((device float *) (dst_base + j0*sizeof(float))) = x0*cos_theta - x1*sin_theta;
|
||||
*((device float *) (dst_base + j1*sizeof(float))) = x0*sin_theta + x1*cos_theta;
|
||||
}}
|
||||
|
||||
kernel void kernel_dsv4_rope_tail_f32_inplace_pair_affine(
|
||||
constant ds4_metal_args_dsv4_rope_affine_pair & args [[buffer(0)]],
|
||||
device const char * src0 [[buffer(1)]],
|
||||
device char * dst [[buffer(4)]],
|
||||
uint tid [[thread_index_in_threadgroup]],
|
||||
ushort3 ntg [[threads_per_threadgroup]],
|
||||
uint3 tgpig [[threadgroup_position_in_grid]]) {
|
||||
const int i1 = tgpig[0];
|
||||
const int i2 = tgpig[1];
|
||||
const int n_nope = args.head_dim - args.n_dims;
|
||||
if (n_nope < 0) {
|
||||
return;
|
||||
}
|
||||
const uint raw_pos = args.pos0 + (uint)i2 * args.pos_step;
|
||||
device const char * src_base =
|
||||
src0 + (uint64_t)i2*args.token_bytes + (uint64_t)i1*args.row_bytes;
|
||||
device char * dst_base =
|
||||
dst + (uint64_t)i2*args.token_bytes + (uint64_t)i1*args.row_bytes;
|
||||
ds4_rope_tail_pair_affine_row(args, src_base, dst_base, n_nope, raw_pos, tid, ntg.x);
|
||||
|
||||
}
|
||||
|
||||
// Decode-only fusion of the KV RoPE tail with the FP8/raw finalizer. Both were
|
||||
// already single 64-thread threadgroups on the same row, back to back, so the
|
||||
// pair cost two dispatches (~12.4 us) to touch 2 KB. The RoPE body below is a
|
||||
// verbatim copy of kernel_dsv4_rope_tail_f32_inplace_pair_affine specialised to
|
||||
// the decode grid (one head, one token, so i1 = i2 = 0) and the finalizer body
|
||||
// is a verbatim copy of kernel_dsv4_kv_fp8_store_f32. The barrier between them
|
||||
// is required because RoPE writes element pairs across lanes while the raw copy
|
||||
// reads them per lane. Arithmetic, order and rounding are unchanged; the header
|
||||
// warning above about trigonometric codegen still applies, so this kernel is
|
||||
// gated and verified against full-vocabulary logits before promotion.
|
||||
kernel void kernel_dsv4_kv_rope_fp8_store_f32(
|
||||
constant ds4_metal_args_dsv4_kv_fp8_store & args,
|
||||
constant ds4_metal_args_dsv4_rope_affine_pair & rope,
|
||||
device float * kv,
|
||||
device float * raw_cache,
|
||||
threadgroup float * scratch [[threadgroup(0)]],
|
||||
uint tid [[thread_index_in_threadgroup]]) {
|
||||
{
|
||||
const int rope_n_nope = rope.head_dim - rope.n_dims;
|
||||
if (rope_n_nope < 0) {
|
||||
return;
|
||||
}
|
||||
ds4_rope_tail_pair_affine_row(rope,
|
||||
(device const char *)kv,
|
||||
(device char *)kv,
|
||||
rope_n_nope,
|
||||
rope.pos0,
|
||||
tid,
|
||||
64u);
|
||||
}
|
||||
/* The RoPE helper writes device-memory pairs that different lanes read
|
||||
* below. A threadgroup-only fence does not make those cross-lane device
|
||||
* writes visible. */
|
||||
threadgroup_barrier(mem_flags::mem_device_and_threadgroup);
|
||||
{
|
||||
|
||||
const int head_dim = args.head_dim;
|
||||
const int n_rot = args.n_rot;
|
||||
const int n_nope = head_dim - n_rot;
|
||||
if (head_dim <= 0 || n_rot < 0 || n_nope < 0 || tid >= 64) {
|
||||
return;
|
||||
}
|
||||
|
||||
device float * raw = raw_cache + (int64_t)args.raw_row * head_dim;
|
||||
|
||||
for (int off = 0; off < n_nope; off += 64) {
|
||||
float v = 0.0f;
|
||||
if (off + (int)tid < n_nope) {
|
||||
v = kv[off + tid];
|
||||
scratch[tid] = abs(v);
|
||||
} else {
|
||||
scratch[tid] = 0.0f;
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
for (uint stride = 32; stride > 0; stride >>= 1) {
|
||||
if (tid < stride) {
|
||||
scratch[tid] = max(scratch[tid], scratch[tid + stride]);
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
}
|
||||
|
||||
const float amax = max(scratch[0], 1.0e-4f);
|
||||
const float fp8_scale = exp2(ceil(log2(amax / 448.0f)));
|
||||
if (off + (int)tid < n_nope) {
|
||||
const float q = dsv4_e4m3fn_dequant(clamp(v / fp8_scale, -448.0f, 448.0f)) * fp8_scale;
|
||||
kv[off + tid] = q;
|
||||
// Diagnostic only: skip the FP16 round-trip that normally matches the
|
||||
// half-typed FlashAttention KV buffer's precision. With this enabled the
|
||||
// indexer will see higher-precision raw values than FlashAttention does,
|
||||
// which is informative but not a production-ready setting.
|
||||
#ifdef DS4_METAL_KV_RAW_F32
|
||||
raw[off + tid] = q;
|
||||
#else
|
||||
raw[off + tid] = (float)((half)q);
|
||||
#endif
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
}
|
||||
|
||||
for (int i = n_nope + tid; i < head_dim; i += 64) {
|
||||
#ifdef DS4_METAL_KV_RAW_F32
|
||||
raw[i] = kv[i];
|
||||
#else
|
||||
raw[i] = (float)((half)kv[i]);
|
||||
#endif
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/* Decode-only sibling of kernel_flash_attn_ext_vec_reduce that also applies the
|
||||
* inverse RoPE tail to the row it just produced, removing a whole dispatch per
|
||||
* layer. Each threadgroup owns one head's entire 512-float row, so the RoPE is
|
||||
* an intra-threadgroup dependency: reduce, barrier, rotate. Both halves call the
|
||||
* same shared noinline helpers the standalone kernels use, so the arithmetic and
|
||||
* its codegen are identical to running the two dispatches back to back. */
|
||||
kernel void kernel_flash_attn_ext_vec_reduce_rope(
|
||||
constant ds4_metal_args_flash_attn_ext_vec_reduce & args,
|
||||
device const char * htmp,
|
||||
device char * dst,
|
||||
constant ds4_metal_args_dsv4_rope_affine_pair & rope,
|
||||
uint tgpig[[threadgroup_position_in_grid]],
|
||||
ushort tiitg[[thread_index_in_threadgroup]],
|
||||
ushort tiisg[[thread_index_in_simdgroup]],
|
||||
ushort sgitg[[simdgroup_index_in_threadgroup]]) {
|
||||
ds4_flash_attn_vec_reduce_row(args, htmp, dst, tgpig, tiisg, sgitg,
|
||||
(short)FC_flash_attn_ext_vec_reduce_NWG,
|
||||
(short)FC_flash_attn_ext_vec_reduce_DV);
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_device);
|
||||
|
||||
const int n_nope = rope.head_dim - rope.n_dims;
|
||||
if (n_nope < 0) {
|
||||
return;
|
||||
}
|
||||
device char * row = dst + (uint64_t)tgpig * rope.row_bytes;
|
||||
ds4_rope_tail_pair_affine_row(rope,
|
||||
(device const char *)row,
|
||||
row,
|
||||
n_nope,
|
||||
rope.pos0,
|
||||
tiitg,
|
||||
(uint)(32 * FC_flash_attn_ext_vec_reduce_NWG));
|
||||
}
|
||||
|
||||
struct ds4_metal_args_dsv4_comp_finalize {
|
||||
ds4_metal_args_dsv4_rope_affine_pair rope;
|
||||
float rms_eps;
|
||||
uint32_t pad0;
|
||||
};
|
||||
|
||||
/* Decode-only emit-path fusion. Every ratio-th token, each layer finalizes
|
||||
* one freshly pooled compressor row per compressor: RMS norm, RoPE tail, and
|
||||
* then the FP8 round-trip + F16 commit copy (attention, 512 floats) or the
|
||||
* Hadamard+FP4 QAT (indexer, 128 floats). Those were seven single-row
|
||||
* dispatches; this kernel is one dispatch with two threadgroups.
|
||||
*
|
||||
* Each phase reproduces its standalone kernel bit-exactly:
|
||||
* - norm: kernel_rms_norm_mul_f32_4's tree (float4 lanes, simd_sum, zero-
|
||||
* padded 32-slot cross-simdgroup reduce); 512 uses 128 virtual threads on
|
||||
* simdgroups 0-3, 128 uses 32 virtual threads on simdgroup 0.
|
||||
* - rope: ds4_rope_tail_pair_affine_row verbatim (lanes 0-63, nthreads=64).
|
||||
* - fp8: kernel_dsv4_fp8_kv_quantize_f32's 64-lane shmem max tree and
|
||||
* round-trip, src==dst so the verbatim tail copy is a no-op and dropped.
|
||||
* - commit: per-element f32->f16 conversion (value-wise exact).
|
||||
* - qat: kernel_dsv4_indexer_hadamard_fp4_f32's butterfly and per-32 amax
|
||||
* tree on lanes 0-127.
|
||||
* Threads outside a phase's virtual width still execute every barrier, so
|
||||
* threadgroup barriers stay uniform across the 256-thread threadgroup. */
|
||||
kernel void kernel_dsv4_comp_row_finalize_f32(
|
||||
constant ds4_metal_args_dsv4_comp_finalize & args [[buffer(0)]],
|
||||
device float * attn_row [[buffer(1)]],
|
||||
device const float * attn_norm_w [[buffer(2)]],
|
||||
device char * attn_cache [[buffer(3)]],
|
||||
device float * index_row [[buffer(4)]],
|
||||
device const float * index_norm_w [[buffer(5)]],
|
||||
device float * attn_state_kv [[buffer(6)]],
|
||||
device float * attn_state_score [[buffer(7)]],
|
||||
device float * index_state_kv [[buffer(8)]],
|
||||
device float * index_state_score [[buffer(9)]],
|
||||
threadgroup float * shmem [[threadgroup(0)]],
|
||||
uint tgpig [[threadgroup_position_in_grid]],
|
||||
ushort tiitg [[thread_index_in_threadgroup]],
|
||||
ushort tiisg [[thread_index_in_simdgroup]],
|
||||
ushort sgitg [[simdgroup_index_in_threadgroup]]) {
|
||||
constant ds4_metal_args_dsv4_rope_affine_pair & rope_args = args.rope;
|
||||
|
||||
if (tgpig == 0) {
|
||||
/* -------- attention compressor row (512 floats) -------- */
|
||||
{
|
||||
device float4 * y4 = (device float4 *)attn_row;
|
||||
device const float4 * x4 = (device const float4 *)attn_row;
|
||||
device const float4 * w4 = (device const float4 *)attn_norm_w;
|
||||
if (sgitg == 0) {
|
||||
shmem[tiisg] = 0.0f;
|
||||
}
|
||||
float sumf = 0.0f;
|
||||
if (tiitg < 128) {
|
||||
sumf = dot(x4[tiitg], x4[tiitg]);
|
||||
}
|
||||
sumf = simd_sum(sumf);
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
if (tiitg < 128 && tiisg == 0) {
|
||||
shmem[sgitg] = sumf;
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
float total = 0.0f;
|
||||
if (tiitg < 128) {
|
||||
total = simd_sum(shmem[tiisg]);
|
||||
}
|
||||
const float mean = total / 512.0f;
|
||||
const float scale = 1.0f/sqrt(mean + args.rms_eps);
|
||||
if (tiitg < 128) {
|
||||
y4[tiitg] = (x4[tiitg]*scale)*w4[tiitg];
|
||||
}
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_device);
|
||||
ds4_rope_tail_pair_affine_row(rope_args,
|
||||
(device const char *)attn_row,
|
||||
(device char *)attn_row,
|
||||
512 - rope_args.n_dims,
|
||||
rope_args.pos0,
|
||||
tiitg,
|
||||
64u);
|
||||
threadgroup_barrier(mem_flags::mem_device);
|
||||
for (int off = 0; off < 512 - rope_args.n_dims; off += 64) {
|
||||
float v = 0.0f;
|
||||
if (tiitg < 64) {
|
||||
v = attn_row[off + tiitg];
|
||||
shmem[tiitg] = abs(v);
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
for (uint stride = 32; stride > 0; stride >>= 1) {
|
||||
if (tiitg < stride) {
|
||||
shmem[tiitg] = max(shmem[tiitg], shmem[tiitg + stride]);
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
}
|
||||
const float amax = max(shmem[0], 1.0e-4f);
|
||||
const float scale = exp2(ceil(log2(amax / 448.0f)));
|
||||
if (tiitg < 64) {
|
||||
const float q = dsv4_e4m3fn_dequant(clamp(v / scale, -448.0f, 448.0f)) * scale;
|
||||
attn_row[off + tiitg] = q;
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_device);
|
||||
if (tiitg < 128) {
|
||||
device const float4 * x4 = (device const float4 *)attn_row;
|
||||
device half4 * o4 = (device half4 *)attn_cache;
|
||||
const float4 v = x4[tiitg];
|
||||
o4[tiitg] = half4(v);
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
if (tgpig >= 2u) {
|
||||
/* Ratio-4 state shifts for both compressors (elementwise row move,
|
||||
* so the flat gid mapping is bit-exact): 4*1024 attention elements
|
||||
* then 4*256 indexer elements. */
|
||||
const uint gid = (tgpig - 2u) * 256u + tiitg;
|
||||
const uint n0 = 4u * 1024u;
|
||||
if (gid < n0) {
|
||||
attn_state_kv[gid] = attn_state_kv[n0 + gid];
|
||||
attn_state_score[gid] = attn_state_score[n0 + gid];
|
||||
return;
|
||||
}
|
||||
const uint gid1 = gid - n0;
|
||||
const uint n1 = 4u * 256u;
|
||||
if (gid1 >= n1) return;
|
||||
index_state_kv[gid1] = index_state_kv[n1 + gid1];
|
||||
index_state_score[gid1] = index_state_score[n1 + gid1];
|
||||
return;
|
||||
}
|
||||
|
||||
/* -------- indexer compressor row (128 floats) -------- */
|
||||
{
|
||||
device float4 * y4 = (device float4 *)index_row;
|
||||
device const float4 * x4 = (device const float4 *)index_row;
|
||||
device const float4 * w4 = (device const float4 *)index_norm_w;
|
||||
if (sgitg == 0) {
|
||||
shmem[tiisg] = 0.0f;
|
||||
}
|
||||
float sumf = 0.0f;
|
||||
if (tiitg < 32) {
|
||||
sumf = dot(x4[tiitg], x4[tiitg]);
|
||||
}
|
||||
sumf = simd_sum(sumf);
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
if (tiisg == 0) {
|
||||
shmem[sgitg] = sumf;
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
float total = 0.0f;
|
||||
if (tiitg < 32) {
|
||||
total = simd_sum(shmem[tiisg]);
|
||||
}
|
||||
const float mean = total / 128.0f;
|
||||
const float scale = 1.0f/sqrt(mean + args.rms_eps);
|
||||
if (tiitg < 32) {
|
||||
y4[tiitg] = (x4[tiitg]*scale)*w4[tiitg];
|
||||
}
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_device);
|
||||
ds4_rope_tail_pair_affine_row(rope_args,
|
||||
(device const char *)index_row,
|
||||
(device char *)index_row,
|
||||
128 - rope_args.n_dims,
|
||||
rope_args.pos0,
|
||||
tiitg,
|
||||
64u);
|
||||
threadgroup_barrier(mem_flags::mem_device);
|
||||
{
|
||||
threadgroup float *vals = shmem;
|
||||
threadgroup float *absbuf = shmem + 128;
|
||||
if (tiitg < 128) {
|
||||
vals[tiitg] = index_row[tiitg];
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
for (uint stride = 1u; stride < 128u; stride <<= 1u) {
|
||||
if (tiitg < 128 && (tiitg & stride) == 0u) {
|
||||
const uint base = (tiitg & ~(2u * stride - 1u)) + (tiitg & (stride - 1u));
|
||||
const float a = vals[base];
|
||||
const float b = vals[base + stride];
|
||||
vals[base] = a + b;
|
||||
vals[base + stride] = a - b;
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
}
|
||||
float v = 0.0f;
|
||||
if (tiitg < 128) {
|
||||
v = vals[tiitg] * 0.08838834764831845f;
|
||||
absbuf[tiitg] = abs(v);
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
const uint block = tiitg >> 5u;
|
||||
const uint lane = tiitg & 31u;
|
||||
const uint block_base = block * 32u;
|
||||
for (uint stride = 16u; stride > 0u; stride >>= 1u) {
|
||||
if (tiitg < 128 && lane < stride) {
|
||||
absbuf[block_base + lane] = max(absbuf[block_base + lane],
|
||||
absbuf[block_base + lane + stride]);
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
}
|
||||
if (tiitg < 128) {
|
||||
const float amax = max(absbuf[block_base], 7.052966104933725e-38f);
|
||||
const float scale = exp2(ceil(log2(amax / 6.0f)));
|
||||
index_row[tiitg] = dsv4_e2m1fn_dequant(clamp(v / scale, -6.0f, 6.0f)) * scale;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Host-visible packed FlashAttention + exact inverse-RoPE decode kernel.
|
||||
kernel void kernel_dsv4_flash_attn_vec_packed32_reduce_rope_f16_dk512_dv512(
|
||||
constant ds4_metal_args_flash_attn_ext_vec & args [[buffer(0)]],
|
||||
device const char * q [[buffer(1)]],
|
||||
device const char * k [[buffer(2)]],
|
||||
device const char * v [[buffer(3)]],
|
||||
device const char * mask [[buffer(4)]],
|
||||
device const char * sinks [[buffer(5)]],
|
||||
device const char * pad [[buffer(6)]],
|
||||
device char * dst [[buffer(7)]],
|
||||
constant ds4_metal_args_dsv4_rope_affine_pair & rope
|
||||
[[buffer(8)]],
|
||||
threadgroup char * shmem [[threadgroup(0)]],
|
||||
uint head [[threadgroup_position_in_grid]],
|
||||
ushort tiitg [[thread_index_in_threadgroup]],
|
||||
ushort tiisg [[thread_index_in_simdgroup]],
|
||||
ushort sgitg [[simdgroup_index_in_threadgroup]]) {
|
||||
/* Uniform specialization guard; host applies the same eligibility gate. */
|
||||
if (!FC_flash_attn_ext_vec_has_mask ||
|
||||
!FC_flash_attn_ext_vec_has_sinks ||
|
||||
FC_flash_attn_ext_vec_has_bias ||
|
||||
FC_flash_attn_ext_vec_has_scap ||
|
||||
FC_flash_attn_ext_vec_nsg != 1 ||
|
||||
FC_flash_attn_ext_vec_nwg != 32 ||
|
||||
FC_flash_attn_ext_vec_ns10 != 512 ||
|
||||
FC_flash_attn_ext_vec_ns20 != 512 ||
|
||||
args.ne01 != 1 || args.ne02 != 64 || args.ne03 != 1 ||
|
||||
args.ne_12_2 != 1 || args.ne_12_3 != 1 ||
|
||||
args.ne31 != 1 || args.ne32 != 1 || args.ne33 != 1 ||
|
||||
args.ne11 <= 0 || args.ne11 > 1024 || head >= (uint)args.ne02 ||
|
||||
args.nb02 != 2048 || args.nb11 != 1024 || args.nb21 != 1024 ||
|
||||
rope.head_dim != 512 || rope.n_dims != 64 ||
|
||||
rope.row_bytes != 2048 || rope.inverse == 0) {
|
||||
return;
|
||||
}
|
||||
|
||||
ds4_flash_attn_vec_packed8_reduce_f16_512(
|
||||
args, q, k, v, mask, sinks, pad, dst, shmem,
|
||||
head, tiisg, sgitg);
|
||||
|
||||
/* Same producer/consumer boundary as the current reduce+RoPE kernel. */
|
||||
threadgroup_barrier(mem_flags::mem_device);
|
||||
|
||||
const int n_nope = rope.head_dim - rope.n_dims;
|
||||
device char * row = dst + (uint64_t)head * rope.row_bytes;
|
||||
ds4_rope_tail_pair_affine_row(rope,
|
||||
(device const char *)row,
|
||||
row,
|
||||
n_nope,
|
||||
rope.pos0,
|
||||
tiitg,
|
||||
32u * 32u);
|
||||
}
|
||||
|
||||
@@ -1398,24 +1398,26 @@ constant int32_t FC_flash_attn_ext_vec_reduce_NWG [[function_constant(FC_FLASH_A
|
||||
|
||||
// Reduces split-K decode FlashAttention partials. It combines each workgroup's
|
||||
// output vector and softmax (sum,max) pair into the final attention result.
|
||||
kernel void kernel_flash_attn_ext_vec_reduce(
|
||||
/* Shared and deliberately noinline so the split-K reduction is compiled once and
|
||||
* every caller gets identical codegen. The RoPE-fused sibling in dsv4_rope.metal
|
||||
* calls this same body, which is what keeps the fusion bit-exact. */
|
||||
static __attribute__((noinline)) void ds4_flash_attn_vec_reduce_row(
|
||||
constant ds4_metal_args_flash_attn_ext_vec_reduce & args,
|
||||
device const char * htmp,
|
||||
device char * dst,
|
||||
uint tgpig[[threadgroup_position_in_grid]],
|
||||
ushort tiisg[[thread_index_in_simdgroup]],
|
||||
ushort sgitg[[simdgroup_index_in_threadgroup]]) {
|
||||
#define NWG (FC_flash_attn_ext_vec_reduce_NWG)
|
||||
#define DV (FC_flash_attn_ext_vec_reduce_DV)
|
||||
|
||||
uint tgpig,
|
||||
ushort tiisg,
|
||||
ushort sgitg,
|
||||
short NWG_,
|
||||
short DV_) {
|
||||
const uint64_t rid = tgpig;
|
||||
|
||||
const short iwg = tiisg;
|
||||
|
||||
device const float * ss = (device const float *) htmp + (uint64_t)args.nrows*DV*NWG;
|
||||
device const float * ss = (device const float *) htmp + (uint64_t)args.nrows*DV_*NWG_;
|
||||
|
||||
float S = ss[rid*(2*NWG) + 2*iwg + 0];
|
||||
float M = ss[rid*(2*NWG) + 2*iwg + 1];
|
||||
float S = ss[rid*(2*NWG_) + 2*iwg + 0];
|
||||
float M = ss[rid*(2*NWG_) + 2*iwg + 1];
|
||||
|
||||
const float m = simd_max(M);
|
||||
const float ms = exp(M - m);
|
||||
@@ -1423,19 +1425,268 @@ kernel void kernel_flash_attn_ext_vec_reduce(
|
||||
S = simd_sum(S*ms);
|
||||
S = S == 0.0f ? 0.0f : 1.0f/S;
|
||||
|
||||
const short DV4 = DV/4;
|
||||
const short DV4 = DV_/4;
|
||||
|
||||
device const float4 * htmp4 = (device const float4 *) htmp + rid*DV4*NWG;
|
||||
device const float4 * htmp4 = (device const float4 *) htmp + rid*DV4*NWG_;
|
||||
device float4 * dst4 = (device float4 *) dst + rid*DV4;
|
||||
|
||||
for (short i = sgitg; i < DV4; i += NWG) {
|
||||
const float4 v = simd_sum(htmp4[i*NWG + iwg]*ms);
|
||||
for (short i = sgitg; i < DV4; i += NWG_) {
|
||||
const float4 v = simd_sum(htmp4[i*NWG_ + iwg]*ms);
|
||||
|
||||
if (iwg == 0) {
|
||||
dst4[i] = v*S;
|
||||
}
|
||||
}
|
||||
|
||||
#undef NWG
|
||||
#undef DV
|
||||
}
|
||||
|
||||
kernel void kernel_flash_attn_ext_vec_reduce(
|
||||
constant ds4_metal_args_flash_attn_ext_vec_reduce & args,
|
||||
device const char * htmp,
|
||||
device char * dst,
|
||||
uint tgpig[[threadgroup_position_in_grid]],
|
||||
ushort tiisg[[thread_index_in_simdgroup]],
|
||||
ushort sgitg[[simdgroup_index_in_threadgroup]]) {
|
||||
ds4_flash_attn_vec_reduce_row(args, htmp, dst, tgpig, tiisg, sgitg,
|
||||
(short)FC_flash_attn_ext_vec_reduce_NWG,
|
||||
(short)FC_flash_attn_ext_vec_reduce_DV);
|
||||
}
|
||||
|
||||
// M5 decode specialization: time-slice all 32 split-K workgroups through eight
|
||||
// physical simdgroups, then reduce through the same 32-lane topology without a
|
||||
// device partial buffer. The host gate fixes the exact F16 512-wide geometry.
|
||||
static inline void ds4_flash_attn_vec_packed8_reduce_f16_512(
|
||||
constant ds4_metal_args_flash_attn_ext_vec & args,
|
||||
device const char * q,
|
||||
device const char * k,
|
||||
device const char * v,
|
||||
device const char * mask,
|
||||
device const char * sinks,
|
||||
device const char * pad,
|
||||
device char * dst,
|
||||
threadgroup char * shmem,
|
||||
uint head,
|
||||
ushort tiisg,
|
||||
ushort sgitg) {
|
||||
constexpr short NW = 32;
|
||||
constexpr short C = 32;
|
||||
constexpr short NSG = 8;
|
||||
constexpr short NWG = 32;
|
||||
constexpr short DK4 = 128;
|
||||
constexpr short DV4 = 128;
|
||||
constexpr short SH = 128;
|
||||
|
||||
/* 24,448 dynamic bytes: shared Q, eight score/mask banks, all 32
|
||||
* split-local weights and stats, sink scales, and a padded 32x33 F32
|
||||
* float4 partial plane. */
|
||||
threadgroup half4 *q_shared = (threadgroup half4 *)shmem;
|
||||
threadgroup half *score_banks =
|
||||
(threadgroup half *)(q_shared + DK4);
|
||||
threadgroup volatile float *weights =
|
||||
(threadgroup volatile float *)(score_banks + NSG * SH);
|
||||
threadgroup volatile float *stats = weights + NWG * C;
|
||||
threadgroup volatile float *sink_scale = stats + 2 * NWG;
|
||||
threadgroup volatile float4 *partial_plane =
|
||||
(threadgroup volatile float4 *)(sink_scale + NWG);
|
||||
|
||||
const short lane = (short)tiisg;
|
||||
threadgroup half *bank = score_banks + (short)sgitg * SH;
|
||||
threadgroup float *ss = (threadgroup float *)bank;
|
||||
threadgroup half *sm = bank + 2 * C;
|
||||
|
||||
device const float4 *q4 =
|
||||
(device const float4 *)(q + (uint64_t)head * args.nb02);
|
||||
if (sgitg == 0) {
|
||||
for (short i = lane; i < DK4; i += NW) {
|
||||
q_shared[i] = (half4)q4[i];
|
||||
}
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
/* Eight physical simdgroups time-slice the exact 32 legacy split-K
|
||||
* workgroups. The official <=1024-key gate gives each virtual split at
|
||||
* most one 32-row block, so its value partial can be formed later from
|
||||
* these materialized weights without changing online-softmax order. */
|
||||
for (short iwg = (short)sgitg; iwg < NWG; iwg += NSG) {
|
||||
float S = 0.0f;
|
||||
float M = -FLT_MAX / 2;
|
||||
float out_scale = 1.0f;
|
||||
const int ic_original = (int)iwg * C;
|
||||
|
||||
weights[(uint)iwg * C + (uint)lane] = 0.0f;
|
||||
ss[lane] = 0.0f;
|
||||
sm[lane] = (half)0.0h;
|
||||
simdgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
if (ic_original < args.ne11) {
|
||||
device const char *k_block = k;
|
||||
device const char *v_block = v;
|
||||
device const half *pm = (device const half *)mask;
|
||||
int ic = ic_original;
|
||||
|
||||
if (FC_flash_attn_ext_vec_has_kvpad && ic + C > args.ne11) {
|
||||
k_block = pad;
|
||||
const uint64_t k_pad_bytes =
|
||||
args.nb11 * (uint64_t)C *
|
||||
(uint64_t)args.ne_12_2 * (uint64_t)args.ne_12_3;
|
||||
const uint64_t v_pad_bytes =
|
||||
args.nb21 * (uint64_t)C *
|
||||
(uint64_t)args.ne_12_2 * (uint64_t)args.ne_12_3;
|
||||
if (FC_flash_attn_ext_vec_shared_kvpad) {
|
||||
v_block = k_block;
|
||||
pm = (device const half *)(k_block +
|
||||
k_pad_bytes + v_pad_bytes);
|
||||
} else {
|
||||
v_block = k_block + k_pad_bytes;
|
||||
pm = (device const half *)(v_block + v_pad_bytes);
|
||||
}
|
||||
ic = 0;
|
||||
}
|
||||
|
||||
sm[lane] = pm[ic + lane];
|
||||
if (simd_max(sm[lane]) > -MAXHALF) {
|
||||
device const half4 *pk4 =
|
||||
(device const half4 *)(k_block +
|
||||
(uint64_t)ic * args.nb11);
|
||||
threadgroup const half4 *pq4 = q_shared;
|
||||
pk4 += lane;
|
||||
pq4 += lane;
|
||||
|
||||
float lane_mqk = 0.0f;
|
||||
FOR_UNROLL (short cc = 0; cc < C; ++cc) {
|
||||
float mqk = 0.0f;
|
||||
FOR_UNROLL (short ii = 0; ii < DK4 / NW; ++ii) {
|
||||
mqk += dot((float4)pk4[cc * DK4 + ii * NW],
|
||||
(float4)pq4[ii * NW]);
|
||||
}
|
||||
mqk = simd_sum(mqk);
|
||||
if (lane == cc) {
|
||||
lane_mqk = mqk;
|
||||
}
|
||||
}
|
||||
|
||||
ss[lane] = fma(lane_mqk, args.scale,
|
||||
(float)sm[lane]);
|
||||
simdgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
const float old_m = M;
|
||||
const float score = ss[lane];
|
||||
M = simd_max(max(M, score));
|
||||
const float ms = exp(old_m - M);
|
||||
const float vs = exp(score - M);
|
||||
S = S * ms + simd_sum(vs);
|
||||
ss[lane] = vs;
|
||||
simdgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
weights[(uint)iwg * C + (uint)lane] = ss[lane];
|
||||
}
|
||||
|
||||
if (FC_flash_attn_ext_vec_has_sinks && iwg == 0) {
|
||||
const float old_m = M;
|
||||
const float sink = lane == 0
|
||||
? ((device const float *)sinks)[head]
|
||||
: -FLT_MAX / 2;
|
||||
M = simd_max(max(M, sink));
|
||||
const float ms = exp(old_m - M);
|
||||
const float vs = exp(sink - M);
|
||||
S = S * ms + simd_sum(vs);
|
||||
out_scale = ms;
|
||||
}
|
||||
} else if (FC_flash_attn_ext_vec_has_sinks && iwg == 0) {
|
||||
const float old_m = M;
|
||||
const float sink = lane == 0
|
||||
? ((device const float *)sinks)[head]
|
||||
: -FLT_MAX / 2;
|
||||
M = simd_max(max(M, sink));
|
||||
const float ms = exp(old_m - M);
|
||||
const float vs = exp(sink - M);
|
||||
S = S * ms + simd_sum(vs);
|
||||
out_scale = ms;
|
||||
}
|
||||
|
||||
if (lane == 0) {
|
||||
stats[2 * (uint)iwg + 0] = S;
|
||||
stats[2 * (uint)iwg + 1] = M;
|
||||
sink_scale[(uint)iwg] = out_scale;
|
||||
}
|
||||
simdgroup_barrier(mem_flags::mem_threadgroup);
|
||||
}
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
/* Recreate the legacy 32-lane reducer exactly: lane is the virtual split
|
||||
* index, including neutral idle splits in their original tree positions. */
|
||||
const short split = lane;
|
||||
float reduce_S = stats[2 * (uint)split + 0];
|
||||
float reduce_M = stats[2 * (uint)split + 1];
|
||||
const float reduce_max = simd_max(reduce_M);
|
||||
const float reduce_ms = exp(reduce_M - reduce_max);
|
||||
reduce_S = simd_sum(reduce_S * reduce_ms);
|
||||
const float reduce_inv =
|
||||
reduce_S == 0.0f ? 0.0f : 1.0f / reduce_S;
|
||||
|
||||
device float4 *dst4 =
|
||||
(device float4 *)(dst +
|
||||
(uint64_t)head * 512u * sizeof(float));
|
||||
|
||||
/* Form one 32-float4 output quadrant at a time. During production each
|
||||
* physical simdgroup time-slices four virtual splits while SIMD lanes are
|
||||
* contiguous output columns, exactly matching the legacy V loads and
|
||||
* cc-major accumulation. A padded 33-column plane avoids a 32-way TG-bank
|
||||
* conflict when the reducer transposes lanes back to virtual splits. */
|
||||
for (short quadrant = 0; quadrant < 4; ++quadrant) {
|
||||
for (short iwg = (short)sgitg; iwg < NWG; iwg += NSG) {
|
||||
float4 lo = float4(0.0f);
|
||||
const int ic_original = (int)iwg * C;
|
||||
if (ic_original < args.ne11) {
|
||||
device const char *v_block = v;
|
||||
int ic = ic_original;
|
||||
if (FC_flash_attn_ext_vec_has_kvpad && ic + C > args.ne11) {
|
||||
device const char *k_block = pad;
|
||||
const uint64_t k_pad_bytes =
|
||||
args.nb11 * (uint64_t)C *
|
||||
(uint64_t)args.ne_12_2 * (uint64_t)args.ne_12_3;
|
||||
if (FC_flash_attn_ext_vec_shared_kvpad) {
|
||||
v_block = k_block;
|
||||
} else {
|
||||
v_block = k_block + k_pad_bytes;
|
||||
}
|
||||
ic = 0;
|
||||
}
|
||||
|
||||
device const half4 *pv4 =
|
||||
(device const half4 *)(v_block +
|
||||
(uint64_t)ic * args.nb21);
|
||||
threadgroup volatile float *split_weights =
|
||||
weights + (uint)iwg * C;
|
||||
const short oc = quadrant * NW + lane;
|
||||
FOR_UNROLL (short cc = 0; cc < C; ++cc) {
|
||||
lo += float4(pv4[cc * DV4 + oc]) *
|
||||
float4(split_weights[cc]);
|
||||
}
|
||||
|
||||
float4 acc = float4(0.0f);
|
||||
acc += lo;
|
||||
if (iwg == 0) {
|
||||
acc *= sink_scale[0];
|
||||
}
|
||||
lo = acc;
|
||||
}
|
||||
partial_plane[(uint)iwg * 33u + (uint)lane] = lo;
|
||||
}
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
/* lane is now the legacy split index and each physical simdgroup
|
||||
* reduces four output columns through the identical simd_sum tree. */
|
||||
for (short out_lane = (short)sgitg; out_lane < NW; out_lane += NSG) {
|
||||
const float4 materialized =
|
||||
(float4)partial_plane[(uint)lane * 33u + (uint)out_lane];
|
||||
const float4 reduced = simd_sum(materialized * reduce_ms);
|
||||
if (lane == 0) {
|
||||
dst4[quadrant * NW + out_lane] = reduced * reduce_inv;
|
||||
}
|
||||
}
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
}
|
||||
}
|
||||
|
||||
2877
metal/moe.metal
2877
metal/moe.metal
File diff suppressed because it is too large
Load Diff
142
metal/norm.metal
142
metal/norm.metal
@@ -241,3 +241,145 @@ kernel void kernel_dsv4_qkv_rms_norm_f32_4(
|
||||
y[i] = (x[i] * scale) * w[i];
|
||||
}
|
||||
}
|
||||
|
||||
// Decode-only triple fusion: the q/kv RMS norm, the KV RoPE tail, and the
|
||||
// FP8/raw finalizer were three back-to-back dispatches on the same rows.
|
||||
// The q threadgroup is byte-identical to kernel_dsv4_qkv_rms_norm_f32_4.
|
||||
// The kv threadgroup continues with the shared affine-row RoPE helper (lane
|
||||
// mapping preserved: r == lane on the first 64 lanes) and a verbatim copy of
|
||||
// kernel_dsv4_kv_fp8_store_f32 with its work predicated to the first 64
|
||||
// lanes (barriers stay uniform across the whole threadgroup). Arithmetic,
|
||||
// order and rounding are unchanged; gated and verified against
|
||||
// full-vocabulary logits before promotion.
|
||||
kernel void kernel_dsv4_qkv_rms_norm_kv_rope_fp8_store_f32(
|
||||
constant ds4_metal_args_qkv_rms_norm & args,
|
||||
constant ds4_metal_args_dsv4_rope_affine_pair & rope,
|
||||
constant ds4_metal_args_dsv4_kv_fp8_store & store,
|
||||
device const float4 * q_src,
|
||||
device const float4 * q_weight,
|
||||
device float4 * q_dst,
|
||||
device const float4 * kv_src,
|
||||
device const float4 * kv_weight,
|
||||
device float4 * kv_dst,
|
||||
device float * raw_cache,
|
||||
threadgroup float * shmem_f32 [[threadgroup(0)]],
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
ushort3 tpitg[[thread_position_in_threadgroup]],
|
||||
ushort sgitg[[simdgroup_index_in_threadgroup]],
|
||||
ushort tiisg[[thread_index_in_simdgroup]],
|
||||
ushort3 ntg[[threads_per_threadgroup]]) {
|
||||
if (sgitg == 0) {
|
||||
shmem_f32[tiisg] = 0.0f;
|
||||
}
|
||||
|
||||
const uint row = tgpig.x;
|
||||
const bool kv_task = tgpig.y != 0;
|
||||
const int n = kv_task ? args.kv_n : args.q_n;
|
||||
const int n4 = kv_task ? args.kv_n4 : args.q_n4;
|
||||
const uint64_t row_stride4 = (kv_task ? args.kv_row_stride : args.q_row_stride) / sizeof(float4);
|
||||
|
||||
device const float4 * x = kv_task ? kv_src + row * row_stride4 : q_src + row * row_stride4;
|
||||
device const float4 * w = kv_task ? kv_weight : q_weight;
|
||||
device float4 * y = kv_task ? kv_dst + row * row_stride4 : q_dst + row * row_stride4;
|
||||
|
||||
float sumf = 0.0f;
|
||||
for (int i = tpitg.x; i < n4; i += ntg.x) {
|
||||
const float4 v = x[i];
|
||||
sumf += dot(v, v);
|
||||
}
|
||||
sumf = simd_sum(sumf);
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
if (tiisg == 0) {
|
||||
shmem_f32[sgitg] = sumf;
|
||||
}
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
sumf = shmem_f32[tiisg];
|
||||
sumf = simd_sum(sumf);
|
||||
|
||||
#ifdef DS4_METAL_NORM_RSQRT_DISABLE
|
||||
const float scale = 1.0f / sqrt(sumf / float(n) + args.eps);
|
||||
#else
|
||||
const float scale = rsqrt(sumf / float(n) + args.eps);
|
||||
#endif
|
||||
|
||||
for (int i = tpitg.x; i < n4; i += ntg.x) {
|
||||
y[i] = (x[i] * scale) * w[i];
|
||||
}
|
||||
|
||||
if (!kv_task) {
|
||||
return;
|
||||
}
|
||||
|
||||
// KV RoPE tail in place, then the FP8/raw finalizer (verbatim bodies).
|
||||
threadgroup_barrier(mem_flags::mem_device_and_threadgroup);
|
||||
|
||||
device char *kv_row = (device char *)(kv_dst + row * row_stride4);
|
||||
const int rope_n_nope = rope.head_dim - rope.n_dims;
|
||||
if (rope_n_nope < 0) {
|
||||
return;
|
||||
}
|
||||
ds4_rope_tail_pair_affine_row(rope,
|
||||
(device const char *)kv_row,
|
||||
kv_row,
|
||||
rope_n_nope,
|
||||
rope.pos0,
|
||||
tpitg.x,
|
||||
ntg.x);
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_device_and_threadgroup);
|
||||
|
||||
const int head_dim = store.head_dim;
|
||||
const int n_rot = store.n_rot;
|
||||
const int n_nope = head_dim - n_rot;
|
||||
if (head_dim <= 0 || n_rot < 0 || n_nope < 0) {
|
||||
return;
|
||||
}
|
||||
const uint tid = tpitg.x;
|
||||
|
||||
device float *kv = (device float *)kv_row;
|
||||
device float *raw = raw_cache + (int64_t)store.raw_row * head_dim;
|
||||
threadgroup float *scratch = shmem_f32 + 32;
|
||||
|
||||
for (int off = 0; off < n_nope; off += 64) {
|
||||
float v = 0.0f;
|
||||
if (tid < 64u && off + (int)tid < n_nope) {
|
||||
v = kv[off + tid];
|
||||
scratch[tid] = abs(v);
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
for (uint stride = 32; stride > 0; stride >>= 1) {
|
||||
if (tid < stride) {
|
||||
scratch[tid] = max(scratch[tid], scratch[tid + stride]);
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
}
|
||||
|
||||
const float amax = max(scratch[0], 1.0e-4f);
|
||||
const float fp8_scale = exp2(ceil(log2(amax / 448.0f)));
|
||||
if (tid < 64u && off + (int)tid < n_nope) {
|
||||
const float q = dsv4_e4m3fn_dequant(clamp(v / fp8_scale, -448.0f, 448.0f)) * fp8_scale;
|
||||
kv[off + tid] = q;
|
||||
#ifdef DS4_METAL_KV_RAW_F32
|
||||
raw[off + tid] = q;
|
||||
#else
|
||||
raw[off + tid] = (float)((half)q);
|
||||
#endif
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
}
|
||||
|
||||
if (tid < 64u) {
|
||||
for (int i = n_nope + tid; i < head_dim; i += 64) {
|
||||
#ifdef DS4_METAL_KV_RAW_F32
|
||||
raw[i] = kv[i];
|
||||
#else
|
||||
raw[i] = (float)((half)kv[i]);
|
||||
#endif
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user