386 lines
15 KiB
Metal
386 lines
15 KiB
Metal
struct ds4_metal_args_dsv4_rope_tail {
|
|
int64_t ne00;
|
|
int64_t ne01;
|
|
int64_t ne02;
|
|
int64_t ne03;
|
|
uint64_t nb00;
|
|
uint64_t nb01;
|
|
uint64_t nb02;
|
|
uint64_t nb03;
|
|
uint64_t nb0;
|
|
uint64_t nb1;
|
|
uint64_t nb2;
|
|
uint64_t nb3;
|
|
int32_t n_dims;
|
|
int32_t mode;
|
|
int32_t n_ctx_orig;
|
|
int32_t inverse;
|
|
float freq_base;
|
|
float freq_scale;
|
|
float ext_factor;
|
|
float attn_factor;
|
|
float beta_fast;
|
|
float beta_slow;
|
|
bool src2;
|
|
};
|
|
|
|
struct ds4_metal_args_dsv4_rope_affine_pair {
|
|
uint64_t row_bytes;
|
|
uint64_t token_bytes;
|
|
int32_t head_dim;
|
|
int32_t n_dims;
|
|
int32_t n_ctx_orig;
|
|
int32_t inverse;
|
|
uint32_t pos0;
|
|
uint32_t pos_step;
|
|
float freq_base;
|
|
float freq_scale;
|
|
float ext_factor;
|
|
float attn_factor;
|
|
float beta_fast;
|
|
float beta_slow;
|
|
};
|
|
|
|
static float rope_yarn_ramp(const float low, const float high, const int i0) {
|
|
const float y = (i0 / 2 - low) / max(0.001f, high - low);
|
|
return 1.0f - min(1.0f, max(0.0f, y));
|
|
}
|
|
|
|
// YaRN algorithm based on LlamaYaRNScaledRotaryEmbedding.py from https://github.com/jquesnelle/yarn
|
|
// MIT licensed. Copyright (c) 2023 Jeffrey Quesnelle and Bowen Peng.
|
|
static void rope_yarn(
|
|
float theta_extrap, float freq_scale, float corr_dims[2], int i0, float ext_factor, float mscale,
|
|
thread float * cos_theta, thread float * sin_theta) {
|
|
// Get n-d rotational scaling corrected for extrapolation
|
|
float theta_interp = freq_scale * theta_extrap;
|
|
float theta = theta_interp;
|
|
if (ext_factor != 0.0f) {
|
|
float ramp_mix = rope_yarn_ramp(corr_dims[0], corr_dims[1], i0) * ext_factor;
|
|
theta = theta_interp * (1 - ramp_mix) + theta_extrap * ramp_mix;
|
|
|
|
// Get n-d magnitude scaling corrected for interpolation
|
|
mscale *= 1.0f + 0.1f * log(1.0f / freq_scale);
|
|
}
|
|
*cos_theta = cos(theta) * mscale;
|
|
*sin_theta = sin(theta) * mscale;
|
|
}
|
|
|
|
// Apparently solving `n_rot = 2pi * x * base^((2 * max_pos_emb) / n_dims)` for x, we get
|
|
// `corr_fac(n_rot) = n_dims * log(max_pos_emb / (n_rot * 2pi)) / (2 * log(base))`
|
|
static float rope_yarn_corr_factor(int n_dims, int n_ctx_orig, float n_rot, float base) {
|
|
return n_dims * log(n_ctx_orig / (n_rot * 2 * M_PI_F)) / (2 * log(base));
|
|
}
|
|
|
|
static void rope_yarn_corr_dims(
|
|
int n_dims, int n_ctx_orig, float freq_base, float beta_fast, float beta_slow, float dims[2]
|
|
) {
|
|
// start and end correction dims
|
|
dims[0] = max(0.0f, floor(rope_yarn_corr_factor(n_dims, n_ctx_orig, beta_fast, freq_base)));
|
|
dims[1] = min(n_dims - 1.0f, ceil(rope_yarn_corr_factor(n_dims, n_ctx_orig, beta_slow, freq_base)));
|
|
}
|
|
|
|
// Applies DeepSeek V4's partial RoPE: the no-position prefix is copied and only
|
|
// the rotated tail is transformed. This is used for Q/K after their projections
|
|
// and before writing/reading the attention KV state.
|
|
kernel void kernel_dsv4_rope_tail_f32(
|
|
constant ds4_metal_args_dsv4_rope_tail & args,
|
|
device const char * src0,
|
|
device const char * src1,
|
|
device const char * src2,
|
|
device char * dst,
|
|
uint tid [[thread_index_in_threadgroup]],
|
|
ushort3 ntg [[threads_per_threadgroup]],
|
|
uint3 tgpig [[threadgroup_position_in_grid]]) {
|
|
const int i1 = tgpig[0];
|
|
const int i2 = tgpig[1];
|
|
const int i3 = tgpig[2];
|
|
|
|
const int n_nope = args.ne00 - args.n_dims;
|
|
if (n_nope < 0) {
|
|
return;
|
|
}
|
|
|
|
device const int32_t * pos = (device const int32_t *) src1;
|
|
|
|
float corr_dims[2];
|
|
rope_yarn_corr_dims(args.n_dims, args.n_ctx_orig, args.freq_base, args.beta_fast, args.beta_slow, corr_dims);
|
|
|
|
const float theta_base = (float) pos[i2];
|
|
const float inv_ndims = -1.f/args.n_dims;
|
|
const bool is_neox = args.mode == 2;
|
|
|
|
for (int i0 = tid; i0 < args.ne00; i0 += ntg.x) {
|
|
device const char * src_base = src0 + i3*args.nb03 + i2*args.nb02 + i1*args.nb01;
|
|
device char * dst_base = dst + i3*args.nb3 + i2*args.nb2 + i1*args.nb1;
|
|
|
|
if (i0 < n_nope) {
|
|
*((device float *) (dst_base + i0*args.nb0)) = *((device const float *) (src_base + i0*args.nb00));
|
|
continue;
|
|
}
|
|
|
|
const int r = i0 - n_nope;
|
|
if (is_neox) {
|
|
const int n_half = args.n_dims/2;
|
|
if (r >= n_half) {
|
|
continue;
|
|
}
|
|
|
|
const int ic = r;
|
|
const int rel_i0 = 2*ic;
|
|
#ifdef DS4_METAL_ROPE_EXP2_LOG2
|
|
// Equivalent to pow(freq_base, k) but expressed through IEEE-754
|
|
// primitives that have tighter precision guarantees than Metal's pow().
|
|
const float theta = theta_base * exp2(inv_ndims * (float)rel_i0 * log2(args.freq_base));
|
|
#else
|
|
const float theta = theta_base * pow(args.freq_base, inv_ndims*rel_i0);
|
|
#endif
|
|
const float freq_factor = args.src2 ? ((device const float *) src2)[ic] : 1.0f;
|
|
|
|
float cos_theta;
|
|
float sin_theta;
|
|
rope_yarn(theta/freq_factor, args.freq_scale, corr_dims, rel_i0, args.ext_factor, args.attn_factor, &cos_theta, &sin_theta);
|
|
if (args.inverse) {
|
|
sin_theta = -sin_theta;
|
|
}
|
|
|
|
const int j0 = n_nope + ic;
|
|
const int j1 = n_nope + ic + n_half;
|
|
const float x0 = *((device const float *) (src_base + j0*args.nb00));
|
|
const float x1 = *((device const float *) (src_base + j1*args.nb00));
|
|
|
|
*((device float *) (dst_base + j0*args.nb0)) = x0*cos_theta - x1*sin_theta;
|
|
*((device float *) (dst_base + j1*args.nb0)) = x0*sin_theta + x1*cos_theta;
|
|
} else {
|
|
if ((r & 1) != 0) {
|
|
continue;
|
|
}
|
|
|
|
const int ic = r/2;
|
|
#ifdef DS4_METAL_ROPE_EXP2_LOG2
|
|
const float theta = theta_base * exp2(inv_ndims * (float)r * log2(args.freq_base));
|
|
#else
|
|
const float theta = theta_base * pow(args.freq_base, inv_ndims*r);
|
|
#endif
|
|
const float freq_factor = args.src2 ? ((device const float *) src2)[ic] : 1.0f;
|
|
|
|
float cos_theta;
|
|
float sin_theta;
|
|
rope_yarn(theta/freq_factor, args.freq_scale, corr_dims, r, args.ext_factor, args.attn_factor, &cos_theta, &sin_theta);
|
|
if (args.inverse) {
|
|
sin_theta = -sin_theta;
|
|
}
|
|
|
|
const int j0 = n_nope + r;
|
|
const int j1 = j0 + 1;
|
|
const float x0 = *((device const float *) (src_base + j0*args.nb00));
|
|
const float x1 = *((device const float *) (src_base + j1*args.nb00));
|
|
|
|
*((device float *) (dst_base + j0*args.nb0)) = x0*cos_theta - x1*sin_theta;
|
|
*((device float *) (dst_base + j1*args.nb0)) = x0*sin_theta + x1*cos_theta;
|
|
}
|
|
}
|
|
}
|
|
|
|
// DS4 only calls the Metal RoPE helper in-place and uses the adjacent-pair
|
|
// layout (mode 0). The generic kernel above still copies the no-position
|
|
// prefix, even though source and destination alias, and dispatches enough
|
|
// lanes for the full head. This specialization maps lanes directly to the
|
|
// rotated pairs and deliberately never reads or writes the unchanged prefix.
|
|
// Keep the arithmetic below in the same order as the mode-0 branch above so
|
|
// the optimized and reference paths remain bit-identical.
|
|
kernel void kernel_dsv4_rope_tail_f32_inplace_pair(
|
|
constant ds4_metal_args_dsv4_rope_tail & args,
|
|
device const char * src0,
|
|
device const char * src1,
|
|
device const char * src2,
|
|
device char * dst,
|
|
uint tid [[thread_index_in_threadgroup]],
|
|
ushort3 ntg [[threads_per_threadgroup]],
|
|
uint3 tgpig [[threadgroup_position_in_grid]]) {
|
|
if (args.mode != 0) {
|
|
return;
|
|
}
|
|
|
|
const int i1 = tgpig[0];
|
|
const int i2 = tgpig[1];
|
|
const int i3 = tgpig[2];
|
|
const int n_nope = args.ne00 - args.n_dims;
|
|
if (n_nope < 0) {
|
|
return;
|
|
}
|
|
|
|
device const int32_t * pos = (device const int32_t *) src1;
|
|
|
|
float corr_dims[2];
|
|
rope_yarn_corr_dims(args.n_dims, args.n_ctx_orig, args.freq_base, args.beta_fast, args.beta_slow, corr_dims);
|
|
|
|
const float theta_base = (float) pos[i2];
|
|
const float inv_ndims = -1.f/args.n_dims;
|
|
device const char * src_base = src0 + i3*args.nb03 + i2*args.nb02 + i1*args.nb01;
|
|
device char * dst_base = dst + i3*args.nb3 + i2*args.nb2 + i1*args.nb1;
|
|
|
|
// Keep each pair on the same SIMD lane as the generic in-place dispatch.
|
|
// n_nope is 32-aligned for the supported DS4 shapes, so logical tail index
|
|
// r ran on lane r%32 in the reference kernel. Compacting pairs would move
|
|
// them to lane (r/2)%32 and can perturb fast-math code generation.
|
|
for (int r = tid; r < args.n_dims; r += ntg.x) {
|
|
if ((r & 1) != 0) {
|
|
continue;
|
|
}
|
|
const int ic = r/2;
|
|
#ifdef DS4_METAL_ROPE_EXP2_LOG2
|
|
const float theta = theta_base * exp2(inv_ndims * (float)r * log2(args.freq_base));
|
|
#else
|
|
const float theta = theta_base * pow(args.freq_base, inv_ndims*r);
|
|
#endif
|
|
const float freq_factor = args.src2 ? ((device const float *) src2)[ic] : 1.0f;
|
|
|
|
float cos_theta;
|
|
float sin_theta;
|
|
rope_yarn(theta/freq_factor, args.freq_scale, corr_dims, r, args.ext_factor, args.attn_factor, &cos_theta, &sin_theta);
|
|
if (args.inverse) {
|
|
sin_theta = -sin_theta;
|
|
}
|
|
|
|
const int j0 = n_nope + r;
|
|
const int j1 = j0 + 1;
|
|
const float x0 = *((device const float *) (src_base + j0*args.nb00));
|
|
const float x1 = *((device const float *) (src_base + j1*args.nb00));
|
|
|
|
*((device float *) (dst_base + j0*args.nb0)) = x0*cos_theta - x1*sin_theta;
|
|
*((device float *) (dst_base + j1*args.nb0)) = x0*sin_theta + x1*cos_theta;
|
|
}
|
|
}
|
|
|
|
// The Q/K RoPE calls use the same position and scaling parameters for every
|
|
// head. Group four heads into one threadgroup so one 64-thread cohort computes
|
|
// the 32 adjacent-pair coefficients and the other cohorts reuse them. Keeping
|
|
// r on the same r%32 SIMD lane as the per-head specialization preserves the
|
|
// fast-math instruction mapping; only the redundant coefficient work changes.
|
|
kernel void kernel_dsv4_rope_tail_f32_inplace_pair_shared4(
|
|
constant ds4_metal_args_dsv4_rope_tail & args,
|
|
device const char * src0,
|
|
device const char * src1,
|
|
device const char * src2,
|
|
device char * dst,
|
|
uint tid [[thread_index_in_threadgroup]],
|
|
uint3 tgpig [[threadgroup_position_in_grid]]) {
|
|
if (args.mode != 0 || args.n_dims != 64) {
|
|
return;
|
|
}
|
|
|
|
const int n_nope = args.ne00 - args.n_dims;
|
|
if (n_nope < 0) {
|
|
return;
|
|
}
|
|
|
|
const uint cohort = tid >> 6;
|
|
const int r = (int)(tid & 63u);
|
|
threadgroup float cos_shared[32];
|
|
threadgroup float sin_shared[32];
|
|
|
|
device const int32_t * pos = (device const int32_t *) src1;
|
|
float corr_dims[2];
|
|
rope_yarn_corr_dims(args.n_dims, args.n_ctx_orig, args.freq_base, args.beta_fast, args.beta_slow, corr_dims);
|
|
|
|
const int i2 = tgpig[1];
|
|
const float theta_base = (float) pos[i2];
|
|
const float inv_ndims = -1.f/args.n_dims;
|
|
|
|
if (cohort == 0 && (r & 1) == 0) {
|
|
const int ic = r/2;
|
|
#ifdef DS4_METAL_ROPE_EXP2_LOG2
|
|
const float theta = theta_base * exp2(inv_ndims * (float)r * log2(args.freq_base));
|
|
#else
|
|
const float theta = theta_base * pow(args.freq_base, inv_ndims*r);
|
|
#endif
|
|
const float freq_factor = args.src2 ? ((device const float *) src2)[ic] : 1.0f;
|
|
|
|
float cos_theta;
|
|
float sin_theta;
|
|
rope_yarn(theta/freq_factor, args.freq_scale, corr_dims, r, args.ext_factor, args.attn_factor, &cos_theta, &sin_theta);
|
|
if (args.inverse) {
|
|
sin_theta = -sin_theta;
|
|
}
|
|
cos_shared[ic] = cos_theta;
|
|
sin_shared[ic] = sin_theta;
|
|
}
|
|
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
const int i1 = (int)tgpig[0]*4 + (int)cohort;
|
|
if (i1 >= args.ne01 || (r & 1) != 0) {
|
|
return;
|
|
}
|
|
|
|
const int i3 = tgpig[2];
|
|
device const char * src_base = src0 + i3*args.nb03 + i2*args.nb02 + i1*args.nb01;
|
|
device char * dst_base = dst + i3*args.nb3 + i2*args.nb2 + i1*args.nb1;
|
|
const int j0 = n_nope + r;
|
|
const int j1 = j0 + 1;
|
|
const float x0 = *((device const float *) (src_base + j0*args.nb00));
|
|
const float x1 = *((device const float *) (src_base + j1*args.nb00));
|
|
const float cos_theta = cos_shared[r/2];
|
|
const float sin_theta = sin_shared[r/2];
|
|
|
|
*((device float *) (dst_base + j0*args.nb0)) = x0*cos_theta - x1*sin_theta;
|
|
*((device float *) (dst_base + j1*args.nb0)) = x0*sin_theta + x1*cos_theta;
|
|
}
|
|
|
|
// DS4 positions are always affine within one RoPE dispatch. This variant
|
|
// reconstructs the same wrapped int32 position in-kernel, avoiding the host
|
|
// position array and its buffer binding while preserving the pair lane mapping
|
|
// and all floating-point operations of the specialization above.
|
|
kernel void kernel_dsv4_rope_tail_f32_inplace_pair_affine(
|
|
constant ds4_metal_args_dsv4_rope_affine_pair & args [[buffer(0)]],
|
|
device const char * src0 [[buffer(1)]],
|
|
device char * dst [[buffer(4)]],
|
|
uint tid [[thread_index_in_threadgroup]],
|
|
ushort3 ntg [[threads_per_threadgroup]],
|
|
uint3 tgpig [[threadgroup_position_in_grid]]) {
|
|
const int i1 = tgpig[0];
|
|
const int i2 = tgpig[1];
|
|
const int n_nope = args.head_dim - args.n_dims;
|
|
if (n_nope < 0) {
|
|
return;
|
|
}
|
|
|
|
float corr_dims[2];
|
|
rope_yarn_corr_dims(args.n_dims, args.n_ctx_orig, args.freq_base, args.beta_fast, args.beta_slow, corr_dims);
|
|
|
|
const uint raw_pos = args.pos0 + (uint)i2 * args.pos_step;
|
|
const float theta_base = (float)as_type<int>(raw_pos);
|
|
const float inv_ndims = -1.f/args.n_dims;
|
|
device const char * src_base =
|
|
src0 + (uint64_t)i2*args.token_bytes + (uint64_t)i1*args.row_bytes;
|
|
device char * dst_base =
|
|
dst + (uint64_t)i2*args.token_bytes + (uint64_t)i1*args.row_bytes;
|
|
|
|
for (int r = tid; r < args.n_dims; r += ntg.x) {
|
|
if ((r & 1) != 0) {
|
|
continue;
|
|
}
|
|
#ifdef DS4_METAL_ROPE_EXP2_LOG2
|
|
const float theta = theta_base * exp2(inv_ndims * (float)r * log2(args.freq_base));
|
|
#else
|
|
const float theta = theta_base * pow(args.freq_base, inv_ndims*r);
|
|
#endif
|
|
const float freq_factor = 1.0f;
|
|
|
|
float cos_theta;
|
|
float sin_theta;
|
|
rope_yarn(theta/freq_factor, args.freq_scale, corr_dims, r, args.ext_factor, args.attn_factor, &cos_theta, &sin_theta);
|
|
if (args.inverse) {
|
|
sin_theta = -sin_theta;
|
|
}
|
|
|
|
const int j0 = n_nope + r;
|
|
const int j1 = j0 + 1;
|
|
const float x0 = *((device const float *) (src_base + j0*sizeof(float)));
|
|
const float x1 = *((device const float *) (src_base + j1*sizeof(float)));
|
|
|
|
*((device float *) (dst_base + j0*sizeof(float))) = x0*cos_theta - x1*sin_theta;
|
|
*((device float *) (dst_base + j1*sizeof(float))) = x0*sin_theta + x1*cos_theta;
|
|
}
|
|
}
|