diff --git a/metal/qwen4.metal b/metal/qwen4.metal index 5da908a1..8cfcf8a2 100644 --- a/metal/qwen4.metal +++ b/metal/qwen4.metal @@ -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(uchar4(qw[0])) >> shift; + const uint q1 = as_type(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]; diff --git a/tests/test_metal_qwen4_moe_down.m b/tests/test_metal_qwen4_moe_down.m index a34678c8..3e844702 100644 --- a/tests/test_metal_qwen4_moe_down.m +++ b/tests/test_metal_qwen4_moe_down.m @@ -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) {