hugo/DS4Server
1
0
Fork 0
Code Issues 2 Pull Requests Actions Packages Projects Releases Wiki Activity

Stream Qwen PLE embeddings

Browse Source
This commit is contained in:
Georg Bauer
2026-09-03 21:12:20 +02:00
parent c414640050
commit 87ccf67d0c
4 changed files with 893 additions and 78 deletions
Download Patch File Download Diff File Expand all files Collapse all files

180
metal/qwen38.metal
View File

@@ -16,6 +16,52 @@ static inline ushort qwen_to_bf16(float value) {
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,
@@ -39,17 +85,18 @@ 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)]],
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_value(packed, scales, biases, row, column,
in_dim, args.u[2], args.u[3]), x[column], sum);
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;
}
@@ -57,25 +104,43 @@ kernel void kernel_qwen_affine_mv(
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)]],
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_value(packed, scales, biases, args.u[4], column,
args.u[0], args.u[2], args.u[3]);
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 ushort *weights [[buffer(5)]],
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(weights[(ulong)row * args.u[0] + column]), x[column], sum);
sum = fma(qwen_bf16(qwen_weight_u16(weights, args.u[13],
(ulong)row * args.u[0] + column)),
x[column], sum);
}
out[row] = sum;
}
@@ -92,7 +157,7 @@ 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)]],
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];
@@ -103,7 +168,8 @@ kernel void kernel_qwen_zero_rms(
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]));
out[index] = x[index] * scale *
(1.0f + qwen_bf16(qwen_weight_u16(weight, args.u[13], index)));
}
}
@@ -160,19 +226,73 @@ kernel void kernel_qwen_hyper_inject(
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 ushort *weight [[buffer(5)]],
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(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]))));
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]);
@@ -185,8 +305,8 @@ kernel void kernel_qwen_gdn_step(
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)]],
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;
@@ -207,9 +327,10 @@ kernel void kernel_qwen_gdn_step(
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 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(a_log[head])) * softplus);
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++) {
@@ -232,7 +353,7 @@ kernel void kernel_qwen_gdn_norm_gate(
device float *out [[buffer(1)]],
device const float *x [[buffer(2)]],
device const float *controls [[buffer(3)]],
device const ushort *weight [[buffer(5)]],
device const uchar *weight [[buffer(5)]],
uint head [[thread_position_in_grid]]) {
const uint dim = args.u[0];
if (head >= args.u[1]) return;
@@ -242,7 +363,8 @@ kernel void kernel_qwen_gdn_norm_gate(
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]) /
out[index] = row[i] * scale *
qwen_bf16(qwen_weight_u16(weight, args.u[13], i)) /
(1.0f + exp(-controls[index]));
}
}
@@ -324,7 +446,7 @@ 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)]],
device const uchar *weight [[buffer(5)]],
uint2 gid [[thread_position_in_grid]]) {
const uint column = gid.x;
const uint head = gid.y;
@@ -335,11 +457,13 @@ kernel void kernel_qwen_head_norm_rope(
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]));
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(weight[pair]));
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);
}