Stream Qwen PLE embeddings
This commit is contained in:
@@ -16,6 +16,52 @@ static inline ushort qwen_to_bf16(float value) {
|
|||||||
return (ushort)(bits >> 16);
|
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(
|
static inline float qwen_quant_value(
|
||||||
device const uint *packed,
|
device const uint *packed,
|
||||||
device const ushort *scales,
|
device const ushort *scales,
|
||||||
@@ -39,17 +85,18 @@ kernel void kernel_qwen_affine_mv(
|
|||||||
constant qwen_kernel_args &args [[buffer(0)]],
|
constant qwen_kernel_args &args [[buffer(0)]],
|
||||||
device float *out [[buffer(1)]],
|
device float *out [[buffer(1)]],
|
||||||
device const float *x [[buffer(2)]],
|
device const float *x [[buffer(2)]],
|
||||||
device const uint *packed [[buffer(5)]],
|
device const uchar *packed [[buffer(5)]],
|
||||||
device const ushort *scales [[buffer(6)]],
|
device const uchar *scales [[buffer(6)]],
|
||||||
device const ushort *biases [[buffer(7)]],
|
device const uchar *biases [[buffer(7)]],
|
||||||
uint row [[thread_position_in_grid]]) {
|
uint row [[thread_position_in_grid]]) {
|
||||||
const uint in_dim = args.u[0];
|
const uint in_dim = args.u[0];
|
||||||
const uint out_dim = args.u[1];
|
const uint out_dim = args.u[1];
|
||||||
if (row >= out_dim) return;
|
if (row >= out_dim) return;
|
||||||
float sum = 0.0f;
|
float sum = 0.0f;
|
||||||
for (uint column = 0; column < in_dim; column++) {
|
for (uint column = 0; column < in_dim; column++) {
|
||||||
sum = fma(qwen_quant_value(packed, scales, biases, row, column,
|
sum = fma(qwen_quant_weight(packed, scales, biases,
|
||||||
in_dim, args.u[2], args.u[3]), x[column], sum);
|
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;
|
out[row] = sum;
|
||||||
}
|
}
|
||||||
@@ -57,25 +104,43 @@ kernel void kernel_qwen_affine_mv(
|
|||||||
kernel void kernel_qwen_affine_embedding(
|
kernel void kernel_qwen_affine_embedding(
|
||||||
constant qwen_kernel_args &args [[buffer(0)]],
|
constant qwen_kernel_args &args [[buffer(0)]],
|
||||||
device float *out [[buffer(1)]],
|
device float *out [[buffer(1)]],
|
||||||
device const uint *packed [[buffer(5)]],
|
device const uchar *packed [[buffer(5)]],
|
||||||
device const ushort *scales [[buffer(6)]],
|
device const uchar *scales [[buffer(6)]],
|
||||||
device const ushort *biases [[buffer(7)]],
|
device const uchar *biases [[buffer(7)]],
|
||||||
uint column [[thread_position_in_grid]]) {
|
uint column [[thread_position_in_grid]]) {
|
||||||
if (column >= args.u[0]) return;
|
if (column >= args.u[0]) return;
|
||||||
out[column] = qwen_quant_value(packed, scales, biases, args.u[4], column,
|
out[column] = qwen_quant_weight(packed, scales, biases,
|
||||||
args.u[0], args.u[2], args.u[3]);
|
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(
|
kernel void kernel_qwen_bf16_mv(
|
||||||
constant qwen_kernel_args &args [[buffer(0)]],
|
constant qwen_kernel_args &args [[buffer(0)]],
|
||||||
device float *out [[buffer(1)]],
|
device float *out [[buffer(1)]],
|
||||||
device const float *x [[buffer(2)]],
|
device const float *x [[buffer(2)]],
|
||||||
device const ushort *weights [[buffer(5)]],
|
device const uchar *weights [[buffer(5)]],
|
||||||
uint row [[thread_position_in_grid]]) {
|
uint row [[thread_position_in_grid]]) {
|
||||||
if (row >= args.u[1]) return;
|
if (row >= args.u[1]) return;
|
||||||
float sum = 0.0f;
|
float sum = 0.0f;
|
||||||
for (uint column = 0; column < args.u[0]; column++) {
|
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;
|
out[row] = sum;
|
||||||
}
|
}
|
||||||
@@ -92,7 +157,7 @@ kernel void kernel_qwen_zero_rms(
|
|||||||
constant qwen_kernel_args &args [[buffer(0)]],
|
constant qwen_kernel_args &args [[buffer(0)]],
|
||||||
device float *out [[buffer(1)]],
|
device float *out [[buffer(1)]],
|
||||||
device const float *x [[buffer(2)]],
|
device const float *x [[buffer(2)]],
|
||||||
device const ushort *weight [[buffer(5)]],
|
device const uchar *weight [[buffer(5)]],
|
||||||
uint group [[thread_position_in_grid]]) {
|
uint group [[thread_position_in_grid]]) {
|
||||||
const uint width = args.u[0];
|
const uint width = args.u[0];
|
||||||
const uint group_size = args.u[1];
|
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]);
|
const float scale = rsqrt(variance / (float)group_size + args.f[0]);
|
||||||
for (uint i = 0; i < group_size; i++) {
|
for (uint i = 0; i < group_size; i++) {
|
||||||
const uint index = start + 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];
|
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(
|
kernel void kernel_qwen_conv_silu(
|
||||||
constant qwen_kernel_args &args [[buffer(0)]],
|
constant qwen_kernel_args &args [[buffer(0)]],
|
||||||
device float *out [[buffer(1)]],
|
device float *out [[buffer(1)]],
|
||||||
device const float *x [[buffer(2)]],
|
device const float *x [[buffer(2)]],
|
||||||
device ushort *state [[buffer(3)]],
|
device ushort *state [[buffer(3)]],
|
||||||
device const ushort *weight [[buffer(5)]],
|
device const uchar *weight [[buffer(5)]],
|
||||||
uint channel [[thread_position_in_grid]]) {
|
uint channel [[thread_position_in_grid]]) {
|
||||||
if (channel >= args.u[0]) return;
|
if (channel >= args.u[0]) return;
|
||||||
device ushort *history = state + (ulong)channel * 3u;
|
device ushort *history = state + (ulong)channel * 3u;
|
||||||
float value = fma(qwen_bf16(history[0]), qwen_bf16(weight[(ulong)channel * 4u]),
|
float value = fma(qwen_bf16(history[0]),
|
||||||
fma(qwen_bf16(history[1]), qwen_bf16(weight[(ulong)channel * 4u + 1u]),
|
qwen_bf16(qwen_weight_u16(weight, args.u[13], (ulong)channel * 4u)),
|
||||||
fma(qwen_bf16(history[2]), qwen_bf16(weight[(ulong)channel * 4u + 2u]),
|
fma(qwen_bf16(history[1]),
|
||||||
x[channel] * qwen_bf16(weight[(ulong)channel * 4u + 3u]))));
|
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[0] = history[1];
|
||||||
history[1] = history[2];
|
history[1] = history[2];
|
||||||
history[2] = qwen_to_bf16(x[channel]);
|
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 *qkv [[buffer(2)]],
|
||||||
device const float *controls [[buffer(3)]],
|
device const float *controls [[buffer(3)]],
|
||||||
device float *state [[buffer(4)]],
|
device float *state [[buffer(4)]],
|
||||||
device const ushort *a_log [[buffer(5)]],
|
device const uchar *a_log [[buffer(5)]],
|
||||||
device const ushort *dt_bias [[buffer(6)]],
|
device const uchar *dt_bias [[buffer(6)]],
|
||||||
uint2 gid [[thread_position_in_grid]]) {
|
uint2 gid [[thread_position_in_grid]]) {
|
||||||
const uint value_index = gid.x;
|
const uint value_index = gid.x;
|
||||||
const uint head = gid.y;
|
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 qscale = rsqrt(qsum + args.f[0]) * rsqrt((float)dim);
|
||||||
const float kscale = rsqrt(ksum + args.f[0]);
|
const float kscale = rsqrt(ksum + args.f[0]);
|
||||||
const float beta = 1.0f / (1.0f + exp(-controls[args.u[3] + head]));
|
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 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;
|
device float *column = state + ((ulong)head * dim * dim) + value_index;
|
||||||
float prediction = 0.0f;
|
float prediction = 0.0f;
|
||||||
for (uint i = 0; i < dim; i++) {
|
for (uint i = 0; i < dim; i++) {
|
||||||
@@ -232,7 +353,7 @@ kernel void kernel_qwen_gdn_norm_gate(
|
|||||||
device float *out [[buffer(1)]],
|
device float *out [[buffer(1)]],
|
||||||
device const float *x [[buffer(2)]],
|
device const float *x [[buffer(2)]],
|
||||||
device const float *controls [[buffer(3)]],
|
device const float *controls [[buffer(3)]],
|
||||||
device const ushort *weight [[buffer(5)]],
|
device const uchar *weight [[buffer(5)]],
|
||||||
uint head [[thread_position_in_grid]]) {
|
uint head [[thread_position_in_grid]]) {
|
||||||
const uint dim = args.u[0];
|
const uint dim = args.u[0];
|
||||||
if (head >= args.u[1]) return;
|
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]);
|
const float scale = rsqrt(variance / (float)dim + args.f[0]);
|
||||||
for (uint i = 0; i < dim; i++) {
|
for (uint i = 0; i < dim; i++) {
|
||||||
const ulong index = (ulong)head * 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]));
|
(1.0f + exp(-controls[index]));
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -324,7 +446,7 @@ kernel void kernel_qwen_head_norm_rope(
|
|||||||
constant qwen_kernel_args &args [[buffer(0)]],
|
constant qwen_kernel_args &args [[buffer(0)]],
|
||||||
device float *out [[buffer(1)]],
|
device float *out [[buffer(1)]],
|
||||||
device const float *x [[buffer(2)]],
|
device const float *x [[buffer(2)]],
|
||||||
device const ushort *weight [[buffer(5)]],
|
device const uchar *weight [[buffer(5)]],
|
||||||
uint2 gid [[thread_position_in_grid]]) {
|
uint2 gid [[thread_position_in_grid]]) {
|
||||||
const uint column = gid.x;
|
const uint column = gid.x;
|
||||||
const uint head = gid.y;
|
const uint head = gid.y;
|
||||||
@@ -335,11 +457,13 @@ kernel void kernel_qwen_head_norm_rope(
|
|||||||
float variance = 0.0f;
|
float variance = 0.0f;
|
||||||
for (uint i = 0; i < dim; i++) variance = fma(row[i], row[i], variance);
|
for (uint i = 0; i < dim; i++) variance = fma(row[i], row[i], variance);
|
||||||
const float scale = rsqrt(variance / (float)dim + args.f[0]);
|
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) {
|
if (column < rotary) {
|
||||||
const uint rotary_half = rotary / 2u;
|
const uint rotary_half = rotary / 2u;
|
||||||
const uint pair = column < rotary_half ? column + rotary_half : column - rotary_half;
|
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);
|
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);
|
value = value * cos(theta) + (column < rotary_half ? -paired : paired) * sin(theta);
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -11762,7 +11762,7 @@ int ds4_gpu_qwen_dispatch(
|
|||||||
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
|
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
|
||||||
if (!enc) return 0;
|
if (!enc) return 0;
|
||||||
[enc setComputePipelineState:pipeline];
|
[enc setComputePipelineState:pipeline];
|
||||||
[enc setBytes:args length:sizeof(*args) atIndex:0];
|
ds4_gpu_qwen_kernel_args bound_args = *args;
|
||||||
|
|
||||||
const DS4MetalTensor *tensors[4] = {
|
const DS4MetalTensor *tensors[4] = {
|
||||||
ds4_gpu_tensor_const_obj(out),
|
ds4_gpu_tensor_const_obj(out),
|
||||||
@@ -11788,8 +11788,13 @@ int ds4_gpu_qwen_dispatch(
|
|||||||
if (owned) [cb commit];
|
if (owned) [cb commit];
|
||||||
return 0;
|
return 0;
|
||||||
}
|
}
|
||||||
[enc setBuffer:weight offset:(NSUInteger)inner atIndex:5 + i];
|
/* Metal resource offsets are four-byte aligned. Preserve an arbitrary
|
||||||
|
* safetensors data offset for the Qwen kernels to decode explicitly. */
|
||||||
|
const uint64_t aligned_inner = inner & ~3ull;
|
||||||
|
bound_args.u[13 + i] = (uint32_t)(inner - aligned_inner);
|
||||||
|
[enc setBuffer:weight offset:(NSUInteger)aligned_inner atIndex:5 + i];
|
||||||
}
|
}
|
||||||
|
[enc setBytes:&bound_args length:sizeof(bound_args) atIndex:0];
|
||||||
|
|
||||||
const NSUInteger width = pipeline.threadExecutionWidth;
|
const NSUInteger width = pipeline.threadExecutionWidth;
|
||||||
const NSUInteger max_threads = pipeline.maxTotalThreadsPerThreadgroup;
|
const NSUInteger max_threads = pipeline.maxTotalThreadsPerThreadgroup;
|
||||||
|
|||||||
@@ -24,8 +24,14 @@ const EXPERTS_USED: usize = 10;
|
|||||||
const EXPERT_WIDTH: u32 = 640;
|
const EXPERT_WIDTH: u32 = 640;
|
||||||
const VOCAB: u32 = 248_320;
|
const VOCAB: u32 = 248_320;
|
||||||
const DENSE_BUDGET: u32 = 2_048;
|
const DENSE_BUDGET: u32 = 2_048;
|
||||||
|
const PLE_HEADS: usize = 16;
|
||||||
|
const PLE_HEAD_DIM: u32 = 160;
|
||||||
|
const PLE_HISTORY: usize = 2;
|
||||||
|
const PLE_CONV_STATE: u32 = 9;
|
||||||
|
const EOS_TOKEN: i32 = 248_044;
|
||||||
|
const PLE_ROW_BYTES: usize = 100;
|
||||||
const CHECKPOINT_MAGIC: &[u8; 8] = b"DS4QWN01";
|
const CHECKPOINT_MAGIC: &[u8; 8] = b"DS4QWN01";
|
||||||
const CHECKPOINT_VERSION: u32 = 1;
|
const CHECKPOINT_VERSION: u32 = 2;
|
||||||
const CHECKPOINT_CHUNK: usize = 8 * 1024 * 1024;
|
const CHECKPOINT_CHUNK: usize = 8 * 1024 * 1024;
|
||||||
|
|
||||||
#[derive(Clone, Copy)]
|
#[derive(Clone, Copy)]
|
||||||
@@ -75,6 +81,15 @@ struct Scratch {
|
|||||||
k_rope: Buffer,
|
k_rope: Buffer,
|
||||||
v: Buffer,
|
v: Buffer,
|
||||||
attention: Buffer,
|
attention: Buffer,
|
||||||
|
ple_packed: Buffer,
|
||||||
|
ple_scales: Buffer,
|
||||||
|
ple_biases: Buffer,
|
||||||
|
ple_embedding: Buffer,
|
||||||
|
ple_key: Buffer,
|
||||||
|
ple_value: Buffer,
|
||||||
|
ple_gated: Buffer,
|
||||||
|
ple_norm: Buffer,
|
||||||
|
ple_output: Buffer,
|
||||||
logits: Buffer,
|
logits: Buffer,
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -108,14 +123,36 @@ impl Scratch {
|
|||||||
k_rope: Buffer::floats(ATTN_KV_WIDTH.into())?,
|
k_rope: Buffer::floats(ATTN_KV_WIDTH.into())?,
|
||||||
v: Buffer::floats(ATTN_KV_WIDTH.into())?,
|
v: Buffer::floats(ATTN_KV_WIDTH.into())?,
|
||||||
attention: Buffer::floats(ATTN_WIDTH.into())?,
|
attention: Buffer::floats(ATTN_WIDTH.into())?,
|
||||||
|
ple_packed: Buffer::bytes((PLE_HEADS as u64) * 80)?,
|
||||||
|
ple_scales: Buffer::bytes((PLE_HEADS as u64) * 10)?,
|
||||||
|
ple_biases: Buffer::bytes((PLE_HEADS as u64) * 10)?,
|
||||||
|
ple_embedding: Buffer::floats(HIDDEN.into())?,
|
||||||
|
ple_key: Buffer::floats(HC_WIDTH.into())?,
|
||||||
|
ple_value: Buffer::floats(HIDDEN.into())?,
|
||||||
|
ple_gated: Buffer::floats(HC_WIDTH.into())?,
|
||||||
|
ple_norm: Buffer::floats(HC_WIDTH.into())?,
|
||||||
|
ple_output: Buffer::floats(HC_WIDTH.into())?,
|
||||||
logits: Buffer::floats(VOCAB.into())?,
|
logits: Buffer::floats(VOCAB.into())?,
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
struct PleContract {
|
||||||
|
multipliers: [i64; 3],
|
||||||
|
sizes: [i64; PLE_HEADS],
|
||||||
|
offsets: [i64; PLE_HEADS],
|
||||||
|
}
|
||||||
|
|
||||||
|
struct PleState {
|
||||||
|
history: [i32; PLE_HISTORY],
|
||||||
|
conv: Buffer,
|
||||||
|
}
|
||||||
|
|
||||||
pub(in crate::engine) struct QwenExecutor {
|
pub(in crate::engine) struct QwenExecutor {
|
||||||
model: QwenModel,
|
model: QwenModel,
|
||||||
states: Vec<LayerState>,
|
states: Vec<LayerState>,
|
||||||
|
ple_contract: PleContract,
|
||||||
|
ple_state: PleState,
|
||||||
scratch: Scratch,
|
scratch: Scratch,
|
||||||
logits: Vec<f32>,
|
logits: Vec<f32>,
|
||||||
tokens: Vec<i32>,
|
tokens: Vec<i32>,
|
||||||
@@ -127,6 +164,7 @@ pub(in crate::engine) struct QwenExecutor {
|
|||||||
|
|
||||||
pub(in crate::engine) struct QwenResidentState {
|
pub(in crate::engine) struct QwenResidentState {
|
||||||
states: Vec<LayerState>,
|
states: Vec<LayerState>,
|
||||||
|
ple_state: PleState,
|
||||||
logits: Vec<f32>,
|
logits: Vec<f32>,
|
||||||
tokens: Vec<i32>,
|
tokens: Vec<i32>,
|
||||||
position: u32,
|
position: u32,
|
||||||
@@ -135,11 +173,14 @@ pub(in crate::engine) struct QwenResidentState {
|
|||||||
|
|
||||||
impl QwenExecutor {
|
impl QwenExecutor {
|
||||||
pub(super) fn open(model: QwenModel, context: u32) -> Result<Self, String> {
|
pub(super) fn open(model: QwenModel, context: u32) -> Result<Self, String> {
|
||||||
|
let ple_contract = ple_contract(&model)?;
|
||||||
let native = Context::open_qwen(model.memory().admission)?;
|
let native = Context::open_qwen(model.memory().admission)?;
|
||||||
let states = allocate_states(context)?;
|
let states = allocate_states(context)?;
|
||||||
Ok(Self {
|
Ok(Self {
|
||||||
model,
|
model,
|
||||||
states,
|
states,
|
||||||
|
ple_contract,
|
||||||
|
ple_state: allocate_ple_state()?,
|
||||||
scratch: Scratch::new()?,
|
scratch: Scratch::new()?,
|
||||||
logits: vec![0.0; VOCAB as usize],
|
logits: vec![0.0; VOCAB as usize],
|
||||||
tokens: Vec::new(),
|
tokens: Vec::new(),
|
||||||
@@ -160,11 +201,19 @@ impl QwenExecutor {
|
|||||||
self.context
|
self.context
|
||||||
));
|
));
|
||||||
}
|
}
|
||||||
self.require_ple()?;
|
if self.position + 1 > DENSE_BUDGET {
|
||||||
|
return Err(
|
||||||
|
"Qwen sparse QSA selection is required beyond 2048 tokens; native QSA belongs to issue #97"
|
||||||
|
.into(),
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
self.begin_token(token)?;
|
self.begin_token(token)?;
|
||||||
|
|
||||||
for layer in 0..LAYERS {
|
for layer in 0..LAYERS {
|
||||||
|
if layer == 1 {
|
||||||
|
self.ple(token)?;
|
||||||
|
}
|
||||||
self.encode_layer(layer)?;
|
self.encode_layer(layer)?;
|
||||||
}
|
}
|
||||||
self.final_output()?;
|
self.final_output()?;
|
||||||
@@ -211,11 +260,143 @@ impl QwenExecutor {
|
|||||||
commands.finish()
|
commands.finish()
|
||||||
}
|
}
|
||||||
|
|
||||||
fn require_ple(&self) -> Result<(), String> {
|
fn ple(&mut self, token: i32) -> Result<(), String> {
|
||||||
Err(
|
let rows = ple_rows(&self.ple_contract, self.ple_state.history, token)?;
|
||||||
"Qwen PLE injection is required at layer 2; native mapped PLE lookup belongs to issue #96"
|
let packed = self.model.tensor("ngram.weight")?;
|
||||||
.into(),
|
let scales = self.model.tensor("ngram.scales")?;
|
||||||
)
|
let biases = self.model.tensor("ngram.biases")?;
|
||||||
|
let mut packed_stage = [0_u8; PLE_HEADS * 80];
|
||||||
|
let mut scales_stage = [0_u8; PLE_HEADS * 10];
|
||||||
|
let mut biases_stage = [0_u8; PLE_HEADS * 10];
|
||||||
|
let packed_bytes = self.model.tensor_bytes(packed)?;
|
||||||
|
let scales_bytes = self.model.tensor_bytes(scales)?;
|
||||||
|
let biases_bytes = self.model.tensor_bytes(biases)?;
|
||||||
|
for (head, &row) in rows.iter().enumerate() {
|
||||||
|
let bytes = gather_ple_row(row, packed_bytes, scales_bytes, biases_bytes)?;
|
||||||
|
packed_stage[head * 80..][..80].copy_from_slice(&bytes[..80]);
|
||||||
|
scales_stage[head * 10..][..10].copy_from_slice(&bytes[80..90]);
|
||||||
|
biases_stage[head * 10..][..10].copy_from_slice(&bytes[90..]);
|
||||||
|
}
|
||||||
|
self.scratch.ple_packed.write(0, &packed_stage)?;
|
||||||
|
self.scratch.ple_scales.write(0, &scales_stage)?;
|
||||||
|
self.scratch.ple_biases.write(0, &biases_stage)?;
|
||||||
|
|
||||||
|
let commands = Commands::begin()?;
|
||||||
|
let mut dequant = args();
|
||||||
|
dequant.u[0] = PLE_HEAD_DIM;
|
||||||
|
dequant.u[1] = PLE_HEADS as u32;
|
||||||
|
dequant.u[2] = 4;
|
||||||
|
dequant.u[3] = 32;
|
||||||
|
self.dispatch(
|
||||||
|
c"kernel_qwen_ple_dequant",
|
||||||
|
&self.scratch.ple_embedding,
|
||||||
|
Some(&self.scratch.ple_packed),
|
||||||
|
Some(&self.scratch.ple_scales),
|
||||||
|
Some(&self.scratch.ple_biases),
|
||||||
|
&[],
|
||||||
|
&dequant,
|
||||||
|
HIDDEN,
|
||||||
|
1,
|
||||||
|
)?;
|
||||||
|
let prefix = "language_model.model.layers.1.ple";
|
||||||
|
self.bf16_mv(
|
||||||
|
self.weight(&format!("{prefix}.key_proj.weight"))?,
|
||||||
|
&self.scratch.ple_embedding,
|
||||||
|
&self.scratch.ple_key,
|
||||||
|
HIDDEN,
|
||||||
|
HC_WIDTH,
|
||||||
|
)?;
|
||||||
|
self.bf16_mv(
|
||||||
|
self.weight(&format!("{prefix}.value_proj.weight"))?,
|
||||||
|
&self.scratch.ple_embedding,
|
||||||
|
&self.scratch.ple_value,
|
||||||
|
HIDDEN,
|
||||||
|
HIDDEN,
|
||||||
|
)?;
|
||||||
|
for (input, output, name) in [
|
||||||
|
(&self.scratch.ple_key, &self.scratch.ple_key, "norm_key"),
|
||||||
|
(&self.scratch.hc, &self.scratch.hc_norm, "norm_query"),
|
||||||
|
] {
|
||||||
|
let mut norm = args();
|
||||||
|
norm.u[0] = HC_WIDTH;
|
||||||
|
norm.u[1] = HIDDEN;
|
||||||
|
norm.f[0] = 1.0e-6;
|
||||||
|
self.dispatch(
|
||||||
|
c"kernel_qwen_zero_rms",
|
||||||
|
output,
|
||||||
|
Some(input),
|
||||||
|
None,
|
||||||
|
None,
|
||||||
|
&[self.view(self.weight(&format!("{prefix}.{name}.weight"))?)],
|
||||||
|
&norm,
|
||||||
|
HC,
|
||||||
|
1,
|
||||||
|
)?;
|
||||||
|
}
|
||||||
|
let mut gate = args();
|
||||||
|
gate.u[0] = HIDDEN;
|
||||||
|
self.dispatch(
|
||||||
|
c"kernel_qwen_ple_gate",
|
||||||
|
&self.scratch.ple_gated,
|
||||||
|
Some(&self.scratch.ple_key),
|
||||||
|
Some(&self.scratch.hc_norm),
|
||||||
|
Some(&self.scratch.ple_value),
|
||||||
|
&[],
|
||||||
|
&gate,
|
||||||
|
HC,
|
||||||
|
1,
|
||||||
|
)?;
|
||||||
|
let mut norm = args();
|
||||||
|
norm.u[0] = HC_WIDTH;
|
||||||
|
norm.u[1] = HIDDEN;
|
||||||
|
norm.f[0] = 1.0e-6;
|
||||||
|
self.dispatch(
|
||||||
|
c"kernel_qwen_zero_rms",
|
||||||
|
&self.scratch.ple_norm,
|
||||||
|
Some(&self.scratch.ple_gated),
|
||||||
|
None,
|
||||||
|
None,
|
||||||
|
&[self.view(self.weight(&format!("{prefix}.norm_conv.weight"))?)],
|
||||||
|
&norm,
|
||||||
|
HC,
|
||||||
|
1,
|
||||||
|
)?;
|
||||||
|
let mut conv = args();
|
||||||
|
conv.u[0] = HC_WIDTH;
|
||||||
|
self.dispatch(
|
||||||
|
c"kernel_qwen_ple_conv",
|
||||||
|
&self.scratch.ple_output,
|
||||||
|
Some(&self.scratch.ple_gated),
|
||||||
|
Some(&self.scratch.ple_norm),
|
||||||
|
Some(&self.ple_state.conv),
|
||||||
|
&[self.view(self.weight(&format!("{prefix}.conv_weight"))?)],
|
||||||
|
&conv,
|
||||||
|
HC_WIDTH,
|
||||||
|
1,
|
||||||
|
)?;
|
||||||
|
let mut add = args();
|
||||||
|
add.u[0] = HC_WIDTH;
|
||||||
|
self.dispatch(
|
||||||
|
c"kernel_qwen_add",
|
||||||
|
&self.scratch.hc_norm,
|
||||||
|
Some(&self.scratch.hc),
|
||||||
|
Some(&self.scratch.ple_output),
|
||||||
|
None,
|
||||||
|
&[],
|
||||||
|
&add,
|
||||||
|
HC_WIDTH,
|
||||||
|
1,
|
||||||
|
)?;
|
||||||
|
self.scratch.hc.copy_from(
|
||||||
|
0,
|
||||||
|
&self.scratch.hc_norm,
|
||||||
|
0,
|
||||||
|
u64::from(HC_WIDTH) * 4,
|
||||||
|
"committing Qwen PLE injection",
|
||||||
|
)?;
|
||||||
|
commands.finish()?;
|
||||||
|
self.ple_state.history = advance_ple_history(self.ple_state.history, token);
|
||||||
|
Ok(())
|
||||||
}
|
}
|
||||||
|
|
||||||
fn encode_layer(&mut self, layer: usize) -> Result<(), String> {
|
fn encode_layer(&mut self, layer: usize) -> Result<(), String> {
|
||||||
@@ -258,7 +439,15 @@ impl QwenExecutor {
|
|||||||
let commands = Commands::begin()?;
|
let commands = Commands::begin()?;
|
||||||
for (&expert, &weight) in ids.iter().zip(&weights) {
|
for (&expert, &weight) in ids.iter().zip(&weights) {
|
||||||
if !(0..EXPERTS as i32).contains(&expert) || !weight.is_finite() || weight < 0.0 {
|
if !(0..EXPERTS as i32).contains(&expert) || !weight.is_finite() || weight < 0.0 {
|
||||||
return Err("Qwen router produced an invalid top-10 selection".into());
|
let mut router = vec![0.0; EXPERTS as usize];
|
||||||
|
self.scratch.router.read_f32(&mut router)?;
|
||||||
|
let non_finite = router.iter().filter(|value| !value.is_finite()).count();
|
||||||
|
let mut block = vec![0.0; HIDDEN as usize];
|
||||||
|
self.scratch.block.read_f32(&mut block)?;
|
||||||
|
let block_non_finite = block.iter().filter(|value| !value.is_finite()).count();
|
||||||
|
return Err(format!(
|
||||||
|
"Qwen layer {layer} router produced invalid expert {expert} with weight {weight} ({non_finite} non-finite logits, {block_non_finite} non-finite inputs)"
|
||||||
|
));
|
||||||
}
|
}
|
||||||
self.expert(&prefix, expert as u32, weight)?;
|
self.expert(&prefix, expert as u32, weight)?;
|
||||||
}
|
}
|
||||||
@@ -1071,6 +1260,7 @@ impl QwenExecutor {
|
|||||||
|
|
||||||
pub(super) fn reset(&mut self) -> Result<(), String> {
|
pub(super) fn reset(&mut self) -> Result<(), String> {
|
||||||
self.states = allocate_states(self.context)?;
|
self.states = allocate_states(self.context)?;
|
||||||
|
self.ple_state = allocate_ple_state()?;
|
||||||
self.logits.fill(0.0);
|
self.logits.fill(0.0);
|
||||||
self.tokens.clear();
|
self.tokens.clear();
|
||||||
self.position = 0;
|
self.position = 0;
|
||||||
@@ -1088,6 +1278,7 @@ impl QwenExecutor {
|
|||||||
fn blank_resident(&self) -> Result<QwenResidentState, String> {
|
fn blank_resident(&self) -> Result<QwenResidentState, String> {
|
||||||
Ok(QwenResidentState {
|
Ok(QwenResidentState {
|
||||||
states: allocate_states(self.context)?,
|
states: allocate_states(self.context)?,
|
||||||
|
ple_state: allocate_ple_state()?,
|
||||||
logits: vec![0.0; VOCAB as usize],
|
logits: vec![0.0; VOCAB as usize],
|
||||||
tokens: Vec::new(),
|
tokens: Vec::new(),
|
||||||
position: 0,
|
position: 0,
|
||||||
@@ -1101,6 +1292,7 @@ impl QwenExecutor {
|
|||||||
) -> Result<(), String> {
|
) -> Result<(), String> {
|
||||||
let mut incoming = state.take().map_or_else(|| self.blank_resident(), Ok)?;
|
let mut incoming = state.take().map_or_else(|| self.blank_resident(), Ok)?;
|
||||||
std::mem::swap(&mut self.states, &mut incoming.states);
|
std::mem::swap(&mut self.states, &mut incoming.states);
|
||||||
|
std::mem::swap(&mut self.ple_state, &mut incoming.ple_state);
|
||||||
std::mem::swap(&mut self.logits, &mut incoming.logits);
|
std::mem::swap(&mut self.logits, &mut incoming.logits);
|
||||||
std::mem::swap(&mut self.tokens, &mut incoming.tokens);
|
std::mem::swap(&mut self.tokens, &mut incoming.tokens);
|
||||||
std::mem::swap(&mut self.position, &mut incoming.position);
|
std::mem::swap(&mut self.position, &mut incoming.position);
|
||||||
@@ -1134,6 +1326,9 @@ impl QwenExecutor {
|
|||||||
file.write_all(&self.model.checkpoint_identity())
|
file.write_all(&self.model.checkpoint_identity())
|
||||||
.map_err(|error| error.to_string())?;
|
.map_err(|error| error.to_string())?;
|
||||||
file.write_all(&tag).map_err(|error| error.to_string())?;
|
file.write_all(&tag).map_err(|error| error.to_string())?;
|
||||||
|
for token in self.ple_state.history {
|
||||||
|
write_u32(&mut file, token as u32)?;
|
||||||
|
}
|
||||||
for &token in &self.tokens {
|
for &token in &self.tokens {
|
||||||
write_u32(&mut file, token as u32)?;
|
write_u32(&mut file, token as u32)?;
|
||||||
}
|
}
|
||||||
@@ -1141,6 +1336,14 @@ impl QwenExecutor {
|
|||||||
write_u32(&mut file, logit.to_bits())?;
|
write_u32(&mut file, logit.to_bits())?;
|
||||||
}
|
}
|
||||||
let mut chunk = vec![0; CHECKPOINT_CHUNK];
|
let mut chunk = vec![0; CHECKPOINT_CHUNK];
|
||||||
|
write_buffer(
|
||||||
|
&mut file,
|
||||||
|
&self.ple_state.conv,
|
||||||
|
0,
|
||||||
|
u64::from(HC_WIDTH) * PLE_CONV_STATE as u64 * 2,
|
||||||
|
&mut chunk,
|
||||||
|
progress,
|
||||||
|
)?;
|
||||||
for state in &self.states {
|
for state in &self.states {
|
||||||
match state {
|
match state {
|
||||||
LayerState::Gdn { conv, recurrent } => {
|
LayerState::Gdn { conv, recurrent } => {
|
||||||
@@ -1216,6 +1419,14 @@ impl QwenExecutor {
|
|||||||
let mut tag = [0; 32];
|
let mut tag = [0; 32];
|
||||||
file.read_exact(&mut tag)
|
file.read_exact(&mut tag)
|
||||||
.map_err(|error| error.to_string())?;
|
.map_err(|error| error.to_string())?;
|
||||||
|
let mut ple_history = [0; PLE_HISTORY];
|
||||||
|
for token in &mut ple_history {
|
||||||
|
let value = read_u32(&mut file)?;
|
||||||
|
if value >= VOCAB {
|
||||||
|
return Err("Qwen checkpoint PLE history is invalid".into());
|
||||||
|
}
|
||||||
|
*token = value as i32;
|
||||||
|
}
|
||||||
let mut tokens = Vec::with_capacity(position as usize);
|
let mut tokens = Vec::with_capacity(position as usize);
|
||||||
for _ in 0..position {
|
for _ in 0..position {
|
||||||
let token = read_u32(&mut file)?;
|
let token = read_u32(&mut file)?;
|
||||||
@@ -1230,6 +1441,14 @@ impl QwenExecutor {
|
|||||||
}
|
}
|
||||||
self.reset()?;
|
self.reset()?;
|
||||||
let mut chunk = vec![0; CHECKPOINT_CHUNK];
|
let mut chunk = vec![0; CHECKPOINT_CHUNK];
|
||||||
|
read_buffer(
|
||||||
|
&mut file,
|
||||||
|
&self.ple_state.conv,
|
||||||
|
0,
|
||||||
|
u64::from(HC_WIDTH) * PLE_CONV_STATE as u64 * 2,
|
||||||
|
&mut chunk,
|
||||||
|
progress,
|
||||||
|
)?;
|
||||||
for state in &self.states {
|
for state in &self.states {
|
||||||
match state {
|
match state {
|
||||||
LayerState::Gdn { conv, recurrent } => {
|
LayerState::Gdn { conv, recurrent } => {
|
||||||
@@ -1272,6 +1491,7 @@ impl QwenExecutor {
|
|||||||
self.position = position;
|
self.position = position;
|
||||||
self.tokens = tokens;
|
self.tokens = tokens;
|
||||||
self.logits = logits;
|
self.logits = logits;
|
||||||
|
self.ple_state.history = ple_history;
|
||||||
self.checkpoint_tag = tag;
|
self.checkpoint_tag = tag;
|
||||||
Ok(true)
|
Ok(true)
|
||||||
}
|
}
|
||||||
@@ -1299,6 +1519,150 @@ fn allocate_states(context: u32) -> Result<Vec<LayerState>, String> {
|
|||||||
.collect()
|
.collect()
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn allocate_ple_state() -> Result<PleState, String> {
|
||||||
|
let conv = Buffer::bytes(u64::from(HC_WIDTH) * PLE_CONV_STATE as u64 * 2)?;
|
||||||
|
conv.fill(0.0, u64::from(HC_WIDTH) * PLE_CONV_STATE as u64 / 2)?;
|
||||||
|
Ok(PleState {
|
||||||
|
history: [EOS_TOKEN; PLE_HISTORY],
|
||||||
|
conv,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
fn ple_contract(model: &QwenModel) -> Result<PleContract, String> {
|
||||||
|
let multipliers = read_i64_array::<3>(
|
||||||
|
model,
|
||||||
|
"language_model.model.layers.1.ple.ple_embedding.layer_multipliers",
|
||||||
|
)?;
|
||||||
|
let sizes = read_i64_array::<PLE_HEADS>(
|
||||||
|
model,
|
||||||
|
"language_model.model.layers.1.ple.ple_embedding.ngram_heads_vocab_sizes",
|
||||||
|
)?;
|
||||||
|
let offsets = read_i64_array::<PLE_HEADS>(
|
||||||
|
model,
|
||||||
|
"language_model.model.layers.1.ple.ple_embedding.ngram_heads_offsets",
|
||||||
|
)?;
|
||||||
|
let expected_multipliers = official_ple_multipliers();
|
||||||
|
let mut expected_sizes = [0; PLE_HEADS];
|
||||||
|
let mut expected_offsets = [0; PLE_HEADS];
|
||||||
|
let mut total = 0_i64;
|
||||||
|
let mut prime = 19_999_999_i64;
|
||||||
|
for head in 0..PLE_HEADS {
|
||||||
|
prime = next_prime(prime);
|
||||||
|
expected_sizes[head] = prime;
|
||||||
|
expected_offsets[head] = total;
|
||||||
|
total += prime;
|
||||||
|
}
|
||||||
|
if multipliers != expected_multipliers || sizes != expected_sizes || offsets != expected_offsets
|
||||||
|
{
|
||||||
|
return Err("Qwen PLE hash parameters do not match the official contract".into());
|
||||||
|
}
|
||||||
|
for (name, shape, bits, group) in [
|
||||||
|
("ngram.weight", [320_001_536, 20], Some(4), Some(32)),
|
||||||
|
("ngram.scales", [320_001_536, 5], Some(4), Some(32)),
|
||||||
|
("ngram.biases", [320_001_536, 5], Some(4), Some(32)),
|
||||||
|
] {
|
||||||
|
let tensor = model.tensor(name)?;
|
||||||
|
if tensor.shape != shape || tensor.quant_bits != bits || tensor.group_size != group {
|
||||||
|
return Err(format!("{name} does not match the Qwen PLE row layout"));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Ok(PleContract {
|
||||||
|
multipliers,
|
||||||
|
sizes,
|
||||||
|
offsets,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
fn read_i64_array<const N: usize>(model: &QwenModel, name: &str) -> Result<[i64; N], String> {
|
||||||
|
let tensor = model.tensor(name)?;
|
||||||
|
if tensor.dtype != "I64" || tensor.shape != [N as u64] {
|
||||||
|
return Err(format!("{name} does not match the Qwen PLE integer layout"));
|
||||||
|
}
|
||||||
|
let bytes = model.tensor_bytes(tensor)?;
|
||||||
|
if bytes.len() != N * 8 {
|
||||||
|
return Err(format!("{name} has an invalid byte length"));
|
||||||
|
}
|
||||||
|
Ok(std::array::from_fn(|index| {
|
||||||
|
i64::from_le_bytes(bytes[index * 8..index * 8 + 8].try_into().unwrap())
|
||||||
|
}))
|
||||||
|
}
|
||||||
|
|
||||||
|
fn official_ple_multipliers() -> [i64; 3] {
|
||||||
|
const GAMMA: u64 = 0x9e37_79b9_7f4a_7c15;
|
||||||
|
const M1: u64 = 0xbf58_476d_1ce4_e5b9;
|
||||||
|
const M2: u64 = 0x94d0_49bb_1331_11eb;
|
||||||
|
let bound = (i64::MAX / VOCAB as i64 / 2) as u64;
|
||||||
|
std::array::from_fn(|index| {
|
||||||
|
let mut value = 1234_u64.wrapping_add(GAMMA.wrapping_mul(index as u64 + 1));
|
||||||
|
value = value.wrapping_add(GAMMA);
|
||||||
|
value = (value ^ (value >> 30)).wrapping_mul(M1);
|
||||||
|
value = (value ^ (value >> 27)).wrapping_mul(M2);
|
||||||
|
value ^= value >> 31;
|
||||||
|
(2 * (value % bound) + 1) as i64
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
fn next_prime(mut value: i64) -> i64 {
|
||||||
|
loop {
|
||||||
|
value += 1;
|
||||||
|
if value % 2 != 0
|
||||||
|
&& (3..=((value as f64).sqrt() as i64))
|
||||||
|
.step_by(2)
|
||||||
|
.all(|divisor| value % divisor != 0)
|
||||||
|
{
|
||||||
|
return value;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn ple_rows(
|
||||||
|
contract: &PleContract,
|
||||||
|
history: [i32; PLE_HISTORY],
|
||||||
|
token: i32,
|
||||||
|
) -> Result<[u64; PLE_HEADS], String> {
|
||||||
|
if token < 0 || token as u32 >= VOCAB {
|
||||||
|
return Err(format!("token {token} is outside the Qwen vocabulary"));
|
||||||
|
}
|
||||||
|
let shifted = [token as i64, history[1] as i64, history[0] as i64];
|
||||||
|
let two = shifted[0].wrapping_mul(contract.multipliers[0])
|
||||||
|
^ shifted[1].wrapping_mul(contract.multipliers[1]);
|
||||||
|
let three = two ^ shifted[2].wrapping_mul(contract.multipliers[2]);
|
||||||
|
Ok(std::array::from_fn(|head| {
|
||||||
|
let mixed = if head < 8 { two } else { three };
|
||||||
|
(contract.offsets[head] + mixed.rem_euclid(contract.sizes[head])) as u64
|
||||||
|
}))
|
||||||
|
}
|
||||||
|
|
||||||
|
fn advance_ple_history(history: [i32; PLE_HISTORY], token: i32) -> [i32; PLE_HISTORY] {
|
||||||
|
if token == EOS_TOKEN {
|
||||||
|
[EOS_TOKEN; PLE_HISTORY]
|
||||||
|
} else {
|
||||||
|
[history[1], token]
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn gather_ple_row(
|
||||||
|
row: u64,
|
||||||
|
packed: &[u8],
|
||||||
|
scales: &[u8],
|
||||||
|
biases: &[u8],
|
||||||
|
) -> Result<[u8; PLE_ROW_BYTES], String> {
|
||||||
|
let row = usize::try_from(row).map_err(|_| "Qwen PLE row exceeds this platform")?;
|
||||||
|
let mut value = [0; PLE_ROW_BYTES];
|
||||||
|
value[..80].copy_from_slice(ple_row_slice(packed, row, 80)?);
|
||||||
|
value[80..90].copy_from_slice(ple_row_slice(scales, row, 10)?);
|
||||||
|
value[90..].copy_from_slice(ple_row_slice(biases, row, 10)?);
|
||||||
|
Ok(value)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn ple_row_slice(data: &[u8], row: usize, width: usize) -> Result<&[u8], String> {
|
||||||
|
let start = row
|
||||||
|
.checked_mul(width)
|
||||||
|
.ok_or_else(|| "Qwen PLE row offset overflows".to_owned())?;
|
||||||
|
data.get(start..start + width)
|
||||||
|
.ok_or_else(|| format!("Qwen PLE row {row} is truncated"))
|
||||||
|
}
|
||||||
|
|
||||||
fn args() -> QwenKernelArgs {
|
fn args() -> QwenKernelArgs {
|
||||||
QwenKernelArgs::default()
|
QwenKernelArgs::default()
|
||||||
}
|
}
|
||||||
@@ -1338,6 +1702,7 @@ fn dispatch_qwen(
|
|||||||
mod tests {
|
mod tests {
|
||||||
use super::*;
|
use super::*;
|
||||||
use memmap2::MmapOptions;
|
use memmap2::MmapOptions;
|
||||||
|
use sha2::{Digest, Sha256};
|
||||||
use std::fs;
|
use std::fs;
|
||||||
use std::path::PathBuf;
|
use std::path::PathBuf;
|
||||||
|
|
||||||
@@ -1353,6 +1718,145 @@ mod tests {
|
|||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn qwen_ple_hash_contract_matches_golden_boundaries() {
|
||||||
|
assert_eq!(
|
||||||
|
official_ple_multipliers(),
|
||||||
|
[23_703_573_157_769, 20_109_073_645_365, 8_052_911_324_071]
|
||||||
|
);
|
||||||
|
let mut contract = PleContract {
|
||||||
|
multipliers: official_ple_multipliers(),
|
||||||
|
sizes: [0; PLE_HEADS],
|
||||||
|
offsets: [0; PLE_HEADS],
|
||||||
|
};
|
||||||
|
let mut prime = 19_999_999;
|
||||||
|
let mut offset = 0;
|
||||||
|
for head in 0..PLE_HEADS {
|
||||||
|
prime = next_prime(prime);
|
||||||
|
contract.sizes[head] = prime;
|
||||||
|
contract.offsets[head] = offset;
|
||||||
|
offset += prime;
|
||||||
|
}
|
||||||
|
let initial = ple_rows(&contract, [EOS_TOKEN; 2], 1).unwrap();
|
||||||
|
let repeated = ple_rows(&contract, [1, 1], 1).unwrap();
|
||||||
|
let boundary = ple_rows(&contract, [EOS_TOKEN, 42], 43).unwrap();
|
||||||
|
assert_eq!(
|
||||||
|
initial,
|
||||||
|
[
|
||||||
|
16_121_432,
|
||||||
|
28_938_500,
|
||||||
|
59_087_997,
|
||||||
|
73_487_090,
|
||||||
|
81_148_277,
|
||||||
|
104_500_129,
|
||||||
|
120_276_032,
|
||||||
|
149_373_875,
|
||||||
|
176_283_436,
|
||||||
|
184_305_849,
|
||||||
|
216_528_839,
|
||||||
|
231_080_079,
|
||||||
|
257_961_536,
|
||||||
|
266_068_568,
|
||||||
|
289_043_455,
|
||||||
|
305_959_965,
|
||||||
|
]
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
repeated,
|
||||||
|
[
|
||||||
|
6_868_091,
|
||||||
|
38_325_817,
|
||||||
|
54_054_700,
|
||||||
|
68_075_137,
|
||||||
|
82_949_816,
|
||||||
|
101_241_419,
|
||||||
|
138_678_867,
|
||||||
|
155_262_032,
|
||||||
|
176_541_251,
|
||||||
|
196_154_476,
|
||||||
|
215_703_237,
|
||||||
|
234_413_824,
|
||||||
|
254_220_543,
|
||||||
|
274_027_268,
|
||||||
|
293_962_951,
|
||||||
|
313_640_732,
|
||||||
|
]
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
boundary,
|
||||||
|
[
|
||||||
|
18_529_343,
|
||||||
|
23_547_650,
|
||||||
|
56_056_978,
|
||||||
|
73_570_159,
|
||||||
|
88_581_601,
|
||||||
|
113_585_506,
|
||||||
|
121_091_299,
|
||||||
|
151_099_148,
|
||||||
|
175_585_266,
|
||||||
|
184_439_538,
|
||||||
|
216_587_431,
|
||||||
|
222_082_137,
|
||||||
|
250_284_866,
|
||||||
|
278_847_169,
|
||||||
|
281_781_121,
|
||||||
|
317_050_322,
|
||||||
|
]
|
||||||
|
);
|
||||||
|
assert_eq!(contract.offsets[0], 0);
|
||||||
|
assert_eq!(contract.offsets[8], 160_000_374);
|
||||||
|
assert_eq!(contract.offsets[15] + contract.sizes[15], 320_001_446);
|
||||||
|
for (head, &row) in initial.iter().enumerate() {
|
||||||
|
assert!(
|
||||||
|
(contract.offsets[head]..contract.offsets[head] + contract.sizes[head])
|
||||||
|
.contains(&(row as i64))
|
||||||
|
);
|
||||||
|
}
|
||||||
|
assert_eq!(320_001_536 % 128, 0);
|
||||||
|
let history = advance_ple_history([EOS_TOKEN; 2], 1);
|
||||||
|
assert_eq!(advance_ple_history(history, 2), [1, 2]);
|
||||||
|
assert_eq!(advance_ple_history([1, 2], EOS_TOKEN), [EOS_TOKEN; 2]);
|
||||||
|
|
||||||
|
let hash_chunks = |chunk: usize| {
|
||||||
|
let mut history = [EOS_TOKEN; PLE_HISTORY];
|
||||||
|
[1, 2, EOS_TOKEN, 3, 4]
|
||||||
|
.chunks(chunk)
|
||||||
|
.flat_map(|tokens| {
|
||||||
|
tokens
|
||||||
|
.iter()
|
||||||
|
.map(|&token| {
|
||||||
|
let rows = ple_rows(&contract, history, token).unwrap();
|
||||||
|
history = advance_ple_history(history, token);
|
||||||
|
rows
|
||||||
|
})
|
||||||
|
.collect::<Vec<_>>()
|
||||||
|
})
|
||||||
|
.collect::<Vec<_>>()
|
||||||
|
};
|
||||||
|
assert_eq!(hash_chunks(1), hash_chunks(2));
|
||||||
|
assert_eq!(hash_chunks(1), hash_chunks(5));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn qwen_ple_gather_copies_only_the_requested_row() {
|
||||||
|
let packed = (0..240).map(|value| value as u8).collect::<Vec<_>>();
|
||||||
|
let scales = (0..30).map(|value| (value + 17) as u8).collect::<Vec<_>>();
|
||||||
|
let biases = (0..30).map(|value| (value + 47) as u8).collect::<Vec<_>>();
|
||||||
|
let direct = gather_ple_row(1, &packed, &scales, &biases).unwrap();
|
||||||
|
assert_eq!(&direct[..80], &packed[80..160]);
|
||||||
|
assert_eq!(&direct[80..90], &scales[10..20]);
|
||||||
|
assert_eq!(&direct[90..], &biases[10..20]);
|
||||||
|
assert_eq!(
|
||||||
|
gather_ple_row(1, &packed, &scales, &biases).unwrap(),
|
||||||
|
direct
|
||||||
|
);
|
||||||
|
assert!(
|
||||||
|
gather_ple_row(3, &packed, &scales, &biases)
|
||||||
|
.unwrap_err()
|
||||||
|
.contains("truncated")
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
#[ignore = "requires Apple Metal"]
|
#[ignore = "requires Apple Metal"]
|
||||||
fn qwen_metal_primitives_match_reference_vectors() {
|
fn qwen_metal_primitives_match_reference_vectors() {
|
||||||
@@ -1363,11 +1867,12 @@ mod tests {
|
|||||||
input.write_f32(&vec![1.0; 64]).unwrap();
|
input.write_f32(&vec![1.0; 64]).unwrap();
|
||||||
let output = Buffer::floats(64).unwrap();
|
let output = Buffer::floats(64).unwrap();
|
||||||
let path = std::env::temp_dir().join(format!(
|
let path = std::env::temp_dir().join(format!(
|
||||||
"ds4-qwen95-weights-{}-{}",
|
"ds4-qwen96-weights-{}-{}",
|
||||||
std::process::id(),
|
std::process::id(),
|
||||||
std::thread::current().name().unwrap_or("test")
|
std::thread::current().name().unwrap_or("test")
|
||||||
));
|
));
|
||||||
let mut bytes = Vec::new();
|
// Safetensors headers are not guaranteed to align the data region.
|
||||||
|
let mut bytes = vec![0];
|
||||||
for _ in 0..8 {
|
for _ in 0..8 {
|
||||||
bytes.extend_from_slice(&0x3333_3333_u32.to_le_bytes());
|
bytes.extend_from_slice(&0x3333_3333_u32.to_le_bytes());
|
||||||
}
|
}
|
||||||
@@ -1398,7 +1903,7 @@ mod tests {
|
|||||||
Some(&input),
|
Some(&input),
|
||||||
None,
|
None,
|
||||||
None,
|
None,
|
||||||
&[view(0, 32), view(32, 2), view(34, 2)],
|
&[view(1, 32), view(33, 2), view(35, 2)],
|
||||||
&affine,
|
&affine,
|
||||||
1,
|
1,
|
||||||
1,
|
1,
|
||||||
@@ -1421,7 +1926,7 @@ mod tests {
|
|||||||
Some(&conv_input),
|
Some(&conv_input),
|
||||||
Some(&conv_state),
|
Some(&conv_state),
|
||||||
None,
|
None,
|
||||||
&[view(40, 8)],
|
&[view(41, 8)],
|
||||||
&conv,
|
&conv,
|
||||||
1,
|
1,
|
||||||
1,
|
1,
|
||||||
@@ -1459,7 +1964,7 @@ mod tests {
|
|||||||
Some(&qkv),
|
Some(&qkv),
|
||||||
Some(&controls),
|
Some(&controls),
|
||||||
Some(&recurrent),
|
Some(&recurrent),
|
||||||
&[view(36, 2), view(38, 2)],
|
&[view(37, 2), view(39, 2)],
|
||||||
&step,
|
&step,
|
||||||
2,
|
2,
|
||||||
1,
|
1,
|
||||||
@@ -1479,7 +1984,7 @@ mod tests {
|
|||||||
Some(&raw),
|
Some(&raw),
|
||||||
Some(&controls),
|
Some(&controls),
|
||||||
None,
|
None,
|
||||||
&[view(48, 4)],
|
&[view(49, 4)],
|
||||||
&norm,
|
&norm,
|
||||||
1,
|
1,
|
||||||
1,
|
1,
|
||||||
@@ -1533,7 +2038,7 @@ mod tests {
|
|||||||
Some(&norm_input),
|
Some(&norm_input),
|
||||||
None,
|
None,
|
||||||
None,
|
None,
|
||||||
&[view(52, 8)],
|
&[view(53, 8)],
|
||||||
&zero_norm,
|
&zero_norm,
|
||||||
2,
|
2,
|
||||||
1,
|
1,
|
||||||
@@ -1670,7 +2175,7 @@ mod tests {
|
|||||||
Some(&rope_input),
|
Some(&rope_input),
|
||||||
None,
|
None,
|
||||||
None,
|
None,
|
||||||
&[view(52, 8)],
|
&[view(53, 8)],
|
||||||
&rope,
|
&rope,
|
||||||
4,
|
4,
|
||||||
1,
|
1,
|
||||||
@@ -1749,6 +2254,94 @@ mod tests {
|
|||||||
probability * 3.0 + (1.0 - probability) * 7.0,
|
probability * 3.0 + (1.0 - probability) * 7.0,
|
||||||
);
|
);
|
||||||
|
|
||||||
|
let ple_packed = Buffer::bytes(80).unwrap();
|
||||||
|
ple_packed.write(0, &[0x33; 80]).unwrap();
|
||||||
|
let ple_scales = Buffer::bytes(10).unwrap();
|
||||||
|
ple_scales
|
||||||
|
.write(0, &bf16(0.5).to_le_bytes().repeat(5))
|
||||||
|
.unwrap();
|
||||||
|
let ple_biases = Buffer::bytes(10).unwrap();
|
||||||
|
ple_biases
|
||||||
|
.write(0, &bf16(-1.0).to_le_bytes().repeat(5))
|
||||||
|
.unwrap();
|
||||||
|
let ple_embedding = Buffer::floats(160).unwrap();
|
||||||
|
let mut dequant = args();
|
||||||
|
dequant.u[0] = 160;
|
||||||
|
dequant.u[1] = 1;
|
||||||
|
dequant.u[2] = 4;
|
||||||
|
dequant.u[3] = 32;
|
||||||
|
dispatch_qwen(
|
||||||
|
c"kernel_qwen_ple_dequant",
|
||||||
|
&ple_embedding,
|
||||||
|
Some(&ple_packed),
|
||||||
|
Some(&ple_scales),
|
||||||
|
Some(&ple_biases),
|
||||||
|
&[],
|
||||||
|
&dequant,
|
||||||
|
160,
|
||||||
|
1,
|
||||||
|
)
|
||||||
|
.unwrap();
|
||||||
|
let mut embedding = [0.0; 160];
|
||||||
|
ple_embedding.read_f32(&mut embedding).unwrap();
|
||||||
|
assert!(embedding.into_iter().all(|value| value == 0.5));
|
||||||
|
|
||||||
|
let ple_key = Buffer::floats(8).unwrap();
|
||||||
|
ple_key.write_f32(&[1.0; 8]).unwrap();
|
||||||
|
let ple_query = Buffer::floats(8).unwrap();
|
||||||
|
ple_query.write_f32(&[1.0; 8]).unwrap();
|
||||||
|
let ple_value = Buffer::floats(2).unwrap();
|
||||||
|
ple_value.write_f32(&[2.0, 3.0]).unwrap();
|
||||||
|
let ple_gated = Buffer::floats(8).unwrap();
|
||||||
|
let mut gate = args();
|
||||||
|
gate.u[0] = 2;
|
||||||
|
dispatch_qwen(
|
||||||
|
c"kernel_qwen_ple_gate",
|
||||||
|
&ple_gated,
|
||||||
|
Some(&ple_key),
|
||||||
|
Some(&ple_query),
|
||||||
|
Some(&ple_value),
|
||||||
|
&[],
|
||||||
|
&gate,
|
||||||
|
4,
|
||||||
|
1,
|
||||||
|
)
|
||||||
|
.unwrap();
|
||||||
|
let expected_gate = 1.0 / (1.0 + (-2.0_f32.sqrt().sqrt()).exp());
|
||||||
|
let mut gated = [0.0; 8];
|
||||||
|
ple_gated.read_f32(&mut gated).unwrap();
|
||||||
|
for stream in 0..4 {
|
||||||
|
close(gated[stream * 2], 2.0 * expected_gate);
|
||||||
|
close(gated[stream * 2 + 1], 3.0 * expected_gate);
|
||||||
|
}
|
||||||
|
|
||||||
|
let ple_normalized = Buffer::floats(1).unwrap();
|
||||||
|
ple_normalized.write_f32(&[2.0]).unwrap();
|
||||||
|
let ple_gate_value = Buffer::floats(1).unwrap();
|
||||||
|
ple_gate_value.write_f32(&[0.5]).unwrap();
|
||||||
|
let ple_state = Buffer::bytes(18).unwrap();
|
||||||
|
ple_state.write(0, &[0; 18]).unwrap();
|
||||||
|
let ple_output = Buffer::floats(1).unwrap();
|
||||||
|
let mut conv = args();
|
||||||
|
conv.u[0] = 1;
|
||||||
|
dispatch_qwen(
|
||||||
|
c"kernel_qwen_ple_conv",
|
||||||
|
&ple_output,
|
||||||
|
Some(&ple_gate_value),
|
||||||
|
Some(&ple_normalized),
|
||||||
|
Some(&ple_state),
|
||||||
|
&[view(41, 8)],
|
||||||
|
&conv,
|
||||||
|
1,
|
||||||
|
1,
|
||||||
|
)
|
||||||
|
.unwrap();
|
||||||
|
ple_output.read_f32(&mut scalar).unwrap();
|
||||||
|
close(scalar[0], 0.5 + 8.0 / (1.0 + (-8.0_f32).exp()));
|
||||||
|
let mut ple_history = [0; 18];
|
||||||
|
ple_state.read(0, &mut ple_history).unwrap();
|
||||||
|
assert_eq!(&ple_history[16..], &bf16(2.0).to_le_bytes());
|
||||||
|
|
||||||
drop(map);
|
drop(map);
|
||||||
drop(file);
|
drop(file);
|
||||||
fs::remove_file(path).unwrap();
|
fs::remove_file(path).unwrap();
|
||||||
@@ -1762,35 +2355,53 @@ mod tests {
|
|||||||
.map(PathBuf::from)
|
.map(PathBuf::from)
|
||||||
.expect("set DS4SERVER_QWEN38_SOURCE to the pinned artifact directory");
|
.expect("set DS4SERVER_QWEN38_SOURCE to the pinned artifact directory");
|
||||||
let model = QwenModel::open(&root, 4).unwrap();
|
let model = QwenModel::open(&root, 4).unwrap();
|
||||||
|
let residency_before = model.mapped_residency().unwrap();
|
||||||
let mut executor = QwenExecutor::open(model, 4).unwrap();
|
let mut executor = QwenExecutor::open(model, 4).unwrap();
|
||||||
let direct = executor.eval(1).unwrap_err();
|
executor.eval(1).unwrap();
|
||||||
let batch = executor.prefill(&[1], |_| true).unwrap_err();
|
let residency_after = executor.model().mapped_residency().unwrap();
|
||||||
assert_eq!(direct, batch);
|
assert!(residency_after.0 <= executor.model().memory().resident_core);
|
||||||
assert_eq!(executor.position, 0);
|
assert!(residency_after.1 <= executor.model().memory().mapped_ple);
|
||||||
assert!(executor.tokens.is_empty());
|
eprintln!(
|
||||||
|
"Qwen mapped residency core/PLE before {:?}, after {:?}",
|
||||||
executor.begin_token(1).unwrap();
|
residency_before, residency_after
|
||||||
executor.encode_layer(0).unwrap();
|
);
|
||||||
executor.final_output().unwrap();
|
|
||||||
assert!(executor.logits.iter().all(|value| value.is_finite()));
|
assert!(executor.logits.iter().all(|value| value.is_finite()));
|
||||||
|
let mut digest = Sha256::new();
|
||||||
|
for value in &executor.logits {
|
||||||
|
digest.update(value.to_bits().to_le_bytes());
|
||||||
|
}
|
||||||
|
let digest: [u8; 32] = digest.finalize().into();
|
||||||
|
assert_eq!(
|
||||||
|
digest,
|
||||||
|
[
|
||||||
|
137, 244, 133, 253, 201, 214, 196, 144, 249, 130, 28, 63, 124, 75, 32, 40, 16, 148,
|
||||||
|
148, 123, 5, 50, 90, 165, 101, 44, 223, 164, 62, 54, 164, 86,
|
||||||
|
]
|
||||||
|
);
|
||||||
let reference = [
|
let reference = [
|
||||||
executor.logits[0],
|
executor.logits[0],
|
||||||
executor.logits[1],
|
executor.logits[1],
|
||||||
executor.logits[1000],
|
executor.logits[1000],
|
||||||
executor.logits[VOCAB as usize - 1],
|
executor.logits[VOCAB as usize - 1],
|
||||||
];
|
];
|
||||||
|
let next = executor
|
||||||
|
.logits
|
||||||
|
.iter()
|
||||||
|
.enumerate()
|
||||||
|
.max_by(|a, b| a.1.total_cmp(b.1))
|
||||||
|
.unwrap()
|
||||||
|
.0 as i32;
|
||||||
for (actual, expected) in
|
for (actual, expected) in
|
||||||
reference
|
reference
|
||||||
.into_iter()
|
.into_iter()
|
||||||
.zip([1.483_976_1, 0.575_980_66, -0.008_828_48, 0.047_039_207])
|
.zip([6.406_557, 2.082_818_3, -3.058_045_6, -0.121_010_3])
|
||||||
{
|
{
|
||||||
close(actual, expected);
|
close(actual, expected);
|
||||||
}
|
}
|
||||||
|
assert_eq!(next, 89_648);
|
||||||
|
|
||||||
executor.position = 1;
|
|
||||||
executor.tokens = vec![1];
|
|
||||||
let checkpoint = std::env::temp_dir().join(format!(
|
let checkpoint = std::env::temp_dir().join(format!(
|
||||||
"ds4-qwen95-checkpoint-{}-{}",
|
"ds4-qwen96-checkpoint-{}-{}",
|
||||||
std::process::id(),
|
std::process::id(),
|
||||||
std::thread::current().name().unwrap_or("test")
|
std::thread::current().name().unwrap_or("test")
|
||||||
));
|
));
|
||||||
@@ -1798,9 +2409,7 @@ mod tests {
|
|||||||
.save_checkpoint(&checkpoint, [7; 32], &mut |_| {})
|
.save_checkpoint(&checkpoint, [7; 32], &mut |_| {})
|
||||||
.unwrap();
|
.unwrap();
|
||||||
|
|
||||||
executor.begin_token(2).unwrap();
|
executor.eval(next).unwrap();
|
||||||
executor.encode_layer(0).unwrap();
|
|
||||||
executor.final_output().unwrap();
|
|
||||||
let continued = [
|
let continued = [
|
||||||
executor.logits[0],
|
executor.logits[0],
|
||||||
executor.logits[1],
|
executor.logits[1],
|
||||||
@@ -1814,10 +2423,13 @@ mod tests {
|
|||||||
conv.read(0, &mut continued_conv).unwrap();
|
conv.read(0, &mut continued_conv).unwrap();
|
||||||
let mut continued_recurrent = vec![0; 4 * 1024];
|
let mut continued_recurrent = vec![0; 4 * 1024];
|
||||||
recurrent.read(0, &mut continued_recurrent).unwrap();
|
recurrent.read(0, &mut continued_recurrent).unwrap();
|
||||||
|
let mut continued_ple = vec![0; HC_WIDTH as usize * PLE_CONV_STATE as usize * 2];
|
||||||
|
executor.ple_state.conv.read(0, &mut continued_ple).unwrap();
|
||||||
|
|
||||||
assert!(executor.load_checkpoint(&checkpoint, &mut |_| {}).unwrap());
|
assert!(executor.load_checkpoint(&checkpoint, &mut |_| {}).unwrap());
|
||||||
assert_eq!(executor.position, 1);
|
assert_eq!(executor.position, 1);
|
||||||
assert_eq!(executor.tokens, [1]);
|
assert_eq!(executor.tokens, [1]);
|
||||||
|
assert_eq!(executor.ple_state.history, [EOS_TOKEN, 1]);
|
||||||
assert_eq!(executor.checkpoint_tag, [7; 32]);
|
assert_eq!(executor.checkpoint_tag, [7; 32]);
|
||||||
assert_eq!(
|
assert_eq!(
|
||||||
[
|
[
|
||||||
@@ -1828,9 +2440,7 @@ mod tests {
|
|||||||
],
|
],
|
||||||
reference
|
reference
|
||||||
);
|
);
|
||||||
executor.begin_token(2).unwrap();
|
executor.eval(next).unwrap();
|
||||||
executor.encode_layer(0).unwrap();
|
|
||||||
executor.final_output().unwrap();
|
|
||||||
for (actual, expected) in [
|
for (actual, expected) in [
|
||||||
executor.logits[0],
|
executor.logits[0],
|
||||||
executor.logits[1],
|
executor.logits[1],
|
||||||
@@ -1851,11 +2461,36 @@ mod tests {
|
|||||||
let mut resumed_recurrent = vec![0; continued_recurrent.len()];
|
let mut resumed_recurrent = vec![0; continued_recurrent.len()];
|
||||||
recurrent.read(0, &mut resumed_recurrent).unwrap();
|
recurrent.read(0, &mut resumed_recurrent).unwrap();
|
||||||
assert_eq!(resumed_recurrent, continued_recurrent);
|
assert_eq!(resumed_recurrent, continued_recurrent);
|
||||||
|
let mut resumed_ple = vec![0; continued_ple.len()];
|
||||||
|
executor.ple_state.conv.read(0, &mut resumed_ple).unwrap();
|
||||||
|
assert_eq!(resumed_ple, continued_ple);
|
||||||
|
|
||||||
|
executor.reset().unwrap();
|
||||||
|
assert_eq!(executor.position, 0);
|
||||||
|
assert_eq!(executor.ple_state.history, [EOS_TOKEN; PLE_HISTORY]);
|
||||||
|
let mut reset_ple = vec![1; continued_ple.len()];
|
||||||
|
executor.ple_state.conv.read(0, &mut reset_ple).unwrap();
|
||||||
|
assert!(reset_ple.into_iter().all(|byte| byte == 0));
|
||||||
|
executor.eval(1).unwrap();
|
||||||
|
for (actual, expected) in [
|
||||||
|
executor.logits[0],
|
||||||
|
executor.logits[1],
|
||||||
|
executor.logits[1000],
|
||||||
|
executor.logits[VOCAB as usize - 1],
|
||||||
|
]
|
||||||
|
.into_iter()
|
||||||
|
.zip(reference)
|
||||||
|
{
|
||||||
|
close(actual, expected);
|
||||||
|
}
|
||||||
|
|
||||||
|
executor.reset().unwrap();
|
||||||
|
executor.eval(EOS_TOKEN).unwrap();
|
||||||
|
assert_eq!(executor.ple_state.history, [EOS_TOKEN; PLE_HISTORY]);
|
||||||
|
|
||||||
executor.position = DENSE_BUDGET;
|
executor.position = DENSE_BUDGET;
|
||||||
let error = executor
|
executor.context = DENSE_BUDGET + 1;
|
||||||
.attention("language_model.model.layers.3", 3)
|
let error = executor.eval(1).unwrap_err();
|
||||||
.unwrap_err();
|
|
||||||
assert!(error.contains("issue #97"));
|
assert!(error.contains("issue #97"));
|
||||||
fs::remove_file(checkpoint).unwrap();
|
fs::remove_file(checkpoint).unwrap();
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -23,6 +23,7 @@ const MTP_BYTES: u64 = 1_672_575_532;
|
|||||||
const KV_BYTES_PER_TOKEN: u64 = 24_576;
|
const KV_BYTES_PER_TOKEN: u64 = 24_576;
|
||||||
const GDN_STATE_BYTES: u64 = 113_246_208;
|
const GDN_STATE_BYTES: u64 = 113_246_208;
|
||||||
const GDN_CONV_BYTES: u64 = 2_211_840;
|
const GDN_CONV_BYTES: u64 = 2_211_840;
|
||||||
|
const PLE_CONV_BYTES: u64 = 184_320;
|
||||||
|
|
||||||
#[derive(Deserialize)]
|
#[derive(Deserialize)]
|
||||||
struct Manifest {
|
struct Manifest {
|
||||||
@@ -120,9 +121,9 @@ pub(super) struct QwenModel {
|
|||||||
impl QwenModel {
|
impl QwenModel {
|
||||||
pub(super) fn open(root: &Path, context: u32) -> Result<Self, String> {
|
pub(super) fn open(root: &Path, context: u32) -> Result<Self, String> {
|
||||||
let loaded = load(root, context, false)?;
|
let loaded = load(root, context, false)?;
|
||||||
let mut paths = loaded
|
let mut bindings = loaded.bindings.core;
|
||||||
.bindings
|
bindings.extend(loaded.bindings.ple);
|
||||||
.core
|
let mut paths = bindings
|
||||||
.iter()
|
.iter()
|
||||||
.map(|binding| binding.file.clone())
|
.map(|binding| binding.file.clone())
|
||||||
.collect::<Vec<_>>();
|
.collect::<Vec<_>>();
|
||||||
@@ -138,9 +139,7 @@ impl QwenModel {
|
|||||||
map_indices.insert(path.clone(), maps.len());
|
map_indices.insert(path.clone(), maps.len());
|
||||||
maps.push(QwenMap { path, map });
|
maps.push(QwenMap { path, map });
|
||||||
}
|
}
|
||||||
let tensors = loaded
|
let tensors = bindings
|
||||||
.bindings
|
|
||||||
.core
|
|
||||||
.into_iter()
|
.into_iter()
|
||||||
.map(|binding| {
|
.map(|binding| {
|
||||||
let tensor = QwenTensor {
|
let tensor = QwenTensor {
|
||||||
@@ -178,6 +177,16 @@ impl QwenModel {
|
|||||||
(&self.maps[index].map, &self.maps[index].path)
|
(&self.maps[index].map, &self.maps[index].path)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
pub(super) fn tensor_bytes<'a>(&'a self, tensor: &QwenTensor) -> Result<&'a [u8], String> {
|
||||||
|
let map = &self.maps[tensor.map].map;
|
||||||
|
let start = usize::try_from(tensor.range.start)
|
||||||
|
.map_err(|_| format!("{} starts beyond this platform", tensor.name))?;
|
||||||
|
let end = usize::try_from(tensor.range.end)
|
||||||
|
.map_err(|_| format!("{} ends beyond this platform", tensor.name))?;
|
||||||
|
map.get(start..end)
|
||||||
|
.ok_or_else(|| format!("{} is outside its mapped artifact", tensor.name))
|
||||||
|
}
|
||||||
|
|
||||||
pub(super) fn checkpoint_identity(&self) -> [u8; 32] {
|
pub(super) fn checkpoint_identity(&self) -> [u8; 32] {
|
||||||
self.identity
|
self.identity
|
||||||
}
|
}
|
||||||
@@ -243,6 +252,48 @@ impl QwenModel {
|
|||||||
pub(super) fn memory(&self) -> &MemoryPlan {
|
pub(super) fn memory(&self) -> &MemoryPlan {
|
||||||
&self.memory
|
&self.memory
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
pub(super) fn mapped_residency(&self) -> Result<(u64, u64), String> {
|
||||||
|
// SAFETY: sysconf is read-only and has no pointer preconditions.
|
||||||
|
let page = unsafe { libc::sysconf(libc::_SC_PAGESIZE) };
|
||||||
|
if page <= 0 {
|
||||||
|
return Err("macOS did not report its virtual-memory page size".into());
|
||||||
|
}
|
||||||
|
let page = page as usize;
|
||||||
|
let mut core = 0_u64;
|
||||||
|
let mut ple = 0_u64;
|
||||||
|
for item in &self.maps {
|
||||||
|
let mut pages = vec![0_i8; item.map.len().div_ceil(page)];
|
||||||
|
// SAFETY: each read-only mmap and residency vector remain valid for this call.
|
||||||
|
if unsafe {
|
||||||
|
libc::mincore(
|
||||||
|
item.map.as_ptr().cast_mut().cast(),
|
||||||
|
item.map.len(),
|
||||||
|
pages.as_mut_ptr(),
|
||||||
|
)
|
||||||
|
} != 0
|
||||||
|
{
|
||||||
|
return Err(format!(
|
||||||
|
"cannot inspect residency for {}: {}",
|
||||||
|
item.path.display(),
|
||||||
|
std::io::Error::last_os_error()
|
||||||
|
));
|
||||||
|
}
|
||||||
|
let bytes = (pages.iter().filter(|value| **value & 1 != 0).count() * page)
|
||||||
|
.min(item.map.len()) as u64;
|
||||||
|
if item
|
||||||
|
.path
|
||||||
|
.file_name()
|
||||||
|
.is_some_and(|name| name == "ngram-table.safetensors")
|
||||||
|
{
|
||||||
|
ple += bytes;
|
||||||
|
} else {
|
||||||
|
core += bytes;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Ok((core, ple))
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
pub(crate) fn validate_artifacts(root: &Path) -> Result<(), String> {
|
pub(crate) fn validate_artifacts(root: &Path) -> Result<(), String> {
|
||||||
@@ -587,7 +638,7 @@ pub(super) fn memory_plan(
|
|||||||
.checked_mul(u64::from(context))
|
.checked_mul(u64::from(context))
|
||||||
.ok_or_else(|| "Qwen KV memory size overflows".to_owned())?;
|
.ok_or_else(|| "Qwen KV memory size overflows".to_owned())?;
|
||||||
let kv_and_recurrent = kv
|
let kv_and_recurrent = kv
|
||||||
.checked_add(GDN_STATE_BYTES + GDN_CONV_BYTES)
|
.checked_add(GDN_STATE_BYTES + GDN_CONV_BYTES + PLE_CONV_BYTES)
|
||||||
.ok_or_else(|| "Qwen recurrent memory size overflows".to_owned())?;
|
.ok_or_else(|| "Qwen recurrent memory size overflows".to_owned())?;
|
||||||
let prefill_transient = u64::from(prefill_chunk)
|
let prefill_transient = u64::from(prefill_chunk)
|
||||||
.checked_mul((4 * 2_560 + 2_048 + 2_048 + 6_144 + 6_144) * 2)
|
.checked_mul((4 * 2_560 + 2_048 + 2_048 + 6_144 + 6_144) * 2)
|
||||||
@@ -630,12 +681,12 @@ mod tests {
|
|||||||
assert_eq!(plan.resident_core, CORE_BYTES);
|
assert_eq!(plan.resident_core, CORE_BYTES);
|
||||||
assert_eq!(plan.mapped_ple, PLE_BYTES);
|
assert_eq!(plan.mapped_ple, PLE_BYTES);
|
||||||
assert_eq!(plan.optional_mtp, MTP_BYTES);
|
assert_eq!(plan.optional_mtp, MTP_BYTES);
|
||||||
assert_eq!(plan.kv_and_recurrent, 6_557_908_992);
|
assert_eq!(plan.kv_and_recurrent, 6_558_093_312);
|
||||||
assert_eq!(plan.prefill_transient, 27_262_976);
|
assert_eq!(plan.prefill_transient, 27_262_976);
|
||||||
assert_eq!(plan.admission, 80_000_430_099);
|
assert_eq!(plan.admission, 80_000_614_419);
|
||||||
let without_mtp = memory_plan(262_144, false, 512).unwrap();
|
let without_mtp = memory_plan(262_144, false, 512).unwrap();
|
||||||
assert_eq!(without_mtp.optional_mtp, MTP_BYTES);
|
assert_eq!(without_mtp.optional_mtp, MTP_BYTES);
|
||||||
assert_eq!(without_mtp.admission, 78_327_854_567);
|
assert_eq!(without_mtp.admission, 78_328_038_887);
|
||||||
assert!(memory_plan(0, false, 512).is_err());
|
assert!(memory_plan(0, false, 512).is_err());
|
||||||
assert!(memory_plan(262_145, false, 512).is_err());
|
assert!(memory_plan(262_145, false, 512).is_err());
|
||||||
assert!(memory_plan(1, false, 0).is_err());
|
assert!(memory_plan(1, false, 0).is_err());
|
||||||
|
|||||||
Reference in New Issue
Block a user