wehub-resource-sync
242 lines
7.5 KiB
Metal
242 lines
7.5 KiB
Metal
// DS4 Metal softmax kernel used by the compressor pooling compatibility path.
|
|
// The single-compressed-row path is intentionally left as soft_max -> mul ->
|
|
// sum_rows instead of using the fused dsv4_softmax_pool kernel.
|
|
|
|
struct ds4_metal_args_soft_max {
|
|
int32_t ne00;
|
|
int32_t ne01;
|
|
int32_t ne02;
|
|
uint64_t nb01;
|
|
uint64_t nb02;
|
|
uint64_t nb03;
|
|
int32_t ne11;
|
|
int32_t ne12;
|
|
int32_t ne13;
|
|
uint64_t nb11;
|
|
uint64_t nb12;
|
|
uint64_t nb13;
|
|
uint64_t nb1;
|
|
uint64_t nb2;
|
|
uint64_t nb3;
|
|
float scale;
|
|
float max_bias;
|
|
float m0;
|
|
float m1;
|
|
int32_t n_head_log2;
|
|
};
|
|
|
|
// Row softmax for score matrices. DS4 uses it in the literal one-compressor-row
|
|
// path where preserving the original graph operation boundary avoids drift.
|
|
template<typename T>
|
|
kernel void kernel_soft_max(
|
|
constant ds4_metal_args_soft_max & args,
|
|
device const char * src0,
|
|
device const char * src1,
|
|
device const char * src2,
|
|
device char * dst,
|
|
threadgroup float * buf [[threadgroup(0)]],
|
|
uint3 tgpig[[threadgroup_position_in_grid]],
|
|
uint3 tpitg[[thread_position_in_threadgroup]],
|
|
uint sgitg[[simdgroup_index_in_threadgroup]],
|
|
uint tiisg[[thread_index_in_simdgroup]],
|
|
uint3 tptg[[threads_per_threadgroup]]) {
|
|
const int32_t i03 = tgpig.z;
|
|
const int32_t i02 = tgpig.y;
|
|
const int32_t i01 = tgpig.x;
|
|
|
|
const int32_t i13 = i03%args.ne13;
|
|
const int32_t i12 = i02%args.ne12;
|
|
const int32_t i11 = i01;
|
|
|
|
device const float * psrc0 = (device const float *) (src0 + i01*args.nb01 + i02*args.nb02 + i03*args.nb03);
|
|
device const T * pmask = src1 != src0 ? (device const T * ) (src1 + i11*args.nb11 + i12*args.nb12 + i13*args.nb13) : nullptr;
|
|
device const float * psrc2 = src2 != src0 ? (device const float *) (src2) : nullptr;
|
|
device float * pdst = (device float *) (dst + i01*args.nb1 + i02*args.nb2 + i03*args.nb3);
|
|
|
|
float slope = 1.0f;
|
|
|
|
if (args.max_bias > 0.0f) {
|
|
const int32_t h = i02;
|
|
|
|
const float base = h < args.n_head_log2 ? args.m0 : args.m1;
|
|
const int exp = h < args.n_head_log2 ? h + 1 : 2*(h - args.n_head_log2) + 1;
|
|
|
|
slope = pow(base, exp);
|
|
}
|
|
|
|
float lmax = psrc2 ? psrc2[i02] : -INFINITY;
|
|
|
|
for (int i00 = tpitg.x; i00 < args.ne00; i00 += tptg.x) {
|
|
lmax = MAX(lmax, psrc0[i00]*args.scale + (pmask ? slope*pmask[i00] : 0.0f));
|
|
}
|
|
|
|
float max_val = simd_max(lmax);
|
|
if (tptg.x > N_SIMDWIDTH) {
|
|
if (sgitg == 0) {
|
|
buf[tiisg] = -INFINITY;
|
|
}
|
|
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
if (tiisg == 0) {
|
|
buf[sgitg] = max_val;
|
|
}
|
|
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
max_val = buf[tiisg];
|
|
max_val = simd_max(max_val);
|
|
}
|
|
|
|
float lsum = 0.0f;
|
|
for (int i00 = tpitg.x; i00 < args.ne00; i00 += tptg.x) {
|
|
const float exp_psrc0 = exp((psrc0[i00]*args.scale + (pmask ? slope*pmask[i00] : 0.0f)) - max_val);
|
|
lsum += exp_psrc0;
|
|
pdst[i00] = exp_psrc0;
|
|
}
|
|
|
|
threadgroup_barrier(mem_flags::mem_none);
|
|
|
|
float sum = simd_sum(lsum);
|
|
|
|
if (tptg.x > N_SIMDWIDTH) {
|
|
if (sgitg == 0) {
|
|
buf[tiisg] = 0.0f;
|
|
}
|
|
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
if (tiisg == 0) {
|
|
buf[sgitg] = sum;
|
|
}
|
|
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
sum = buf[tiisg];
|
|
sum = simd_sum(sum);
|
|
}
|
|
|
|
if (psrc2) {
|
|
sum += exp(psrc2[i02] - max_val);
|
|
}
|
|
|
|
const float inv_sum = 1.0f/sum;
|
|
|
|
for (int i00 = tpitg.x; i00 < args.ne00; i00 += tptg.x) {
|
|
pdst[i00] *= inv_sum;
|
|
}
|
|
}
|
|
|
|
// Vectorized float4 row softmax for contiguous score rows whose length is a
|
|
// multiple of four; used by the same DS4 compressor/indexer graph path.
|
|
template<typename T>
|
|
kernel void kernel_soft_max_4(
|
|
constant ds4_metal_args_soft_max & args,
|
|
device const char * src0,
|
|
device const char * src1,
|
|
device const char * src2,
|
|
device char * dst,
|
|
threadgroup float * buf [[threadgroup(0)]],
|
|
uint3 tgpig[[threadgroup_position_in_grid]],
|
|
uint3 tpitg[[thread_position_in_threadgroup]],
|
|
uint sgitg[[simdgroup_index_in_threadgroup]],
|
|
uint tiisg[[thread_index_in_simdgroup]],
|
|
uint3 tptg[[threads_per_threadgroup]]) {
|
|
const int32_t i03 = tgpig.z;
|
|
const int32_t i02 = tgpig.y;
|
|
const int32_t i01 = tgpig.x;
|
|
|
|
const int32_t i13 = i03%args.ne13;
|
|
const int32_t i12 = i02%args.ne12;
|
|
const int32_t i11 = i01;
|
|
|
|
device const float4 * psrc4 = (device const float4 *) (src0 + i01*args.nb01 + i02*args.nb02 + i03*args.nb03);
|
|
device const T * pmask = src1 != src0 ? (device const T * ) (src1 + i11*args.nb11 + i12*args.nb12 + i13*args.nb13) : nullptr;
|
|
device const float * psrc2 = src2 != src0 ? (device const float * ) (src2) : nullptr;
|
|
device float4 * pdst4 = (device float4 *) (dst + i01*args.nb1 + i02*args.nb2 + i03*args.nb3);
|
|
|
|
float slope = 1.0f;
|
|
|
|
if (args.max_bias > 0.0f) {
|
|
const int32_t h = i02;
|
|
|
|
const float base = h < args.n_head_log2 ? args.m0 : args.m1;
|
|
const int exp = h < args.n_head_log2 ? h + 1 : 2*(h - args.n_head_log2) + 1;
|
|
|
|
slope = pow(base, exp);
|
|
}
|
|
|
|
float4 lmax4 = psrc2 ? psrc2[i02] : -INFINITY;
|
|
|
|
for (int i00 = tpitg.x; i00 < args.ne00/4; i00 += tptg.x) {
|
|
lmax4 = fmax(lmax4, psrc4[i00]*args.scale + (float4)((pmask ? slope*pmask[i00] : 0.0f)));
|
|
}
|
|
|
|
const float lmax = MAX(MAX(lmax4[0], lmax4[1]), MAX(lmax4[2], lmax4[3]));
|
|
|
|
float max_val = simd_max(lmax);
|
|
if (tptg.x > N_SIMDWIDTH) {
|
|
if (sgitg == 0) {
|
|
buf[tiisg] = -INFINITY;
|
|
}
|
|
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
if (tiisg == 0) {
|
|
buf[sgitg] = max_val;
|
|
}
|
|
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
max_val = buf[tiisg];
|
|
max_val = simd_max(max_val);
|
|
}
|
|
|
|
float4 lsum4 = 0.0f;
|
|
for (int i00 = tpitg.x; i00 < args.ne00/4; i00 += tptg.x) {
|
|
const float4 exp_psrc4 = exp((psrc4[i00]*args.scale + (float4)((pmask ? slope*pmask[i00] : 0.0f))) - max_val);
|
|
lsum4 += exp_psrc4;
|
|
pdst4[i00] = exp_psrc4;
|
|
}
|
|
|
|
const float lsum = lsum4[0] + lsum4[1] + lsum4[2] + lsum4[3];
|
|
|
|
threadgroup_barrier(mem_flags::mem_none);
|
|
|
|
float sum = simd_sum(lsum);
|
|
|
|
if (tptg.x > N_SIMDWIDTH) {
|
|
if (sgitg == 0) {
|
|
buf[tiisg] = 0.0f;
|
|
}
|
|
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
if (tiisg == 0) {
|
|
buf[sgitg] = sum;
|
|
}
|
|
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
sum = buf[tiisg];
|
|
sum = simd_sum(sum);
|
|
}
|
|
|
|
if (psrc2) {
|
|
sum += exp(psrc2[i02] - max_val);
|
|
}
|
|
|
|
const float inv_sum = 1.0f/sum;
|
|
|
|
for (int i00 = tpitg.x; i00 < args.ne00/4; i00 += tptg.x) {
|
|
pdst4[i00] *= inv_sum;
|
|
}
|
|
}
|
|
|
|
typedef decltype(kernel_soft_max<float>) kernel_soft_max_t;
|
|
typedef decltype(kernel_soft_max_4<float4>) kernel_soft_max_4_t;
|
|
|
|
// Host-visible F32 softmax variants used by compressor pooling.
|
|
template [[host_name("kernel_soft_max_f32")]] kernel kernel_soft_max_t kernel_soft_max<float>;
|
|
template [[host_name("kernel_soft_max_f32_4")]] kernel kernel_soft_max_4_t kernel_soft_max_4<float4>;
|