9133 lines
393 KiB
Metal
9133 lines
393 KiB
Metal
// DS4 Metal routed-MoE matvec kernels.
|
|
|
|
#ifndef QK_K
|
|
#define QK_K 256
|
|
#endif
|
|
#define QK_MXFP4 32
|
|
#define N_R0_Q2_K 4
|
|
#define N_R0_GLM_Q2_PAIR2_K 1
|
|
#define N_R0_Q4_K 2
|
|
#define N_R0_Q8_K 2
|
|
#define N_R0_GLM_Q4_PAIR2_K 1
|
|
#define N_R0_GLM_Q4_PAIR_K 4
|
|
#define N_R0_Q5_PAIR_K 4
|
|
#define N_R0_Q5_K 4
|
|
#define N_R0_Q6_K 2
|
|
#define N_R0_IQ2_XXS 4
|
|
#define N_R0_MXFP4 2
|
|
|
|
static constant float ds4_metal_mxfp4_values[16] = {
|
|
0.0f, 0.5f, 1.0f, 1.5f, 2.0f, 3.0f, 4.0f, 6.0f,
|
|
-0.0f, -0.5f, -1.0f, -1.5f, -2.0f, -3.0f, -4.0f, -6.0f,
|
|
};
|
|
|
|
static inline float ds4_metal_e8m0_to_f32(uchar e) {
|
|
const uint bits = e == 0 ? 0x00400000u : (uint)e << 23;
|
|
return as_type<float>(bits);
|
|
}
|
|
|
|
// BEGIN GENERATED MXFP4 HALF LUT
|
|
// Generated by metal/generate_mxfp4_half_lut.py; do not edit by hand.
|
|
// Index = (E8M0 exponent byte << 4) | E2M1 code.
|
|
static constant ushort ds4_metal_mxfp4_half_lut[256][16] = {
|
|
{ 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u },
|
|
{ 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u },
|
|
{ 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u },
|
|
{ 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u },
|
|
{ 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u },
|
|
{ 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u },
|
|
{ 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u },
|
|
{ 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u },
|
|
{ 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u },
|
|
{ 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u },
|
|
{ 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u },
|
|
{ 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u },
|
|
{ 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u },
|
|
{ 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u },
|
|
{ 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u },
|
|
{ 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u },
|
|
{ 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u },
|
|
{ 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u },
|
|
{ 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u },
|
|
{ 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u },
|
|
{ 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u },
|
|
{ 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u },
|
|
{ 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u },
|
|
{ 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u },
|
|
{ 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u },
|
|
{ 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u },
|
|
{ 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u },
|
|
{ 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u },
|
|
{ 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u },
|
|
{ 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u },
|
|
{ 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u },
|
|
{ 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u },
|
|
{ 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u },
|
|
{ 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u },
|
|
{ 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u },
|
|
{ 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u },
|
|
{ 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u },
|
|
{ 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u },
|
|
{ 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u },
|
|
{ 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u },
|
|
{ 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u },
|
|
{ 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u },
|
|
{ 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u },
|
|
{ 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u },
|
|
{ 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u },
|
|
{ 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u },
|
|
{ 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u },
|
|
{ 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u },
|
|
{ 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u },
|
|
{ 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u },
|
|
{ 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u },
|
|
{ 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u },
|
|
{ 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u },
|
|
{ 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u },
|
|
{ 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u },
|
|
{ 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u },
|
|
{ 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u },
|
|
{ 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u },
|
|
{ 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u },
|
|
{ 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u },
|
|
{ 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u },
|
|
{ 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u },
|
|
{ 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u },
|
|
{ 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u },
|
|
{ 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u },
|
|
{ 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u },
|
|
{ 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u },
|
|
{ 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u },
|
|
{ 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u },
|
|
{ 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u },
|
|
{ 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u },
|
|
{ 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u },
|
|
{ 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u },
|
|
{ 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u },
|
|
{ 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u },
|
|
{ 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u },
|
|
{ 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u },
|
|
{ 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u },
|
|
{ 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u },
|
|
{ 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u },
|
|
{ 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u },
|
|
{ 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u },
|
|
{ 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u },
|
|
{ 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u },
|
|
{ 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u },
|
|
{ 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u },
|
|
{ 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u },
|
|
{ 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u },
|
|
{ 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u },
|
|
{ 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u },
|
|
{ 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u },
|
|
{ 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u },
|
|
{ 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u },
|
|
{ 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u },
|
|
{ 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u },
|
|
{ 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u },
|
|
{ 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u },
|
|
{ 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u },
|
|
{ 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u },
|
|
{ 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u },
|
|
{ 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0001u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8001u },
|
|
{ 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0000u, 0x0001u, 0x0001u, 0x0002u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8000u, 0x8001u, 0x8001u, 0x8002u },
|
|
{ 0x0000u, 0x0000u, 0x0000u, 0x0001u, 0x0001u, 0x0002u, 0x0002u, 0x0003u, 0x8000u, 0x8000u, 0x8000u, 0x8001u, 0x8001u, 0x8002u, 0x8002u, 0x8003u },
|
|
{ 0x0000u, 0x0000u, 0x0001u, 0x0002u, 0x0002u, 0x0003u, 0x0004u, 0x0006u, 0x8000u, 0x8000u, 0x8001u, 0x8002u, 0x8002u, 0x8003u, 0x8004u, 0x8006u },
|
|
{ 0x0000u, 0x0001u, 0x0002u, 0x0003u, 0x0004u, 0x0006u, 0x0008u, 0x000cu, 0x8000u, 0x8001u, 0x8002u, 0x8003u, 0x8004u, 0x8006u, 0x8008u, 0x800cu },
|
|
{ 0x0000u, 0x0002u, 0x0004u, 0x0006u, 0x0008u, 0x000cu, 0x0010u, 0x0018u, 0x8000u, 0x8002u, 0x8004u, 0x8006u, 0x8008u, 0x800cu, 0x8010u, 0x8018u },
|
|
{ 0x0000u, 0x0004u, 0x0008u, 0x000cu, 0x0010u, 0x0018u, 0x0020u, 0x0030u, 0x8000u, 0x8004u, 0x8008u, 0x800cu, 0x8010u, 0x8018u, 0x8020u, 0x8030u },
|
|
{ 0x0000u, 0x0008u, 0x0010u, 0x0018u, 0x0020u, 0x0030u, 0x0040u, 0x0060u, 0x8000u, 0x8008u, 0x8010u, 0x8018u, 0x8020u, 0x8030u, 0x8040u, 0x8060u },
|
|
{ 0x0000u, 0x0010u, 0x0020u, 0x0030u, 0x0040u, 0x0060u, 0x0080u, 0x00c0u, 0x8000u, 0x8010u, 0x8020u, 0x8030u, 0x8040u, 0x8060u, 0x8080u, 0x80c0u },
|
|
{ 0x0000u, 0x0020u, 0x0040u, 0x0060u, 0x0080u, 0x00c0u, 0x0100u, 0x0180u, 0x8000u, 0x8020u, 0x8040u, 0x8060u, 0x8080u, 0x80c0u, 0x8100u, 0x8180u },
|
|
{ 0x0000u, 0x0040u, 0x0080u, 0x00c0u, 0x0100u, 0x0180u, 0x0200u, 0x0300u, 0x8000u, 0x8040u, 0x8080u, 0x80c0u, 0x8100u, 0x8180u, 0x8200u, 0x8300u },
|
|
{ 0x0000u, 0x0080u, 0x0100u, 0x0180u, 0x0200u, 0x0300u, 0x0400u, 0x0600u, 0x8000u, 0x8080u, 0x8100u, 0x8180u, 0x8200u, 0x8300u, 0x8400u, 0x8600u },
|
|
{ 0x0000u, 0x0100u, 0x0200u, 0x0300u, 0x0400u, 0x0600u, 0x0800u, 0x0a00u, 0x8000u, 0x8100u, 0x8200u, 0x8300u, 0x8400u, 0x8600u, 0x8800u, 0x8a00u },
|
|
{ 0x0000u, 0x0200u, 0x0400u, 0x0600u, 0x0800u, 0x0a00u, 0x0c00u, 0x0e00u, 0x8000u, 0x8200u, 0x8400u, 0x8600u, 0x8800u, 0x8a00u, 0x8c00u, 0x8e00u },
|
|
{ 0x0000u, 0x0400u, 0x0800u, 0x0a00u, 0x0c00u, 0x0e00u, 0x1000u, 0x1200u, 0x8000u, 0x8400u, 0x8800u, 0x8a00u, 0x8c00u, 0x8e00u, 0x9000u, 0x9200u },
|
|
{ 0x0000u, 0x0800u, 0x0c00u, 0x0e00u, 0x1000u, 0x1200u, 0x1400u, 0x1600u, 0x8000u, 0x8800u, 0x8c00u, 0x8e00u, 0x9000u, 0x9200u, 0x9400u, 0x9600u },
|
|
{ 0x0000u, 0x0c00u, 0x1000u, 0x1200u, 0x1400u, 0x1600u, 0x1800u, 0x1a00u, 0x8000u, 0x8c00u, 0x9000u, 0x9200u, 0x9400u, 0x9600u, 0x9800u, 0x9a00u },
|
|
{ 0x0000u, 0x1000u, 0x1400u, 0x1600u, 0x1800u, 0x1a00u, 0x1c00u, 0x1e00u, 0x8000u, 0x9000u, 0x9400u, 0x9600u, 0x9800u, 0x9a00u, 0x9c00u, 0x9e00u },
|
|
{ 0x0000u, 0x1400u, 0x1800u, 0x1a00u, 0x1c00u, 0x1e00u, 0x2000u, 0x2200u, 0x8000u, 0x9400u, 0x9800u, 0x9a00u, 0x9c00u, 0x9e00u, 0xa000u, 0xa200u },
|
|
{ 0x0000u, 0x1800u, 0x1c00u, 0x1e00u, 0x2000u, 0x2200u, 0x2400u, 0x2600u, 0x8000u, 0x9800u, 0x9c00u, 0x9e00u, 0xa000u, 0xa200u, 0xa400u, 0xa600u },
|
|
{ 0x0000u, 0x1c00u, 0x2000u, 0x2200u, 0x2400u, 0x2600u, 0x2800u, 0x2a00u, 0x8000u, 0x9c00u, 0xa000u, 0xa200u, 0xa400u, 0xa600u, 0xa800u, 0xaa00u },
|
|
{ 0x0000u, 0x2000u, 0x2400u, 0x2600u, 0x2800u, 0x2a00u, 0x2c00u, 0x2e00u, 0x8000u, 0xa000u, 0xa400u, 0xa600u, 0xa800u, 0xaa00u, 0xac00u, 0xae00u },
|
|
{ 0x0000u, 0x2400u, 0x2800u, 0x2a00u, 0x2c00u, 0x2e00u, 0x3000u, 0x3200u, 0x8000u, 0xa400u, 0xa800u, 0xaa00u, 0xac00u, 0xae00u, 0xb000u, 0xb200u },
|
|
{ 0x0000u, 0x2800u, 0x2c00u, 0x2e00u, 0x3000u, 0x3200u, 0x3400u, 0x3600u, 0x8000u, 0xa800u, 0xac00u, 0xae00u, 0xb000u, 0xb200u, 0xb400u, 0xb600u },
|
|
{ 0x0000u, 0x2c00u, 0x3000u, 0x3200u, 0x3400u, 0x3600u, 0x3800u, 0x3a00u, 0x8000u, 0xac00u, 0xb000u, 0xb200u, 0xb400u, 0xb600u, 0xb800u, 0xba00u },
|
|
{ 0x0000u, 0x3000u, 0x3400u, 0x3600u, 0x3800u, 0x3a00u, 0x3c00u, 0x3e00u, 0x8000u, 0xb000u, 0xb400u, 0xb600u, 0xb800u, 0xba00u, 0xbc00u, 0xbe00u },
|
|
{ 0x0000u, 0x3400u, 0x3800u, 0x3a00u, 0x3c00u, 0x3e00u, 0x4000u, 0x4200u, 0x8000u, 0xb400u, 0xb800u, 0xba00u, 0xbc00u, 0xbe00u, 0xc000u, 0xc200u },
|
|
{ 0x0000u, 0x3800u, 0x3c00u, 0x3e00u, 0x4000u, 0x4200u, 0x4400u, 0x4600u, 0x8000u, 0xb800u, 0xbc00u, 0xbe00u, 0xc000u, 0xc200u, 0xc400u, 0xc600u },
|
|
{ 0x0000u, 0x3c00u, 0x4000u, 0x4200u, 0x4400u, 0x4600u, 0x4800u, 0x4a00u, 0x8000u, 0xbc00u, 0xc000u, 0xc200u, 0xc400u, 0xc600u, 0xc800u, 0xca00u },
|
|
{ 0x0000u, 0x4000u, 0x4400u, 0x4600u, 0x4800u, 0x4a00u, 0x4c00u, 0x4e00u, 0x8000u, 0xc000u, 0xc400u, 0xc600u, 0xc800u, 0xca00u, 0xcc00u, 0xce00u },
|
|
{ 0x0000u, 0x4400u, 0x4800u, 0x4a00u, 0x4c00u, 0x4e00u, 0x5000u, 0x5200u, 0x8000u, 0xc400u, 0xc800u, 0xca00u, 0xcc00u, 0xce00u, 0xd000u, 0xd200u },
|
|
{ 0x0000u, 0x4800u, 0x4c00u, 0x4e00u, 0x5000u, 0x5200u, 0x5400u, 0x5600u, 0x8000u, 0xc800u, 0xcc00u, 0xce00u, 0xd000u, 0xd200u, 0xd400u, 0xd600u },
|
|
{ 0x0000u, 0x4c00u, 0x5000u, 0x5200u, 0x5400u, 0x5600u, 0x5800u, 0x5a00u, 0x8000u, 0xcc00u, 0xd000u, 0xd200u, 0xd400u, 0xd600u, 0xd800u, 0xda00u },
|
|
{ 0x0000u, 0x5000u, 0x5400u, 0x5600u, 0x5800u, 0x5a00u, 0x5c00u, 0x5e00u, 0x8000u, 0xd000u, 0xd400u, 0xd600u, 0xd800u, 0xda00u, 0xdc00u, 0xde00u },
|
|
{ 0x0000u, 0x5400u, 0x5800u, 0x5a00u, 0x5c00u, 0x5e00u, 0x6000u, 0x6200u, 0x8000u, 0xd400u, 0xd800u, 0xda00u, 0xdc00u, 0xde00u, 0xe000u, 0xe200u },
|
|
{ 0x0000u, 0x5800u, 0x5c00u, 0x5e00u, 0x6000u, 0x6200u, 0x6400u, 0x6600u, 0x8000u, 0xd800u, 0xdc00u, 0xde00u, 0xe000u, 0xe200u, 0xe400u, 0xe600u },
|
|
{ 0x0000u, 0x5c00u, 0x6000u, 0x6200u, 0x6400u, 0x6600u, 0x6800u, 0x6a00u, 0x8000u, 0xdc00u, 0xe000u, 0xe200u, 0xe400u, 0xe600u, 0xe800u, 0xea00u },
|
|
{ 0x0000u, 0x6000u, 0x6400u, 0x6600u, 0x6800u, 0x6a00u, 0x6c00u, 0x6e00u, 0x8000u, 0xe000u, 0xe400u, 0xe600u, 0xe800u, 0xea00u, 0xec00u, 0xee00u },
|
|
{ 0x0000u, 0x6400u, 0x6800u, 0x6a00u, 0x6c00u, 0x6e00u, 0x7000u, 0x7200u, 0x8000u, 0xe400u, 0xe800u, 0xea00u, 0xec00u, 0xee00u, 0xf000u, 0xf200u },
|
|
{ 0x0000u, 0x6800u, 0x6c00u, 0x6e00u, 0x7000u, 0x7200u, 0x7400u, 0x7600u, 0x8000u, 0xe800u, 0xec00u, 0xee00u, 0xf000u, 0xf200u, 0xf400u, 0xf600u },
|
|
{ 0x0000u, 0x6c00u, 0x7000u, 0x7200u, 0x7400u, 0x7600u, 0x7800u, 0x7a00u, 0x8000u, 0xec00u, 0xf000u, 0xf200u, 0xf400u, 0xf600u, 0xf800u, 0xfa00u },
|
|
{ 0x0000u, 0x7000u, 0x7400u, 0x7600u, 0x7800u, 0x7a00u, 0x7c00u, 0x7c00u, 0x8000u, 0xf000u, 0xf400u, 0xf600u, 0xf800u, 0xfa00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7400u, 0x7800u, 0x7a00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xf400u, 0xf800u, 0xfa00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7800u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xf800u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x0000u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x8000u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
{ 0x7e00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7c00u, 0x7e00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u, 0xfc00u },
|
|
};
|
|
// END GENERATED MXFP4 HALF LUT
|
|
|
|
static constant uchar ds4_metal_kmask_iq2xs[8] = {
|
|
1, 2, 4, 8, 16, 32, 64, 128
|
|
};
|
|
|
|
static constant uchar ds4_metal_ksigns_iq2xs[128] = {
|
|
0, 129, 130, 3, 132, 5, 6, 135, 136, 9, 10, 139, 12, 141, 142, 15,
|
|
144, 17, 18, 147, 20, 149, 150, 23, 24, 153, 154, 27, 156, 29, 30, 159,
|
|
160, 33, 34, 163, 36, 165, 166, 39, 40, 169, 170, 43, 172, 45, 46, 175,
|
|
48, 177, 178, 51, 180, 53, 54, 183, 184, 57, 58, 187, 60, 189, 190, 63,
|
|
192, 65, 66, 195, 68, 197, 198, 71, 72, 201, 202, 75, 204, 77, 78, 207,
|
|
80, 209, 210, 83, 212, 85, 86, 215, 216, 89, 90, 219, 92, 221, 222, 95,
|
|
96, 225, 226, 99, 228, 101, 102, 231, 232, 105, 106, 235, 108, 237, 238, 111,
|
|
240, 113, 114, 243, 116, 245, 246, 119, 120, 249, 250, 123, 252, 125, 126, 255,
|
|
};
|
|
|
|
static constant ulong ds4_metal_iq2xxs_grid[256] = {
|
|
0x0808080808080808, 0x080808080808082b, 0x0808080808081919, 0x0808080808082b08,
|
|
0x0808080808082b2b, 0x0808080808190819, 0x0808080808191908, 0x08080808082b0808,
|
|
0x08080808082b082b, 0x08080808082b2b08, 0x08080808082b2b2b, 0x0808080819080819,
|
|
0x0808080819081908, 0x0808080819190808, 0x0808080819192b08, 0x08080808192b0819,
|
|
0x08080808192b1908, 0x080808082b080808, 0x080808082b08082b, 0x080808082b082b2b,
|
|
0x080808082b2b082b, 0x0808081908080819, 0x0808081908081908, 0x0808081908190808,
|
|
0x0808081908191919, 0x0808081919080808, 0x080808192b081908, 0x080808192b192b08,
|
|
0x0808082b08080808, 0x0808082b0808082b, 0x0808082b082b082b, 0x0808082b2b08082b,
|
|
0x0808190808080819, 0x0808190808081908, 0x0808190808190808, 0x08081908082b0819,
|
|
0x08081908082b1908, 0x0808190819080808, 0x080819081908082b, 0x0808190819082b08,
|
|
0x08081908192b0808, 0x080819082b080819, 0x080819082b081908, 0x080819082b190808,
|
|
0x080819082b2b1908, 0x0808191908080808, 0x080819190808082b, 0x0808191908082b08,
|
|
0x08081919082b0808, 0x080819191908192b, 0x08081919192b2b19, 0x080819192b080808,
|
|
0x080819192b190819, 0x0808192b08082b19, 0x0808192b08190808, 0x0808192b19080808,
|
|
0x0808192b2b081908, 0x0808192b2b2b1908, 0x08082b0808080808, 0x08082b0808081919,
|
|
0x08082b0808082b08, 0x08082b0808191908, 0x08082b08082b2b08, 0x08082b0819080819,
|
|
0x08082b0819081908, 0x08082b0819190808, 0x08082b081919082b, 0x08082b082b082b08,
|
|
0x08082b1908081908, 0x08082b1919080808, 0x08082b2b0808082b, 0x08082b2b08191908,
|
|
0x0819080808080819, 0x0819080808081908, 0x0819080808190808, 0x08190808082b0819,
|
|
0x0819080819080808, 0x08190808192b0808, 0x081908082b081908, 0x081908082b190808,
|
|
0x081908082b191919, 0x0819081908080808, 0x0819081908082b08, 0x08190819082b0808,
|
|
0x0819081919190808, 0x0819081919192b2b, 0x081908192b080808, 0x0819082b082b1908,
|
|
0x0819082b19081919, 0x0819190808080808, 0x0819190808082b08, 0x08191908082b0808,
|
|
0x08191908082b1919, 0x0819190819082b19, 0x081919082b080808, 0x0819191908192b08,
|
|
0x08191919192b082b, 0x0819192b08080808, 0x0819192b0819192b, 0x08192b0808080819,
|
|
0x08192b0808081908, 0x08192b0808190808, 0x08192b0819080808, 0x08192b082b080819,
|
|
0x08192b1908080808, 0x08192b1908081919, 0x08192b192b2b0808, 0x08192b2b19190819,
|
|
0x082b080808080808, 0x082b08080808082b, 0x082b080808082b2b, 0x082b080819081908,
|
|
0x082b0808192b0819, 0x082b08082b080808, 0x082b08082b08082b, 0x082b0819082b2b19,
|
|
0x082b081919082b08, 0x082b082b08080808, 0x082b082b0808082b, 0x082b190808080819,
|
|
0x082b190808081908, 0x082b190808190808, 0x082b190819080808, 0x082b19081919192b,
|
|
0x082b191908080808, 0x082b191919080819, 0x082b1919192b1908, 0x082b192b2b190808,
|
|
0x082b2b0808082b08, 0x082b2b08082b0808, 0x082b2b082b191908, 0x082b2b2b19081908,
|
|
0x1908080808080819, 0x1908080808081908, 0x1908080808190808, 0x1908080808192b08,
|
|
0x19080808082b0819, 0x19080808082b1908, 0x1908080819080808, 0x1908080819082b08,
|
|
0x190808081919192b, 0x19080808192b0808, 0x190808082b080819, 0x190808082b081908,
|
|
0x190808082b190808, 0x1908081908080808, 0x19080819082b0808, 0x19080819192b0819,
|
|
0x190808192b080808, 0x190808192b081919, 0x1908082b08080819, 0x1908082b08190808,
|
|
0x1908082b19082b08, 0x1908082b1919192b, 0x1908082b192b2b08, 0x1908190808080808,
|
|
0x1908190808082b08, 0x19081908082b0808, 0x190819082b080808, 0x190819082b192b19,
|
|
0x190819190819082b, 0x19081919082b1908, 0x1908192b08080808, 0x19082b0808080819,
|
|
0x19082b0808081908, 0x19082b0808190808, 0x19082b0819080808, 0x19082b0819081919,
|
|
0x19082b1908080808, 0x19082b1919192b08, 0x19082b19192b0819, 0x19082b192b08082b,
|
|
0x19082b2b19081919, 0x19082b2b2b190808, 0x1919080808080808, 0x1919080808082b08,
|
|
0x1919080808190819, 0x1919080808192b19, 0x19190808082b0808, 0x191908082b080808,
|
|
0x191908082b082b08, 0x1919081908081908, 0x191908191908082b, 0x191908192b2b1908,
|
|
0x1919082b2b190819, 0x191919082b190808, 0x191919082b19082b, 0x1919191908082b2b,
|
|
0x1919192b08080819, 0x1919192b19191908, 0x19192b0808080808, 0x19192b0808190819,
|
|
0x19192b0808192b19, 0x19192b08192b1908, 0x19192b1919080808, 0x19192b2b08082b08,
|
|
0x192b080808081908, 0x192b080808190808, 0x192b080819080808, 0x192b0808192b2b08,
|
|
0x192b081908080808, 0x192b081919191919, 0x192b082b08192b08, 0x192b082b192b0808,
|
|
0x192b190808080808, 0x192b190808081919, 0x192b191908190808, 0x192b19190819082b,
|
|
0x192b19192b081908, 0x192b2b081908082b, 0x2b08080808080808, 0x2b0808080808082b,
|
|
0x2b08080808082b2b, 0x2b08080819080819, 0x2b0808082b08082b, 0x2b08081908081908,
|
|
0x2b08081908192b08, 0x2b08081919080808, 0x2b08082b08190819, 0x2b08190808080819,
|
|
0x2b08190808081908, 0x2b08190808190808, 0x2b08190808191919, 0x2b08190819080808,
|
|
0x2b081908192b0808, 0x2b08191908080808, 0x2b0819191908192b, 0x2b0819192b191908,
|
|
0x2b08192b08082b19, 0x2b08192b19080808, 0x2b08192b192b0808, 0x2b082b080808082b,
|
|
0x2b082b1908081908, 0x2b082b2b08190819, 0x2b19080808081908, 0x2b19080808190808,
|
|
0x2b190808082b1908, 0x2b19080819080808, 0x2b1908082b2b0819, 0x2b1908190819192b,
|
|
0x2b1908192b080808, 0x2b19082b19081919, 0x2b19190808080808, 0x2b191908082b082b,
|
|
0x2b19190819081908, 0x2b19191919190819, 0x2b192b082b080819, 0x2b192b19082b0808,
|
|
0x2b2b08080808082b, 0x2b2b080819190808, 0x2b2b08082b081919, 0x2b2b081908082b19,
|
|
0x2b2b082b08080808, 0x2b2b190808192b08, 0x2b2b2b0819190808, 0x2b2b2b1908081908,
|
|
};
|
|
|
|
#define kmask_iq2xs ds4_metal_kmask_iq2xs
|
|
#define ksigns_iq2xs ds4_metal_ksigns_iq2xs
|
|
#define iq2xxs_grid ds4_metal_iq2xxs_grid
|
|
|
|
struct block_q2_K {
|
|
uchar scales[QK_K/16];
|
|
uchar qs[QK_K/4];
|
|
half d;
|
|
half dmin;
|
|
};
|
|
|
|
struct block_q4_K {
|
|
half d;
|
|
half dmin;
|
|
uchar scales[12];
|
|
uchar qs[QK_K/2];
|
|
};
|
|
|
|
struct block_q5_K {
|
|
half d;
|
|
half dmin;
|
|
uchar scales[12];
|
|
uchar qh[QK_K/8];
|
|
uchar qs[QK_K/2];
|
|
};
|
|
|
|
struct block_q6_K {
|
|
uchar ql[QK_K/2];
|
|
uchar qh[QK_K/4];
|
|
char scales[QK_K/16];
|
|
half d;
|
|
};
|
|
|
|
struct block_iq2_xxs {
|
|
half d;
|
|
ushort qs[QK_K/8];
|
|
};
|
|
|
|
struct block_mxfp4 {
|
|
uchar e;
|
|
uchar qs[QK_MXFP4/2];
|
|
};
|
|
|
|
struct ds4_metal_glm_routed_moe_args {
|
|
uint32_t in_dim;
|
|
uint32_t mid_dim;
|
|
uint32_t out_dim;
|
|
uint32_t n_total_expert;
|
|
uint32_t n_expert_used;
|
|
uint32_t n_tokens;
|
|
uint32_t mid_token_stride;
|
|
uint32_t down_type;
|
|
float swiglu_clamp;
|
|
/* Expert ownership under tensor parallelism: tp_world 0/1 = full
|
|
* compute; otherwise each rank owns a contiguous expert range. */
|
|
int32_t tp_rank;
|
|
int32_t tp_world;
|
|
int32_t tp_expert_base;
|
|
uint64_t gate_expert_bytes;
|
|
uint64_t gate_row_bytes;
|
|
uint64_t up_expert_bytes;
|
|
uint64_t up_row_bytes;
|
|
uint64_t down_expert_bytes;
|
|
uint64_t down_row_bytes;
|
|
};
|
|
|
|
static inline float ds4_glm_swiglu(float gate, float up, float limit) {
|
|
if (limit > 1.0e-6f) {
|
|
gate = min(gate, limit);
|
|
up = clamp(up, -limit, limit);
|
|
}
|
|
return (gate / (1.0f + exp(-gate))) * up;
|
|
}
|
|
|
|
|
|
static inline bool ds4_tp_owns_expert(int expert, int n_total,
|
|
int tp_rank, int tp_world) {
|
|
if (tp_world <= 1) return true;
|
|
const int first = tp_rank * (n_total / tp_world);
|
|
const int last = tp_rank + 1 == tp_world ? n_total : first + n_total / tp_world;
|
|
return expert >= first && expert < last;
|
|
}
|
|
|
|
struct ds4_metal_dsv4_moe_swiglu_weight_args {
|
|
uint32_t width;
|
|
uint32_t rows;
|
|
uint64_t gate_row_stride;
|
|
uint64_t up_row_stride;
|
|
uint64_t mid_row_stride;
|
|
uint64_t weight_stride;
|
|
uint32_t write_clamped;
|
|
float clamp_value;
|
|
};
|
|
|
|
struct ds4_metal_dsv4_moe_sum6_args {
|
|
uint32_t width;
|
|
uint32_t tokens;
|
|
uint64_t src_token_stride;
|
|
uint64_t dst_token_stride;
|
|
};
|
|
|
|
// Routed-MoE activation for the selected experts:
|
|
// clamp(gate), clamp(up), silu(gate) * up * route_weight. Normal inference
|
|
// does not consume gate/up after this point, so the fast path avoids writing the
|
|
// clamped intermediates back. A diagnostic env switch can restore those writes
|
|
// when comparing the old multi-kernel intermediate tensors.
|
|
kernel void kernel_dsv4_moe_swiglu_weight(
|
|
constant ds4_metal_dsv4_moe_swiglu_weight_args &args,
|
|
device char *gate,
|
|
device char *up,
|
|
device char *mid,
|
|
device const char *weights,
|
|
uint row [[threadgroup_position_in_grid]],
|
|
uint tid [[thread_position_in_threadgroup]],
|
|
uint ntg [[threads_per_threadgroup]]) {
|
|
if (row >= args.rows) return;
|
|
|
|
device float *gate_row = (device float *)(gate + (uint64_t)row * args.gate_row_stride);
|
|
device float *up_row = (device float *)(up + (uint64_t)row * args.up_row_stride);
|
|
device float *mid_row = (device float *)(mid + (uint64_t)row * args.mid_row_stride);
|
|
device const float *w = (device const float *)(weights + (uint64_t)row * args.weight_stride);
|
|
const float route_weight = w[0];
|
|
const float c = args.clamp_value;
|
|
|
|
for (uint i = tid; i < args.width; i += ntg) {
|
|
float g = gate_row[i];
|
|
float u = up_row[i];
|
|
if (c > 1.0e-6f) {
|
|
g = min(g, c);
|
|
u = clamp(u, -c, c);
|
|
if (args.write_clamped != 0) {
|
|
gate_row[i] = g;
|
|
up_row[i] = u;
|
|
}
|
|
}
|
|
const float silu = g / (1.0f + exp(-g));
|
|
mid_row[i] = silu * u * route_weight;
|
|
}
|
|
}
|
|
|
|
// Same routed-MoE activation as above, but stores the down-projection input in
|
|
// half precision. The grouped matmul path converts F32 activations to half
|
|
// before MMA anyway, so this cuts the large mid write/read traffic without
|
|
// changing the effective matmul input precision.
|
|
kernel void kernel_dsv4_moe_swiglu_weight_f16(
|
|
constant ds4_metal_dsv4_moe_swiglu_weight_args &args,
|
|
device char *gate,
|
|
device char *up,
|
|
device char *mid,
|
|
device const char *weights,
|
|
uint row [[threadgroup_position_in_grid]],
|
|
uint tid [[thread_position_in_threadgroup]],
|
|
uint ntg [[threads_per_threadgroup]]) {
|
|
if (row >= args.rows) return;
|
|
|
|
device float *gate_row = (device float *)(gate + (uint64_t)row * args.gate_row_stride);
|
|
device float *up_row = (device float *)(up + (uint64_t)row * args.up_row_stride);
|
|
device half *mid_row = (device half *)(mid + (uint64_t)row * args.mid_row_stride);
|
|
device const float *w = (device const float *)(weights + (uint64_t)row * args.weight_stride);
|
|
const float route_weight = w[0];
|
|
const float c = args.clamp_value;
|
|
|
|
for (uint i = tid; i < args.width; i += ntg) {
|
|
float g = gate_row[i];
|
|
float u = up_row[i];
|
|
if (c > 1.0e-6f) {
|
|
g = min(g, c);
|
|
u = clamp(u, -c, c);
|
|
if (args.write_clamped != 0) {
|
|
gate_row[i] = g;
|
|
up_row[i] = u;
|
|
}
|
|
}
|
|
const float silu = g / (1.0f + exp(-g));
|
|
mid_row[i] = (half)(silu * u * route_weight);
|
|
}
|
|
}
|
|
|
|
kernel void kernel_dsv4_moe_sum6_f32(
|
|
constant ds4_metal_dsv4_moe_sum6_args &args,
|
|
device const char *src,
|
|
device char *dst,
|
|
uint token[[threadgroup_position_in_grid]],
|
|
uint tid[[thread_position_in_threadgroup]],
|
|
uint ntg[[threads_per_threadgroup]]) {
|
|
if (token >= args.tokens) return;
|
|
|
|
device const float *s =
|
|
(device const float *)(src + (uint64_t)token * args.src_token_stride);
|
|
device float *d =
|
|
(device float *)(dst + (uint64_t)token * args.dst_token_stride);
|
|
|
|
for (uint col = tid; col < args.width; col += ntg) {
|
|
float v = s[col];
|
|
v += s[args.width + col];
|
|
v += s[2u * args.width + col];
|
|
v += s[3u * args.width + col];
|
|
v += s[4u * args.width + col];
|
|
v += s[5u * args.width + col];
|
|
d[col] = v;
|
|
}
|
|
}
|
|
|
|
kernel void kernel_dsv4_moe_sum8_f32(
|
|
constant ds4_metal_dsv4_moe_sum6_args &args,
|
|
device const char *src,
|
|
device char *dst,
|
|
uint token[[threadgroup_position_in_grid]],
|
|
uint tid[[thread_position_in_threadgroup]],
|
|
uint ntg[[threads_per_threadgroup]]) {
|
|
if (token >= args.tokens) return;
|
|
|
|
device const float *s =
|
|
(device const float *)(src + (uint64_t)token * args.src_token_stride);
|
|
device float *d =
|
|
(device float *)(dst + (uint64_t)token * args.dst_token_stride);
|
|
|
|
for (uint col = tid; col < args.width; col += ntg) {
|
|
float v = s[col];
|
|
v += s[args.width + col];
|
|
v += s[2u * args.width + col];
|
|
v += s[3u * args.width + col];
|
|
v += s[4u * args.width + col];
|
|
v += s[5u * args.width + col];
|
|
v += s[6u * args.width + col];
|
|
v += s[7u * args.width + col];
|
|
d[col] = v;
|
|
}
|
|
}
|
|
|
|
template <typename type4x4>
|
|
void dequantize_q2_K(device const block_q2_K *xb, short il, thread type4x4 & reg) {
|
|
const float d = xb->d;
|
|
const float min = xb->dmin;
|
|
device const uint8_t * q = (device const uint8_t *)xb->qs;
|
|
float dl, ml;
|
|
uint8_t sc = xb->scales[il];
|
|
|
|
q = q + 32*(il/8) + 16*(il&1);
|
|
il = (il/2)%4;
|
|
|
|
half coef = il>1 ? (il>2 ? 1/64.h : 1/16.h) : (il>0 ? 1/4.h : 1.h);
|
|
uchar mask = il>1 ? (il>2 ? 192 : 48) : (il>0 ? 12 : 3);
|
|
dl = d * (sc & 0xF) * coef, ml = min * (sc >> 4);
|
|
for (int i = 0; i < 16; ++i) {
|
|
reg[i/4][i%4] = dl * (q[i] & mask) - ml;
|
|
}
|
|
}
|
|
|
|
static inline float ds4_glm_q2_K_value(device const block_q2_K *blocks, uint k) {
|
|
const uint block = k / QK_K;
|
|
const uint idx = k - block * QK_K;
|
|
device const block_q2_K *xb = blocks + block;
|
|
const uint group = idx / 16u;
|
|
const uint l = idx - group * 16u;
|
|
const uint q_base = 32u * (group / 8u) + 16u * (group & 1u);
|
|
const uint shift = ((group / 2u) & 3u) * 2u;
|
|
const uint q = ((uint)xb->qs[q_base + l] >> shift) & 0x03u;
|
|
const uint sc = (uint)xb->scales[group];
|
|
return (float)xb->d * (float)(sc & 0x0fu) * (float)q -
|
|
(float)xb->dmin * (float)(sc >> 4u);
|
|
}
|
|
|
|
static inline uchar2 get_scale_min_k4_just2(int j, int k, device const uchar * q) {
|
|
return j < 4 ? uchar2{uchar(q[j+0+k] & 63), uchar(q[j+4+k] & 63)}
|
|
: uchar2{uchar((q[j+4+k] & 0xF) | ((q[j-4+k] & 0xc0) >> 2)),
|
|
uchar((q[j+4+k] >> 4) | ((q[j-0+k] & 0xc0) >> 2))};
|
|
}
|
|
|
|
static inline float ds4_glm_q4_K_value(device const block_q4_K *blocks, uint k) {
|
|
const uint block = k / QK_K;
|
|
const uint idx = k - block * QK_K;
|
|
device const block_q4_K *xb = blocks + block;
|
|
const uint group = idx / 32u;
|
|
const uint l = idx - group * 32u;
|
|
const uchar2 sm = get_scale_min_k4_just2((int)group, 0, xb->scales);
|
|
const uint byte_off = (group >> 1u) * 32u + l;
|
|
const uint shift = (group & 1u) * 4u;
|
|
const uint q = (xb->qs[byte_off] >> shift) & 0x0Fu;
|
|
return (float)xb->d * (float)sm.x * (float)q -
|
|
(float)xb->dmin * (float)sm.y;
|
|
}
|
|
|
|
static inline float ds4_glm_q5_K_value(device const block_q5_K *blocks, uint k) {
|
|
const uint block = k / QK_K;
|
|
const uint idx = k - block * QK_K;
|
|
device const block_q5_K *xb = blocks + block;
|
|
const uint group = idx / 32u;
|
|
const uint l = idx - group * 32u;
|
|
const uchar2 sm = get_scale_min_k4_just2((int)group, 0, xb->scales);
|
|
const uint ql_base = (group >> 1u) * 32u + l;
|
|
const uint shift = (group & 1u) * 4u;
|
|
uint q = (xb->qs[ql_base] >> shift) & 0x0Fu;
|
|
q += (xb->qh[l] & (uchar)(1u << group)) ? 16u : 0u;
|
|
return (float)xb->d * (float)sm.x * (float)q -
|
|
(float)xb->dmin * (float)sm.y;
|
|
}
|
|
|
|
static inline float ds4_glm_q6_K_value(device const block_q6_K *blocks, uint k) {
|
|
const uint block = k / QK_K;
|
|
const uint idx = k - block * QK_K;
|
|
device const block_q6_K *xb = blocks + block;
|
|
const uint n128 = idx >> 7u;
|
|
const uint r = idx & 127u;
|
|
const uint l = r & 31u;
|
|
const uint quarter = r >> 5u;
|
|
const uint ql_base = n128 * 64u;
|
|
const uint qh_base = n128 * 32u;
|
|
const uint sc_base = n128 * 8u;
|
|
uint q;
|
|
int sc;
|
|
|
|
if (quarter == 0u) {
|
|
q = (xb->ql[ql_base + l] & 0x0Fu) | (((xb->qh[qh_base + l] >> 0u) & 3u) << 4u);
|
|
sc = (int)xb->scales[sc_base + l / 16u + 0u];
|
|
} else if (quarter == 1u) {
|
|
q = (xb->ql[ql_base + 32u + l] & 0x0Fu) | (((xb->qh[qh_base + l] >> 2u) & 3u) << 4u);
|
|
sc = (int)xb->scales[sc_base + l / 16u + 2u];
|
|
} else if (quarter == 2u) {
|
|
q = (xb->ql[ql_base + l] >> 4u) | (((xb->qh[qh_base + l] >> 4u) & 3u) << 4u);
|
|
sc = (int)xb->scales[sc_base + l / 16u + 4u];
|
|
} else {
|
|
q = (xb->ql[ql_base + 32u + l] >> 4u) | (((xb->qh[qh_base + l] >> 6u) & 3u) << 4u);
|
|
sc = (int)xb->scales[sc_base + l / 16u + 6u];
|
|
}
|
|
|
|
return (float)xb->d * (float)sc * (float)((int)q - 32);
|
|
}
|
|
|
|
kernel void kernel_glm_q4_K_pair_swiglu_f32(
|
|
constant ds4_metal_glm_routed_moe_args &args,
|
|
device const char *gate,
|
|
device const char *up,
|
|
device const float *x,
|
|
device const int32_t *selected,
|
|
device const float *weights,
|
|
device float *mid,
|
|
threadgroup float *scratch [[threadgroup(0)]],
|
|
uint3 tgpig [[threadgroup_position_in_grid]],
|
|
uint tid [[thread_index_in_threadgroup]]) {
|
|
const uint ntg = 256u;
|
|
const uint row = tgpig.x;
|
|
const uint slot = tgpig.y;
|
|
const uint token = tgpig.z;
|
|
if (row >= args.mid_dim || slot >= args.n_expert_used || token >= args.n_tokens) return;
|
|
|
|
const uint64_t selected_off = (uint64_t)token * args.n_expert_used + slot;
|
|
const uint64_t mid_off = (uint64_t)token * args.mid_token_stride +
|
|
(uint64_t)slot * args.mid_dim + row;
|
|
const int expert = selected[selected_off];
|
|
if (!ds4_tp_owns_expert(expert, args.n_total_expert,
|
|
args.tp_rank, args.tp_world)) {
|
|
if (tid == 0u) mid[mid_off] = 0.0f;
|
|
return;
|
|
}
|
|
const int local_expert = expert - args.tp_expert_base;
|
|
|
|
device const block_q4_K *gate_row =
|
|
(device const block_q4_K *)(gate +
|
|
(uint64_t)(uint)local_expert * args.gate_expert_bytes +
|
|
(uint64_t)row * args.gate_row_bytes);
|
|
device const block_q4_K *up_row =
|
|
(device const block_q4_K *)(up +
|
|
(uint64_t)(uint)local_expert * args.up_expert_bytes +
|
|
(uint64_t)row * args.up_row_bytes);
|
|
|
|
float acc_gate = 0.0f;
|
|
float acc_up = 0.0f;
|
|
device const float *token_x = x + (uint64_t)token * args.in_dim;
|
|
for (uint k = tid; k < args.in_dim; k += ntg) {
|
|
const float xv = token_x[k];
|
|
acc_gate += ds4_glm_q4_K_value(gate_row, k) * xv;
|
|
acc_up += ds4_glm_q4_K_value(up_row, k) * xv;
|
|
}
|
|
|
|
scratch[tid] = acc_gate;
|
|
scratch[ntg + tid] = acc_up;
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
for (uint stride = ntg >> 1u; stride > 0u; stride >>= 1u) {
|
|
if (tid < stride) {
|
|
scratch[tid] += scratch[tid + stride];
|
|
scratch[ntg + tid] += scratch[ntg + tid + stride];
|
|
}
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
}
|
|
|
|
if (tid == 0u) {
|
|
mid[mid_off] = ds4_glm_swiglu(scratch[0], scratch[ntg],
|
|
args.swiglu_clamp) * weights[selected_off];
|
|
}
|
|
}
|
|
|
|
template <short N_R0>
|
|
static inline void glm_q2_K_pair_swiglu_simd_f32_impl(
|
|
ds4_metal_glm_routed_moe_args args,
|
|
device const char *gate,
|
|
device const char *up,
|
|
device const float *x,
|
|
device const float *weights,
|
|
device float *mid,
|
|
threadgroup float *scratch,
|
|
uint3 tgpig,
|
|
uint slot,
|
|
uint token,
|
|
uint64_t selected_off,
|
|
int expert,
|
|
ushort tiisg,
|
|
ushort sgitg) {
|
|
const short NSG = 2;
|
|
const uint row0 = ((uint)tgpig.x * (uint)NSG + (uint)sgitg) * N_R0;
|
|
if (row0 >= args.mid_dim || slot >= args.n_expert_used || token >= args.n_tokens) return;
|
|
|
|
const uint64_t mid_base = (uint64_t)token * args.mid_token_stride +
|
|
(uint64_t)slot * args.mid_dim;
|
|
if (expert < 0 || (uint)expert >= args.n_total_expert) {
|
|
if (tiisg == 0u) {
|
|
for (short row = 0;
|
|
row < N_R0 && row0 + (uint)row < args.mid_dim;
|
|
row++) {
|
|
mid[mid_base + row0 + (uint)row] = 0.0f;
|
|
}
|
|
}
|
|
return;
|
|
}
|
|
|
|
const short ix = tiisg / 8;
|
|
const short it = tiisg % 8;
|
|
const short iq = it / 4;
|
|
const short ir = it % 4;
|
|
const short is = (8 * ir) / 16;
|
|
const int nb = args.in_dim / QK_K;
|
|
const uint64_t expert_gate = (uint64_t)(uint)expert * args.gate_expert_bytes;
|
|
const uint64_t expert_up = (uint64_t)(uint)expert * args.up_expert_bytes;
|
|
device const block_q2_K *xg =
|
|
(device const block_q2_K *)(gate + expert_gate +
|
|
(uint64_t)row0 * args.gate_row_bytes);
|
|
device const block_q2_K *xu =
|
|
(device const block_q2_K *)(up + expert_up +
|
|
(uint64_t)row0 * args.up_row_bytes);
|
|
device const float *y = x + (uint64_t)token * args.in_dim;
|
|
device const float *y4 = y + ix * QK_K + 128 * iq + 8 * ir;
|
|
|
|
float yl[32];
|
|
float sumg[N_R0] = {0.f};
|
|
float sumu[N_R0] = {0.f};
|
|
|
|
for (int ib = ix; ib < nb; ib += 4) {
|
|
float4 sumy = {0.f, 0.f, 0.f, 0.f};
|
|
for (short i = 0; i < 8; ++i) {
|
|
yl[i + 0] = y4[i + 0]; sumy[0] += yl[i + 0];
|
|
yl[i + 8] = y4[i + 32]; sumy[1] += yl[i + 8];
|
|
yl[i + 16] = y4[i + 64]; sumy[2] += yl[i + 16];
|
|
yl[i + 24] = y4[i + 96]; sumy[3] += yl[i + 24];
|
|
}
|
|
|
|
device const uint8_t *scg = (device const uint8_t *)xg[ib].scales + 8 * iq + is;
|
|
device const uint16_t *qsg = (device const uint16_t *)xg[ib].qs + 16 * iq + 4 * ir;
|
|
device const half *dhg = &xg[ib].d;
|
|
device const uint8_t *scu = (device const uint8_t *)xu[ib].scales + 8 * iq + is;
|
|
device const uint16_t *qsu = (device const uint16_t *)xu[ib].qs + 16 * iq + 4 * ir;
|
|
device const half *dhu = &xu[ib].d;
|
|
|
|
for (short row = 0;
|
|
row < N_R0 && row0 + (uint)row < args.mid_dim;
|
|
row++) {
|
|
float4 acc1g = {0.f, 0.f, 0.f, 0.f};
|
|
float4 acc2g = {0.f, 0.f, 0.f, 0.f};
|
|
float4 acc1u = {0.f, 0.f, 0.f, 0.f};
|
|
float4 acc2u = {0.f, 0.f, 0.f, 0.f};
|
|
|
|
for (int i = 0; i < 8; i += 2) {
|
|
acc1g[0] += yl[i + 0] * (qsg[i / 2] & 0x0003);
|
|
acc2g[0] += yl[i + 1] * (qsg[i / 2] & 0x0300);
|
|
acc1g[1] += yl[i + 8] * (qsg[i / 2] & 0x000c);
|
|
acc2g[1] += yl[i + 9] * (qsg[i / 2] & 0x0c00);
|
|
acc1g[2] += yl[i + 16] * (qsg[i / 2] & 0x0030);
|
|
acc2g[2] += yl[i + 17] * (qsg[i / 2] & 0x3000);
|
|
acc1g[3] += yl[i + 24] * (qsg[i / 2] & 0x00c0);
|
|
acc2g[3] += yl[i + 25] * (qsg[i / 2] & 0xc000);
|
|
|
|
acc1u[0] += yl[i + 0] * (qsu[i / 2] & 0x0003);
|
|
acc2u[0] += yl[i + 1] * (qsu[i / 2] & 0x0300);
|
|
acc1u[1] += yl[i + 8] * (qsu[i / 2] & 0x000c);
|
|
acc2u[1] += yl[i + 9] * (qsu[i / 2] & 0x0c00);
|
|
acc1u[2] += yl[i + 16] * (qsu[i / 2] & 0x0030);
|
|
acc2u[2] += yl[i + 17] * (qsu[i / 2] & 0x3000);
|
|
acc1u[3] += yl[i + 24] * (qsu[i / 2] & 0x00c0);
|
|
acc2u[3] += yl[i + 25] * (qsu[i / 2] & 0xc000);
|
|
}
|
|
|
|
const float dg = dhg[0];
|
|
const float mg = dhg[1] * 1.f / 16.f;
|
|
sumg[row] += dg * ((acc1g[0] + 1.f / 256.f * acc2g[0]) * (scg[0] & 0xF) * 1.f / 1.f +
|
|
(acc1g[1] + 1.f / 256.f * acc2g[1]) * (scg[2] & 0xF) * 1.f / 4.f +
|
|
(acc1g[2] + 1.f / 256.f * acc2g[2]) * (scg[4] & 0xF) * 1.f / 16.f +
|
|
(acc1g[3] + 1.f / 256.f * acc2g[3]) * (scg[6] & 0xF) * 1.f / 64.f) -
|
|
mg * (sumy[0] * (scg[0] & 0xF0) + sumy[1] * (scg[2] & 0xF0) +
|
|
sumy[2] * (scg[4] & 0xF0) + sumy[3] * (scg[6] & 0xF0));
|
|
|
|
const float du = dhu[0];
|
|
const float mu = dhu[1] * 1.f / 16.f;
|
|
sumu[row] += du * ((acc1u[0] + 1.f / 256.f * acc2u[0]) * (scu[0] & 0xF) * 1.f / 1.f +
|
|
(acc1u[1] + 1.f / 256.f * acc2u[1]) * (scu[2] & 0xF) * 1.f / 4.f +
|
|
(acc1u[2] + 1.f / 256.f * acc2u[2]) * (scu[4] & 0xF) * 1.f / 16.f +
|
|
(acc1u[3] + 1.f / 256.f * acc2u[3]) * (scu[6] & 0xF) * 1.f / 64.f) -
|
|
mu * (sumy[0] * (scu[0] & 0xF0) + sumy[1] * (scu[2] & 0xF0) +
|
|
sumy[2] * (scu[4] & 0xF0) + sumy[3] * (scu[6] & 0xF0));
|
|
|
|
qsg += args.gate_row_bytes / 2;
|
|
scg += args.gate_row_bytes;
|
|
dhg += args.gate_row_bytes / 2;
|
|
qsu += args.up_row_bytes / 2;
|
|
scu += args.up_row_bytes;
|
|
dhu += args.up_row_bytes / 2;
|
|
}
|
|
|
|
y4 += 4 * QK_K;
|
|
}
|
|
|
|
for (short row = 0;
|
|
row < N_R0 && row0 + (uint)row < args.mid_dim;
|
|
row++) {
|
|
const float g = simd_sum(sumg[row]);
|
|
const float u = simd_sum(sumu[row]);
|
|
if (tiisg == 0u) {
|
|
mid[mid_base + row0 + (uint)row] =
|
|
ds4_glm_swiglu(g, u, args.swiglu_clamp) * weights[selected_off];
|
|
}
|
|
}
|
|
|
|
(void)scratch;
|
|
}
|
|
|
|
kernel void kernel_glm_q2_K_pair_swiglu_f32(
|
|
constant ds4_metal_glm_routed_moe_args &args,
|
|
device const char *gate,
|
|
device const char *up,
|
|
device const float *x,
|
|
device const int32_t *selected,
|
|
device const float *weights,
|
|
device float *mid,
|
|
threadgroup float *scratch [[threadgroup(0)]],
|
|
uint3 tgpig [[threadgroup_position_in_grid]],
|
|
ushort tiisg [[thread_index_in_simdgroup]],
|
|
ushort sgitg [[simdgroup_index_in_threadgroup]]) {
|
|
const uint slot = tgpig.y;
|
|
const uint token = tgpig.z;
|
|
if (slot >= args.n_expert_used || token >= args.n_tokens) return;
|
|
const uint64_t selected_off = (uint64_t)token * args.n_expert_used + slot;
|
|
const int expert = selected[selected_off];
|
|
if (!ds4_tp_owns_expert(expert, args.n_total_expert,
|
|
args.tp_rank, args.tp_world)) return;
|
|
glm_q2_K_pair_swiglu_simd_f32_impl<N_R0_Q2_K>(
|
|
args, gate, up, x, weights, mid, scratch,
|
|
tgpig, slot, token, selected_off, expert - args.tp_expert_base,
|
|
tiisg, sgitg);
|
|
}
|
|
|
|
kernel void kernel_glm_q2_K_addr_pair_swiglu2_f32(
|
|
constant ds4_metal_glm_routed_moe_args &args,
|
|
device const uint64_t *gate_addrs,
|
|
device const uint64_t *up_addrs,
|
|
device const float *x,
|
|
device const int32_t *selected,
|
|
device const float *weights,
|
|
device float *mid,
|
|
threadgroup float *scratch [[threadgroup(0)]],
|
|
uint3 tgpig [[threadgroup_position_in_grid]],
|
|
ushort tiisg [[thread_index_in_simdgroup]],
|
|
ushort sgitg [[simdgroup_index_in_threadgroup]]) {
|
|
const uint slot = tgpig.y;
|
|
const uint token = tgpig.z;
|
|
if (slot >= args.n_expert_used || token >= args.n_tokens) return;
|
|
|
|
const uint64_t selected_off = (uint64_t)token * args.n_expert_used + slot;
|
|
const int expert = selected[selected_off];
|
|
const short NSG = 2;
|
|
const uint row0 = ((uint)tgpig.x * (uint)NSG + (uint)sgitg) * N_R0_GLM_Q2_PAIR2_K;
|
|
const uint64_t mid_base = (uint64_t)token * args.mid_token_stride +
|
|
(uint64_t)slot * args.mid_dim;
|
|
if (row0 >= args.mid_dim) return;
|
|
|
|
if (expert < 0 || (uint)expert >= args.n_total_expert ||
|
|
!ds4_tp_owns_expert(expert, args.n_total_expert,
|
|
args.tp_rank, args.tp_world)) {
|
|
if (tiisg == 0u) {
|
|
for (short row = 0;
|
|
row < N_R0_GLM_Q2_PAIR2_K && row0 + (uint)row < args.mid_dim;
|
|
row++) {
|
|
mid[mid_base + row0 + (uint)row] = 0.0f;
|
|
}
|
|
}
|
|
return;
|
|
}
|
|
const uint64_t gate_addr = gate_addrs[(uint)expert];
|
|
const uint64_t up_addr = up_addrs[(uint)expert];
|
|
if (gate_addr == 0 || up_addr == 0) {
|
|
if (tiisg == 0u) {
|
|
for (short row = 0;
|
|
row < N_R0_GLM_Q2_PAIR2_K && row0 + (uint)row < args.mid_dim;
|
|
row++) {
|
|
mid[mid_base + row0 + (uint)row] = 0.0f;
|
|
}
|
|
}
|
|
return;
|
|
}
|
|
|
|
ds4_metal_glm_routed_moe_args local = args;
|
|
local.n_total_expert = 1;
|
|
local.gate_expert_bytes = 0;
|
|
local.up_expert_bytes = 0;
|
|
glm_q2_K_pair_swiglu_simd_f32_impl<N_R0_GLM_Q2_PAIR2_K>(
|
|
local,
|
|
reinterpret_cast<device const char *>(gate_addr),
|
|
reinterpret_cast<device const char *>(up_addr),
|
|
x, weights, mid, scratch,
|
|
tgpig, slot, token, selected_off, 0, tiisg, sgitg);
|
|
}
|
|
|
|
kernel void kernel_glm_q2_K_addr_pair_swiglu2_f32_masked(
|
|
constant ds4_metal_glm_routed_moe_args &args,
|
|
constant uint32_t &active_mask,
|
|
device const uint64_t *gate_addrs,
|
|
device const uint64_t *up_addrs,
|
|
device const float *x,
|
|
device const int32_t *selected,
|
|
device const float *weights,
|
|
device float *mid,
|
|
threadgroup float *scratch [[threadgroup(0)]],
|
|
uint3 tgpig [[threadgroup_position_in_grid]],
|
|
ushort tiisg [[thread_index_in_simdgroup]],
|
|
ushort sgitg [[simdgroup_index_in_threadgroup]]) {
|
|
const uint slot = tgpig.y;
|
|
const uint token = tgpig.z;
|
|
if (slot >= args.n_expert_used || token >= args.n_tokens) return;
|
|
if ((active_mask & (1u << slot)) == 0u) return;
|
|
|
|
const uint64_t selected_off = (uint64_t)token * args.n_expert_used + slot;
|
|
const int expert = selected[selected_off];
|
|
const short NSG = 2;
|
|
const uint row0 = ((uint)tgpig.x * (uint)NSG + (uint)sgitg) * N_R0_GLM_Q2_PAIR2_K;
|
|
const uint64_t mid_base = (uint64_t)token * args.mid_token_stride +
|
|
(uint64_t)slot * args.mid_dim;
|
|
if (row0 >= args.mid_dim) return;
|
|
|
|
if (expert < 0 || (uint)expert >= args.n_total_expert ||
|
|
!ds4_tp_owns_expert(expert, args.n_total_expert,
|
|
args.tp_rank, args.tp_world)) {
|
|
if (tiisg == 0u) {
|
|
for (short row = 0;
|
|
row < N_R0_GLM_Q2_PAIR2_K && row0 + (uint)row < args.mid_dim;
|
|
row++) {
|
|
mid[mid_base + row0 + (uint)row] = 0.0f;
|
|
}
|
|
}
|
|
return;
|
|
}
|
|
const uint64_t gate_addr = gate_addrs[(uint)expert];
|
|
const uint64_t up_addr = up_addrs[(uint)expert];
|
|
if (gate_addr == 0 || up_addr == 0) {
|
|
if (tiisg == 0u) {
|
|
for (short row = 0;
|
|
row < N_R0_GLM_Q2_PAIR2_K && row0 + (uint)row < args.mid_dim;
|
|
row++) {
|
|
mid[mid_base + row0 + (uint)row] = 0.0f;
|
|
}
|
|
}
|
|
return;
|
|
}
|
|
|
|
ds4_metal_glm_routed_moe_args local = args;
|
|
local.n_total_expert = 1;
|
|
local.gate_expert_bytes = 0;
|
|
local.up_expert_bytes = 0;
|
|
glm_q2_K_pair_swiglu_simd_f32_impl<N_R0_GLM_Q2_PAIR2_K>(
|
|
local,
|
|
reinterpret_cast<device const char *>(gate_addr),
|
|
reinterpret_cast<device const char *>(up_addr),
|
|
x, weights, mid, scratch,
|
|
tgpig, slot, token, selected_off, 0, tiisg, sgitg);
|
|
}
|
|
|
|
template <short N_R0>
|
|
static inline void glm_q4_K_pair_swiglu_simd_f32_impl(
|
|
ds4_metal_glm_routed_moe_args args,
|
|
device const char *gate,
|
|
device const char *up,
|
|
device const float *x,
|
|
device const float *weights,
|
|
device float *mid,
|
|
threadgroup float *scratch,
|
|
uint3 tgpig,
|
|
uint slot,
|
|
uint token,
|
|
uint64_t selected_off,
|
|
int expert,
|
|
ushort tiisg,
|
|
ushort sgitg) {
|
|
const short NSG = 2;
|
|
constexpr uint16_t kmask1 = 0x3f3f;
|
|
constexpr uint16_t kmask2 = 0x0f0f;
|
|
constexpr uint16_t kmask3 = 0xc0c0;
|
|
|
|
const uint row0 = ((uint)tgpig.x * (uint)NSG + (uint)sgitg) * N_R0;
|
|
if (row0 >= args.mid_dim || slot >= args.n_expert_used || token >= args.n_tokens) return;
|
|
|
|
const uint64_t mid_base = (uint64_t)token * args.mid_token_stride +
|
|
(uint64_t)slot * args.mid_dim;
|
|
if (expert < 0 || (uint)expert >= args.n_total_expert) {
|
|
if (tiisg == 0u) {
|
|
for (short row = 0;
|
|
row < N_R0 && row0 + (uint)row < args.mid_dim;
|
|
row++) {
|
|
mid[mid_base + row0 + (uint)row] = 0.0f;
|
|
}
|
|
}
|
|
return;
|
|
}
|
|
|
|
const short ix = tiisg / 8;
|
|
const short it = tiisg % 8;
|
|
const short iq = it / 4;
|
|
const short ir = it % 4;
|
|
const int nb = args.in_dim / QK_K;
|
|
const uint64_t expert_gate = (uint64_t)(uint)expert * args.gate_expert_bytes;
|
|
const uint64_t expert_up = (uint64_t)(uint)expert * args.up_expert_bytes;
|
|
device const block_q4_K *xg =
|
|
(device const block_q4_K *)(gate + expert_gate +
|
|
(uint64_t)row0 * args.gate_row_bytes);
|
|
device const block_q4_K *xu =
|
|
(device const block_q4_K *)(up + expert_up +
|
|
(uint64_t)row0 * args.up_row_bytes);
|
|
device const float *y = x + (uint64_t)token * args.in_dim;
|
|
device const float *y4 = y + ix * QK_K + 64 * iq + 8 * ir;
|
|
|
|
float sumg[N_R0] = {0.f};
|
|
float sumu[N_R0] = {0.f};
|
|
uint16_t sc16[4];
|
|
thread const uint8_t *sc8 = (thread const uint8_t *)sc16;
|
|
|
|
for (int ib = ix; ib < nb; ib += 4) {
|
|
float yl[16];
|
|
float yh[16];
|
|
float4 sumy = {0.f, 0.f, 0.f, 0.f};
|
|
|
|
for (short i = 0; i < 8; ++i) {
|
|
yl[i + 0] = y4[i + 0]; sumy[0] += yl[i + 0];
|
|
yl[i + 8] = y4[i + 32]; sumy[1] += yl[i + 8];
|
|
yh[i + 0] = y4[i + 128]; sumy[2] += yh[i + 0];
|
|
yh[i + 8] = y4[i + 160]; sumy[3] += yh[i + 8];
|
|
}
|
|
|
|
device const uint16_t *scg = (device const uint16_t *)xg[ib].scales + iq;
|
|
device const uint16_t *qg1 = (device const uint16_t *)xg[ib].qs + 16 * iq + 4 * ir;
|
|
device const half *dhg = &xg[ib].d;
|
|
device const uint16_t *scu = (device const uint16_t *)xu[ib].scales + iq;
|
|
device const uint16_t *qu1 = (device const uint16_t *)xu[ib].qs + 16 * iq + 4 * ir;
|
|
device const half *dhu = &xu[ib].d;
|
|
|
|
for (short row = 0;
|
|
row < N_R0 && row0 + (uint)row < args.mid_dim;
|
|
row++) {
|
|
sc16[0] = scg[0] & kmask1;
|
|
sc16[1] = scg[2] & kmask1;
|
|
sc16[2] = ((scg[4] >> 0) & kmask2) | ((scg[0] & kmask3) >> 2);
|
|
sc16[3] = ((scg[4] >> 4) & kmask2) | ((scg[2] & kmask3) >> 2);
|
|
|
|
device const uint16_t *qg2 = qg1 + 32;
|
|
float4 acc1g = {0.f, 0.f, 0.f, 0.f};
|
|
float4 acc2g = {0.f, 0.f, 0.f, 0.f};
|
|
|
|
FOR_UNROLL (short i = 0; i < 4; ++i) {
|
|
acc1g[0] += yl[2 * i + 0] * (qg1[i] & 0x000F);
|
|
acc1g[1] += yl[2 * i + 1] * (qg1[i] & 0x0F00);
|
|
acc1g[2] += yl[2 * i + 8] * (qg1[i] & 0x00F0);
|
|
acc1g[3] += yl[2 * i + 9] * (qg1[i] & 0xF000);
|
|
acc2g[0] += yh[2 * i + 0] * (qg2[i] & 0x000F);
|
|
acc2g[1] += yh[2 * i + 1] * (qg2[i] & 0x0F00);
|
|
acc2g[2] += yh[2 * i + 8] * (qg2[i] & 0x00F0);
|
|
acc2g[3] += yh[2 * i + 9] * (qg2[i] & 0xF000);
|
|
}
|
|
|
|
sumg[row] += dhg[0] * ((acc1g[0] + 1.f / 256.f * acc1g[1]) * sc8[0] +
|
|
(acc1g[2] + 1.f / 256.f * acc1g[3]) * sc8[1] * 1.f / 16.f +
|
|
(acc2g[0] + 1.f / 256.f * acc2g[1]) * sc8[4] +
|
|
(acc2g[2] + 1.f / 256.f * acc2g[3]) * sc8[5] * 1.f / 16.f) -
|
|
dhg[1] * (sumy[0] * sc8[2] + sumy[1] * sc8[3] +
|
|
sumy[2] * sc8[6] + sumy[3] * sc8[7]);
|
|
|
|
sc16[0] = scu[0] & kmask1;
|
|
sc16[1] = scu[2] & kmask1;
|
|
sc16[2] = ((scu[4] >> 0) & kmask2) | ((scu[0] & kmask3) >> 2);
|
|
sc16[3] = ((scu[4] >> 4) & kmask2) | ((scu[2] & kmask3) >> 2);
|
|
|
|
device const uint16_t *qu2 = qu1 + 32;
|
|
float4 acc1u = {0.f, 0.f, 0.f, 0.f};
|
|
float4 acc2u = {0.f, 0.f, 0.f, 0.f};
|
|
|
|
FOR_UNROLL (short i = 0; i < 4; ++i) {
|
|
acc1u[0] += yl[2 * i + 0] * (qu1[i] & 0x000F);
|
|
acc1u[1] += yl[2 * i + 1] * (qu1[i] & 0x0F00);
|
|
acc1u[2] += yl[2 * i + 8] * (qu1[i] & 0x00F0);
|
|
acc1u[3] += yl[2 * i + 9] * (qu1[i] & 0xF000);
|
|
acc2u[0] += yh[2 * i + 0] * (qu2[i] & 0x000F);
|
|
acc2u[1] += yh[2 * i + 1] * (qu2[i] & 0x0F00);
|
|
acc2u[2] += yh[2 * i + 8] * (qu2[i] & 0x00F0);
|
|
acc2u[3] += yh[2 * i + 9] * (qu2[i] & 0xF000);
|
|
}
|
|
|
|
sumu[row] += dhu[0] * ((acc1u[0] + 1.f / 256.f * acc1u[1]) * sc8[0] +
|
|
(acc1u[2] + 1.f / 256.f * acc1u[3]) * sc8[1] * 1.f / 16.f +
|
|
(acc2u[0] + 1.f / 256.f * acc2u[1]) * sc8[4] +
|
|
(acc2u[2] + 1.f / 256.f * acc2u[3]) * sc8[5] * 1.f / 16.f) -
|
|
dhu[1] * (sumy[0] * sc8[2] + sumy[1] * sc8[3] +
|
|
sumy[2] * sc8[6] + sumy[3] * sc8[7]);
|
|
|
|
qg1 += args.gate_row_bytes / 2;
|
|
scg += args.gate_row_bytes / 2;
|
|
dhg += args.gate_row_bytes / 2;
|
|
qu1 += args.up_row_bytes / 2;
|
|
scu += args.up_row_bytes / 2;
|
|
dhu += args.up_row_bytes / 2;
|
|
}
|
|
|
|
y4 += 4 * QK_K;
|
|
}
|
|
|
|
for (int row = 0;
|
|
row < N_R0 && row0 + (uint)row < args.mid_dim;
|
|
++row) {
|
|
const float g = simd_sum(sumg[row]);
|
|
const float u = simd_sum(sumu[row]);
|
|
if (tiisg == 0u) {
|
|
mid[mid_base + row0 + (uint)row] =
|
|
ds4_glm_swiglu(g, u, args.swiglu_clamp) * weights[selected_off];
|
|
}
|
|
}
|
|
|
|
(void)scratch;
|
|
}
|
|
|
|
kernel void kernel_glm_q4_K_pair_swiglu2_f32(
|
|
constant ds4_metal_glm_routed_moe_args &args,
|
|
device const char *gate,
|
|
device const char *up,
|
|
device const float *x,
|
|
device const int32_t *selected,
|
|
device const float *weights,
|
|
device float *mid,
|
|
threadgroup float *scratch [[threadgroup(0)]],
|
|
uint3 tgpig [[threadgroup_position_in_grid]],
|
|
ushort tiisg [[thread_index_in_simdgroup]],
|
|
ushort sgitg [[simdgroup_index_in_threadgroup]]) {
|
|
const uint slot = tgpig.y;
|
|
const uint token = tgpig.z;
|
|
if (slot >= args.n_expert_used || token >= args.n_tokens) return;
|
|
const uint64_t selected_off = (uint64_t)token * args.n_expert_used + slot;
|
|
const int expert = selected[selected_off];
|
|
if (!ds4_tp_owns_expert(expert, args.n_total_expert,
|
|
args.tp_rank, args.tp_world)) return;
|
|
glm_q4_K_pair_swiglu_simd_f32_impl<N_R0_GLM_Q4_PAIR2_K>(
|
|
args, gate, up, x, weights, mid, scratch,
|
|
tgpig, slot, token, selected_off,
|
|
expert - args.tp_expert_base, tiisg, sgitg);
|
|
}
|
|
|
|
kernel void kernel_glm_q4_K_addr_pair_swiglu_f32(
|
|
constant ds4_metal_glm_routed_moe_args &args,
|
|
device const uint64_t *gate_addrs,
|
|
device const uint64_t *up_addrs,
|
|
device const float *x,
|
|
device const int32_t *selected,
|
|
device const float *weights,
|
|
device float *mid,
|
|
threadgroup float *scratch [[threadgroup(0)]],
|
|
uint3 tgpig [[threadgroup_position_in_grid]],
|
|
ushort tiisg [[thread_index_in_simdgroup]],
|
|
ushort sgitg [[simdgroup_index_in_threadgroup]]) {
|
|
const uint slot = tgpig.y;
|
|
const uint token = tgpig.z;
|
|
if (slot >= args.n_expert_used || token >= args.n_tokens) return;
|
|
|
|
const uint64_t selected_off = (uint64_t)token * args.n_expert_used + slot;
|
|
const int expert = selected[selected_off];
|
|
const short NSG = 2;
|
|
const uint row0 = ((uint)tgpig.x * (uint)NSG + (uint)sgitg) * N_R0_Q4_K;
|
|
const uint64_t mid_base = (uint64_t)token * args.mid_token_stride +
|
|
(uint64_t)slot * args.mid_dim;
|
|
if (row0 >= args.mid_dim) return;
|
|
|
|
if (expert < 0 || (uint)expert >= args.n_total_expert ||
|
|
!ds4_tp_owns_expert(expert, args.n_total_expert,
|
|
args.tp_rank, args.tp_world)) {
|
|
if (tiisg == 0u) {
|
|
for (short row = 0;
|
|
row < N_R0_Q4_K && row0 + (uint)row < args.mid_dim;
|
|
row++) {
|
|
mid[mid_base + row0 + (uint)row] = 0.0f;
|
|
}
|
|
}
|
|
return;
|
|
}
|
|
const uint64_t gate_addr = gate_addrs[(uint)expert];
|
|
const uint64_t up_addr = up_addrs[(uint)expert];
|
|
if (gate_addr == 0 || up_addr == 0) {
|
|
if (tiisg == 0u) {
|
|
for (short row = 0;
|
|
row < N_R0_Q4_K && row0 + (uint)row < args.mid_dim;
|
|
row++) {
|
|
mid[mid_base + row0 + (uint)row] = 0.0f;
|
|
}
|
|
}
|
|
return;
|
|
}
|
|
|
|
ds4_metal_glm_routed_moe_args local = args;
|
|
local.n_total_expert = 1;
|
|
local.gate_expert_bytes = 0;
|
|
local.up_expert_bytes = 0;
|
|
glm_q4_K_pair_swiglu_simd_f32_impl<N_R0_Q4_K>(
|
|
local,
|
|
reinterpret_cast<device const char *>(gate_addr),
|
|
reinterpret_cast<device const char *>(up_addr),
|
|
x, weights, mid, scratch,
|
|
tgpig, slot, token, selected_off, 0, tiisg, sgitg);
|
|
}
|
|
|
|
kernel void kernel_glm_q4_K_addr_pair_swiglu_f32_masked(
|
|
constant ds4_metal_glm_routed_moe_args &args,
|
|
constant uint32_t &active_mask,
|
|
device const uint64_t *gate_addrs,
|
|
device const uint64_t *up_addrs,
|
|
device const float *x,
|
|
device const int32_t *selected,
|
|
device const float *weights,
|
|
device float *mid,
|
|
threadgroup float *scratch [[threadgroup(0)]],
|
|
uint3 tgpig [[threadgroup_position_in_grid]],
|
|
ushort tiisg [[thread_index_in_simdgroup]],
|
|
ushort sgitg [[simdgroup_index_in_threadgroup]]) {
|
|
const uint slot = tgpig.y;
|
|
const uint token = tgpig.z;
|
|
if (slot >= args.n_expert_used || token >= args.n_tokens) return;
|
|
if ((active_mask & (1u << slot)) == 0u) return;
|
|
|
|
const uint64_t selected_off = (uint64_t)token * args.n_expert_used + slot;
|
|
const int expert = selected[selected_off];
|
|
const short NSG = 2;
|
|
const uint row0 = ((uint)tgpig.x * (uint)NSG + (uint)sgitg) * N_R0_Q4_K;
|
|
const uint64_t mid_base = (uint64_t)token * args.mid_token_stride +
|
|
(uint64_t)slot * args.mid_dim;
|
|
if (row0 >= args.mid_dim) return;
|
|
|
|
if (expert < 0 || (uint)expert >= args.n_total_expert ||
|
|
!ds4_tp_owns_expert(expert, args.n_total_expert,
|
|
args.tp_rank, args.tp_world)) {
|
|
if (tiisg == 0u) {
|
|
for (short row = 0;
|
|
row < N_R0_Q4_K && row0 + (uint)row < args.mid_dim;
|
|
row++) {
|
|
mid[mid_base + row0 + (uint)row] = 0.0f;
|
|
}
|
|
}
|
|
return;
|
|
}
|
|
const uint64_t gate_addr = gate_addrs[(uint)expert];
|
|
const uint64_t up_addr = up_addrs[(uint)expert];
|
|
if (gate_addr == 0 || up_addr == 0) {
|
|
if (tiisg == 0u) {
|
|
for (short row = 0;
|
|
row < N_R0_Q4_K && row0 + (uint)row < args.mid_dim;
|
|
row++) {
|
|
mid[mid_base + row0 + (uint)row] = 0.0f;
|
|
}
|
|
}
|
|
return;
|
|
}
|
|
|
|
ds4_metal_glm_routed_moe_args local = args;
|
|
local.n_total_expert = 1;
|
|
local.gate_expert_bytes = 0;
|
|
local.up_expert_bytes = 0;
|
|
glm_q4_K_pair_swiglu_simd_f32_impl<N_R0_Q4_K>(
|
|
local,
|
|
reinterpret_cast<device const char *>(gate_addr),
|
|
reinterpret_cast<device const char *>(up_addr),
|
|
x, weights, mid, scratch,
|
|
tgpig, slot, token, selected_off, 0, tiisg, sgitg);
|
|
}
|
|
|
|
kernel void kernel_glm_q4_K_pair_swiglu4_f32(
|
|
constant ds4_metal_glm_routed_moe_args &args,
|
|
device const char *gate,
|
|
device const char *up,
|
|
device const float *x,
|
|
device const int32_t *selected,
|
|
device const float *weights,
|
|
device float *mid,
|
|
threadgroup float *scratch [[threadgroup(0)]],
|
|
uint3 tgpig [[threadgroup_position_in_grid]],
|
|
ushort tiisg [[thread_index_in_simdgroup]],
|
|
ushort sgitg [[simdgroup_index_in_threadgroup]]) {
|
|
const uint slot = tgpig.y;
|
|
const uint token = tgpig.z;
|
|
if (slot >= args.n_expert_used || token >= args.n_tokens) return;
|
|
const uint64_t selected_off = (uint64_t)token * args.n_expert_used + slot;
|
|
const int expert = selected[selected_off];
|
|
if (!ds4_tp_owns_expert(expert, args.n_total_expert,
|
|
args.tp_rank, args.tp_world)) return;
|
|
glm_q4_K_pair_swiglu_simd_f32_impl<N_R0_GLM_Q4_PAIR_K>(
|
|
args, gate, up, x, weights, mid, scratch,
|
|
tgpig, slot, token, selected_off,
|
|
expert - args.tp_expert_base, tiisg, sgitg);
|
|
}
|
|
|
|
kernel void kernel_glm_q4_K_pair_swiglu2_mapped_f32(
|
|
constant ds4_metal_glm_routed_moe_args &args,
|
|
device const char *gate,
|
|
device const char *up,
|
|
device const float *x,
|
|
device const uint32_t *htpe,
|
|
device const int32_t *hids,
|
|
device const float *weights,
|
|
device float *mid,
|
|
threadgroup float *scratch [[threadgroup(0)]],
|
|
uint3 tgpig [[threadgroup_position_in_grid]],
|
|
ushort tiisg [[thread_index_in_simdgroup]],
|
|
ushort sgitg [[simdgroup_index_in_threadgroup]]) {
|
|
const uint expert = tgpig.z;
|
|
if (expert >= args.n_total_expert) return;
|
|
if (!ds4_tp_owns_expert((int)expert, args.n_total_expert,
|
|
args.tp_rank, args.tp_world)) return;
|
|
const uint count = htpe[expert];
|
|
const uint map_base = tgpig.y * 32u;
|
|
for (uint i = 0; i < 32u; i++) {
|
|
const uint map_row = map_base + i;
|
|
if (map_row >= count) break;
|
|
const int id = hids[(uint64_t)expert * args.n_tokens + map_row];
|
|
if (id < 0) continue;
|
|
const uint token = (uint)id / args.n_expert_used;
|
|
const uint slot = (uint)id - token * args.n_expert_used;
|
|
if (slot >= args.n_expert_used || token >= args.n_tokens) continue;
|
|
const uint64_t selected_off = (uint64_t)token * args.n_expert_used + slot;
|
|
glm_q4_K_pair_swiglu_simd_f32_impl<N_R0_Q4_K>(
|
|
args, gate, up, x, weights, mid, scratch,
|
|
tgpig, slot, token, selected_off,
|
|
(int)expert - args.tp_expert_base, tiisg, sgitg);
|
|
}
|
|
}
|
|
|
|
kernel void kernel_glm_q4_K_pair_swiglu2_mapped_row_f32(
|
|
constant ds4_metal_glm_routed_moe_args &args,
|
|
device const char *gate,
|
|
device const char *up,
|
|
device const float *x,
|
|
device const uint32_t *htpe,
|
|
device const int32_t *hids,
|
|
device const float *weights,
|
|
device float *mid,
|
|
threadgroup float *scratch [[threadgroup(0)]],
|
|
uint3 tgpig [[threadgroup_position_in_grid]],
|
|
ushort tiisg [[thread_index_in_simdgroup]],
|
|
ushort sgitg [[simdgroup_index_in_threadgroup]]) {
|
|
const uint expert = tgpig.z;
|
|
const uint map_row = tgpig.y;
|
|
if (expert >= args.n_total_expert || map_row >= htpe[expert]) return;
|
|
if (!ds4_tp_owns_expert((int)expert, args.n_total_expert,
|
|
args.tp_rank, args.tp_world)) return;
|
|
const int id = hids[(uint64_t)expert * args.n_tokens + map_row];
|
|
if (id < 0) return;
|
|
const uint token = (uint)id / args.n_expert_used;
|
|
const uint slot = (uint)id - token * args.n_expert_used;
|
|
if (slot >= args.n_expert_used || token >= args.n_tokens) return;
|
|
const uint64_t selected_off = (uint64_t)token * args.n_expert_used + slot;
|
|
glm_q4_K_pair_swiglu_simd_f32_impl<N_R0_Q4_K>(
|
|
args, gate, up, x, weights, mid, scratch,
|
|
tgpig, slot, token, selected_off,
|
|
(int)expert - args.tp_expert_base, tiisg, sgitg);
|
|
}
|
|
|
|
static inline void glm_q5_K_pair_swiglu_f32_impl(
|
|
constant ds4_metal_glm_routed_moe_args &args,
|
|
device const char *gate,
|
|
device const char *up,
|
|
device const float *x,
|
|
device const float *weights,
|
|
device float *mid,
|
|
threadgroup float *scratch,
|
|
uint3 tgpig,
|
|
uint slot,
|
|
uint token,
|
|
uint64_t selected_off,
|
|
int expert,
|
|
ushort tiisg,
|
|
ushort sgitg) {
|
|
const short NSG = 2;
|
|
constexpr uint16_t kmask1 = 0x3f3f;
|
|
constexpr uint16_t kmask2 = 0x0f0f;
|
|
constexpr uint16_t kmask3 = 0xc0c0;
|
|
|
|
const uint row0 = ((uint)tgpig.x * (uint)NSG + (uint)sgitg) * N_R0_Q5_PAIR_K;
|
|
if (row0 >= args.mid_dim || slot >= args.n_expert_used || token >= args.n_tokens) return;
|
|
|
|
const uint64_t mid_base = (uint64_t)token * args.mid_token_stride +
|
|
(uint64_t)slot * args.mid_dim;
|
|
if (expert < 0 || (uint)expert >= args.n_total_expert) {
|
|
if (tiisg == 0u) {
|
|
for (short row = 0; row < N_R0_Q5_PAIR_K && row0 + (uint)row < args.mid_dim; row++) {
|
|
mid[mid_base + row0 + (uint)row] = 0.0f;
|
|
}
|
|
}
|
|
return;
|
|
}
|
|
|
|
const short ix = tiisg / 8;
|
|
const short it = tiisg % 8;
|
|
const short iq = it / 4;
|
|
const short ir = it % 4;
|
|
const uint bit0 = 2u * (uint)iq;
|
|
const uint bit1 = bit0 + 1u;
|
|
const uint bit2 = bit0 + 4u;
|
|
const uint bit3 = bit0 + 5u;
|
|
const int nb = args.in_dim / QK_K;
|
|
|
|
const uint64_t expert_gate = (uint64_t)(uint)expert * args.gate_expert_bytes;
|
|
const uint64_t expert_up = (uint64_t)(uint)expert * args.up_expert_bytes;
|
|
device const block_q5_K *gate_rows =
|
|
(device const block_q5_K *)(gate +
|
|
expert_gate + (uint64_t)row0 * args.gate_row_bytes);
|
|
device const block_q5_K *up_rows =
|
|
(device const block_q5_K *)(up +
|
|
expert_up + (uint64_t)row0 * args.up_row_bytes);
|
|
device const float *y = x + (uint64_t)token * args.in_dim;
|
|
device const float *y4 = y + ix * QK_K + 64 * iq + 8 * ir;
|
|
|
|
float sumg[N_R0_Q5_PAIR_K] = {0.f};
|
|
float sumu[N_R0_Q5_PAIR_K] = {0.f};
|
|
uint16_t sc16[4];
|
|
thread const uint8_t *sc8 = (thread const uint8_t *)sc16;
|
|
|
|
for (int ib = ix; ib < nb; ib += 4) {
|
|
float yl[16];
|
|
float yh[16];
|
|
float4 sumy = {0.f, 0.f, 0.f, 0.f};
|
|
|
|
for (short i = 0; i < 8; ++i) {
|
|
yl[i + 0] = y4[i + 0]; sumy[0] += yl[i + 0];
|
|
yl[i + 8] = y4[i + 32]; sumy[1] += yl[i + 8];
|
|
yh[i + 0] = y4[i + 128]; sumy[2] += yh[i + 0];
|
|
yh[i + 8] = y4[i + 160]; sumy[3] += yh[i + 8];
|
|
}
|
|
|
|
device const uint16_t *scg = (device const uint16_t *)gate_rows[ib].scales + iq;
|
|
device const uint16_t *qg1 = (device const uint16_t *)gate_rows[ib].qs + 16 * iq + 4 * ir;
|
|
device const uint16_t *qhg = (device const uint16_t *)gate_rows[ib].qh + 4 * ir;
|
|
device const half *dhg = &gate_rows[ib].d;
|
|
device const uint16_t *scu = (device const uint16_t *)up_rows[ib].scales + iq;
|
|
device const uint16_t *qu1 = (device const uint16_t *)up_rows[ib].qs + 16 * iq + 4 * ir;
|
|
device const uint16_t *qhu = (device const uint16_t *)up_rows[ib].qh + 4 * ir;
|
|
device const half *dhu = &up_rows[ib].d;
|
|
|
|
for (short row = 0; row < N_R0_Q5_PAIR_K && row0 + (uint)row < args.mid_dim; row++) {
|
|
sc16[0] = scg[0] & kmask1;
|
|
sc16[1] = scg[2] & kmask1;
|
|
sc16[2] = ((scg[4] >> 0) & kmask2) | ((scg[0] & kmask3) >> 2);
|
|
sc16[3] = ((scg[4] >> 4) & kmask2) | ((scg[2] & kmask3) >> 2);
|
|
|
|
device const uint16_t *qg2 = qg1 + 32;
|
|
float4 accg = {0.f, 0.f, 0.f, 0.f};
|
|
|
|
FOR_UNROLL (short i = 0; i < 4; ++i) {
|
|
const uint ql1 = (uint)qg1[i];
|
|
const uint ql2 = (uint)qg2[i];
|
|
const uint hb = (uint)qhg[i];
|
|
accg[0] += yl[2 * i + 0] *
|
|
(float)((ql1 & 0x000Fu) + (((hb >> bit0) & 1u) << 4u)) +
|
|
yl[2 * i + 1] *
|
|
(float)(((ql1 >> 8u) & 0x000Fu) + (((hb >> (bit0 + 8u)) & 1u) << 4u));
|
|
accg[1] += yl[2 * i + 8] *
|
|
(float)(((ql1 >> 4u) & 0x000Fu) + (((hb >> bit1) & 1u) << 4u)) +
|
|
yl[2 * i + 9] *
|
|
(float)(((ql1 >> 12u) & 0x000Fu) + (((hb >> (bit1 + 8u)) & 1u) << 4u));
|
|
accg[2] += yh[2 * i + 0] *
|
|
(float)((ql2 & 0x000Fu) + (((hb >> bit2) & 1u) << 4u)) +
|
|
yh[2 * i + 1] *
|
|
(float)(((ql2 >> 8u) & 0x000Fu) + (((hb >> (bit2 + 8u)) & 1u) << 4u));
|
|
accg[3] += yh[2 * i + 8] *
|
|
(float)(((ql2 >> 4u) & 0x000Fu) + (((hb >> bit3) & 1u) << 4u)) +
|
|
yh[2 * i + 9] *
|
|
(float)(((ql2 >> 12u) & 0x000Fu) + (((hb >> (bit3 + 8u)) & 1u) << 4u));
|
|
}
|
|
|
|
sumg[row] += dhg[0] * (accg[0] * sc8[0] + accg[1] * sc8[1] +
|
|
accg[2] * sc8[4] + accg[3] * sc8[5]) -
|
|
dhg[1] * (sumy[0] * sc8[2] + sumy[1] * sc8[3] +
|
|
sumy[2] * sc8[6] + sumy[3] * sc8[7]);
|
|
|
|
sc16[0] = scu[0] & kmask1;
|
|
sc16[1] = scu[2] & kmask1;
|
|
sc16[2] = ((scu[4] >> 0) & kmask2) | ((scu[0] & kmask3) >> 2);
|
|
sc16[3] = ((scu[4] >> 4) & kmask2) | ((scu[2] & kmask3) >> 2);
|
|
|
|
device const uint16_t *qu2 = qu1 + 32;
|
|
float4 accu = {0.f, 0.f, 0.f, 0.f};
|
|
|
|
FOR_UNROLL (short i = 0; i < 4; ++i) {
|
|
const uint ql1 = (uint)qu1[i];
|
|
const uint ql2 = (uint)qu2[i];
|
|
const uint hb = (uint)qhu[i];
|
|
accu[0] += yl[2 * i + 0] *
|
|
(float)((ql1 & 0x000Fu) + (((hb >> bit0) & 1u) << 4u)) +
|
|
yl[2 * i + 1] *
|
|
(float)(((ql1 >> 8u) & 0x000Fu) + (((hb >> (bit0 + 8u)) & 1u) << 4u));
|
|
accu[1] += yl[2 * i + 8] *
|
|
(float)(((ql1 >> 4u) & 0x000Fu) + (((hb >> bit1) & 1u) << 4u)) +
|
|
yl[2 * i + 9] *
|
|
(float)(((ql1 >> 12u) & 0x000Fu) + (((hb >> (bit1 + 8u)) & 1u) << 4u));
|
|
accu[2] += yh[2 * i + 0] *
|
|
(float)((ql2 & 0x000Fu) + (((hb >> bit2) & 1u) << 4u)) +
|
|
yh[2 * i + 1] *
|
|
(float)(((ql2 >> 8u) & 0x000Fu) + (((hb >> (bit2 + 8u)) & 1u) << 4u));
|
|
accu[3] += yh[2 * i + 8] *
|
|
(float)(((ql2 >> 4u) & 0x000Fu) + (((hb >> bit3) & 1u) << 4u)) +
|
|
yh[2 * i + 9] *
|
|
(float)(((ql2 >> 12u) & 0x000Fu) + (((hb >> (bit3 + 8u)) & 1u) << 4u));
|
|
}
|
|
|
|
sumu[row] += dhu[0] * (accu[0] * sc8[0] + accu[1] * sc8[1] +
|
|
accu[2] * sc8[4] + accu[3] * sc8[5]) -
|
|
dhu[1] * (sumy[0] * sc8[2] + sumy[1] * sc8[3] +
|
|
sumy[2] * sc8[6] + sumy[3] * sc8[7]);
|
|
|
|
qg1 += args.gate_row_bytes / 2;
|
|
qhg += args.gate_row_bytes / 2;
|
|
scg += args.gate_row_bytes / 2;
|
|
dhg += args.gate_row_bytes / 2;
|
|
qu1 += args.up_row_bytes / 2;
|
|
qhu += args.up_row_bytes / 2;
|
|
scu += args.up_row_bytes / 2;
|
|
dhu += args.up_row_bytes / 2;
|
|
}
|
|
|
|
y4 += 4 * QK_K;
|
|
}
|
|
|
|
for (short row = 0; row < N_R0_Q5_PAIR_K && row0 + (uint)row < args.mid_dim; row++) {
|
|
const float g = simd_sum(sumg[row]);
|
|
const float u = simd_sum(sumu[row]);
|
|
if (tiisg == 0u) {
|
|
mid[mid_base + row0 + (uint)row] =
|
|
ds4_glm_swiglu(g, u, args.swiglu_clamp) * weights[selected_off];
|
|
}
|
|
}
|
|
|
|
(void)scratch;
|
|
}
|
|
|
|
kernel void kernel_glm_q5_K_pair_swiglu_f32(
|
|
constant ds4_metal_glm_routed_moe_args &args,
|
|
device const char *gate,
|
|
device const char *up,
|
|
device const float *x,
|
|
device const int32_t *selected,
|
|
device const float *weights,
|
|
device float *mid,
|
|
threadgroup float *scratch [[threadgroup(0)]],
|
|
uint3 tgpig [[threadgroup_position_in_grid]],
|
|
ushort tiisg [[thread_index_in_simdgroup]],
|
|
ushort sgitg [[simdgroup_index_in_threadgroup]]) {
|
|
const uint slot = tgpig.y;
|
|
const uint token = tgpig.z;
|
|
if (slot >= args.n_expert_used || token >= args.n_tokens) return;
|
|
const uint64_t selected_off = (uint64_t)token * args.n_expert_used + slot;
|
|
const int expert = selected[selected_off];
|
|
glm_q5_K_pair_swiglu_f32_impl(
|
|
args, gate, up, x, weights, mid, scratch,
|
|
tgpig, slot, token, selected_off, expert, tiisg, sgitg);
|
|
}
|
|
|
|
kernel void kernel_glm_q5_K_pair_swiglu_mapped_f32(
|
|
constant ds4_metal_glm_routed_moe_args &args,
|
|
device const char *gate,
|
|
device const char *up,
|
|
device const float *x,
|
|
device const uint32_t *htpe,
|
|
device const int32_t *hids,
|
|
device const float *weights,
|
|
device float *mid,
|
|
threadgroup float *scratch [[threadgroup(0)]],
|
|
uint3 tgpig [[threadgroup_position_in_grid]],
|
|
ushort tiisg [[thread_index_in_simdgroup]],
|
|
ushort sgitg [[simdgroup_index_in_threadgroup]]) {
|
|
const uint expert = tgpig.z;
|
|
if (expert >= args.n_total_expert) return;
|
|
const uint count = htpe[expert];
|
|
const uint map_base = tgpig.y * 32u;
|
|
for (uint i = 0; i < 32u; i++) {
|
|
const uint map_row = map_base + i;
|
|
if (map_row >= count) break;
|
|
const int id = hids[(uint64_t)expert * args.n_tokens + map_row];
|
|
if (id < 0) continue;
|
|
const uint token = (uint)id / args.n_expert_used;
|
|
const uint slot = (uint)id - token * args.n_expert_used;
|
|
if (slot >= args.n_expert_used || token >= args.n_tokens) continue;
|
|
const uint64_t selected_off = (uint64_t)token * args.n_expert_used + slot;
|
|
glm_q5_K_pair_swiglu_f32_impl(
|
|
args, gate, up, x, weights, mid, scratch,
|
|
tgpig, slot, token, selected_off, (int)expert, tiisg, sgitg);
|
|
}
|
|
}
|
|
|
|
kernel void kernel_glm_q5_K_pair_swiglu_mapped_row_f32(
|
|
constant ds4_metal_glm_routed_moe_args &args,
|
|
device const char *gate,
|
|
device const char *up,
|
|
device const float *x,
|
|
device const uint32_t *htpe,
|
|
device const int32_t *hids,
|
|
device const float *weights,
|
|
device float *mid,
|
|
threadgroup float *scratch [[threadgroup(0)]],
|
|
uint3 tgpig [[threadgroup_position_in_grid]],
|
|
ushort tiisg [[thread_index_in_simdgroup]],
|
|
ushort sgitg [[simdgroup_index_in_threadgroup]]) {
|
|
const uint expert = tgpig.z;
|
|
const uint map_row = tgpig.y;
|
|
if (expert >= args.n_total_expert || map_row >= htpe[expert]) return;
|
|
const int id = hids[(uint64_t)expert * args.n_tokens + map_row];
|
|
if (id < 0) return;
|
|
const uint token = (uint)id / args.n_expert_used;
|
|
const uint slot = (uint)id - token * args.n_expert_used;
|
|
if (slot >= args.n_expert_used || token >= args.n_tokens) return;
|
|
const uint64_t selected_off = (uint64_t)token * args.n_expert_used + slot;
|
|
glm_q5_K_pair_swiglu_f32_impl(
|
|
args, gate, up, x, weights, mid, scratch,
|
|
tgpig, slot, token, selected_off, (int)expert, tiisg, sgitg);
|
|
}
|
|
|
|
kernel void kernel_glm_q5_K_down_f32(
|
|
constant ds4_metal_glm_routed_moe_args &args,
|
|
device const char *down,
|
|
device const int32_t *selected,
|
|
device const float *mid,
|
|
device float *out,
|
|
uint3 tgpig [[threadgroup_position_in_grid]],
|
|
ushort tiisg [[thread_index_in_simdgroup]],
|
|
ushort sgitg [[simdgroup_index_in_threadgroup]]) {
|
|
const short NSG = 2;
|
|
constexpr uint16_t kmask1 = 0x3f3f;
|
|
constexpr uint16_t kmask2 = 0x0f0f;
|
|
constexpr uint16_t kmask3 = 0xc0c0;
|
|
|
|
const uint row0 = ((uint)tgpig.x * (uint)NSG + (uint)sgitg) * N_R0_Q5_K;
|
|
const uint token = tgpig.y;
|
|
if (row0 >= args.out_dim || token >= args.n_tokens) return;
|
|
|
|
const short ix = tiisg / 8;
|
|
const short it = tiisg % 8;
|
|
const short iq = it / 4;
|
|
const short ir = it % 4;
|
|
const uint bit0 = 2u * (uint)iq;
|
|
const uint bit1 = bit0 + 1u;
|
|
const uint bit2 = bit0 + 4u;
|
|
const uint bit3 = bit0 + 5u;
|
|
const int nb = args.mid_dim / QK_K;
|
|
|
|
float sumf[N_R0_Q5_K] = {0.f};
|
|
uint16_t sc16[4];
|
|
thread const uint8_t *sc8 = (thread const uint8_t *)sc16;
|
|
const uint64_t selected_base = (uint64_t)token * args.n_expert_used;
|
|
const uint64_t mid_base = (uint64_t)token * args.mid_token_stride;
|
|
for (uint slot = 0; slot < args.n_expert_used; slot++) {
|
|
const int expert = selected[selected_base + slot];
|
|
if (expert < 0 || (uint)expert >= args.n_total_expert) continue;
|
|
device const block_q5_K *x =
|
|
(device const block_q5_K *)(down +
|
|
(uint64_t)(uint)expert * args.down_expert_bytes +
|
|
(uint64_t)row0 * args.down_row_bytes);
|
|
device const float *y = mid + mid_base + (uint64_t)slot * args.mid_dim;
|
|
device const float *y4 = y + ix * QK_K + 64 * iq + 8 * ir;
|
|
|
|
for (int ib = ix; ib < nb; ib += 4) {
|
|
float yl[16];
|
|
float yh[16];
|
|
float4 sumy = {0.f, 0.f, 0.f, 0.f};
|
|
|
|
for (short i = 0; i < 8; ++i) {
|
|
yl[i + 0] = y4[i + 0]; sumy[0] += yl[i + 0];
|
|
yl[i + 8] = y4[i + 32]; sumy[1] += yl[i + 8];
|
|
yh[i + 0] = y4[i + 128]; sumy[2] += yh[i + 0];
|
|
yh[i + 8] = y4[i + 160]; sumy[3] += yh[i + 8];
|
|
}
|
|
|
|
device const uint16_t *sc = (device const uint16_t *)x[ib].scales + iq;
|
|
device const uint16_t *q1 = (device const uint16_t *)x[ib].qs + 16 * iq + 4 * ir;
|
|
device const uint16_t *qh = (device const uint16_t *)x[ib].qh + 4 * ir;
|
|
device const half *dh = &x[ib].d;
|
|
|
|
for (short row = 0; row < N_R0_Q5_K && row0 + (uint)row < args.out_dim; row++) {
|
|
sc16[0] = sc[0] & kmask1;
|
|
sc16[1] = sc[2] & kmask1;
|
|
sc16[2] = ((sc[4] >> 0) & kmask2) | ((sc[0] & kmask3) >> 2);
|
|
sc16[3] = ((sc[4] >> 4) & kmask2) | ((sc[2] & kmask3) >> 2);
|
|
|
|
device const uint16_t *q2 = q1 + 32;
|
|
float4 acc = {0.f, 0.f, 0.f, 0.f};
|
|
|
|
FOR_UNROLL (short i = 0; i < 4; ++i) {
|
|
const uint ql1 = (uint)q1[i];
|
|
const uint ql2 = (uint)q2[i];
|
|
const uint hb = (uint)qh[i];
|
|
acc[0] += yl[2 * i + 0] *
|
|
(float)((ql1 & 0x000Fu) + (((hb >> bit0) & 1u) << 4u)) +
|
|
yl[2 * i + 1] *
|
|
(float)(((ql1 >> 8u) & 0x000Fu) + (((hb >> (bit0 + 8u)) & 1u) << 4u));
|
|
acc[1] += yl[2 * i + 8] *
|
|
(float)(((ql1 >> 4u) & 0x000Fu) + (((hb >> bit1) & 1u) << 4u)) +
|
|
yl[2 * i + 9] *
|
|
(float)(((ql1 >> 12u) & 0x000Fu) + (((hb >> (bit1 + 8u)) & 1u) << 4u));
|
|
acc[2] += yh[2 * i + 0] *
|
|
(float)((ql2 & 0x000Fu) + (((hb >> bit2) & 1u) << 4u)) +
|
|
yh[2 * i + 1] *
|
|
(float)(((ql2 >> 8u) & 0x000Fu) + (((hb >> (bit2 + 8u)) & 1u) << 4u));
|
|
acc[3] += yh[2 * i + 8] *
|
|
(float)(((ql2 >> 4u) & 0x000Fu) + (((hb >> bit3) & 1u) << 4u)) +
|
|
yh[2 * i + 9] *
|
|
(float)(((ql2 >> 12u) & 0x000Fu) + (((hb >> (bit3 + 8u)) & 1u) << 4u));
|
|
}
|
|
|
|
sumf[row] += dh[0] * (acc[0] * sc8[0] + acc[1] * sc8[1] +
|
|
acc[2] * sc8[4] + acc[3] * sc8[5]) -
|
|
dh[1] * (sumy[0] * sc8[2] + sumy[1] * sc8[3] +
|
|
sumy[2] * sc8[6] + sumy[3] * sc8[7]);
|
|
|
|
q1 += args.down_row_bytes / 2;
|
|
qh += args.down_row_bytes / 2;
|
|
sc += args.down_row_bytes / 2;
|
|
dh += args.down_row_bytes / 2;
|
|
}
|
|
|
|
y4 += 4 * QK_K;
|
|
}
|
|
}
|
|
|
|
for (short row = 0; row < N_R0_Q5_K && row0 + (uint)row < args.out_dim; row++) {
|
|
const float sum_all = simd_sum(sumf[row]);
|
|
if (tiisg == 0u) {
|
|
out[(uint64_t)token * args.out_dim + row0 + (uint)row] = sum_all;
|
|
}
|
|
}
|
|
}
|
|
|
|
kernel void kernel_glm_q2_K_down_f32(
|
|
constant ds4_metal_glm_routed_moe_args &args,
|
|
device const char *down,
|
|
device const int32_t *selected,
|
|
device const float *mid,
|
|
device float *out,
|
|
threadgroup float *scratch [[threadgroup(0)]],
|
|
uint3 tgpig [[threadgroup_position_in_grid]],
|
|
ushort tiisg [[thread_index_in_simdgroup]],
|
|
ushort sgitg [[simdgroup_index_in_threadgroup]]) {
|
|
const short NSG = 2;
|
|
const uint row0 = ((uint)tgpig.x * (uint)NSG + (uint)sgitg) * N_R0_Q2_K;
|
|
const uint token = tgpig.y;
|
|
if (row0 >= args.out_dim || token >= args.n_tokens) return;
|
|
|
|
const short ix = tiisg / 8;
|
|
const short it = tiisg % 8;
|
|
const short iq = it / 4;
|
|
const short ir = it % 4;
|
|
const short is = (8 * ir) / 16;
|
|
const int nb = args.mid_dim / QK_K;
|
|
float sumf[N_R0_Q2_K] = {0.f};
|
|
const uint64_t selected_base = (uint64_t)token * args.n_expert_used;
|
|
const uint64_t mid_base = (uint64_t)token * args.mid_token_stride;
|
|
for (uint slot = 0; slot < args.n_expert_used; slot++) {
|
|
const int expert = selected[selected_base + slot];
|
|
if (expert < 0 || (uint)expert >= args.n_total_expert) continue;
|
|
if (!ds4_tp_owns_expert(expert, args.n_total_expert,
|
|
args.tp_rank, args.tp_world)) continue;
|
|
device const block_q2_K *x =
|
|
(device const block_q2_K *)(down +
|
|
(uint64_t)(uint)(expert - args.tp_expert_base) * args.down_expert_bytes +
|
|
(uint64_t)row0 * args.down_row_bytes);
|
|
device const float *y = mid + mid_base + (uint64_t)slot * args.mid_dim;
|
|
device const float *y4 = y + ix * QK_K + 128 * iq + 8 * ir;
|
|
|
|
for (int ib = ix; ib < nb; ib += 4) {
|
|
float yl[32];
|
|
float4 sumy = {0.f, 0.f, 0.f, 0.f};
|
|
for (short i = 0; i < 8; ++i) {
|
|
yl[i + 0] = y4[i + 0]; sumy[0] += yl[i + 0];
|
|
yl[i + 8] = y4[i + 32]; sumy[1] += yl[i + 8];
|
|
yl[i + 16] = y4[i + 64]; sumy[2] += yl[i + 16];
|
|
yl[i + 24] = y4[i + 96]; sumy[3] += yl[i + 24];
|
|
}
|
|
|
|
device const uint8_t *sc = (device const uint8_t *)x[ib].scales + 8 * iq + is;
|
|
device const uint16_t *qs = (device const uint16_t *)x[ib].qs + 16 * iq + 4 * ir;
|
|
device const half *dh = &x[ib].d;
|
|
for (short row = 0; row < N_R0_Q2_K && row0 + (uint)row < args.out_dim; row++) {
|
|
float4 acc1 = {0.f, 0.f, 0.f, 0.f};
|
|
float4 acc2 = {0.f, 0.f, 0.f, 0.f};
|
|
|
|
for (int i = 0; i < 8; i += 2) {
|
|
acc1[0] += yl[i + 0] * (qs[i / 2] & 0x0003);
|
|
acc2[0] += yl[i + 1] * (qs[i / 2] & 0x0300);
|
|
acc1[1] += yl[i + 8] * (qs[i / 2] & 0x000c);
|
|
acc2[1] += yl[i + 9] * (qs[i / 2] & 0x0c00);
|
|
acc1[2] += yl[i + 16] * (qs[i / 2] & 0x0030);
|
|
acc2[2] += yl[i + 17] * (qs[i / 2] & 0x3000);
|
|
acc1[3] += yl[i + 24] * (qs[i / 2] & 0x00c0);
|
|
acc2[3] += yl[i + 25] * (qs[i / 2] & 0xc000);
|
|
}
|
|
|
|
const float d = dh[0];
|
|
const float m = dh[1] * 1.f / 16.f;
|
|
sumf[row] += d * ((acc1[0] + 1.f / 256.f * acc2[0]) * (sc[0] & 0xF) * 1.f / 1.f +
|
|
(acc1[1] + 1.f / 256.f * acc2[1]) * (sc[2] & 0xF) * 1.f / 4.f +
|
|
(acc1[2] + 1.f / 256.f * acc2[2]) * (sc[4] & 0xF) * 1.f / 16.f +
|
|
(acc1[3] + 1.f / 256.f * acc2[3]) * (sc[6] & 0xF) * 1.f / 64.f) -
|
|
m * (sumy[0] * (sc[0] & 0xF0) + sumy[1] * (sc[2] & 0xF0) +
|
|
sumy[2] * (sc[4] & 0xF0) + sumy[3] * (sc[6] & 0xF0));
|
|
|
|
qs += args.down_row_bytes / 2;
|
|
sc += args.down_row_bytes;
|
|
dh += args.down_row_bytes / 2;
|
|
}
|
|
|
|
y4 += 4 * QK_K;
|
|
}
|
|
}
|
|
|
|
for (short row = 0; row < N_R0_Q2_K && row0 + (uint)row < args.out_dim; row++) {
|
|
const float sum_all = simd_sum(sumf[row]);
|
|
if (tiisg == 0u) {
|
|
out[(uint64_t)token * args.out_dim + row0 + (uint)row] = sum_all;
|
|
}
|
|
}
|
|
|
|
(void)scratch;
|
|
}
|
|
|
|
kernel void kernel_glm_q2_K_addr_down_f32(
|
|
constant ds4_metal_glm_routed_moe_args &args,
|
|
device const uint64_t *down_addrs,
|
|
device const int32_t *selected,
|
|
device const float *mid,
|
|
device float *out,
|
|
threadgroup float *scratch [[threadgroup(0)]],
|
|
uint3 tgpig [[threadgroup_position_in_grid]],
|
|
ushort tiisg [[thread_index_in_simdgroup]],
|
|
ushort sgitg [[simdgroup_index_in_threadgroup]]) {
|
|
const short NSG = 2;
|
|
const uint row0 = ((uint)tgpig.x * (uint)NSG + (uint)sgitg) * N_R0_Q2_K;
|
|
const uint token = tgpig.y;
|
|
if (row0 >= args.out_dim || token >= args.n_tokens) return;
|
|
|
|
const short ix = tiisg / 8;
|
|
const short it = tiisg % 8;
|
|
const short iq = it / 4;
|
|
const short ir = it % 4;
|
|
const short is = (8 * ir) / 16;
|
|
const int nb = args.mid_dim / QK_K;
|
|
float sumf[N_R0_Q2_K] = {0.f};
|
|
const uint64_t selected_base = (uint64_t)token * args.n_expert_used;
|
|
const uint64_t mid_base = (uint64_t)token * args.mid_token_stride;
|
|
for (uint slot = 0; slot < args.n_expert_used; slot++) {
|
|
const int expert = selected[selected_base + slot];
|
|
if (expert < 0 || (uint)expert >= args.n_total_expert) continue;
|
|
if (!ds4_tp_owns_expert(expert, args.n_total_expert,
|
|
args.tp_rank, args.tp_world)) continue;
|
|
const uint64_t down_addr = down_addrs[(uint)expert];
|
|
if (down_addr == 0) continue;
|
|
device const block_q2_K *x =
|
|
(device const block_q2_K *)(reinterpret_cast<device const char *>(down_addr) +
|
|
(uint64_t)row0 * args.down_row_bytes);
|
|
device const float *y = mid + mid_base + (uint64_t)slot * args.mid_dim;
|
|
device const float *y4 = y + ix * QK_K + 128 * iq + 8 * ir;
|
|
|
|
for (int ib = ix; ib < nb; ib += 4) {
|
|
float yl[32];
|
|
float4 sumy = {0.f, 0.f, 0.f, 0.f};
|
|
for (short i = 0; i < 8; ++i) {
|
|
yl[i + 0] = y4[i + 0]; sumy[0] += yl[i + 0];
|
|
yl[i + 8] = y4[i + 32]; sumy[1] += yl[i + 8];
|
|
yl[i + 16] = y4[i + 64]; sumy[2] += yl[i + 16];
|
|
yl[i + 24] = y4[i + 96]; sumy[3] += yl[i + 24];
|
|
}
|
|
|
|
device const uint8_t *sc = (device const uint8_t *)x[ib].scales + 8 * iq + is;
|
|
device const uint16_t *qs = (device const uint16_t *)x[ib].qs + 16 * iq + 4 * ir;
|
|
device const half *dh = &x[ib].d;
|
|
for (short row = 0; row < N_R0_Q2_K && row0 + (uint)row < args.out_dim; row++) {
|
|
float4 acc1 = {0.f, 0.f, 0.f, 0.f};
|
|
float4 acc2 = {0.f, 0.f, 0.f, 0.f};
|
|
|
|
for (int i = 0; i < 8; i += 2) {
|
|
acc1[0] += yl[i + 0] * (qs[i / 2] & 0x0003);
|
|
acc2[0] += yl[i + 1] * (qs[i / 2] & 0x0300);
|
|
acc1[1] += yl[i + 8] * (qs[i / 2] & 0x000c);
|
|
acc2[1] += yl[i + 9] * (qs[i / 2] & 0x0c00);
|
|
acc1[2] += yl[i + 16] * (qs[i / 2] & 0x0030);
|
|
acc2[2] += yl[i + 17] * (qs[i / 2] & 0x3000);
|
|
acc1[3] += yl[i + 24] * (qs[i / 2] & 0x00c0);
|
|
acc2[3] += yl[i + 25] * (qs[i / 2] & 0xc000);
|
|
}
|
|
|
|
const float d = dh[0];
|
|
const float m = dh[1] * 1.f / 16.f;
|
|
sumf[row] += d * ((acc1[0] + 1.f / 256.f * acc2[0]) * (sc[0] & 0xF) * 1.f / 1.f +
|
|
(acc1[1] + 1.f / 256.f * acc2[1]) * (sc[2] & 0xF) * 1.f / 4.f +
|
|
(acc1[2] + 1.f / 256.f * acc2[2]) * (sc[4] & 0xF) * 1.f / 16.f +
|
|
(acc1[3] + 1.f / 256.f * acc2[3]) * (sc[6] & 0xF) * 1.f / 64.f) -
|
|
m * (sumy[0] * (sc[0] & 0xF0) + sumy[1] * (sc[2] & 0xF0) +
|
|
sumy[2] * (sc[4] & 0xF0) + sumy[3] * (sc[6] & 0xF0));
|
|
|
|
qs += args.down_row_bytes / 2;
|
|
sc += args.down_row_bytes;
|
|
dh += args.down_row_bytes / 2;
|
|
}
|
|
|
|
y4 += 4 * QK_K;
|
|
}
|
|
}
|
|
|
|
for (short row = 0; row < N_R0_Q2_K && row0 + (uint)row < args.out_dim; row++) {
|
|
const float sum_all = simd_sum(sumf[row]);
|
|
if (tiisg == 0u) {
|
|
out[(uint64_t)token * args.out_dim + row0 + (uint)row] = sum_all;
|
|
}
|
|
}
|
|
|
|
(void)scratch;
|
|
}
|
|
|
|
kernel void kernel_glm_q4_K_down_f32(
|
|
constant ds4_metal_glm_routed_moe_args &args,
|
|
device const char *down,
|
|
device const int32_t *selected,
|
|
device const float *mid,
|
|
device float *out,
|
|
threadgroup float *scratch [[threadgroup(0)]],
|
|
uint3 tgpig [[threadgroup_position_in_grid]],
|
|
uint tid [[thread_index_in_threadgroup]]) {
|
|
const uint ntg = 256u;
|
|
const uint row = tgpig.x;
|
|
const uint token = tgpig.y;
|
|
if (row >= args.out_dim || token >= args.n_tokens) return;
|
|
|
|
float acc = 0.0f;
|
|
const uint64_t selected_base = (uint64_t)token * args.n_expert_used;
|
|
const uint64_t mid_base = (uint64_t)token * args.mid_token_stride;
|
|
for (uint slot = 0; slot < args.n_expert_used; slot++) {
|
|
const int expert = selected[selected_base + slot];
|
|
if (!ds4_tp_owns_expert(expert, args.n_total_expert,
|
|
args.tp_rank, args.tp_world)) continue;
|
|
device const block_q4_K *down_row =
|
|
(device const block_q4_K *)(down +
|
|
(uint64_t)(uint)(expert - args.tp_expert_base) *
|
|
args.down_expert_bytes +
|
|
(uint64_t)row * args.down_row_bytes);
|
|
device const float *slot_mid = mid + mid_base + (uint64_t)slot * args.mid_dim;
|
|
for (uint k = tid; k < args.mid_dim; k += ntg) {
|
|
acc += ds4_glm_q4_K_value(down_row, k) * slot_mid[k];
|
|
}
|
|
}
|
|
|
|
scratch[tid] = acc;
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
for (uint stride = ntg >> 1u; stride > 0u; stride >>= 1u) {
|
|
if (tid < stride) scratch[tid] += scratch[tid + stride];
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
}
|
|
if (tid == 0u) {
|
|
out[(uint64_t)token * args.out_dim + row] = scratch[0];
|
|
}
|
|
}
|
|
|
|
kernel void kernel_glm_q4_K_addr_down_f32(
|
|
constant ds4_metal_glm_routed_moe_args &args,
|
|
device const uint64_t *down_addrs,
|
|
device const int32_t *selected,
|
|
device const float *mid,
|
|
device float *out,
|
|
threadgroup float *scratch [[threadgroup(0)]],
|
|
uint3 tgpig [[threadgroup_position_in_grid]],
|
|
uint tid [[thread_index_in_threadgroup]]) {
|
|
const uint ntg = 256u;
|
|
const uint row = tgpig.x;
|
|
const uint token = tgpig.y;
|
|
if (row >= args.out_dim || token >= args.n_tokens) return;
|
|
|
|
float acc = 0.0f;
|
|
const uint64_t selected_base = (uint64_t)token * args.n_expert_used;
|
|
const uint64_t mid_base = (uint64_t)token * args.mid_token_stride;
|
|
for (uint slot = 0; slot < args.n_expert_used; slot++) {
|
|
const int expert = selected[selected_base + slot];
|
|
if (expert < 0 || (uint)expert >= args.n_total_expert) continue;
|
|
if (!ds4_tp_owns_expert(expert, args.n_total_expert,
|
|
args.tp_rank, args.tp_world)) continue;
|
|
const uint64_t down_addr = down_addrs[(uint)expert];
|
|
if (down_addr == 0) continue;
|
|
device const block_q4_K *down_row =
|
|
(device const block_q4_K *)(reinterpret_cast<device const char *>(down_addr) +
|
|
(uint64_t)row * args.down_row_bytes);
|
|
device const float *slot_mid = mid + mid_base + (uint64_t)slot * args.mid_dim;
|
|
for (uint k = tid; k < args.mid_dim; k += ntg) {
|
|
acc += ds4_glm_q4_K_value(down_row, k) * slot_mid[k];
|
|
}
|
|
}
|
|
|
|
scratch[tid] = acc;
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
for (uint stride = ntg >> 1u; stride > 0u; stride >>= 1u) {
|
|
if (tid < stride) scratch[tid] += scratch[tid + stride];
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
}
|
|
if (tid == 0u) {
|
|
out[(uint64_t)token * args.out_dim + row] = scratch[0];
|
|
}
|
|
}
|
|
|
|
kernel void kernel_glm_q4_K_down_simd_f32(
|
|
constant ds4_metal_glm_routed_moe_args &args,
|
|
device const char *down,
|
|
device const int32_t *selected,
|
|
device const float *mid,
|
|
device float *out,
|
|
uint3 tgpig [[threadgroup_position_in_grid]],
|
|
ushort tiisg [[thread_index_in_simdgroup]],
|
|
ushort sgitg [[simdgroup_index_in_threadgroup]]) {
|
|
const short NSG = 2;
|
|
const short nr0 = N_R0_Q4_K;
|
|
const int nb = args.mid_dim / QK_K;
|
|
const uint row0 = ((uint)tgpig.x * (uint)NSG + (uint)sgitg) * (uint)nr0;
|
|
const uint token = tgpig.y;
|
|
if (row0 >= args.out_dim || token >= args.n_tokens) return;
|
|
|
|
constexpr uint16_t kmask1 = 0x3f3f;
|
|
constexpr uint16_t kmask2 = 0x0f0f;
|
|
constexpr uint16_t kmask3 = 0xc0c0;
|
|
|
|
const short ix = tiisg / 8;
|
|
const short it = tiisg % 8;
|
|
const short iq = it / 4;
|
|
const short ir = it % 4;
|
|
|
|
float sumf[N_R0_Q4_K] = {0.f};
|
|
uint16_t sc16[4];
|
|
thread const uint8_t *sc8 = (thread const uint8_t *)sc16;
|
|
|
|
const uint64_t selected_base = (uint64_t)token * args.n_expert_used;
|
|
const uint64_t mid_base = (uint64_t)token * args.mid_token_stride;
|
|
for (uint slot = 0; slot < args.n_expert_used; slot++) {
|
|
const int expert = selected[selected_base + slot];
|
|
if (!ds4_tp_owns_expert(expert, args.n_total_expert,
|
|
args.tp_rank, args.tp_world)) continue;
|
|
|
|
device const block_q4_K *x =
|
|
(device const block_q4_K *)(down +
|
|
(uint64_t)(uint)(expert - args.tp_expert_base) *
|
|
args.down_expert_bytes +
|
|
(uint64_t)row0 * args.down_row_bytes);
|
|
device const float *y = mid + mid_base + (uint64_t)slot * args.mid_dim;
|
|
device const float *y4 = y + ix * QK_K + 64 * iq + 8 * ir;
|
|
|
|
for (int ib = ix; ib < nb; ib += 4) {
|
|
float yl[16];
|
|
float yh[16];
|
|
float4 sumy = {0.f, 0.f, 0.f, 0.f};
|
|
|
|
for (short i = 0; i < 8; ++i) {
|
|
yl[i + 0] = y4[i + 0]; sumy[0] += yl[i + 0];
|
|
yl[i + 8] = y4[i + 32]; sumy[1] += yl[i + 8];
|
|
yh[i + 0] = y4[i + 128]; sumy[2] += yh[i + 0];
|
|
yh[i + 8] = y4[i + 160]; sumy[3] += yh[i + 8];
|
|
}
|
|
|
|
device const uint16_t *sc = (device const uint16_t *)x[ib].scales + iq;
|
|
device const uint16_t *q1 = (device const uint16_t *)x[ib].qs + 16 * iq + 4 * ir;
|
|
device const half *dh = &x[ib].d;
|
|
|
|
for (short row = 0; row < nr0 && row0 + (uint)row < args.out_dim; row++) {
|
|
sc16[0] = sc[0] & kmask1;
|
|
sc16[1] = sc[2] & kmask1;
|
|
sc16[2] = ((sc[4] >> 0) & kmask2) | ((sc[0] & kmask3) >> 2);
|
|
sc16[3] = ((sc[4] >> 4) & kmask2) | ((sc[2] & kmask3) >> 2);
|
|
|
|
device const uint16_t *q2 = q1 + 32;
|
|
float4 acc1 = {0.f, 0.f, 0.f, 0.f};
|
|
float4 acc2 = {0.f, 0.f, 0.f, 0.f};
|
|
|
|
FOR_UNROLL (short i = 0; i < 4; ++i) {
|
|
acc1[0] += yl[2 * i + 0] * (q1[i] & 0x000F);
|
|
acc1[1] += yl[2 * i + 1] * (q1[i] & 0x0F00);
|
|
acc1[2] += yl[2 * i + 8] * (q1[i] & 0x00F0);
|
|
acc1[3] += yl[2 * i + 9] * (q1[i] & 0xF000);
|
|
acc2[0] += yh[2 * i + 0] * (q2[i] & 0x000F);
|
|
acc2[1] += yh[2 * i + 1] * (q2[i] & 0x0F00);
|
|
acc2[2] += yh[2 * i + 8] * (q2[i] & 0x00F0);
|
|
acc2[3] += yh[2 * i + 9] * (q2[i] & 0xF000);
|
|
}
|
|
|
|
sumf[row] += dh[0] * ((acc1[0] + 1.f / 256.f * acc1[1]) * sc8[0] +
|
|
(acc1[2] + 1.f / 256.f * acc1[3]) * sc8[1] * 1.f / 16.f +
|
|
(acc2[0] + 1.f / 256.f * acc2[1]) * sc8[4] +
|
|
(acc2[2] + 1.f / 256.f * acc2[3]) * sc8[5] * 1.f / 16.f) -
|
|
dh[1] * (sumy[0] * sc8[2] + sumy[1] * sc8[3] +
|
|
sumy[2] * sc8[6] + sumy[3] * sc8[7]);
|
|
|
|
q1 += args.down_row_bytes / 2;
|
|
sc += args.down_row_bytes / 2;
|
|
dh += args.down_row_bytes / 2;
|
|
}
|
|
|
|
y4 += 4 * QK_K;
|
|
}
|
|
}
|
|
|
|
for (short row = 0; row < nr0 && row0 + (uint)row < args.out_dim; row++) {
|
|
const float sum_all = simd_sum(sumf[row]);
|
|
if (tiisg == 0u) {
|
|
out[(uint64_t)token * args.out_dim + row0 + (uint)row] = sum_all;
|
|
}
|
|
}
|
|
}
|
|
|
|
kernel void kernel_glm_q4_K_addr_down_simd_f32(
|
|
constant ds4_metal_glm_routed_moe_args &args,
|
|
device const uint64_t *down_addrs,
|
|
device const int32_t *selected,
|
|
device const float *mid,
|
|
device float *out,
|
|
uint3 tgpig [[threadgroup_position_in_grid]],
|
|
ushort tiisg [[thread_index_in_simdgroup]],
|
|
ushort sgitg [[simdgroup_index_in_threadgroup]]) {
|
|
const short NSG = 2;
|
|
const short nr0 = N_R0_Q4_K;
|
|
const int nb = args.mid_dim / QK_K;
|
|
const uint row0 = ((uint)tgpig.x * (uint)NSG + (uint)sgitg) * (uint)nr0;
|
|
const uint token = tgpig.y;
|
|
if (row0 >= args.out_dim || token >= args.n_tokens) return;
|
|
|
|
constexpr uint16_t kmask1 = 0x3f3f;
|
|
constexpr uint16_t kmask2 = 0x0f0f;
|
|
constexpr uint16_t kmask3 = 0xc0c0;
|
|
|
|
const short ix = tiisg / 8;
|
|
const short it = tiisg % 8;
|
|
const short iq = it / 4;
|
|
const short ir = it % 4;
|
|
|
|
float sumf[N_R0_Q4_K] = {0.f};
|
|
uint16_t sc16[4];
|
|
thread const uint8_t *sc8 = (thread const uint8_t *)sc16;
|
|
|
|
const uint64_t selected_base = (uint64_t)token * args.n_expert_used;
|
|
const uint64_t mid_base = (uint64_t)token * args.mid_token_stride;
|
|
for (uint slot = 0; slot < args.n_expert_used; slot++) {
|
|
const int expert = selected[selected_base + slot];
|
|
if (expert < 0 || (uint)expert >= args.n_total_expert) continue;
|
|
if (!ds4_tp_owns_expert(expert, args.n_total_expert,
|
|
args.tp_rank, args.tp_world)) continue;
|
|
const uint64_t down_addr = down_addrs[(uint)expert];
|
|
if (down_addr == 0) continue;
|
|
|
|
device const block_q4_K *x =
|
|
(device const block_q4_K *)(reinterpret_cast<device const char *>(down_addr) +
|
|
(uint64_t)row0 * args.down_row_bytes);
|
|
device const float *y = mid + mid_base + (uint64_t)slot * args.mid_dim;
|
|
device const float *y4 = y + ix * QK_K + 64 * iq + 8 * ir;
|
|
|
|
for (int ib = ix; ib < nb; ib += 4) {
|
|
float yl[16];
|
|
float yh[16];
|
|
float4 sumy = {0.f, 0.f, 0.f, 0.f};
|
|
|
|
for (short i = 0; i < 8; ++i) {
|
|
yl[i + 0] = y4[i + 0]; sumy[0] += yl[i + 0];
|
|
yl[i + 8] = y4[i + 32]; sumy[1] += yl[i + 8];
|
|
yh[i + 0] = y4[i + 128]; sumy[2] += yh[i + 0];
|
|
yh[i + 8] = y4[i + 160]; sumy[3] += yh[i + 8];
|
|
}
|
|
|
|
device const uint16_t *sc = (device const uint16_t *)x[ib].scales + iq;
|
|
device const uint16_t *q1 = (device const uint16_t *)x[ib].qs + 16 * iq + 4 * ir;
|
|
device const half *dh = &x[ib].d;
|
|
|
|
for (short row = 0; row < nr0 && row0 + (uint)row < args.out_dim; row++) {
|
|
sc16[0] = sc[0] & kmask1;
|
|
sc16[1] = sc[2] & kmask1;
|
|
sc16[2] = ((sc[4] >> 0) & kmask2) | ((sc[0] & kmask3) >> 2);
|
|
sc16[3] = ((sc[4] >> 4) & kmask2) | ((sc[2] & kmask3) >> 2);
|
|
|
|
device const uint16_t *q2 = q1 + 32;
|
|
float4 acc1 = {0.f, 0.f, 0.f, 0.f};
|
|
float4 acc2 = {0.f, 0.f, 0.f, 0.f};
|
|
|
|
FOR_UNROLL (short i = 0; i < 4; ++i) {
|
|
acc1[0] += yl[2 * i + 0] * (q1[i] & 0x000F);
|
|
acc1[1] += yl[2 * i + 1] * (q1[i] & 0x0F00);
|
|
acc1[2] += yl[2 * i + 8] * (q1[i] & 0x00F0);
|
|
acc1[3] += yl[2 * i + 9] * (q1[i] & 0xF000);
|
|
acc2[0] += yh[2 * i + 0] * (q2[i] & 0x000F);
|
|
acc2[1] += yh[2 * i + 1] * (q2[i] & 0x0F00);
|
|
acc2[2] += yh[2 * i + 8] * (q2[i] & 0x00F0);
|
|
acc2[3] += yh[2 * i + 9] * (q2[i] & 0xF000);
|
|
}
|
|
|
|
sumf[row] += dh[0] * ((acc1[0] + 1.f / 256.f * acc1[1]) * sc8[0] +
|
|
(acc1[2] + 1.f / 256.f * acc1[3]) * sc8[1] * 1.f / 16.f +
|
|
(acc2[0] + 1.f / 256.f * acc2[1]) * sc8[4] +
|
|
(acc2[2] + 1.f / 256.f * acc2[3]) * sc8[5] * 1.f / 16.f) -
|
|
dh[1] * (sumy[0] * sc8[2] + sumy[1] * sc8[3] +
|
|
sumy[2] * sc8[6] + sumy[3] * sc8[7]);
|
|
|
|
q1 += args.down_row_bytes / 2;
|
|
sc += args.down_row_bytes / 2;
|
|
dh += args.down_row_bytes / 2;
|
|
}
|
|
|
|
y4 += 4 * QK_K;
|
|
}
|
|
}
|
|
|
|
for (short row = 0; row < nr0 && row0 + (uint)row < args.out_dim; row++) {
|
|
const float sum_all = simd_sum(sumf[row]);
|
|
if (tiisg == 0u) {
|
|
out[(uint64_t)token * args.out_dim + row0 + (uint)row] = sum_all;
|
|
}
|
|
}
|
|
}
|
|
|
|
kernel void kernel_glm_q6_K_down_f32(
|
|
constant ds4_metal_glm_routed_moe_args &args,
|
|
device const char *down,
|
|
device const int32_t *selected,
|
|
device const float *mid,
|
|
device float *out,
|
|
uint3 tgpig [[threadgroup_position_in_grid]],
|
|
ushort tiisg [[thread_index_in_simdgroup]],
|
|
ushort sgitg [[simdgroup_index_in_threadgroup]]) {
|
|
const short NSG = 2;
|
|
constexpr uint kmask1 = 0x03u;
|
|
constexpr uint kmask2 = 0x0Cu;
|
|
constexpr uint kmask3 = 0x30u;
|
|
constexpr uint kmask4 = 0xC0u;
|
|
|
|
const uint row0 = ((uint)tgpig.x * (uint)NSG + (uint)sgitg) * N_R0_Q6_K;
|
|
const uint token = tgpig.y;
|
|
if (row0 >= args.out_dim || token >= args.n_tokens) return;
|
|
|
|
const int nb = args.mid_dim / QK_K;
|
|
float sumf[N_R0_Q6_K] = {0.f};
|
|
float yl[16];
|
|
const short tid = tiisg / 2;
|
|
const short ix = tiisg % 2;
|
|
const short ip = tid / 8;
|
|
const short il = tid % 8;
|
|
const short l0 = 4 * il;
|
|
const short is = 8 * ip + l0 / 16;
|
|
const short y_offset = 128 * ip + l0;
|
|
const short q_offset_l = 64 * ip + l0;
|
|
const short q_offset_h = 32 * ip + l0;
|
|
|
|
const uint64_t selected_base = (uint64_t)token * args.n_expert_used;
|
|
const uint64_t mid_base = (uint64_t)token * args.mid_token_stride;
|
|
for (uint slot = 0; slot < args.n_expert_used; slot++) {
|
|
const int expert = selected[selected_base + slot];
|
|
if (expert < 0 || (uint)expert >= args.n_total_expert) continue;
|
|
device const block_q6_K *x =
|
|
(device const block_q6_K *)(down +
|
|
(uint64_t)(uint)expert * args.down_expert_bytes +
|
|
(uint64_t)row0 * args.down_row_bytes);
|
|
device const float *yy = mid + mid_base + (uint64_t)slot * args.mid_dim;
|
|
|
|
for (int ib = ix; ib < nb; ib += 2) {
|
|
device const uchar *q1 = x[ib].ql + q_offset_l;
|
|
device const uchar *q2 = q1 + 32;
|
|
device const uchar *qh = x[ib].qh + q_offset_h;
|
|
device const char *sc = x[ib].scales + is;
|
|
device const half *dh = &x[ib].d;
|
|
device const float *y = yy + (uint64_t)ib * QK_K + y_offset;
|
|
|
|
for (short l = 0; l < 4; ++l) {
|
|
yl[4 * l + 0] = y[l + 0];
|
|
yl[4 * l + 1] = y[l + 32];
|
|
yl[4 * l + 2] = y[l + 64];
|
|
yl[4 * l + 3] = y[l + 96];
|
|
}
|
|
|
|
for (short row = 0; row < N_R0_Q6_K && row0 + (uint)row < args.out_dim; row++) {
|
|
float4 sums = {0.f, 0.f, 0.f, 0.f};
|
|
|
|
FOR_UNROLL (short l = 0; l < 4; ++l) {
|
|
const uint h = (uint)qh[l];
|
|
sums[0] += yl[4 * l + 0] *
|
|
(float)((int)((q1[l] & 0x0Fu) | ((h & kmask1) << 4u)) - 32);
|
|
sums[1] += yl[4 * l + 1] *
|
|
(float)((int)((q2[l] & 0x0Fu) | ((h & kmask2) << 2u)) - 32);
|
|
sums[2] += yl[4 * l + 2] *
|
|
(float)((int)((q1[l] >> 4u) | (h & kmask3)) - 32);
|
|
sums[3] += yl[4 * l + 3] *
|
|
(float)((int)((q2[l] >> 4u) | ((h & kmask4) >> 2u)) - 32);
|
|
}
|
|
|
|
sumf[row] += dh[0] * (sums[0] * sc[0] + sums[1] * sc[2] +
|
|
sums[2] * sc[4] + sums[3] * sc[6]);
|
|
|
|
q1 += args.down_row_bytes;
|
|
q2 += args.down_row_bytes;
|
|
qh += args.down_row_bytes;
|
|
sc += args.down_row_bytes;
|
|
dh += args.down_row_bytes / 2;
|
|
}
|
|
}
|
|
}
|
|
|
|
for (short row = 0; row < N_R0_Q6_K && row0 + (uint)row < args.out_dim; row++) {
|
|
const float sum_all = simd_sum(sumf[row]);
|
|
if (tiisg == 0u) {
|
|
out[(uint64_t)token * args.out_dim + row0 + (uint)row] = sum_all;
|
|
}
|
|
}
|
|
}
|
|
|
|
template <typename type4x4>
|
|
void dequantize_q4_K(device const block_q4_K *xb, short il, thread type4x4 ®) {
|
|
device const uchar *q = xb->qs;
|
|
|
|
short is = (il / 4) * 2;
|
|
q = q + (il / 4) * 32 + 16 * (il & 1);
|
|
il = il & 3;
|
|
const uchar2 sc = get_scale_min_k4_just2(is, il / 2, xb->scales);
|
|
const float d = il < 2 ?
|
|
(float)xb->d :
|
|
(float)xb->d * (1.0f / 16.0f);
|
|
const float min = (float)xb->dmin;
|
|
const float dl = d * sc[0];
|
|
const float ml = min * sc[1];
|
|
|
|
const ushort mask = il < 2 ? 0x0F : 0xF0;
|
|
for (int i = 0; i < 16; ++i) {
|
|
reg[i / 4][i % 4] = dl * (q[i] & mask) - ml;
|
|
}
|
|
}
|
|
|
|
template <typename type4x4>
|
|
void dequantize_mxfp4(device const block_mxfp4 *xb, short il, thread type4x4 ®) {
|
|
const float d = ds4_metal_e8m0_to_f32(xb->e);
|
|
const uint shift = il == 0 ? 0u : 4u;
|
|
FOR_UNROLL (short i = 0; i < QK_MXFP4/2; i++) {
|
|
const uint q = ((uint)xb->qs[i] >> shift) & 0x0fu;
|
|
reg[i/4][i%4] = d * ds4_metal_mxfp4_values[q];
|
|
}
|
|
}
|
|
|
|
// The resident MXFP4 prefill down projection stages its weight tile as half.
|
|
// All finite E8M0 x E2M1 products can therefore be rounded once in this exact
|
|
// binary16 table. E8M0 0xff falls back to the established arithmetic because
|
|
// infinity-times-zero produces a GPU-canonicalized NaN rather than a portable
|
|
// table payload.
|
|
static __attribute__((noinline)) half ds4_metal_mxfp4_half_fallback_value(
|
|
uchar e,
|
|
uint q) {
|
|
half result;
|
|
result = ds4_metal_e8m0_to_f32(e) * ds4_metal_mxfp4_values[q];
|
|
return result;
|
|
}
|
|
|
|
static inline half ds4_metal_mxfp4_half_lut_value(uchar e, uint q) {
|
|
if (e == 0xffu) {
|
|
return ds4_metal_mxfp4_half_fallback_value(e, q);
|
|
}
|
|
return as_type<half>(ds4_metal_mxfp4_half_lut[(uint)e][q]);
|
|
}
|
|
|
|
void dequantize_mxfp4_half_lut(
|
|
device const block_mxfp4 *xb,
|
|
short il,
|
|
thread half4x4 ®) {
|
|
const uchar e = xb->e;
|
|
const uint shift = il == 0 ? 0u : 4u;
|
|
if (e == 0xffu) {
|
|
FOR_UNROLL (short i = 0; i < QK_MXFP4/2; i++) {
|
|
const uint q = ((uint)xb->qs[i] >> shift) & 0x0fu;
|
|
reg[i/4][i%4] = ds4_metal_mxfp4_half_fallback_value(e, q);
|
|
}
|
|
return;
|
|
}
|
|
FOR_UNROLL (short i = 0; i < QK_MXFP4/2; i++) {
|
|
const uint q = ((uint)xb->qs[i] >> shift) & 0x0fu;
|
|
reg[i/4][i%4] = as_type<half>(
|
|
ds4_metal_mxfp4_half_lut[(uint)e][q]);
|
|
}
|
|
}
|
|
|
|
// Exact half-domain scaling for the resident MXFP4 prefill pair tile. E8M0 is
|
|
// a pure power of two and every E2M1 magnitude is exact in binary16, so for
|
|
// E8M0 bytes in [103, 142] the half product is the once-rounded value of the
|
|
// float product: half(d) is exact there and half(d) * half(v) rounds exactly
|
|
// once, exactly like half(float(d) * float(v)). Outside that band (subnormal
|
|
// or overflowing halves, where a second rounding could appear) the established
|
|
// float arithmetic is kept. The 32-byte table replaces both the 64-byte float
|
|
// table and the per-element f32->f16 narrowing.
|
|
static constant half ds4_metal_mxfp4_half_values[16] = {
|
|
0.0h, 0.5h, 1.0h, 1.5h, 2.0h, 3.0h, 4.0h, 6.0h,
|
|
-0.0h, -0.5h, -1.0h, -1.5h, -2.0h, -3.0h, -4.0h, -6.0h,
|
|
};
|
|
|
|
void dequantize_mxfp4_half_scale(
|
|
device const block_mxfp4 *xb,
|
|
short il,
|
|
thread half4x4 ®) {
|
|
const uchar e = xb->e;
|
|
const uint shift = il == 0 ? 0u : 4u;
|
|
if (e < 103u || e > 142u) {
|
|
const float d = ds4_metal_e8m0_to_f32(e);
|
|
FOR_UNROLL (short i = 0; i < QK_MXFP4/2; i++) {
|
|
const uint q = ((uint)xb->qs[i] >> shift) & 0x0fu;
|
|
reg[i/4][i%4] = d * ds4_metal_mxfp4_values[q];
|
|
}
|
|
return;
|
|
}
|
|
const half dh = (half)ds4_metal_e8m0_to_f32(e);
|
|
FOR_UNROLL (short i = 0; i < QK_MXFP4/2; i++) {
|
|
const uint q = ((uint)xb->qs[i] >> shift) & 0x0fu;
|
|
reg[i/4][i%4] = dh * ds4_metal_mxfp4_half_values[q];
|
|
}
|
|
}
|
|
|
|
// Test-only raw-bit oracle for the half-scale path: covers all 4096
|
|
// exponent/code pairs including both sides of the band boundary.
|
|
// Test-only raw-bit oracle. Keeping both conversions in one GPU invocation
|
|
// covers all 4096 exponent/code pairs, including the 0xff fallback row,
|
|
// without depending on host floating-point conversion behavior.
|
|
kernel void kernel_test_mxfp4_down_half_lut(
|
|
device ushort *legacy [[buffer(0)]],
|
|
device ushort *lut [[buffer(1)]],
|
|
uint tid [[thread_position_in_grid]]) {
|
|
if (tid >= 4096u) {
|
|
return;
|
|
}
|
|
const uchar e = (uchar)(tid >> 4u);
|
|
const uint q = tid & 0x0fu;
|
|
half legacy_value;
|
|
legacy_value = ds4_metal_e8m0_to_f32(e) *
|
|
ds4_metal_mxfp4_values[q];
|
|
legacy[tid] = as_type<ushort>(legacy_value);
|
|
lut[tid] = as_type<ushort>(ds4_metal_mxfp4_half_lut_value(e, q));
|
|
}
|
|
|
|
template <typename type4x4>
|
|
void dequantize_q5_K(device const block_q5_K *xb, short il, thread type4x4 ®) {
|
|
const short group = il / 2;
|
|
const short l0 = (il & 1) * 16;
|
|
const uchar2 sm = get_scale_min_k4_just2(group, 0, xb->scales);
|
|
const float dl = (float)xb->d * (float)sm.x;
|
|
const float ml = (float)xb->dmin * (float)sm.y;
|
|
device const uchar *q = xb->qs + (group / 2) * 32 + l0;
|
|
device const uchar *qh = xb->qh + l0;
|
|
const uint shift = (uint)(group & 1) * 4u;
|
|
const uchar hmask = (uchar)(1u << (uint)group);
|
|
|
|
for (int i = 0; i < 16; ++i) {
|
|
uint v = ((uint)q[i] >> shift) & 0x0Fu;
|
|
v += (qh[i] & hmask) ? 16u : 0u;
|
|
reg[i / 4][i % 4] = dl * (float)v - ml;
|
|
}
|
|
}
|
|
|
|
template <typename type4x4>
|
|
void dequantize_q6_K(device const block_q6_K *xb, short il, thread type4x4 ®) {
|
|
const short n128 = il / 8;
|
|
const short il128 = il & 7;
|
|
const short quarter = il128 / 2;
|
|
const short l0 = (il128 & 1) * 16;
|
|
const uint ql_base = (uint)n128 * 64u;
|
|
const uint qh_base = (uint)n128 * 32u;
|
|
const uint sc_base = (uint)n128 * 8u + (uint)quarter * 2u + (uint)(il128 & 1);
|
|
const float d = (float)xb->d * (float)((int)xb->scales[sc_base]);
|
|
|
|
for (int i = 0; i < 16; ++i) {
|
|
const uint l = (uint)l0 + (uint)i;
|
|
uint v;
|
|
if (quarter == 0) {
|
|
v = ((uint)xb->ql[ql_base + l] & 0x0Fu) |
|
|
((((uint)xb->qh[qh_base + l] >> 0u) & 3u) << 4u);
|
|
} else if (quarter == 1) {
|
|
v = ((uint)xb->ql[ql_base + 32u + l] & 0x0Fu) |
|
|
((((uint)xb->qh[qh_base + l] >> 2u) & 3u) << 4u);
|
|
} else if (quarter == 2) {
|
|
v = ((uint)xb->ql[ql_base + l] >> 4u) |
|
|
((((uint)xb->qh[qh_base + l] >> 4u) & 3u) << 4u);
|
|
} else {
|
|
v = ((uint)xb->ql[ql_base + 32u + l] >> 4u) |
|
|
((((uint)xb->qh[qh_base + l] >> 6u) & 3u) << 4u);
|
|
}
|
|
reg[i / 4][i % 4] = d * (float)((int)v - 32);
|
|
}
|
|
}
|
|
|
|
template <typename type4x4>
|
|
void dequantize_iq2_xxs(device const block_iq2_xxs * xb, short il, thread type4x4 & reg) {
|
|
const float d = xb->d;
|
|
const int ib32 = il/2;
|
|
il = il%2;
|
|
device const uint16_t * q2 = xb->qs + 4*ib32;
|
|
const uint32_t aux32_g = q2[0] | (q2[1] << 16);
|
|
const uint32_t aux32_s = q2[2] | (q2[3] << 16);
|
|
thread const uint8_t * aux8 = (thread const uint8_t *)&aux32_g;
|
|
const float dl = d * (0.5f + (aux32_s >> 28)) * 0.25f;
|
|
constant uint8_t * grid = (constant uint8_t *)(iq2xxs_grid + aux8[2*il+0]);
|
|
uint8_t signs = ksigns_iq2xs[(aux32_s >> 14*il) & 127];
|
|
for (int i = 0; i < 8; ++i) {
|
|
reg[i/4][i%4] = dl * grid[i] * (signs & kmask_iq2xs[i] ? -1.f : 1.f);
|
|
}
|
|
grid = (constant uint8_t *)(iq2xxs_grid + aux8[2*il+1]);
|
|
signs = ksigns_iq2xs[(aux32_s >> (14*il+7)) & 127];
|
|
for (int i = 0; i < 8; ++i) {
|
|
reg[2+i/4][i%4] = dl * grid[i] * (signs & kmask_iq2xs[i] ? -1.f : 1.f);
|
|
}
|
|
}
|
|
|
|
struct ds4_metal_args_mul_mv_id {
|
|
int32_t nei0;
|
|
int32_t nei1;
|
|
uint64_t nbi1;
|
|
int32_t ne00;
|
|
int32_t ne01;
|
|
int32_t ne02;
|
|
uint64_t nb00;
|
|
uint64_t nb01;
|
|
uint64_t nb02;
|
|
int32_t ne10;
|
|
int32_t ne11;
|
|
int32_t ne12;
|
|
int32_t ne13;
|
|
uint64_t nb10;
|
|
uint64_t nb11;
|
|
uint64_t nb12;
|
|
int32_t ne0;
|
|
int32_t ne1;
|
|
uint64_t nb1;
|
|
int32_t nr0;
|
|
/* Tensor-parallel expert ownership: when tp_world > 1, routed-MoE
|
|
* kernels skip experts outside this rank's contiguous id range so this
|
|
* rank never touches their weights. Zero-initialized (tp_world 0) means
|
|
* no split; only the routed id kernels read these. */
|
|
int32_t tp_rank;
|
|
int32_t tp_world;
|
|
/* When set, the sum6 kernels add this per-row vector (the shared-expert
|
|
* partial) into their accumulation, folding the TP local-add. */
|
|
int32_t tp_addend;
|
|
/* First expert id present at the bound blob base (TP mappings bind only
|
|
* the owned range; kernels rebase before pointer math). */
|
|
int32_t tp_expert_base;
|
|
};
|
|
|
|
struct ds4_metal_moe_expert_group_args {
|
|
uint32_t expert_base;
|
|
uint32_t expert_count;
|
|
uint32_t accumulate;
|
|
uint32_t pad0;
|
|
};
|
|
|
|
struct ds4_metal_q4_gather_slots6_args {
|
|
uint64_t expert_bytes;
|
|
uint32_t group_size;
|
|
uint32_t n_slots;
|
|
};
|
|
|
|
struct ds4_metal_q4_expert_table {
|
|
array<device const char *, 384> experts [[id(0)]];
|
|
};
|
|
|
|
struct ds4_metal_expert_address_table {
|
|
device const uint64_t *addrs;
|
|
};
|
|
|
|
struct ds4_metal_stream_expert_validate_args {
|
|
uint32_t n_total_expert;
|
|
uint32_t n_expert;
|
|
};
|
|
|
|
struct ds4_metal_stream_expert_split_args {
|
|
uint32_t active_mask;
|
|
uint32_t accumulate;
|
|
};
|
|
|
|
struct ds4_metal_args_mul_mm_id_map0 {
|
|
int32_t ne02;
|
|
int32_t ne10;
|
|
int32_t ne11;
|
|
uint64_t nb11;
|
|
uint64_t nb12;
|
|
int32_t ne21;
|
|
int32_t ne20;
|
|
uint64_t nb21;
|
|
};
|
|
|
|
struct ds4_metal_args_mul_mm_id {
|
|
int32_t ne00;
|
|
int32_t ne02;
|
|
uint64_t nb01;
|
|
uint64_t nb02;
|
|
uint64_t nb03;
|
|
int32_t ne11;
|
|
uint64_t nb10;
|
|
uint64_t nb11;
|
|
uint64_t nb12;
|
|
uint64_t nb13;
|
|
int32_t ne20;
|
|
int32_t ne21;
|
|
int32_t ne0;
|
|
int32_t ne1;
|
|
int16_t r2;
|
|
int16_t r3;
|
|
/* Tensor-parallel expert ownership; tp_world<=1 disables the split.
|
|
* tp_expert_base rebases expert ids when the bind covers only the
|
|
* owned contiguous range of the expert blob. */
|
|
int32_t tp_rank;
|
|
int32_t tp_world;
|
|
int32_t tp_expert_base;
|
|
};
|
|
|
|
template<int nr0, typename args_t>
|
|
void kernel_mul_mv_q2_K_f32_impl(
|
|
args_t args,
|
|
device const char * src0,
|
|
device const char * src1,
|
|
device char * dst,
|
|
threadgroup char * shmem,
|
|
uint3 tgpig,
|
|
ushort tiisg,
|
|
ushort sgitg) {
|
|
const short NSG = FC_mul_mv_nsg;
|
|
|
|
const int nb = args.ne00/QK_K;
|
|
|
|
const int r0 = tgpig.x;
|
|
const int r1 = tgpig.y;
|
|
const int im = tgpig.z;
|
|
|
|
const int first_row = (r0 * NSG + sgitg) * nr0;
|
|
|
|
const uint i12 = im%args.ne12;
|
|
const uint i13 = im/args.ne12;
|
|
|
|
const uint64_t offset0 = first_row*args.nb01 + (i12/args.r2)*args.nb02 + (i13/args.r3)*args.nb03;
|
|
const uint64_t offset1 = r1*args.nb11 + (i12 )*args.nb12 + (i13 )*args.nb13;
|
|
|
|
device const block_q2_K * x = (device const block_q2_K *) (src0 + offset0);
|
|
device const float * y = (device const float *) (src1 + offset1);
|
|
|
|
float yl[32];
|
|
float sumf[nr0]={0.f};
|
|
|
|
const short ix = tiisg/8; // 0...3
|
|
const short it = tiisg%8; // 0...7
|
|
const short iq = it/4; // 0 or 1
|
|
const short ir = it%4; // 0...3
|
|
const short is = (8*ir)/16;// 0 or 1
|
|
|
|
device const float * y4 = y + ix * QK_K + 128 * iq + 8 * ir;
|
|
|
|
for (int ib = ix; ib < nb; ib += 4) {
|
|
float4 sumy = {0.f, 0.f, 0.f, 0.f};
|
|
for (short i = 0; i < 8; ++i) {
|
|
yl[i+ 0] = y4[i+ 0]; sumy[0] += yl[i+ 0];
|
|
yl[i+ 8] = y4[i+32]; sumy[1] += yl[i+ 8];
|
|
yl[i+16] = y4[i+64]; sumy[2] += yl[i+16];
|
|
yl[i+24] = y4[i+96]; sumy[3] += yl[i+24];
|
|
}
|
|
|
|
device const uint8_t * sc = (device const uint8_t *)x[ib].scales + 8*iq + is;
|
|
device const uint16_t * qs = (device const uint16_t *)x[ib].qs + 16 * iq + 4 * ir;
|
|
device const half * dh = &x[ib].d;
|
|
|
|
for (short row = 0; row < nr0; row++) {
|
|
float4 acc1 = {0.f, 0.f, 0.f, 0.f};
|
|
float4 acc2 = {0.f, 0.f, 0.f, 0.f};
|
|
for (int i = 0; i < 8; i += 2) {
|
|
acc1[0] += yl[i+ 0] * (qs[i/2] & 0x0003);
|
|
acc2[0] += yl[i+ 1] * (qs[i/2] & 0x0300);
|
|
acc1[1] += yl[i+ 8] * (qs[i/2] & 0x000c);
|
|
acc2[1] += yl[i+ 9] * (qs[i/2] & 0x0c00);
|
|
acc1[2] += yl[i+16] * (qs[i/2] & 0x0030);
|
|
acc2[2] += yl[i+17] * (qs[i/2] & 0x3000);
|
|
acc1[3] += yl[i+24] * (qs[i/2] & 0x00c0);
|
|
acc2[3] += yl[i+25] * (qs[i/2] & 0xc000);
|
|
}
|
|
float dall = dh[0];
|
|
float dmin = dh[1] * 1.f/16.f;
|
|
sumf[row] += dall * ((acc1[0] + 1.f/256.f * acc2[0]) * (sc[0] & 0xF) * 1.f/ 1.f +
|
|
(acc1[1] + 1.f/256.f * acc2[1]) * (sc[2] & 0xF) * 1.f/ 4.f +
|
|
(acc1[2] + 1.f/256.f * acc2[2]) * (sc[4] & 0xF) * 1.f/16.f +
|
|
(acc1[3] + 1.f/256.f * acc2[3]) * (sc[6] & 0xF) * 1.f/64.f) -
|
|
dmin * (sumy[0] * (sc[0] & 0xF0) + sumy[1] * (sc[2] & 0xF0) + sumy[2] * (sc[4] & 0xF0) + sumy[3] * (sc[6] & 0xF0));
|
|
|
|
qs += args.nb01/2;
|
|
sc += args.nb01;
|
|
dh += args.nb01/2;
|
|
}
|
|
|
|
y4 += 4 * QK_K;
|
|
}
|
|
|
|
device float * dst_f32 = (device float *) dst + (uint64_t)im*args.ne0*args.ne1 + (uint64_t)r1*args.ne0;
|
|
|
|
for (int row = 0; row < nr0 && first_row + row < args.ne0; ++row) {
|
|
float sum_all = simd_sum(sumf[row]);
|
|
if (tiisg == 0) {
|
|
dst_f32[first_row + row] = sum_all;
|
|
}
|
|
}
|
|
}
|
|
|
|
template<int nr0, typename args_t>
|
|
void kernel_mul_mv_q4_K_f32_impl(
|
|
args_t args,
|
|
device const char *src0,
|
|
device const char *src1,
|
|
device char *dst,
|
|
threadgroup char *shmem,
|
|
uint3 tgpig,
|
|
ushort tiisg,
|
|
ushort sgitg) {
|
|
const short NSG = FC_mul_mv_nsg;
|
|
|
|
constexpr uint16_t kmask1 = 0x3f3f;
|
|
constexpr uint16_t kmask2 = 0x0f0f;
|
|
constexpr uint16_t kmask3 = 0xc0c0;
|
|
|
|
const short ix = tiisg / 8;
|
|
const short it = tiisg % 8;
|
|
const short iq = it / 4;
|
|
const short ir = it % 4;
|
|
|
|
const int nb = args.ne00 / QK_K;
|
|
|
|
const int r0 = tgpig.x;
|
|
const int r1 = tgpig.y;
|
|
const int im = tgpig.z;
|
|
|
|
const int first_row = (r0 * NSG + sgitg) * nr0;
|
|
|
|
const uint i12 = im % args.ne12;
|
|
const uint i13 = im / args.ne12;
|
|
|
|
const uint64_t offset0 = first_row * args.nb01 + (i12 / args.r2) * args.nb02 + (i13 / args.r3) * args.nb03;
|
|
const uint64_t offset1 = r1 * args.nb11 + i12 * args.nb12 + i13 * args.nb13;
|
|
|
|
device const block_q4_K *x = (device const block_q4_K *)(src0 + offset0);
|
|
device const float *y = (device const float *)(src1 + offset1);
|
|
|
|
float yl[16];
|
|
float yh[16];
|
|
float sumf[nr0] = {0.f};
|
|
|
|
device const float *y4 = y + ix * QK_K + 64 * iq + 8 * ir;
|
|
|
|
uint16_t sc16[4];
|
|
thread const uint8_t *sc8 = (thread const uint8_t *)sc16;
|
|
|
|
for (int ib = ix; ib < nb; ib += 4) {
|
|
float4 sumy = {0.f, 0.f, 0.f, 0.f};
|
|
|
|
for (short i = 0; i < 8; ++i) {
|
|
yl[i + 0] = y4[i + 0]; sumy[0] += yl[i + 0];
|
|
yl[i + 8] = y4[i + 32]; sumy[1] += yl[i + 8];
|
|
yh[i + 0] = y4[i + 128]; sumy[2] += yh[i + 0];
|
|
yh[i + 8] = y4[i + 160]; sumy[3] += yh[i + 8];
|
|
}
|
|
|
|
device const uint16_t *sc = (device const uint16_t *)x[ib].scales + iq;
|
|
device const uint16_t *q1 = (device const uint16_t *)x[ib].qs + 16 * iq + 4 * ir;
|
|
device const half *dh = &x[ib].d;
|
|
|
|
for (short row = 0; row < nr0; row++) {
|
|
sc16[0] = sc[0] & kmask1;
|
|
sc16[1] = sc[2] & kmask1;
|
|
sc16[2] = ((sc[4] >> 0) & kmask2) | ((sc[0] & kmask3) >> 2);
|
|
sc16[3] = ((sc[4] >> 4) & kmask2) | ((sc[2] & kmask3) >> 2);
|
|
|
|
device const uint16_t *q2 = q1 + 32;
|
|
|
|
float4 acc1 = {0.f, 0.f, 0.f, 0.f};
|
|
float4 acc2 = {0.f, 0.f, 0.f, 0.f};
|
|
|
|
FOR_UNROLL (short i = 0; i < 4; ++i) {
|
|
acc1[0] += yl[2 * i + 0] * (q1[i] & 0x000F);
|
|
acc1[1] += yl[2 * i + 1] * (q1[i] & 0x0F00);
|
|
acc1[2] += yl[2 * i + 8] * (q1[i] & 0x00F0);
|
|
acc1[3] += yl[2 * i + 9] * (q1[i] & 0xF000);
|
|
acc2[0] += yh[2 * i + 0] * (q2[i] & 0x000F);
|
|
acc2[1] += yh[2 * i + 1] * (q2[i] & 0x0F00);
|
|
acc2[2] += yh[2 * i + 8] * (q2[i] & 0x00F0);
|
|
acc2[3] += yh[2 * i + 9] * (q2[i] & 0xF000);
|
|
}
|
|
|
|
sumf[row] += dh[0] * ((acc1[0] + 1.f / 256.f * acc1[1]) * sc8[0] +
|
|
(acc1[2] + 1.f / 256.f * acc1[3]) * sc8[1] * 1.f / 16.f +
|
|
(acc2[0] + 1.f / 256.f * acc2[1]) * sc8[4] +
|
|
(acc2[2] + 1.f / 256.f * acc2[3]) * sc8[5] * 1.f / 16.f) -
|
|
dh[1] * (sumy[0] * sc8[2] + sumy[1] * sc8[3] + sumy[2] * sc8[6] + sumy[3] * sc8[7]);
|
|
|
|
q1 += args.nb01 / 2;
|
|
sc += args.nb01 / 2;
|
|
dh += args.nb01 / 2;
|
|
}
|
|
|
|
y4 += 4 * QK_K;
|
|
}
|
|
|
|
device float *dst_f32 = (device float *)dst + (uint64_t)im * args.ne0 * args.ne1 + (uint64_t)r1 * args.ne0;
|
|
|
|
for (int row = 0; row < nr0 && first_row + row < args.ne0; ++row) {
|
|
float sum_all = simd_sum(sumf[row]);
|
|
if (tiisg == 0) {
|
|
dst_f32[first_row + row] = sum_all;
|
|
}
|
|
}
|
|
|
|
(void)shmem;
|
|
}
|
|
|
|
template<int nr0, typename args_t>
|
|
void kernel_mul_mv_mxfp4_f32_impl(
|
|
args_t args,
|
|
device const char *src0,
|
|
device const char *src1,
|
|
device char *dst,
|
|
threadgroup char *shmem,
|
|
uint3 tgpig,
|
|
ushort tiisg,
|
|
ushort sgitg) {
|
|
const short NSG = FC_mul_mv_nsg;
|
|
const int nb = args.ne00 / QK_MXFP4;
|
|
const int first_row = (tgpig.x * NSG + sgitg) * nr0;
|
|
const uint i12 = tgpig.z % args.ne12;
|
|
const uint i13 = tgpig.z / args.ne12;
|
|
const uint64_t offset0 = (uint64_t)first_row * args.nb01 +
|
|
(uint64_t)(i12 / args.r2) * args.nb02 + (uint64_t)(i13 / args.r3) * args.nb03;
|
|
const uint64_t offset1 = (uint64_t)tgpig.y * args.nb11 +
|
|
(uint64_t)i12 * args.nb12 + (uint64_t)i13 * args.nb13;
|
|
|
|
device const block_mxfp4 *x = (device const block_mxfp4 *)(src0 + offset0);
|
|
device const float *y = (device const float *)(src1 + offset1);
|
|
const int row_blocks = (int)(args.nb01 / sizeof(block_mxfp4));
|
|
const short ix = tiisg / 2;
|
|
const short it = tiisg & 1;
|
|
|
|
threadgroup float *lut = (threadgroup float *)shmem;
|
|
if (sgitg == 0) lut[tiisg] = ds4_metal_mxfp4_values[tiisg & 15];
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
float sumf[nr0] = {0.f};
|
|
device const float *yb = y + ix * QK_MXFP4 + it * 8;
|
|
for (int ib = ix; ib < nb; ib += 16) {
|
|
device const float4 *y4 = (device const float4 *)yb;
|
|
const float4 yl0 = y4[0];
|
|
const float4 yl1 = y4[4];
|
|
const float4 yl2 = y4[1];
|
|
const float4 yl3 = y4[5];
|
|
|
|
FOR_UNROLL (short row = 0; row < nr0; row++) {
|
|
device const block_mxfp4 &xb = x[row * row_blocks + ib];
|
|
device const uchar *q = xb.qs + 8 * it;
|
|
float4 acc = yl0 * float4(lut[q[0] & 15], lut[q[1] & 15],
|
|
lut[q[2] & 15], lut[q[3] & 15]);
|
|
acc += yl1 * float4(lut[q[0] >> 4], lut[q[1] >> 4],
|
|
lut[q[2] >> 4], lut[q[3] >> 4]);
|
|
acc += yl2 * float4(lut[q[4] & 15], lut[q[5] & 15],
|
|
lut[q[6] & 15], lut[q[7] & 15]);
|
|
acc += yl3 * float4(lut[q[4] >> 4], lut[q[5] >> 4],
|
|
lut[q[6] >> 4], lut[q[7] >> 4]);
|
|
sumf[row] += ds4_metal_e8m0_to_f32(xb.e) *
|
|
((acc.x + acc.y) + (acc.z + acc.w));
|
|
}
|
|
yb += 16 * QK_MXFP4;
|
|
}
|
|
|
|
device float *dst_f32 = (device float *)dst +
|
|
(uint64_t)tgpig.z * args.ne0 * args.ne1 + (uint64_t)tgpig.y * args.ne0;
|
|
FOR_UNROLL (short row = 0; row < nr0; row++) {
|
|
if (first_row + row < args.ne0) {
|
|
const float sum = simd_sum(sumf[row]);
|
|
if (tiisg == 0) dst_f32[first_row + row] = sum;
|
|
}
|
|
}
|
|
}
|
|
|
|
template<int nr0, typename args_t>
|
|
void kernel_mul_mv_q8_K_f32_impl(
|
|
args_t args,
|
|
device const char *src0,
|
|
device const char *src1,
|
|
device char *dst,
|
|
threadgroup char *shmem,
|
|
uint3 tgpig,
|
|
ushort tiisg,
|
|
ushort sgitg) {
|
|
const short NSG = FC_mul_mv_nsg;
|
|
|
|
const int nb = args.ne00 / QK_K;
|
|
|
|
const int r0 = tgpig.x;
|
|
const int r1 = tgpig.y;
|
|
const int im = tgpig.z;
|
|
|
|
const int first_row = (r0 * NSG + sgitg) * nr0;
|
|
|
|
const uint i12 = im % args.ne12;
|
|
const uint i13 = im / args.ne12;
|
|
|
|
const uint64_t offset0 = first_row * args.nb01 + (i12 / args.r2) * args.nb02 + (i13 / args.r3) * args.nb03;
|
|
const uint64_t offset1 = r1 * args.nb11 + i12 * args.nb12 + i13 * args.nb13;
|
|
|
|
device const block_q8_K *x = (device const block_q8_K *)(src0 + offset0);
|
|
device const float *y = (device const float *)(src1 + offset1);
|
|
|
|
float sumf[nr0] = {0.f};
|
|
|
|
const short ix = tiisg / 8;
|
|
const short il = tiisg % 8;
|
|
const int ib0 = ix;
|
|
|
|
device const float *yb = y + ib0 * QK_K + il * 32;
|
|
|
|
for (int ib = ib0; ib < nb; ib += 4) {
|
|
float yl[32];
|
|
for (short i = 0; i < 32; ++i) {
|
|
yl[i] = yb[i];
|
|
}
|
|
|
|
for (short row = 0; row < nr0; row++) {
|
|
device const block_q8_K *xr =
|
|
(device const block_q8_K *)((device const char *)x + (uint64_t)row * args.nb01);
|
|
device const int8_t *qs = xr[ib].qs + il * 32;
|
|
|
|
float sumq = 0.f;
|
|
for (short i = 0; i < 32; ++i) {
|
|
sumq += (float)qs[i] * yl[i];
|
|
}
|
|
|
|
sumf[row] += sumq * xr[ib].d;
|
|
}
|
|
|
|
yb += 4 * QK_K;
|
|
}
|
|
|
|
device float *dst_f32 =
|
|
(device float *)dst + (uint64_t)im * args.ne0 * args.ne1 + (uint64_t)r1 * args.ne0;
|
|
|
|
for (int row = 0; row < nr0 && first_row + row < args.ne0; ++row) {
|
|
float sum_all = simd_sum(sumf[row]);
|
|
if (tiisg == 0) {
|
|
dst_f32[first_row + row] = sum_all;
|
|
}
|
|
}
|
|
|
|
(void)shmem;
|
|
}
|
|
|
|
template<int nr0, typename args_t>
|
|
void kernel_mul_mv_iq2_xxs_f32_impl(
|
|
args_t args,
|
|
device const char * src0,
|
|
device const char * src1,
|
|
device char * dst,
|
|
threadgroup char * shmem,
|
|
uint3 tgpig,
|
|
ushort tiisg,
|
|
ushort sgitg) {
|
|
const short NSG = FC_mul_mv_nsg;
|
|
|
|
const int nb = args.ne00/QK_K;
|
|
|
|
const int r0 = tgpig.x;
|
|
const int r1 = tgpig.y;
|
|
const int im = tgpig.z;
|
|
|
|
const int first_row = (r0 * NSG + sgitg) * nr0;
|
|
|
|
const uint i12 = im%args.ne12;
|
|
const uint i13 = im/args.ne12;
|
|
|
|
const uint64_t offset0 = first_row*args.nb01 + (i12/args.r2)*args.nb02 + (i13/args.r3)*args.nb03;
|
|
const uint64_t offset1 = r1*args.nb11 + (i12 )*args.nb12 + (i13 )*args.nb13;
|
|
|
|
device const block_iq2_xxs * x = (device const block_iq2_xxs *) (src0 + offset0);
|
|
device const float * y = (device const float *) (src1 + offset1);
|
|
|
|
float yl[32];
|
|
float sumf[nr0]={0.f};
|
|
|
|
const int nb32 = nb * (QK_K / 32);
|
|
|
|
threadgroup uint64_t * svalues = (threadgroup uint64_t *)(shmem);
|
|
threadgroup uint8_t * ssigns = (threadgroup uint8_t *)(svalues + 256);
|
|
{
|
|
int nval = 4;
|
|
int pos = (32*sgitg + tiisg)*nval;
|
|
for (int i = 0; i < nval; ++i) svalues[pos + i] = ds4_metal_iq2xxs_grid[pos + i];
|
|
nval = 2;
|
|
pos = (32*sgitg + tiisg)*nval;
|
|
for (int i = 0; i < nval; ++i) ssigns[pos+i] = ds4_metal_ksigns_iq2xs[pos+i];
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
}
|
|
|
|
const int ix = tiisg;
|
|
|
|
device const float * y4 = y + 32 * ix;
|
|
|
|
for (int ib32 = ix; ib32 < nb32; ib32 += 32) {
|
|
for (short i = 0; i < 32; ++i) {
|
|
yl[i] = y4[i];
|
|
}
|
|
|
|
const int ibl = ib32 / (QK_K / 32);
|
|
const int ib = ib32 % (QK_K / 32);
|
|
|
|
device const block_iq2_xxs * xr = x + ibl;
|
|
device const uint16_t * q2 = xr->qs + 4 * ib;
|
|
device const half * dh = &xr->d;
|
|
|
|
for (short row = 0; row < nr0; row++) {
|
|
const float db = dh[0];
|
|
device const uint8_t * aux8 = (device const uint8_t *)q2;
|
|
const uint32_t aux32 = q2[2] | (q2[3] << 16);
|
|
const float d = db * (0.5f + (aux32 >> 28));
|
|
|
|
float sum = 0;
|
|
for (short l = 0; l < 4; ++l) {
|
|
const threadgroup uint8_t * grid = (const threadgroup uint8_t *)(svalues + aux8[l]);
|
|
const uint8_t signs = ssigns[(aux32 >> 7*l) & 127];
|
|
for (short j = 0; j < 8; ++j) {
|
|
sum += yl[8*l + j] * grid[j] * (signs & ds4_metal_kmask_iq2xs[j] ? -1.f : 1.f);
|
|
}
|
|
}
|
|
sumf[row] += d * sum;
|
|
|
|
dh += args.nb01/2;
|
|
q2 += args.nb01/2;
|
|
}
|
|
|
|
y4 += 32 * 32;
|
|
}
|
|
|
|
device float * dst_f32 = (device float *) dst + (uint64_t)im*args.ne0*args.ne1 + (uint64_t)r1*args.ne0;
|
|
|
|
for (int row = 0; row < nr0 && first_row + row < args.ne0; ++row) {
|
|
float sum_all = simd_sum(sumf[row]);
|
|
if (tiisg == 0) {
|
|
dst_f32[first_row + row] = sum_all * 0.25f;
|
|
}
|
|
}
|
|
}
|
|
|
|
template<int nr0>
|
|
void kernel_mul_mv_iq2_xxs_pair_f32_impl(
|
|
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,
|
|
threadgroup char * shmem,
|
|
uint3 tgpig,
|
|
ushort tiisg,
|
|
ushort sgitg) {
|
|
const short NSG = FC_mul_mv_nsg;
|
|
|
|
const int nb = args.ne00/QK_K;
|
|
|
|
const int r0 = tgpig.x;
|
|
const int r1 = tgpig.y;
|
|
const int im = tgpig.z;
|
|
|
|
const int first_row = (r0 * NSG + sgitg) * nr0;
|
|
|
|
const uint i12 = im%args.ne12;
|
|
const uint i13 = im/args.ne12;
|
|
|
|
const uint64_t offset0 = first_row*args.nb01 + (i12/args.r2)*args.nb02 + (i13/args.r3)*args.nb03;
|
|
const uint64_t offset1 = r1*args.nb11 + (i12 )*args.nb12 + (i13 )*args.nb13;
|
|
|
|
device const block_iq2_xxs * xg = (device const block_iq2_xxs *) (src0_gate + offset0);
|
|
device const block_iq2_xxs * xu = (device const block_iq2_xxs *) (src0_up + offset0);
|
|
device const float * y = (device const float *) (src1 + offset1);
|
|
|
|
float yl[32];
|
|
float sumg[nr0]={0.f};
|
|
float sumu[nr0]={0.f};
|
|
|
|
const int nb32 = nb * (QK_K / 32);
|
|
|
|
threadgroup uint64_t * svalues = (threadgroup uint64_t *)(shmem);
|
|
threadgroup uint8_t * ssigns = (threadgroup uint8_t *)(svalues + 256);
|
|
{
|
|
int nval = 4;
|
|
int pos = (32*sgitg + tiisg)*nval;
|
|
for (int i = 0; i < nval; ++i) svalues[pos + i] = ds4_metal_iq2xxs_grid[pos + i];
|
|
nval = 2;
|
|
pos = (32*sgitg + tiisg)*nval;
|
|
for (int i = 0; i < nval; ++i) ssigns[pos+i] = ds4_metal_ksigns_iq2xs[pos+i];
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
}
|
|
|
|
const int ix = tiisg;
|
|
device const float * y4 = y + 32 * ix;
|
|
|
|
for (int ib32 = ix; ib32 < nb32; ib32 += 32) {
|
|
for (short i = 0; i < 32; ++i) {
|
|
yl[i] = y4[i];
|
|
}
|
|
|
|
const int ibl = ib32 / (QK_K / 32);
|
|
const int ib = ib32 % (QK_K / 32);
|
|
|
|
device const block_iq2_xxs * xgr = xg + ibl;
|
|
device const block_iq2_xxs * xur = xu + ibl;
|
|
device const uint16_t * qg = xgr->qs + 4 * ib;
|
|
device const uint16_t * qu = xur->qs + 4 * ib;
|
|
device const half * dhg = &xgr->d;
|
|
device const half * dhu = &xur->d;
|
|
|
|
for (short row = 0; row < nr0; row++) {
|
|
device const uint8_t * aux8g = (device const uint8_t *)qg;
|
|
device const uint8_t * aux8u = (device const uint8_t *)qu;
|
|
const uint32_t aux32g = qg[2] | (qg[3] << 16);
|
|
const uint32_t aux32u = qu[2] | (qu[3] << 16);
|
|
const float dg = (float)dhg[0] * (0.5f + (aux32g >> 28));
|
|
const float du = (float)dhu[0] * (0.5f + (aux32u >> 28));
|
|
|
|
float sg = 0;
|
|
float su = 0;
|
|
for (short l = 0; l < 4; ++l) {
|
|
const threadgroup uint8_t * gridg = (const threadgroup uint8_t *)(svalues + aux8g[l]);
|
|
const threadgroup uint8_t * gridu = (const threadgroup uint8_t *)(svalues + aux8u[l]);
|
|
const uint8_t signg = ssigns[(aux32g >> 7*l) & 127];
|
|
const uint8_t signu = ssigns[(aux32u >> 7*l) & 127];
|
|
for (short j = 0; j < 8; ++j) {
|
|
const float v = yl[8*l + j];
|
|
sg += v * gridg[j] * (signg & ds4_metal_kmask_iq2xs[j] ? -1.f : 1.f);
|
|
su += v * gridu[j] * (signu & ds4_metal_kmask_iq2xs[j] ? -1.f : 1.f);
|
|
}
|
|
}
|
|
sumg[row] += dg * sg;
|
|
sumu[row] += du * su;
|
|
|
|
dhg += args.nb01/2;
|
|
dhu += args.nb01/2;
|
|
qg += args.nb01/2;
|
|
qu += args.nb01/2;
|
|
}
|
|
|
|
y4 += 32 * 32;
|
|
}
|
|
|
|
device float * dst_gate_f32 = (device float *) dst_gate + (uint64_t)im*args.ne0*args.ne1 + (uint64_t)r1*args.ne0;
|
|
device float * dst_up_f32 = (device float *) dst_up + (uint64_t)im*args.ne0*args.ne1 + (uint64_t)r1*args.ne0;
|
|
|
|
for (int row = 0; row < nr0 && first_row + row < args.ne0; ++row) {
|
|
const float sum_gate = simd_sum(sumg[row]);
|
|
const float sum_up = simd_sum(sumu[row]);
|
|
if (tiisg == 0) {
|
|
dst_gate_f32[first_row + row] = sum_gate * 0.25f;
|
|
dst_up_f32[first_row + row] = sum_up * 0.25f;
|
|
}
|
|
}
|
|
}
|
|
|
|
typedef void (kernel_mul_mv2_disp_t)(
|
|
ds4_metal_args_mul_mv args,
|
|
device const char * src0,
|
|
device const char * src1,
|
|
device char * dst,
|
|
threadgroup char * shmem,
|
|
uint3 tgpig,
|
|
ushort tiisg,
|
|
ushort sgitg);
|
|
|
|
template<kernel_mul_mv2_disp_t disp_fn>
|
|
void mmv_fn(
|
|
ds4_metal_args_mul_mv args,
|
|
device const char * src0,
|
|
device const char * src1,
|
|
device char * dst,
|
|
threadgroup char * shmem,
|
|
uint3 tgpig,
|
|
ushort tiitg,
|
|
ushort tiisg,
|
|
ushort sgitg) {
|
|
disp_fn(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
|
|
}
|
|
|
|
typedef decltype(mmv_fn<kernel_mul_mv_q2_K_f32_impl<N_R0_Q2_K>>) mul_mv_id_disp_fn_t;
|
|
|
|
// Decode-time expert matvec. The ids tensor selects the routed expert for each
|
|
// slot, then this wrapper invokes the quantized row kernel for Q8_0, Q2_K, or
|
|
// IQ2_XXS weights without materializing per-expert dispatches on the CPU.
|
|
template<mul_mv_id_disp_fn_t disp_fn>
|
|
kernel void kernel_mul_mv_id(
|
|
constant ds4_metal_args_mul_mv_id & args,
|
|
device const char * src0s,
|
|
device const char * src1,
|
|
device char * dst,
|
|
device const char * ids,
|
|
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]]) {
|
|
(void)tiitg;
|
|
|
|
const int iid1 = tgpig.z/args.nei0;
|
|
const int idx = tgpig.z%args.nei0;
|
|
|
|
tgpig.z = 0;
|
|
|
|
const int32_t i02 = ((device const int32_t *) (ids + iid1*args.nbi1))[idx];
|
|
|
|
const int64_t i11 = idx % args.ne11;
|
|
const int64_t i12 = iid1;
|
|
|
|
const int64_t i1 = idx;
|
|
const int64_t i2 = i12;
|
|
|
|
device char * dst_cur = dst + (i1*args.ne0 + i2*args.ne1*args.ne0)*sizeof(float);
|
|
|
|
if (!ds4_tp_owns_expert(i02, args.ne02, args.tp_rank, args.tp_world)) {
|
|
/* Unowned expert under the TP split: zero this threadgroup's output
|
|
* rows so the downstream expert-sum stages stay unchanged. */
|
|
const short NSG = FC_mul_mv_nsg;
|
|
const int row0 = (tgpig.x * NSG + sgitg) * args.nr0;
|
|
device float *dst_f32 = (device float *)dst_cur;
|
|
for (int r = 0; r < args.nr0 && row0 + r < args.ne0; r++) {
|
|
if (tiisg == 0) dst_f32[row0 + r] = 0.0f;
|
|
}
|
|
return;
|
|
}
|
|
|
|
device const char * src0_cur =
|
|
src0s + (int64_t)(i02 - args.tp_expert_base)*args.nb02;
|
|
device const char * src1_cur = src1 + i11*args.nb11 + i12*args.nb12;
|
|
|
|
ds4_metal_args_mul_mv args0 = {
|
|
/*.ne00 =*/ args.ne00,
|
|
/*.ne01 =*/ args.ne01,
|
|
/*.ne02 =*/ 1,
|
|
/*.nb00 =*/ args.nb00,
|
|
/*.nb01 =*/ args.nb01,
|
|
/*.nb02 =*/ args.nb02,
|
|
/*.nb03 =*/ args.nb02,
|
|
/*.ne10 =*/ args.ne10,
|
|
/*.ne11 =*/ 1,
|
|
/*.ne12 =*/ 1,
|
|
/*.nb10 =*/ args.nb10,
|
|
/*.nb11 =*/ args.nb11,
|
|
/*.nb12 =*/ args.nb12,
|
|
/*.nb13 =*/ args.nb12,
|
|
/*.ne0 =*/ args.ne0,
|
|
/*.ne1 =*/ 1,
|
|
/*.nr0 =*/ args.nr0,
|
|
/*.r2 =*/ 1,
|
|
/*.r3 =*/ 1,
|
|
};
|
|
|
|
disp_fn(
|
|
args0,
|
|
/* src0 */ src0_cur,
|
|
/* src1 */ src1_cur,
|
|
/* dst */ dst_cur,
|
|
shmem,
|
|
tgpig,
|
|
tiitg,
|
|
tiisg,
|
|
sgitg);
|
|
}
|
|
|
|
typedef decltype(kernel_mul_mv_id<mmv_fn<kernel_mul_mv_q2_K_f32_impl<N_R0_Q2_K>>>) kernel_mul_mv_id_q_t;
|
|
typedef decltype(kernel_mul_mv_id<mmv_fn<kernel_mul_mv_q8_0_f32_impl<N_R0_Q8_0>>>) kernel_mul_mv_id_q8_0_t;
|
|
|
|
// Host-visible decode MoE matvec variants for the DS4 quant formats.
|
|
template [[host_name("kernel_mul_mv_id_q8_0_f32")]] kernel kernel_mul_mv_id_q8_0_t kernel_mul_mv_id<mmv_fn<kernel_mul_mv_q8_0_f32_impl<N_R0_Q8_0>>>;
|
|
template [[host_name("kernel_mul_mv_id_q2_K_f32")]] kernel kernel_mul_mv_id_q_t kernel_mul_mv_id<mmv_fn<kernel_mul_mv_q2_K_f32_impl<N_R0_Q2_K>>>;
|
|
template [[host_name("kernel_mul_mv_id_q4_K_f32")]] kernel kernel_mul_mv_id_q_t kernel_mul_mv_id<mmv_fn<kernel_mul_mv_q4_K_f32_impl<N_R0_Q4_K>>>;
|
|
template [[host_name("kernel_mul_mv_id_q8_K_f32")]] kernel kernel_mul_mv_id_q_t kernel_mul_mv_id<mmv_fn<kernel_mul_mv_q8_K_f32_impl<N_R0_Q8_K>>>;
|
|
template [[host_name("kernel_mul_mv_id_iq2_xxs_f32")]] kernel kernel_mul_mv_id_q_t kernel_mul_mv_id<mmv_fn<kernel_mul_mv_iq2_xxs_f32_impl<N_R0_IQ2_XXS>>>;
|
|
template [[host_name("kernel_mul_mv_id_mxfp4_f32")]] kernel kernel_mul_mv_id_q_t kernel_mul_mv_id<mmv_fn<kernel_mul_mv_mxfp4_f32_impl<N_R0_MXFP4>>>;
|
|
|
|
// Plain dense Q4_K matvec on the classic impl. The mul_mv_ext family used
|
|
// by the generic dense path tops out around 220 GB/s on M5 for the GLM
|
|
// DenseQ4 decode shapes; this impl streams the same rows at 530-650 GB/s.
|
|
kernel void kernel_mul_mv_q4_K_dense_f32(
|
|
constant ds4_metal_args_mul_mv & args,
|
|
device const char * src0,
|
|
device const char * src1,
|
|
device char * dst,
|
|
threadgroup char * shmem [[threadgroup(0)]],
|
|
uint3 tgpig[[threadgroup_position_in_grid]],
|
|
ushort tiisg[[thread_index_in_simdgroup]],
|
|
ushort sgitg[[simdgroup_index_in_threadgroup]]) {
|
|
kernel_mul_mv_q4_K_f32_impl<N_R0_Q4_K>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
|
|
}
|
|
|
|
// DS4 attention output low projection, specialized for the fixed block
|
|
// diagonal mapping used by the model:
|
|
//
|
|
// low[token, group, rank] = heads[token, group, :] * Woa[group, rank, :]
|
|
//
|
|
// The generic GGML-style id matvec supports arbitrary routed expert ids. Here
|
|
// the id is always equal to the group number, so this wrapper keeps the exact
|
|
// Q8_0 dot kernel but removes the id-buffer load and the CPU-side id table.
|
|
kernel void kernel_dsv4_attn_out_low_q8_0_f32(
|
|
constant ds4_metal_args_mul_mv_id & args,
|
|
device const char * src0s,
|
|
device const char * src1,
|
|
device char * dst,
|
|
threadgroup char * shmem [[threadgroup(0)]],
|
|
uint3 tgpig[[threadgroup_position_in_grid]],
|
|
ushort tiitg[[thread_index_in_threadgroup]],
|
|
ushort tiisg[[thread_index_in_simdgroup]],
|
|
ushort sgitg[[simdgroup_index_in_threadgroup]]) {
|
|
const int iid1 = tgpig.z/args.nei0;
|
|
const int idx = tgpig.z%args.nei0;
|
|
|
|
tgpig.z = 0;
|
|
|
|
const int64_t i11 = idx % args.ne11;
|
|
const int64_t i12 = iid1;
|
|
|
|
device const char * src0_cur = src0s + idx*args.nb02;
|
|
device const char * src1_cur = src1 + i11*args.nb11 + i12*args.nb12;
|
|
device char * dst_cur = dst + (idx*args.ne0 + i12*args.ne1*args.ne0)*sizeof(float);
|
|
|
|
ds4_metal_args_mul_mv args0 = {
|
|
/*.ne00 =*/ args.ne00,
|
|
/*.ne01 =*/ args.ne01,
|
|
/*.ne02 =*/ 1,
|
|
/*.nb00 =*/ args.nb00,
|
|
/*.nb01 =*/ args.nb01,
|
|
/*.nb02 =*/ args.nb02,
|
|
/*.nb03 =*/ args.nb02,
|
|
/*.ne10 =*/ args.ne10,
|
|
/*.ne11 =*/ 1,
|
|
/*.ne12 =*/ 1,
|
|
/*.nb10 =*/ args.nb10,
|
|
/*.nb11 =*/ args.nb11,
|
|
/*.nb12 =*/ args.nb12,
|
|
/*.nb13 =*/ args.nb12,
|
|
/*.ne0 =*/ args.ne0,
|
|
/*.ne1 =*/ 1,
|
|
/*.nr0 =*/ args.nr0,
|
|
/*.r2 =*/ 1,
|
|
/*.r3 =*/ 1,
|
|
};
|
|
|
|
kernel_mul_mv_q8_0_f32_impl<N_R0_Q8_0, thread ds4_metal_args_mul_mv &>(
|
|
args0,
|
|
src0_cur,
|
|
src1_cur,
|
|
dst_cur,
|
|
shmem,
|
|
tgpig,
|
|
tiisg,
|
|
sgitg);
|
|
}
|
|
|
|
kernel void kernel_dsv4_attn_out_low_q4_K_f32(
|
|
constant ds4_metal_args_mul_mv_id & args,
|
|
device const char * src0s,
|
|
device const char * src1,
|
|
device char * dst,
|
|
threadgroup char * shmem [[threadgroup(0)]],
|
|
uint3 tgpig[[threadgroup_position_in_grid]],
|
|
ushort tiitg[[thread_index_in_threadgroup]],
|
|
ushort tiisg[[thread_index_in_simdgroup]],
|
|
ushort sgitg[[simdgroup_index_in_threadgroup]]) {
|
|
(void)tiitg;
|
|
|
|
const int iid1 = tgpig.z / args.nei0;
|
|
const int idx = tgpig.z % args.nei0;
|
|
|
|
tgpig.z = 0;
|
|
|
|
const int64_t i11 = idx % args.ne11;
|
|
const int64_t i12 = iid1;
|
|
|
|
device const char * src0_cur = src0s + idx * args.nb02;
|
|
device const char * src1_cur = src1 + i11 * args.nb11 + i12 * args.nb12;
|
|
device char * dst_cur = dst + (idx * args.ne0 + i12 * args.ne1 * args.ne0) * sizeof(float);
|
|
|
|
ds4_metal_args_mul_mv args0 = {
|
|
/*.ne00 =*/ args.ne00,
|
|
/*.ne01 =*/ args.ne01,
|
|
/*.ne02 =*/ 1,
|
|
/*.nb00 =*/ args.nb00,
|
|
/*.nb01 =*/ args.nb01,
|
|
/*.nb02 =*/ args.nb02,
|
|
/*.nb03 =*/ args.nb02,
|
|
/*.ne10 =*/ args.ne10,
|
|
/*.ne11 =*/ 1,
|
|
/*.ne12 =*/ 1,
|
|
/*.nb10 =*/ args.nb10,
|
|
/*.nb11 =*/ args.nb11,
|
|
/*.nb12 =*/ args.nb12,
|
|
/*.nb13 =*/ args.nb12,
|
|
/*.ne0 =*/ args.ne0,
|
|
/*.ne1 =*/ 1,
|
|
/*.nr0 =*/ args.nr0,
|
|
/*.r2 =*/ 1,
|
|
/*.r3 =*/ 1,
|
|
};
|
|
|
|
kernel_mul_mv_q4_K_f32_impl<N_R0_Q4_K>(
|
|
args0,
|
|
src0_cur,
|
|
src1_cur,
|
|
dst_cur,
|
|
shmem,
|
|
tgpig,
|
|
tiisg,
|
|
sgitg);
|
|
}
|
|
|
|
kernel void kernel_mul_mv_id_iq2_xxs_pair_f32(
|
|
constant ds4_metal_args_mul_mv_id & args,
|
|
device const char * src0_gate,
|
|
device const char * src0_up,
|
|
device const char * src1,
|
|
device char * dst_gate,
|
|
device char * dst_up,
|
|
device const char * ids,
|
|
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]]) {
|
|
const int iid1 = tgpig.z/args.nei0;
|
|
const int idx = tgpig.z%args.nei0;
|
|
|
|
tgpig.z = 0;
|
|
|
|
const int32_t i02 = ((device const int32_t *) (ids + iid1*args.nbi1))[idx];
|
|
|
|
const int64_t i11 = idx % args.ne11;
|
|
const int64_t i12 = iid1;
|
|
|
|
device char * dst_gate_cur = dst_gate + (idx*args.ne0 + i12*args.ne1*args.ne0)*sizeof(float);
|
|
device char * dst_up_cur = dst_up + (idx*args.ne0 + i12*args.ne1*args.ne0)*sizeof(float);
|
|
|
|
if (!ds4_tp_owns_expert(i02, args.ne02, args.tp_rank, args.tp_world)) {
|
|
/* Unowned expert: zero this threadgroup's gate/up rows so the
|
|
* separate swiglu-weight stage yields zero mid rows. */
|
|
const short NSG_z = FC_mul_mv_nsg;
|
|
const int row0_z = (tgpig.x * NSG_z + sgitg) * args.nr0;
|
|
device float *zg = (device float *)dst_gate_cur;
|
|
device float *zu = (device float *)dst_up_cur;
|
|
for (int r = 0; r < args.nr0 && row0_z + r < args.ne0; r++) {
|
|
if (tiisg == 0) {
|
|
zg[row0_z + r] = 0.0f;
|
|
zu[row0_z + r] = 0.0f;
|
|
}
|
|
}
|
|
return;
|
|
}
|
|
|
|
device const char * src0_gate_cur = src0_gate + (int64_t)(i02 - args.tp_expert_base)*args.nb02;
|
|
device const char * src0_up_cur = src0_up + (int64_t)(i02 - args.tp_expert_base)*args.nb02;
|
|
device const char * src1_cur = src1 + i11*args.nb11 + i12*args.nb12;
|
|
|
|
ds4_metal_args_mul_mv args0 = {
|
|
args.ne00, args.ne01, 1,
|
|
args.nb00, args.nb01, args.nb02, args.nb02,
|
|
args.ne10, 1, 1,
|
|
args.nb10, args.nb11, args.nb12, args.nb12,
|
|
args.ne0, 1, args.nr0, 1, 1,
|
|
};
|
|
|
|
(void)tiitg;
|
|
kernel_mul_mv_iq2_xxs_pair_f32_impl<N_R0_IQ2_XXS>(
|
|
args0,
|
|
src0_gate_cur,
|
|
src0_up_cur,
|
|
src1_cur,
|
|
dst_gate_cur,
|
|
dst_up_cur,
|
|
shmem,
|
|
tgpig,
|
|
tiisg,
|
|
sgitg);
|
|
}
|
|
|
|
// Decode-only routed expert gate/up projection fused with the DS4 activation:
|
|
//
|
|
// mid = silu(clamp(gate)) * clamp(up) * route_weight
|
|
//
|
|
// The quantized dot products are intentionally the same IQ2_XXS paired path as
|
|
// `kernel_mul_mv_id_iq2_xxs_pair_f32`. The only extra work is done by lane 0
|
|
// after each exact reduced row has been produced. This removes the separate
|
|
// routed activation dispatch and avoids rereading the gate/up rows before the
|
|
// down projection. The host uses this only for the normal release path where
|
|
// diagnostics do not request clamped gate/up intermediates.
|
|
kernel void kernel_mul_mv_id_iq2_xxs_pair_swiglu_f32(
|
|
constant ds4_metal_args_mul_mv_id & args,
|
|
constant ds4_metal_dsv4_moe_swiglu_weight_args & act,
|
|
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,
|
|
device const char * ids,
|
|
device const char * weights,
|
|
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]]) {
|
|
const short NSG = FC_mul_mv_nsg;
|
|
const int iid1 = tgpig.z / args.nei0;
|
|
const int idx = tgpig.z % args.nei0;
|
|
|
|
tgpig.z = 0;
|
|
|
|
const int32_t i02 = ((device const int32_t *) (ids + iid1 * args.nbi1))[idx];
|
|
if (!ds4_tp_owns_expert(i02, args.ne02, args.tp_rank, args.tp_world)) return;
|
|
const int i02b = i02 - args.tp_expert_base;
|
|
const int64_t i11 = idx % args.ne11;
|
|
const int64_t i12 = iid1;
|
|
|
|
const int nb = args.ne00 / QK_K;
|
|
const int first_row = (tgpig.x * NSG + sgitg) * N_R0_IQ2_XXS;
|
|
const int nb32 = nb * (QK_K / 32);
|
|
|
|
device const block_iq2_xxs *xg =
|
|
(device const block_iq2_xxs *)(src0_gate + (int64_t)i02b * args.nb02 + (uint64_t)first_row * args.nb01);
|
|
device const block_iq2_xxs *xu =
|
|
(device const block_iq2_xxs *)(src0_up + (int64_t)i02b * args.nb02 + (uint64_t)first_row * args.nb01);
|
|
device const float *y =
|
|
(device const float *)(src1 + i11 * args.nb11 + i12 * args.nb12);
|
|
|
|
float yl[32];
|
|
float sumg[N_R0_IQ2_XXS] = {0.f};
|
|
float sumu[N_R0_IQ2_XXS] = {0.f};
|
|
|
|
threadgroup uint64_t *svalues = (threadgroup uint64_t *)(shmem);
|
|
threadgroup uint8_t *ssigns = (threadgroup uint8_t *)(svalues + 256);
|
|
{
|
|
int nval = 4;
|
|
int pos = (32 * sgitg + tiisg) * nval;
|
|
for (int i = 0; i < nval; ++i) svalues[pos + i] = ds4_metal_iq2xxs_grid[pos + i];
|
|
nval = 2;
|
|
pos = (32 * sgitg + tiisg) * nval;
|
|
for (int i = 0; i < nval; ++i) ssigns[pos + i] = ds4_metal_ksigns_iq2xs[pos + i];
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
}
|
|
|
|
const int ix = tiisg;
|
|
device const float *y4 = y + 32 * ix;
|
|
|
|
for (int ib32 = ix; ib32 < nb32; ib32 += 32) {
|
|
for (short i = 0; i < 32; ++i) {
|
|
yl[i] = y4[i];
|
|
}
|
|
|
|
const int ibl = ib32 / (QK_K / 32);
|
|
const int ib = ib32 % (QK_K / 32);
|
|
|
|
device const block_iq2_xxs *xgr = xg + ibl;
|
|
device const block_iq2_xxs *xur = xu + ibl;
|
|
device const uint16_t *qg = xgr->qs + 4 * ib;
|
|
device const uint16_t *qu = xur->qs + 4 * ib;
|
|
device const half *dhg = &xgr->d;
|
|
device const half *dhu = &xur->d;
|
|
|
|
for (short row = 0; row < N_R0_IQ2_XXS; row++) {
|
|
device const uint8_t *aux8g = (device const uint8_t *)qg;
|
|
device const uint8_t *aux8u = (device const uint8_t *)qu;
|
|
const uint32_t aux32g = qg[2] | (qg[3] << 16);
|
|
const uint32_t aux32u = qu[2] | (qu[3] << 16);
|
|
const float dg = (float)dhg[0] * (0.5f + (aux32g >> 28));
|
|
const float du = (float)dhu[0] * (0.5f + (aux32u >> 28));
|
|
|
|
float sg = 0;
|
|
float su = 0;
|
|
for (short l = 0; l < 4; ++l) {
|
|
const threadgroup uint8_t *gridg = (const threadgroup uint8_t *)(svalues + aux8g[l]);
|
|
const threadgroup uint8_t *gridu = (const threadgroup uint8_t *)(svalues + aux8u[l]);
|
|
const uint8_t signg = ssigns[(aux32g >> 7 * l) & 127];
|
|
const uint8_t signu = ssigns[(aux32u >> 7 * l) & 127];
|
|
for (short j = 0; j < 8; ++j) {
|
|
const float v = yl[8 * l + j];
|
|
sg += v * gridg[j] * (signg & ds4_metal_kmask_iq2xs[j] ? -1.f : 1.f);
|
|
su += v * gridu[j] * (signu & ds4_metal_kmask_iq2xs[j] ? -1.f : 1.f);
|
|
}
|
|
}
|
|
sumg[row] += dg * sg;
|
|
sumu[row] += du * su;
|
|
|
|
dhg += args.nb01 / 2;
|
|
dhu += args.nb01 / 2;
|
|
qg += args.nb01 / 2;
|
|
qu += args.nb01 / 2;
|
|
}
|
|
|
|
y4 += 32 * 32;
|
|
}
|
|
|
|
device float *dst_gate_f32 =
|
|
(device float *)dst_gate + (uint64_t)i12 * args.ne0 * args.ne1 + (uint64_t)i11 * args.ne0;
|
|
device float *dst_up_f32 =
|
|
(device float *)dst_up + (uint64_t)i12 * args.ne0 * args.ne1 + (uint64_t)i11 * args.ne0;
|
|
const uint64_t pair_row = (uint64_t)i12 * (uint64_t)args.nei0 + (uint64_t)idx;
|
|
device float *dst_mid_f32 =
|
|
(device float *)(dst_mid + pair_row * act.mid_row_stride);
|
|
device const float *route_w =
|
|
(device const float *)(weights + pair_row * act.weight_stride);
|
|
|
|
const float c = act.clamp_value;
|
|
const float route_weight = route_w[0];
|
|
float4 reduced_gate;
|
|
float4 reduced_up;
|
|
for (int row = 0; row < N_R0_IQ2_XXS; ++row) {
|
|
reduced_gate[row] = simd_sum(sumg[row]);
|
|
reduced_up[row] = simd_sum(sumu[row]);
|
|
}
|
|
if (tiisg < N_R0_IQ2_XXS && first_row + tiisg < args.ne0) {
|
|
const uint out_row = first_row + tiisg;
|
|
const float gate = reduced_gate[tiisg] * 0.25f;
|
|
const float up = reduced_up[tiisg] * 0.25f;
|
|
float g = gate;
|
|
float u = up;
|
|
if (c > 1.0e-6f) {
|
|
g = min(g, c);
|
|
u = clamp(u, -c, c);
|
|
}
|
|
dst_gate_f32[out_row] = gate;
|
|
dst_up_f32[out_row] = up;
|
|
const float silu = g / (1.0f + exp(-g));
|
|
dst_mid_f32[out_row] = silu * u * route_weight;
|
|
}
|
|
|
|
(void)tiitg;
|
|
}
|
|
|
|
kernel void kernel_mul_mv_id_iq2_xxs_pair_swiglu_pack2_overlap_f32(
|
|
constant ds4_metal_args_mul_mv_id & args,
|
|
constant ds4_metal_dsv4_moe_swiglu_weight_args & act,
|
|
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,
|
|
device const char * ids,
|
|
device const char * weights,
|
|
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 NSG = 4;
|
|
const int iid1 = tgpig.z / args.nei0;
|
|
const int idx = tgpig.z % args.nei0;
|
|
|
|
tgpig.z = 0;
|
|
|
|
const int32_t i02 = ((device const int32_t *) (ids + iid1 * args.nbi1))[idx];
|
|
if (!ds4_tp_owns_expert(i02, args.ne02, args.tp_rank, args.tp_world)) return;
|
|
const int i02b = i02 - args.tp_expert_base;
|
|
const int64_t i11 = idx % args.ne11;
|
|
const int64_t i12 = iid1;
|
|
|
|
const int nb = args.ne00 / QK_K;
|
|
const int first_row = (tgpig.x * NSG + sgitg) * N_R0_IQ2_XXS;
|
|
const int nb32 = nb * (QK_K / 32);
|
|
|
|
device const block_iq2_xxs *xg =
|
|
(device const block_iq2_xxs *)(src0_gate + (int64_t)i02b * args.nb02 + (uint64_t)first_row * args.nb01);
|
|
device const block_iq2_xxs *xu =
|
|
(device const block_iq2_xxs *)(src0_up + (int64_t)i02b * args.nb02 + (uint64_t)first_row * args.nb01);
|
|
device const float *y =
|
|
(device const float *)(src1 + i11 * args.nb11 + i12 * args.nb12);
|
|
|
|
float yl[32];
|
|
float sumg[N_R0_IQ2_XXS] = {0.f};
|
|
float sumu[N_R0_IQ2_XXS] = {0.f};
|
|
|
|
const int ix = tiisg;
|
|
device const float *y4 = y + 32 * ix;
|
|
bool yl_preloaded = sgitg >= 2;
|
|
if (yl_preloaded) {
|
|
for (short i = 0; i < 32; ++i) {
|
|
yl[i] = y4[i];
|
|
}
|
|
}
|
|
|
|
threadgroup uint64_t *svalues = (threadgroup uint64_t *)(shmem);
|
|
threadgroup uint8_t *ssigns = (threadgroup uint8_t *)(svalues + 256);
|
|
if (sgitg < 2) {
|
|
int nval = 4;
|
|
int pos = (32 * sgitg + tiisg) * nval;
|
|
for (int i = 0; i < nval; ++i) svalues[pos + i] = ds4_metal_iq2xxs_grid[pos + i];
|
|
nval = 2;
|
|
pos = (32 * sgitg + tiisg) * nval;
|
|
for (int i = 0; i < nval; ++i) ssigns[pos + i] = ds4_metal_ksigns_iq2xs[pos + i];
|
|
}
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
for (int ib32 = ix; ib32 < nb32; ib32 += 32) {
|
|
if (!yl_preloaded) {
|
|
for (short i = 0; i < 32; ++i) {
|
|
yl[i] = y4[i];
|
|
}
|
|
}
|
|
yl_preloaded = false;
|
|
|
|
const int ibl = ib32 / (QK_K / 32);
|
|
const int ib = ib32 % (QK_K / 32);
|
|
|
|
device const block_iq2_xxs *xgr = xg + ibl;
|
|
device const block_iq2_xxs *xur = xu + ibl;
|
|
device const uint16_t *qg = xgr->qs + 4 * ib;
|
|
device const uint16_t *qu = xur->qs + 4 * ib;
|
|
device const half *dhg = &xgr->d;
|
|
device const half *dhu = &xur->d;
|
|
|
|
for (short row = 0; row < N_R0_IQ2_XXS; row++) {
|
|
device const uint8_t *aux8g = (device const uint8_t *)qg;
|
|
device const uint8_t *aux8u = (device const uint8_t *)qu;
|
|
const uint32_t aux32g = qg[2] | (qg[3] << 16);
|
|
const uint32_t aux32u = qu[2] | (qu[3] << 16);
|
|
const float dg = (float)dhg[0] * (0.5f + (aux32g >> 28));
|
|
const float du = (float)dhu[0] * (0.5f + (aux32u >> 28));
|
|
|
|
float sg = 0;
|
|
float su = 0;
|
|
for (short l = 0; l < 4; ++l) {
|
|
const threadgroup uint8_t *gridg = (const threadgroup uint8_t *)(svalues + aux8g[l]);
|
|
const threadgroup uint8_t *gridu = (const threadgroup uint8_t *)(svalues + aux8u[l]);
|
|
const uint8_t signg = ssigns[(aux32g >> 7 * l) & 127];
|
|
const uint8_t signu = ssigns[(aux32u >> 7 * l) & 127];
|
|
for (short j = 0; j < 8; ++j) {
|
|
const float v = yl[8 * l + j];
|
|
sg += v * gridg[j] * (signg & ds4_metal_kmask_iq2xs[j] ? -1.f : 1.f);
|
|
su += v * gridu[j] * (signu & ds4_metal_kmask_iq2xs[j] ? -1.f : 1.f);
|
|
}
|
|
}
|
|
sumg[row] += dg * sg;
|
|
sumu[row] += du * su;
|
|
|
|
dhg += args.nb01 / 2;
|
|
dhu += args.nb01 / 2;
|
|
qg += args.nb01 / 2;
|
|
qu += args.nb01 / 2;
|
|
}
|
|
|
|
y4 += 32 * 32;
|
|
}
|
|
|
|
device float *dst_gate_f32 =
|
|
(device float *)dst_gate + (uint64_t)i12 * args.ne0 * args.ne1 + (uint64_t)i11 * args.ne0;
|
|
device float *dst_up_f32 =
|
|
(device float *)dst_up + (uint64_t)i12 * args.ne0 * args.ne1 + (uint64_t)i11 * args.ne0;
|
|
const uint64_t pair_row = (uint64_t)i12 * (uint64_t)args.nei0 + (uint64_t)idx;
|
|
device float *dst_mid_f32 =
|
|
(device float *)(dst_mid + pair_row * act.mid_row_stride);
|
|
device const float *route_w =
|
|
(device const float *)(weights + pair_row * act.weight_stride);
|
|
|
|
const float c = act.clamp_value;
|
|
const float route_weight = route_w[0];
|
|
float4 reduced_gate;
|
|
float4 reduced_up;
|
|
for (int row = 0; row < N_R0_IQ2_XXS; ++row) {
|
|
reduced_gate[row] = simd_sum(sumg[row]);
|
|
reduced_up[row] = simd_sum(sumu[row]);
|
|
}
|
|
if (tiisg < N_R0_IQ2_XXS && first_row + tiisg < args.ne0) {
|
|
const uint out_row = first_row + tiisg;
|
|
const float gate = reduced_gate[tiisg] * 0.25f;
|
|
const float up = reduced_up[tiisg] * 0.25f;
|
|
float g = gate;
|
|
float u = up;
|
|
if (c > 1.0e-6f) {
|
|
g = min(g, c);
|
|
u = clamp(u, -c, c);
|
|
}
|
|
dst_gate_f32[out_row] = gate;
|
|
dst_up_f32[out_row] = up;
|
|
const float silu = g / (1.0f + exp(-g));
|
|
dst_mid_f32[out_row] = silu * u * route_weight;
|
|
}
|
|
|
|
(void)tiitg;
|
|
}
|
|
|
|
|
|
kernel void kernel_mul_mv_slots6_iq2_xxs_pair_swiglu_f32(
|
|
constant ds4_metal_args_mul_mv_id & args,
|
|
constant ds4_metal_dsv4_moe_swiglu_weight_args & act,
|
|
device const char * src0_gate0,
|
|
device const char * src0_gate1,
|
|
device const char * src0_gate2,
|
|
device const char * src0_gate3,
|
|
device const char * src0_gate4,
|
|
device const char * src0_gate5,
|
|
device const char * src0_up0,
|
|
device const char * src0_up1,
|
|
device const char * src0_up2,
|
|
device const char * src0_up3,
|
|
device const char * src0_up4,
|
|
device const char * src0_up5,
|
|
device const char * src1,
|
|
device char * dst_gate,
|
|
device char * dst_up,
|
|
device char * dst_mid,
|
|
device const char * weights,
|
|
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]]) {
|
|
const int iid1 = tgpig.z / args.nei0;
|
|
const int idx = tgpig.z % args.nei0;
|
|
|
|
tgpig.z = 0;
|
|
|
|
const int64_t i11 = idx % args.ne11;
|
|
const int64_t i12 = iid1;
|
|
|
|
device const char *src0_gate_cur = src0_gate0;
|
|
device const char *src0_up_cur = src0_up0;
|
|
switch (idx) {
|
|
case 1: src0_gate_cur = src0_gate1; src0_up_cur = src0_up1; break;
|
|
case 2: src0_gate_cur = src0_gate2; src0_up_cur = src0_up2; break;
|
|
case 3: src0_gate_cur = src0_gate3; src0_up_cur = src0_up3; break;
|
|
case 4: src0_gate_cur = src0_gate4; src0_up_cur = src0_up4; break;
|
|
case 5: src0_gate_cur = src0_gate5; src0_up_cur = src0_up5; break;
|
|
default: break;
|
|
}
|
|
|
|
device const char *src1_cur = src1 + i11 * args.nb11 + i12 * args.nb12;
|
|
|
|
device char *dst_gate_cur = dst_gate + (idx * args.ne0 + i12 * args.ne1 * args.ne0) * sizeof(float);
|
|
device char *dst_up_cur = dst_up + (idx * args.ne0 + i12 * args.ne1 * args.ne0) * sizeof(float);
|
|
|
|
ds4_metal_args_mul_mv args0 = {
|
|
args.ne00, args.ne01, 1,
|
|
args.nb00, args.nb01, args.nb02, args.nb02,
|
|
args.ne10, 1, 1,
|
|
args.nb10, args.nb11, args.nb12, args.nb12,
|
|
args.ne0, 1, args.nr0, 1, 1,
|
|
};
|
|
|
|
kernel_mul_mv_iq2_xxs_pair_f32_impl<N_R0_IQ2_XXS>(
|
|
args0,
|
|
src0_gate_cur,
|
|
src0_up_cur,
|
|
src1_cur,
|
|
dst_gate_cur,
|
|
dst_up_cur,
|
|
shmem,
|
|
tgpig,
|
|
tiisg,
|
|
sgitg);
|
|
|
|
const short NSG = FC_mul_mv_nsg;
|
|
const int first_row = (tgpig.x * NSG + sgitg) * N_R0_IQ2_XXS;
|
|
device float *gate_f32 = (device float *)dst_gate_cur;
|
|
device float *up_f32 = (device float *)dst_up_cur;
|
|
const uint64_t pair_row = (uint64_t)i12 * (uint64_t)args.nei0 + (uint64_t)idx;
|
|
device float *mid_f32 = (device float *)(dst_mid + pair_row * act.mid_row_stride);
|
|
device const float *route_w = (device const float *)(weights + pair_row * act.weight_stride);
|
|
const float c = act.clamp_value;
|
|
const float route_weight = route_w[0];
|
|
|
|
if (tiisg == 0) {
|
|
for (int row = 0; row < N_R0_IQ2_XXS && first_row + row < args.ne0; ++row) {
|
|
const uint out_row = first_row + row;
|
|
float g = gate_f32[out_row];
|
|
float u = up_f32[out_row];
|
|
if (c > 1.0e-6f) {
|
|
g = min(g, c);
|
|
u = clamp(u, -c, c);
|
|
}
|
|
const float silu = g / (1.0f + exp(-g));
|
|
mid_f32[out_row] = silu * u * route_weight;
|
|
}
|
|
}
|
|
|
|
(void)tiitg;
|
|
}
|
|
|
|
kernel void kernel_mul_mv_addr_iq2_xxs_pair_swiglu_f32(
|
|
constant ds4_metal_args_mul_mv_id & args,
|
|
constant ds4_metal_dsv4_moe_swiglu_weight_args & act,
|
|
device const uint64_t * gate_addrs,
|
|
device const uint64_t * up_addrs,
|
|
device const char * src1,
|
|
device char * dst_gate,
|
|
device char * dst_up,
|
|
device char * dst_mid,
|
|
device const char * ids,
|
|
device const char * weights,
|
|
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]]) {
|
|
const int iid1 = tgpig.z / args.nei0;
|
|
const int idx = tgpig.z % args.nei0;
|
|
|
|
tgpig.z = 0;
|
|
|
|
const int32_t i02 = ((device const int32_t *)(ids + iid1 * args.nbi1))[idx];
|
|
if (i02 < 0 || i02 >= args.ne02 || i02 >= 384) {
|
|
return;
|
|
}
|
|
const uint64_t gate_addr = gate_addrs[(uint)i02];
|
|
const uint64_t up_addr = up_addrs[(uint)i02];
|
|
if (gate_addr == 0 || up_addr == 0) {
|
|
return;
|
|
}
|
|
|
|
const int64_t i11 = idx % args.ne11;
|
|
const int64_t i12 = iid1;
|
|
|
|
device const char *src0_gate_cur =
|
|
reinterpret_cast<device const char *>(gate_addr);
|
|
device const char *src0_up_cur =
|
|
reinterpret_cast<device const char *>(up_addr);
|
|
device const char *src1_cur = src1 + i11 * args.nb11 + i12 * args.nb12;
|
|
|
|
device char *dst_gate_cur = dst_gate + (idx * args.ne0 + i12 * args.ne1 * args.ne0) * sizeof(float);
|
|
device char *dst_up_cur = dst_up + (idx * args.ne0 + i12 * args.ne1 * args.ne0) * sizeof(float);
|
|
|
|
ds4_metal_args_mul_mv args0 = {
|
|
args.ne00, args.ne01, 1,
|
|
args.nb00, args.nb01, args.nb02, args.nb02,
|
|
args.ne10, 1, 1,
|
|
args.nb10, args.nb11, args.nb12, args.nb12,
|
|
args.ne0, 1, args.nr0, 1, 1,
|
|
};
|
|
|
|
kernel_mul_mv_iq2_xxs_pair_f32_impl<N_R0_IQ2_XXS>(
|
|
args0,
|
|
src0_gate_cur,
|
|
src0_up_cur,
|
|
src1_cur,
|
|
dst_gate_cur,
|
|
dst_up_cur,
|
|
shmem,
|
|
tgpig,
|
|
tiisg,
|
|
sgitg);
|
|
|
|
const short NSG = FC_mul_mv_nsg;
|
|
const int first_row = (tgpig.x * NSG + sgitg) * N_R0_IQ2_XXS;
|
|
device float *gate_f32 = (device float *)dst_gate_cur;
|
|
device float *up_f32 = (device float *)dst_up_cur;
|
|
const uint64_t pair_row = (uint64_t)i12 * (uint64_t)args.nei0 + (uint64_t)idx;
|
|
device float *mid_f32 = (device float *)(dst_mid + pair_row * act.mid_row_stride);
|
|
device const float *route_w = (device const float *)(weights + pair_row * act.weight_stride);
|
|
const float c = act.clamp_value;
|
|
const float route_weight = route_w[0];
|
|
|
|
if (tiisg == 0) {
|
|
for (int row = 0; row < N_R0_IQ2_XXS && first_row + row < args.ne0; ++row) {
|
|
const uint out_row = first_row + row;
|
|
float g = gate_f32[out_row];
|
|
float u = up_f32[out_row];
|
|
if (c > 1.0e-6f) {
|
|
g = min(g, c);
|
|
u = clamp(u, -c, c);
|
|
}
|
|
const float silu = g / (1.0f + exp(-g));
|
|
mid_f32[out_row] = silu * u * route_weight;
|
|
}
|
|
}
|
|
|
|
(void)tiitg;
|
|
}
|
|
|
|
kernel void kernel_mul_mv_addr_iq2_xxs_f32(
|
|
constant ds4_metal_args_mul_mv_id & args,
|
|
device const uint64_t * addrs,
|
|
device const char * src1,
|
|
device char * dst,
|
|
device const char * ids,
|
|
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]]) {
|
|
(void)tiitg;
|
|
|
|
const int iid1 = tgpig.z / args.nei0;
|
|
const int idx = tgpig.z % args.nei0;
|
|
|
|
tgpig.z = 0;
|
|
|
|
const int32_t i02 = ((device const int32_t *)(ids + iid1 * args.nbi1))[idx];
|
|
if (i02 < 0 || i02 >= args.ne02 || i02 >= 384) {
|
|
return;
|
|
}
|
|
const uint64_t addr = addrs[(uint)i02];
|
|
if (addr == 0) {
|
|
return;
|
|
}
|
|
|
|
const int64_t i11 = idx % args.ne11;
|
|
const int64_t i12 = iid1;
|
|
|
|
device const char *src0_cur = reinterpret_cast<device const char *>(addr);
|
|
device const char *src1_cur = src1 + i11 * args.nb11 + i12 * args.nb12;
|
|
device char *dst_cur = dst + (idx * args.ne0 + i12 * args.ne1 * args.ne0) * sizeof(float);
|
|
|
|
ds4_metal_args_mul_mv args0 = {
|
|
args.ne00, args.ne01, 1,
|
|
args.nb00, args.nb01, args.nb02, args.nb02,
|
|
args.ne10, 1, 1,
|
|
args.nb10, args.nb11, args.nb12, args.nb12,
|
|
args.ne0, 1, args.nr0, 1, 1,
|
|
};
|
|
|
|
kernel_mul_mv_iq2_xxs_f32_impl<N_R0_IQ2_XXS>(
|
|
args0,
|
|
src0_cur,
|
|
src1_cur,
|
|
dst_cur,
|
|
shmem,
|
|
tgpig,
|
|
tiisg,
|
|
sgitg);
|
|
}
|
|
|
|
kernel void kernel_mul_mv_addr_iq2_xxs_pair_swiglu_masked_f32(
|
|
constant ds4_metal_args_mul_mv_id & args,
|
|
constant ds4_metal_dsv4_moe_swiglu_weight_args & act,
|
|
constant ds4_metal_stream_expert_split_args & split,
|
|
device const uint64_t * gate_addrs,
|
|
device const uint64_t * up_addrs,
|
|
device const char * src1,
|
|
device char * dst_gate,
|
|
device char * dst_up,
|
|
device char * dst_mid,
|
|
device const char * ids,
|
|
device const char * weights,
|
|
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]]) {
|
|
const int iid1 = tgpig.z / args.nei0;
|
|
const int idx = tgpig.z % args.nei0;
|
|
if ((split.active_mask & (1u << (uint)idx)) == 0) {
|
|
return;
|
|
}
|
|
|
|
tgpig.z = 0;
|
|
|
|
const int32_t i02 = ((device const int32_t *)(ids + iid1 * args.nbi1))[idx];
|
|
if (i02 < 0 || i02 >= args.ne02 || i02 >= 384) {
|
|
return;
|
|
}
|
|
const uint64_t gate_addr = gate_addrs[(uint)i02];
|
|
const uint64_t up_addr = up_addrs[(uint)i02];
|
|
if (gate_addr == 0 || up_addr == 0) {
|
|
return;
|
|
}
|
|
|
|
const int64_t i11 = idx % args.ne11;
|
|
const int64_t i12 = iid1;
|
|
|
|
device const char *src0_gate_cur =
|
|
reinterpret_cast<device const char *>(gate_addr);
|
|
device const char *src0_up_cur =
|
|
reinterpret_cast<device const char *>(up_addr);
|
|
device const char *src1_cur = src1 + i11 * args.nb11 + i12 * args.nb12;
|
|
|
|
device char *dst_gate_cur = dst_gate + (idx * args.ne0 + i12 * args.ne1 * args.ne0) * sizeof(float);
|
|
device char *dst_up_cur = dst_up + (idx * args.ne0 + i12 * args.ne1 * args.ne0) * sizeof(float);
|
|
|
|
ds4_metal_args_mul_mv args0 = {
|
|
args.ne00, args.ne01, 1,
|
|
args.nb00, args.nb01, args.nb02, args.nb02,
|
|
args.ne10, 1, 1,
|
|
args.nb10, args.nb11, args.nb12, args.nb12,
|
|
args.ne0, 1, args.nr0, 1, 1,
|
|
};
|
|
|
|
kernel_mul_mv_iq2_xxs_pair_f32_impl<N_R0_IQ2_XXS>(
|
|
args0,
|
|
src0_gate_cur,
|
|
src0_up_cur,
|
|
src1_cur,
|
|
dst_gate_cur,
|
|
dst_up_cur,
|
|
shmem,
|
|
tgpig,
|
|
tiisg,
|
|
sgitg);
|
|
|
|
const short NSG = FC_mul_mv_nsg;
|
|
const int first_row = (tgpig.x * NSG + sgitg) * N_R0_IQ2_XXS;
|
|
device float *gate_f32 = (device float *)dst_gate_cur;
|
|
device float *up_f32 = (device float *)dst_up_cur;
|
|
const uint64_t pair_row = (uint64_t)i12 * (uint64_t)args.nei0 + (uint64_t)idx;
|
|
device float *mid_f32 = (device float *)(dst_mid + pair_row * act.mid_row_stride);
|
|
device const float *route_w = (device const float *)(weights + pair_row * act.weight_stride);
|
|
const float c = act.clamp_value;
|
|
const float route_weight = route_w[0];
|
|
|
|
if (tiisg == 0) {
|
|
for (int row = 0; row < N_R0_IQ2_XXS && first_row + row < args.ne0; ++row) {
|
|
const uint out_row = first_row + row;
|
|
float g = gate_f32[out_row];
|
|
float u = up_f32[out_row];
|
|
if (c > 1.0e-6f) {
|
|
g = min(g, c);
|
|
u = clamp(u, -c, c);
|
|
}
|
|
const float silu = g / (1.0f + exp(-g));
|
|
mid_f32[out_row] = silu * u * route_weight;
|
|
}
|
|
}
|
|
|
|
(void)tiitg;
|
|
}
|
|
|
|
kernel void kernel_stream_expert_cache_validate(
|
|
constant ds4_metal_stream_expert_validate_args & args,
|
|
device const char * ids,
|
|
device const uint64_t * gate_addrs,
|
|
device const uint64_t * up_addrs,
|
|
device const uint64_t * down_addrs,
|
|
device uint32_t * status,
|
|
uint tid [[thread_position_in_grid]]) {
|
|
if (tid != 0) return;
|
|
|
|
uint32_t miss_mask = 0;
|
|
uint32_t invalid_mask = 0;
|
|
const uint32_t n_expert = min(args.n_expert, (uint32_t)6);
|
|
device const int32_t *selected = (device const int32_t *)ids;
|
|
|
|
status[3] = n_expert;
|
|
for (uint32_t i = 0; i < 6; i++) {
|
|
const int32_t expert = i < n_expert ? selected[i] : -1;
|
|
status[4 + i] = as_type<uint32_t>(expert);
|
|
if (i >= n_expert) continue;
|
|
if (expert < 0 ||
|
|
(uint32_t)expert >= args.n_total_expert ||
|
|
(uint32_t)expert >= 384) {
|
|
invalid_mask |= (1u << i);
|
|
continue;
|
|
}
|
|
const uint32_t e = (uint32_t)expert;
|
|
if (gate_addrs[e] == 0 || up_addrs[e] == 0 || down_addrs[e] == 0) {
|
|
miss_mask |= (1u << i);
|
|
}
|
|
}
|
|
|
|
status[0] = (miss_mask == 0 && invalid_mask == 0) ? 1u : 0u;
|
|
status[1] = miss_mask;
|
|
status[2] = invalid_mask;
|
|
}
|
|
|
|
kernel void kernel_mul_mv_id_q4_K_pair_f32(
|
|
constant ds4_metal_args_mul_mv_id & args,
|
|
device const char * src0_gate,
|
|
device const char * src0_up,
|
|
device const char * src1,
|
|
device char * dst_gate,
|
|
device char * dst_up,
|
|
device const char * ids,
|
|
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]]) {
|
|
const int iid1 = tgpig.z / args.nei0;
|
|
const int idx = tgpig.z % args.nei0;
|
|
|
|
tgpig.z = 0;
|
|
|
|
const int32_t i02 = ((device const int32_t *)(ids + iid1 * args.nbi1))[idx];
|
|
const int64_t i11 = idx % args.ne11;
|
|
const int64_t i12 = iid1;
|
|
|
|
device char *dst_gate_cur = dst_gate + (idx * args.ne0 + i12 * args.ne1 * args.ne0) * sizeof(float);
|
|
device char *dst_up_cur = dst_up + (idx * args.ne0 + i12 * args.ne1 * args.ne0) * sizeof(float);
|
|
|
|
if (!ds4_tp_owns_expert(i02, args.ne02, args.tp_rank, args.tp_world)) {
|
|
const short NSG_z = FC_mul_mv_nsg;
|
|
const int row0_z = (tgpig.x * NSG_z + sgitg) * args.nr0;
|
|
device float *zg = (device float *)dst_gate_cur;
|
|
device float *zu = (device float *)dst_up_cur;
|
|
for (int r = 0; r < args.nr0 && row0_z + r < args.ne0; r++) {
|
|
if (tiisg == 0) {
|
|
zg[row0_z + r] = 0.0f;
|
|
zu[row0_z + r] = 0.0f;
|
|
}
|
|
}
|
|
return;
|
|
}
|
|
|
|
device const char *src0_gate_cur = src0_gate + (int64_t)(i02 - args.tp_expert_base) * args.nb02;
|
|
device const char *src0_up_cur = src0_up + (int64_t)(i02 - args.tp_expert_base) * args.nb02;
|
|
device const char *src1_cur = src1 + i11 * args.nb11 + i12 * args.nb12;
|
|
|
|
ds4_metal_args_mul_mv args0 = {
|
|
args.ne00, args.ne01, 1,
|
|
args.nb00, args.nb01, args.nb02, args.nb02,
|
|
args.ne10, 1, 1,
|
|
args.nb10, args.nb11, args.nb12, args.nb12,
|
|
args.ne0, 1, args.nr0, 1, 1,
|
|
};
|
|
|
|
(void)tiitg;
|
|
kernel_mul_mv_q4_K_f32_impl<N_R0_Q4_K>(
|
|
args0,
|
|
src0_gate_cur,
|
|
src1_cur,
|
|
dst_gate_cur,
|
|
shmem,
|
|
tgpig,
|
|
tiisg,
|
|
sgitg);
|
|
kernel_mul_mv_q4_K_f32_impl<N_R0_Q4_K>(
|
|
args0,
|
|
src0_up_cur,
|
|
src1_cur,
|
|
dst_up_cur,
|
|
shmem,
|
|
tgpig,
|
|
tiisg,
|
|
sgitg);
|
|
}
|
|
|
|
// Same release-path fusion as the IQ2_XXS kernel above for the Q4_K expert
|
|
// variant. The Q4 pair path reuses the existing exact matvec implementation
|
|
// for gate and up, then the same lane that wrote each row derives the routed
|
|
// SwiGLU input. This keeps Q4 behavior aligned with the Q2 optimization while
|
|
// preserving the old pair projection arithmetic.
|
|
kernel void kernel_mul_mv_id_q4_K_pair_swiglu_f32(
|
|
constant ds4_metal_args_mul_mv_id & args,
|
|
constant ds4_metal_dsv4_moe_swiglu_weight_args & act,
|
|
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,
|
|
device const char * ids,
|
|
device const char * weights,
|
|
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]]) {
|
|
const int iid1 = tgpig.z / args.nei0;
|
|
const int idx = tgpig.z % args.nei0;
|
|
|
|
tgpig.z = 0;
|
|
|
|
const int32_t i02 = ((device const int32_t *)(ids + iid1 * args.nbi1))[idx];
|
|
if (!ds4_tp_owns_expert(i02, args.ne02, args.tp_rank, args.tp_world)) return;
|
|
const int i02b = i02 - args.tp_expert_base;
|
|
const int64_t i11 = idx % args.ne11;
|
|
const int64_t i12 = iid1;
|
|
|
|
device const char *src0_gate_cur = src0_gate + (int64_t)i02b * args.nb02;
|
|
device const char *src0_up_cur = src0_up + (int64_t)i02b * args.nb02;
|
|
device const char *src1_cur = src1 + i11 * args.nb11 + i12 * args.nb12;
|
|
|
|
device char *dst_gate_cur = dst_gate + (idx * args.ne0 + i12 * args.ne1 * args.ne0) * sizeof(float);
|
|
device char *dst_up_cur = dst_up + (idx * args.ne0 + i12 * args.ne1 * args.ne0) * sizeof(float);
|
|
|
|
const short NSG = FC_mul_mv_nsg;
|
|
constexpr uint16_t kmask1 = 0x3f3f;
|
|
constexpr uint16_t kmask2 = 0x0f0f;
|
|
constexpr uint16_t kmask3 = 0xc0c0;
|
|
|
|
const short ix = tiisg / 8;
|
|
const short it = tiisg % 8;
|
|
const short iq = it / 4;
|
|
const short ir = it % 4;
|
|
const int nb = args.ne00 / QK_K;
|
|
const int first_row = (tgpig.x * NSG + sgitg) * N_R0_Q4_K;
|
|
device float *gate_f32 = (device float *)dst_gate_cur;
|
|
device float *up_f32 = (device float *)dst_up_cur;
|
|
const uint64_t pair_row = (uint64_t)i12 * (uint64_t)args.nei0 + (uint64_t)idx;
|
|
device float *mid_f32 = (device float *)(dst_mid + pair_row * act.mid_row_stride);
|
|
device const float *route_w = (device const float *)(weights + pair_row * act.weight_stride);
|
|
const float c = act.clamp_value;
|
|
const float route_weight = route_w[0];
|
|
|
|
device const block_q4_K *xg =
|
|
(device const block_q4_K *)(src0_gate_cur + (uint64_t)first_row * args.nb01);
|
|
device const block_q4_K *xu =
|
|
(device const block_q4_K *)(src0_up_cur + (uint64_t)first_row * args.nb01);
|
|
device const float *y = (device const float *)src1_cur;
|
|
device const float *y4 = y + ix * QK_K + 64 * iq + 8 * ir;
|
|
|
|
float sumg[N_R0_Q4_K] = {0.f};
|
|
float sumu[N_R0_Q4_K] = {0.f};
|
|
uint16_t sc16[4];
|
|
thread const uint8_t *sc8 = (thread const uint8_t *)sc16;
|
|
|
|
for (int ib = ix; ib < nb; ib += 4) {
|
|
float yl[16];
|
|
float yh[16];
|
|
float4 sumy = {0.f, 0.f, 0.f, 0.f};
|
|
|
|
for (short i = 0; i < 8; ++i) {
|
|
yl[i + 0] = y4[i + 0]; sumy[0] += yl[i + 0];
|
|
yl[i + 8] = y4[i + 32]; sumy[1] += yl[i + 8];
|
|
yh[i + 0] = y4[i + 128]; sumy[2] += yh[i + 0];
|
|
yh[i + 8] = y4[i + 160]; sumy[3] += yh[i + 8];
|
|
}
|
|
|
|
device const uint16_t *scg = (device const uint16_t *)xg[ib].scales + iq;
|
|
device const uint16_t *qg1 = (device const uint16_t *)xg[ib].qs + 16 * iq + 4 * ir;
|
|
device const half *dhg = &xg[ib].d;
|
|
device const uint16_t *scu = (device const uint16_t *)xu[ib].scales + iq;
|
|
device const uint16_t *qu1 = (device const uint16_t *)xu[ib].qs + 16 * iq + 4 * ir;
|
|
device const half *dhu = &xu[ib].d;
|
|
|
|
for (short row = 0; row < N_R0_Q4_K; row++) {
|
|
sc16[0] = scg[0] & kmask1;
|
|
sc16[1] = scg[2] & kmask1;
|
|
sc16[2] = ((scg[4] >> 0) & kmask2) | ((scg[0] & kmask3) >> 2);
|
|
sc16[3] = ((scg[4] >> 4) & kmask2) | ((scg[2] & kmask3) >> 2);
|
|
|
|
device const uint16_t *qg2 = qg1 + 32;
|
|
float4 acc1g = {0.f, 0.f, 0.f, 0.f};
|
|
float4 acc2g = {0.f, 0.f, 0.f, 0.f};
|
|
|
|
FOR_UNROLL (short i = 0; i < 4; ++i) {
|
|
acc1g[0] += yl[2 * i + 0] * (qg1[i] & 0x000F);
|
|
acc1g[1] += yl[2 * i + 1] * (qg1[i] & 0x0F00);
|
|
acc1g[2] += yl[2 * i + 8] * (qg1[i] & 0x00F0);
|
|
acc1g[3] += yl[2 * i + 9] * (qg1[i] & 0xF000);
|
|
acc2g[0] += yh[2 * i + 0] * (qg2[i] & 0x000F);
|
|
acc2g[1] += yh[2 * i + 1] * (qg2[i] & 0x0F00);
|
|
acc2g[2] += yh[2 * i + 8] * (qg2[i] & 0x00F0);
|
|
acc2g[3] += yh[2 * i + 9] * (qg2[i] & 0xF000);
|
|
}
|
|
|
|
sumg[row] += dhg[0] * ((acc1g[0] + 1.f / 256.f * acc1g[1]) * sc8[0] +
|
|
(acc1g[2] + 1.f / 256.f * acc1g[3]) * sc8[1] * 1.f / 16.f +
|
|
(acc2g[0] + 1.f / 256.f * acc2g[1]) * sc8[4] +
|
|
(acc2g[2] + 1.f / 256.f * acc2g[3]) * sc8[5] * 1.f / 16.f) -
|
|
dhg[1] * (sumy[0] * sc8[2] + sumy[1] * sc8[3] +
|
|
sumy[2] * sc8[6] + sumy[3] * sc8[7]);
|
|
|
|
sc16[0] = scu[0] & kmask1;
|
|
sc16[1] = scu[2] & kmask1;
|
|
sc16[2] = ((scu[4] >> 0) & kmask2) | ((scu[0] & kmask3) >> 2);
|
|
sc16[3] = ((scu[4] >> 4) & kmask2) | ((scu[2] & kmask3) >> 2);
|
|
|
|
device const uint16_t *qu2 = qu1 + 32;
|
|
float4 acc1u = {0.f, 0.f, 0.f, 0.f};
|
|
float4 acc2u = {0.f, 0.f, 0.f, 0.f};
|
|
|
|
FOR_UNROLL (short i = 0; i < 4; ++i) {
|
|
acc1u[0] += yl[2 * i + 0] * (qu1[i] & 0x000F);
|
|
acc1u[1] += yl[2 * i + 1] * (qu1[i] & 0x0F00);
|
|
acc1u[2] += yl[2 * i + 8] * (qu1[i] & 0x00F0);
|
|
acc1u[3] += yl[2 * i + 9] * (qu1[i] & 0xF000);
|
|
acc2u[0] += yh[2 * i + 0] * (qu2[i] & 0x000F);
|
|
acc2u[1] += yh[2 * i + 1] * (qu2[i] & 0x0F00);
|
|
acc2u[2] += yh[2 * i + 8] * (qu2[i] & 0x00F0);
|
|
acc2u[3] += yh[2 * i + 9] * (qu2[i] & 0xF000);
|
|
}
|
|
|
|
sumu[row] += dhu[0] * ((acc1u[0] + 1.f / 256.f * acc1u[1]) * sc8[0] +
|
|
(acc1u[2] + 1.f / 256.f * acc1u[3]) * sc8[1] * 1.f / 16.f +
|
|
(acc2u[0] + 1.f / 256.f * acc2u[1]) * sc8[4] +
|
|
(acc2u[2] + 1.f / 256.f * acc2u[3]) * sc8[5] * 1.f / 16.f) -
|
|
dhu[1] * (sumy[0] * sc8[2] + sumy[1] * sc8[3] +
|
|
sumy[2] * sc8[6] + sumy[3] * sc8[7]);
|
|
|
|
qg1 += args.nb01 / 2;
|
|
scg += args.nb01 / 2;
|
|
dhg += args.nb01 / 2;
|
|
qu1 += args.nb01 / 2;
|
|
scu += args.nb01 / 2;
|
|
dhu += args.nb01 / 2;
|
|
}
|
|
|
|
y4 += 4 * QK_K;
|
|
}
|
|
|
|
for (int row = 0; row < N_R0_Q4_K && first_row + row < args.ne0; ++row) {
|
|
const float gate = simd_sum(sumg[row]);
|
|
const float up = simd_sum(sumu[row]);
|
|
if (tiisg == 0) {
|
|
const uint out_row = first_row + row;
|
|
float g = gate;
|
|
float u = up;
|
|
if (c > 1.0e-6f) {
|
|
g = min(g, c);
|
|
u = clamp(u, -c, c);
|
|
}
|
|
gate_f32[out_row] = gate;
|
|
up_f32[out_row] = up;
|
|
const float silu = g / (1.0f + exp(-g));
|
|
mid_f32[out_row] = silu * u * route_weight;
|
|
}
|
|
}
|
|
|
|
(void)tiitg;
|
|
}
|
|
|
|
template<typename args_t>
|
|
void kernel_mul_mv_mxfp4_pair_swiglu_impl(
|
|
args_t args,
|
|
constant ds4_metal_dsv4_moe_swiglu_weight_args &act,
|
|
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,
|
|
float route_weight,
|
|
threadgroup char *shmem,
|
|
uint3 tgpig,
|
|
ushort tiisg,
|
|
ushort sgitg) {
|
|
const short NSG = FC_mul_mv_nsg;
|
|
const int first_row = (tgpig.x * NSG + sgitg) * N_R0_MXFP4;
|
|
const int nb = args.ne00 / QK_MXFP4;
|
|
const int row_blocks = (int)(args.nb01 / sizeof(block_mxfp4));
|
|
const short ix = tiisg / 2;
|
|
const short it = tiisg & 1;
|
|
|
|
threadgroup float *lut = (threadgroup float *)shmem;
|
|
if (sgitg == 0) lut[tiisg] = ds4_metal_mxfp4_values[tiisg & 15];
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
device const block_mxfp4 *xg =
|
|
(device const block_mxfp4 *)(src0_gate + (uint64_t)first_row * args.nb01);
|
|
device const block_mxfp4 *xu =
|
|
(device const block_mxfp4 *)(src0_up + (uint64_t)first_row * args.nb01);
|
|
device const float *yb = (device const float *)src1 + ix * QK_MXFP4 + it * 8;
|
|
float sumg[N_R0_MXFP4] = {0.f};
|
|
float sumu[N_R0_MXFP4] = {0.f};
|
|
|
|
for (int ib = ix; ib < nb; ib += 16) {
|
|
device const float4 *y4 = (device const float4 *)yb;
|
|
const float4 yl0 = y4[0];
|
|
const float4 yl1 = y4[4];
|
|
const float4 yl2 = y4[1];
|
|
const float4 yl3 = y4[5];
|
|
|
|
FOR_UNROLL (short row = 0; row < N_R0_MXFP4; row++) {
|
|
device const block_mxfp4 &bg = xg[row * row_blocks + ib];
|
|
device const block_mxfp4 &bu = xu[row * row_blocks + ib];
|
|
device const uchar *qg = bg.qs + 8 * it;
|
|
device const uchar *qu = bu.qs + 8 * it;
|
|
|
|
float4 ag = yl0 * float4(lut[qg[0] & 15], lut[qg[1] & 15],
|
|
lut[qg[2] & 15], lut[qg[3] & 15]);
|
|
ag += yl1 * float4(lut[qg[0] >> 4], lut[qg[1] >> 4],
|
|
lut[qg[2] >> 4], lut[qg[3] >> 4]);
|
|
ag += yl2 * float4(lut[qg[4] & 15], lut[qg[5] & 15],
|
|
lut[qg[6] & 15], lut[qg[7] & 15]);
|
|
ag += yl3 * float4(lut[qg[4] >> 4], lut[qg[5] >> 4],
|
|
lut[qg[6] >> 4], lut[qg[7] >> 4]);
|
|
|
|
float4 au = yl0 * float4(lut[qu[0] & 15], lut[qu[1] & 15],
|
|
lut[qu[2] & 15], lut[qu[3] & 15]);
|
|
au += yl1 * float4(lut[qu[0] >> 4], lut[qu[1] >> 4],
|
|
lut[qu[2] >> 4], lut[qu[3] >> 4]);
|
|
au += yl2 * float4(lut[qu[4] & 15], lut[qu[5] & 15],
|
|
lut[qu[6] & 15], lut[qu[7] & 15]);
|
|
au += yl3 * float4(lut[qu[4] >> 4], lut[qu[5] >> 4],
|
|
lut[qu[6] >> 4], lut[qu[7] >> 4]);
|
|
|
|
sumg[row] += ds4_metal_e8m0_to_f32(bg.e) *
|
|
((ag.x + ag.y) + (ag.z + ag.w));
|
|
sumu[row] += ds4_metal_e8m0_to_f32(bu.e) *
|
|
((au.x + au.y) + (au.z + au.w));
|
|
}
|
|
yb += 16 * QK_MXFP4;
|
|
}
|
|
|
|
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 < N_R0_MXFP4; row++) {
|
|
if (first_row + row < args.ne0) {
|
|
const float gate = simd_sum(sumg[row]);
|
|
const float up = simd_sum(sumu[row]);
|
|
if (tiisg == 0) {
|
|
const uint out_row = first_row + row;
|
|
float g = gate;
|
|
float u = up;
|
|
if (act.clamp_value > 1.0e-6f) {
|
|
g = min(g, act.clamp_value);
|
|
u = clamp(u, -act.clamp_value, act.clamp_value);
|
|
}
|
|
gate_f32[out_row] = gate;
|
|
up_f32[out_row] = up;
|
|
mid_f32[out_row] = (g / (1.0f + exp(-g))) * u * route_weight;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
kernel void kernel_mul_mv_id_mxfp4_pair_swiglu_f32(
|
|
constant ds4_metal_args_mul_mv_id &args,
|
|
constant ds4_metal_dsv4_moe_swiglu_weight_args &act,
|
|
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,
|
|
device const char *ids,
|
|
device const char *weights,
|
|
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]]) {
|
|
const int iid1 = tgpig.z / args.nei0;
|
|
const int idx = tgpig.z % args.nei0;
|
|
const int32_t expert = ((device const int32_t *)(ids + (uint64_t)iid1 * args.nbi1))[idx];
|
|
if (!ds4_tp_owns_expert(expert, args.ne02, args.tp_rank, args.tp_world)) return;
|
|
|
|
const uint64_t pair_row = (uint64_t)iid1 * args.nei0 + (uint64_t)idx;
|
|
device const float *route =
|
|
(device const float *)(weights + pair_row * act.weight_stride);
|
|
device char *gate_cur = dst_gate + pair_row * args.ne0 * sizeof(float);
|
|
device char *up_cur = dst_up + pair_row * args.ne0 * sizeof(float);
|
|
device char *mid_cur = dst_mid + pair_row * act.mid_row_stride;
|
|
device const char *x_cur = src1 + (uint64_t)(idx % args.ne11) * args.nb11 +
|
|
(uint64_t)iid1 * args.nb12;
|
|
device const char *gate_expert = src0_gate +
|
|
(int64_t)(expert - args.tp_expert_base) * args.nb02;
|
|
device const char *up_expert = src0_up +
|
|
(int64_t)(expert - args.tp_expert_base) * args.nb02;
|
|
tgpig.z = 0;
|
|
kernel_mul_mv_mxfp4_pair_swiglu_impl(args, act, gate_expert, up_expert,
|
|
x_cur, gate_cur, up_cur, mid_cur,
|
|
route[0], shmem, tgpig, tiisg, sgitg);
|
|
(void)tiitg;
|
|
}
|
|
|
|
kernel void kernel_mul_mv_id_mxfp4_pair_swiglu_fixed_route_f32(
|
|
constant ds4_metal_args_mul_mv_id &args,
|
|
constant ds4_metal_dsv4_moe_swiglu_weight_args &act,
|
|
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,
|
|
device const char *ids,
|
|
device const char *weights,
|
|
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]]) {
|
|
const int idx = (int)tgpig.z;
|
|
const int32_t expert = ((device const int32_t *)ids)[idx];
|
|
const uint64_t pair_row = (uint64_t)idx;
|
|
device const float *route =
|
|
(device const float *)(weights + pair_row * act.weight_stride);
|
|
device char *gate_cur = dst_gate + pair_row * args.ne0 * sizeof(float);
|
|
device char *up_cur = dst_up + pair_row * args.ne0 * sizeof(float);
|
|
device char *mid_cur = dst_mid + pair_row * act.mid_row_stride;
|
|
device const char *gate_expert =
|
|
src0_gate + (int64_t)expert * args.nb02;
|
|
device const char *up_expert =
|
|
src0_up + (int64_t)expert * args.nb02;
|
|
tgpig.z = 0;
|
|
kernel_mul_mv_mxfp4_pair_swiglu_impl(args, act, gate_expert, up_expert,
|
|
src1, gate_cur, up_cur, mid_cur,
|
|
route[0], shmem, tgpig, tiisg, sgitg);
|
|
(void)tiitg;
|
|
}
|
|
|
|
/* Exact-shape sibling of the fixed-route decode pair-SwiGLU kernel. The host
|
|
* selects it only after proving expert_in_dim == 4096 and gate/up row bytes ==
|
|
* 2176, so the per-lane K walk has a compile-time trip count (128 blocks / 16
|
|
* lane pairs = 8 steps) and the row stride is a literal. That lets the
|
|
* compiler unroll the walk and keep several independent expert loads in
|
|
* flight; the per-lane block order, the per-block accumulate order and the
|
|
* simd_sum tree are byte-identical to kernel_mul_mv_mxfp4_pair_swiglu_impl. */
|
|
#define DS4_MXFP4_PAIR_STATIC_NB 128
|
|
#define DS4_MXFP4_PAIR_STATIC_ROW_BLOCKS 128
|
|
|
|
template<typename args_t>
|
|
void kernel_mul_mv_mxfp4_pair_swiglu_static_impl(
|
|
args_t args,
|
|
constant ds4_metal_dsv4_moe_swiglu_weight_args &act,
|
|
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,
|
|
float route_weight,
|
|
threadgroup char *shmem,
|
|
uint3 tgpig,
|
|
ushort tiisg,
|
|
ushort sgitg) {
|
|
const short NSG = FC_mul_mv_nsg;
|
|
const int first_row = (tgpig.x * NSG + sgitg) * N_R0_MXFP4;
|
|
const short ix = tiisg / 2;
|
|
const short it = tiisg & 1;
|
|
|
|
threadgroup float *lut = (threadgroup float *)shmem;
|
|
if (sgitg == 0) lut[tiisg] = ds4_metal_mxfp4_values[tiisg & 15];
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
device const block_mxfp4 *xg =
|
|
(device const block_mxfp4 *)(src0_gate + (uint64_t)first_row * args.nb01);
|
|
device const block_mxfp4 *xu =
|
|
(device const block_mxfp4 *)(src0_up + (uint64_t)first_row * args.nb01);
|
|
device const float *yb = (device const float *)src1 + ix * QK_MXFP4 + it * 8;
|
|
float sumg[N_R0_MXFP4] = {0.f};
|
|
float sumu[N_R0_MXFP4] = {0.f};
|
|
|
|
for (int ib = ix; ib < DS4_MXFP4_PAIR_STATIC_NB; ib += 16) {
|
|
device const float4 *y4 = (device const float4 *)yb;
|
|
const float4 yl0 = y4[0];
|
|
const float4 yl1 = y4[4];
|
|
const float4 yl2 = y4[1];
|
|
const float4 yl3 = y4[5];
|
|
|
|
FOR_UNROLL (short row = 0; row < N_R0_MXFP4; row++) {
|
|
device const block_mxfp4 &bg =
|
|
xg[row * DS4_MXFP4_PAIR_STATIC_ROW_BLOCKS + ib];
|
|
device const block_mxfp4 &bu =
|
|
xu[row * DS4_MXFP4_PAIR_STATIC_ROW_BLOCKS + ib];
|
|
device const uchar *qg = bg.qs + 8 * it;
|
|
device const uchar *qu = bu.qs + 8 * it;
|
|
|
|
float4 ag = yl0 * float4(lut[qg[0] & 15], lut[qg[1] & 15],
|
|
lut[qg[2] & 15], lut[qg[3] & 15]);
|
|
ag += yl1 * float4(lut[qg[0] >> 4], lut[qg[1] >> 4],
|
|
lut[qg[2] >> 4], lut[qg[3] >> 4]);
|
|
ag += yl2 * float4(lut[qg[4] & 15], lut[qg[5] & 15],
|
|
lut[qg[6] & 15], lut[qg[7] & 15]);
|
|
ag += yl3 * float4(lut[qg[4] >> 4], lut[qg[5] >> 4],
|
|
lut[qg[6] >> 4], lut[qg[7] >> 4]);
|
|
|
|
float4 au = yl0 * float4(lut[qu[0] & 15], lut[qu[1] & 15],
|
|
lut[qu[2] & 15], lut[qu[3] & 15]);
|
|
au += yl1 * float4(lut[qu[0] >> 4], lut[qu[1] >> 4],
|
|
lut[qu[2] >> 4], lut[qu[3] >> 4]);
|
|
au += yl2 * float4(lut[qu[4] & 15], lut[qu[5] & 15],
|
|
lut[qu[6] & 15], lut[qu[7] & 15]);
|
|
au += yl3 * float4(lut[qu[4] >> 4], lut[qu[5] >> 4],
|
|
lut[qu[6] >> 4], lut[qu[7] >> 4]);
|
|
|
|
sumg[row] += ds4_metal_e8m0_to_f32(bg.e) *
|
|
((ag.x + ag.y) + (ag.z + ag.w));
|
|
sumu[row] += ds4_metal_e8m0_to_f32(bu.e) *
|
|
((au.x + au.y) + (au.z + au.w));
|
|
}
|
|
yb += 16 * QK_MXFP4;
|
|
}
|
|
|
|
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 < N_R0_MXFP4; row++) {
|
|
if (first_row + row < args.ne0) {
|
|
const float gate = simd_sum(sumg[row]);
|
|
const float up = simd_sum(sumu[row]);
|
|
if (tiisg == 0) {
|
|
const uint out_row = first_row + row;
|
|
float g = gate;
|
|
float u = up;
|
|
if (act.clamp_value > 1.0e-6f) {
|
|
g = min(g, act.clamp_value);
|
|
u = clamp(u, -act.clamp_value, act.clamp_value);
|
|
}
|
|
gate_f32[out_row] = gate;
|
|
up_f32[out_row] = up;
|
|
mid_f32[out_row] = (g / (1.0f + exp(-g))) * u * route_weight;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
kernel void kernel_mul_mv_id_mxfp4_pair_swiglu_fixed_route_static_f32(
|
|
constant ds4_metal_args_mul_mv_id &args,
|
|
constant ds4_metal_dsv4_moe_swiglu_weight_args &act,
|
|
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,
|
|
device const char *ids,
|
|
device const char *weights,
|
|
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]]) {
|
|
const int idx = (int)tgpig.z;
|
|
const int32_t expert = ((device const int32_t *)ids)[idx];
|
|
const uint64_t pair_row = (uint64_t)idx;
|
|
device const float *route =
|
|
(device const float *)(weights + pair_row * act.weight_stride);
|
|
device char *gate_cur = dst_gate + pair_row * args.ne0 * sizeof(float);
|
|
device char *up_cur = dst_up + pair_row * args.ne0 * sizeof(float);
|
|
device char *mid_cur = dst_mid + pair_row * act.mid_row_stride;
|
|
device const char *gate_expert =
|
|
src0_gate + (int64_t)expert * args.nb02;
|
|
device const char *up_expert =
|
|
src0_up + (int64_t)expert * args.nb02;
|
|
tgpig.z = 0;
|
|
/* Defensive: the host proves this shape before selecting the pipeline, so
|
|
* the dynamic sibling here is unreachable in production. Keeping it makes
|
|
* a mis-selected pipeline produce identical results instead of garbage. */
|
|
if (args.ne00 == DS4_MXFP4_PAIR_STATIC_NB * QK_MXFP4 &&
|
|
args.nb01 == (uint64_t)DS4_MXFP4_PAIR_STATIC_ROW_BLOCKS *
|
|
sizeof(block_mxfp4)) {
|
|
kernel_mul_mv_mxfp4_pair_swiglu_static_impl(
|
|
args, act, gate_expert, up_expert, src1, gate_cur, up_cur, mid_cur,
|
|
route[0], shmem, tgpig, tiisg, sgitg);
|
|
} else {
|
|
kernel_mul_mv_mxfp4_pair_swiglu_impl(
|
|
args, act, gate_expert, up_expert, src1, gate_cur, up_cur, mid_cur,
|
|
route[0], shmem, tgpig, tiisg, sgitg);
|
|
}
|
|
(void)tiitg;
|
|
}
|
|
|
|
kernel void kernel_mul_mv_slots6_mxfp4_pair_swiglu_f32(
|
|
constant ds4_metal_args_mul_mv_id &args,
|
|
constant ds4_metal_dsv4_moe_swiglu_weight_args &act,
|
|
device const char *gate0, device const char *gate1,
|
|
device const char *gate2, device const char *gate3,
|
|
device const char *gate4, device const char *gate5,
|
|
device const char *up0, device const char *up1,
|
|
device const char *up2, device const char *up3,
|
|
device const char *up4, device const char *up5,
|
|
device const char *src1,
|
|
device char *dst_gate,
|
|
device char *dst_up,
|
|
device char *dst_mid,
|
|
device const char *weights,
|
|
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]]) {
|
|
const int iid1 = tgpig.z / args.nei0;
|
|
const int idx = tgpig.z % args.nei0;
|
|
device const char *gate_expert = gate0;
|
|
device const char *up_expert = up0;
|
|
switch (idx) {
|
|
case 1: gate_expert = gate1; up_expert = up1; break;
|
|
case 2: gate_expert = gate2; up_expert = up2; break;
|
|
case 3: gate_expert = gate3; up_expert = up3; break;
|
|
case 4: gate_expert = gate4; up_expert = up4; break;
|
|
case 5: gate_expert = gate5; up_expert = up5; break;
|
|
default: break;
|
|
}
|
|
|
|
const uint64_t pair_row = (uint64_t)iid1 * args.nei0 + (uint64_t)idx;
|
|
device const float *route =
|
|
(device const float *)(weights + pair_row * act.weight_stride);
|
|
device char *gate_cur = dst_gate + pair_row * args.ne0 * sizeof(float);
|
|
device char *up_cur = dst_up + pair_row * args.ne0 * sizeof(float);
|
|
device char *mid_cur = dst_mid + pair_row * act.mid_row_stride;
|
|
device const char *x_cur = src1 + (uint64_t)(idx % args.ne11) * args.nb11 +
|
|
(uint64_t)iid1 * args.nb12;
|
|
tgpig.z = 0;
|
|
kernel_mul_mv_mxfp4_pair_swiglu_impl(args, act, gate_expert, up_expert,
|
|
x_cur, gate_cur, up_cur, mid_cur,
|
|
route[0], shmem, tgpig, tiisg, sgitg);
|
|
(void)tiitg;
|
|
}
|
|
|
|
kernel void kernel_mul_mv_table_q4_K_pair_swiglu_f32(
|
|
constant ds4_metal_args_mul_mv_id & args,
|
|
constant ds4_metal_dsv4_moe_swiglu_weight_args & act,
|
|
device const ds4_metal_q4_expert_table & gate_table,
|
|
device const ds4_metal_q4_expert_table & up_table,
|
|
device const char * src1,
|
|
device char * dst_gate,
|
|
device char * dst_up,
|
|
device char * dst_mid,
|
|
device const char * ids,
|
|
device const char * weights,
|
|
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]]) {
|
|
const int iid1 = tgpig.z / args.nei0;
|
|
const int idx = tgpig.z % args.nei0;
|
|
|
|
tgpig.z = 0;
|
|
|
|
const int32_t i02 = ((device const int32_t *)(ids + iid1 * args.nbi1))[idx];
|
|
if (i02 < 0 || i02 >= args.ne02 || i02 >= 384) {
|
|
return;
|
|
}
|
|
const int64_t i11 = idx % args.ne11;
|
|
const int64_t i12 = iid1;
|
|
|
|
device const char *src0_gate_cur = gate_table.experts[(uint)i02];
|
|
device const char *src0_up_cur = up_table.experts[(uint)i02];
|
|
device const char *src1_cur = src1 + i11 * args.nb11 + i12 * args.nb12;
|
|
|
|
device char *dst_gate_cur = dst_gate + (idx * args.ne0 + i12 * args.ne1 * args.ne0) * sizeof(float);
|
|
device char *dst_up_cur = dst_up + (idx * args.ne0 + i12 * args.ne1 * args.ne0) * sizeof(float);
|
|
|
|
const short NSG = FC_mul_mv_nsg;
|
|
constexpr uint16_t kmask1 = 0x3f3f;
|
|
constexpr uint16_t kmask2 = 0x0f0f;
|
|
constexpr uint16_t kmask3 = 0xc0c0;
|
|
|
|
const short ix = tiisg / 8;
|
|
const short it = tiisg % 8;
|
|
const short iq = it / 4;
|
|
const short ir = it % 4;
|
|
const int nb = args.ne00 / QK_K;
|
|
const int first_row = (tgpig.x * NSG + sgitg) * N_R0_Q4_K;
|
|
device float *gate_f32 = (device float *)dst_gate_cur;
|
|
device float *up_f32 = (device float *)dst_up_cur;
|
|
const uint64_t pair_row = (uint64_t)i12 * (uint64_t)args.nei0 + (uint64_t)idx;
|
|
device float *mid_f32 = (device float *)(dst_mid + pair_row * act.mid_row_stride);
|
|
device const float *route_w = (device const float *)(weights + pair_row * act.weight_stride);
|
|
const float c = act.clamp_value;
|
|
const float route_weight = route_w[0];
|
|
|
|
device const block_q4_K *xg =
|
|
(device const block_q4_K *)(src0_gate_cur + (uint64_t)first_row * args.nb01);
|
|
device const block_q4_K *xu =
|
|
(device const block_q4_K *)(src0_up_cur + (uint64_t)first_row * args.nb01);
|
|
device const float *y = (device const float *)src1_cur;
|
|
device const float *y4 = y + ix * QK_K + 64 * iq + 8 * ir;
|
|
|
|
float sumg[N_R0_Q4_K] = {0.f};
|
|
float sumu[N_R0_Q4_K] = {0.f};
|
|
uint16_t sc16[4];
|
|
thread const uint8_t *sc8 = (thread const uint8_t *)sc16;
|
|
|
|
for (int ib = ix; ib < nb; ib += 4) {
|
|
float yl[16];
|
|
float yh[16];
|
|
float4 sumy = {0.f, 0.f, 0.f, 0.f};
|
|
|
|
for (short i = 0; i < 8; ++i) {
|
|
yl[i + 0] = y4[i + 0]; sumy[0] += yl[i + 0];
|
|
yl[i + 8] = y4[i + 32]; sumy[1] += yl[i + 8];
|
|
yh[i + 0] = y4[i + 128]; sumy[2] += yh[i + 0];
|
|
yh[i + 8] = y4[i + 160]; sumy[3] += yh[i + 8];
|
|
}
|
|
|
|
device const uint16_t *scg = (device const uint16_t *)xg[ib].scales + iq;
|
|
device const uint16_t *qg1 = (device const uint16_t *)xg[ib].qs + 16 * iq + 4 * ir;
|
|
device const half *dhg = &xg[ib].d;
|
|
device const uint16_t *scu = (device const uint16_t *)xu[ib].scales + iq;
|
|
device const uint16_t *qu1 = (device const uint16_t *)xu[ib].qs + 16 * iq + 4 * ir;
|
|
device const half *dhu = &xu[ib].d;
|
|
|
|
for (short row = 0; row < N_R0_Q4_K; row++) {
|
|
sc16[0] = scg[0] & kmask1;
|
|
sc16[1] = scg[2] & kmask1;
|
|
sc16[2] = ((scg[4] >> 0) & kmask2) | ((scg[0] & kmask3) >> 2);
|
|
sc16[3] = ((scg[4] >> 4) & kmask2) | ((scg[2] & kmask3) >> 2);
|
|
|
|
device const uint16_t *qg2 = qg1 + 32;
|
|
float4 acc1g = {0.f, 0.f, 0.f, 0.f};
|
|
float4 acc2g = {0.f, 0.f, 0.f, 0.f};
|
|
|
|
FOR_UNROLL (short i = 0; i < 4; ++i) {
|
|
acc1g[0] += yl[2 * i + 0] * (qg1[i] & 0x000F);
|
|
acc1g[1] += yl[2 * i + 1] * (qg1[i] & 0x0F00);
|
|
acc1g[2] += yl[2 * i + 8] * (qg1[i] & 0x00F0);
|
|
acc1g[3] += yl[2 * i + 9] * (qg1[i] & 0xF000);
|
|
acc2g[0] += yh[2 * i + 0] * (qg2[i] & 0x000F);
|
|
acc2g[1] += yh[2 * i + 1] * (qg2[i] & 0x0F00);
|
|
acc2g[2] += yh[2 * i + 8] * (qg2[i] & 0x00F0);
|
|
acc2g[3] += yh[2 * i + 9] * (qg2[i] & 0xF000);
|
|
}
|
|
|
|
sumg[row] += dhg[0] * ((acc1g[0] + 1.f / 256.f * acc1g[1]) * sc8[0] +
|
|
(acc1g[2] + 1.f / 256.f * acc1g[3]) * sc8[1] * 1.f / 16.f +
|
|
(acc2g[0] + 1.f / 256.f * acc2g[1]) * sc8[4] +
|
|
(acc2g[2] + 1.f / 256.f * acc2g[3]) * sc8[5] * 1.f / 16.f) -
|
|
dhg[1] * (sumy[0] * sc8[2] + sumy[1] * sc8[3] +
|
|
sumy[2] * sc8[6] + sumy[3] * sc8[7]);
|
|
|
|
sc16[0] = scu[0] & kmask1;
|
|
sc16[1] = scu[2] & kmask1;
|
|
sc16[2] = ((scu[4] >> 0) & kmask2) | ((scu[0] & kmask3) >> 2);
|
|
sc16[3] = ((scu[4] >> 4) & kmask2) | ((scu[2] & kmask3) >> 2);
|
|
|
|
device const uint16_t *qu2 = qu1 + 32;
|
|
float4 acc1u = {0.f, 0.f, 0.f, 0.f};
|
|
float4 acc2u = {0.f, 0.f, 0.f, 0.f};
|
|
|
|
FOR_UNROLL (short i = 0; i < 4; ++i) {
|
|
acc1u[0] += yl[2 * i + 0] * (qu1[i] & 0x000F);
|
|
acc1u[1] += yl[2 * i + 1] * (qu1[i] & 0x0F00);
|
|
acc1u[2] += yl[2 * i + 8] * (qu1[i] & 0x00F0);
|
|
acc1u[3] += yl[2 * i + 9] * (qu1[i] & 0xF000);
|
|
acc2u[0] += yh[2 * i + 0] * (qu2[i] & 0x000F);
|
|
acc2u[1] += yh[2 * i + 1] * (qu2[i] & 0x0F00);
|
|
acc2u[2] += yh[2 * i + 8] * (qu2[i] & 0x00F0);
|
|
acc2u[3] += yh[2 * i + 9] * (qu2[i] & 0xF000);
|
|
}
|
|
|
|
sumu[row] += dhu[0] * ((acc1u[0] + 1.f / 256.f * acc1u[1]) * sc8[0] +
|
|
(acc1u[2] + 1.f / 256.f * acc1u[3]) * sc8[1] * 1.f / 16.f +
|
|
(acc2u[0] + 1.f / 256.f * acc2u[1]) * sc8[4] +
|
|
(acc2u[2] + 1.f / 256.f * acc2u[3]) * sc8[5] * 1.f / 16.f) -
|
|
dhu[1] * (sumy[0] * sc8[2] + sumy[1] * sc8[3] +
|
|
sumy[2] * sc8[6] + sumy[3] * sc8[7]);
|
|
|
|
qg1 += args.nb01 / 2;
|
|
scg += args.nb01 / 2;
|
|
dhg += args.nb01 / 2;
|
|
qu1 += args.nb01 / 2;
|
|
scu += args.nb01 / 2;
|
|
dhu += args.nb01 / 2;
|
|
}
|
|
|
|
y4 += 4 * QK_K;
|
|
}
|
|
|
|
for (int row = 0; row < N_R0_Q4_K && first_row + row < args.ne0; ++row) {
|
|
const float gate = simd_sum(sumg[row]);
|
|
const float up = simd_sum(sumu[row]);
|
|
if (tiisg == 0) {
|
|
const uint out_row = first_row + row;
|
|
float g = gate;
|
|
float u = up;
|
|
if (c > 1.0e-6f) {
|
|
g = min(g, c);
|
|
u = clamp(u, -c, c);
|
|
}
|
|
gate_f32[out_row] = gate;
|
|
up_f32[out_row] = up;
|
|
const float silu = g / (1.0f + exp(-g));
|
|
mid_f32[out_row] = silu * u * route_weight;
|
|
}
|
|
}
|
|
|
|
(void)tiitg;
|
|
}
|
|
|
|
kernel void kernel_mul_mv_addr_q4_K_pair_swiglu_f32(
|
|
constant ds4_metal_args_mul_mv_id & args,
|
|
constant ds4_metal_dsv4_moe_swiglu_weight_args & act,
|
|
device const ulong * gate_addrs,
|
|
device const ulong * up_addrs,
|
|
device const char * src1,
|
|
device char * dst_gate,
|
|
device char * dst_up,
|
|
device char * dst_mid,
|
|
device const char * ids,
|
|
device const char * weights,
|
|
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]]) {
|
|
const int iid1 = tgpig.z / args.nei0;
|
|
const int idx = tgpig.z % args.nei0;
|
|
|
|
tgpig.z = 0;
|
|
|
|
const int32_t i02 = ((device const int32_t *)(ids + iid1 * args.nbi1))[idx];
|
|
if (i02 < 0 || i02 >= args.ne02 || i02 >= 384) {
|
|
return;
|
|
}
|
|
const int64_t i11 = idx % args.ne11;
|
|
const int64_t i12 = iid1;
|
|
|
|
device const char *src0_gate_cur =
|
|
reinterpret_cast<device const char *>(gate_addrs[(uint)i02]);
|
|
device const char *src0_up_cur =
|
|
reinterpret_cast<device const char *>(up_addrs[(uint)i02]);
|
|
device const char *src1_cur = src1 + i11 * args.nb11 + i12 * args.nb12;
|
|
|
|
device char *dst_gate_cur = dst_gate + (idx * args.ne0 + i12 * args.ne1 * args.ne0) * sizeof(float);
|
|
device char *dst_up_cur = dst_up + (idx * args.ne0 + i12 * args.ne1 * args.ne0) * sizeof(float);
|
|
|
|
ds4_metal_args_mul_mv args0 = {
|
|
args.ne00, args.ne01, 1,
|
|
args.nb00, args.nb01, args.nb02, args.nb02,
|
|
args.ne10, 1, 1,
|
|
args.nb10, args.nb11, args.nb12, args.nb12,
|
|
args.ne0, 1, args.nr0, 1, 1,
|
|
};
|
|
|
|
kernel_mul_mv_q4_K_f32_impl<N_R0_Q4_K>(
|
|
args0,
|
|
src0_gate_cur,
|
|
src1_cur,
|
|
dst_gate_cur,
|
|
shmem,
|
|
tgpig,
|
|
tiisg,
|
|
sgitg);
|
|
kernel_mul_mv_q4_K_f32_impl<N_R0_Q4_K>(
|
|
args0,
|
|
src0_up_cur,
|
|
src1_cur,
|
|
dst_up_cur,
|
|
shmem,
|
|
tgpig,
|
|
tiisg,
|
|
sgitg);
|
|
|
|
const short NSG = FC_mul_mv_nsg;
|
|
const int first_row = (tgpig.x * NSG + sgitg) * N_R0_Q4_K;
|
|
device float *gate_f32 = (device float *)dst_gate_cur;
|
|
device float *up_f32 = (device float *)dst_up_cur;
|
|
const uint64_t pair_row = (uint64_t)i12 * (uint64_t)args.nei0 + (uint64_t)idx;
|
|
device float *mid_f32 = (device float *)(dst_mid + pair_row * act.mid_row_stride);
|
|
device const float *route_w = (device const float *)(weights + pair_row * act.weight_stride);
|
|
const float c = act.clamp_value;
|
|
const float route_weight = route_w[0];
|
|
|
|
if (tiisg == 0) {
|
|
for (int row = 0; row < N_R0_Q4_K && first_row + row < args.ne0; ++row) {
|
|
const uint out_row = first_row + row;
|
|
float g = gate_f32[out_row];
|
|
float u = up_f32[out_row];
|
|
if (c > 1.0e-6f) {
|
|
g = min(g, c);
|
|
u = clamp(u, -c, c);
|
|
}
|
|
const float silu = g / (1.0f + exp(-g));
|
|
mid_f32[out_row] = silu * u * route_weight;
|
|
}
|
|
}
|
|
|
|
(void)tiitg;
|
|
}
|
|
|
|
kernel void kernel_q4_gather_slots6(
|
|
constant ds4_metal_q4_gather_slots6_args &args,
|
|
device const char *src_group0,
|
|
device const char *src_group1,
|
|
device const char *src_group2,
|
|
device const char *src_group3,
|
|
device const char *src_group4,
|
|
device const char *src_group5,
|
|
device const int32_t *ids,
|
|
device char *dst,
|
|
uint3 tgpig [[threadgroup_position_in_grid]],
|
|
uint tiitg [[thread_index_in_threadgroup]]) {
|
|
const uint slot = tgpig.y;
|
|
if (slot >= args.n_slots || args.group_size == 0) return;
|
|
|
|
const int32_t expert = ids[slot];
|
|
if (expert < 0) return;
|
|
|
|
const uint expert_u = (uint)expert;
|
|
const uint group = expert_u / args.group_size;
|
|
if (group >= 6) return;
|
|
|
|
const uint local_expert = expert_u - group * args.group_size;
|
|
device const char *src_group = src_group0;
|
|
switch (group) {
|
|
case 1: src_group = src_group1; break;
|
|
case 2: src_group = src_group2; break;
|
|
case 3: src_group = src_group3; break;
|
|
case 4: src_group = src_group4; break;
|
|
case 5: src_group = src_group5; break;
|
|
default: break;
|
|
}
|
|
|
|
const uint64_t chunk = (uint64_t)tgpig.x * 256ul + (uint64_t)tiitg;
|
|
const uint64_t n_chunks = args.expert_bytes >> 4;
|
|
if (chunk >= n_chunks) return;
|
|
|
|
device const uint4 *src =
|
|
(device const uint4 *)(src_group + (uint64_t)local_expert * args.expert_bytes);
|
|
device uint4 *out =
|
|
(device uint4 *)(dst + (uint64_t)slot * args.expert_bytes);
|
|
out[chunk] = src[chunk];
|
|
}
|
|
|
|
kernel void kernel_mul_mv_slots6_q4_K_pair_swiglu_f32(
|
|
constant ds4_metal_args_mul_mv_id & args,
|
|
constant ds4_metal_dsv4_moe_swiglu_weight_args & act,
|
|
device const char * src0_gate0,
|
|
device const char * src0_gate1,
|
|
device const char * src0_gate2,
|
|
device const char * src0_gate3,
|
|
device const char * src0_gate4,
|
|
device const char * src0_gate5,
|
|
device const char * src0_up0,
|
|
device const char * src0_up1,
|
|
device const char * src0_up2,
|
|
device const char * src0_up3,
|
|
device const char * src0_up4,
|
|
device const char * src0_up5,
|
|
device const char * src1,
|
|
device char * dst_gate,
|
|
device char * dst_up,
|
|
device char * dst_mid,
|
|
device const char * weights,
|
|
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]]) {
|
|
const int iid1 = tgpig.z / args.nei0;
|
|
const int idx = tgpig.z % args.nei0;
|
|
|
|
tgpig.z = 0;
|
|
|
|
const int64_t i11 = idx % args.ne11;
|
|
const int64_t i12 = iid1;
|
|
|
|
device const char *src0_gate_cur = src0_gate0;
|
|
device const char *src0_up_cur = src0_up0;
|
|
switch (idx) {
|
|
case 1: src0_gate_cur = src0_gate1; src0_up_cur = src0_up1; break;
|
|
case 2: src0_gate_cur = src0_gate2; src0_up_cur = src0_up2; break;
|
|
case 3: src0_gate_cur = src0_gate3; src0_up_cur = src0_up3; break;
|
|
case 4: src0_gate_cur = src0_gate4; src0_up_cur = src0_up4; break;
|
|
case 5: src0_gate_cur = src0_gate5; src0_up_cur = src0_up5; break;
|
|
default: break;
|
|
}
|
|
|
|
device const char *src1_cur = src1 + i11 * args.nb11 + i12 * args.nb12;
|
|
|
|
device char *dst_gate_cur = dst_gate + (idx * args.ne0 + i12 * args.ne1 * args.ne0) * sizeof(float);
|
|
device char *dst_up_cur = dst_up + (idx * args.ne0 + i12 * args.ne1 * args.ne0) * sizeof(float);
|
|
|
|
ds4_metal_args_mul_mv args0 = {
|
|
args.ne00, args.ne01, 1,
|
|
args.nb00, args.nb01, args.nb02, args.nb02,
|
|
args.ne10, 1, 1,
|
|
args.nb10, args.nb11, args.nb12, args.nb12,
|
|
args.ne0, 1, args.nr0, 1, 1,
|
|
};
|
|
|
|
kernel_mul_mv_q4_K_f32_impl<N_R0_Q4_K>(
|
|
args0,
|
|
src0_gate_cur,
|
|
src1_cur,
|
|
dst_gate_cur,
|
|
shmem,
|
|
tgpig,
|
|
tiisg,
|
|
sgitg);
|
|
kernel_mul_mv_q4_K_f32_impl<N_R0_Q4_K>(
|
|
args0,
|
|
src0_up_cur,
|
|
src1_cur,
|
|
dst_up_cur,
|
|
shmem,
|
|
tgpig,
|
|
tiisg,
|
|
sgitg);
|
|
|
|
const short NSG = FC_mul_mv_nsg;
|
|
const int first_row = (tgpig.x * NSG + sgitg) * N_R0_Q4_K;
|
|
device float *gate_f32 = (device float *)dst_gate_cur;
|
|
device float *up_f32 = (device float *)dst_up_cur;
|
|
const uint64_t pair_row = (uint64_t)i12 * (uint64_t)args.nei0 + (uint64_t)idx;
|
|
device float *mid_f32 = (device float *)(dst_mid + pair_row * act.mid_row_stride);
|
|
device const float *route_w = (device const float *)(weights + pair_row * act.weight_stride);
|
|
const float c = act.clamp_value;
|
|
const float route_weight = route_w[0];
|
|
|
|
if (tiisg == 0) {
|
|
for (int row = 0; row < N_R0_Q4_K && first_row + row < args.ne0; ++row) {
|
|
const uint out_row = first_row + row;
|
|
float g = gate_f32[out_row];
|
|
float u = up_f32[out_row];
|
|
if (c > 1.0e-6f) {
|
|
g = min(g, c);
|
|
u = clamp(u, -c, c);
|
|
}
|
|
const float silu = g / (1.0f + exp(-g));
|
|
mid_f32[out_row] = silu * u * route_weight;
|
|
}
|
|
}
|
|
|
|
(void)tiitg;
|
|
}
|
|
|
|
static inline device const char *ds4_q4_group24_select(
|
|
uint32_t group_id,
|
|
device const char *src00,
|
|
device const char *src01,
|
|
device const char *src02,
|
|
device const char *src03,
|
|
device const char *src04,
|
|
device const char *src05,
|
|
device const char *src06,
|
|
device const char *src07,
|
|
device const char *src08,
|
|
device const char *src09,
|
|
device const char *src10,
|
|
device const char *src11,
|
|
device const char *src12,
|
|
device const char *src13,
|
|
device const char *src14,
|
|
device const char *src15,
|
|
device const char *src16,
|
|
device const char *src17,
|
|
device const char *src18,
|
|
device const char *src19,
|
|
device const char *src20,
|
|
device const char *src21,
|
|
device const char *src22,
|
|
device const char *src23) {
|
|
switch (group_id) {
|
|
case 1: return src01;
|
|
case 2: return src02;
|
|
case 3: return src03;
|
|
case 4: return src04;
|
|
case 5: return src05;
|
|
case 6: return src06;
|
|
case 7: return src07;
|
|
case 8: return src08;
|
|
case 9: return src09;
|
|
case 10: return src10;
|
|
case 11: return src11;
|
|
case 12: return src12;
|
|
case 13: return src13;
|
|
case 14: return src14;
|
|
case 15: return src15;
|
|
case 16: return src16;
|
|
case 17: return src17;
|
|
case 18: return src18;
|
|
case 19: return src19;
|
|
case 20: return src20;
|
|
case 21: return src21;
|
|
case 22: return src22;
|
|
case 23: return src23;
|
|
default: return src00;
|
|
}
|
|
}
|
|
|
|
kernel void kernel_mul_mv_group6_q4_K_pair_swiglu_f32(
|
|
constant ds4_metal_args_mul_mv_id & args,
|
|
constant ds4_metal_dsv4_moe_swiglu_weight_args & act,
|
|
device const char * src0_gate0,
|
|
device const char * src0_gate1,
|
|
device const char * src0_gate2,
|
|
device const char * src0_gate3,
|
|
device const char * src0_gate4,
|
|
device const char * src0_gate5,
|
|
device const char * src0_up0,
|
|
device const char * src0_up1,
|
|
device const char * src0_up2,
|
|
device const char * src0_up3,
|
|
device const char * src0_up4,
|
|
device const char * src0_up5,
|
|
device const char * src1,
|
|
device char * dst_gate,
|
|
device char * dst_up,
|
|
device char * dst_mid,
|
|
device const char * ids,
|
|
device const char * weights,
|
|
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 uint32_t expert_group_size = 64;
|
|
const int iid1 = tgpig.z / args.nei0;
|
|
const int idx = tgpig.z % args.nei0;
|
|
|
|
tgpig.z = 0;
|
|
|
|
const int32_t expert = ((device const int32_t *)(ids + iid1 * args.nbi1))[idx];
|
|
if (expert < 0) {
|
|
return;
|
|
}
|
|
const uint32_t expert_u = (uint32_t)expert;
|
|
const uint32_t group_id = expert_u / expert_group_size;
|
|
if (group_id >= 6) {
|
|
return;
|
|
}
|
|
const uint32_t expert_local = expert_u - group_id * expert_group_size;
|
|
|
|
device const char *src0_gate_cur = src0_gate0;
|
|
device const char *src0_up_cur = src0_up0;
|
|
switch (group_id) {
|
|
case 1: src0_gate_cur = src0_gate1; src0_up_cur = src0_up1; break;
|
|
case 2: src0_gate_cur = src0_gate2; src0_up_cur = src0_up2; break;
|
|
case 3: src0_gate_cur = src0_gate3; src0_up_cur = src0_up3; break;
|
|
case 4: src0_gate_cur = src0_gate4; src0_up_cur = src0_up4; break;
|
|
case 5: src0_gate_cur = src0_gate5; src0_up_cur = src0_up5; break;
|
|
default: break;
|
|
}
|
|
|
|
const int64_t i11 = idx % args.ne11;
|
|
const int64_t i12 = iid1;
|
|
|
|
src0_gate_cur += (uint64_t)expert_local * args.nb02;
|
|
src0_up_cur += (uint64_t)expert_local * args.nb02;
|
|
device const char *src1_cur = src1 + i11 * args.nb11 + i12 * args.nb12;
|
|
|
|
device char *dst_gate_cur = dst_gate + (idx * args.ne0 + i12 * args.ne1 * args.ne0) * sizeof(float);
|
|
device char *dst_up_cur = dst_up + (idx * args.ne0 + i12 * args.ne1 * args.ne0) * sizeof(float);
|
|
|
|
ds4_metal_args_mul_mv args0 = {
|
|
args.ne00, args.ne01, 1,
|
|
args.nb00, args.nb01, args.nb02, args.nb02,
|
|
args.ne10, 1, 1,
|
|
args.nb10, args.nb11, args.nb12, args.nb12,
|
|
args.ne0, 1, args.nr0, 1, 1,
|
|
};
|
|
|
|
kernel_mul_mv_q4_K_f32_impl<N_R0_Q4_K>(
|
|
args0,
|
|
src0_gate_cur,
|
|
src1_cur,
|
|
dst_gate_cur,
|
|
shmem,
|
|
tgpig,
|
|
tiisg,
|
|
sgitg);
|
|
kernel_mul_mv_q4_K_f32_impl<N_R0_Q4_K>(
|
|
args0,
|
|
src0_up_cur,
|
|
src1_cur,
|
|
dst_up_cur,
|
|
shmem,
|
|
tgpig,
|
|
tiisg,
|
|
sgitg);
|
|
|
|
const short NSG = FC_mul_mv_nsg;
|
|
const int first_row = (tgpig.x * NSG + sgitg) * N_R0_Q4_K;
|
|
device float *gate_f32 = (device float *)dst_gate_cur;
|
|
device float *up_f32 = (device float *)dst_up_cur;
|
|
const uint64_t pair_row = (uint64_t)i12 * (uint64_t)args.nei0 + (uint64_t)idx;
|
|
device float *mid_f32 = (device float *)(dst_mid + pair_row * act.mid_row_stride);
|
|
device const float *route_w = (device const float *)(weights + pair_row * act.weight_stride);
|
|
const float c = act.clamp_value;
|
|
const float route_weight = route_w[0];
|
|
|
|
if (tiisg == 0) {
|
|
for (int row = 0; row < N_R0_Q4_K && first_row + row < args.ne0; ++row) {
|
|
const uint out_row = first_row + row;
|
|
float g = gate_f32[out_row];
|
|
float u = up_f32[out_row];
|
|
if (c > 1.0e-6f) {
|
|
g = min(g, c);
|
|
u = clamp(u, -c, c);
|
|
}
|
|
const float silu = g / (1.0f + exp(-g));
|
|
mid_f32[out_row] = silu * u * route_weight;
|
|
}
|
|
}
|
|
|
|
(void)tiitg;
|
|
}
|
|
|
|
kernel void kernel_mul_mv_group8_q4_K_pair_swiglu_f32(
|
|
constant ds4_metal_args_mul_mv_id & args,
|
|
constant ds4_metal_dsv4_moe_swiglu_weight_args & act,
|
|
device const char * src0_gate0,
|
|
device const char * src0_gate1,
|
|
device const char * src0_gate2,
|
|
device const char * src0_gate3,
|
|
device const char * src0_gate4,
|
|
device const char * src0_gate5,
|
|
device const char * src0_gate6,
|
|
device const char * src0_gate7,
|
|
device const char * src0_up0,
|
|
device const char * src0_up1,
|
|
device const char * src0_up2,
|
|
device const char * src0_up3,
|
|
device const char * src0_up4,
|
|
device const char * src0_up5,
|
|
device const char * src0_up6,
|
|
device const char * src0_up7,
|
|
device const char * src1,
|
|
device char * dst_gate,
|
|
device char * dst_up,
|
|
device char * dst_mid,
|
|
device const char * ids,
|
|
device const char * weights,
|
|
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 uint32_t expert_group_size = 48;
|
|
const int iid1 = tgpig.z / args.nei0;
|
|
const int idx = tgpig.z % args.nei0;
|
|
|
|
tgpig.z = 0;
|
|
|
|
const int32_t expert = ((device const int32_t *)(ids + iid1 * args.nbi1))[idx];
|
|
if (expert < 0) {
|
|
return;
|
|
}
|
|
const uint32_t expert_u = (uint32_t)expert;
|
|
const uint32_t group_id = expert_u / expert_group_size;
|
|
if (group_id >= 8) {
|
|
return;
|
|
}
|
|
const uint32_t expert_local = expert_u - group_id * expert_group_size;
|
|
|
|
device const char *src0_gate_cur = src0_gate0;
|
|
device const char *src0_up_cur = src0_up0;
|
|
switch (group_id) {
|
|
case 1: src0_gate_cur = src0_gate1; src0_up_cur = src0_up1; break;
|
|
case 2: src0_gate_cur = src0_gate2; src0_up_cur = src0_up2; break;
|
|
case 3: src0_gate_cur = src0_gate3; src0_up_cur = src0_up3; break;
|
|
case 4: src0_gate_cur = src0_gate4; src0_up_cur = src0_up4; break;
|
|
case 5: src0_gate_cur = src0_gate5; src0_up_cur = src0_up5; break;
|
|
case 6: src0_gate_cur = src0_gate6; src0_up_cur = src0_up6; break;
|
|
case 7: src0_gate_cur = src0_gate7; src0_up_cur = src0_up7; break;
|
|
default: break;
|
|
}
|
|
|
|
const int64_t i11 = idx % args.ne11;
|
|
const int64_t i12 = iid1;
|
|
|
|
src0_gate_cur += (uint64_t)expert_local * args.nb02;
|
|
src0_up_cur += (uint64_t)expert_local * args.nb02;
|
|
device const char *src1_cur = src1 + i11 * args.nb11 + i12 * args.nb12;
|
|
|
|
device char *dst_gate_cur = dst_gate + (idx * args.ne0 + i12 * args.ne1 * args.ne0) * sizeof(float);
|
|
device char *dst_up_cur = dst_up + (idx * args.ne0 + i12 * args.ne1 * args.ne0) * sizeof(float);
|
|
|
|
ds4_metal_args_mul_mv args0 = {
|
|
args.ne00, args.ne01, 1,
|
|
args.nb00, args.nb01, args.nb02, args.nb02,
|
|
args.ne10, 1, 1,
|
|
args.nb10, args.nb11, args.nb12, args.nb12,
|
|
args.ne0, 1, args.nr0, 1, 1,
|
|
};
|
|
|
|
kernel_mul_mv_q4_K_f32_impl<N_R0_Q4_K>(
|
|
args0,
|
|
src0_gate_cur,
|
|
src1_cur,
|
|
dst_gate_cur,
|
|
shmem,
|
|
tgpig,
|
|
tiisg,
|
|
sgitg);
|
|
kernel_mul_mv_q4_K_f32_impl<N_R0_Q4_K>(
|
|
args0,
|
|
src0_up_cur,
|
|
src1_cur,
|
|
dst_up_cur,
|
|
shmem,
|
|
tgpig,
|
|
tiisg,
|
|
sgitg);
|
|
|
|
const short NSG = FC_mul_mv_nsg;
|
|
const int first_row = (tgpig.x * NSG + sgitg) * N_R0_Q4_K;
|
|
device float *gate_f32 = (device float *)dst_gate_cur;
|
|
device float *up_f32 = (device float *)dst_up_cur;
|
|
const uint64_t pair_row = (uint64_t)i12 * (uint64_t)args.nei0 + (uint64_t)idx;
|
|
device float *mid_f32 = (device float *)(dst_mid + pair_row * act.mid_row_stride);
|
|
device const float *route_w = (device const float *)(weights + pair_row * act.weight_stride);
|
|
const float c = act.clamp_value;
|
|
const float route_weight = route_w[0];
|
|
|
|
if (tiisg == 0) {
|
|
for (int row = 0; row < N_R0_Q4_K && first_row + row < args.ne0; ++row) {
|
|
const uint out_row = first_row + row;
|
|
float g = gate_f32[out_row];
|
|
float u = up_f32[out_row];
|
|
if (c > 1.0e-6f) {
|
|
g = min(g, c);
|
|
u = clamp(u, -c, c);
|
|
}
|
|
const float silu = g / (1.0f + exp(-g));
|
|
mid_f32[out_row] = silu * u * route_weight;
|
|
}
|
|
}
|
|
|
|
(void)tiitg;
|
|
}
|
|
|
|
kernel void kernel_mul_mv_group24_q4_K_id_f32(
|
|
constant ds4_metal_args_mul_mv_id & args,
|
|
device const char * src00,
|
|
device const char * src01,
|
|
device const char * src02,
|
|
device const char * src03,
|
|
device const char * src04,
|
|
device const char * src05,
|
|
device const char * src06,
|
|
device const char * src07,
|
|
device const char * src08,
|
|
device const char * src09,
|
|
device const char * src10,
|
|
device const char * src11,
|
|
device const char * src12,
|
|
device const char * src13,
|
|
device const char * src14,
|
|
device const char * src15,
|
|
device const char * src16,
|
|
device const char * src17,
|
|
device const char * src18,
|
|
device const char * src19,
|
|
device const char * src20,
|
|
device const char * src21,
|
|
device const char * src22,
|
|
device const char * src23,
|
|
device const char * src1,
|
|
device char * dst,
|
|
device const char * ids,
|
|
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 uint32_t expert_group_size = 16;
|
|
const int iid1 = tgpig.z / args.nei0;
|
|
const int idx = tgpig.z % args.nei0;
|
|
|
|
tgpig.z = 0;
|
|
|
|
const int32_t expert = ((device const int32_t *)(ids + iid1 * args.nbi1))[idx];
|
|
if (expert < 0) {
|
|
return;
|
|
}
|
|
const uint32_t expert_u = (uint32_t)expert;
|
|
const uint32_t group_id = expert_u / expert_group_size;
|
|
if (group_id >= 24) {
|
|
return;
|
|
}
|
|
const uint32_t expert_local = expert_u - group_id * expert_group_size;
|
|
|
|
device const char *src0_cur = ds4_q4_group24_select(group_id,
|
|
src00, src01, src02, src03,
|
|
src04, src05, src06, src07,
|
|
src08, src09, src10, src11,
|
|
src12, src13, src14, src15,
|
|
src16, src17, src18, src19,
|
|
src20, src21, src22, src23);
|
|
const int64_t i11 = idx % args.ne11;
|
|
const int64_t i12 = iid1;
|
|
|
|
src0_cur += (uint64_t)expert_local * args.nb02;
|
|
device const char *src1_cur = src1 + i11 * args.nb11 + i12 * args.nb12;
|
|
device char *dst_cur = dst + (idx * args.ne0 + i12 * args.ne1 * args.ne0) * sizeof(float);
|
|
|
|
ds4_metal_args_mul_mv args0 = {
|
|
args.ne00, args.ne01, 1,
|
|
args.nb00, args.nb01, args.nb02, args.nb02,
|
|
args.ne10, 1, 1,
|
|
args.nb10, args.nb11, args.nb12, args.nb12,
|
|
args.ne0, 1, args.nr0, 1, 1,
|
|
};
|
|
|
|
kernel_mul_mv_q4_K_f32_impl<N_R0_Q4_K>(
|
|
args0,
|
|
src0_cur,
|
|
src1_cur,
|
|
dst_cur,
|
|
shmem,
|
|
tgpig,
|
|
tiisg,
|
|
sgitg);
|
|
|
|
(void)tiitg;
|
|
}
|
|
|
|
kernel void kernel_mul_mv_group_q4_K_pair_swiglu_f32(
|
|
constant ds4_metal_args_mul_mv_id & args,
|
|
constant ds4_metal_dsv4_moe_swiglu_weight_args & act,
|
|
constant ds4_metal_moe_expert_group_args & group,
|
|
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,
|
|
device const char * ids,
|
|
device const char * weights,
|
|
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]]) {
|
|
const int iid1 = tgpig.z / args.nei0;
|
|
const int idx = tgpig.z % args.nei0;
|
|
|
|
tgpig.z = 0;
|
|
|
|
const int32_t expert_global = ((device const int32_t *)(ids + iid1 * args.nbi1))[idx];
|
|
if (expert_global < 0) {
|
|
return;
|
|
}
|
|
const uint32_t expert_u = (uint32_t)expert_global;
|
|
if (expert_u < group.expert_base ||
|
|
expert_u >= group.expert_base + group.expert_count) {
|
|
return;
|
|
}
|
|
const uint32_t expert_local = expert_u - group.expert_base;
|
|
|
|
const int64_t i11 = idx % args.ne11;
|
|
const int64_t i12 = iid1;
|
|
|
|
device const char *src0_gate_cur = src0_gate + (uint64_t)expert_local * args.nb02;
|
|
device const char *src0_up_cur = src0_up + (uint64_t)expert_local * args.nb02;
|
|
device const char *src1_cur = src1 + i11 * args.nb11 + i12 * args.nb12;
|
|
|
|
device char *dst_gate_cur = dst_gate + (idx * args.ne0 + i12 * args.ne1 * args.ne0) * sizeof(float);
|
|
device char *dst_up_cur = dst_up + (idx * args.ne0 + i12 * args.ne1 * args.ne0) * sizeof(float);
|
|
|
|
ds4_metal_args_mul_mv args0 = {
|
|
args.ne00, args.ne01, 1,
|
|
args.nb00, args.nb01, args.nb02, args.nb02,
|
|
args.ne10, 1, 1,
|
|
args.nb10, args.nb11, args.nb12, args.nb12,
|
|
args.ne0, 1, args.nr0, 1, 1,
|
|
};
|
|
|
|
kernel_mul_mv_q4_K_f32_impl<N_R0_Q4_K>(
|
|
args0,
|
|
src0_gate_cur,
|
|
src1_cur,
|
|
dst_gate_cur,
|
|
shmem,
|
|
tgpig,
|
|
tiisg,
|
|
sgitg);
|
|
kernel_mul_mv_q4_K_f32_impl<N_R0_Q4_K>(
|
|
args0,
|
|
src0_up_cur,
|
|
src1_cur,
|
|
dst_up_cur,
|
|
shmem,
|
|
tgpig,
|
|
tiisg,
|
|
sgitg);
|
|
|
|
const short NSG = FC_mul_mv_nsg;
|
|
const int first_row = (tgpig.x * NSG + sgitg) * N_R0_Q4_K;
|
|
device float *gate_f32 = (device float *)dst_gate_cur;
|
|
device float *up_f32 = (device float *)dst_up_cur;
|
|
const uint64_t pair_row = (uint64_t)i12 * (uint64_t)args.nei0 + (uint64_t)idx;
|
|
device float *mid_f32 = (device float *)(dst_mid + pair_row * act.mid_row_stride);
|
|
device const float *route_w = (device const float *)(weights + pair_row * act.weight_stride);
|
|
const float c = act.clamp_value;
|
|
const float route_weight = route_w[0];
|
|
|
|
if (tiisg == 0) {
|
|
for (int row = 0; row < N_R0_Q4_K && first_row + row < args.ne0; ++row) {
|
|
const uint out_row = first_row + row;
|
|
float g = gate_f32[out_row];
|
|
float u = up_f32[out_row];
|
|
if (c > 1.0e-6f) {
|
|
g = min(g, c);
|
|
u = clamp(u, -c, c);
|
|
}
|
|
const float silu = g / (1.0f + exp(-g));
|
|
mid_f32[out_row] = silu * u * route_weight;
|
|
}
|
|
}
|
|
|
|
(void)tiitg;
|
|
}
|
|
|
|
|
|
/* IQ2_XXS down projection summed over the token's selected experts (the
|
|
* GLM 5.2 routed blobs are IQ2_XXS end to end, unlike DS4 Flash's Q2_K
|
|
* down). Same contract as the q2_K sum kernel: mid rows already carry
|
|
* silu(gate)*up*route_weight, expert ownership honors the TP split, and
|
|
* add_in folds the shared-expert partial when tp_addend is set. */
|
|
kernel void kernel_mul_mv_id_iq2_xxs_sum6_f32(
|
|
constant ds4_metal_args_mul_mv_id & args,
|
|
device const char * src0s,
|
|
device const char * src1,
|
|
device char * dst,
|
|
device const char * ids,
|
|
device const char * add_in,
|
|
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]]) {
|
|
const short NSG = FC_mul_mv_nsg;
|
|
const short nr0 = N_R0_IQ2_XXS;
|
|
const int nb = args.ne00/QK_K;
|
|
const int nb32 = nb * (QK_K / 32);
|
|
const int first_row = (tgpig.x * NSG + sgitg) * nr0;
|
|
const uint token = tgpig.y;
|
|
device const int32_t *token_ids = (device const int32_t *)(ids + (uint64_t)token * args.nbi1);
|
|
device const char *token_src1 = src1 + (uint64_t)token * args.nb12;
|
|
|
|
float sumf[nr0] = {0.f};
|
|
float yl[32];
|
|
|
|
threadgroup uint64_t *svalues = (threadgroup uint64_t *)(shmem);
|
|
threadgroup uint8_t *ssigns = (threadgroup uint8_t *)(svalues + 256);
|
|
{
|
|
int nval = 4;
|
|
int pos = (32 * sgitg + tiisg) * nval;
|
|
for (int i = 0; i < nval; ++i) svalues[pos + i] = ds4_metal_iq2xxs_grid[pos + i];
|
|
nval = 2;
|
|
pos = (32 * sgitg + tiisg) * nval;
|
|
for (int i = 0; i < nval; ++i) ssigns[pos + i] = ds4_metal_ksigns_iq2xs[pos + i];
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
}
|
|
|
|
const int ix = tiisg;
|
|
|
|
for (int expert_slot = 0; expert_slot < args.nei0; expert_slot++) {
|
|
const int32_t expert = token_ids[expert_slot];
|
|
if (!ds4_tp_owns_expert(expert, args.ne02, args.tp_rank, args.tp_world)) continue;
|
|
device const block_iq2_xxs *x =
|
|
(device const block_iq2_xxs *)(src0s + (int64_t)(expert - args.tp_expert_base)*args.nb02 + first_row*args.nb01);
|
|
device const float *y = (device const float *)(token_src1 + expert_slot*args.nb11);
|
|
device const float *y4 = y + 32 * ix;
|
|
|
|
for (int ib32 = ix; ib32 < nb32; ib32 += 32) {
|
|
for (short i = 0; i < 32; ++i) {
|
|
yl[i] = y4[i];
|
|
}
|
|
|
|
const int ibl = ib32 / (QK_K / 32);
|
|
const int ib = ib32 % (QK_K / 32);
|
|
|
|
device const block_iq2_xxs *xr = x + ibl;
|
|
device const uint16_t *q2 = xr->qs + 4 * ib;
|
|
device const half *dh = &xr->d;
|
|
|
|
for (short row = 0; row < nr0; row++) {
|
|
if (first_row + row < args.ne0) {
|
|
device const uint8_t *aux8 = (device const uint8_t *)q2;
|
|
const uint32_t aux32 = q2[2] | (q2[3] << 16);
|
|
const float d = (float)dh[0] * (0.5f + (aux32 >> 28));
|
|
|
|
float s = 0;
|
|
for (short l = 0; l < 4; ++l) {
|
|
const threadgroup uint8_t *grid = (const threadgroup uint8_t *)(svalues + aux8[l]);
|
|
const uint8_t sign = ssigns[(aux32 >> 7 * l) & 127];
|
|
for (short j = 0; j < 8; ++j) {
|
|
s += yl[8 * l + j] * grid[j] *
|
|
(sign & ds4_metal_kmask_iq2xs[j] ? -1.f : 1.f);
|
|
}
|
|
}
|
|
sumf[row] += d * s;
|
|
}
|
|
dh += args.nb01 / 2;
|
|
q2 += args.nb01 / 2;
|
|
}
|
|
|
|
y4 += 32 * 32;
|
|
}
|
|
}
|
|
|
|
device float * dst_f32 = (device float *)(dst + (uint64_t)token * args.nb1);
|
|
for (int row = 0; row < nr0 && first_row + row < args.ne0; row++) {
|
|
const float sum_all = simd_sum(sumf[row]) * 0.25f;
|
|
if (tiisg == 0) {
|
|
float outv = sum_all;
|
|
if (args.tp_addend) {
|
|
outv += ((device const float *) add_in)[first_row + row];
|
|
}
|
|
dst_f32[first_row + row] = outv;
|
|
}
|
|
}
|
|
|
|
(void)shmem;
|
|
(void)tiitg;
|
|
(void)tgpig;
|
|
}
|
|
|
|
kernel void kernel_mul_mv_id_q2_K_sum6_f32(
|
|
constant ds4_metal_args_mul_mv_id & args,
|
|
device const char * src0s,
|
|
device const char * src1,
|
|
device char * dst,
|
|
device const char * ids,
|
|
device const char * add_in,
|
|
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]]) {
|
|
const short NSG = FC_mul_mv_nsg;
|
|
const short nr0 = N_R0_Q2_K;
|
|
const int nb = args.ne00/QK_K;
|
|
const int first_row = (tgpig.x * NSG + sgitg) * nr0;
|
|
const uint token = tgpig.y;
|
|
device const int32_t *token_ids = (device const int32_t *)(ids + (uint64_t)token * args.nbi1);
|
|
device const char *token_src1 = src1 + (uint64_t)token * args.nb12;
|
|
|
|
float sumf[nr0] = {0.f};
|
|
|
|
const short ix = tiisg/8;
|
|
const short it = tiisg%8;
|
|
const short iq = it/4;
|
|
const short ir = it%4;
|
|
const short is = (8*ir)/16;
|
|
|
|
for (int expert_slot = 0; expert_slot < args.nei0; expert_slot++) {
|
|
const int32_t expert = token_ids[expert_slot];
|
|
if (!ds4_tp_owns_expert(expert, args.ne02, args.tp_rank, args.tp_world)) continue;
|
|
device const block_q2_K * x = (device const block_q2_K *)(src0s + (int64_t)(expert - args.tp_expert_base)*args.nb02 + first_row*args.nb01);
|
|
device const float * y = (device const float *)(token_src1 + expert_slot*args.nb11);
|
|
device const float * y4 = y + ix * QK_K + 128 * iq + 8 * ir;
|
|
|
|
for (int ib = ix; ib < nb; ib += 4) {
|
|
float yl[32];
|
|
float4 sumy = {0.f, 0.f, 0.f, 0.f};
|
|
for (short i = 0; i < 8; ++i) {
|
|
yl[i+ 0] = y4[i+ 0]; sumy[0] += yl[i+ 0];
|
|
yl[i+ 8] = y4[i+32]; sumy[1] += yl[i+ 8];
|
|
yl[i+16] = y4[i+64]; sumy[2] += yl[i+16];
|
|
yl[i+24] = y4[i+96]; sumy[3] += yl[i+24];
|
|
}
|
|
|
|
device const uint8_t * sc = (device const uint8_t *)x[ib].scales + 8*iq + is;
|
|
device const uint16_t * qs = (device const uint16_t *)x[ib].qs + 16 * iq + 4 * ir;
|
|
device const half * dh = &x[ib].d;
|
|
|
|
for (short row = 0; row < nr0; row++) {
|
|
if (first_row + row < args.ne0) {
|
|
float4 acc1 = {0.f, 0.f, 0.f, 0.f};
|
|
float4 acc2 = {0.f, 0.f, 0.f, 0.f};
|
|
for (int i = 0; i < 8; i += 2) {
|
|
acc1[0] += yl[i+ 0] * (qs[i/2] & 0x0003);
|
|
acc2[0] += yl[i+ 1] * (qs[i/2] & 0x0300);
|
|
acc1[1] += yl[i+ 8] * (qs[i/2] & 0x000c);
|
|
acc2[1] += yl[i+ 9] * (qs[i/2] & 0x0c00);
|
|
acc1[2] += yl[i+16] * (qs[i/2] & 0x0030);
|
|
acc2[2] += yl[i+17] * (qs[i/2] & 0x3000);
|
|
acc1[3] += yl[i+24] * (qs[i/2] & 0x00c0);
|
|
acc2[3] += yl[i+25] * (qs[i/2] & 0xc000);
|
|
}
|
|
float dall = dh[0];
|
|
float dmin = dh[1] * 1.f/16.f;
|
|
sumf[row] += dall * ((acc1[0] + 1.f/256.f * acc2[0]) * (sc[0] & 0xF) * 1.f/ 1.f +
|
|
(acc1[1] + 1.f/256.f * acc2[1]) * (sc[2] & 0xF) * 1.f/ 4.f +
|
|
(acc1[2] + 1.f/256.f * acc2[2]) * (sc[4] & 0xF) * 1.f/16.f +
|
|
(acc1[3] + 1.f/256.f * acc2[3]) * (sc[6] & 0xF) * 1.f/64.f) -
|
|
dmin * (sumy[0] * (sc[0] & 0xF0) + sumy[1] * (sc[2] & 0xF0) +
|
|
sumy[2] * (sc[4] & 0xF0) + sumy[3] * (sc[6] & 0xF0));
|
|
}
|
|
|
|
qs += args.nb01/2;
|
|
sc += args.nb01;
|
|
dh += args.nb01/2;
|
|
}
|
|
|
|
y4 += 4 * QK_K;
|
|
}
|
|
}
|
|
|
|
device float * dst_f32 = (device float *)(dst + (uint64_t)token * args.nb1);
|
|
for (int row = 0; row < nr0 && first_row + row < args.ne0; row++) {
|
|
const float sum_all = simd_sum(sumf[row]);
|
|
if (tiisg == 0) {
|
|
float outv = sum_all;
|
|
if (args.tp_addend) {
|
|
outv += ((device const float *) add_in)[first_row + row];
|
|
}
|
|
dst_f32[first_row + row] = outv;
|
|
}
|
|
}
|
|
|
|
(void)shmem;
|
|
(void)tiitg;
|
|
(void)tgpig;
|
|
}
|
|
|
|
kernel void kernel_mul_mv_slots6_q2_K_sum6_f32(
|
|
constant ds4_metal_args_mul_mv_id & args,
|
|
device const char * src00,
|
|
device const char * src01,
|
|
device const char * src02,
|
|
device const char * src03,
|
|
device const char * src04,
|
|
device const char * src05,
|
|
device const char * src1,
|
|
device char * dst,
|
|
threadgroup char * shmem [[threadgroup(0)]],
|
|
uint3 tgpig[[threadgroup_position_in_grid]],
|
|
ushort tiitg[[thread_index_in_threadgroup]],
|
|
ushort tiisg[[thread_index_in_simdgroup]],
|
|
ushort sgitg[[simdgroup_index_in_threadgroup]]) {
|
|
const short NSG = FC_mul_mv_nsg;
|
|
const short nr0 = N_R0_Q2_K;
|
|
const int nb = args.ne00 / QK_K;
|
|
const int first_row = (tgpig.x * NSG + sgitg) * nr0;
|
|
const uint token = tgpig.y;
|
|
device const char *token_src1 = src1 + (uint64_t)token * args.nb12;
|
|
|
|
float sumf[nr0] = {0.f};
|
|
|
|
const short ix = tiisg / 8;
|
|
const short it = tiisg % 8;
|
|
const short iq = it / 4;
|
|
const short ir = it % 4;
|
|
const short is = (8 * ir) / 16;
|
|
|
|
for (int expert_slot = 0; expert_slot < 6; expert_slot++) {
|
|
device const char *src0_cur = src00;
|
|
switch (expert_slot) {
|
|
case 1: src0_cur = src01; break;
|
|
case 2: src0_cur = src02; break;
|
|
case 3: src0_cur = src03; break;
|
|
case 4: src0_cur = src04; break;
|
|
case 5: src0_cur = src05; break;
|
|
default: break;
|
|
}
|
|
device const block_q2_K *x =
|
|
(device const block_q2_K *)(src0_cur + first_row * args.nb01);
|
|
device const float *y = (device const float *)(token_src1 + expert_slot * args.nb11);
|
|
device const float *y4 = y + ix * QK_K + 128 * iq + 8 * ir;
|
|
|
|
for (int ib = ix; ib < nb; ib += 4) {
|
|
float yl[32];
|
|
float4 sumy = {0.f, 0.f, 0.f, 0.f};
|
|
for (short i = 0; i < 8; ++i) {
|
|
yl[i + 0] = y4[i + 0]; sumy[0] += yl[i + 0];
|
|
yl[i + 8] = y4[i + 32]; sumy[1] += yl[i + 8];
|
|
yl[i + 16] = y4[i + 64]; sumy[2] += yl[i + 16];
|
|
yl[i + 24] = y4[i + 96]; sumy[3] += yl[i + 24];
|
|
}
|
|
|
|
device const uint8_t *sc = (device const uint8_t *)x[ib].scales + 8 * iq + is;
|
|
device const uint16_t *qs = (device const uint16_t *)x[ib].qs + 16 * iq + 4 * ir;
|
|
device const half *dh = &x[ib].d;
|
|
|
|
for (short row = 0; row < nr0; row++) {
|
|
if (first_row + row < args.ne0) {
|
|
float4 acc1 = {0.f, 0.f, 0.f, 0.f};
|
|
float4 acc2 = {0.f, 0.f, 0.f, 0.f};
|
|
for (int i = 0; i < 8; i += 2) {
|
|
acc1[0] += yl[i + 0] * (qs[i / 2] & 0x0003);
|
|
acc2[0] += yl[i + 1] * (qs[i / 2] & 0x0300);
|
|
acc1[1] += yl[i + 8] * (qs[i / 2] & 0x000c);
|
|
acc2[1] += yl[i + 9] * (qs[i / 2] & 0x0c00);
|
|
acc1[2] += yl[i + 16] * (qs[i / 2] & 0x0030);
|
|
acc2[2] += yl[i + 17] * (qs[i / 2] & 0x3000);
|
|
acc1[3] += yl[i + 24] * (qs[i / 2] & 0x00c0);
|
|
acc2[3] += yl[i + 25] * (qs[i / 2] & 0xc000);
|
|
}
|
|
float dall = dh[0];
|
|
float dmin = dh[1] * 1.f / 16.f;
|
|
sumf[row] += dall * ((acc1[0] + 1.f / 256.f * acc2[0]) * (sc[0] & 0xF) * 1.f / 1.f +
|
|
(acc1[1] + 1.f / 256.f * acc2[1]) * (sc[2] & 0xF) * 1.f / 4.f +
|
|
(acc1[2] + 1.f / 256.f * acc2[2]) * (sc[4] & 0xF) * 1.f / 16.f +
|
|
(acc1[3] + 1.f / 256.f * acc2[3]) * (sc[6] & 0xF) * 1.f / 64.f) -
|
|
dmin * (sumy[0] * (sc[0] & 0xF0) + sumy[1] * (sc[2] & 0xF0) +
|
|
sumy[2] * (sc[4] & 0xF0) + sumy[3] * (sc[6] & 0xF0));
|
|
}
|
|
|
|
qs += args.nb01 / 2;
|
|
sc += args.nb01;
|
|
dh += args.nb01 / 2;
|
|
}
|
|
|
|
y4 += 4 * QK_K;
|
|
}
|
|
}
|
|
|
|
device float *dst_f32 = (device float *)(dst + (uint64_t)token * args.nb1);
|
|
for (int row = 0; row < nr0 && first_row + row < args.ne0; row++) {
|
|
const float sum_all = simd_sum(sumf[row]);
|
|
if (tiisg == 0) dst_f32[first_row + row] = sum_all;
|
|
}
|
|
|
|
(void)shmem;
|
|
(void)tiitg;
|
|
(void)tgpig;
|
|
}
|
|
|
|
kernel void kernel_mul_mv_addr_q2_K_sum6_f32(
|
|
constant ds4_metal_args_mul_mv_id & args,
|
|
device const uint64_t * addrs,
|
|
device const char * src1,
|
|
device char * dst,
|
|
device const char * ids,
|
|
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]]) {
|
|
const short NSG = FC_mul_mv_nsg;
|
|
const short nr0 = N_R0_Q2_K;
|
|
const int nb = args.ne00 / QK_K;
|
|
const int first_row = (tgpig.x * NSG + sgitg) * nr0;
|
|
const uint token = tgpig.y;
|
|
device const char *token_src1 = src1 + (uint64_t)token * args.nb12;
|
|
device const int32_t *token_ids =
|
|
(device const int32_t *)(ids + (uint64_t)token * args.nbi1);
|
|
|
|
float sumf[nr0] = {0.f};
|
|
|
|
const short ix = tiisg / 8;
|
|
const short it = tiisg % 8;
|
|
const short iq = it / 4;
|
|
const short ir = it % 4;
|
|
const short is = (8 * ir) / 16;
|
|
|
|
for (int expert_slot = 0; expert_slot < 6; expert_slot++) {
|
|
const int32_t expert = token_ids[expert_slot];
|
|
if (expert < 0 || expert >= args.ne02 || expert >= 384) {
|
|
continue;
|
|
}
|
|
const uint64_t addr = addrs[(uint)expert];
|
|
if (addr == 0) {
|
|
continue;
|
|
}
|
|
device const char *src0_cur =
|
|
reinterpret_cast<device const char *>(addr);
|
|
device const block_q2_K *x =
|
|
(device const block_q2_K *)(src0_cur + first_row * args.nb01);
|
|
device const float *y = (device const float *)(token_src1 + expert_slot * args.nb11);
|
|
device const float *y4 = y + ix * QK_K + 128 * iq + 8 * ir;
|
|
|
|
for (int ib = ix; ib < nb; ib += 4) {
|
|
float yl[32];
|
|
float4 sumy = {0.f, 0.f, 0.f, 0.f};
|
|
for (short i = 0; i < 8; ++i) {
|
|
yl[i + 0] = y4[i + 0]; sumy[0] += yl[i + 0];
|
|
yl[i + 8] = y4[i + 32]; sumy[1] += yl[i + 8];
|
|
yl[i + 16] = y4[i + 64]; sumy[2] += yl[i + 16];
|
|
yl[i + 24] = y4[i + 96]; sumy[3] += yl[i + 24];
|
|
}
|
|
|
|
device const uint8_t *sc = (device const uint8_t *)x[ib].scales + 8 * iq + is;
|
|
device const uint16_t *qs = (device const uint16_t *)x[ib].qs + 16 * iq + 4 * ir;
|
|
device const half *dh = &x[ib].d;
|
|
|
|
for (short row = 0; row < nr0; row++) {
|
|
if (first_row + row < args.ne0) {
|
|
float4 acc1 = {0.f, 0.f, 0.f, 0.f};
|
|
float4 acc2 = {0.f, 0.f, 0.f, 0.f};
|
|
for (int i = 0; i < 8; i += 2) {
|
|
acc1[0] += yl[i + 0] * (qs[i / 2] & 0x0003);
|
|
acc2[0] += yl[i + 1] * (qs[i / 2] & 0x0300);
|
|
acc1[1] += yl[i + 8] * (qs[i / 2] & 0x000c);
|
|
acc2[1] += yl[i + 9] * (qs[i / 2] & 0x0c00);
|
|
acc1[2] += yl[i + 16] * (qs[i / 2] & 0x0030);
|
|
acc2[2] += yl[i + 17] * (qs[i / 2] & 0x3000);
|
|
acc1[3] += yl[i + 24] * (qs[i / 2] & 0x00c0);
|
|
acc2[3] += yl[i + 25] * (qs[i / 2] & 0xc000);
|
|
}
|
|
float dall = dh[0];
|
|
float dmin = dh[1] * 1.f / 16.f;
|
|
sumf[row] += dall * ((acc1[0] + 1.f / 256.f * acc2[0]) * (sc[0] & 0xF) * 1.f / 1.f +
|
|
(acc1[1] + 1.f / 256.f * acc2[1]) * (sc[2] & 0xF) * 1.f / 4.f +
|
|
(acc1[2] + 1.f / 256.f * acc2[2]) * (sc[4] & 0xF) * 1.f / 16.f +
|
|
(acc1[3] + 1.f / 256.f * acc2[3]) * (sc[6] & 0xF) * 1.f / 64.f) -
|
|
dmin * (sumy[0] * (sc[0] & 0xF0) + sumy[1] * (sc[2] & 0xF0) +
|
|
sumy[2] * (sc[4] & 0xF0) + sumy[3] * (sc[6] & 0xF0));
|
|
}
|
|
qs += args.nb01 / 2;
|
|
sc += args.nb01;
|
|
dh += args.nb01 / 2;
|
|
}
|
|
y4 += 4 * QK_K;
|
|
}
|
|
}
|
|
|
|
device float * dst_f32 = (device float *)(dst + (uint64_t)token * args.nb1);
|
|
for (int row = 0; row < nr0 && first_row + row < args.ne0; row++) {
|
|
const float sum_all = simd_sum(sumf[row]);
|
|
if (tiisg == 0) dst_f32[first_row + row] = sum_all;
|
|
}
|
|
|
|
(void)shmem;
|
|
(void)tiitg;
|
|
(void)tgpig;
|
|
}
|
|
|
|
kernel void kernel_mul_mv_addr_q2_K_sum6_masked_f32(
|
|
constant ds4_metal_args_mul_mv_id & args,
|
|
constant ds4_metal_stream_expert_split_args & split,
|
|
device const uint64_t * addrs,
|
|
device const char * src1,
|
|
device char * dst,
|
|
device const char * ids,
|
|
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]]) {
|
|
const short NSG = FC_mul_mv_nsg;
|
|
const short nr0 = N_R0_Q2_K;
|
|
const int nb = args.ne00 / QK_K;
|
|
const int first_row = (tgpig.x * NSG + sgitg) * nr0;
|
|
const uint token = tgpig.y;
|
|
device const char *token_src1 = src1 + (uint64_t)token * args.nb12;
|
|
device const int32_t *token_ids =
|
|
(device const int32_t *)(ids + (uint64_t)token * args.nbi1);
|
|
|
|
float sumf[nr0] = {0.f};
|
|
|
|
const short ix = tiisg / 8;
|
|
const short it = tiisg % 8;
|
|
const short iq = it / 4;
|
|
const short ir = it % 4;
|
|
const short is = (8 * ir) / 16;
|
|
|
|
for (int expert_slot = 0; expert_slot < 6; expert_slot++) {
|
|
if ((split.active_mask & (1u << (uint)expert_slot)) == 0) {
|
|
continue;
|
|
}
|
|
const int32_t expert = token_ids[expert_slot];
|
|
if (expert < 0 || expert >= args.ne02 || expert >= 384) {
|
|
continue;
|
|
}
|
|
const uint64_t addr = addrs[(uint)expert];
|
|
if (addr == 0) {
|
|
continue;
|
|
}
|
|
device const char *src0_cur =
|
|
reinterpret_cast<device const char *>(addr);
|
|
device const block_q2_K *x =
|
|
(device const block_q2_K *)(src0_cur + first_row * args.nb01);
|
|
device const float *y = (device const float *)(token_src1 + expert_slot * args.nb11);
|
|
device const float *y4 = y + ix * QK_K + 128 * iq + 8 * ir;
|
|
|
|
for (int ib = ix; ib < nb; ib += 4) {
|
|
float yl[32];
|
|
float4 sumy = {0.f, 0.f, 0.f, 0.f};
|
|
for (short i = 0; i < 8; ++i) {
|
|
yl[i + 0] = y4[i + 0]; sumy[0] += yl[i + 0];
|
|
yl[i + 8] = y4[i + 32]; sumy[1] += yl[i + 8];
|
|
yl[i + 16] = y4[i + 64]; sumy[2] += yl[i + 16];
|
|
yl[i + 24] = y4[i + 96]; sumy[3] += yl[i + 24];
|
|
}
|
|
|
|
device const uint8_t *sc = (device const uint8_t *)x[ib].scales + 8 * iq + is;
|
|
device const uint16_t *qs = (device const uint16_t *)x[ib].qs + 16 * iq + 4 * ir;
|
|
device const half *dh = &x[ib].d;
|
|
|
|
for (short row = 0; row < nr0; row++) {
|
|
if (first_row + row < args.ne0) {
|
|
float4 acc1 = {0.f, 0.f, 0.f, 0.f};
|
|
float4 acc2 = {0.f, 0.f, 0.f, 0.f};
|
|
for (int i = 0; i < 8; i += 2) {
|
|
acc1[0] += yl[i + 0] * (qs[i / 2] & 0x0003);
|
|
acc2[0] += yl[i + 1] * (qs[i / 2] & 0x0300);
|
|
acc1[1] += yl[i + 8] * (qs[i / 2] & 0x000c);
|
|
acc2[1] += yl[i + 9] * (qs[i / 2] & 0x0c00);
|
|
acc1[2] += yl[i + 16] * (qs[i / 2] & 0x0030);
|
|
acc2[2] += yl[i + 17] * (qs[i / 2] & 0x3000);
|
|
acc1[3] += yl[i + 24] * (qs[i / 2] & 0x00c0);
|
|
acc2[3] += yl[i + 25] * (qs[i / 2] & 0xc000);
|
|
}
|
|
float dall = dh[0];
|
|
float dmin = dh[1] * 1.f / 16.f;
|
|
sumf[row] += dall * ((acc1[0] + 1.f / 256.f * acc2[0]) * (sc[0] & 0xF) * 1.f / 1.f +
|
|
(acc1[1] + 1.f / 256.f * acc2[1]) * (sc[2] & 0xF) * 1.f / 4.f +
|
|
(acc1[2] + 1.f / 256.f * acc2[2]) * (sc[4] & 0xF) * 1.f / 16.f +
|
|
(acc1[3] + 1.f / 256.f * acc2[3]) * (sc[6] & 0xF) * 1.f / 64.f) -
|
|
dmin * (sumy[0] * (sc[0] & 0xF0) + sumy[1] * (sc[2] & 0xF0) +
|
|
sumy[2] * (sc[4] & 0xF0) + sumy[3] * (sc[6] & 0xF0));
|
|
}
|
|
qs += args.nb01 / 2;
|
|
sc += args.nb01;
|
|
dh += args.nb01 / 2;
|
|
}
|
|
y4 += 4 * QK_K;
|
|
}
|
|
}
|
|
|
|
device float * dst_f32 = (device float *)(dst + (uint64_t)token * args.nb1);
|
|
for (int row = 0; row < nr0 && first_row + row < args.ne0; row++) {
|
|
const float sum_all = simd_sum(sumf[row]);
|
|
if (tiisg == 0) {
|
|
if (split.accumulate) {
|
|
dst_f32[first_row + row] += sum_all;
|
|
} else {
|
|
dst_f32[first_row + row] = sum_all;
|
|
}
|
|
}
|
|
}
|
|
|
|
(void)shmem;
|
|
(void)tiitg;
|
|
(void)tgpig;
|
|
}
|
|
|
|
static inline float2 ds4_mxfp4_accumulate_rows(
|
|
device const char *src0,
|
|
uint64_t row_bytes,
|
|
device const float *y,
|
|
uint32_t n_cols,
|
|
uint32_t first_row,
|
|
uint32_t n_rows,
|
|
threadgroup const float *lut,
|
|
ushort tiisg) {
|
|
const int nb = (int)(n_cols / QK_MXFP4);
|
|
const int row_blocks = (int)(row_bytes / sizeof(block_mxfp4));
|
|
const short ix = tiisg / 2;
|
|
const short it = tiisg & 1;
|
|
device const block_mxfp4 *x =
|
|
(device const block_mxfp4 *)(src0 + (uint64_t)first_row * row_bytes);
|
|
device const float *yb = y + ix * QK_MXFP4 + it * 8;
|
|
float2 sums = 0.0f;
|
|
|
|
for (int ib = ix; ib < nb; ib += 16) {
|
|
device const float4 *y4 = (device const float4 *)yb;
|
|
const float4 yl0 = y4[0];
|
|
const float4 yl1 = y4[4];
|
|
const float4 yl2 = y4[1];
|
|
const float4 yl3 = y4[5];
|
|
FOR_UNROLL (short row = 0; row < N_R0_MXFP4; row++) {
|
|
if (first_row + row < n_rows) {
|
|
device const block_mxfp4 &b = x[row * row_blocks + ib];
|
|
device const uchar *q = b.qs + 8 * it;
|
|
float4 acc = yl0 * float4(lut[q[0] & 15], lut[q[1] & 15],
|
|
lut[q[2] & 15], lut[q[3] & 15]);
|
|
acc += yl1 * float4(lut[q[0] >> 4], lut[q[1] >> 4],
|
|
lut[q[2] >> 4], lut[q[3] >> 4]);
|
|
acc += yl2 * float4(lut[q[4] & 15], lut[q[5] & 15],
|
|
lut[q[6] & 15], lut[q[7] & 15]);
|
|
acc += yl3 * float4(lut[q[4] >> 4], lut[q[5] >> 4],
|
|
lut[q[6] >> 4], lut[q[7] >> 4]);
|
|
sums[row] += ds4_metal_e8m0_to_f32(b.e) *
|
|
((acc.x + acc.y) + (acc.z + acc.w));
|
|
}
|
|
}
|
|
yb += 16 * QK_MXFP4;
|
|
}
|
|
return sums;
|
|
}
|
|
|
|
// Exact-shape sibling for the resident Flash decode down projection. Its
|
|
// NR2/nsg1 grid covers all 4096 rows without a tail, so retain the established
|
|
// per-lane K walk and arithmetic while omitting only the redundant row guard.
|
|
static inline float2 ds4_mxfp4_accumulate_full_rows(
|
|
device const char *src0,
|
|
uint64_t row_bytes,
|
|
device const float *y,
|
|
uint32_t n_cols,
|
|
uint32_t first_row,
|
|
threadgroup const float *lut,
|
|
ushort tiisg) {
|
|
const int nb = (int)(n_cols / QK_MXFP4);
|
|
const int row_blocks = (int)(row_bytes / sizeof(block_mxfp4));
|
|
const short ix = tiisg / 2;
|
|
const short it = tiisg & 1;
|
|
device const block_mxfp4 *x =
|
|
(device const block_mxfp4 *)(src0 + (uint64_t)first_row * row_bytes);
|
|
device const float *yb = y + ix * QK_MXFP4 + it * 8;
|
|
float2 sums = 0.0f;
|
|
|
|
for (int ib = ix; ib < nb; ib += 16) {
|
|
device const float4 *y4 = (device const float4 *)yb;
|
|
const float4 yl0 = y4[0];
|
|
const float4 yl1 = y4[4];
|
|
const float4 yl2 = y4[1];
|
|
const float4 yl3 = y4[5];
|
|
FOR_UNROLL (short row = 0; row < N_R0_MXFP4; row++) {
|
|
device const block_mxfp4 &b = x[row * row_blocks + ib];
|
|
device const uchar *q = b.qs + 8 * it;
|
|
float4 acc = yl0 * float4(lut[q[0] & 15], lut[q[1] & 15],
|
|
lut[q[2] & 15], lut[q[3] & 15]);
|
|
acc += yl1 * float4(lut[q[0] >> 4], lut[q[1] >> 4],
|
|
lut[q[2] >> 4], lut[q[3] >> 4]);
|
|
acc += yl2 * float4(lut[q[4] & 15], lut[q[5] & 15],
|
|
lut[q[6] & 15], lut[q[7] & 15]);
|
|
acc += yl3 * float4(lut[q[4] >> 4], lut[q[5] >> 4],
|
|
lut[q[6] >> 4], lut[q[7] >> 4]);
|
|
sums[row] += ds4_metal_e8m0_to_f32(b.e) *
|
|
((acc.x + acc.y) + (acc.z + acc.w));
|
|
}
|
|
yb += 16 * QK_MXFP4;
|
|
}
|
|
return sums;
|
|
}
|
|
|
|
kernel void kernel_mul_mv_id_mxfp4_sum6_f32(
|
|
constant ds4_metal_args_mul_mv_id &args,
|
|
device const char *src0s,
|
|
device const char *src1,
|
|
device char *dst,
|
|
device const char *ids,
|
|
device const char *add_in,
|
|
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]]) {
|
|
const short NSG = FC_mul_mv_nsg;
|
|
const uint32_t first_row = (uint32_t)((tgpig.x * NSG + sgitg) * N_R0_MXFP4);
|
|
const uint32_t token = tgpig.y;
|
|
device const int32_t *token_ids =
|
|
(device const int32_t *)(ids + (uint64_t)token * args.nbi1);
|
|
device const char *token_src1 = src1 + (uint64_t)token * args.nb12;
|
|
threadgroup float *lut = (threadgroup float *)shmem;
|
|
if (sgitg == 0) lut[tiisg] = ds4_metal_mxfp4_values[tiisg & 15];
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
float2 sumf = 0.0f;
|
|
for (int slot = 0; slot < args.nei0; slot++) {
|
|
const int32_t expert = token_ids[slot];
|
|
if (!ds4_tp_owns_expert(expert, args.ne02, args.tp_rank, args.tp_world)) continue;
|
|
device const char *expert_base = src0s +
|
|
(int64_t)(expert - args.tp_expert_base) * args.nb02;
|
|
device const float *y =
|
|
(device const float *)(token_src1 + (uint64_t)slot * args.nb11);
|
|
sumf += ds4_mxfp4_accumulate_rows(expert_base, args.nb01, y,
|
|
args.ne00, first_row, args.ne0,
|
|
lut, tiisg);
|
|
}
|
|
|
|
device float *out = (device float *)(dst + (uint64_t)token * args.nb1);
|
|
FOR_UNROLL (short row = 0; row < N_R0_MXFP4; row++) {
|
|
if (first_row + row < (uint32_t)args.ne0) {
|
|
const float value = simd_sum(sumf[row]);
|
|
if (tiisg == 0) {
|
|
out[first_row + row] = value +
|
|
(args.tp_addend ? ((device const float *)add_in)[first_row + row] : 0.0f);
|
|
}
|
|
}
|
|
}
|
|
(void)tiitg;
|
|
}
|
|
|
|
kernel void kernel_mul_mv_id_mxfp4_sum6_fixed_route_f32(
|
|
constant ds4_metal_args_mul_mv_id &args,
|
|
device const char *src0s,
|
|
device const char *src1,
|
|
device char *dst,
|
|
device const char *ids,
|
|
device const char *add_in,
|
|
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]]) {
|
|
const short NSG = FC_mul_mv_nsg;
|
|
const uint32_t first_row = (uint32_t)((tgpig.x * NSG + sgitg) * N_R0_MXFP4);
|
|
device const int32_t *token_ids = (device const int32_t *)ids;
|
|
device const char *token_src1 = src1;
|
|
threadgroup float *lut = (threadgroup float *)shmem;
|
|
if (sgitg == 0) lut[tiisg] = ds4_metal_mxfp4_values[tiisg & 15];
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
float2 sumf = 0.0f;
|
|
for (int slot = 0; slot < args.nei0; slot++) {
|
|
const int32_t expert = token_ids[slot];
|
|
device const char *expert_base =
|
|
src0s + (int64_t)expert * args.nb02;
|
|
device const float *y =
|
|
(device const float *)(token_src1 + (uint64_t)slot * args.nb11);
|
|
sumf += ds4_mxfp4_accumulate_rows(expert_base, args.nb01, y,
|
|
args.ne00, first_row, args.ne0,
|
|
lut, tiisg);
|
|
}
|
|
|
|
device float *out = (device float *)dst;
|
|
FOR_UNROLL (short row = 0; row < N_R0_MXFP4; row++) {
|
|
if (first_row + row < (uint32_t)args.ne0) {
|
|
const float value = simd_sum(sumf[row]);
|
|
if (tiisg == 0) {
|
|
out[first_row + row] = value +
|
|
(args.tp_addend ? ((device const float *)add_in)[first_row + row] : 0.0f);
|
|
}
|
|
}
|
|
}
|
|
(void)tiitg;
|
|
}
|
|
|
|
kernel void kernel_mul_mv_id_mxfp4_sum6_fixed_route_full_rows_f32(
|
|
constant ds4_metal_args_mul_mv_id &args,
|
|
device const char *src0s,
|
|
device const char *src1,
|
|
device char *dst,
|
|
device const char *ids,
|
|
device const char *add_in,
|
|
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]]) {
|
|
const short NSG = FC_mul_mv_nsg;
|
|
const uint32_t first_row = (uint32_t)((tgpig.x * NSG + sgitg) * N_R0_MXFP4);
|
|
device const int32_t *token_ids = (device const int32_t *)ids;
|
|
device const char *token_src1 = src1;
|
|
threadgroup float *lut = (threadgroup float *)shmem;
|
|
if (sgitg == 0) lut[tiisg] = ds4_metal_mxfp4_values[tiisg & 15];
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
float2 sumf = 0.0f;
|
|
for (int slot = 0; slot < args.nei0; slot++) {
|
|
const int32_t expert = token_ids[slot];
|
|
device const char *expert_base =
|
|
src0s + (int64_t)expert * args.nb02;
|
|
device const float *y =
|
|
(device const float *)(token_src1 + (uint64_t)slot * args.nb11);
|
|
sumf += ds4_mxfp4_accumulate_full_rows(
|
|
expert_base, args.nb01, y, args.ne00, first_row, lut, tiisg);
|
|
}
|
|
|
|
device float *out = (device float *)dst;
|
|
FOR_UNROLL (short row = 0; row < N_R0_MXFP4; row++) {
|
|
const float value = simd_sum(sumf[row]);
|
|
if (tiisg == 0) {
|
|
out[first_row + row] = value +
|
|
(args.tp_addend ? ((device const float *)add_in)[first_row + row] : 0.0f);
|
|
}
|
|
}
|
|
(void)tiitg;
|
|
}
|
|
|
|
/* Exact-shape sibling of the fixed-route full-rows decode down projection.
|
|
* The host proves ne00 == 2048, row bytes == 1088, nei0 == 6 and ne0 == 4096
|
|
* before selecting it, so the K walk (64 blocks / 16 lane pairs = 4 steps),
|
|
* the row stride and the routed-slot count are all literals here. The lane
|
|
* mapping, the block visit order, the per-block accumulate order and the
|
|
* simd_sum tree are byte-identical to ds4_mxfp4_accumulate_full_rows. */
|
|
#define DS4_MXFP4_DOWN_STATIC_NB 64
|
|
#define DS4_MXFP4_DOWN_STATIC_ROW_BLOCKS 64
|
|
#define DS4_MXFP4_DOWN_STATIC_SLOTS 6
|
|
|
|
static inline float2 ds4_mxfp4_accumulate_full_rows_static(
|
|
device const char *src0,
|
|
device const float *y,
|
|
uint32_t first_row,
|
|
threadgroup const float *lut,
|
|
ushort tiisg) {
|
|
const short ix = tiisg / 2;
|
|
const short it = tiisg & 1;
|
|
device const block_mxfp4 *x =
|
|
(device const block_mxfp4 *)(src0 +
|
|
(uint64_t)first_row * (DS4_MXFP4_DOWN_STATIC_ROW_BLOCKS *
|
|
sizeof(block_mxfp4)));
|
|
device const float *yb = y + ix * QK_MXFP4 + it * 8;
|
|
float2 sums = 0.0f;
|
|
|
|
for (int ib = ix; ib < DS4_MXFP4_DOWN_STATIC_NB; ib += 16) {
|
|
device const float4 *y4 = (device const float4 *)yb;
|
|
const float4 yl0 = y4[0];
|
|
const float4 yl1 = y4[4];
|
|
const float4 yl2 = y4[1];
|
|
const float4 yl3 = y4[5];
|
|
FOR_UNROLL (short row = 0; row < N_R0_MXFP4; row++) {
|
|
device const block_mxfp4 &b =
|
|
x[row * DS4_MXFP4_DOWN_STATIC_ROW_BLOCKS + ib];
|
|
device const uchar *q = b.qs + 8 * it;
|
|
float4 acc = yl0 * float4(lut[q[0] & 15], lut[q[1] & 15],
|
|
lut[q[2] & 15], lut[q[3] & 15]);
|
|
acc += yl1 * float4(lut[q[0] >> 4], lut[q[1] >> 4],
|
|
lut[q[2] >> 4], lut[q[3] >> 4]);
|
|
acc += yl2 * float4(lut[q[4] & 15], lut[q[5] & 15],
|
|
lut[q[6] & 15], lut[q[7] & 15]);
|
|
acc += yl3 * float4(lut[q[4] >> 4], lut[q[5] >> 4],
|
|
lut[q[6] >> 4], lut[q[7] >> 4]);
|
|
sums[row] += ds4_metal_e8m0_to_f32(b.e) *
|
|
((acc.x + acc.y) + (acc.z + acc.w));
|
|
}
|
|
yb += 16 * QK_MXFP4;
|
|
}
|
|
return sums;
|
|
}
|
|
|
|
kernel void kernel_mul_mv_id_mxfp4_sum6_fixed_route_full_rows_static_f32(
|
|
constant ds4_metal_args_mul_mv_id &args,
|
|
device const char *src0s,
|
|
device const char *src1,
|
|
device char *dst,
|
|
device const char *ids,
|
|
device const char *add_in,
|
|
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]]) {
|
|
const short NSG = FC_mul_mv_nsg;
|
|
const uint32_t first_row = (uint32_t)((tgpig.x * NSG + sgitg) * N_R0_MXFP4);
|
|
device const int32_t *token_ids = (device const int32_t *)ids;
|
|
device const char *token_src1 = src1;
|
|
threadgroup float *lut = (threadgroup float *)shmem;
|
|
if (sgitg == 0) lut[tiisg] = ds4_metal_mxfp4_values[tiisg & 15];
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
/* Defensive: unreachable in production because the host proves the shape
|
|
* before selecting this pipeline; identical results either way. */
|
|
const bool static_shape =
|
|
args.ne00 == DS4_MXFP4_DOWN_STATIC_NB * QK_MXFP4 &&
|
|
args.nb01 == (uint64_t)DS4_MXFP4_DOWN_STATIC_ROW_BLOCKS *
|
|
sizeof(block_mxfp4) &&
|
|
args.nei0 == DS4_MXFP4_DOWN_STATIC_SLOTS;
|
|
|
|
float2 sumf = 0.0f;
|
|
if (static_shape) {
|
|
for (short slot = 0; slot < DS4_MXFP4_DOWN_STATIC_SLOTS; slot++) {
|
|
const int32_t expert = token_ids[slot];
|
|
device const char *expert_base =
|
|
src0s + (int64_t)expert * args.nb02;
|
|
device const float *y =
|
|
(device const float *)(token_src1 + (uint64_t)slot * args.nb11);
|
|
sumf += ds4_mxfp4_accumulate_full_rows_static(
|
|
expert_base, y, first_row, lut, tiisg);
|
|
}
|
|
} else {
|
|
for (int slot = 0; slot < args.nei0; slot++) {
|
|
const int32_t expert = token_ids[slot];
|
|
device const char *expert_base =
|
|
src0s + (int64_t)expert * args.nb02;
|
|
device const float *y =
|
|
(device const float *)(token_src1 + (uint64_t)slot * args.nb11);
|
|
sumf += ds4_mxfp4_accumulate_full_rows(
|
|
expert_base, args.nb01, y, args.ne00, first_row, lut, tiisg);
|
|
}
|
|
}
|
|
|
|
device float *out = (device float *)dst;
|
|
FOR_UNROLL (short row = 0; row < N_R0_MXFP4; row++) {
|
|
const float value = simd_sum(sumf[row]);
|
|
if (tiisg == 0) {
|
|
out[first_row + row] = value +
|
|
(args.tp_addend ? ((device const float *)add_in)[first_row + row] : 0.0f);
|
|
}
|
|
}
|
|
(void)tiitg;
|
|
}
|
|
|
|
kernel void kernel_mul_mv_slots6_mxfp4_sum6_f32(
|
|
constant ds4_metal_args_mul_mv_id &args,
|
|
device const char *src00, device const char *src01,
|
|
device const char *src02, device const char *src03,
|
|
device const char *src04, device const char *src05,
|
|
device const char *src1,
|
|
device char *dst,
|
|
threadgroup char *shmem [[threadgroup(0)]],
|
|
uint3 tgpig [[threadgroup_position_in_grid]],
|
|
ushort tiitg [[thread_index_in_threadgroup]],
|
|
ushort tiisg [[thread_index_in_simdgroup]],
|
|
ushort sgitg [[simdgroup_index_in_threadgroup]]) {
|
|
const short NSG = FC_mul_mv_nsg;
|
|
const uint32_t first_row = (uint32_t)((tgpig.x * NSG + sgitg) * N_R0_MXFP4);
|
|
const uint32_t token = tgpig.y;
|
|
device const char *token_src1 = src1 + (uint64_t)token * args.nb12;
|
|
threadgroup float *lut = (threadgroup float *)shmem;
|
|
if (sgitg == 0) lut[tiisg] = ds4_metal_mxfp4_values[tiisg & 15];
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
float2 sumf = 0.0f;
|
|
for (int slot = 0; slot < 6; slot++) {
|
|
device const char *expert_base = src00;
|
|
switch (slot) {
|
|
case 1: expert_base = src01; break;
|
|
case 2: expert_base = src02; break;
|
|
case 3: expert_base = src03; break;
|
|
case 4: expert_base = src04; break;
|
|
case 5: expert_base = src05; break;
|
|
default: break;
|
|
}
|
|
device const float *y =
|
|
(device const float *)(token_src1 + (uint64_t)slot * args.nb11);
|
|
sumf += ds4_mxfp4_accumulate_rows(expert_base, args.nb01, y,
|
|
args.ne00, first_row, args.ne0,
|
|
lut, tiisg);
|
|
}
|
|
|
|
device float *out = (device float *)(dst + (uint64_t)token * args.nb1);
|
|
FOR_UNROLL (short row = 0; row < N_R0_MXFP4; row++) {
|
|
if (first_row + row < (uint32_t)args.ne0) {
|
|
const float value = simd_sum(sumf[row]);
|
|
if (tiisg == 0) out[first_row + row] = value;
|
|
}
|
|
}
|
|
(void)tiitg;
|
|
}
|
|
|
|
kernel void kernel_mul_mv_id_q4_K_sum6_f32(
|
|
constant ds4_metal_args_mul_mv_id & args,
|
|
device const char * src0s,
|
|
device const char * src1,
|
|
device char * dst,
|
|
device const char * ids,
|
|
device const char * add_in,
|
|
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]]) {
|
|
const short NSG = FC_mul_mv_nsg;
|
|
const short nr0 = N_R0_Q4_K;
|
|
const int nb = args.ne00 / QK_K;
|
|
const int first_row = (tgpig.x * NSG + sgitg) * nr0;
|
|
const uint token = tgpig.y;
|
|
device const int32_t *token_ids = (device const int32_t *)(ids + (uint64_t)token * args.nbi1);
|
|
device const char *token_src1 = src1 + (uint64_t)token * args.nb12;
|
|
|
|
constexpr uint16_t kmask1 = 0x3f3f;
|
|
constexpr uint16_t kmask2 = 0x0f0f;
|
|
constexpr uint16_t kmask3 = 0xc0c0;
|
|
|
|
const short ix = tiisg / 8;
|
|
const short it = tiisg % 8;
|
|
const short iq = it / 4;
|
|
const short ir = it % 4;
|
|
|
|
float sumf[nr0] = {0.f};
|
|
uint16_t sc16[4];
|
|
thread const uint8_t *sc8 = (thread const uint8_t *)sc16;
|
|
|
|
for (int expert_slot = 0; expert_slot < args.nei0; expert_slot++) {
|
|
const int32_t expert = token_ids[expert_slot];
|
|
if (!ds4_tp_owns_expert(expert, args.ne02, args.tp_rank, args.tp_world)) continue;
|
|
device const block_q4_K *x =
|
|
(device const block_q4_K *)(src0s + (int64_t)(expert - args.tp_expert_base) * args.nb02 + first_row * args.nb01);
|
|
device const float *y = (device const float *)(token_src1 + expert_slot * args.nb11);
|
|
device const float *y4 = y + ix * QK_K + 64 * iq + 8 * ir;
|
|
|
|
for (int ib = ix; ib < nb; ib += 4) {
|
|
float yl[16];
|
|
float yh[16];
|
|
float4 sumy = {0.f, 0.f, 0.f, 0.f};
|
|
|
|
for (short i = 0; i < 8; ++i) {
|
|
yl[i + 0] = y4[i + 0]; sumy[0] += yl[i + 0];
|
|
yl[i + 8] = y4[i + 32]; sumy[1] += yl[i + 8];
|
|
yh[i + 0] = y4[i + 128]; sumy[2] += yh[i + 0];
|
|
yh[i + 8] = y4[i + 160]; sumy[3] += yh[i + 8];
|
|
}
|
|
|
|
device const uint16_t *sc = (device const uint16_t *)x[ib].scales + iq;
|
|
device const uint16_t *q1 = (device const uint16_t *)x[ib].qs + 16 * iq + 4 * ir;
|
|
device const half *dh = &x[ib].d;
|
|
|
|
for (short row = 0; row < nr0; row++) {
|
|
if (first_row + row < args.ne0) {
|
|
sc16[0] = sc[0] & kmask1;
|
|
sc16[1] = sc[2] & kmask1;
|
|
sc16[2] = ((sc[4] >> 0) & kmask2) | ((sc[0] & kmask3) >> 2);
|
|
sc16[3] = ((sc[4] >> 4) & kmask2) | ((sc[2] & kmask3) >> 2);
|
|
|
|
device const uint16_t *q2 = q1 + 32;
|
|
|
|
float4 acc1 = {0.f, 0.f, 0.f, 0.f};
|
|
float4 acc2 = {0.f, 0.f, 0.f, 0.f};
|
|
|
|
FOR_UNROLL (short i = 0; i < 4; ++i) {
|
|
acc1[0] += yl[2 * i + 0] * (q1[i] & 0x000F);
|
|
acc1[1] += yl[2 * i + 1] * (q1[i] & 0x0F00);
|
|
acc1[2] += yl[2 * i + 8] * (q1[i] & 0x00F0);
|
|
acc1[3] += yl[2 * i + 9] * (q1[i] & 0xF000);
|
|
acc2[0] += yh[2 * i + 0] * (q2[i] & 0x000F);
|
|
acc2[1] += yh[2 * i + 1] * (q2[i] & 0x0F00);
|
|
acc2[2] += yh[2 * i + 8] * (q2[i] & 0x00F0);
|
|
acc2[3] += yh[2 * i + 9] * (q2[i] & 0xF000);
|
|
}
|
|
|
|
sumf[row] += dh[0] * ((acc1[0] + 1.f / 256.f * acc1[1]) * sc8[0] +
|
|
(acc1[2] + 1.f / 256.f * acc1[3]) * sc8[1] * 1.f / 16.f +
|
|
(acc2[0] + 1.f / 256.f * acc2[1]) * sc8[4] +
|
|
(acc2[2] + 1.f / 256.f * acc2[3]) * sc8[5] * 1.f / 16.f) -
|
|
dh[1] * (sumy[0] * sc8[2] + sumy[1] * sc8[3] +
|
|
sumy[2] * sc8[6] + sumy[3] * sc8[7]);
|
|
}
|
|
|
|
q1 += args.nb01 / 2;
|
|
sc += args.nb01 / 2;
|
|
dh += args.nb01 / 2;
|
|
}
|
|
|
|
y4 += 4 * QK_K;
|
|
}
|
|
}
|
|
|
|
device float *dst_f32 = (device float *)(dst + (uint64_t)token * args.nb1);
|
|
for (int row = 0; row < nr0 && first_row + row < args.ne0; row++) {
|
|
const float sum_all = simd_sum(sumf[row]);
|
|
if (tiisg == 0) {
|
|
float outv = sum_all;
|
|
if (args.tp_addend) {
|
|
outv += ((device const float *) add_in)[first_row + row];
|
|
}
|
|
dst_f32[first_row + row] = outv;
|
|
}
|
|
}
|
|
|
|
(void)shmem;
|
|
(void)tiitg;
|
|
(void)tgpig;
|
|
}
|
|
|
|
kernel void kernel_mul_mv_group_q4_K_sum6_f32(
|
|
constant ds4_metal_args_mul_mv_id & args,
|
|
constant ds4_metal_moe_expert_group_args & group,
|
|
device const char * src0s,
|
|
device const char * src1,
|
|
device char * dst,
|
|
device const char * ids,
|
|
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]]) {
|
|
const short NSG = FC_mul_mv_nsg;
|
|
const short nr0 = N_R0_Q4_K;
|
|
const int nb = args.ne00 / QK_K;
|
|
const int first_row = (tgpig.x * NSG + sgitg) * nr0;
|
|
const uint token = tgpig.y;
|
|
device const int32_t *token_ids = (device const int32_t *)(ids + (uint64_t)token * args.nbi1);
|
|
device const char *token_src1 = src1 + (uint64_t)token * args.nb12;
|
|
|
|
constexpr uint16_t kmask1 = 0x3f3f;
|
|
constexpr uint16_t kmask2 = 0x0f0f;
|
|
constexpr uint16_t kmask3 = 0xc0c0;
|
|
|
|
const short ix = tiisg / 8;
|
|
const short it = tiisg % 8;
|
|
const short iq = it / 4;
|
|
const short ir = it % 4;
|
|
|
|
float sumf[nr0] = {0.f};
|
|
uint16_t sc16[4];
|
|
thread const uint8_t *sc8 = (thread const uint8_t *)sc16;
|
|
|
|
for (int expert_slot = 0; expert_slot < 6; expert_slot++) {
|
|
const int32_t expert = token_ids[expert_slot];
|
|
if (expert < 0) {
|
|
continue;
|
|
}
|
|
const uint32_t expert_u = (uint32_t)expert;
|
|
if (expert_u < group.expert_base ||
|
|
expert_u >= group.expert_base + group.expert_count) {
|
|
continue;
|
|
}
|
|
const uint32_t expert_local = expert_u - group.expert_base;
|
|
|
|
device const block_q4_K *x =
|
|
(device const block_q4_K *)(src0s + (uint64_t)expert_local * args.nb02 + first_row * args.nb01);
|
|
device const float *y = (device const float *)(token_src1 + expert_slot * args.nb11);
|
|
device const float *y4 = y + ix * QK_K + 64 * iq + 8 * ir;
|
|
|
|
for (int ib = ix; ib < nb; ib += 4) {
|
|
float yl[16];
|
|
float yh[16];
|
|
float4 sumy = {0.f, 0.f, 0.f, 0.f};
|
|
|
|
for (short i = 0; i < 8; ++i) {
|
|
yl[i + 0] = y4[i + 0]; sumy[0] += yl[i + 0];
|
|
yl[i + 8] = y4[i + 32]; sumy[1] += yl[i + 8];
|
|
yh[i + 0] = y4[i + 128]; sumy[2] += yh[i + 0];
|
|
yh[i + 8] = y4[i + 160]; sumy[3] += yh[i + 8];
|
|
}
|
|
|
|
device const uint16_t *sc = (device const uint16_t *)x[ib].scales + iq;
|
|
device const uint16_t *q1 = (device const uint16_t *)x[ib].qs + 16 * iq + 4 * ir;
|
|
device const half *dh = &x[ib].d;
|
|
|
|
for (short row = 0; row < nr0; row++) {
|
|
if (first_row + row < args.ne0) {
|
|
sc16[0] = sc[0] & kmask1;
|
|
sc16[1] = sc[2] & kmask1;
|
|
sc16[2] = ((sc[4] >> 0) & kmask2) | ((sc[0] & kmask3) >> 2);
|
|
sc16[3] = ((sc[4] >> 4) & kmask2) | ((sc[2] & kmask3) >> 2);
|
|
|
|
device const uint16_t *q2 = q1 + 32;
|
|
|
|
float4 acc1 = {0.f, 0.f, 0.f, 0.f};
|
|
float4 acc2 = {0.f, 0.f, 0.f, 0.f};
|
|
|
|
FOR_UNROLL (short i = 0; i < 4; ++i) {
|
|
acc1[0] += yl[2 * i + 0] * (q1[i] & 0x000F);
|
|
acc1[1] += yl[2 * i + 1] * (q1[i] & 0x0F00);
|
|
acc1[2] += yl[2 * i + 8] * (q1[i] & 0x00F0);
|
|
acc1[3] += yl[2 * i + 9] * (q1[i] & 0xF000);
|
|
acc2[0] += yh[2 * i + 0] * (q2[i] & 0x000F);
|
|
acc2[1] += yh[2 * i + 1] * (q2[i] & 0x0F00);
|
|
acc2[2] += yh[2 * i + 8] * (q2[i] & 0x00F0);
|
|
acc2[3] += yh[2 * i + 9] * (q2[i] & 0xF000);
|
|
}
|
|
|
|
sumf[row] += dh[0] * ((acc1[0] + 1.f / 256.f * acc1[1]) * sc8[0] +
|
|
(acc1[2] + 1.f / 256.f * acc1[3]) * sc8[1] * 1.f / 16.f +
|
|
(acc2[0] + 1.f / 256.f * acc2[1]) * sc8[4] +
|
|
(acc2[2] + 1.f / 256.f * acc2[3]) * sc8[5] * 1.f / 16.f) -
|
|
dh[1] * (sumy[0] * sc8[2] + sumy[1] * sc8[3] +
|
|
sumy[2] * sc8[6] + sumy[3] * sc8[7]);
|
|
}
|
|
|
|
q1 += args.nb01 / 2;
|
|
sc += args.nb01 / 2;
|
|
dh += args.nb01 / 2;
|
|
}
|
|
|
|
y4 += 4 * QK_K;
|
|
}
|
|
}
|
|
|
|
device float *dst_f32 = (device float *)(dst + (uint64_t)token * args.nb1);
|
|
for (int row = 0; row < nr0 && first_row + row < args.ne0; row++) {
|
|
const float sum_all = simd_sum(sumf[row]);
|
|
if (tiisg == 0) {
|
|
if (group.accumulate) {
|
|
dst_f32[first_row + row] += sum_all;
|
|
} else {
|
|
dst_f32[first_row + row] = sum_all;
|
|
}
|
|
}
|
|
}
|
|
|
|
(void)shmem;
|
|
(void)tiitg;
|
|
(void)tgpig;
|
|
}
|
|
|
|
kernel void kernel_mul_mv_table_q4_K_sum6_f32(
|
|
constant ds4_metal_args_mul_mv_id & args,
|
|
device const ds4_metal_q4_expert_table & table,
|
|
device const char * src1,
|
|
device char * dst,
|
|
device const char * ids,
|
|
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]]) {
|
|
const short NSG = FC_mul_mv_nsg;
|
|
const short nr0 = N_R0_Q4_K;
|
|
const int nb = args.ne00 / QK_K;
|
|
const int first_row = (tgpig.x * NSG + sgitg) * nr0;
|
|
const uint token = tgpig.y;
|
|
device const int32_t *token_ids = (device const int32_t *)(ids + (uint64_t)token * args.nbi1);
|
|
device const char *token_src1 = src1 + (uint64_t)token * args.nb12;
|
|
|
|
constexpr uint16_t kmask1 = 0x3f3f;
|
|
constexpr uint16_t kmask2 = 0x0f0f;
|
|
constexpr uint16_t kmask3 = 0xc0c0;
|
|
|
|
const short ix = tiisg / 8;
|
|
const short it = tiisg % 8;
|
|
const short iq = it / 4;
|
|
const short ir = it % 4;
|
|
|
|
float sumf[nr0] = {0.f};
|
|
uint16_t sc16[4];
|
|
thread const uint8_t *sc8 = (thread const uint8_t *)sc16;
|
|
|
|
for (int expert_slot = 0; expert_slot < 6; expert_slot++) {
|
|
const int32_t expert = token_ids[expert_slot];
|
|
if (expert < 0 || expert >= args.ne02 || expert >= 384) {
|
|
return;
|
|
}
|
|
device const block_q4_K *x =
|
|
(device const block_q4_K *)(table.experts[(uint)expert] + first_row * args.nb01);
|
|
device const float *y = (device const float *)(token_src1 + expert_slot * args.nb11);
|
|
device const float *y4 = y + ix * QK_K + 64 * iq + 8 * ir;
|
|
|
|
for (int ib = ix; ib < nb; ib += 4) {
|
|
float yl[16];
|
|
float yh[16];
|
|
float4 sumy = {0.f, 0.f, 0.f, 0.f};
|
|
|
|
for (short i = 0; i < 8; ++i) {
|
|
yl[i + 0] = y4[i + 0]; sumy[0] += yl[i + 0];
|
|
yl[i + 8] = y4[i + 32]; sumy[1] += yl[i + 8];
|
|
yh[i + 0] = y4[i + 128]; sumy[2] += yh[i + 0];
|
|
yh[i + 8] = y4[i + 160]; sumy[3] += yh[i + 8];
|
|
}
|
|
|
|
device const uint16_t *sc = (device const uint16_t *)x[ib].scales + iq;
|
|
device const uint16_t *q1 = (device const uint16_t *)x[ib].qs + 16 * iq + 4 * ir;
|
|
device const half *dh = &x[ib].d;
|
|
|
|
for (short row = 0; row < nr0; row++) {
|
|
if (first_row + row < args.ne0) {
|
|
sc16[0] = sc[0] & kmask1;
|
|
sc16[1] = sc[2] & kmask1;
|
|
sc16[2] = ((sc[4] >> 0) & kmask2) | ((sc[0] & kmask3) >> 2);
|
|
sc16[3] = ((sc[4] >> 4) & kmask2) | ((sc[2] & kmask3) >> 2);
|
|
|
|
device const uint16_t *q2 = q1 + 32;
|
|
|
|
float4 acc1 = {0.f, 0.f, 0.f, 0.f};
|
|
float4 acc2 = {0.f, 0.f, 0.f, 0.f};
|
|
|
|
FOR_UNROLL (short i = 0; i < 4; ++i) {
|
|
acc1[0] += yl[2 * i + 0] * (q1[i] & 0x000F);
|
|
acc1[1] += yl[2 * i + 1] * (q1[i] & 0x0F00);
|
|
acc1[2] += yl[2 * i + 8] * (q1[i] & 0x00F0);
|
|
acc1[3] += yl[2 * i + 9] * (q1[i] & 0xF000);
|
|
acc2[0] += yh[2 * i + 0] * (q2[i] & 0x000F);
|
|
acc2[1] += yh[2 * i + 1] * (q2[i] & 0x0F00);
|
|
acc2[2] += yh[2 * i + 8] * (q2[i] & 0x00F0);
|
|
acc2[3] += yh[2 * i + 9] * (q2[i] & 0xF000);
|
|
}
|
|
|
|
sumf[row] += dh[0] * ((acc1[0] + 1.f / 256.f * acc1[1]) * sc8[0] +
|
|
(acc1[2] + 1.f / 256.f * acc1[3]) * sc8[1] * 1.f / 16.f +
|
|
(acc2[0] + 1.f / 256.f * acc2[1]) * sc8[4] +
|
|
(acc2[2] + 1.f / 256.f * acc2[3]) * sc8[5] * 1.f / 16.f) -
|
|
dh[1] * (sumy[0] * sc8[2] + sumy[1] * sc8[3] +
|
|
sumy[2] * sc8[6] + sumy[3] * sc8[7]);
|
|
}
|
|
|
|
q1 += args.nb01 / 2;
|
|
sc += args.nb01 / 2;
|
|
dh += args.nb01 / 2;
|
|
}
|
|
|
|
y4 += 4 * QK_K;
|
|
}
|
|
}
|
|
|
|
device float *dst_f32 = (device float *)(dst + (uint64_t)token * args.nb1);
|
|
for (int row = 0; row < nr0 && first_row + row < args.ne0; row++) {
|
|
const float sum_all = simd_sum(sumf[row]);
|
|
if (tiisg == 0) dst_f32[first_row + row] = sum_all;
|
|
}
|
|
|
|
(void)shmem;
|
|
(void)tiitg;
|
|
(void)tgpig;
|
|
}
|
|
|
|
kernel void kernel_mul_mv_addr_q4_K_sum6_f32(
|
|
constant ds4_metal_args_mul_mv_id & args,
|
|
device const ulong * addrs,
|
|
device const char * src1,
|
|
device char * dst,
|
|
device const char * ids,
|
|
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]]) {
|
|
const short NSG = FC_mul_mv_nsg;
|
|
const short nr0 = N_R0_Q4_K;
|
|
const int nb = args.ne00 / QK_K;
|
|
const int first_row = (tgpig.x * NSG + sgitg) * nr0;
|
|
const uint token = tgpig.y;
|
|
device const int32_t *token_ids = (device const int32_t *)(ids + (uint64_t)token * args.nbi1);
|
|
device const char *token_src1 = src1 + (uint64_t)token * args.nb12;
|
|
|
|
constexpr uint16_t kmask1 = 0x3f3f;
|
|
constexpr uint16_t kmask2 = 0x0f0f;
|
|
constexpr uint16_t kmask3 = 0xc0c0;
|
|
|
|
const short ix = tiisg / 8;
|
|
const short it = tiisg % 8;
|
|
const short iq = it / 4;
|
|
const short ir = it % 4;
|
|
|
|
float sumf[nr0] = {0.f};
|
|
uint16_t sc16[4];
|
|
thread const uint8_t *sc8 = (thread const uint8_t *)sc16;
|
|
|
|
for (int expert_slot = 0; expert_slot < 6; expert_slot++) {
|
|
const int32_t expert = token_ids[expert_slot];
|
|
if (expert < 0 || expert >= args.ne02 || expert >= 384) {
|
|
return;
|
|
}
|
|
device const char *expert_base =
|
|
reinterpret_cast<device const char *>(addrs[(uint)expert]);
|
|
device const block_q4_K *x =
|
|
(device const block_q4_K *)(expert_base + first_row * args.nb01);
|
|
device const float *y = (device const float *)(token_src1 + expert_slot * args.nb11);
|
|
device const float *y4 = y + ix * QK_K + 64 * iq + 8 * ir;
|
|
|
|
for (int ib = ix; ib < nb; ib += 4) {
|
|
float yl[16];
|
|
float yh[16];
|
|
float4 sumy = {0.f, 0.f, 0.f, 0.f};
|
|
|
|
for (short i = 0; i < 8; ++i) {
|
|
yl[i + 0] = y4[i + 0]; sumy[0] += yl[i + 0];
|
|
yl[i + 8] = y4[i + 32]; sumy[1] += yl[i + 8];
|
|
yh[i + 0] = y4[i + 128]; sumy[2] += yh[i + 0];
|
|
yh[i + 8] = y4[i + 160]; sumy[3] += yh[i + 8];
|
|
}
|
|
|
|
device const uint16_t *sc = (device const uint16_t *)x[ib].scales + iq;
|
|
device const uint16_t *q1 = (device const uint16_t *)x[ib].qs + 16 * iq + 4 * ir;
|
|
device const half *dh = &x[ib].d;
|
|
|
|
for (short row = 0; row < nr0; row++) {
|
|
if (first_row + row < args.ne0) {
|
|
sc16[0] = sc[0] & kmask1;
|
|
sc16[1] = sc[2] & kmask1;
|
|
sc16[2] = ((sc[4] >> 0) & kmask2) | ((sc[0] & kmask3) >> 2);
|
|
sc16[3] = ((sc[4] >> 4) & kmask2) | ((sc[2] & kmask3) >> 2);
|
|
|
|
device const uint16_t *q2 = q1 + 32;
|
|
|
|
float4 acc1 = {0.f, 0.f, 0.f, 0.f};
|
|
float4 acc2 = {0.f, 0.f, 0.f, 0.f};
|
|
|
|
FOR_UNROLL (short i = 0; i < 4; ++i) {
|
|
acc1[0] += yl[2 * i + 0] * (q1[i] & 0x000F);
|
|
acc1[1] += yl[2 * i + 1] * (q1[i] & 0x0F00);
|
|
acc1[2] += yl[2 * i + 8] * (q1[i] & 0x00F0);
|
|
acc1[3] += yl[2 * i + 9] * (q1[i] & 0xF000);
|
|
acc2[0] += yh[2 * i + 0] * (q2[i] & 0x000F);
|
|
acc2[1] += yh[2 * i + 1] * (q2[i] & 0x0F00);
|
|
acc2[2] += yh[2 * i + 8] * (q2[i] & 0x00F0);
|
|
acc2[3] += yh[2 * i + 9] * (q2[i] & 0xF000);
|
|
}
|
|
|
|
sumf[row] += dh[0] * ((acc1[0] + 1.f / 256.f * acc1[1]) * sc8[0] +
|
|
(acc1[2] + 1.f / 256.f * acc1[3]) * sc8[1] * 1.f / 16.f +
|
|
(acc2[0] + 1.f / 256.f * acc2[1]) * sc8[4] +
|
|
(acc2[2] + 1.f / 256.f * acc2[3]) * sc8[5] * 1.f / 16.f) -
|
|
dh[1] * (sumy[0] * sc8[2] + sumy[1] * sc8[3] +
|
|
sumy[2] * sc8[6] + sumy[3] * sc8[7]);
|
|
}
|
|
|
|
q1 += args.nb01 / 2;
|
|
sc += args.nb01 / 2;
|
|
dh += args.nb01 / 2;
|
|
}
|
|
|
|
y4 += 4 * QK_K;
|
|
}
|
|
}
|
|
|
|
device float *dst_f32 = (device float *)(dst + (uint64_t)token * args.nb1);
|
|
for (int row = 0; row < nr0 && first_row + row < args.ne0; row++) {
|
|
const float sum_all = simd_sum(sumf[row]);
|
|
if (tiisg == 0) dst_f32[first_row + row] = sum_all;
|
|
}
|
|
|
|
(void)shmem;
|
|
(void)tiitg;
|
|
(void)tgpig;
|
|
}
|
|
|
|
kernel void kernel_mul_mv_slots6_q4_K_sum6_f32(
|
|
constant ds4_metal_args_mul_mv_id & args,
|
|
device const char * src00,
|
|
device const char * src01,
|
|
device const char * src02,
|
|
device const char * src03,
|
|
device const char * src04,
|
|
device const char * src05,
|
|
device const char * src1,
|
|
device char * dst,
|
|
threadgroup char * shmem [[threadgroup(0)]],
|
|
uint3 tgpig[[threadgroup_position_in_grid]],
|
|
ushort tiitg[[thread_index_in_threadgroup]],
|
|
ushort tiisg[[thread_index_in_simdgroup]],
|
|
ushort sgitg[[simdgroup_index_in_threadgroup]]) {
|
|
const short NSG = FC_mul_mv_nsg;
|
|
const short nr0 = N_R0_Q4_K;
|
|
const int nb = args.ne00 / QK_K;
|
|
const int first_row = (tgpig.x * NSG + sgitg) * nr0;
|
|
const uint token = tgpig.y;
|
|
device const char *token_src1 = src1 + (uint64_t)token * args.nb12;
|
|
|
|
constexpr uint16_t kmask1 = 0x3f3f;
|
|
constexpr uint16_t kmask2 = 0x0f0f;
|
|
constexpr uint16_t kmask3 = 0xc0c0;
|
|
|
|
const short ix = tiisg / 8;
|
|
const short it = tiisg % 8;
|
|
const short iq = it / 4;
|
|
const short ir = it % 4;
|
|
|
|
float sumf[nr0] = {0.f};
|
|
uint16_t sc16[4];
|
|
thread const uint8_t *sc8 = (thread const uint8_t *)sc16;
|
|
|
|
for (int expert_slot = 0; expert_slot < 6; expert_slot++) {
|
|
device const char *src0_cur = src00;
|
|
switch (expert_slot) {
|
|
case 1: src0_cur = src01; break;
|
|
case 2: src0_cur = src02; break;
|
|
case 3: src0_cur = src03; break;
|
|
case 4: src0_cur = src04; break;
|
|
case 5: src0_cur = src05; break;
|
|
default: break;
|
|
}
|
|
device const block_q4_K *x =
|
|
(device const block_q4_K *)(src0_cur + first_row * args.nb01);
|
|
device const float *y = (device const float *)(token_src1 + expert_slot * args.nb11);
|
|
device const float *y4 = y + ix * QK_K + 64 * iq + 8 * ir;
|
|
|
|
for (int ib = ix; ib < nb; ib += 4) {
|
|
float yl[16];
|
|
float yh[16];
|
|
float4 sumy = {0.f, 0.f, 0.f, 0.f};
|
|
|
|
for (short i = 0; i < 8; ++i) {
|
|
yl[i + 0] = y4[i + 0]; sumy[0] += yl[i + 0];
|
|
yl[i + 8] = y4[i + 32]; sumy[1] += yl[i + 8];
|
|
yh[i + 0] = y4[i + 128]; sumy[2] += yh[i + 0];
|
|
yh[i + 8] = y4[i + 160]; sumy[3] += yh[i + 8];
|
|
}
|
|
|
|
device const uint16_t *sc = (device const uint16_t *)x[ib].scales + iq;
|
|
device const uint16_t *q1 = (device const uint16_t *)x[ib].qs + 16 * iq + 4 * ir;
|
|
device const half *dh = &x[ib].d;
|
|
|
|
for (short row = 0; row < nr0; row++) {
|
|
if (first_row + row < args.ne0) {
|
|
sc16[0] = sc[0] & kmask1;
|
|
sc16[1] = sc[2] & kmask1;
|
|
sc16[2] = ((sc[4] >> 0) & kmask2) | ((sc[0] & kmask3) >> 2);
|
|
sc16[3] = ((sc[4] >> 4) & kmask2) | ((sc[2] & kmask3) >> 2);
|
|
|
|
device const uint16_t *q2 = q1 + 32;
|
|
|
|
float4 acc1 = {0.f, 0.f, 0.f, 0.f};
|
|
float4 acc2 = {0.f, 0.f, 0.f, 0.f};
|
|
|
|
FOR_UNROLL (short i = 0; i < 4; ++i) {
|
|
acc1[0] += yl[2 * i + 0] * (q1[i] & 0x000F);
|
|
acc1[1] += yl[2 * i + 1] * (q1[i] & 0x0F00);
|
|
acc1[2] += yl[2 * i + 8] * (q1[i] & 0x00F0);
|
|
acc1[3] += yl[2 * i + 9] * (q1[i] & 0xF000);
|
|
acc2[0] += yh[2 * i + 0] * (q2[i] & 0x000F);
|
|
acc2[1] += yh[2 * i + 1] * (q2[i] & 0x0F00);
|
|
acc2[2] += yh[2 * i + 8] * (q2[i] & 0x00F0);
|
|
acc2[3] += yh[2 * i + 9] * (q2[i] & 0xF000);
|
|
}
|
|
|
|
sumf[row] += dh[0] * ((acc1[0] + 1.f / 256.f * acc1[1]) * sc8[0] +
|
|
(acc1[2] + 1.f / 256.f * acc1[3]) * sc8[1] * 1.f / 16.f +
|
|
(acc2[0] + 1.f / 256.f * acc2[1]) * sc8[4] +
|
|
(acc2[2] + 1.f / 256.f * acc2[3]) * sc8[5] * 1.f / 16.f) -
|
|
dh[1] * (sumy[0] * sc8[2] + sumy[1] * sc8[3] +
|
|
sumy[2] * sc8[6] + sumy[3] * sc8[7]);
|
|
}
|
|
|
|
q1 += args.nb01 / 2;
|
|
sc += args.nb01 / 2;
|
|
dh += args.nb01 / 2;
|
|
}
|
|
|
|
y4 += 4 * QK_K;
|
|
}
|
|
}
|
|
|
|
device float *dst_f32 = (device float *)(dst + (uint64_t)token * args.nb1);
|
|
for (int row = 0; row < nr0 && first_row + row < args.ne0; row++) {
|
|
const float sum_all = simd_sum(sumf[row]);
|
|
if (tiisg == 0) dst_f32[first_row + row] = sum_all;
|
|
}
|
|
|
|
(void)shmem;
|
|
(void)tiitg;
|
|
(void)tgpig;
|
|
}
|
|
|
|
kernel void kernel_mul_mv_group6_q4_K_sum6_f32(
|
|
constant ds4_metal_args_mul_mv_id & args,
|
|
device const char * src00,
|
|
device const char * src01,
|
|
device const char * src02,
|
|
device const char * src03,
|
|
device const char * src04,
|
|
device const char * src05,
|
|
device const char * src1,
|
|
device char * dst,
|
|
device const char * ids,
|
|
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 uint32_t expert_group_size = 64;
|
|
const short NSG = FC_mul_mv_nsg;
|
|
const short nr0 = N_R0_Q4_K;
|
|
const int nb = args.ne00 / QK_K;
|
|
const int first_row = (tgpig.x * NSG + sgitg) * nr0;
|
|
const uint token = tgpig.y;
|
|
device const int32_t *token_ids = (device const int32_t *)(ids + (uint64_t)token * args.nbi1);
|
|
device const char *token_src1 = src1 + (uint64_t)token * args.nb12;
|
|
|
|
constexpr uint16_t kmask1 = 0x3f3f;
|
|
constexpr uint16_t kmask2 = 0x0f0f;
|
|
constexpr uint16_t kmask3 = 0xc0c0;
|
|
|
|
const short ix = tiisg / 8;
|
|
const short it = tiisg % 8;
|
|
const short iq = it / 4;
|
|
const short ir = it % 4;
|
|
|
|
float sumf[nr0] = {0.f};
|
|
uint16_t sc16[4];
|
|
thread const uint8_t *sc8 = (thread const uint8_t *)sc16;
|
|
|
|
for (int expert_slot = 0; expert_slot < 6; expert_slot++) {
|
|
const int32_t expert = token_ids[expert_slot];
|
|
if (expert < 0) {
|
|
continue;
|
|
}
|
|
const uint32_t expert_u = (uint32_t)expert;
|
|
const uint32_t group_id = expert_u / expert_group_size;
|
|
if (group_id >= 6) {
|
|
continue;
|
|
}
|
|
const uint32_t expert_local = expert_u - group_id * expert_group_size;
|
|
|
|
device const char *src0_cur = src00;
|
|
switch (group_id) {
|
|
case 1: src0_cur = src01; break;
|
|
case 2: src0_cur = src02; break;
|
|
case 3: src0_cur = src03; break;
|
|
case 4: src0_cur = src04; break;
|
|
case 5: src0_cur = src05; break;
|
|
default: break;
|
|
}
|
|
|
|
device const block_q4_K *x =
|
|
(device const block_q4_K *)(src0_cur + (uint64_t)expert_local * args.nb02 + first_row * args.nb01);
|
|
device const float *y = (device const float *)(token_src1 + expert_slot * args.nb11);
|
|
device const float *y4 = y + ix * QK_K + 64 * iq + 8 * ir;
|
|
|
|
for (int ib = ix; ib < nb; ib += 4) {
|
|
float yl[16];
|
|
float yh[16];
|
|
float4 sumy = {0.f, 0.f, 0.f, 0.f};
|
|
|
|
for (short i = 0; i < 8; ++i) {
|
|
yl[i + 0] = y4[i + 0]; sumy[0] += yl[i + 0];
|
|
yl[i + 8] = y4[i + 32]; sumy[1] += yl[i + 8];
|
|
yh[i + 0] = y4[i + 128]; sumy[2] += yh[i + 0];
|
|
yh[i + 8] = y4[i + 160]; sumy[3] += yh[i + 8];
|
|
}
|
|
|
|
device const uint16_t *sc = (device const uint16_t *)x[ib].scales + iq;
|
|
device const uint16_t *q1 = (device const uint16_t *)x[ib].qs + 16 * iq + 4 * ir;
|
|
device const half *dh = &x[ib].d;
|
|
|
|
for (short row = 0; row < nr0; row++) {
|
|
if (first_row + row < args.ne0) {
|
|
sc16[0] = sc[0] & kmask1;
|
|
sc16[1] = sc[2] & kmask1;
|
|
sc16[2] = ((sc[4] >> 0) & kmask2) | ((sc[0] & kmask3) >> 2);
|
|
sc16[3] = ((sc[4] >> 4) & kmask2) | ((sc[2] & kmask3) >> 2);
|
|
|
|
device const uint16_t *q2 = q1 + 32;
|
|
|
|
float4 acc1 = {0.f, 0.f, 0.f, 0.f};
|
|
float4 acc2 = {0.f, 0.f, 0.f, 0.f};
|
|
|
|
FOR_UNROLL (short i = 0; i < 4; ++i) {
|
|
acc1[0] += yl[2 * i + 0] * (q1[i] & 0x000F);
|
|
acc1[1] += yl[2 * i + 1] * (q1[i] & 0x0F00);
|
|
acc1[2] += yl[2 * i + 8] * (q1[i] & 0x00F0);
|
|
acc1[3] += yl[2 * i + 9] * (q1[i] & 0xF000);
|
|
acc2[0] += yh[2 * i + 0] * (q2[i] & 0x000F);
|
|
acc2[1] += yh[2 * i + 1] * (q2[i] & 0x0F00);
|
|
acc2[2] += yh[2 * i + 8] * (q2[i] & 0x00F0);
|
|
acc2[3] += yh[2 * i + 9] * (q2[i] & 0xF000);
|
|
}
|
|
|
|
sumf[row] += dh[0] * ((acc1[0] + 1.f / 256.f * acc1[1]) * sc8[0] +
|
|
(acc1[2] + 1.f / 256.f * acc1[3]) * sc8[1] * 1.f / 16.f +
|
|
(acc2[0] + 1.f / 256.f * acc2[1]) * sc8[4] +
|
|
(acc2[2] + 1.f / 256.f * acc2[3]) * sc8[5] * 1.f / 16.f) -
|
|
dh[1] * (sumy[0] * sc8[2] + sumy[1] * sc8[3] +
|
|
sumy[2] * sc8[6] + sumy[3] * sc8[7]);
|
|
}
|
|
|
|
q1 += args.nb01 / 2;
|
|
sc += args.nb01 / 2;
|
|
dh += args.nb01 / 2;
|
|
}
|
|
|
|
y4 += 4 * QK_K;
|
|
}
|
|
}
|
|
|
|
device float *dst_f32 = (device float *)(dst + (uint64_t)token * args.nb1);
|
|
for (int row = 0; row < nr0 && first_row + row < args.ne0; row++) {
|
|
const float sum_all = simd_sum(sumf[row]);
|
|
if (tiisg == 0) dst_f32[first_row + row] = sum_all;
|
|
}
|
|
|
|
(void)shmem;
|
|
(void)tiitg;
|
|
(void)tgpig;
|
|
}
|
|
|
|
kernel void kernel_mul_mv_group8_q4_K_sum6_f32(
|
|
constant ds4_metal_args_mul_mv_id & args,
|
|
device const char * src00,
|
|
device const char * src01,
|
|
device const char * src02,
|
|
device const char * src03,
|
|
device const char * src04,
|
|
device const char * src05,
|
|
device const char * src06,
|
|
device const char * src07,
|
|
device const char * src1,
|
|
device char * dst,
|
|
device const char * ids,
|
|
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 uint32_t expert_group_size = 48;
|
|
const short NSG = FC_mul_mv_nsg;
|
|
const short nr0 = N_R0_Q4_K;
|
|
const int nb = args.ne00 / QK_K;
|
|
const int first_row = (tgpig.x * NSG + sgitg) * nr0;
|
|
const uint token = tgpig.y;
|
|
device const int32_t *token_ids = (device const int32_t *)(ids + (uint64_t)token * args.nbi1);
|
|
device const char *token_src1 = src1 + (uint64_t)token * args.nb12;
|
|
|
|
constexpr uint16_t kmask1 = 0x3f3f;
|
|
constexpr uint16_t kmask2 = 0x0f0f;
|
|
constexpr uint16_t kmask3 = 0xc0c0;
|
|
|
|
const short ix = tiisg / 8;
|
|
const short it = tiisg % 8;
|
|
const short iq = it / 4;
|
|
const short ir = it % 4;
|
|
|
|
float sumf[nr0] = {0.f};
|
|
uint16_t sc16[4];
|
|
thread const uint8_t *sc8 = (thread const uint8_t *)sc16;
|
|
|
|
for (int expert_slot = 0; expert_slot < 6; expert_slot++) {
|
|
const int32_t expert = token_ids[expert_slot];
|
|
if (expert < 0) {
|
|
continue;
|
|
}
|
|
const uint32_t expert_u = (uint32_t)expert;
|
|
const uint32_t group_id = expert_u / expert_group_size;
|
|
if (group_id >= 8) {
|
|
continue;
|
|
}
|
|
const uint32_t expert_local = expert_u - group_id * expert_group_size;
|
|
|
|
device const char *src0_cur = src00;
|
|
switch (group_id) {
|
|
case 1: src0_cur = src01; break;
|
|
case 2: src0_cur = src02; break;
|
|
case 3: src0_cur = src03; break;
|
|
case 4: src0_cur = src04; break;
|
|
case 5: src0_cur = src05; break;
|
|
case 6: src0_cur = src06; break;
|
|
case 7: src0_cur = src07; break;
|
|
default: break;
|
|
}
|
|
|
|
device const block_q4_K *x =
|
|
(device const block_q4_K *)(src0_cur + (uint64_t)expert_local * args.nb02 + first_row * args.nb01);
|
|
device const float *y = (device const float *)(token_src1 + expert_slot * args.nb11);
|
|
device const float *y4 = y + ix * QK_K + 64 * iq + 8 * ir;
|
|
|
|
for (int ib = ix; ib < nb; ib += 4) {
|
|
float yl[16];
|
|
float yh[16];
|
|
float4 sumy = {0.f, 0.f, 0.f, 0.f};
|
|
|
|
for (short i = 0; i < 8; ++i) {
|
|
yl[i + 0] = y4[i + 0]; sumy[0] += yl[i + 0];
|
|
yl[i + 8] = y4[i + 32]; sumy[1] += yl[i + 8];
|
|
yh[i + 0] = y4[i + 128]; sumy[2] += yh[i + 0];
|
|
yh[i + 8] = y4[i + 160]; sumy[3] += yh[i + 8];
|
|
}
|
|
|
|
device const uint16_t *sc = (device const uint16_t *)x[ib].scales + iq;
|
|
device const uint16_t *q1 = (device const uint16_t *)x[ib].qs + 16 * iq + 4 * ir;
|
|
device const half *dh = &x[ib].d;
|
|
|
|
for (short row = 0; row < nr0; row++) {
|
|
if (first_row + row < args.ne0) {
|
|
sc16[0] = sc[0] & kmask1;
|
|
sc16[1] = sc[2] & kmask1;
|
|
sc16[2] = ((sc[4] >> 0) & kmask2) | ((sc[0] & kmask3) >> 2);
|
|
sc16[3] = ((sc[4] >> 4) & kmask2) | ((sc[2] & kmask3) >> 2);
|
|
|
|
device const uint16_t *q2 = q1 + 32;
|
|
|
|
float4 acc1 = {0.f, 0.f, 0.f, 0.f};
|
|
float4 acc2 = {0.f, 0.f, 0.f, 0.f};
|
|
|
|
FOR_UNROLL (short i = 0; i < 4; ++i) {
|
|
acc1[0] += yl[2 * i + 0] * (q1[i] & 0x000F);
|
|
acc1[1] += yl[2 * i + 1] * (q1[i] & 0x0F00);
|
|
acc1[2] += yl[2 * i + 8] * (q1[i] & 0x00F0);
|
|
acc1[3] += yl[2 * i + 9] * (q1[i] & 0xF000);
|
|
acc2[0] += yh[2 * i + 0] * (q2[i] & 0x000F);
|
|
acc2[1] += yh[2 * i + 1] * (q2[i] & 0x0F00);
|
|
acc2[2] += yh[2 * i + 8] * (q2[i] & 0x00F0);
|
|
acc2[3] += yh[2 * i + 9] * (q2[i] & 0xF000);
|
|
}
|
|
|
|
sumf[row] += dh[0] * ((acc1[0] + 1.f / 256.f * acc1[1]) * sc8[0] +
|
|
(acc1[2] + 1.f / 256.f * acc1[3]) * sc8[1] * 1.f / 16.f +
|
|
(acc2[0] + 1.f / 256.f * acc2[1]) * sc8[4] +
|
|
(acc2[2] + 1.f / 256.f * acc2[3]) * sc8[5] * 1.f / 16.f) -
|
|
dh[1] * (sumy[0] * sc8[2] + sumy[1] * sc8[3] +
|
|
sumy[2] * sc8[6] + sumy[3] * sc8[7]);
|
|
}
|
|
|
|
q1 += args.nb01 / 2;
|
|
sc += args.nb01 / 2;
|
|
dh += args.nb01 / 2;
|
|
}
|
|
|
|
y4 += 4 * QK_K;
|
|
}
|
|
}
|
|
|
|
device float *dst_f32 = (device float *)(dst + (uint64_t)token * args.nb1);
|
|
for (int row = 0; row < nr0 && first_row + row < args.ne0; row++) {
|
|
const float sum_all = simd_sum(sumf[row]);
|
|
if (tiisg == 0) dst_f32[first_row + row] = sum_all;
|
|
}
|
|
|
|
(void)shmem;
|
|
(void)tiitg;
|
|
(void)tgpig;
|
|
}
|
|
|
|
kernel void kernel_mul_mv_group24_q4_K_sum6_f32(
|
|
constant ds4_metal_args_mul_mv_id & args,
|
|
device const char * src00,
|
|
device const char * src01,
|
|
device const char * src02,
|
|
device const char * src03,
|
|
device const char * src04,
|
|
device const char * src05,
|
|
device const char * src06,
|
|
device const char * src07,
|
|
device const char * src08,
|
|
device const char * src09,
|
|
device const char * src10,
|
|
device const char * src11,
|
|
device const char * src12,
|
|
device const char * src13,
|
|
device const char * src14,
|
|
device const char * src15,
|
|
device const char * src16,
|
|
device const char * src17,
|
|
device const char * src18,
|
|
device const char * src19,
|
|
device const char * src20,
|
|
device const char * src21,
|
|
device const char * src22,
|
|
device const char * src23,
|
|
device const char * src1,
|
|
device char * dst,
|
|
device const char * ids,
|
|
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 uint32_t expert_group_size = 16;
|
|
const short NSG = FC_mul_mv_nsg;
|
|
const short nr0 = N_R0_Q4_K;
|
|
const int nb = args.ne00 / QK_K;
|
|
const int first_row = (tgpig.x * NSG + sgitg) * nr0;
|
|
const uint token = tgpig.y;
|
|
device const int32_t *token_ids = (device const int32_t *)(ids + (uint64_t)token * args.nbi1);
|
|
device const char *token_src1 = src1 + (uint64_t)token * args.nb12;
|
|
|
|
constexpr uint16_t kmask1 = 0x3f3f;
|
|
constexpr uint16_t kmask2 = 0x0f0f;
|
|
constexpr uint16_t kmask3 = 0xc0c0;
|
|
|
|
const short ix = tiisg / 8;
|
|
const short it = tiisg % 8;
|
|
const short iq = it / 4;
|
|
const short ir = it % 4;
|
|
|
|
float sumf[nr0] = {0.f};
|
|
uint16_t sc16[4];
|
|
thread const uint8_t *sc8 = (thread const uint8_t *)sc16;
|
|
|
|
for (int expert_slot = 0; expert_slot < 6; expert_slot++) {
|
|
const int32_t expert = token_ids[expert_slot];
|
|
if (expert < 0) {
|
|
continue;
|
|
}
|
|
const uint32_t expert_u = (uint32_t)expert;
|
|
const uint32_t group_id = expert_u / expert_group_size;
|
|
if (group_id >= 24) {
|
|
continue;
|
|
}
|
|
const uint32_t expert_local = expert_u - group_id * expert_group_size;
|
|
|
|
device const char *src0_cur = ds4_q4_group24_select(group_id,
|
|
src00, src01, src02, src03,
|
|
src04, src05, src06, src07,
|
|
src08, src09, src10, src11,
|
|
src12, src13, src14, src15,
|
|
src16, src17, src18, src19,
|
|
src20, src21, src22, src23);
|
|
device const block_q4_K *x =
|
|
(device const block_q4_K *)(src0_cur + (uint64_t)expert_local * args.nb02 + first_row * args.nb01);
|
|
device const float *y = (device const float *)(token_src1 + expert_slot * args.nb11);
|
|
device const float *y4 = y + ix * QK_K + 64 * iq + 8 * ir;
|
|
|
|
for (int ib = ix; ib < nb; ib += 4) {
|
|
float yl[16];
|
|
float yh[16];
|
|
float4 sumy = {0.f, 0.f, 0.f, 0.f};
|
|
|
|
for (short i = 0; i < 8; ++i) {
|
|
yl[i + 0] = y4[i + 0]; sumy[0] += yl[i + 0];
|
|
yl[i + 8] = y4[i + 32]; sumy[1] += yl[i + 8];
|
|
yh[i + 0] = y4[i + 128]; sumy[2] += yh[i + 0];
|
|
yh[i + 8] = y4[i + 160]; sumy[3] += yh[i + 8];
|
|
}
|
|
|
|
device const uint16_t *sc = (device const uint16_t *)x[ib].scales + iq;
|
|
device const uint16_t *q1 = (device const uint16_t *)x[ib].qs + 16 * iq + 4 * ir;
|
|
device const half *dh = &x[ib].d;
|
|
|
|
for (short row = 0; row < nr0; row++) {
|
|
if (first_row + row < args.ne0) {
|
|
sc16[0] = sc[0] & kmask1;
|
|
sc16[1] = sc[2] & kmask1;
|
|
sc16[2] = ((sc[4] >> 0) & kmask2) | ((sc[0] & kmask3) >> 2);
|
|
sc16[3] = ((sc[4] >> 4) & kmask2) | ((sc[2] & kmask3) >> 2);
|
|
|
|
device const uint16_t *q2 = q1 + 32;
|
|
|
|
float4 acc1 = {0.f, 0.f, 0.f, 0.f};
|
|
float4 acc2 = {0.f, 0.f, 0.f, 0.f};
|
|
|
|
FOR_UNROLL (short i = 0; i < 4; ++i) {
|
|
acc1[0] += yl[2 * i + 0] * (q1[i] & 0x000F);
|
|
acc1[1] += yl[2 * i + 1] * (q1[i] & 0x0F00);
|
|
acc1[2] += yl[2 * i + 8] * (q1[i] & 0x00F0);
|
|
acc1[3] += yl[2 * i + 9] * (q1[i] & 0xF000);
|
|
acc2[0] += yh[2 * i + 0] * (q2[i] & 0x000F);
|
|
acc2[1] += yh[2 * i + 1] * (q2[i] & 0x0F00);
|
|
acc2[2] += yh[2 * i + 8] * (q2[i] & 0x00F0);
|
|
acc2[3] += yh[2 * i + 9] * (q2[i] & 0xF000);
|
|
}
|
|
|
|
sumf[row] += dh[0] * ((acc1[0] + 1.f / 256.f * acc1[1]) * sc8[0] +
|
|
(acc1[2] + 1.f / 256.f * acc1[3]) * sc8[1] * 1.f / 16.f +
|
|
(acc2[0] + 1.f / 256.f * acc2[1]) * sc8[4] +
|
|
(acc2[2] + 1.f / 256.f * acc2[3]) * sc8[5] * 1.f / 16.f) -
|
|
dh[1] * (sumy[0] * sc8[2] + sumy[1] * sc8[3] +
|
|
sumy[2] * sc8[6] + sumy[3] * sc8[7]);
|
|
}
|
|
|
|
q1 += args.nb01 / 2;
|
|
sc += args.nb01 / 2;
|
|
dh += args.nb01 / 2;
|
|
}
|
|
|
|
y4 += 4 * QK_K;
|
|
}
|
|
}
|
|
|
|
device float *dst_f32 = (device float *)(dst + (uint64_t)token * args.nb1);
|
|
for (int row = 0; row < nr0 && first_row + row < args.ne0; row++) {
|
|
const float sum_all = simd_sum(sumf[row]);
|
|
if (tiisg == 0) dst_f32[first_row + row] = sum_all;
|
|
}
|
|
|
|
(void)shmem;
|
|
(void)tiitg;
|
|
(void)tgpig;
|
|
}
|
|
|
|
#define QK_NL 16
|
|
|
|
// Builds the compact per-expert work map used by batched MoE matmul. DS4 routes
|
|
// each token to a small fixed top-k list, so this turns token-major ids into
|
|
// expert-major slices that the tiled matmul can consume.
|
|
template<short ne20>
|
|
kernel void kernel_mul_mm_id_map0(
|
|
constant ds4_metal_args_mul_mm_id_map0 & args,
|
|
device const char * src2,
|
|
device char * htpe,
|
|
device char * hids,
|
|
device char * work,
|
|
threadgroup char * shmem [[threadgroup(0)]],
|
|
ushort tpitg[[thread_position_in_threadgroup]],
|
|
ushort ntg[[threads_per_threadgroup]]) {
|
|
const short ide = tpitg;
|
|
|
|
uint32_t n_all = 0;
|
|
|
|
device int32_t * ids_i32 = (device int32_t *) hids + ide*args.ne21;
|
|
|
|
for (int i21 = 0; i21 < args.ne21; i21 += ntg) {
|
|
if (i21 + tpitg < args.ne21) {
|
|
device const int32_t * src2_i32 = (device const int32_t *) (src2 + (i21 + tpitg)*args.nb21);
|
|
|
|
threadgroup uint16_t * sids = (threadgroup uint16_t *) shmem + tpitg*ne20;
|
|
|
|
#pragma unroll(ne20)
|
|
for (short i20 = 0; i20 < ne20; i20++) {
|
|
sids[i20] = src2_i32[i20];
|
|
}
|
|
}
|
|
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
for (short t = 0; t < ntg; t++) {
|
|
if (i21 + t >= args.ne21) {
|
|
break;
|
|
}
|
|
|
|
threadgroup const uint16_t * sids = (threadgroup const uint16_t *) shmem + t*ne20;
|
|
|
|
short sel = 0;
|
|
#pragma unroll(ne20)
|
|
for (short i20 = 0; i20 < ne20; i20++) {
|
|
sel += (sids[i20] == ide)*(i20 + 1);
|
|
}
|
|
|
|
ids_i32[n_all] = (i21 + t)*ne20 + sel - 1;
|
|
|
|
n_all += sel > 0;
|
|
}
|
|
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
}
|
|
|
|
device uint32_t * tpe_u32 = (device uint32_t *) (htpe);
|
|
tpe_u32[ide] = n_all;
|
|
|
|
// Reuse the route-id staging memory after the map is complete to build a
|
|
// compact list of non-empty 32-row matmul tiles. The old dispatch covered
|
|
// every possible token tile for every expert, even though most experts
|
|
// receive only a small fraction of the prompt rows.
|
|
threadgroup uint16_t * tile_counts = (threadgroup uint16_t *) shmem;
|
|
const uint16_t n_tiles = (uint16_t)((n_all + 31u) / 32u);
|
|
tile_counts[ide] = n_tiles;
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
uint32_t tile_base = 0;
|
|
for (ushort i = 0; i < ide; i++) {
|
|
tile_base += tile_counts[i];
|
|
}
|
|
|
|
device uint32_t * work_count = (device uint32_t *) work;
|
|
device uint2 * work_items = (device uint2 *)(work + 8);
|
|
for (uint32_t tile = 0; tile < n_tiles; tile++) {
|
|
work_items[tile_base + tile] = uint2((uint32_t)ide, tile * 32u);
|
|
}
|
|
if (ide + 1u == ntg) {
|
|
work_count[0] = tile_base + n_tiles;
|
|
}
|
|
}
|
|
|
|
typedef decltype(kernel_mul_mm_id_map0<1>) kernel_mul_mm_id_map0_t;
|
|
|
|
// Host-visible map builders for the routed-expert counts used by DS4 graph
|
|
// shapes. Some arities are generic leftovers retained for nearby batch sizes.
|
|
template [[host_name("kernel_mul_mm_id_map0_ne20_1" )]] kernel kernel_mul_mm_id_map0_t kernel_mul_mm_id_map0<1>;
|
|
template [[host_name("kernel_mul_mm_id_map0_ne20_2" )]] kernel kernel_mul_mm_id_map0_t kernel_mul_mm_id_map0<2>;
|
|
template [[host_name("kernel_mul_mm_id_map0_ne20_4" )]] kernel kernel_mul_mm_id_map0_t kernel_mul_mm_id_map0<4>;
|
|
template [[host_name("kernel_mul_mm_id_map0_ne20_5" )]] kernel kernel_mul_mm_id_map0_t kernel_mul_mm_id_map0<5>;
|
|
template [[host_name("kernel_mul_mm_id_map0_ne20_6" )]] kernel kernel_mul_mm_id_map0_t kernel_mul_mm_id_map0<6>;
|
|
template [[host_name("kernel_mul_mm_id_map0_ne20_8" )]] kernel kernel_mul_mm_id_map0_t kernel_mul_mm_id_map0<8>;
|
|
template [[host_name("kernel_mul_mm_id_map0_ne20_10")]] kernel kernel_mul_mm_id_map0_t kernel_mul_mm_id_map0<10>;
|
|
template [[host_name("kernel_mul_mm_id_map0_ne20_16")]] kernel kernel_mul_mm_id_map0_t kernel_mul_mm_id_map0<16>;
|
|
template [[host_name("kernel_mul_mm_id_map0_ne20_22")]] kernel kernel_mul_mm_id_map0_t kernel_mul_mm_id_map0<22>;
|
|
|
|
// Token-centric map builder for the common six-route MXFP4 prefill shape.
|
|
// Each selected route reserves one expert-local row instead of making every
|
|
// expert scan every token. Expert-local order is intentionally unspecified:
|
|
// hids carries the original token/slot id, so every downstream row still
|
|
// writes its distinct fixed destination. The work trailer exactly matches the
|
|
// established direct-dispatch ABI used by the mapped pair and down kernels.
|
|
template<short ne20>
|
|
kernel void kernel_mul_mm_id_map_scatter_work(
|
|
constant ds4_metal_args_mul_mm_id_map0 & args,
|
|
device const char * src2,
|
|
device char * htpe,
|
|
device char * hids,
|
|
device char * work,
|
|
threadgroup char * shmem [[threadgroup(0)]],
|
|
ushort tpitg[[thread_position_in_threadgroup]],
|
|
ushort ntg[[threads_per_threadgroup]]) {
|
|
threadgroup atomic_uint * counts =
|
|
(threadgroup atomic_uint *) shmem;
|
|
threadgroup uint16_t * tile_counts =
|
|
(threadgroup uint16_t *)(shmem +
|
|
(uint32_t)args.ne02*sizeof(uint32_t));
|
|
device uint32_t * tpe_u32 = (device uint32_t *) htpe;
|
|
device int32_t * ids_i32 = (device int32_t *) hids;
|
|
|
|
if (tpitg < args.ne02) {
|
|
atomic_store_explicit(counts + tpitg, 0u, memory_order_relaxed);
|
|
}
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
for (int i21 = tpitg; i21 < args.ne21; i21 += ntg) {
|
|
device const int32_t * src2_i32 =
|
|
(device const int32_t *)(src2 + i21*args.nb21);
|
|
|
|
#pragma unroll(ne20)
|
|
for (short i20 = 0; i20 < ne20; i20++) {
|
|
const int32_t expert = src2_i32[i20];
|
|
if ((uint32_t)expert >= (uint32_t)args.ne02) {
|
|
continue;
|
|
}
|
|
|
|
const uint32_t row = atomic_fetch_add_explicit(
|
|
counts + expert, 1u, memory_order_relaxed);
|
|
// Production top-k selections are unique. Keep malformed or
|
|
// synthetic duplicates from exceeding the fixed expert slice.
|
|
if (row < (uint32_t)args.ne21) {
|
|
ids_i32[(uint32_t)expert*args.ne21 + row] =
|
|
i21*ne20 + i20;
|
|
}
|
|
}
|
|
}
|
|
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
const short ide = tpitg;
|
|
const uint32_t n_all = min(
|
|
atomic_load_explicit(counts + ide, memory_order_relaxed),
|
|
(uint32_t)args.ne21);
|
|
tpe_u32[ide] = n_all;
|
|
|
|
const uint16_t n_tiles = (uint16_t)((n_all + 31u) / 32u);
|
|
tile_counts[ide] = n_tiles;
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
uint32_t tile_base = 0;
|
|
for (ushort i = 0; i < ide; i++) {
|
|
tile_base += tile_counts[i];
|
|
}
|
|
|
|
device uint32_t * work_count = (device uint32_t *) work;
|
|
device uint2 * work_items = (device uint2 *)(work + 8);
|
|
for (uint32_t tile = 0; tile < n_tiles; tile++) {
|
|
work_items[tile_base + tile] =
|
|
uint2((uint32_t)ide, tile * 32u);
|
|
}
|
|
if (ide + 1u == ntg) {
|
|
// Direct consumers launch the padded work capacity and cull against
|
|
// this exact count. Zero work must overwrite any stale prior value.
|
|
work_count[0] = tile_base + n_tiles;
|
|
}
|
|
}
|
|
|
|
typedef decltype(kernel_mul_mm_id_map_scatter_work<6>)
|
|
kernel_mul_mm_id_map_scatter_work_t;
|
|
template [[host_name("kernel_mul_mm_id_map_scatter_work_ne20_6")]]
|
|
kernel kernel_mul_mm_id_map_scatter_work_t
|
|
kernel_mul_mm_id_map_scatter_work<6>;
|
|
|
|
// Batched routed-expert matmul. It reads the expert-major map produced above,
|
|
// loads selected expert weights, and writes results back to token-major slots
|
|
// so the DS4 FFN can apply SwiGLU, weighting, and the down projection.
|
|
template<short NR1, typename S0, typename S0_4x4, typename S0_8x8, typename S1, typename S1_2x4, typename S1_8x8, typename block_q, short nl, void (*dequantize_func)(device const block_q *, short, thread S0_4x4 &), typename T0, typename T0_4x4, typename T1, typename T1_2x4, bool CULL_TAIL_SIMDGROUPS = false>
|
|
kernel void kernel_mul_mm_id(
|
|
constant ds4_metal_args_mul_mm_id & args,
|
|
device const char * src0,
|
|
device const char * src1,
|
|
device const char * htpe,
|
|
device const char * hids,
|
|
device char * dst,
|
|
device const char * work,
|
|
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 int NR0 = 64;
|
|
static_assert(NR1 == 32, "kernel_mul_mm_id accumulator layout supports only 32 routed rows");
|
|
|
|
constexpr int NK = 32;
|
|
constexpr int NL0 = NK/16;
|
|
constexpr int NL1 = NK/8;
|
|
constexpr int SA_BYTES = NR0 * NR1 * (int)sizeof(S0);
|
|
|
|
threadgroup S0 * sa = (threadgroup S0 *)(shmem);
|
|
threadgroup S1 * sb = (threadgroup S1 *)(shmem + SA_BYTES);
|
|
|
|
device const uint32_t * work_count = (device const uint32_t *) work;
|
|
const uint32_t work_index = tgpig.x;
|
|
if (work_index >= work_count[0]) {
|
|
return;
|
|
}
|
|
device const uint2 * work_items = (device const uint2 *)(work + 8);
|
|
const uint2 item = work_items[work_index];
|
|
const int im = (int)item.x;
|
|
const int r0 = tgpig.y*NR0;
|
|
const int r1 = (int)item.y;
|
|
|
|
device const uint32_t * tpe_u32 = (device const uint32_t *) (htpe);
|
|
device const int32_t * ids_i32 = (device const int32_t *) (hids);
|
|
|
|
const int32_t neh1 = tpe_u32[im];
|
|
|
|
if (r1 >= neh1) {
|
|
return;
|
|
}
|
|
|
|
const short nr0 = (args.ne0 - r0 < NR0) ? (args.ne0 - r0) : NR0;
|
|
const short nr1 = ( neh1 - r1 < NR1) ? ( neh1 - r1) : NR1;
|
|
// SIMDgroups 0/1 own routed rows 0..15 and SIMDgroups 2/3 own rows
|
|
// 16..31. Keep all threads in staging and at every threadgroup barrier,
|
|
// but let the second row-half skip MMA on short final expert tiles.
|
|
const bool mma_active =
|
|
!CULL_TAIL_SIMDGROUPS || 16*(short)(sgitg/2) < nr1;
|
|
|
|
if (!ds4_tp_owns_expert(im, args.ne02, args.tp_rank, args.tp_world)) {
|
|
/* Unowned expert under the TP split: zero this tile's output rows so
|
|
* the downstream swiglu/sum stages stay unchanged. Each (token,slot)
|
|
* row belongs to exactly one expert, so nothing else writes them. */
|
|
for (short j = sgitg; j < nr1; j += 4) {
|
|
const int idj = ids_i32[im*args.ne21 + r1 + j];
|
|
|
|
const short ide = idj % args.ne20;
|
|
const short idt = idj / args.ne20;
|
|
|
|
device float * D = (device float *) dst + r0 + ide*args.ne0 + idt*args.ne1*args.ne0;
|
|
|
|
for (int i = tiisg; i < nr0; i += 32) {
|
|
D[i] = 0.0f;
|
|
}
|
|
}
|
|
return;
|
|
}
|
|
|
|
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 id = ids_i32[im*args.ne21 + r1 + lr1];
|
|
|
|
const short i11 = (id % args.ne20) % args.ne11;
|
|
const short i12 = (id / args.ne20);
|
|
const short i13 = 0;
|
|
|
|
const uint64_t offset0 = (uint64_t)(im - args.tp_expert_base)*args.nb02 + i13*args.nb03;
|
|
const short offset1 = il0/nl;
|
|
|
|
device const block_q * x = (device const block_q *)(src0 + args.nb01*(r0 + lr0) + offset0) + offset1;
|
|
|
|
const short iy = 8*(tiitg % NL1);
|
|
|
|
device const T1 * y = (device const T1 *)(src1
|
|
+ args.nb13*i13
|
|
+ args.nb12*i12
|
|
+ args.nb11*i11
|
|
+ args.nb10*iy);
|
|
|
|
S0_8x8 ma[4];
|
|
S1_8x8 mb[2];
|
|
|
|
simdgroup_float8x8 mc[8];
|
|
|
|
for (short i = 0; i < 8; i++){
|
|
mc[i] = make_filled_simdgroup_matrix<float, 8>(0.f);
|
|
}
|
|
|
|
for (int loop_k = 0; loop_k < args.ne00; loop_k += NK) {
|
|
if (is_same<T0_4x4, block_q>::value && FC_mul_mm_bc_inp) {
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
for (short i = 0; i < 16; i++) {
|
|
const short sx = 2*il0 + i/8;
|
|
const short sy = (tiitg/NL0)/8;
|
|
|
|
const short lx = (tiitg/NL0)%8;
|
|
const short ly = i%8;
|
|
|
|
const short ib = 8*sx + sy;
|
|
|
|
*(sa + 64*ib + 8*ly + lx) = loop_k + 16*il + i < args.ne00 ? *((device T0 *) x + i) : 0;
|
|
}
|
|
} else {
|
|
S0_4x4 temp_a;
|
|
dequantize_func(x, il, temp_a);
|
|
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
FOR_UNROLL (short i = 0; i < 16; i++) {
|
|
const short sx = 2*il0 + i/8;
|
|
const short sy = (tiitg/NL0)/8;
|
|
|
|
const short lx = (tiitg/NL0)%8;
|
|
const short ly = i%8;
|
|
|
|
const short ib = 8*sx + sy;
|
|
|
|
*(sa + 64*ib + 8*ly + lx) = temp_a[i/4][i%4];
|
|
}
|
|
}
|
|
|
|
if (FC_mul_mm_bc_inp) {
|
|
for (short i = 0; i < 8; ++i) {
|
|
const short sx = (tiitg%NL1);
|
|
const short sy = (tiitg/NL1)/8;
|
|
|
|
const short lx = i;
|
|
const short ly = (tiitg/NL1)%8;
|
|
|
|
const short ib = 4*sx + sy;
|
|
|
|
*(sb + 64*ib + 8*ly + lx) = loop_k + iy + i < args.ne00 ? (S1) *((device T1 *) y + i) : 0;
|
|
}
|
|
} else {
|
|
const short sx = (tiitg%NL1);
|
|
const short sy = (tiitg/NL1)/8;
|
|
|
|
const short ly = (tiitg/NL1)%8;
|
|
|
|
const short ib = 4*sx + sy;
|
|
|
|
*(threadgroup S1_2x4 *)(sb + 64*ib + 8*ly) = (S1_2x4)(*((device T1_2x4 *) y));
|
|
}
|
|
|
|
il = (il + 2 < nl) ? il + 2 : il % 2;
|
|
x = (il < 2) ? x + (2 + nl - 1)/nl : x;
|
|
|
|
y += NK;
|
|
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
threadgroup const S0 * lsma = (sa + 4*64*(sgitg%2));
|
|
threadgroup const S1 * lsmb = (sb + 2*64*(sgitg/2));
|
|
|
|
if (mma_active) {
|
|
FOR_UNROLL (short ik = 0; ik < NK/8; ik++) {
|
|
simdgroup_barrier(mem_flags::mem_none);
|
|
|
|
FOR_UNROLL (short i = 0; i < 4; i++) {
|
|
simdgroup_load(ma[i], lsma + 64*i, 8, 0, false);
|
|
}
|
|
|
|
simdgroup_barrier(mem_flags::mem_none);
|
|
|
|
FOR_UNROLL (short i = 0; i < 2; i++) {
|
|
simdgroup_load(mb[i], lsmb + 64*i, 8, 0, false);
|
|
}
|
|
|
|
simdgroup_barrier(mem_flags::mem_none);
|
|
|
|
FOR_UNROLL (short i = 0; i < 8; i++){
|
|
simdgroup_multiply_accumulate(mc[i], mb[i/4], ma[i%4], mc[i]);
|
|
}
|
|
|
|
lsma += 8*64;
|
|
lsmb += 4*64;
|
|
}
|
|
}
|
|
}
|
|
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
threadgroup float * temp_str = ((threadgroup float *) shmem) + 32*(sgitg&1) + (16*(sgitg >> 1))*NR0;
|
|
|
|
if (mma_active) {
|
|
for (short i = 0; i < 8; i++) {
|
|
simdgroup_store(mc[i], temp_str + 8*(i%4) + 8*NR0*(i/4), NR0, 0, false);
|
|
}
|
|
}
|
|
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
for (short j = sgitg; j < nr1; j += 4) {
|
|
const int idj = ids_i32[im*args.ne21 + r1 + j];
|
|
|
|
const short ide = idj % args.ne20;
|
|
const short idt = idj / args.ne20;
|
|
|
|
device float * D = (device float *) dst + r0 + ide*args.ne0 + idt*args.ne1*args.ne0;
|
|
device float4 * D4 = (device float4 *) D;
|
|
|
|
threadgroup float * C = (threadgroup float *) shmem + j*NR0;
|
|
threadgroup float4 * C4 = (threadgroup float4 *) C;
|
|
|
|
int i = tiisg;
|
|
for (; i < nr0/4; i += 32) {
|
|
*(D4 + i) = *(C4 + i);
|
|
}
|
|
|
|
i = (4*(nr0/4)) + tiisg;
|
|
for (; i < nr0; i += 32) {
|
|
*(D + i) = *(C + i);
|
|
}
|
|
}
|
|
}
|
|
|
|
// Address-table variant used by SSD streaming. The routing ids remain the
|
|
// model's original expert ids, but each expert's resident buffer is found via a
|
|
// GPU-address table instead of a contiguous full-layer tensor.
|
|
template<short NR1, typename S0, typename S0_4x4, typename S0_8x8, typename S1, typename S1_2x4, typename S1_8x8, typename block_q, short nl, void (*dequantize_func)(device const block_q *, short, thread S0_4x4 &), typename T0, typename T0_4x4, typename T1, typename T1_2x4>
|
|
kernel void kernel_mul_mm_id_addr(
|
|
constant ds4_metal_args_mul_mm_id & args,
|
|
device const uint64_t * src0_addrs,
|
|
device const char * src1,
|
|
device const char * htpe,
|
|
device const char * hids,
|
|
device char * dst,
|
|
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 int NR0 = 64;
|
|
static_assert(NR1 == 32, "kernel_mul_mm_id_addr accumulator layout supports only 32 routed rows");
|
|
|
|
constexpr int NK = 32;
|
|
constexpr int NL0 = NK/16;
|
|
constexpr int NL1 = NK/8;
|
|
constexpr int SA_BYTES = NR0 * NR1 * (int)sizeof(S0);
|
|
|
|
threadgroup S0 * sa = (threadgroup S0 *)(shmem);
|
|
threadgroup S1 * sb = (threadgroup S1 *)(shmem + SA_BYTES);
|
|
|
|
const int im = tgpig.z;
|
|
const int r0 = tgpig.y*NR0;
|
|
const int r1 = tgpig.x*NR1;
|
|
|
|
device const uint32_t * tpe_u32 = (device const uint32_t *) (htpe);
|
|
device const int32_t * ids_i32 = (device const int32_t *) (hids);
|
|
|
|
const int32_t neh1 = tpe_u32[im];
|
|
|
|
if (r1 >= neh1) {
|
|
return;
|
|
}
|
|
|
|
const uint64_t base_addr = src0_addrs[im];
|
|
if (base_addr == 0) {
|
|
return;
|
|
}
|
|
device const char * src0 = reinterpret_cast<device const char *>(base_addr);
|
|
|
|
const short nr0 = (args.ne0 - r0 < NR0) ? (args.ne0 - r0) : NR0;
|
|
const short nr1 = ( neh1 - r1 < NR1) ? ( neh1 - 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 id = ids_i32[im*args.ne21 + r1 + lr1];
|
|
|
|
const short i11 = (id % args.ne20) % args.ne11;
|
|
const short i12 = (id / args.ne20);
|
|
const short i13 = 0;
|
|
|
|
const short offset1 = il0/nl;
|
|
|
|
device const block_q * x = (device const block_q *)(src0 + args.nb01*(r0 + lr0)) + offset1;
|
|
|
|
const short iy = 8*(tiitg % NL1);
|
|
|
|
device const T1 * y = (device const T1 *)(src1
|
|
+ args.nb13*i13
|
|
+ args.nb12*i12
|
|
+ args.nb11*i11
|
|
+ args.nb10*iy);
|
|
|
|
S0_8x8 ma[4];
|
|
S1_8x8 mb[2];
|
|
|
|
simdgroup_float8x8 mc[8];
|
|
|
|
for (short i = 0; i < 8; i++){
|
|
mc[i] = make_filled_simdgroup_matrix<float, 8>(0.f);
|
|
}
|
|
|
|
for (int loop_k = 0; loop_k < args.ne00; loop_k += NK) {
|
|
if (is_same<T0_4x4, block_q>::value && FC_mul_mm_bc_inp) {
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
for (short i = 0; i < 16; i++) {
|
|
const short sx = 2*il0 + i/8;
|
|
const short sy = (tiitg/NL0)/8;
|
|
|
|
const short lx = (tiitg/NL0)%8;
|
|
const short ly = i%8;
|
|
|
|
const short ib = 8*sx + sy;
|
|
|
|
*(sa + 64*ib + 8*ly + lx) = loop_k + 16*il + i < args.ne00 ? *((device T0 *) x + i) : 0;
|
|
}
|
|
} else {
|
|
S0_4x4 temp_a;
|
|
dequantize_func(x, il, temp_a);
|
|
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
FOR_UNROLL (short i = 0; i < 16; i++) {
|
|
const short sx = 2*il0 + i/8;
|
|
const short sy = (tiitg/NL0)/8;
|
|
|
|
const short lx = (tiitg/NL0)%8;
|
|
const short ly = i%8;
|
|
|
|
const short ib = 8*sx + sy;
|
|
|
|
*(sa + 64*ib + 8*ly + lx) = temp_a[i/4][i%4];
|
|
}
|
|
}
|
|
|
|
if (FC_mul_mm_bc_inp) {
|
|
for (short i = 0; i < 8; ++i) {
|
|
const short sx = (tiitg%NL1);
|
|
const short sy = (tiitg/NL1)/8;
|
|
|
|
const short lx = i;
|
|
const short ly = (tiitg/NL1)%8;
|
|
|
|
const short ib = 4*sx + sy;
|
|
|
|
*(sb + 64*ib + 8*ly + lx) = loop_k + iy + i < args.ne00 ? (S1) *((device T1 *) y + i) : 0;
|
|
}
|
|
} else {
|
|
const short sx = (tiitg%NL1);
|
|
const short sy = (tiitg/NL1)/8;
|
|
|
|
const short ly = (tiitg/NL1)%8;
|
|
|
|
const short ib = 4*sx + sy;
|
|
|
|
*(threadgroup S1_2x4 *)(sb + 64*ib + 8*ly) = (S1_2x4)(*((device T1_2x4 *) y));
|
|
}
|
|
|
|
il = (il + 2 < nl) ? il + 2 : il % 2;
|
|
x = (il < 2) ? x + (2 + nl - 1)/nl : x;
|
|
|
|
y += NK;
|
|
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
threadgroup const S0 * lsma = (sa + 4*64*(sgitg%2));
|
|
threadgroup const S1 * lsmb = (sb + 2*64*(sgitg/2));
|
|
|
|
FOR_UNROLL (short ik = 0; ik < NK/8; ik++) {
|
|
simdgroup_barrier(mem_flags::mem_none);
|
|
|
|
FOR_UNROLL (short i = 0; i < 4; i++) {
|
|
simdgroup_load(ma[i], lsma + 64*i, 8, 0, false);
|
|
}
|
|
|
|
simdgroup_barrier(mem_flags::mem_none);
|
|
|
|
FOR_UNROLL (short i = 0; i < 2; i++) {
|
|
simdgroup_load(mb[i], lsmb + 64*i, 8, 0, false);
|
|
}
|
|
|
|
simdgroup_barrier(mem_flags::mem_none);
|
|
|
|
FOR_UNROLL (short i = 0; i < 8; i++){
|
|
simdgroup_multiply_accumulate(mc[i], mb[i/4], ma[i%4], mc[i]);
|
|
}
|
|
|
|
lsma += 8*64;
|
|
lsmb += 4*64;
|
|
}
|
|
}
|
|
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
threadgroup float * temp_str = ((threadgroup float *) shmem) + 32*(sgitg&1) + (16*(sgitg >> 1))*NR0;
|
|
|
|
for (short i = 0; i < 8; i++) {
|
|
simdgroup_store(mc[i], temp_str + 8*(i%4) + 8*NR0*(i/4), NR0, 0, false);
|
|
}
|
|
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
for (short j = sgitg; j < nr1; j += 4) {
|
|
const int idj = ids_i32[im*args.ne21 + r1 + j];
|
|
|
|
const short ide = idj % args.ne20;
|
|
const short idt = idj / args.ne20;
|
|
|
|
device float * D = (device float *) dst + r0 + ide*args.ne0 + idt*args.ne1*args.ne0;
|
|
device float4 * D4 = (device float4 *) D;
|
|
|
|
threadgroup float * C = (threadgroup float *) shmem + j*NR0;
|
|
threadgroup float4 * C4 = (threadgroup float4 *) C;
|
|
|
|
int i = tiisg;
|
|
for (; i < nr0/4; i += 32) {
|
|
*(D4 + i) = *(C4 + i);
|
|
}
|
|
|
|
i = (4*(nr0/4)) + tiisg;
|
|
for (; i < nr0; i += 32) {
|
|
*(D + i) = *(C + i);
|
|
}
|
|
}
|
|
}
|
|
|
|
// Fused routed gate+up grouped matmul with the SwiGLU epilogue, generic over
|
|
// the expert quant block. Both A streams share one staged B tile per k-step;
|
|
// each output keeps the exact MMA accumulation order of the separate GEMMs,
|
|
// and the epilogue matches kernel_dsv4_moe_swiglu_weight_f16, so the fused
|
|
// result is bit-identical to the unfused path.
|
|
template<typename block_q, short nl, void (*dequantize_func)(device const block_q *, short, thread half4x4 &), bool CULL_TAIL_SIMDGROUPS = false>
|
|
kernel void kernel_mul_mm_id_pair_swiglu_f16_impl(
|
|
constant ds4_metal_args_mul_mm_id & args,
|
|
constant ds4_metal_dsv4_moe_swiglu_weight_args & act,
|
|
device const char * src0_gate,
|
|
device const char * src0_up,
|
|
device const char * src1,
|
|
device const char * htpe,
|
|
device const char * hids,
|
|
device char * dst_mid,
|
|
device const char * weights,
|
|
device const char * work,
|
|
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]]) {
|
|
threadgroup half *sa_gate = (threadgroup half *)(shmem);
|
|
threadgroup half *sa_up = (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;
|
|
|
|
device const uint32_t * work_count = (device const uint32_t *) work;
|
|
const uint32_t work_index = tgpig.x;
|
|
if (work_index >= work_count[0]) {
|
|
return;
|
|
}
|
|
device const uint2 * work_items = (device const uint2 *)(work + 8);
|
|
const uint2 item = work_items[work_index];
|
|
const int im = (int)item.x;
|
|
const int r0 = tgpig.y*NR0;
|
|
const int r1 = (int)item.y;
|
|
|
|
device const uint32_t * tpe_u32 = (device const uint32_t *) (htpe);
|
|
device const int32_t * ids_i32 = (device const int32_t *) (hids);
|
|
|
|
const int32_t neh1 = tpe_u32[im];
|
|
|
|
if (r1 >= neh1) {
|
|
return;
|
|
}
|
|
|
|
const short nr0 = (args.ne0 - r0 < NR0) ? (args.ne0 - r0) : NR0;
|
|
const short nr1 = ( neh1 - r1 < NR1) ? ( neh1 - r1) : NR1;
|
|
// SIMDgroups 0/1 own routed rows 0..15 and SIMDgroups 2/3 own rows
|
|
// 16..31. Keep every thread in staging and at every threadgroup barrier,
|
|
// but let the second row-half skip MMA on short final expert tiles.
|
|
const bool mma_active =
|
|
!CULL_TAIL_SIMDGROUPS || 16*(short)(sgitg/2) < 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 id = ids_i32[im*args.ne21 + r1 + lr1];
|
|
|
|
const short i11 = (id % args.ne20) % args.ne11;
|
|
const short i12 = (id / args.ne20);
|
|
const short i13 = 0;
|
|
|
|
const uint64_t offset0 = im*args.nb02 + i13*args.nb03;
|
|
const short offset1 = il0/nl;
|
|
|
|
device const block_q * xg =
|
|
(device const block_q *)(src0_gate + args.nb01*(r0 + lr0) + offset0) + offset1;
|
|
device const block_q * xu =
|
|
(device const block_q *)(src0_up + 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*i11
|
|
+ args.nb10*iy);
|
|
|
|
simdgroup_half8x8 ma_g[4];
|
|
simdgroup_half8x8 ma_u[4];
|
|
simdgroup_half8x8 mb[2];
|
|
|
|
simdgroup_float8x8 mc_gate[8];
|
|
simdgroup_float8x8 mc_up[8];
|
|
|
|
for (short i = 0; i < 8; i++) {
|
|
mc_gate[i] = make_filled_simdgroup_matrix<float, 8>(0.f);
|
|
mc_up[i] = make_filled_simdgroup_matrix<float, 8>(0.f);
|
|
}
|
|
|
|
for (int loop_k = 0; loop_k < args.ne00; loop_k += NK) {
|
|
half4x4 temp_gate;
|
|
dequantize_func(xg, il, temp_gate);
|
|
half4x4 temp_up;
|
|
dequantize_func(xu, il, temp_up);
|
|
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
// Stage the shared B tile and both expert A tiles in one pass: two
|
|
// threadgroup barriers per k-step total instead of two per matrix.
|
|
{
|
|
const short sx_b = (tiitg%NL1);
|
|
const short sy_b = (tiitg/NL1)/8;
|
|
const short ly_b = (tiitg/NL1)%8;
|
|
const short ib_b = 4*sx_b + sy_b;
|
|
*(threadgroup half2x4 *)(sb + 64*ib_b + 8*ly_b) =
|
|
(half2x4)(*((device float2x4 *) y));
|
|
}
|
|
|
|
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_gate + 64*ib + 8*ly + lx) = temp_gate[i/4][i%4];
|
|
*(sa_up + 64*ib + 8*ly + lx) = temp_up[i/4][i%4];
|
|
}
|
|
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
threadgroup const half * lsma_gate = (sa_gate + 4*64*(sgitg%2));
|
|
threadgroup const half * lsma_up = (sa_up + 4*64*(sgitg%2));
|
|
threadgroup const half * lsmb = (sb + 2*64*(sgitg/2));
|
|
|
|
if (mma_active) {
|
|
FOR_UNROLL (short ik = 0; ik < NK/8; ik++) {
|
|
simdgroup_barrier(mem_flags::mem_none);
|
|
|
|
FOR_UNROLL (short i = 0; i < 4; i++) {
|
|
simdgroup_load(ma_g[i], lsma_gate + 64*i, 8, 0, false);
|
|
simdgroup_load(ma_u[i], lsma_up + 64*i, 8, 0, false);
|
|
}
|
|
FOR_UNROLL (short i = 0; i < 2; i++) {
|
|
simdgroup_load(mb[i], lsmb + 64*i, 8, 0, false);
|
|
}
|
|
|
|
simdgroup_barrier(mem_flags::mem_none);
|
|
|
|
FOR_UNROLL (short i = 0; i < 8; i++) {
|
|
simdgroup_multiply_accumulate(mc_gate[i], mb[i/4], ma_g[i%4], mc_gate[i]);
|
|
simdgroup_multiply_accumulate(mc_up[i], mb[i/4], ma_u[i%4], mc_up[i]);
|
|
}
|
|
|
|
lsma_gate += 8*64;
|
|
lsma_up += 8*64;
|
|
lsmb += 4*64;
|
|
}
|
|
}
|
|
|
|
il = (il + 2 < nl) ? il + 2 : il % 2;
|
|
xg = (il < 2) ? xg + (2 + nl - 1)/nl : xg;
|
|
xu = (il < 2) ? xu + (2 + nl - 1)/nl : xu;
|
|
y += NK;
|
|
}
|
|
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
threadgroup float * temp_gate = (threadgroup float *) shmem;
|
|
threadgroup float * temp_up = temp_gate + NR0*NR1;
|
|
threadgroup float * temp_gate_str =
|
|
temp_gate + 32*(sgitg&1) + (16*(sgitg >> 1))*NR0;
|
|
threadgroup float * temp_up_str =
|
|
temp_up + 32*(sgitg&1) + (16*(sgitg >> 1))*NR0;
|
|
|
|
if (mma_active) {
|
|
for (short i = 0; i < 8; i++) {
|
|
simdgroup_store(mc_gate[i], temp_gate_str + 8*(i%4) + 8*NR0*(i/4), NR0, 0, false);
|
|
simdgroup_store(mc_up[i], temp_up_str + 8*(i%4) + 8*NR0*(i/4), NR0, 0, false);
|
|
}
|
|
}
|
|
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
const float c = act.clamp_value;
|
|
for (short j = sgitg; j < nr1; j += 4) {
|
|
const int idj = ids_i32[im*args.ne21 + r1 + j];
|
|
|
|
const short ide = idj % args.ne20;
|
|
const short idt = idj / args.ne20;
|
|
|
|
device half *D = (device half *)(dst_mid +
|
|
((uint64_t)idt*args.ne1 + (uint64_t)ide)*act.mid_row_stride) + r0;
|
|
device const float *w = (device const float *)(weights + (uint64_t)idj*act.weight_stride);
|
|
const float route_weight = w[0];
|
|
|
|
threadgroup float *Cg = temp_gate + j*NR0;
|
|
threadgroup float *Cu = temp_up + j*NR0;
|
|
|
|
int i = tiisg;
|
|
for (; i < nr0; i += 32) {
|
|
float g = Cg[i];
|
|
float u = Cu[i];
|
|
if (c > 1.0e-6f) {
|
|
g = min(g, c);
|
|
u = clamp(u, -c, c);
|
|
}
|
|
const float silu = g / (1.0f + exp(-g));
|
|
D[i] = (half)(silu * u * route_weight);
|
|
}
|
|
}
|
|
}
|
|
|
|
// MXFP4 resident-prefill specialization with a compact 32x32
|
|
// output/routed-row tile. Two SIMDgroups share the same four 8-column A
|
|
// blocks and each owns one 16-row half of B. The per-output dequantization,
|
|
// MMA K sequence, and SwiGLU epilogue are identical to the 64x32 kernel
|
|
// above. On a short final work item, SIMDgroup 1 remains in staging and at
|
|
// every threadgroup barrier but skips its MMA/store when rows 16..31 are
|
|
// empty.
|
|
template<typename block_q, short nl, void (*dequantize_func)(device const block_q *, short, thread half4x4 &)>
|
|
kernel void kernel_mul_mm_id_pair_swiglu_f16_compact_tail_impl(
|
|
constant ds4_metal_args_mul_mm_id & args,
|
|
constant ds4_metal_dsv4_moe_swiglu_weight_args & act,
|
|
device const char * src0_gate,
|
|
device const char * src0_up,
|
|
device const char * src1,
|
|
device const char * htpe,
|
|
device const char * hids,
|
|
device char * dst_mid,
|
|
device const char * weights,
|
|
device const char * work,
|
|
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]]) {
|
|
threadgroup half *sa_gate = (threadgroup half *)(shmem);
|
|
threadgroup half *sa_up = (threadgroup half *)(shmem + 2048);
|
|
threadgroup half *sb = (threadgroup half *)(shmem + 4096);
|
|
|
|
constexpr int NR0 = 32;
|
|
constexpr int NR1 = 32;
|
|
constexpr int NK = 32;
|
|
constexpr int NL0 = NK/16;
|
|
constexpr int NL1 = NK/8;
|
|
|
|
device const uint32_t * work_count = (device const uint32_t *) work;
|
|
const uint32_t work_index = tgpig.x;
|
|
if (work_index >= work_count[0]) {
|
|
return;
|
|
}
|
|
device const uint2 * work_items = (device const uint2 *)(work + 8);
|
|
const uint2 item = work_items[work_index];
|
|
const int im = (int)item.x;
|
|
const int r0 = tgpig.y*NR0;
|
|
const int r1 = (int)item.y;
|
|
|
|
device const uint32_t * tpe_u32 = (device const uint32_t *) (htpe);
|
|
device const int32_t * ids_i32 = (device const int32_t *) (hids);
|
|
|
|
const int32_t neh1 = tpe_u32[im];
|
|
if (r1 >= neh1) {
|
|
return;
|
|
}
|
|
|
|
const short nr0 = (args.ne0 - r0 < NR0) ? (args.ne0 - r0) : NR0;
|
|
const short nr1 = ( neh1 - r1 < NR1) ? ( neh1 - r1) : NR1;
|
|
const bool mma_active = 16*(short)sgitg < nr1;
|
|
|
|
const short lr0 = ((short)tiitg/NL0) < nr0 ? ((short)tiitg/NL0) : nr0 - 1;
|
|
const short row_b0 = (short)tiitg/NL1;
|
|
const short row_b1 = row_b0 + 16;
|
|
const short lr1_b0 = row_b0 < nr1 ? row_b0 : nr1 - 1;
|
|
const short lr1_b1 = row_b1 < nr1 ? row_b1 : nr1 - 1;
|
|
|
|
const short il0 = (tiitg % NL0);
|
|
short il = il0;
|
|
|
|
const int id_b0 = ids_i32[im*args.ne21 + r1 + lr1_b0];
|
|
const int id_b1 = ids_i32[im*args.ne21 + r1 + lr1_b1];
|
|
|
|
const short i11_b0 = (id_b0 % args.ne20) % args.ne11;
|
|
const short i12_b0 = (id_b0 / args.ne20);
|
|
const short i11_b1 = (id_b1 % args.ne20) % args.ne11;
|
|
const short i12_b1 = (id_b1 / args.ne20);
|
|
const short i13 = 0;
|
|
|
|
const uint64_t offset0 = im*args.nb02 + i13*args.nb03;
|
|
const short offset1 = il0/nl;
|
|
|
|
device const block_q * xg =
|
|
(device const block_q *)(src0_gate + args.nb01*(r0 + lr0) + offset0) + offset1;
|
|
device const block_q * xu =
|
|
(device const block_q *)(src0_up + args.nb01*(r0 + lr0) + offset0) + offset1;
|
|
|
|
const short iy = 8*(tiitg % NL1);
|
|
|
|
device const float * y_b0 = (device const float *)(src1
|
|
+ args.nb13*i13
|
|
+ args.nb12*i12_b0
|
|
+ args.nb11*i11_b0
|
|
+ args.nb10*iy);
|
|
device const float * y_b1 = (device const float *)(src1
|
|
+ args.nb13*i13
|
|
+ args.nb12*i12_b1
|
|
+ args.nb11*i11_b1
|
|
+ args.nb10*iy);
|
|
|
|
simdgroup_half8x8 ma_g[4];
|
|
simdgroup_half8x8 ma_u[4];
|
|
simdgroup_half8x8 mb[2];
|
|
|
|
simdgroup_float8x8 mc_gate[8];
|
|
simdgroup_float8x8 mc_up[8];
|
|
|
|
for (short i = 0; i < 8; i++) {
|
|
mc_gate[i] = make_filled_simdgroup_matrix<float, 8>(0.f);
|
|
mc_up[i] = make_filled_simdgroup_matrix<float, 8>(0.f);
|
|
}
|
|
|
|
for (int loop_k = 0; loop_k < args.ne00; loop_k += NK) {
|
|
half4x4 temp_gate;
|
|
dequantize_func(xg, il, temp_gate);
|
|
half4x4 temp_up;
|
|
dequantize_func(xu, il, temp_up);
|
|
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
// Sixty-four threads stage all 32 B rows by loading one row in each
|
|
// 16-row half. The shared-memory block order matches the 64x32 path.
|
|
{
|
|
const short sx_b = (tiitg%NL1);
|
|
const short sy_b0 = row_b0/8;
|
|
const short sy_b1 = row_b1/8;
|
|
const short ly_b = row_b0%8;
|
|
const short ib_b0 = 4*sx_b + sy_b0;
|
|
const short ib_b1 = 4*sx_b + sy_b1;
|
|
*(threadgroup half2x4 *)(sb + 64*ib_b0 + 8*ly_b) =
|
|
(half2x4)(*((device float2x4 *) y_b0));
|
|
*(threadgroup half2x4 *)(sb + 64*ib_b1 + 8*ly_b) =
|
|
(half2x4)(*((device float2x4 *) y_b1));
|
|
}
|
|
|
|
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 = 4*sx + sy;
|
|
*(sa_gate + 64*ib + 8*ly + lx) = temp_gate[i/4][i%4];
|
|
*(sa_up + 64*ib + 8*ly + lx) = temp_up[i/4][i%4];
|
|
}
|
|
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
threadgroup const half * lsma_gate = sa_gate;
|
|
threadgroup const half * lsma_up = sa_up;
|
|
threadgroup const half * lsmb = sb + 2*64*sgitg;
|
|
|
|
if (mma_active) {
|
|
FOR_UNROLL (short ik = 0; ik < NK/8; ik++) {
|
|
simdgroup_barrier(mem_flags::mem_none);
|
|
|
|
FOR_UNROLL (short i = 0; i < 4; i++) {
|
|
simdgroup_load(ma_g[i], lsma_gate + 64*i, 8, 0, false);
|
|
simdgroup_load(ma_u[i], lsma_up + 64*i, 8, 0, false);
|
|
}
|
|
FOR_UNROLL (short i = 0; i < 2; i++) {
|
|
simdgroup_load(mb[i], lsmb + 64*i, 8, 0, false);
|
|
}
|
|
|
|
simdgroup_barrier(mem_flags::mem_none);
|
|
|
|
FOR_UNROLL (short i = 0; i < 8; i++) {
|
|
simdgroup_multiply_accumulate(mc_gate[i], mb[i/4], ma_g[i%4], mc_gate[i]);
|
|
simdgroup_multiply_accumulate(mc_up[i], mb[i/4], ma_u[i%4], mc_up[i]);
|
|
}
|
|
|
|
lsma_gate += 4*64;
|
|
lsma_up += 4*64;
|
|
lsmb += 4*64;
|
|
}
|
|
}
|
|
|
|
il = (il + 2 < nl) ? il + 2 : il % 2;
|
|
xg = (il < 2) ? xg + (2 + nl - 1)/nl : xg;
|
|
xu = (il < 2) ? xu + (2 + nl - 1)/nl : xu;
|
|
y_b0 += NK;
|
|
y_b1 += NK;
|
|
}
|
|
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
threadgroup float * temp_gate = (threadgroup float *) shmem;
|
|
threadgroup float * temp_up = temp_gate + NR0*NR1;
|
|
threadgroup float * temp_gate_str = temp_gate + 16*sgitg*NR0;
|
|
threadgroup float * temp_up_str = temp_up + 16*sgitg*NR0;
|
|
|
|
if (mma_active) {
|
|
for (short i = 0; i < 8; i++) {
|
|
simdgroup_store(mc_gate[i], temp_gate_str + 8*(i%4) + 8*NR0*(i/4), NR0, 0, false);
|
|
simdgroup_store(mc_up[i], temp_up_str + 8*(i%4) + 8*NR0*(i/4), NR0, 0, false);
|
|
}
|
|
}
|
|
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
const float c = act.clamp_value;
|
|
for (short j = sgitg; j < nr1; j += 2) {
|
|
const int idj = ids_i32[im*args.ne21 + r1 + j];
|
|
|
|
const short ide = idj % args.ne20;
|
|
const short idt = idj / args.ne20;
|
|
|
|
device half *D = (device half *)(dst_mid +
|
|
((uint64_t)idt*args.ne1 + (uint64_t)ide)*act.mid_row_stride) + r0;
|
|
device const float *w = (device const float *)(weights + (uint64_t)idj*act.weight_stride);
|
|
const float route_weight = w[0];
|
|
|
|
threadgroup float *Cg = temp_gate + j*NR0;
|
|
threadgroup float *Cu = temp_up + j*NR0;
|
|
|
|
int i = tiisg;
|
|
for (; i < nr0; i += 32) {
|
|
float g = Cg[i];
|
|
float u = Cu[i];
|
|
if (c > 1.0e-6f) {
|
|
g = min(g, c);
|
|
u = clamp(u, -c, c);
|
|
}
|
|
const float silu = g / (1.0f + exp(-g));
|
|
D[i] = (half)(silu * u * route_weight);
|
|
}
|
|
}
|
|
}
|
|
|
|
typedef decltype(kernel_mul_mm_id_pair_swiglu_f16_impl<block_iq2_xxs, QK_NL, dequantize_iq2_xxs>) mul_mm_id_pair_swiglu_f16_iq2;
|
|
typedef decltype(kernel_mul_mm_id_pair_swiglu_f16_impl<block_q4_K, QK_NL, dequantize_q4_K>) mul_mm_id_pair_swiglu_f16_q4;
|
|
typedef decltype(kernel_mul_mm_id_pair_swiglu_f16_impl<block_mxfp4, 2, dequantize_mxfp4>) mul_mm_id_pair_swiglu_f16_mxfp4;
|
|
typedef decltype(kernel_mul_mm_id_pair_swiglu_f16_impl<block_mxfp4, 2, dequantize_mxfp4, true>) mul_mm_id_pair_swiglu_f16_mxfp4_tail_cull;
|
|
typedef decltype(kernel_mul_mm_id_pair_swiglu_f16_compact_tail_impl<block_mxfp4, 2, dequantize_mxfp4>) mul_mm_id_pair_swiglu_f16_mxfp4_compact_tail;
|
|
|
|
// Host-visible fused routed pair matmuls for the DS4 expert quant formats.
|
|
template [[host_name("kernel_mul_mm_id_iq2_xxs_pair_swiglu_f16")]] kernel mul_mm_id_pair_swiglu_f16_iq2 kernel_mul_mm_id_pair_swiglu_f16_impl<block_iq2_xxs, QK_NL, dequantize_iq2_xxs>;
|
|
template [[host_name("kernel_mul_mm_id_q4_K_pair_swiglu_f16")]] kernel mul_mm_id_pair_swiglu_f16_q4 kernel_mul_mm_id_pair_swiglu_f16_impl<block_q4_K, QK_NL, dequantize_q4_K>;
|
|
template [[host_name("kernel_mul_mm_id_mxfp4_pair_swiglu_f16")]] kernel mul_mm_id_pair_swiglu_f16_mxfp4 kernel_mul_mm_id_pair_swiglu_f16_impl<block_mxfp4, 2, dequantize_mxfp4>;
|
|
template [[host_name("kernel_mul_mm_id_mxfp4_pair_swiglu_f16_half_scale")]] kernel mul_mm_id_pair_swiglu_f16_mxfp4 kernel_mul_mm_id_pair_swiglu_f16_impl<block_mxfp4, 2, dequantize_mxfp4_half_scale>;
|
|
template [[host_name("kernel_mul_mm_id_mxfp4_pair_swiglu_f16_tail_cull_half_scale")]] kernel mul_mm_id_pair_swiglu_f16_mxfp4_tail_cull kernel_mul_mm_id_pair_swiglu_f16_impl<block_mxfp4, 2, dequantize_mxfp4_half_scale, true>;
|
|
template [[host_name("kernel_mul_mm_id_mxfp4_pair_swiglu_f16_compact_tail_cull_half_scale")]] kernel mul_mm_id_pair_swiglu_f16_mxfp4_compact_tail kernel_mul_mm_id_pair_swiglu_f16_compact_tail_impl<block_mxfp4, 2, dequantize_mxfp4_half_scale>;
|
|
template [[host_name("kernel_mul_mm_id_mxfp4_pair_swiglu_f16_tail_cull")]] kernel mul_mm_id_pair_swiglu_f16_mxfp4_tail_cull kernel_mul_mm_id_pair_swiglu_f16_impl<block_mxfp4, 2, dequantize_mxfp4, true>;
|
|
template [[host_name("kernel_mul_mm_id_mxfp4_pair_swiglu_f16_compact_tail_cull")]] kernel mul_mm_id_pair_swiglu_f16_mxfp4_compact_tail kernel_mul_mm_id_pair_swiglu_f16_compact_tail_impl<block_mxfp4, 2, dequantize_mxfp4>;
|
|
|
|
typedef decltype(kernel_mul_mm_id<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q2_K, QK_NL, dequantize_q2_K, float, float4x4, float, float2x4>) mul_mm_id;
|
|
typedef decltype(kernel_mul_mm_id<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q2_K, QK_NL, dequantize_q2_K, half, half4x4, half, half2x4>) mul_mm_id_f16_rhs;
|
|
typedef decltype(kernel_mul_mm_id<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_mxfp4, 2, dequantize_mxfp4, half, half4x4, half, half2x4, true>) mul_mm_id_mxfp4_f16_rhs_tail_cull;
|
|
typedef decltype(kernel_mul_mm_id<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_mxfp4, 2, dequantize_mxfp4_half_lut, half, half4x4, half, half2x4>) mul_mm_id_mxfp4_f16_rhs_half_lut;
|
|
typedef decltype(kernel_mul_mm_id<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_mxfp4, 2, dequantize_mxfp4_half_lut, half, half4x4, half, half2x4, true>) mul_mm_id_mxfp4_f16_rhs_half_lut_tail_cull;
|
|
typedef decltype(kernel_mul_mm_id_addr<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q2_K, QK_NL, dequantize_q2_K, float, float4x4, float, float2x4>) mul_mm_id_addr;
|
|
typedef decltype(kernel_mul_mm_id_addr<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q2_K, QK_NL, dequantize_q2_K, half, half4x4, half, half2x4>) mul_mm_id_addr_f16_rhs;
|
|
|
|
// Host-visible batched MoE matmul variants for the DS4 quant formats.
|
|
template [[host_name("kernel_mul_mm_id_q8_0_f32")]] kernel mul_mm_id kernel_mul_mm_id<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q8_0, 2, dequantize_q8_0, float, float4x4, float, float2x4>;
|
|
template [[host_name("kernel_mul_mm_id_q2_K_f32")]] kernel mul_mm_id kernel_mul_mm_id<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q2_K, QK_NL, dequantize_q2_K, float, float4x4, float, float2x4>;
|
|
template [[host_name("kernel_mul_mm_id_q4_K_f32")]] kernel mul_mm_id kernel_mul_mm_id<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q4_K, QK_NL, dequantize_q4_K, float, float4x4, float, float2x4>;
|
|
template [[host_name("kernel_mul_mm_id_q5_K_f32")]] kernel mul_mm_id kernel_mul_mm_id<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q5_K, QK_NL, dequantize_q5_K, float, float4x4, float, float2x4>;
|
|
template [[host_name("kernel_mul_mm_id_q6_K_f32")]] kernel mul_mm_id kernel_mul_mm_id<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q6_K, QK_NL, dequantize_q6_K, float, float4x4, float, float2x4>;
|
|
template [[host_name("kernel_mul_mm_id_iq2_xxs_f32")]] kernel mul_mm_id kernel_mul_mm_id<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_iq2_xxs, QK_NL, dequantize_iq2_xxs, float, float4x4, float, float2x4>;
|
|
template [[host_name("kernel_mul_mm_id_mxfp4_f32")]] kernel mul_mm_id kernel_mul_mm_id<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_mxfp4, 2, dequantize_mxfp4, float, float4x4, float, float2x4>;
|
|
template [[host_name("kernel_mul_mm_id_q8_0_f16")]] kernel mul_mm_id_f16_rhs kernel_mul_mm_id<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q8_0, 2, dequantize_q8_0, half, half4x4, half, half2x4>;
|
|
template [[host_name("kernel_mul_mm_id_q2_K_f16")]] kernel mul_mm_id_f16_rhs kernel_mul_mm_id<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q2_K, QK_NL, dequantize_q2_K, half, half4x4, half, half2x4>;
|
|
template [[host_name("kernel_mul_mm_id_q4_K_f16")]] kernel mul_mm_id_f16_rhs kernel_mul_mm_id<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q4_K, QK_NL, dequantize_q4_K, half, half4x4, half, half2x4>;
|
|
template [[host_name("kernel_mul_mm_id_q5_K_f16")]] kernel mul_mm_id_f16_rhs kernel_mul_mm_id<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q5_K, QK_NL, dequantize_q5_K, half, half4x4, half, half2x4>;
|
|
template [[host_name("kernel_mul_mm_id_q6_K_f16")]] kernel mul_mm_id_f16_rhs kernel_mul_mm_id<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q6_K, QK_NL, dequantize_q6_K, half, half4x4, half, half2x4>;
|
|
template [[host_name("kernel_mul_mm_id_iq2_xxs_f16")]] kernel mul_mm_id_f16_rhs kernel_mul_mm_id<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_iq2_xxs, QK_NL, dequantize_iq2_xxs, half, half4x4, half, half2x4>;
|
|
template [[host_name("kernel_mul_mm_id_mxfp4_f16")]] kernel mul_mm_id_f16_rhs kernel_mul_mm_id<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_mxfp4, 2, dequantize_mxfp4, half, half4x4, half, half2x4>;
|
|
template [[host_name("kernel_mul_mm_id_mxfp4_f16_tail_cull")]] kernel mul_mm_id_mxfp4_f16_rhs_tail_cull kernel_mul_mm_id<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_mxfp4, 2, dequantize_mxfp4, half, half4x4, half, half2x4, true>;
|
|
template [[host_name("kernel_mul_mm_id_mxfp4_f16_half_lut")]] kernel mul_mm_id_mxfp4_f16_rhs_half_lut kernel_mul_mm_id<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_mxfp4, 2, dequantize_mxfp4_half_lut, half, half4x4, half, half2x4>;
|
|
template [[host_name("kernel_mul_mm_id_mxfp4_f16_half_lut_tail_cull")]] kernel mul_mm_id_mxfp4_f16_rhs_half_lut_tail_cull kernel_mul_mm_id<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_mxfp4, 2, dequantize_mxfp4_half_lut, half, half4x4, half, half2x4, true>;
|
|
|
|
template [[host_name("kernel_mul_mm_id_addr_q2_K_f32")]] kernel mul_mm_id_addr kernel_mul_mm_id_addr<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q2_K, QK_NL, dequantize_q2_K, float, float4x4, float, float2x4>;
|
|
template [[host_name("kernel_mul_mm_id_addr_q4_K_f32")]] kernel mul_mm_id_addr kernel_mul_mm_id_addr<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q4_K, QK_NL, dequantize_q4_K, float, float4x4, float, float2x4>;
|
|
template [[host_name("kernel_mul_mm_id_addr_mxfp4_f32")]] kernel mul_mm_id_addr kernel_mul_mm_id_addr<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_mxfp4, 2, dequantize_mxfp4, float, float4x4, float, float2x4>;
|
|
template [[host_name("kernel_mul_mm_id_addr_q2_K_f16")]] kernel mul_mm_id_addr_f16_rhs kernel_mul_mm_id_addr<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q2_K, QK_NL, dequantize_q2_K, half, half4x4, half, half2x4>;
|
|
template [[host_name("kernel_mul_mm_id_addr_q4_K_f16")]] kernel mul_mm_id_addr_f16_rhs kernel_mul_mm_id_addr<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q4_K, QK_NL, dequantize_q4_K, half, half4x4, half, half2x4>;
|
|
template [[host_name("kernel_mul_mm_id_addr_mxfp4_f16")]] kernel mul_mm_id_addr_f16_rhs kernel_mul_mm_id_addr<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_mxfp4, 2, dequantize_mxfp4, half, half4x4, half, half2x4>;
|
|
|
|
#ifdef DS4_METAL_HAS_TENSOR
|
|
// Attention-output low-rank projection retained for Metal4 prefill. It uses
|
|
// the same direct-RHS idea as dense matmul: dequantize the Q8_0 or Q4_K low
|
|
// projection weights to a half tile, then let TensorOps read the dense head
|
|
// activations directly. Only the 64-token direct-RHS instantiations are
|
|
// exported because the staged-RHS and 32-token variants were benchmark-only
|
|
// experiments.
|
|
//
|
|
// Full tiles (the host dispatch guarantee for aligned batches) skip all bounds
|
|
// work. The weight tile is double-buffered: the next k-step's dequant
|
|
// overlaps the current cooperative matmul, so the k-loop needs one
|
|
// threadgroup barrier per step instead of two.
|
|
template<typename block_q, short qnl,
|
|
void (*dequantize_func)(device const block_q *, short, thread half4x4 &),
|
|
short NR1>
|
|
kernel void kernel_attn_out_low_mpp_direct_rhs(
|
|
constant ds4_metal_args_mul_mm_id & args,
|
|
device const char * srcA,
|
|
device const char * srcB,
|
|
device char * dst,
|
|
threadgroup char * shmem [[threadgroup(0)]],
|
|
uint3 tgpig [[threadgroup_position_in_grid]],
|
|
ushort tiitg [[thread_index_in_threadgroup]],
|
|
ushort sgitg [[simdgroup_index_in_threadgroup]]) {
|
|
(void) sgitg;
|
|
|
|
constexpr int NR0 = 64;
|
|
constexpr int NK = 32;
|
|
constexpr int NL = NK/16;
|
|
constexpr int NUM_THREADS = 128;
|
|
|
|
const int K = args.ne00;
|
|
const int M = args.ne0;
|
|
const int N = args.ne21;
|
|
const int G = args.ne1;
|
|
const int group = tgpig.z;
|
|
const int r0 = tgpig.y*NR0;
|
|
const int r1 = tgpig.x*NR1;
|
|
const bool full_tile = r0 + NR0 <= M && r1 + NR1 <= N && (K % NK) == 0;
|
|
|
|
threadgroup half *sa = (threadgroup half *)shmem;
|
|
auto tA0 = tensor(sa, dextents<int32_t, 2>(NK, NR0));
|
|
auto tA1 = tensor(sa + NR0*NK, dextents<int32_t, 2>(NK, NR0));
|
|
|
|
device float *ptrB = (device float *)(srcB + args.nb11*group);
|
|
const int strideB = args.nb12/sizeof(float);
|
|
auto tB = tensor(ptrB, dextents<int32_t, 2>(K, N), array<int, 2>({1, strideB}));
|
|
|
|
matmul2d<
|
|
matmul2d_descriptor(NR1, NR0, NK, false, true, true,
|
|
matmul2d_descriptor::mode::multiply_accumulate),
|
|
execution_simdgroups<4>> mm;
|
|
|
|
auto cT = mm.template get_destination_cooperative_tensor<decltype(tB), decltype(tA0), float>();
|
|
|
|
#pragma unroll
|
|
for (uint16_t i = 0; i < cT.get_capacity(); ++i) {
|
|
if (cT.is_valid_element(i)) {
|
|
cT[i] = 0.0f;
|
|
}
|
|
}
|
|
|
|
auto stage_tile = [&](const int loop_k, threadgroup half *buf) {
|
|
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 (full_tile) {
|
|
device const block_q *row_ptr =
|
|
(device const block_q *)(srcA + args.nb01*(r0 + row) + group*args.nb02);
|
|
|
|
half4x4 temp_a;
|
|
dequantize_func(row_ptr + k_pos/(16*qnl), (k_pos/16)%qnl, temp_a);
|
|
threadgroup half4 *dst4 = (threadgroup half4 *)(buf + row*NK + k_base);
|
|
dst4[0] = temp_a[0];
|
|
dst4[1] = temp_a[1];
|
|
dst4[2] = temp_a[2];
|
|
dst4[3] = temp_a[3];
|
|
} else if (r0 + row < M) {
|
|
const int block_idx = k_pos/(16*qnl);
|
|
const short il = (k_pos/16)%qnl;
|
|
device const block_q *row_ptr =
|
|
(device const block_q *)(srcA + args.nb01*(r0 + row) + group*args.nb02);
|
|
|
|
half4x4 temp_a;
|
|
dequantize_func(row_ptr + block_idx, il, temp_a);
|
|
FOR_UNROLL (short i = 0; i < 16; i++) {
|
|
buf[row*NK + k_base + i] = (k_pos + i < K) ? temp_a[i/4][i%4] : (half)0;
|
|
}
|
|
} else {
|
|
FOR_UNROLL (short i = 0; i < 16; i++) {
|
|
buf[row*NK + k_base + i] = (half)0;
|
|
}
|
|
}
|
|
}
|
|
};
|
|
|
|
stage_tile(0, sa);
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
uint buf_sel = 0;
|
|
for (int loop_k = 0; loop_k < K; loop_k += NK) {
|
|
auto mA = (buf_sel ? tA1 : tA0).slice(0, 0);
|
|
auto mB = tB.slice(loop_k, r1);
|
|
mm.run(mB, mA, cT);
|
|
|
|
const int next_k = loop_k + NK;
|
|
if (next_k < K) {
|
|
buf_sel ^= 1u;
|
|
stage_tile(next_k, buf_sel ? sa + NR0*NK : sa);
|
|
}
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
}
|
|
|
|
device float *dst_group = (device float *)dst + group*M;
|
|
if (full_tile) {
|
|
device float *dst_tile = dst_group + r0 + (uint64_t)r1*G*M;
|
|
auto tD = tensor(dst_tile, dextents<int32_t, 2>(NR0, NR1), array<int, 2>({1, G*M}));
|
|
cT.store(tD);
|
|
} else {
|
|
auto tD = tensor(dst_group, dextents<int32_t, 2>(M, N), array<int, 2>({1, G*M}));
|
|
auto mD = tD.slice(r0, r1);
|
|
cT.store(mD);
|
|
}
|
|
}
|
|
|
|
// Routed-expert grouped matmul on the Metal4 TensorOps/MPP pipeline. The
|
|
// barrier after mm.run prevents the next K iteration from replacing staged
|
|
// tiles while the cooperative matmul still reads them.
|
|
template<typename S0, typename S0_4x4, typename S0_8x8, typename S1, typename S1_2x4, typename S1_8x8, typename block_q, short nl, void (*dequantize_func)(device const block_q *, short, thread S0_4x4 &), typename T0, typename T0_4x4, typename T1, typename T1_2x4>
|
|
kernel void kernel_mul_mm_id_mpp(
|
|
constant ds4_metal_args_mul_mm_id & args,
|
|
device const char * src0,
|
|
device const char * src1,
|
|
device const char * htpe,
|
|
device const char * hids,
|
|
device char * dst,
|
|
device const char * work,
|
|
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]]) {
|
|
threadgroup S0 * sa = (threadgroup S0 *)(shmem);
|
|
threadgroup S1 * sb = (threadgroup S1 *)(shmem + 4096);
|
|
threadgroup float *sc = (threadgroup float *)shmem;
|
|
|
|
constexpr int NR0 = 64;
|
|
constexpr int NR1 = 32;
|
|
constexpr int NK = 32;
|
|
constexpr int NL0 = NK/16;
|
|
constexpr int NL1 = NK/8;
|
|
|
|
device const uint32_t *work_count = (device const uint32_t *)work;
|
|
const uint32_t work_index = tgpig.x;
|
|
if (work_index >= work_count[0]) {
|
|
return;
|
|
}
|
|
device const uint2 *work_items = (device const uint2 *)(work + 8);
|
|
const uint2 item = work_items[work_index];
|
|
const int im = (int)item.x;
|
|
const int r0 = tgpig.y*NR0;
|
|
const int r1 = (int)item.y;
|
|
|
|
device const uint32_t * tpe_u32 = (device const uint32_t *) (htpe);
|
|
device const int32_t * ids_i32 = (device const int32_t *) (hids);
|
|
|
|
const int32_t neh1 = tpe_u32[im];
|
|
|
|
if (r1 >= neh1) {
|
|
return;
|
|
}
|
|
|
|
const short nr0 = (args.ne0 - r0 < NR0) ? (args.ne0 - r0) : NR0;
|
|
const short nr1 = ( neh1 - r1 < NR1) ? ( neh1 - r1) : NR1;
|
|
|
|
if (!ds4_tp_owns_expert(im, args.ne02, args.tp_rank, args.tp_world)) {
|
|
for (short j = sgitg; j < nr1; j += 4) {
|
|
const int idj = ids_i32[im*args.ne21 + r1 + j];
|
|
const short ide = idj % args.ne20;
|
|
const short idt = idj / args.ne20;
|
|
device float *D = (device float *)dst + r0 + ide*args.ne0 +
|
|
idt*args.ne1*args.ne0;
|
|
for (int i = tiisg; i < nr0; i += 32) D[i] = 0.0f;
|
|
}
|
|
return;
|
|
}
|
|
|
|
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 id = ids_i32[im*args.ne21 + r1 + lr1];
|
|
|
|
const short i11 = (id % args.ne20) % args.ne11;
|
|
const short i12 = (id / args.ne20);
|
|
const short i13 = 0;
|
|
|
|
const uint64_t offset0 =
|
|
(uint64_t)(im - args.tp_expert_base)*args.nb02 + i13*args.nb03;
|
|
const short offset1 = il0/nl;
|
|
|
|
device const block_q * x = (device const block_q *)(src0 + args.nb01*(r0 + lr0) + offset0) + offset1;
|
|
|
|
const short iy = 8*(tiitg % NL1);
|
|
|
|
device const T1 * y = (device const T1 *)(src1
|
|
+ args.nb13*i13
|
|
+ args.nb12*i12
|
|
+ args.nb11*i11
|
|
+ args.nb10*iy);
|
|
|
|
auto tA = tensor(sa, dextents<int32_t, 2>(NK, NR0));
|
|
auto tB = tensor(sb, dextents<int32_t, 2>(NR1, NK));
|
|
|
|
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 < args.ne00; loop_k += NK) {
|
|
if (is_same<T0_4x4, block_q>::value && FC_mul_mm_bc_inp) {
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
for (short i = 0; i < 16; i++) {
|
|
const short sx = 2*il0 + i/8;
|
|
const short sy = (tiitg/NL0)/8;
|
|
const short lx = i%8;
|
|
const short ly = (tiitg/NL0)%8;
|
|
|
|
*(sa + NK*(8*sy + ly) + 8*sx + lx) =
|
|
loop_k + 16*il + i < args.ne00 ? *((device T0 *) x + i) : 0;
|
|
}
|
|
} else {
|
|
S0_4x4 temp_a;
|
|
dequantize_func(x, il, temp_a);
|
|
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
FOR_UNROLL (short i = 0; i < 16; i++) {
|
|
const short sx = 2*il0 + i/8;
|
|
const short sy = (tiitg/NL0)/8;
|
|
const short lx = i%8;
|
|
const short ly = (tiitg/NL0)%8;
|
|
|
|
*(sa + NK*(8*sy + ly) + 8*sx + lx) = temp_a[i/4][i%4];
|
|
}
|
|
}
|
|
|
|
if (FC_mul_mm_bc_inp) {
|
|
for (short i = 0; i < 8; ++i) {
|
|
const short sx = (tiitg%NL1);
|
|
const short sy = (tiitg/NL1)/8;
|
|
const short lx = i;
|
|
const short ly = (tiitg/NL1)%8;
|
|
|
|
*(sb + NK*(8*sy + ly) + 8*sx + lx) =
|
|
loop_k + iy + i < args.ne00 ? (S1) *((device T1 *) y + i) : 0;
|
|
}
|
|
} else {
|
|
const short sx = (tiitg%NL1);
|
|
const short sy = (tiitg/NL1)/8;
|
|
const short ly = (tiitg/NL1)%8;
|
|
|
|
*(threadgroup S1_2x4 *)(sb + NK*(8*sy + ly) + 8*sx) =
|
|
(S1_2x4)(*((device T1_2x4 *) y));
|
|
}
|
|
|
|
il = (il + 2 < nl) ? il + 2 : il % 2;
|
|
x = (il < 2) ? x + (2 + nl - 1)/nl : x;
|
|
|
|
y += NK;
|
|
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
auto sA = tA.slice(0, 0);
|
|
auto sB = tB.slice(0, 0);
|
|
mm.run(sB, sA, cT);
|
|
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
}
|
|
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
auto tC = tensor(sc, dextents<int32_t, 2>(NR0, NR1));
|
|
cT.store(tC);
|
|
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
for (short j = tiitg/32; j < nr1; j += 4) {
|
|
const int idj = ids_i32[im*args.ne21 + r1 + j];
|
|
|
|
const short ide = idj % args.ne20;
|
|
const short idt = idj / args.ne20;
|
|
|
|
device float * D = (device float *) dst + r0 + ide*args.ne0 + idt*args.ne1*args.ne0;
|
|
device float4 * D4 = (device float4 *) D;
|
|
|
|
threadgroup float * C = (threadgroup float *) shmem + j*NR0;
|
|
threadgroup float4 * C4 = (threadgroup float4 *) C;
|
|
|
|
int i = tiisg;
|
|
for (; i < nr0/4; i += 32) {
|
|
*(D4 + i) = *(C4 + i);
|
|
}
|
|
|
|
i = (4*(nr0/4)) + tiisg;
|
|
for (; i < nr0; i += 32) {
|
|
*(D + i) = *(C + i);
|
|
}
|
|
}
|
|
}
|
|
|
|
|
|
|
|
typedef decltype(kernel_mul_mm_id_mpp<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q2_K, QK_NL, dequantize_q2_K, float, float4x4, float, float2x4>) mul_mm_id_mpp_t;
|
|
typedef decltype(kernel_mul_mm_id_mpp<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q2_K, QK_NL, dequantize_q2_K, half, half4x4, half, half2x4>) mul_mm_id_mpp_f16_rhs_t;
|
|
|
|
template [[host_name("kernel_mul_mm_id_iq2_xxs_f32_mpp")]] kernel mul_mm_id_mpp_t kernel_mul_mm_id_mpp<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_iq2_xxs, QK_NL, dequantize_iq2_xxs, float, float4x4, float, float2x4>;
|
|
template [[host_name("kernel_mul_mm_id_q2_K_f16_mpp")]] kernel mul_mm_id_mpp_f16_rhs_t kernel_mul_mm_id_mpp<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q2_K, QK_NL, dequantize_q2_K, half, half4x4, half, half2x4>;
|
|
template [[host_name("kernel_mul_mm_id_iq2_xxs_f16_mpp")]] kernel mul_mm_id_mpp_f16_rhs_t kernel_mul_mm_id_mpp<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_iq2_xxs, QK_NL, dequantize_iq2_xxs, half, half4x4, half, half2x4>;
|
|
|
|
typedef decltype(kernel_attn_out_low_mpp_direct_rhs<
|
|
block_q8_0, 2, dequantize_q8_0_pairs, 64>)
|
|
attn_out_low_q8_0_mpp_direct_rhs_n64_t;
|
|
typedef decltype(kernel_attn_out_low_mpp_direct_rhs<
|
|
block_q4_K, QK_NL, dequantize_q4_K, 64>)
|
|
attn_out_low_q4_K_mpp_direct_rhs_n64_t;
|
|
|
|
template [[host_name("kernel_attn_out_low_q8_0_mpp_direct_rhs_n64")]]
|
|
kernel attn_out_low_q8_0_mpp_direct_rhs_n64_t
|
|
kernel_attn_out_low_mpp_direct_rhs<
|
|
block_q8_0, 2, dequantize_q8_0_pairs, 64>;
|
|
template [[host_name("kernel_attn_out_low_q4_K_mpp_direct_rhs_n64")]]
|
|
kernel attn_out_low_q4_K_mpp_direct_rhs_n64_t
|
|
kernel_attn_out_low_mpp_direct_rhs<
|
|
block_q4_K, QK_NL, dequantize_q4_K, 64>;
|
|
|
|
#endif
|
|
|
|
#undef QK_NL
|
|
#undef kmask_iq2xs
|
|
#undef ksigns_iq2xs
|
|
#undef iq2xxs_grid
|
|
#undef QK_K
|
|
#undef N_R0_Q2_K
|
|
#undef N_R0_Q4_K
|
|
#undef N_R0_Q8_K
|
|
#undef N_R0_GLM_Q4_PAIR2_K
|
|
#undef N_R0_GLM_Q4_PAIR_K
|
|
#undef N_R0_Q5_PAIR_K
|
|
#undef N_R0_Q5_K
|
|
#undef N_R0_Q6_K
|
|
#undef N_R0_IQ2_XXS
|