Read half-aligned Qwen Q2 weight rows correctly on Metal

The ordinary Q2 row dot treated packed codes as word-aligned. Metal read the wrong codes when the expert row started two bytes into a word. Read the same two words via packed bytes and preserve the accumulation order.

Compile the two production down kernels independently in the comparison test so the intervening fused reduction does not pull in unrelated dependencies. The remote M5 Max passes 19,955,180 comparisons across 1,650 fixtures, including half-aligned and address-table weights.
This commit is contained in:
antirez committed 2026-11-03 16:55:02 +01:00
1 parent 5e3a3449e2
commit b3f50180ec
2 files changed
+18 -4

No files matched your search

+6 -3
View File
@@ -2836,9 +2836,12 @@ static inline float qwen4_row_dot(device const char *row, device const float *x,
const float dmin = (float)(*(device const half *)(blk + 82));
const uint sc = blk[group];
const float ds = d * (float)(sc & 0xFu), dm = dmin * (float)(sc >> 4);
/* the 8 quant bytes are 4-byte aligned (84-byte blocks) */
device const uint *qw = (device const uint *)(blk + 16 + q_base + l);
const uint q0 = qw[0] >> shift, q1 = qw[1] >> shift;
/* Q2_K requires half-aligned scales. A sliced row can start
* two bytes into a word, so read its codes with packed alignment. */
device const packed_uchar4 *qw =
(device const packed_uchar4 *)(blk + 16 + q_base + l);
const uint q0 = as_type<uint>(uchar4(qw[0])) >> shift;
const uint q1 = as_type<uint>(uchar4(qw[1])) >> shift;
device const float *y = x + ib * 256 + group * 16 + l;
for (uint i = 0; i < 4; i++) acc += (ds * (float)((q0 >> (8u * i)) & 3u) - dm) * y[i];
for (uint i = 0; i < 4; i++) acc += (ds * (float)((q1 >> (8u * i)) & 3u) - dm) * y[4 + i];
+12 -1
View File
@@ -51,6 +51,16 @@ static NSString *between(NSString *s, NSString *begin, NSString *end) {
if (b.location<=NSMaxRange(a)) fail(@"invalid source range");
return [s substringWithRange:NSMakeRange(a.location,b.location-a.location)];
}
/* Production kernel definitions close at the start of a line. Copy each
* tested entry point independently of any intervening kernels. */
static NSString *kernel_source(NSString *source, NSString *signature) {
NSRange start = [source rangeOfString:signature];
if (start.location == NSNotFound) fail(@"production kernel signature missing");
NSRange rest = NSMakeRange(NSMaxRange(start), source.length - NSMaxRange(start));
NSRange end = [source rangeOfString:@"\n}\n" options:0 range:rest];
if (end.location == NSNotFound) fail(@"production kernel definition incomplete");
return [source substringWithRange:NSMakeRange(start.location, NSMaxRange(end) - start.location)];
}
static NSString *production_source(NSString *repo, NSString *override) {
NSString *q=read_source(override ?: [repo stringByAppendingPathComponent:@"metal/qwen4.metal"]);
NSString *m=read_source([repo stringByAppendingPathComponent:@"metal/moe.metal"]);
@@ -59,7 +69,8 @@ static NSString *production_source(NSString *repo, NSString *override) {
[s appendString:between(m,@"static constant uchar ds4_metal_ksigns_iq2xs[128]",@"#define kmask_iq2xs")];
[s appendString:between(q,@"constant bool qwen4_expert_addresses ",@"/* --- small multi-output GEMV")];
[s appendString:between(q,@"constant uint qwen4_mv_type ",@"/* mid[t][s][r] =")];
[s appendString:between(q,@"/* part[t][s][r] = down_row",@"/* MXFP4 routed down rows")];
[s appendString:kernel_source(q, @"kernel void kernel_qwen4_moe_down(")];
[s appendString:kernel_source(q, @"kernel void kernel_qwen4_moe_down_q2k(")];
return s;
}
static uint64_t source_hash(NSString *s) {