695 lines
28 KiB
Metal
695 lines
28 KiB
Metal
// Qwen3.8 Flash Next primitives. Rust owns the graph and all state lifetimes;
|
|
// this file contains only the data-parallel kernels executed by Metal.
|
|
|
|
struct qwen_kernel_args {
|
|
uint u[16];
|
|
float f[8];
|
|
};
|
|
|
|
static inline float qwen_bf16(ushort value) {
|
|
return as_type<float>((uint)value << 16);
|
|
}
|
|
|
|
static inline ushort qwen_to_bf16(float value) {
|
|
uint bits = as_type<uint>(value);
|
|
bits += 0x7fffu + ((bits >> 16) & 1u);
|
|
return (ushort)(bits >> 16);
|
|
}
|
|
|
|
static inline ushort qwen_weight_u16(
|
|
device const uchar *data,
|
|
uint byte_offset,
|
|
ulong index) {
|
|
const ulong byte = (ulong)byte_offset + index * 2u;
|
|
if ((byte & 1u) == 0u) return *((device const ushort *)(data + byte));
|
|
return (ushort)data[byte] | ((ushort)data[byte + 1u] << 8u);
|
|
}
|
|
|
|
static inline uint qwen_weight_u32(
|
|
device const uchar *data,
|
|
uint byte_offset,
|
|
ulong index) {
|
|
const ulong byte = (ulong)byte_offset + index * 4u;
|
|
if ((byte & 3u) == 0u) return *((device const uint *)(data + byte));
|
|
return (uint)data[byte] |
|
|
((uint)data[byte + 1u] << 8u) |
|
|
((uint)data[byte + 2u] << 16u) |
|
|
((uint)data[byte + 3u] << 24u);
|
|
}
|
|
|
|
static inline float qwen_quant_weight(
|
|
device const uchar *packed,
|
|
device const uchar *scales,
|
|
device const uchar *biases,
|
|
uint packed_offset,
|
|
uint scales_offset,
|
|
uint biases_offset,
|
|
uint row,
|
|
uint column,
|
|
uint in_dim,
|
|
uint bits,
|
|
uint group_size) {
|
|
const uint per_word = 32u / bits;
|
|
const uint packed_columns = in_dim / per_word;
|
|
const uint groups = in_dim / group_size;
|
|
const uint word = qwen_weight_u32(packed, packed_offset,
|
|
(ulong)row * packed_columns + column / per_word);
|
|
const uint mask = (1u << bits) - 1u;
|
|
const uint quant = (word >> ((column % per_word) * bits)) & mask;
|
|
const uint group = row * groups + column / group_size;
|
|
return fma((float)quant,
|
|
qwen_bf16(qwen_weight_u16(scales, scales_offset, group)),
|
|
qwen_bf16(qwen_weight_u16(biases, biases_offset, group)));
|
|
}
|
|
|
|
static inline float qwen_quant_value(
|
|
device const uint *packed,
|
|
device const ushort *scales,
|
|
device const ushort *biases,
|
|
uint row,
|
|
uint column,
|
|
uint in_dim,
|
|
uint bits,
|
|
uint group_size) {
|
|
const uint per_word = 32u / bits;
|
|
const uint packed_columns = in_dim / per_word;
|
|
const uint groups = in_dim / group_size;
|
|
const uint word = packed[(ulong)row * packed_columns + column / per_word];
|
|
const uint mask = (1u << bits) - 1u;
|
|
const uint quant = (word >> ((column % per_word) * bits)) & mask;
|
|
const uint group = row * groups + column / group_size;
|
|
return fma((float)quant, qwen_bf16(scales[group]), qwen_bf16(biases[group]));
|
|
}
|
|
|
|
kernel void kernel_qwen_affine_mv(
|
|
constant qwen_kernel_args &args [[buffer(0)]],
|
|
device float *out [[buffer(1)]],
|
|
device const float *x [[buffer(2)]],
|
|
device const uchar *packed [[buffer(5)]],
|
|
device const uchar *scales [[buffer(6)]],
|
|
device const uchar *biases [[buffer(7)]],
|
|
uint row [[thread_position_in_grid]]) {
|
|
const uint in_dim = args.u[0];
|
|
const uint out_dim = args.u[1];
|
|
if (row >= out_dim) return;
|
|
float sum = 0.0f;
|
|
for (uint column = 0; column < in_dim; column++) {
|
|
sum = fma(qwen_quant_weight(packed, scales, biases,
|
|
args.u[13], args.u[14], args.u[15], row, column,
|
|
in_dim, args.u[2], args.u[3]), x[column], sum);
|
|
}
|
|
out[row] = sum;
|
|
}
|
|
|
|
kernel void kernel_qwen_affine_embedding(
|
|
constant qwen_kernel_args &args [[buffer(0)]],
|
|
device float *out [[buffer(1)]],
|
|
device const uchar *packed [[buffer(5)]],
|
|
device const uchar *scales [[buffer(6)]],
|
|
device const uchar *biases [[buffer(7)]],
|
|
uint column [[thread_position_in_grid]]) {
|
|
if (column >= args.u[0]) return;
|
|
out[column] = qwen_quant_weight(packed, scales, biases,
|
|
args.u[13], args.u[14], args.u[15],
|
|
args.u[4], column, args.u[0], args.u[2], args.u[3]);
|
|
}
|
|
|
|
kernel void kernel_qwen_ple_dequant(
|
|
constant qwen_kernel_args &args [[buffer(0)]],
|
|
device float *out [[buffer(1)]],
|
|
device const uint *packed [[buffer(2)]],
|
|
device const ushort *scales [[buffer(3)]],
|
|
device const ushort *biases [[buffer(4)]],
|
|
uint index [[thread_position_in_grid]]) {
|
|
const uint dim = args.u[0];
|
|
if (index >= dim * args.u[1]) return;
|
|
const uint row = index / dim;
|
|
const uint column = index % dim;
|
|
out[index] = qwen_quant_value(packed, scales, biases, row, column,
|
|
dim, args.u[2], args.u[3]);
|
|
}
|
|
|
|
kernel void kernel_qwen_bf16_mv(
|
|
constant qwen_kernel_args &args [[buffer(0)]],
|
|
device float *out [[buffer(1)]],
|
|
device const float *x [[buffer(2)]],
|
|
device const uchar *weights [[buffer(5)]],
|
|
uint row [[thread_position_in_grid]]) {
|
|
if (row >= args.u[1]) return;
|
|
float sum = 0.0f;
|
|
for (uint column = 0; column < args.u[0]; column++) {
|
|
sum = fma(qwen_bf16(qwen_weight_u16(weights, args.u[13],
|
|
(ulong)row * args.u[0] + column)),
|
|
x[column], sum);
|
|
}
|
|
out[row] = sum;
|
|
}
|
|
|
|
kernel void kernel_qwen_repeat4(
|
|
constant qwen_kernel_args &args [[buffer(0)]],
|
|
device float *out [[buffer(1)]],
|
|
device const float *x [[buffer(2)]],
|
|
uint index [[thread_position_in_grid]]) {
|
|
if (index < args.u[0] * 4u) out[index] = x[index % args.u[0]];
|
|
}
|
|
|
|
kernel void kernel_qwen_zero_rms(
|
|
constant qwen_kernel_args &args [[buffer(0)]],
|
|
device float *out [[buffer(1)]],
|
|
device const float *x [[buffer(2)]],
|
|
device const uchar *weight [[buffer(5)]],
|
|
uint group [[thread_position_in_grid]]) {
|
|
const uint width = args.u[0];
|
|
const uint group_size = args.u[1];
|
|
if (group >= width / group_size) return;
|
|
const uint start = group * group_size;
|
|
float variance = 0.0f;
|
|
for (uint i = 0; i < group_size; i++) variance = fma(x[start + i], x[start + i], variance);
|
|
const float scale = rsqrt(variance / (float)group_size + args.f[0]);
|
|
for (uint i = 0; i < group_size; i++) {
|
|
const uint index = start + i;
|
|
out[index] = x[index] * scale *
|
|
(1.0f + qwen_bf16(qwen_weight_u16(weight, args.u[13], index)));
|
|
}
|
|
}
|
|
|
|
kernel void kernel_qwen_silu_div4(
|
|
constant qwen_kernel_args &args [[buffer(0)]],
|
|
device float *out [[buffer(1)]],
|
|
device const float *x [[buffer(2)]],
|
|
uint index [[thread_position_in_grid]]) {
|
|
if (index >= args.u[0]) return;
|
|
const float value = x[index] * 0.25f;
|
|
out[index] = value / (1.0f + exp(-value));
|
|
}
|
|
|
|
kernel void kernel_qwen_sigmoid(
|
|
constant qwen_kernel_args &args [[buffer(0)]],
|
|
device float *out [[buffer(1)]],
|
|
device const float *x [[buffer(2)]],
|
|
uint index [[thread_position_in_grid]]) {
|
|
if (index < args.u[0]) out[index] = 1.0f / (1.0f + exp(-x[index]));
|
|
}
|
|
|
|
kernel void kernel_qwen_sigmoid2_div4(
|
|
constant qwen_kernel_args &args [[buffer(0)]],
|
|
device float *out [[buffer(1)]],
|
|
device const float *x [[buffer(2)]],
|
|
uint index [[thread_position_in_grid]]) {
|
|
if (index < args.u[0]) out[index] = 2.0f / (1.0f + exp(-x[index] * 0.25f));
|
|
}
|
|
|
|
kernel void kernel_qwen_hyper_mix(
|
|
constant qwen_kernel_args &args [[buffer(0)]],
|
|
device float *out [[buffer(1)]],
|
|
device const float *normalized [[buffer(2)]],
|
|
device const float *mix [[buffer(3)]],
|
|
uint index [[thread_position_in_grid]]) {
|
|
if (index >= args.u[0]) return;
|
|
float value = 0.0f;
|
|
for (uint stream = 0; stream < 4u; stream++) {
|
|
const uint offset = stream * args.u[0] + index;
|
|
value = fma(normalized[offset], mix[offset], value);
|
|
}
|
|
out[index] = value * 0.25f;
|
|
}
|
|
|
|
kernel void kernel_qwen_hyper_inject(
|
|
constant qwen_kernel_args &args [[buffer(0)]],
|
|
device float *out [[buffer(1)]],
|
|
device const float *residual [[buffer(2)]],
|
|
device const float *block [[buffer(3)]],
|
|
device const float *gate [[buffer(4)]],
|
|
uint index [[thread_position_in_grid]]) {
|
|
const uint hidden = args.u[0];
|
|
if (index >= hidden * 4u) return;
|
|
out[index] = residual[index] + block[index % hidden] * gate[index / hidden];
|
|
}
|
|
|
|
kernel void kernel_qwen_ple_gate(
|
|
constant qwen_kernel_args &args [[buffer(0)]],
|
|
device float *out [[buffer(1)]],
|
|
device const float *key [[buffer(2)]],
|
|
device const float *query [[buffer(3)]],
|
|
device const float *value [[buffer(4)]],
|
|
uint stream [[thread_position_in_grid]]) {
|
|
const uint hidden = args.u[0];
|
|
if (stream >= 4u) return;
|
|
const ulong base = (ulong)stream * hidden;
|
|
float score = 0.0f;
|
|
for (uint i = 0; i < hidden; i++) score = fma(key[base + i], query[base + i], score);
|
|
score *= rsqrt((float)hidden);
|
|
const float transformed = copysign(sqrt(max(abs(score), 1.0e-6f)), score);
|
|
const float gate = 1.0f / (1.0f + exp(-transformed));
|
|
for (uint i = 0; i < hidden; i++) out[base + i] = value[i] * gate;
|
|
}
|
|
|
|
kernel void kernel_qwen_ple_conv(
|
|
constant qwen_kernel_args &args [[buffer(0)]],
|
|
device float *out [[buffer(1)]],
|
|
device const float *gated [[buffer(2)]],
|
|
device const float *normalized [[buffer(3)]],
|
|
device ushort *state [[buffer(4)]],
|
|
device const uchar *weight [[buffer(5)]],
|
|
uint channel [[thread_position_in_grid]]) {
|
|
if (channel >= args.u[0]) return;
|
|
device ushort *history = state + (ulong)channel * 9u;
|
|
float value = fma(qwen_bf16(history[0]),
|
|
qwen_bf16(qwen_weight_u16(weight, args.u[13], (ulong)channel * 4u)),
|
|
fma(qwen_bf16(history[3]),
|
|
qwen_bf16(qwen_weight_u16(weight, args.u[13], (ulong)channel * 4u + 1u)),
|
|
fma(qwen_bf16(history[6]),
|
|
qwen_bf16(qwen_weight_u16(weight, args.u[13], (ulong)channel * 4u + 2u)),
|
|
normalized[channel] *
|
|
qwen_bf16(qwen_weight_u16(weight, args.u[13], (ulong)channel * 4u + 3u)))));
|
|
for (uint i = 0; i < 8u; i++) history[i] = history[i + 1u];
|
|
history[8] = qwen_to_bf16(normalized[channel]);
|
|
out[channel] = gated[channel] + value / (1.0f + exp(-value));
|
|
}
|
|
|
|
kernel void kernel_qwen_add(
|
|
constant qwen_kernel_args &args [[buffer(0)]],
|
|
device float *out [[buffer(1)]],
|
|
device const float *a [[buffer(2)]],
|
|
device const float *b [[buffer(3)]],
|
|
uint index [[thread_position_in_grid]]) {
|
|
if (index < args.u[0]) out[index] = a[index] + b[index];
|
|
}
|
|
|
|
kernel void kernel_qwen_conv_silu(
|
|
constant qwen_kernel_args &args [[buffer(0)]],
|
|
device float *out [[buffer(1)]],
|
|
device const float *x [[buffer(2)]],
|
|
device ushort *state [[buffer(3)]],
|
|
device const uchar *weight [[buffer(5)]],
|
|
uint channel [[thread_position_in_grid]]) {
|
|
if (channel >= args.u[0]) return;
|
|
device ushort *history = state + (ulong)channel * 3u;
|
|
float value = fma(qwen_bf16(history[0]),
|
|
qwen_bf16(qwen_weight_u16(weight, args.u[13], (ulong)channel * 4u)),
|
|
fma(qwen_bf16(history[1]),
|
|
qwen_bf16(qwen_weight_u16(weight, args.u[13], (ulong)channel * 4u + 1u)),
|
|
fma(qwen_bf16(history[2]),
|
|
qwen_bf16(qwen_weight_u16(weight, args.u[13], (ulong)channel * 4u + 2u)),
|
|
x[channel] *
|
|
qwen_bf16(qwen_weight_u16(weight, args.u[13], (ulong)channel * 4u + 3u)))));
|
|
history[0] = history[1];
|
|
history[1] = history[2];
|
|
history[2] = qwen_to_bf16(x[channel]);
|
|
out[channel] = value / (1.0f + exp(-value));
|
|
}
|
|
|
|
kernel void kernel_qwen_gdn_step(
|
|
constant qwen_kernel_args &args [[buffer(0)]],
|
|
device float *out [[buffer(1)]],
|
|
device const float *qkv [[buffer(2)]],
|
|
device const float *controls [[buffer(3)]],
|
|
device float *state [[buffer(4)]],
|
|
device const uchar *a_log [[buffer(5)]],
|
|
device const uchar *dt_bias [[buffer(6)]],
|
|
uint2 gid [[thread_position_in_grid]]) {
|
|
const uint value_index = gid.x;
|
|
const uint head = gid.y;
|
|
const uint dim = args.u[0];
|
|
const uint key_heads = args.u[1];
|
|
const uint value_heads = args.u[2];
|
|
if (value_index >= dim || head >= value_heads) return;
|
|
const uint key_head = head / (value_heads / key_heads);
|
|
device const float *q_raw = qkv + (ulong)key_head * dim;
|
|
device const float *k_raw = qkv + (ulong)key_heads * dim + (ulong)key_head * dim;
|
|
device const float *value = qkv + (ulong)key_heads * dim * 2u + (ulong)head * dim;
|
|
float qsum = 0.0f;
|
|
float ksum = 0.0f;
|
|
for (uint i = 0; i < dim; i++) {
|
|
qsum = fma(q_raw[i], q_raw[i], qsum);
|
|
ksum = fma(k_raw[i], k_raw[i], ksum);
|
|
}
|
|
const float qscale = rsqrt(qsum + args.f[0]) * rsqrt((float)dim);
|
|
const float kscale = rsqrt(ksum + args.f[0]);
|
|
const float beta = 1.0f / (1.0f + exp(-controls[args.u[3] + head]));
|
|
const float step = controls[args.u[4] + head] +
|
|
qwen_bf16(qwen_weight_u16(dt_bias, args.u[14], head));
|
|
const float softplus = max(step, 0.0f) + log(1.0f + exp(-abs(step)));
|
|
const float decay = exp(-exp(qwen_bf16(qwen_weight_u16(a_log, args.u[13], head))) * softplus);
|
|
device float *column = state + ((ulong)head * dim * dim) + value_index;
|
|
float prediction = 0.0f;
|
|
for (uint i = 0; i < dim; i++) {
|
|
prediction = fma(column[(ulong)i * dim] * decay, k_raw[i] * kscale, prediction);
|
|
}
|
|
const float delta = (value[value_index] - prediction) * beta;
|
|
float result = 0.0f;
|
|
for (uint i = 0; i < dim; i++) {
|
|
const ulong offset = (ulong)i * dim;
|
|
const float updated = column[offset] * decay + k_raw[i] * kscale * delta;
|
|
column[offset] = updated;
|
|
result = fma(updated, q_raw[i] * qscale, result);
|
|
}
|
|
device float *head_out = out + (ulong)head * dim;
|
|
head_out[value_index] = result;
|
|
}
|
|
|
|
kernel void kernel_qwen_gdn_norm_gate(
|
|
constant qwen_kernel_args &args [[buffer(0)]],
|
|
device float *out [[buffer(1)]],
|
|
device const float *x [[buffer(2)]],
|
|
device const float *controls [[buffer(3)]],
|
|
device const uchar *weight [[buffer(5)]],
|
|
uint head [[thread_position_in_grid]]) {
|
|
const uint dim = args.u[0];
|
|
if (head >= args.u[1]) return;
|
|
device const float *row = x + (ulong)head * dim;
|
|
float variance = 0.0f;
|
|
for (uint i = 0; i < dim; i++) variance = fma(row[i], row[i], variance);
|
|
const float scale = rsqrt(variance / (float)dim + args.f[0]);
|
|
for (uint i = 0; i < dim; i++) {
|
|
const ulong index = (ulong)head * dim + i;
|
|
out[index] = row[i] * scale *
|
|
qwen_bf16(qwen_weight_u16(weight, args.u[13], i)) /
|
|
(1.0f + exp(-controls[index]));
|
|
}
|
|
}
|
|
|
|
kernel void kernel_qwen_swiglu(
|
|
constant qwen_kernel_args &args [[buffer(0)]],
|
|
device float *out [[buffer(1)]],
|
|
device const float *gate [[buffer(2)]],
|
|
device const float *up [[buffer(3)]],
|
|
uint index [[thread_position_in_grid]]) {
|
|
if (index >= args.u[0]) return;
|
|
out[index] = gate[index] / (1.0f + exp(-gate[index])) * up[index];
|
|
}
|
|
|
|
kernel void kernel_qwen_route_top10(
|
|
constant qwen_kernel_args &args [[buffer(0)]],
|
|
device int *ids [[buffer(1)]],
|
|
device float *weights [[buffer(2)]],
|
|
device const float *logits [[buffer(3)]],
|
|
uint gid [[thread_position_in_grid]]) {
|
|
if (gid != 0u) return;
|
|
float max_value = -INFINITY;
|
|
for (uint i = 0; i < args.u[0]; i++) max_value = max(max_value, logits[i]);
|
|
float sum = 0.0f;
|
|
for (uint i = 0; i < args.u[0]; i++) sum += exp(logits[i] - max_value);
|
|
float selected_sum = 0.0f;
|
|
for (uint slot = 0; slot < 10u; slot++) {
|
|
float best = -1.0f;
|
|
int best_id = -1;
|
|
for (uint i = 0; i < args.u[0]; i++) {
|
|
bool used = false;
|
|
for (uint j = 0; j < slot; j++) used = used || ids[j] == (int)i;
|
|
const float probability = exp(logits[i] - max_value) / sum;
|
|
if (!used && probability > best) {
|
|
best = probability;
|
|
best_id = (int)i;
|
|
}
|
|
}
|
|
ids[slot] = best_id;
|
|
weights[slot] = best;
|
|
selected_sum += best;
|
|
}
|
|
for (uint slot = 0; slot < 10u; slot++) weights[slot] /= selected_sum;
|
|
}
|
|
|
|
kernel void kernel_qwen_accumulate(
|
|
constant qwen_kernel_args &args [[buffer(0)]],
|
|
device float *out [[buffer(1)]],
|
|
device const float *x [[buffer(2)]],
|
|
uint index [[thread_position_in_grid]]) {
|
|
if (index < args.u[0]) out[index] += x[index] * args.f[0];
|
|
}
|
|
|
|
kernel void kernel_qwen_accumulate_sigmoid_scalar(
|
|
constant qwen_kernel_args &args [[buffer(0)]],
|
|
device float *out [[buffer(1)]],
|
|
device const float *x [[buffer(2)]],
|
|
device const float *gate [[buffer(3)]],
|
|
uint index [[thread_position_in_grid]]) {
|
|
if (index < args.u[0]) out[index] += x[index] / (1.0f + exp(-gate[0]));
|
|
}
|
|
|
|
kernel void kernel_qwen_split_q_gate(
|
|
constant qwen_kernel_args &args [[buffer(0)]],
|
|
device float *q [[buffer(1)]],
|
|
device const float *packed [[buffer(2)]],
|
|
device float *gate [[buffer(3)]],
|
|
uint index [[thread_position_in_grid]]) {
|
|
const uint heads = args.u[0];
|
|
const uint dim = args.u[1];
|
|
if (index >= heads * dim) return;
|
|
const uint head = index / dim;
|
|
const uint column = index % dim;
|
|
q[index] = packed[(ulong)head * dim * 2u + column];
|
|
gate[index] = packed[(ulong)head * dim * 2u + dim + column];
|
|
}
|
|
|
|
kernel void kernel_qwen_head_norm_rope(
|
|
constant qwen_kernel_args &args [[buffer(0)]],
|
|
device float *out [[buffer(1)]],
|
|
device const float *x [[buffer(2)]],
|
|
device const uchar *weight [[buffer(5)]],
|
|
uint2 gid [[thread_position_in_grid]]) {
|
|
const uint column = gid.x;
|
|
const uint head = gid.y;
|
|
const uint dim = args.u[0];
|
|
const uint rotary = args.u[1];
|
|
if (column >= dim || head >= args.u[2]) return;
|
|
device const float *row = x + (ulong)head * dim;
|
|
float variance = 0.0f;
|
|
for (uint i = 0; i < dim; i++) variance = fma(row[i], row[i], variance);
|
|
const float scale = rsqrt(variance / (float)dim + args.f[0]);
|
|
float value = row[column] * scale *
|
|
(1.0f + qwen_bf16(qwen_weight_u16(weight, args.u[13], column)));
|
|
if (column < rotary) {
|
|
const uint rotary_half = rotary / 2u;
|
|
const uint pair = column < rotary_half ? column + rotary_half : column - rotary_half;
|
|
const float paired = row[pair] * scale *
|
|
(1.0f + qwen_bf16(qwen_weight_u16(weight, args.u[13], pair)));
|
|
const float theta = (float)args.u[3] * pow(args.f[1], -2.0f * (float)(column % rotary_half) / (float)rotary);
|
|
value = value * cos(theta) + (column < rotary_half ? -paired : paired) * sin(theta);
|
|
}
|
|
out[(ulong)head * dim + column] = value;
|
|
}
|
|
|
|
kernel void kernel_qwen_store_kv_bf16(
|
|
constant qwen_kernel_args &args [[buffer(0)]],
|
|
device ushort *cache [[buffer(1)]],
|
|
device const float *key [[buffer(2)]],
|
|
device const float *value [[buffer(3)]],
|
|
uint index [[thread_position_in_grid]]) {
|
|
const uint width = args.u[0];
|
|
if (index >= width) return;
|
|
const ulong base = (ulong)args.u[1] * width * 2u;
|
|
cache[base + index] = qwen_to_bf16(key[index]);
|
|
cache[base + width + index] = qwen_to_bf16(value[index]);
|
|
}
|
|
|
|
kernel void kernel_qwen_qsa_store_raw(
|
|
constant qwen_kernel_args &args [[buffer(0)]],
|
|
device ushort *raw [[buffer(1)]],
|
|
device const float *projected [[buffer(2)]],
|
|
device float *query [[buffer(3)]],
|
|
uint index [[thread_position_in_grid]]) {
|
|
const uint dim = args.u[0];
|
|
const uint query_width = args.u[1] * dim;
|
|
if (index < query_width) {
|
|
query[index] = projected[index];
|
|
} else if (index < query_width + dim) {
|
|
raw[(ulong)args.u[2] * dim + index - query_width] = qwen_to_bf16(projected[index]);
|
|
}
|
|
}
|
|
|
|
kernel void kernel_qwen_qsa_pool_key(
|
|
constant qwen_kernel_args &args [[buffer(0)]],
|
|
device ushort *pooled [[buffer(1)]],
|
|
device const ushort *raw [[buffer(2)]],
|
|
device const uchar *weight [[buffer(5)]],
|
|
uint column [[thread_position_in_grid]]) {
|
|
const uint dim = args.u[0];
|
|
if (column >= dim) return;
|
|
const ulong raw_start = (ulong)args.u[3] * dim;
|
|
float mean = 0.0f;
|
|
for (uint token = 0; token < args.u[2]; token++) {
|
|
mean += qwen_bf16(raw[raw_start + (ulong)token * dim + column]);
|
|
}
|
|
const ushort mean_bf16 = qwen_to_bf16(mean / (float)args.u[2]);
|
|
float variance = 0.0f;
|
|
for (uint i = 0; i < dim; i++) {
|
|
float item = 0.0f;
|
|
for (uint token = 0; token < args.u[2]; token++) {
|
|
item += qwen_bf16(raw[raw_start + (ulong)token * dim + i]);
|
|
}
|
|
item = qwen_bf16(qwen_to_bf16(item / (float)args.u[2]));
|
|
variance = fma(item, item, variance);
|
|
}
|
|
const float scale = rsqrt(variance / (float)dim + args.f[0]);
|
|
float value = qwen_bf16(mean_bf16) * scale *
|
|
(1.0f + qwen_bf16(qwen_weight_u16(weight, args.u[13], column)));
|
|
const uint rotary = 64u;
|
|
if (column < rotary) {
|
|
const uint rotary_half = rotary / 2u;
|
|
const uint pair = column < rotary_half ? column + rotary_half : column - rotary_half;
|
|
float paired_mean = 0.0f;
|
|
for (uint token = 0; token < args.u[2]; token++) {
|
|
paired_mean += qwen_bf16(raw[raw_start + (ulong)token * dim + pair]);
|
|
}
|
|
paired_mean = qwen_bf16(qwen_to_bf16(paired_mean / (float)args.u[2]));
|
|
const float paired = paired_mean * scale *
|
|
(1.0f + qwen_bf16(qwen_weight_u16(weight, args.u[13], pair)));
|
|
const float theta = (float)args.u[3] *
|
|
pow(args.f[1], -2.0f * (float)(column % rotary_half) / (float)rotary);
|
|
value = value * cos(theta) + (column < rotary_half ? -paired : paired) * sin(theta);
|
|
}
|
|
pooled[(ulong)args.u[1] * dim + column] = qwen_to_bf16(value);
|
|
}
|
|
|
|
kernel void kernel_qwen_qsa_scores(
|
|
constant qwen_kernel_args &args [[buffer(0)]],
|
|
device float *scores [[buffer(1)]],
|
|
device const float *query [[buffer(2)]],
|
|
device const ushort *pooled [[buffer(3)]],
|
|
uint block [[thread_position_in_grid]]) {
|
|
const uint dim = args.u[0];
|
|
if (block >= args.u[2]) return;
|
|
float score = 0.0f;
|
|
for (uint head = 0; head < args.u[1]; head++) {
|
|
float head_score = 0.0f;
|
|
for (uint i = 0; i < dim; i++) {
|
|
head_score = fma(query[(ulong)head * dim + i],
|
|
qwen_bf16(pooled[(ulong)block * dim + i]),
|
|
head_score);
|
|
}
|
|
score += max(head_score, 0.0f);
|
|
}
|
|
scores[block] = score * args.f[0];
|
|
}
|
|
|
|
kernel void kernel_qwen_qsa_sort_blocks(
|
|
constant qwen_kernel_args &args [[buffer(0)]],
|
|
device int *selected [[buffer(1)]],
|
|
uint gid [[thread_position_in_grid]]) {
|
|
if (gid != 0u) return;
|
|
for (uint i = 1; i < args.u[0]; i++) {
|
|
const int value = selected[i];
|
|
uint j = i;
|
|
while (j > 0u && selected[j - 1u] > value) {
|
|
selected[j] = selected[j - 1u];
|
|
j--;
|
|
}
|
|
selected[j] = value;
|
|
}
|
|
}
|
|
|
|
static inline uint qwen_qsa_token(
|
|
device const int *selected,
|
|
uint ordinal,
|
|
uint selected_count,
|
|
uint ratio,
|
|
uint tail_start) {
|
|
const uint selected_tokens = selected_count * ratio;
|
|
return ordinal < selected_tokens
|
|
? (uint)selected[ordinal / ratio] * ratio + ordinal % ratio
|
|
: tail_start + ordinal - selected_tokens;
|
|
}
|
|
|
|
kernel void kernel_qwen_sparse_attention(
|
|
constant qwen_kernel_args &args [[buffer(0)]],
|
|
device float *out [[buffer(1)]],
|
|
device const float *query [[buffer(2)]],
|
|
device const ushort *cache [[buffer(3)]],
|
|
device const int *selected [[buffer(4)]],
|
|
uint2 gid [[thread_position_in_grid]],
|
|
uint lane [[thread_index_in_threadgroup]]) {
|
|
const uint heads = args.u[0];
|
|
const uint kv_heads = args.u[1];
|
|
const uint dim = args.u[2];
|
|
const uint column = gid.x;
|
|
const uint head = gid.y;
|
|
if (column >= dim || head >= heads) return;
|
|
const uint kv_head = head / (heads / kv_heads);
|
|
const uint tokens = args.u[4] * args.u[5] + args.u[7];
|
|
const float attention_scale = rsqrt((float)dim);
|
|
threadgroup float probabilities[2051];
|
|
if (lane == 0u) {
|
|
float max_score = -INFINITY;
|
|
for (uint ordinal = 0; ordinal < tokens; ordinal++) {
|
|
const uint token = qwen_qsa_token(selected, ordinal, args.u[4], args.u[5], args.u[6]);
|
|
const ulong base = (ulong)token * kv_heads * dim * 2u + (ulong)kv_head * dim;
|
|
float score = 0.0f;
|
|
for (uint i = 0; i < dim; i++) {
|
|
score = fma(query[(ulong)head * dim + i], qwen_bf16(cache[base + i]), score);
|
|
}
|
|
probabilities[ordinal] = score * attention_scale;
|
|
max_score = max(max_score, probabilities[ordinal]);
|
|
}
|
|
float denominator = 0.0f;
|
|
for (uint ordinal = 0; ordinal < tokens; ordinal++) {
|
|
probabilities[ordinal] = exp(probabilities[ordinal] - max_score);
|
|
denominator += probabilities[ordinal];
|
|
}
|
|
for (uint ordinal = 0; ordinal < tokens; ordinal++) {
|
|
probabilities[ordinal] /= denominator;
|
|
}
|
|
}
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
float value = 0.0f;
|
|
for (uint ordinal = 0; ordinal < tokens; ordinal++) {
|
|
const uint token = qwen_qsa_token(selected, ordinal, args.u[4], args.u[5], args.u[6]);
|
|
const ulong base = (ulong)token * kv_heads * dim * 2u + (ulong)kv_head * dim;
|
|
value = fma(probabilities[ordinal],
|
|
qwen_bf16(cache[base + kv_heads * dim + column]),
|
|
value);
|
|
}
|
|
out[(ulong)head * dim + column] = value;
|
|
}
|
|
|
|
kernel void kernel_qwen_dense_attention(
|
|
constant qwen_kernel_args &args [[buffer(0)]],
|
|
device float *out [[buffer(1)]],
|
|
device const float *query [[buffer(2)]],
|
|
device const ushort *cache [[buffer(3)]],
|
|
uint2 gid [[thread_position_in_grid]],
|
|
uint lane [[thread_index_in_threadgroup]]) {
|
|
const uint heads = args.u[0];
|
|
const uint kv_heads = args.u[1];
|
|
const uint dim = args.u[2];
|
|
const uint tokens = args.u[3];
|
|
const uint column = gid.x;
|
|
const uint head = gid.y;
|
|
if (column >= dim || head >= heads) return;
|
|
const uint kv_head = head / (heads / kv_heads);
|
|
const float attention_scale = rsqrt((float)dim);
|
|
threadgroup float probabilities[2048];
|
|
if (lane == 0u) {
|
|
float max_score = -INFINITY;
|
|
for (uint token = 0; token < tokens; token++) {
|
|
const ulong base = (ulong)token * kv_heads * dim * 2u + (ulong)kv_head * dim;
|
|
float score = 0.0f;
|
|
for (uint i = 0; i < dim; i++) score = fma(query[(ulong)head * dim + i], qwen_bf16(cache[base + i]), score);
|
|
probabilities[token] = score * attention_scale;
|
|
max_score = max(max_score, probabilities[token]);
|
|
}
|
|
float denominator = 0.0f;
|
|
for (uint token = 0; token < tokens; token++) {
|
|
probabilities[token] = exp(probabilities[token] - max_score);
|
|
denominator += probabilities[token];
|
|
}
|
|
for (uint token = 0; token < tokens; token++) probabilities[token] /= denominator;
|
|
}
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
float value = 0.0f;
|
|
for (uint token = 0; token < tokens; token++) {
|
|
const ulong base = (ulong)token * kv_heads * dim * 2u + (ulong)kv_head * dim;
|
|
value = fma(probabilities[token], qwen_bf16(cache[base + kv_heads * dim + column]), value);
|
|
}
|
|
out[(ulong)head * dim + column] = value;
|
|
}
|
|
|
|
kernel void kernel_qwen_gate_attention(
|
|
constant qwen_kernel_args &args [[buffer(0)]],
|
|
device float *out [[buffer(1)]],
|
|
device const float *attention [[buffer(2)]],
|
|
device const float *gate [[buffer(3)]],
|
|
uint index [[thread_position_in_grid]]) {
|
|
if (index < args.u[0]) out[index] = attention[index] / (1.0f + exp(-gate[index]));
|
|
}
|