1442 lines
50 KiB
Metal
1442 lines
50 KiB
Metal
#define FC_FLASH_ATTN_EXT_PAD 100
|
|
#define FC_FLASH_ATTN_EXT_BLK 200
|
|
#define FC_FLASH_ATTN_EXT 300
|
|
#define FC_FLASH_ATTN_EXT_VEC 400
|
|
#define FC_FLASH_ATTN_EXT_VEC_REDUCE 500
|
|
#define OP_FLASH_ATTN_EXT_NQPSG 8
|
|
#define OP_FLASH_ATTN_EXT_NCPSG 64
|
|
#define OP_FLASH_ATTN_EXT_VEC_NQPSG 1
|
|
#define OP_FLASH_ATTN_EXT_VEC_NCPSG 32
|
|
|
|
#ifndef PAD2
|
|
#define PAD2(x, n) (((x) + (n) - 1) & ~((n) - 1))
|
|
#endif
|
|
|
|
template <typename type4>
|
|
void dequantize_f32_t4(device const float4 * src, short il, thread type4 & reg) {
|
|
reg = (type4)(*src);
|
|
}
|
|
|
|
template <typename type4>
|
|
void dequantize_f16_t4(device const half4 * src, short il, thread type4 & reg) {
|
|
reg = (type4)(*(src));
|
|
}
|
|
|
|
template <typename type4x4>
|
|
void dequantize_f32(device const float4x4 * src, short il, thread type4x4 & reg);
|
|
|
|
template <typename type4x4>
|
|
void dequantize_f16(device const half4x4 * src, short il, thread type4x4 & reg);
|
|
|
|
struct ds4_metal_args_flash_attn_ext_pad {
|
|
int32_t ne11;
|
|
int32_t ne_12_2;
|
|
int32_t ne_12_3;
|
|
uint64_t nb11;
|
|
uint64_t nb12;
|
|
uint64_t nb13;
|
|
uint64_t nb21;
|
|
uint64_t nb22;
|
|
uint64_t nb23;
|
|
int32_t ne31;
|
|
int32_t ne32;
|
|
int32_t ne33;
|
|
uint64_t nb31;
|
|
uint64_t nb32;
|
|
uint64_t nb33;
|
|
};
|
|
|
|
struct ds4_metal_args_flash_attn_ext_blk {
|
|
int32_t ne01;
|
|
int32_t ne30;
|
|
int32_t ne31;
|
|
int32_t ne32;
|
|
int32_t ne33;
|
|
uint64_t nb31;
|
|
uint64_t nb32;
|
|
uint64_t nb33;
|
|
};
|
|
|
|
struct ds4_metal_args_flash_attn_ext {
|
|
int32_t ne01;
|
|
int32_t ne02;
|
|
int32_t ne03;
|
|
uint64_t nb01;
|
|
uint64_t nb02;
|
|
uint64_t nb03;
|
|
int32_t ne11;
|
|
int32_t ne_12_2;
|
|
int32_t ne_12_3;
|
|
int32_t ns10;
|
|
uint64_t nb11;
|
|
uint64_t nb12;
|
|
uint64_t nb13;
|
|
int32_t ns20;
|
|
uint64_t nb21;
|
|
uint64_t nb22;
|
|
uint64_t nb23;
|
|
int32_t ne31;
|
|
int32_t ne32;
|
|
int32_t ne33;
|
|
uint64_t nb31;
|
|
uint64_t nb32;
|
|
uint64_t nb33;
|
|
int32_t ne1;
|
|
int32_t ne2;
|
|
int32_t ne3;
|
|
float scale;
|
|
float max_bias;
|
|
float m0;
|
|
float m1;
|
|
int32_t n_head_log2;
|
|
float logit_softcap;
|
|
};
|
|
|
|
struct ds4_metal_args_flash_attn_ext_vec {
|
|
int32_t ne01;
|
|
int32_t ne02;
|
|
int32_t ne03;
|
|
uint64_t nb01;
|
|
uint64_t nb02;
|
|
uint64_t nb03;
|
|
int32_t ne11;
|
|
int32_t ne_12_2;
|
|
int32_t ne_12_3;
|
|
int32_t ns10;
|
|
uint64_t nb11;
|
|
uint64_t nb12;
|
|
uint64_t nb13;
|
|
int32_t ns20;
|
|
uint64_t nb21;
|
|
uint64_t nb22;
|
|
uint64_t nb23;
|
|
int32_t ne31;
|
|
int32_t ne32;
|
|
int32_t ne33;
|
|
uint64_t nb31;
|
|
uint64_t nb32;
|
|
uint64_t nb33;
|
|
int32_t ne1;
|
|
int32_t ne2;
|
|
int32_t ne3;
|
|
float scale;
|
|
float max_bias;
|
|
float m0;
|
|
float m1;
|
|
int32_t n_head_log2;
|
|
float logit_softcap;
|
|
};
|
|
|
|
struct ds4_metal_args_flash_attn_ext_vec_reduce {
|
|
int32_t nrows;
|
|
};
|
|
|
|
constant bool FC_flash_attn_ext_pad_has_mask [[function_constant(FC_FLASH_ATTN_EXT_PAD + 0)]];
|
|
constant int32_t FC_flash_attn_ext_pad_ncpsg [[function_constant(FC_FLASH_ATTN_EXT_PAD + 25)]];
|
|
|
|
// DS4 FlashAttention padding: pads the final partial K/V/mask cache block so the
|
|
// vector FlashAttention kernel can read full 32-row chunks.
|
|
kernel void kernel_flash_attn_ext_pad(
|
|
constant ds4_metal_args_flash_attn_ext_pad & args,
|
|
device const char * k,
|
|
device const char * v,
|
|
device const char * mask,
|
|
device char * dst,
|
|
uint3 tgpig[[threadgroup_position_in_grid]],
|
|
ushort tiitg[[thread_index_in_threadgroup]],
|
|
ushort3 ntg[[threads_per_threadgroup]]) {
|
|
const int32_t C = FC_flash_attn_ext_pad_ncpsg;
|
|
|
|
device char * k_pad = dst;
|
|
device char * v_pad = k_pad + args.nb11*C*args.ne_12_2*args.ne_12_3;
|
|
device char * mask_pad = v_pad + args.nb21*C*args.ne_12_2*args.ne_12_3;
|
|
|
|
const int32_t icp = args.ne11 % C;
|
|
const int32_t ic0 = args.ne11 - icp;
|
|
|
|
const int32_t i1 = tgpig[0];
|
|
const int32_t i2 = tgpig[1];
|
|
const int32_t i3 = tgpig[2];
|
|
|
|
if (i2 < args.ne_12_2 && i3 < args.ne_12_3) {
|
|
device const char * k_src = k + args.nb11*(ic0 + i1) + args.nb12*i2 + args.nb13*i3;
|
|
device const char * v_src = v + args.nb21*(ic0 + i1) + args.nb22*i2 + args.nb23*i3;
|
|
|
|
device char * k_dst = k_pad + args.nb11*i1 + args.nb11*C*i2 + args.nb11*C*args.ne_12_2*i3;
|
|
device char * v_dst = v_pad + args.nb21*i1 + args.nb21*C*i2 + args.nb21*C*args.ne_12_2*i3;
|
|
|
|
if (i1 >= icp) {
|
|
for (uint64_t i = tiitg; i < args.nb11; i += ntg.x) {
|
|
k_dst[i] = 0;
|
|
}
|
|
for (uint64_t i = tiitg; i < args.nb21; i += ntg.x) {
|
|
v_dst[i] = 0;
|
|
}
|
|
} else {
|
|
for (uint64_t i = tiitg; i < args.nb11; i += ntg.x) {
|
|
k_dst[i] = k_src[i];
|
|
}
|
|
for (uint64_t i = tiitg; i < args.nb21; i += ntg.x) {
|
|
v_dst[i] = v_src[i];
|
|
}
|
|
}
|
|
}
|
|
|
|
if (FC_flash_attn_ext_pad_has_mask) {
|
|
if (i2 < args.ne32 && i3 < args.ne33) {
|
|
for (int ib = i1; ib < args.ne31; ib += C) {
|
|
device const half * mask_src = (device const half *)(mask + args.nb31*ib + args.nb32*i2 + args.nb33*i3) + ic0;
|
|
device half * mask_dst = (device half *)(mask_pad) + C*ib + C*args.ne31*i2 + C*args.ne31*args.ne32*i3;
|
|
|
|
for (int i = tiitg; i < C; i += ntg.x) {
|
|
if (i >= icp) {
|
|
mask_dst[i] = -MAXHALF;
|
|
} else {
|
|
mask_dst[i] = mask_src[i];
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
constant int32_t FC_flash_attn_ext_blk_nqptg [[function_constant(FC_FLASH_ATTN_EXT_BLK + 24)]];
|
|
constant int32_t FC_flash_attn_ext_blk_ncpsg [[function_constant(FC_FLASH_ATTN_EXT_BLK + 25)]];
|
|
|
|
// DS4 FlashAttention mask scan: marks blocks so the non-vector kernel can skip
|
|
// blocks that are entirely masked or entirely zero.
|
|
kernel void kernel_flash_attn_ext_blk(
|
|
constant ds4_metal_args_flash_attn_ext_blk & args,
|
|
device const char * mask,
|
|
device char * dst,
|
|
uint3 tgpig[[threadgroup_position_in_grid]],
|
|
ushort tiisg[[thread_index_in_simdgroup]]) {
|
|
const int32_t Q = FC_flash_attn_ext_blk_nqptg;
|
|
const int32_t C = FC_flash_attn_ext_blk_ncpsg;
|
|
|
|
constexpr short NW = N_SIMDWIDTH;
|
|
|
|
const int32_t i3 = tgpig[2]/args.ne32;
|
|
const int32_t i2 = tgpig[2]%args.ne32;
|
|
const int32_t i1 = tgpig[1];
|
|
const int32_t i0 = tgpig[0];
|
|
|
|
char res = i0*C + C > args.ne30 ? 1 : 0;
|
|
|
|
if ((C > NW || Q > 1) && res == 0) {
|
|
half mmin = MAXHALF;
|
|
half mmax = -MAXHALF;
|
|
const int32_t q0 = i1*Q;
|
|
|
|
FOR_UNROLL (short j = 0; j < Q; ++j) {
|
|
if (q0 + j < args.ne31) {
|
|
device const half * mask_src =
|
|
(device const half *) (mask + (q0 + j)*args.nb31 + i2*args.nb32 + i3*args.nb33) +
|
|
i0*C + tiisg;
|
|
|
|
FOR_UNROLL (short ii = 0; ii < C/NW; ++ii) {
|
|
mmin = min(mmin, mask_src[ii*NW]);
|
|
mmax = max(mmax, mask_src[ii*NW]);
|
|
}
|
|
}
|
|
}
|
|
|
|
mmin = simd_min(mmin);
|
|
mmax = simd_max(mmax);
|
|
|
|
if (mmax > -MAXHALF) {
|
|
if (mmin == 0.0 && mmax == 0.0) {
|
|
res = 2;
|
|
} else {
|
|
res = 1;
|
|
}
|
|
}
|
|
}
|
|
|
|
const int32_t nblk1 = ((args.ne01 + Q - 1)/Q);
|
|
const int32_t nblk0 = ((args.ne30 + C - 1)/C);
|
|
|
|
if (tiisg == 0) {
|
|
dst[((i3*args.ne32 + i2)*nblk1 + i1)*nblk0 + i0] = res;
|
|
}
|
|
}
|
|
|
|
constant bool FC_flash_attn_ext_has_mask [[function_constant(FC_FLASH_ATTN_EXT + 0)]];
|
|
constant bool FC_flash_attn_ext_has_sinks [[function_constant(FC_FLASH_ATTN_EXT + 1)]];
|
|
constant bool FC_flash_attn_ext_has_bias [[function_constant(FC_FLASH_ATTN_EXT + 2)]];
|
|
constant bool FC_flash_attn_ext_has_scap [[function_constant(FC_FLASH_ATTN_EXT + 3)]];
|
|
constant bool FC_flash_attn_ext_has_kvpad [[function_constant(FC_FLASH_ATTN_EXT + 4)]];
|
|
|
|
constant bool FC_flash_attn_ext_bc_mask [[function_constant(FC_FLASH_ATTN_EXT + 10)]];
|
|
|
|
constant int32_t FC_flash_attn_ext_ns10 [[function_constant(FC_FLASH_ATTN_EXT + 20)]];
|
|
constant int32_t FC_flash_attn_ext_ns20 [[function_constant(FC_FLASH_ATTN_EXT + 21)]];
|
|
constant int32_t FC_flash_attn_ext_nsg [[function_constant(FC_FLASH_ATTN_EXT + 22)]];
|
|
|
|
// DS4 non-vector FlashAttention. The only exported instance uses the model's
|
|
// 512-wide F16 K/V rows; keeping the template body generic preserves the same
|
|
// arithmetic for dense and compressed-attention prefill.
|
|
template<
|
|
typename q_t,
|
|
typename q4_t,
|
|
typename q8x8_t,
|
|
typename k_t,
|
|
typename k4x4_t,
|
|
typename k8x8_t,
|
|
typename v_t,
|
|
typename v4x4_t,
|
|
typename v8x8_t,
|
|
typename qk_t,
|
|
typename qk8x8_t,
|
|
typename s_t,
|
|
typename s2_t,
|
|
typename s8x8_t,
|
|
typename o_t,
|
|
typename o4_t,
|
|
typename o8x8_t,
|
|
typename kd4x4_t,
|
|
short nl_k,
|
|
void (*deq_k)(device const kd4x4_t *, short, thread k4x4_t &),
|
|
typename vd4x4_t,
|
|
short nl_v,
|
|
void (*deq_v)(device const vd4x4_t *, short, thread v4x4_t &),
|
|
short DK,
|
|
short DV,
|
|
short Q,
|
|
short C,
|
|
short NSG>
|
|
void kernel_flash_attn_ext_impl(
|
|
constant ds4_metal_args_flash_attn_ext & 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 const char * blk,
|
|
device char * dst,
|
|
threadgroup half * shmem_f16,
|
|
uint3 tgpig,
|
|
ushort tiisg,
|
|
ushort sgitg) {
|
|
const ushort iq3 = tgpig[2];
|
|
const ushort iq2 = tgpig[1];
|
|
const ushort iq1 = tgpig[0]*Q;
|
|
|
|
#define NS10 (FC_flash_attn_ext_ns10)
|
|
#define NS20 (FC_flash_attn_ext_ns20)
|
|
|
|
constexpr short KV = 8;
|
|
|
|
constexpr short DK4 = DK/4;
|
|
constexpr short DK8 = DK/8;
|
|
constexpr short DK16 = DK/16;
|
|
constexpr short DV4 = DV/4;
|
|
constexpr short DV16 = DV/16;
|
|
|
|
constexpr short PV = PAD2(DV, 64);
|
|
constexpr short PV4 = PV/4;
|
|
constexpr short PV8 = PV/8;
|
|
|
|
constexpr short NW = N_SIMDWIDTH;
|
|
constexpr short NQ = Q/NSG;
|
|
constexpr short SH = 2*C;
|
|
|
|
constexpr short TS = 2*SH;
|
|
constexpr short T = DK + 2*PV;
|
|
|
|
threadgroup q_t * sq = (threadgroup q_t *) (shmem_f16 + 0*T);
|
|
threadgroup q4_t * sq4 = (threadgroup q4_t *) (shmem_f16 + 0*T);
|
|
threadgroup o_t * so = (threadgroup o_t *) (shmem_f16 + 0*T + Q*DK);
|
|
threadgroup o4_t * so4 = (threadgroup o4_t *) (shmem_f16 + 0*T + Q*DK);
|
|
threadgroup s_t * ss = (threadgroup s_t *) (shmem_f16 + Q*T);
|
|
threadgroup s2_t * ss2 = (threadgroup s2_t *) (shmem_f16 + Q*T);
|
|
|
|
threadgroup k_t * sk = (threadgroup k_t *) (shmem_f16 + sgitg*(4*16*KV) + Q*T + Q*TS);
|
|
threadgroup k4x4_t * sk4x4 = (threadgroup k4x4_t *) (shmem_f16 + sgitg*(4*16*KV) + Q*T + Q*TS);
|
|
|
|
threadgroup v_t * sv = (threadgroup v_t *) (shmem_f16 + sgitg*(4*16*KV) + Q*T + Q*TS);
|
|
threadgroup v4x4_t * sv4x4 = (threadgroup v4x4_t *) (shmem_f16 + sgitg*(4*16*KV) + Q*T + Q*TS);
|
|
|
|
threadgroup half2 * sm2 = (threadgroup half2 *) (shmem_f16 + Q*T + 2*C);
|
|
|
|
device const half2 * pm2[NQ];
|
|
|
|
FOR_UNROLL (short jj = 0; jj < NQ; ++jj) {
|
|
const short j = jj*NSG + sgitg;
|
|
|
|
pm2[jj] = (device const half2 *) ((device const char *) mask + (iq1 + j)*args.nb31 + (iq2%args.ne32)*args.nb32 + (iq3%args.ne33)*args.nb33);
|
|
}
|
|
|
|
{
|
|
const int32_t nblk1 = ((args.ne01 + Q - 1)/Q);
|
|
const int32_t nblk0 = ((args.ne11 + C - 1)/C);
|
|
|
|
blk += (((iq3%args.ne33)*args.ne32 + (iq2%args.ne32))*nblk1 + iq1/Q)*nblk0;
|
|
}
|
|
|
|
{
|
|
q += iq1*args.nb01 + iq2*args.nb02 + iq3*args.nb03;
|
|
|
|
const short ikv2 = iq2/(args.ne02/args.ne_12_2);
|
|
const short ikv3 = iq3/(args.ne03/args.ne_12_3);
|
|
|
|
k += ikv2*args.nb12 + ikv3*args.nb13;
|
|
v += ikv2*args.nb22 + ikv3*args.nb23;
|
|
}
|
|
|
|
FOR_UNROLL (short jj = 0; jj < NQ; ++jj) {
|
|
const short j = jj*NSG + sgitg;
|
|
|
|
device const float4 * q4 = (device const float4 *) ((device const char *) q + j*args.nb01);
|
|
|
|
for (short i = tiisg; i < DK4; i += NW) {
|
|
if (iq1 + j < args.ne01) {
|
|
sq4[j*DK4 + i] = (q4_t) q4[i];
|
|
} else {
|
|
sq4[j*DK4 + i] = 0;
|
|
}
|
|
}
|
|
}
|
|
|
|
FOR_UNROLL (short jj = 0; jj < NQ; ++jj) {
|
|
const short j = jj*NSG + sgitg;
|
|
|
|
for (short i = tiisg; i < DV4; i += NW) {
|
|
so4[j*PV4 + i] = 0;
|
|
}
|
|
|
|
for (short i = tiisg; i < SH; i += NW) {
|
|
ss[j*SH + i] = 0.0f;
|
|
}
|
|
}
|
|
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
float S[NQ] = { [0 ... NQ-1] = 0.0f };
|
|
|
|
{
|
|
float M[NQ] = { [0 ... NQ-1] = -FLT_MAX/2 };
|
|
|
|
float slope = 1.0f;
|
|
|
|
if (FC_flash_attn_ext_has_bias) {
|
|
const short h = iq2;
|
|
|
|
const float base = h < args.n_head_log2 ? args.m0 : args.m1;
|
|
const short exph = h < args.n_head_log2 ? h + 1 : 2*(h - args.n_head_log2) + 1;
|
|
|
|
slope = pow(base, exph);
|
|
}
|
|
|
|
for (int ic0 = 0; ; ++ic0) {
|
|
int ic = ic0*C;
|
|
if (ic >= args.ne11) {
|
|
break;
|
|
}
|
|
|
|
if (FC_flash_attn_ext_has_kvpad && ic + C > args.ne11) {
|
|
k = pad;
|
|
v = k + args.nb11*C*args.ne_12_2*args.ne_12_3;
|
|
mask = v + args.nb21*C*args.ne_12_2*args.ne_12_3;
|
|
|
|
const short ikv2 = iq2/(args.ne02/args.ne_12_2);
|
|
const short ikv3 = iq3/(args.ne03/args.ne_12_3);
|
|
|
|
k += (ikv2 + ikv3*args.ne_12_2)*args.nb11*C;
|
|
v += (ikv2 + ikv3*args.ne_12_2)*args.nb21*C;
|
|
|
|
if (!FC_flash_attn_ext_has_mask) {
|
|
threadgroup half * sm = (threadgroup half *) (sm2);
|
|
|
|
FOR_UNROLL (short jj = 0; jj < NQ; ++jj) {
|
|
const short j = jj*NSG + sgitg;
|
|
|
|
for (short i = tiisg; i < C; i += NW) {
|
|
if (ic + i >= args.ne11) {
|
|
sm[2*j*SH + i] = -MAXHALF;
|
|
}
|
|
}
|
|
}
|
|
} else {
|
|
FOR_UNROLL (short jj = 0; jj < NQ; ++jj) {
|
|
const short j = jj*NSG + sgitg;
|
|
|
|
pm2[jj] = (device const half2 *) ((device const half *) mask +
|
|
(iq1 + j)*C +
|
|
(iq2%args.ne32)*(C*args.ne31) +
|
|
(iq3%args.ne33)*(C*args.ne31*args.ne32));
|
|
}
|
|
}
|
|
|
|
ic = 0;
|
|
}
|
|
|
|
char blk_cur = 1;
|
|
|
|
if (FC_flash_attn_ext_has_mask) {
|
|
blk_cur = blk[ic0];
|
|
|
|
if (blk_cur == 0) {
|
|
FOR_UNROLL (short jj = 0; jj < NQ; ++jj) {
|
|
pm2[jj] += NW;
|
|
}
|
|
|
|
continue;
|
|
}
|
|
|
|
if (blk_cur == 1) {
|
|
FOR_UNROLL (short jj = 0; jj < NQ; ++jj) {
|
|
const short j = jj*NSG + sgitg;
|
|
|
|
if (FC_flash_attn_ext_bc_mask) {
|
|
sm2[j*SH + tiisg] = (iq1 + j) < args.ne31 ? pm2[jj][tiisg] : half2(-MAXHALF, -MAXHALF);
|
|
} else {
|
|
sm2[j*SH + tiisg] = pm2[jj][tiisg];
|
|
}
|
|
|
|
pm2[jj] += NW;
|
|
}
|
|
} else if (blk_cur == 2) {
|
|
FOR_UNROLL (short jj = 0; jj < NQ; ++jj) {
|
|
pm2[jj] += NW;
|
|
}
|
|
}
|
|
}
|
|
|
|
if (is_same<kd4x4_t, k4x4_t>::value) {
|
|
device const k_t * pk = (device const k_t *) (k + ic*args.nb11);
|
|
threadgroup const q_t * pq = sq;
|
|
threadgroup s_t * ps = ss;
|
|
|
|
pk += sgitg*(8*NS10);
|
|
ps += sgitg*(8*1);
|
|
|
|
static_assert((C/8) % NSG == 0, "");
|
|
|
|
constexpr short NC = (C/8)/NSG;
|
|
|
|
FOR_UNROLL (short cc = 0; cc < NC; ++cc) {
|
|
qk8x8_t mqk = make_filled_simdgroup_matrix<qk_t, 8>((qk_t) 0.0f);
|
|
|
|
if (DK % 16 != 0) {
|
|
k8x8_t mk;
|
|
q8x8_t mq;
|
|
|
|
FOR_UNROLL (short i = 0; i < DK8; ++i) {
|
|
simdgroup_barrier(mem_flags::mem_none);
|
|
|
|
simdgroup_load(mk, pk + 8*i, NS10, 0, true);
|
|
simdgroup_load(mq, pq + 8*i, DK);
|
|
|
|
simdgroup_barrier(mem_flags::mem_none);
|
|
|
|
simdgroup_multiply_accumulate(mqk, mq, mk, mqk);
|
|
}
|
|
} else {
|
|
k8x8_t mk[2];
|
|
q8x8_t mq[2];
|
|
|
|
#pragma unroll (MIN(DK8/2, 4*NSG))
|
|
for (short i = 0; i < DK8/2; ++i) {
|
|
simdgroup_barrier(mem_flags::mem_none);
|
|
|
|
simdgroup_load(mq[0], pq + 0*8 + 16*i, DK);
|
|
simdgroup_load(mq[1], pq + 1*8 + 16*i, DK);
|
|
|
|
simdgroup_load(mk[0], pk + 0*8 + 16*i, NS10, 0, true);
|
|
simdgroup_load(mk[1], pk + 1*8 + 16*i, NS10, 0, true);
|
|
|
|
simdgroup_barrier(mem_flags::mem_none);
|
|
|
|
simdgroup_multiply_accumulate(mqk, mq[0], mk[0], mqk);
|
|
simdgroup_multiply_accumulate(mqk, mq[1], mk[1], mqk);
|
|
}
|
|
}
|
|
|
|
simdgroup_store(mqk, ps, SH, 0, false);
|
|
|
|
pk += 8*(NSG*NS10);
|
|
ps += 8*(NSG);
|
|
}
|
|
} else {
|
|
for (short ccc = 0; ccc < (C/8)/NSG; ++ccc) {
|
|
const short cc = ccc*NSG + sgitg;
|
|
|
|
const short tx = tiisg%4;
|
|
const short ty = tiisg/4;
|
|
|
|
qk8x8_t mqk = make_filled_simdgroup_matrix<qk_t, 8>((qk_t) 0.0f);
|
|
|
|
for (short ii = 0; ii < DK16; ii += 4) {
|
|
device const kd4x4_t * pk4x4 = (device const kd4x4_t *) (k + ((ic + 8*cc + ty)*args.nb11));
|
|
|
|
if (DK16%4 == 0) {
|
|
{
|
|
k4x4_t tmp;
|
|
deq_k(pk4x4 + (ii + tx)/nl_k, (ii + tx)%nl_k, tmp);
|
|
sk4x4[4*ty + tx] = tmp;
|
|
}
|
|
|
|
simdgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
FOR_UNROLL (short k = 0; k < 4; ++k) {
|
|
k8x8_t mk;
|
|
q8x8_t mq;
|
|
|
|
simdgroup_load(mk, sk + 16*k + 0*8, 4*16, 0, true);
|
|
simdgroup_load(mq, sq + (2*(ii + k) + 0)*8, DK);
|
|
simdgroup_multiply_accumulate(mqk, mq, mk, mqk);
|
|
|
|
simdgroup_load(mk, sk + 16*k + 1*8, 4*16, 0, true);
|
|
simdgroup_load(mq, sq + (2*(ii + k) + 1)*8, DK);
|
|
simdgroup_multiply_accumulate(mqk, mq, mk, mqk);
|
|
}
|
|
} else {
|
|
if (ii + tx < DK16) {
|
|
k4x4_t tmp;
|
|
deq_k(pk4x4 + (ii + tx)/nl_k, (ii + tx)%nl_k, tmp);
|
|
sk4x4[4*ty + tx] = tmp;
|
|
}
|
|
|
|
simdgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
for (short k = 0; k < 4 && ii + k < DK16; ++k) {
|
|
k8x8_t mk;
|
|
q8x8_t mq;
|
|
|
|
simdgroup_load(mk, sk + 16*k + 0*8, 4*16, 0, true);
|
|
simdgroup_load(mq, sq + (2*(ii + k) + 0)*8, DK);
|
|
simdgroup_multiply_accumulate(mqk, mq, mk, mqk);
|
|
|
|
simdgroup_load(mk, sk + 16*k + 1*8, 4*16, 0, true);
|
|
simdgroup_load(mq, sq + (2*(ii + k) + 1)*8, DK);
|
|
simdgroup_multiply_accumulate(mqk, mq, mk, mqk);
|
|
}
|
|
}
|
|
}
|
|
|
|
simdgroup_store(mqk, ss + 8*cc, SH, 0, false);
|
|
}
|
|
}
|
|
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
FOR_UNROLL (short jj = 0; jj < NQ; ++jj) {
|
|
const short j = jj*NSG + sgitg;
|
|
|
|
const float m = M[jj];
|
|
|
|
float2 s2 = ss2[j*SH/2 + tiisg]*args.scale;
|
|
|
|
if (FC_flash_attn_ext_has_scap) {
|
|
s2 = args.logit_softcap*precise::tanh(s2);
|
|
}
|
|
|
|
if (blk_cur != 2) {
|
|
if (FC_flash_attn_ext_has_bias) {
|
|
s2 += s2_t(sm2[j*SH + tiisg])*slope;
|
|
} else {
|
|
s2 += s2_t(sm2[j*SH + tiisg]);
|
|
}
|
|
}
|
|
|
|
M[jj] = simd_max(max(M[jj], max(s2[0], s2[1])));
|
|
|
|
const float ms = exp(m - M[jj]);
|
|
const float2 vs2 = exp(s2 - M[jj]);
|
|
|
|
S[jj] = S[jj]*ms + simd_sum(vs2[0] + vs2[1]);
|
|
|
|
ss2[j*SH/2 + tiisg] = vs2;
|
|
|
|
if (DV4 % NW == 0) {
|
|
FOR_UNROLL (short ii = 0; ii < DV4/NW; ++ii) {
|
|
const short i = ii*NW + tiisg;
|
|
|
|
so4[j*PV4 + i] *= ms;
|
|
}
|
|
} else {
|
|
for (short i = tiisg; i < DV4; i += NW) {
|
|
so4[j*PV4 + i] *= ms;
|
|
}
|
|
}
|
|
}
|
|
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
{
|
|
if (is_same<vd4x4_t, v4x4_t>::value) {
|
|
static_assert(PV8 % NSG == 0, "");
|
|
|
|
constexpr short NO = PV8/NSG;
|
|
|
|
o8x8_t lo[NO];
|
|
|
|
{
|
|
auto sot = so + 8*sgitg;
|
|
|
|
FOR_UNROLL (short ii = 0; ii < NO; ++ii) {
|
|
simdgroup_load(lo[ii], sot, PV, 0, false);
|
|
|
|
sot += 8*NSG;
|
|
}
|
|
}
|
|
|
|
{
|
|
device const v_t * pv = (device const v_t *) (v + ic*args.nb21);
|
|
|
|
pv += 8*sgitg;
|
|
|
|
if (DV <= 64) {
|
|
FOR_UNROLL (short cc = 0; cc < C/8; ++cc) {
|
|
s8x8_t vs;
|
|
simdgroup_load(vs, ss + 8*cc, SH, 0, false);
|
|
|
|
FOR_UNROLL (short ii = 0; ii < NO/2; ++ii) {
|
|
v8x8_t mv[2];
|
|
|
|
simdgroup_load(mv[0], pv + 0*NSG + 16*ii*NSG, NS20, 0, false);
|
|
simdgroup_load(mv[1], pv + 8*NSG + 16*ii*NSG, NS20, 0, false);
|
|
|
|
simdgroup_multiply_accumulate(lo[2*ii + 0], vs, mv[0], lo[2*ii + 0]);
|
|
simdgroup_multiply_accumulate(lo[2*ii + 1], vs, mv[1], lo[2*ii + 1]);
|
|
}
|
|
|
|
pv += 8*NS20;
|
|
}
|
|
} else {
|
|
constexpr short NC = (C/8)/2;
|
|
|
|
FOR_UNROLL (short cc = 0; cc < NC; ++cc) {
|
|
s8x8_t vs[2];
|
|
|
|
simdgroup_load(vs[0], ss + 16*cc + 0, SH, 0, false);
|
|
simdgroup_load(vs[1], ss + 16*cc + 8, SH, 0, false);
|
|
|
|
FOR_UNROLL (short ii = 0; ii < NO/2; ++ii) {
|
|
v8x8_t mv[4];
|
|
|
|
simdgroup_load(mv[0], pv + 0*NSG + 16*ii*NSG + 0*8*NS20, NS20, 0, false);
|
|
simdgroup_load(mv[1], pv + 8*NSG + 16*ii*NSG + 0*8*NS20, NS20, 0, false);
|
|
simdgroup_load(mv[2], pv + 0*NSG + 16*ii*NSG + 1*8*NS20, NS20, 0, false);
|
|
simdgroup_load(mv[3], pv + 8*NSG + 16*ii*NSG + 1*8*NS20, NS20, 0, false);
|
|
|
|
simdgroup_multiply_accumulate(lo[2*ii + 0], vs[0], mv[0], lo[2*ii + 0]);
|
|
simdgroup_multiply_accumulate(lo[2*ii + 1], vs[0], mv[1], lo[2*ii + 1]);
|
|
simdgroup_multiply_accumulate(lo[2*ii + 0], vs[1], mv[2], lo[2*ii + 0]);
|
|
simdgroup_multiply_accumulate(lo[2*ii + 1], vs[1], mv[3], lo[2*ii + 1]);
|
|
}
|
|
|
|
pv += 2*8*NS20;
|
|
}
|
|
}
|
|
}
|
|
|
|
{
|
|
auto sot = so + 8*sgitg;
|
|
|
|
FOR_UNROLL (short ii = 0; ii < NO; ++ii) {
|
|
simdgroup_store(lo[ii], sot, PV, 0, false);
|
|
|
|
sot += 8*NSG;
|
|
}
|
|
}
|
|
} else {
|
|
const short tx = tiisg%4;
|
|
const short ty = tiisg/4;
|
|
|
|
for (short cc = 0; cc < C/8; ++cc) {
|
|
s8x8_t vs;
|
|
simdgroup_load(vs, ss + 8*cc, SH, 0, false);
|
|
|
|
for (short ii = 4*sgitg; ii < DV16; ii += 4*NSG) {
|
|
device const vd4x4_t * pv4x4 = (device const vd4x4_t *) (v + ((ic + 8*cc + ty)*args.nb21));
|
|
|
|
if (DV16%4 == 0) {
|
|
{
|
|
v4x4_t tmp;
|
|
deq_v(pv4x4 + (ii + tx)/nl_v, (ii + tx)%nl_v, tmp);
|
|
sv4x4[4*ty + tx] = tmp;
|
|
}
|
|
|
|
simdgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
FOR_UNROLL (short k = 0; k < 4; ++k) {
|
|
v8x8_t mv[2];
|
|
o8x8_t lo[2];
|
|
|
|
simdgroup_load(mv[0], sv + 16*k + 0*8, 4*16, 0, false);
|
|
simdgroup_load(mv[1], sv + 16*k + 1*8, 4*16, 0, false);
|
|
simdgroup_load(lo[0], so + 8*(2*(ii + k) + 0), PV, 0, false);
|
|
simdgroup_load(lo[1], so + 8*(2*(ii + k) + 1), PV, 0, false);
|
|
|
|
simdgroup_multiply_accumulate(lo[0], vs, mv[0], lo[0]);
|
|
simdgroup_multiply_accumulate(lo[1], vs, mv[1], lo[1]);
|
|
|
|
simdgroup_store(lo[0], so + 8*(2*(ii + k) + 0), PV, 0, false);
|
|
simdgroup_store(lo[1], so + 8*(2*(ii + k) + 1), PV, 0, false);
|
|
}
|
|
} else {
|
|
if (ii + tx < DV16) {
|
|
v4x4_t tmp;
|
|
deq_v(pv4x4 + (ii + tx)/nl_v, (ii + tx)%nl_v, tmp);
|
|
sv4x4[4*ty + tx] = tmp;
|
|
}
|
|
|
|
simdgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
for (short k = 0; k < 4 && ii + k < DV16; ++k) {
|
|
v8x8_t mv[2];
|
|
o8x8_t lo[2];
|
|
|
|
simdgroup_load(mv[0], sv + 16*k + 0*8, 4*16, 0, false);
|
|
simdgroup_load(mv[1], sv + 16*k + 1*8, 4*16, 0, false);
|
|
simdgroup_load(lo[0], so + 8*(2*(ii + k) + 0), PV, 0, false);
|
|
simdgroup_load(lo[1], so + 8*(2*(ii + k) + 1), PV, 0, false);
|
|
|
|
simdgroup_multiply_accumulate(lo[0], vs, mv[0], lo[0]);
|
|
simdgroup_multiply_accumulate(lo[1], vs, mv[1], lo[1]);
|
|
|
|
simdgroup_store(lo[0], so + 8*(2*(ii + k) + 0), PV, 0, false);
|
|
simdgroup_store(lo[1], so + 8*(2*(ii + k) + 1), PV, 0, false);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
}
|
|
|
|
if (FC_flash_attn_ext_has_sinks) {
|
|
FOR_UNROLL (short jj = 0; jj < NQ; ++jj) {
|
|
const short j = jj*NSG + sgitg;
|
|
|
|
const float m = M[jj];
|
|
const float s = tiisg == 0 ? ((device const float *) sinks)[iq2] : -FLT_MAX/2;
|
|
|
|
M[jj] = simd_max(max(M[jj], s));
|
|
|
|
const float ms = exp(m - M[jj]);
|
|
const float vs = exp(s - M[jj]);
|
|
|
|
S[jj] = S[jj]*ms + simd_sum(vs);
|
|
|
|
for (short i = tiisg; i < DV4; i += NW) {
|
|
so4[j*PV4 + i] *= ms;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
for (short jj = 0; jj < NQ; ++jj) {
|
|
const short j = jj*NSG + sgitg;
|
|
if (iq1 + j >= args.ne01) {
|
|
break;
|
|
}
|
|
|
|
device float4 * dst4 = (device float4 *) dst + ((uint64_t)iq3*args.ne2*args.ne1 + iq2 + (uint64_t)(iq1 + j)*args.ne1)*DV4;
|
|
|
|
const float scale = S[jj] == 0.0 ? 0.0f : 1.0f/S[jj];
|
|
|
|
if (DV4 % NW == 0) {
|
|
FOR_UNROLL (short ii = 0; ii < DV4/NW; ++ii) {
|
|
const short i = ii*NW + tiisg;
|
|
|
|
dst4[i] = (float4) so4[j*PV4 + i]*scale;
|
|
}
|
|
} else {
|
|
for (short i = tiisg; i < DV4; i += NW) {
|
|
dst4[i] = (float4) so4[j*PV4 + i]*scale;
|
|
}
|
|
}
|
|
}
|
|
|
|
#undef NS10
|
|
#undef NS20
|
|
}
|
|
|
|
// Batched FlashAttention for prompt/prefill rows. It computes QK, applies mask,
|
|
// sinks, ALiBi/softcap options when enabled, and multiplies by V without
|
|
// materializing the full attention matrix.
|
|
template<
|
|
typename q_t,
|
|
typename q4_t,
|
|
typename q8x8_t,
|
|
typename k_t,
|
|
typename k4x4_t,
|
|
typename k8x8_t,
|
|
typename v_t,
|
|
typename v4x4_t,
|
|
typename v8x8_t,
|
|
typename qk_t,
|
|
typename qk8x8_t,
|
|
typename s_t,
|
|
typename s2_t,
|
|
typename s8x8_t,
|
|
typename o_t,
|
|
typename o4_t,
|
|
typename o8x8_t,
|
|
typename kd4x4_t,
|
|
short nl_k,
|
|
void (*deq_k)(device const kd4x4_t *, short, thread k4x4_t &),
|
|
typename vd4x4_t,
|
|
short nl_v,
|
|
void (*deq_v)(device const vd4x4_t *, short, thread v4x4_t &),
|
|
short DK,
|
|
short DV,
|
|
short Q = OP_FLASH_ATTN_EXT_NQPSG,
|
|
short C = OP_FLASH_ATTN_EXT_NCPSG>
|
|
kernel void kernel_flash_attn_ext(
|
|
constant ds4_metal_args_flash_attn_ext & 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 const char * blk,
|
|
device char * dst,
|
|
threadgroup half * shmem_f16 [[threadgroup(0)]],
|
|
uint3 tgpig[[threadgroup_position_in_grid]],
|
|
ushort tiisg[[thread_index_in_simdgroup]],
|
|
ushort sgitg[[simdgroup_index_in_threadgroup]]) {
|
|
#define FWD_TMPL q_t, q4_t, q8x8_t, k_t, k4x4_t, k8x8_t, v_t, v4x4_t, v8x8_t, qk_t, qk8x8_t, s_t, s2_t, s8x8_t, o_t, o4_t, o8x8_t, kd4x4_t, nl_k, deq_k, vd4x4_t, nl_v, deq_v, DK, DV, Q, C
|
|
#define FWD_ARGS args, q, k, v, mask, sinks, pad, blk, dst, shmem_f16, tgpig, tiisg, sgitg
|
|
switch (FC_flash_attn_ext_nsg) {
|
|
case 4: kernel_flash_attn_ext_impl<FWD_TMPL, 4>(FWD_ARGS); break;
|
|
case 8: kernel_flash_attn_ext_impl<FWD_TMPL, 8>(FWD_ARGS); break;
|
|
}
|
|
#undef FWD_TMPL
|
|
#undef FWD_ARGS
|
|
}
|
|
|
|
#define FA_NONVEC_TYPES \
|
|
half, half4, simdgroup_half8x8, \
|
|
half, half4x4, simdgroup_half8x8, \
|
|
half, half4x4, simdgroup_half8x8, \
|
|
float, simdgroup_float8x8, \
|
|
float, float2, simdgroup_float8x8, \
|
|
float, float4, simdgroup_float8x8
|
|
|
|
typedef decltype(kernel_flash_attn_ext<FA_NONVEC_TYPES, half4x4, 1, dequantize_f16, half4x4, 1, dequantize_f16, 512, 512>) flash_attn_ext_dk512_t;
|
|
typedef decltype(kernel_flash_attn_ext<FA_NONVEC_TYPES, half4x4, 1, dequantize_f16, half4x4, 1, dequantize_f16, 256, 256>) flash_attn_ext_dk256_t;
|
|
|
|
// Host-visible prefill FlashAttention variant for DS4's 512-wide F16 K/V rows.
|
|
template [[host_name("kernel_flash_attn_ext_f16_dk512_dv512")]]
|
|
kernel flash_attn_ext_dk512_t kernel_flash_attn_ext<FA_NONVEC_TYPES, half4x4, 1, dequantize_f16, half4x4, 1, dequantize_f16, 512, 512>;
|
|
|
|
// Host-visible prefill FlashAttention variant for GLM's 256-wide F16 K/V rows.
|
|
template [[host_name("kernel_flash_attn_ext_f16_dk256_dv256")]]
|
|
kernel flash_attn_ext_dk256_t kernel_flash_attn_ext<FA_NONVEC_TYPES, half4x4, 1, dequantize_f16, half4x4, 1, dequantize_f16, 256, 256>;
|
|
|
|
#undef FA_NONVEC_TYPES
|
|
|
|
constant bool FC_flash_attn_ext_vec_has_mask [[function_constant(FC_FLASH_ATTN_EXT_VEC + 0)]];
|
|
constant bool FC_flash_attn_ext_vec_has_sinks [[function_constant(FC_FLASH_ATTN_EXT_VEC + 1)]];
|
|
constant bool FC_flash_attn_ext_vec_has_bias [[function_constant(FC_FLASH_ATTN_EXT_VEC + 2)]];
|
|
constant bool FC_flash_attn_ext_vec_has_scap [[function_constant(FC_FLASH_ATTN_EXT_VEC + 3)]];
|
|
constant bool FC_flash_attn_ext_vec_has_kvpad [[function_constant(FC_FLASH_ATTN_EXT_VEC + 4)]];
|
|
constant bool FC_flash_attn_ext_vec_shared_kvpad [[function_constant(FC_FLASH_ATTN_EXT_VEC + 5)]];
|
|
constant int32_t FC_flash_attn_ext_vec_ns10 [[function_constant(FC_FLASH_ATTN_EXT_VEC + 20)]];
|
|
constant int32_t FC_flash_attn_ext_vec_ns20 [[function_constant(FC_FLASH_ATTN_EXT_VEC + 21)]];
|
|
constant int32_t FC_flash_attn_ext_vec_nsg [[function_constant(FC_FLASH_ATTN_EXT_VEC + 22)]];
|
|
constant int32_t FC_flash_attn_ext_vec_nwg [[function_constant(FC_FLASH_ATTN_EXT_VEC + 23)]];
|
|
|
|
// Decode FlashAttention for one query row. DS4 uses this in generation to scan
|
|
// raw and compressed KV cache chunks, optionally splitting long contexts across
|
|
// workgroups and writing partial softmax state for a later reduction.
|
|
template<
|
|
typename q4_t,
|
|
typename k4_t,
|
|
typename v4_t,
|
|
typename qk_t,
|
|
typename s_t,
|
|
typename s4_t,
|
|
typename o4_t,
|
|
typename kd4_t,
|
|
short nl_k,
|
|
void (*deq_k_t4)(device const kd4_t *, short, thread k4_t &),
|
|
typename vd4_t,
|
|
short nl_v,
|
|
void (*deq_v_t4)(device const vd4_t *, short, thread v4_t &),
|
|
short DK,
|
|
short DV,
|
|
short NE = 4,
|
|
short Q = OP_FLASH_ATTN_EXT_VEC_NQPSG,
|
|
short C = OP_FLASH_ATTN_EXT_VEC_NCPSG>
|
|
kernel void kernel_flash_attn_ext_vec(
|
|
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 half * shmem_f16 [[threadgroup(0)]],
|
|
uint3 tgpig[[threadgroup_position_in_grid]],
|
|
ushort tiisg[[thread_index_in_simdgroup]],
|
|
ushort sgitg[[simdgroup_index_in_threadgroup]]) {
|
|
static_assert(DK % 32 == 0, "DK must be divisible by 32");
|
|
static_assert(DV % 32 == 0, "DV must be divisible by 32");
|
|
|
|
#define NWG (FC_flash_attn_ext_vec_nwg)
|
|
#define NSG (FC_flash_attn_ext_vec_nsg)
|
|
#define NS10 (FC_flash_attn_ext_vec_ns10)
|
|
#define NS20 (FC_flash_attn_ext_vec_ns20)
|
|
|
|
const short iwg = tgpig[2]%NWG;
|
|
|
|
const ushort iq3 = tgpig[2]/NWG;
|
|
const ushort iq2 = tgpig[1];
|
|
const ushort iq1 = tgpig[0];
|
|
|
|
constexpr short DK4 = DK/4;
|
|
constexpr short DV4 = DV/4;
|
|
constexpr short PK = PAD2(DK, 128);
|
|
constexpr short PK4 = PK/4;
|
|
constexpr short PV = PAD2(DV, 128);
|
|
constexpr short PV4 = PV/4;
|
|
constexpr short NW = N_SIMDWIDTH;
|
|
constexpr short NL = NW/NE;
|
|
constexpr short SH = 4*C;
|
|
|
|
static_assert(DK4 % NL == 0, "DK4 must be divisible by NL");
|
|
static_assert(DV4 % NL == 0, "DV4 must be divisible by NL");
|
|
|
|
threadgroup q4_t * sq4 = (threadgroup q4_t *) (shmem_f16 + 0*PK);
|
|
threadgroup s_t * ss = (threadgroup s_t *) (shmem_f16 + sgitg*SH + NSG*PK);
|
|
threadgroup s4_t * ss4 = (threadgroup s4_t *) (shmem_f16 + sgitg*SH + NSG*PK);
|
|
threadgroup half * sm = (threadgroup half *) (shmem_f16 + sgitg*SH + 2*C + NSG*PK);
|
|
threadgroup o4_t * so4 = (threadgroup o4_t *) (shmem_f16 + 2*sgitg*PV + NSG*PK + NSG*SH);
|
|
|
|
so4 += tiisg;
|
|
|
|
{
|
|
q += iq1*args.nb01 + iq2*args.nb02 + iq3*args.nb03;
|
|
|
|
const short ikv2 = iq2/(args.ne02/args.ne_12_2);
|
|
const short ikv3 = iq3/(args.ne03/args.ne_12_3);
|
|
|
|
k += ikv2*args.nb12 + ikv3*args.nb13;
|
|
v += ikv2*args.nb22 + ikv3*args.nb23;
|
|
}
|
|
|
|
device const float4 * q4 = (device const float4 *) ((device const char *) q);
|
|
|
|
if (iq1 < args.ne01) {
|
|
for (short i = tiisg; i < PK4; i += NW) {
|
|
if (i < DK4) {
|
|
sq4[i] = (q4_t) q4[i];
|
|
} else {
|
|
sq4[i] = (q4_t) 0.0f;
|
|
}
|
|
}
|
|
}
|
|
|
|
for (short i = 0; i < DV4/NL; ++i) {
|
|
so4[i*NL] = (o4_t) 0.0f;
|
|
}
|
|
|
|
for (short i = tiisg; i < SH/4; i += NW) {
|
|
ss4[i] = (s4_t) 0.0f;
|
|
}
|
|
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
{
|
|
float S = 0.0f;
|
|
float M = -FLT_MAX/2;
|
|
|
|
const short tx = tiisg%NL;
|
|
const short ty = tiisg/NL;
|
|
|
|
device const half * pm = (device const half *) (mask + iq1*args.nb31 + (iq2%args.ne32)*args.nb32 + (iq3%args.ne33)*args.nb33);
|
|
|
|
float slope = 1.0f;
|
|
|
|
if (FC_flash_attn_ext_vec_has_bias) {
|
|
const short h = iq2;
|
|
|
|
const float base = h < args.n_head_log2 ? args.m0 : args.m1;
|
|
const short exph = h < args.n_head_log2 ? h + 1 : 2*(h - args.n_head_log2) + 1;
|
|
|
|
slope = pow(base, exph);
|
|
}
|
|
|
|
for (int ic0 = iwg*NSG + sgitg; ; ic0 += NWG*NSG) {
|
|
int ic = ic0*C;
|
|
if (ic >= args.ne11) {
|
|
break;
|
|
}
|
|
|
|
if (FC_flash_attn_ext_vec_has_kvpad && ic + C > args.ne11) {
|
|
k = pad;
|
|
if (FC_flash_attn_ext_vec_shared_kvpad) {
|
|
v = k;
|
|
mask = k + args.nb11*C*args.ne_12_2*args.ne_12_3 +
|
|
args.nb21*C*args.ne_12_2*args.ne_12_3;
|
|
} else {
|
|
v = k + args.nb11*C*args.ne_12_2*args.ne_12_3;
|
|
mask = v + args.nb21*C*args.ne_12_2*args.ne_12_3;
|
|
}
|
|
|
|
const short ikv2 = iq2/(args.ne02/args.ne_12_2);
|
|
const short ikv3 = iq3/(args.ne03/args.ne_12_3);
|
|
|
|
k += (ikv2 + ikv3*args.ne_12_2)*args.nb11*C;
|
|
v += (ikv2 + ikv3*args.ne_12_2)*args.nb21*C;
|
|
|
|
if (!FC_flash_attn_ext_vec_has_mask) {
|
|
if (ic + tiisg >= args.ne11) {
|
|
sm[tiisg] = -MAXHALF;
|
|
}
|
|
} else {
|
|
pm = (device const half *) (mask) +
|
|
iq1*C +
|
|
(iq2%args.ne32)*(C*args.ne31) +
|
|
(iq3%args.ne33)*(C*args.ne31*args.ne32);
|
|
}
|
|
|
|
ic = 0;
|
|
}
|
|
|
|
if (FC_flash_attn_ext_vec_has_mask) {
|
|
sm[tiisg] = pm[ic + tiisg];
|
|
}
|
|
|
|
if (simd_max(sm[tiisg]) <= -MAXHALF) {
|
|
continue;
|
|
}
|
|
|
|
{
|
|
device const k4_t * pk4 = (device const k4_t *) (k + ic*args.nb11);
|
|
threadgroup const q4_t * pq4 = sq4;
|
|
|
|
pk4 += ty*NS10/4 + tx;
|
|
pq4 += tx;
|
|
|
|
qk_t mqk[C/NE] = { [ 0 ... C/NE - 1] = 0.0f };
|
|
|
|
FOR_UNROLL (short cc = 0; cc < C/NE; ++cc) {
|
|
if (is_same<kd4_t, k4_t>::value) {
|
|
FOR_UNROLL (short ii = 0; ii < DK4/NL; ++ii) {
|
|
mqk[cc] += dot((float4) pk4[cc*NE*NS10/4 + ii*NL], (float4) pq4[ii*NL]);
|
|
}
|
|
} else {
|
|
device const kd4_t * pk = (device const kd4_t *) (k + ((ic + NE*cc + ty)*args.nb11));
|
|
|
|
k4_t mk;
|
|
|
|
FOR_UNROLL (short ii = 0; ii < DK4/NL; ++ii) {
|
|
const short i = ii*NL + tx;
|
|
|
|
deq_k_t4(pk + i/nl_k, i%nl_k, mk);
|
|
|
|
mqk[cc] += dot((float4) mk, (float4) sq4[i]);
|
|
}
|
|
}
|
|
|
|
if (NE == 1) {
|
|
mqk[cc] = simd_sum(mqk[cc]);
|
|
} else {
|
|
if (NE <= 1) {
|
|
mqk[cc] += simd_shuffle_down(mqk[cc], 16);
|
|
}
|
|
if (NE <= 2) {
|
|
mqk[cc] += simd_shuffle_down(mqk[cc], 8);
|
|
}
|
|
if (NE <= 4) {
|
|
mqk[cc] += simd_shuffle_down(mqk[cc], 4);
|
|
}
|
|
if (NE <= 8) {
|
|
mqk[cc] += simd_shuffle_down(mqk[cc], 2);
|
|
}
|
|
if (NE <= 16) {
|
|
mqk[cc] += simd_shuffle_down(mqk[cc], 1);
|
|
}
|
|
|
|
mqk[cc] = simd_shuffle(mqk[cc], NL*ty);
|
|
}
|
|
}
|
|
|
|
if (FC_flash_attn_ext_vec_has_mask &&
|
|
!FC_flash_attn_ext_vec_has_scap &&
|
|
!FC_flash_attn_ext_vec_has_bias) {
|
|
ss[NE*tx + ty] = fma(mqk[tx], args.scale, (qk_t) sm[NE*tx + ty]);
|
|
} else {
|
|
mqk[tx] *= args.scale;
|
|
|
|
if (FC_flash_attn_ext_vec_has_scap) {
|
|
mqk[tx] = args.logit_softcap*precise::tanh(mqk[tx]);
|
|
}
|
|
|
|
if (FC_flash_attn_ext_vec_has_bias) {
|
|
mqk[tx] += (qk_t) sm[NE*tx + ty]*slope;
|
|
} else {
|
|
mqk[tx] += (qk_t) sm[NE*tx + ty];
|
|
}
|
|
|
|
ss[NE*tx + ty] = mqk[tx];
|
|
}
|
|
}
|
|
|
|
simdgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
{
|
|
const float m = M;
|
|
const float s = ss[tiisg];
|
|
|
|
M = simd_max(max(M, s));
|
|
|
|
const float ms = exp(m - M);
|
|
const float vs = exp(s - M);
|
|
|
|
S = S*ms + simd_sum(vs);
|
|
|
|
ss[tiisg] = vs;
|
|
|
|
if ((DV4/NL % NW == 0) || ty == 0) {
|
|
FOR_UNROLL (short ii = 0; ii < DV4/NL; ++ii) {
|
|
so4[ii*NL] *= ms;
|
|
}
|
|
}
|
|
}
|
|
|
|
simdgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
{
|
|
o4_t lo[DV4/NL];
|
|
FOR_UNROLL (short ii = 0; ii < DV4/NL; ++ii) {
|
|
lo[ii] = 0.0f;
|
|
}
|
|
|
|
if (is_same<vd4_t, v4_t>::value) {
|
|
device const v4_t * pv4 = (device const v4_t *) (v + ic*args.nb21);
|
|
|
|
pv4 += ty*NS20/4 + tx;
|
|
|
|
const auto sst = ss + ty;
|
|
|
|
FOR_UNROLL (short cc = 0; cc < C/NE; ++cc) {
|
|
FOR_UNROLL (short ii = 0; ii < DV4/NL; ++ii) {
|
|
lo[ii] += o4_t(float4(pv4[cc*NE*NS20/4 + ii*NL])*float4(sst[cc*NE]));
|
|
}
|
|
}
|
|
} else {
|
|
FOR_UNROLL (short cc = 0; cc < C/NE; ++cc) {
|
|
device const vd4_t * pv4 = (device const vd4_t *) (v + ((ic + NE*cc + ty)*args.nb21));
|
|
|
|
FOR_UNROLL (short ii = 0; ii < DV4/NL; ++ii) {
|
|
const short i = ii*NL + tx;
|
|
|
|
v4_t mv;
|
|
deq_v_t4(pv4 + i/nl_v, i%nl_v, mv);
|
|
|
|
lo[ii] += o4_t(float4(mv)*float4(ss[NE*cc + ty]));
|
|
}
|
|
}
|
|
}
|
|
|
|
FOR_UNROLL (short ii = 0; ii < DV4/NL; ++ii) {
|
|
if (NE > 1) {
|
|
lo[ii][0] += simd_shuffle_down(lo[ii][0], 16);
|
|
lo[ii][1] += simd_shuffle_down(lo[ii][1], 16);
|
|
lo[ii][2] += simd_shuffle_down(lo[ii][2], 16);
|
|
lo[ii][3] += simd_shuffle_down(lo[ii][3], 16);
|
|
}
|
|
|
|
if (NE > 2) {
|
|
lo[ii][0] += simd_shuffle_down(lo[ii][0], 8);
|
|
lo[ii][1] += simd_shuffle_down(lo[ii][1], 8);
|
|
lo[ii][2] += simd_shuffle_down(lo[ii][2], 8);
|
|
lo[ii][3] += simd_shuffle_down(lo[ii][3], 8);
|
|
}
|
|
|
|
if (NE > 4) {
|
|
lo[ii][0] += simd_shuffle_down(lo[ii][0], 4);
|
|
lo[ii][1] += simd_shuffle_down(lo[ii][1], 4);
|
|
lo[ii][2] += simd_shuffle_down(lo[ii][2], 4);
|
|
lo[ii][3] += simd_shuffle_down(lo[ii][3], 4);
|
|
}
|
|
|
|
if (NE > 8) {
|
|
lo[ii][0] += simd_shuffle_down(lo[ii][0], 2);
|
|
lo[ii][1] += simd_shuffle_down(lo[ii][1], 2);
|
|
lo[ii][2] += simd_shuffle_down(lo[ii][2], 2);
|
|
lo[ii][3] += simd_shuffle_down(lo[ii][3], 2);
|
|
}
|
|
|
|
if (NE > 16) {
|
|
lo[ii][0] += simd_shuffle_down(lo[ii][0], 1);
|
|
lo[ii][1] += simd_shuffle_down(lo[ii][1], 1);
|
|
lo[ii][2] += simd_shuffle_down(lo[ii][2], 1);
|
|
lo[ii][3] += simd_shuffle_down(lo[ii][3], 1);
|
|
}
|
|
}
|
|
|
|
if ((DV4/NL % NW == 0) || ty == 0) {
|
|
FOR_UNROLL (short ii = 0; ii < DV4/NL; ++ii) {
|
|
so4[ii*NL] += lo[ii];
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
if (FC_flash_attn_ext_vec_has_sinks && sgitg == 0 && iwg == 0) {
|
|
const float m = M;
|
|
const float s = tiisg == 0 ? ((device const float *) sinks)[iq2] : -FLT_MAX/2;
|
|
|
|
M = simd_max(max(M, s));
|
|
|
|
const float ms = exp(m - M);
|
|
const float vs = exp(s - M);
|
|
|
|
S = S*ms + simd_sum(vs);
|
|
|
|
if ((DV4/NL % NW == 0) || ty == 0) {
|
|
FOR_UNROLL (short ii = 0; ii < DV4/NL; ++ii) {
|
|
so4[ii*NL] *= ms;
|
|
}
|
|
}
|
|
}
|
|
|
|
if (tiisg == 0) {
|
|
ss[0] = (s_t) S;
|
|
ss[1] = (s_t) M;
|
|
}
|
|
}
|
|
|
|
so4 -= tiisg;
|
|
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
for (short r = NSG/2; r > 0; r >>= 1) {
|
|
if (sgitg < r) {
|
|
const float S0 = ss[ 0];
|
|
const float S1 = ss[r*(SH/2) + 0];
|
|
|
|
const float M0 = ss[ 1];
|
|
const float M1 = ss[r*(SH/2) + 1];
|
|
|
|
const float M = max(M0, M1);
|
|
|
|
const float ms0 = exp(M0 - M);
|
|
const float ms1 = exp(M1 - M);
|
|
|
|
const float S = S0*ms0 + S1*ms1;
|
|
|
|
if (tiisg == 0) {
|
|
ss[0] = S;
|
|
ss[1] = M;
|
|
}
|
|
|
|
for (short i = tiisg; i < DV4; i += NW) {
|
|
so4[i] = so4[i]*ms0 + so4[i + r*PV4]*ms1;
|
|
}
|
|
}
|
|
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
}
|
|
|
|
if (sgitg == 0) {
|
|
const int64_t nrows = args.ne3*args.ne2*args.ne1;
|
|
const int64_t rid = iq3*args.ne2*args.ne1 + iq2 + iq1*args.ne1;
|
|
|
|
device float4 * dst4 = (device float4 *) dst;
|
|
device float * dst1 = (device float *) dst + nrows*DV*NWG;
|
|
|
|
const float S = NWG == 1 ? (ss[0] == 0.0f ? 0.0f : 1.0f/ss[0]) : 1.0f;
|
|
|
|
for (short i = tiisg; i < DV4; i += NW) {
|
|
dst4[rid*DV4*NWG + NWG*i + iwg] = (float4) so4[i]*S;
|
|
}
|
|
|
|
if (NWG > 1) {
|
|
if (tiisg == 0) {
|
|
dst1[rid*(2*NWG) + 2*iwg + 0] = ss[0];
|
|
dst1[rid*(2*NWG) + 2*iwg + 1] = ss[1];
|
|
}
|
|
}
|
|
}
|
|
|
|
#undef NWG
|
|
#undef NSG
|
|
#undef NS10
|
|
#undef NS20
|
|
}
|
|
|
|
#define FA_TYPES \
|
|
half4, \
|
|
half4, \
|
|
half4, \
|
|
float, \
|
|
float, float4, \
|
|
float4
|
|
|
|
#define FA_TYPES_F32 \
|
|
half4, \
|
|
float4, \
|
|
float4, \
|
|
float, \
|
|
float, float4, \
|
|
float4
|
|
|
|
typedef decltype(kernel_flash_attn_ext_vec<FA_TYPES, half4, 1, dequantize_f16_t4, half4, 1, dequantize_f16_t4, 128, 128, 4>) flash_attn_ext_vec_t;
|
|
|
|
// Host-visible decode FlashAttention variant for DS4's 512-wide F16 K/V rows.
|
|
template [[host_name("kernel_flash_attn_ext_vec_f16_dk512_dv512")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, half4, 1, dequantize_f16_t4, half4, 1, dequantize_f16_t4, 512, 512, 1>;
|
|
|
|
#undef FA_TYPES
|
|
#undef FA_TYPES_F32
|
|
|
|
constant int32_t FC_flash_attn_ext_vec_reduce_DV [[function_constant(FC_FLASH_ATTN_EXT_VEC_REDUCE + 0)]];
|
|
constant int32_t FC_flash_attn_ext_vec_reduce_NWG [[function_constant(FC_FLASH_ATTN_EXT_VEC_REDUCE + 1)]];
|
|
|
|
// 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(
|
|
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)
|
|
|
|
const uint64_t rid = tgpig;
|
|
|
|
const short iwg = tiisg;
|
|
|
|
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];
|
|
|
|
const float m = simd_max(M);
|
|
const float ms = exp(M - m);
|
|
|
|
S = simd_sum(S*ms);
|
|
S = S == 0.0f ? 0.0f : 1.0f/S;
|
|
|
|
const short DV4 = DV/4;
|
|
|
|
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);
|
|
|
|
if (iwg == 0) {
|
|
dst4[i] = v*S;
|
|
}
|
|
}
|
|
|
|
#undef NWG
|
|
#undef DV
|
|
}
|