Implement Qwen text core
This commit is contained in:
Georg Bauer
408
metal/qwen38.metal
Normal file
408
metal/qwen38.metal
Normal file
@@ -0,0 +1,408 @@
|
||||
// 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 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 uint *packed [[buffer(5)]],
|
||||
device const ushort *scales [[buffer(6)]],
|
||||
device const ushort *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_value(packed, scales, biases, 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 uint *packed [[buffer(5)]],
|
||||
device const ushort *scales [[buffer(6)]],
|
||||
device const ushort *biases [[buffer(7)]],
|
||||
uint column [[thread_position_in_grid]]) {
|
||||
if (column >= args.u[0]) return;
|
||||
out[column] = qwen_quant_value(packed, scales, biases, args.u[4], column,
|
||||
args.u[0], 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 ushort *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(weights[(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 ushort *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(weight[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_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 ushort *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(weight[(ulong)channel * 4u]),
|
||||
fma(qwen_bf16(history[1]), qwen_bf16(weight[(ulong)channel * 4u + 1u]),
|
||||
fma(qwen_bf16(history[2]), qwen_bf16(weight[(ulong)channel * 4u + 2u]),
|
||||
x[channel] * qwen_bf16(weight[(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 ushort *a_log [[buffer(5)]],
|
||||
device const ushort *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(dt_bias[head]);
|
||||
const float softplus = max(step, 0.0f) + log(1.0f + exp(-abs(step)));
|
||||
const float decay = exp(-exp(qwen_bf16(a_log[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 ushort *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(weight[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 ushort *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(weight[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(weight[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_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)]],
|
||||
uint head [[thread_position_in_grid]]) {
|
||||
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];
|
||||
if (head >= heads) return;
|
||||
const uint kv_head = head / (heads / kv_heads);
|
||||
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);
|
||||
max_score = max(max_score, score * rsqrt((float)dim));
|
||||
}
|
||||
float denominator = 0.0f;
|
||||
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);
|
||||
denominator += exp(score * rsqrt((float)dim) - max_score);
|
||||
}
|
||||
for (uint column = 0; column < dim; column++) {
|
||||
float value = 0.0f;
|
||||
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);
|
||||
const float probability = exp(score * rsqrt((float)dim) - max_score) / denominator;
|
||||
value = fma(probability, 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]));
|
||||
}
|
||||
Reference in New Issue
Block a user