Metal: accelerate Qwen MXFP4 prefill and batched down rows
Decode each MXFP4 prefill half-block with one exponent read and four
byte-aligned packed loads. Keep the original LUT multiplication and final
half/float conversion, without changing accumulation or tile geometry.
Reuse activations across two MXFP4 down rows for M1 Max multi-token NR2
dispatches, preserving each row's block, FMA and reduction order. Use the
existing paired Q8 helper for shared experts and guard odd output tails.
Specialize this through function constant 907 and include it in the
pipeline cache key, compiling the extra register footprint out of T1.
Retain the existing NR1/NR4, other-device and safe-math fallback policies.
Add a standalone staging oracle against two unchanged stage8 decoders,
covering all exponent bytes, nibble values and byte alignments with both
storage types and math modes. Extend exact down tests with odd row counts,
block tails, T1/T2/T3, shared Q8/MXFP4 slots, NR1/2/4 and output guards.
Register the oracle build/test/clean targets and document measured gains.
Validation on Apple M1 Max:
- Build all binaries and the affected Metal test targets successfully.
- Staging oracle: 1,048,576 comparisons, zero NaN payload differences.
- Exact MV tests pass with DS4_METAL_MATH_SAFE=0 and 1.
- MoE MM specialization: 132 passing cases per math mode.
- SSD experts: 94 passing cases per math mode, including overlap,
selected-only reads, cache reuse, failure/retry and fallback paths.
- Original/current intermediate output fingerprints match in both modes.
- Eight full Q4 SSD runs, including MTP, produce identical output within
each workload; git diff --check passes.
Against 46425263, alternating resident prefill measurements reduce median
kernel time by 5.2-6.1%. Same-process two-token down measurements improve
throughput by 14.7-17.2%, with bit-identical results and no T1 GPU regression.
Full SSD prefill improves from 22.94 to 23.41 t/s in the measured workload;
median MTP generation stays at 7.65 t/s. These local results do not imply
the same gains for complete inference or other GPUs.
This commit is contained in:
Giorgio Oppo committed
1 parent
4642526317
commit
ee79f22b90
7 files changed
+392
-7
No files matched your search
@@ -99,5 +99,6 @@ __pycache__/
|
||||
/tests/test_metal_qwen4_moe_mid
|
||||
|
||||
/tests/test_metal_qwen4_moe_down
|
||||
/tests/test_metal_qwen4_mxfp4_stage
|
||||
/tests/test_metal_q8_gemv_reference
|
||||
/tests/test_metal_qwen4_moe_half
|
||||
@@ -253,6 +253,13 @@ test-metal-qwen4-moe-down: tests/test_metal_qwen4_moe_down
|
||||
check-metal-qwen4-moe-down: tests/test_metal_qwen4_moe_down
|
||||
./tests/test_metal_qwen4_moe_down --compile-only
|
||||
|
||||
tests/test_metal_qwen4_mxfp4_stage: tests/test_metal_qwen4_mxfp4_stage.m
|
||||
$(CC) $(filter-out -ffast-math,$(OBJCFLAGS)) -o $@ $< $(METAL_LDLIBS)
|
||||
|
||||
.PHONY: test-metal-qwen4-mxfp4-stage
|
||||
test-metal-qwen4-mxfp4-stage: tests/test_metal_qwen4_mxfp4_stage
|
||||
./tests/test_metal_qwen4_mxfp4_stage
|
||||
|
||||
tests/test_qwen4_ssd_experts.o: tests/test_qwen4_ssd_experts.c ds4_gpu.h
|
||||
$(CC) $(CFLAGS) -fno-fast-math -I. -c -o $@ $<
|
||||
|
||||
@@ -1249,6 +1256,7 @@ clean:
|
||||
rm -f tests/test_metal_qwen4_hc
|
||||
rm -f tests/test_metal_qwen4_moe_mid
|
||||
rm -f tests/test_metal_qwen4_moe_down
|
||||
rm -f tests/test_metal_qwen4_mxfp4_stage
|
||||
|
||||
# The active tokenizer includes generated Unicode classes.
|
||||
ds4.o ds4_cpu.o ds4_cpu_test_hooks.o: ds4_qwen4_unicode.inc
|
||||
@@ -99,6 +99,34 @@ and system variation remained substantial. Prefill medians were 6.94 and
|
||||
Q2's smaller sparse-prefill tiles did not show a stable benefit for Q4 in
|
||||
local comparisons. The measurements do not establish gains on other devices.
|
||||
|
||||
MXFP4 prefill also decodes each 16-value half-block with one exponent read
|
||||
and packed byte loads, retaining the original multiplication and final
|
||||
half/float conversion. Compared with `46425263`, alternating resident-kernel
|
||||
benchmarks on M1 Max at 2048 tokens and 10 routed slots reduced median
|
||||
gate/up-plus-down time from 70.65 to 66.31 ms with 32 experts and from 81.57
|
||||
to 77.34 ms with 256 experts (6.1% and 5.2% less time).
|
||||
The staging oracle checks all exponent bytes, nibble codes and byte
|
||||
alignments against two original 8-value decoders in both math modes.
|
||||
|
||||
On M1 Max, batched MXFP4 down rows also reuse activations across two output
|
||||
rows. A function constant enables this only for multi-token NR2 dispatches;
|
||||
single-token decode and NR1/NR4 retain the original row loop. A warmed
|
||||
same-process comparison with 16 alternating samples per variant, 128 calls
|
||||
per sample, width 640, 2560 output rows and 10 routed slots measured:
|
||||
|
||||
| Two-token down kernel | Before GPU time | After GPU time |
|
||||
| --- | --- | --- |
|
||||
| Q8 shared expert | 176.62 us | 150.70 us |
|
||||
| MXFP4 shared expert | 216.20 us | 188.48 us |
|
||||
|
||||
Outputs were bit-identical, and single-token GPU timings did not regress.
|
||||
Full Q4 SSD runs remain sensitive to I/O: two alternating pairs at context
|
||||
32768 and chunk 1024 gave prefill medians of 22.94 and 23.41 t/s on a
|
||||
201-token prompt, with all 100 generated tokens unchanged. With MTP and two
|
||||
draft tokens on the Rome prompt, median generation remained 7.65 t/s despite
|
||||
the faster down kernel. These checks compare against `46425263`; they do
|
||||
not imply the same percentage gain for complete inference or other devices.
|
||||
|
||||
Leave the expert-cache budget automatic initially. A plain count passed to
|
||||
`--ssd-streaming-cache-experts` requests dynamic cache slots; an `NGB` budget
|
||||
also includes the reserved layer-staging space. Startup accounts for the
|
||||
|
||||
+8
-2
@@ -48736,12 +48736,18 @@ static int qwen4_dispatch_resident(int kernel, const void *args, size_t args_len
|
||||
const qwen4_moe_args *a = args;
|
||||
const bool specialize = !q4_mid && !iq2_mid && !q2_down && qwen4_moe_mv_specialize(a->weight_type);
|
||||
const uint32_t values[] = {a->weight_type, a->shared_type, specialize ? a->in_dim : 0u, specialize ? qwen4_moe_mv_rows() : 0u};
|
||||
NSString *key = [NSString stringWithFormat:@"%s_type=%u_shared=%u_dim=%u_rows=%u_addr=%u",
|
||||
qwen4_kernel_names[kernel], values[0], values[1], values[2], values[3], addresses];
|
||||
/* M1 batched MXFP4 benefits from sharing loads across two rows.
|
||||
* Compile it out of single-token decode to retain its occupancy. */
|
||||
const bool pair_rows = kernel == QWEN4_K_MOE_DOWN_MXFP4_PF && a->n_tokens > 1u &&
|
||||
(!specialize || values[3] == 2u) && ds4_gpu_device_name_contains("M1 Max");
|
||||
NSString *key = [NSString stringWithFormat:@"%s_type=%u_shared=%u_dim=%u_rows=%u_addr=%u_pair=%u",
|
||||
qwen4_kernel_names[kernel], values[0], values[1], values[2], values[3], addresses, pair_rows];
|
||||
pipeline = [g_pipeline_cache objectForKey:key];
|
||||
if (!pipeline) {
|
||||
MTLFunctionConstantValues *constants = [[MTLFunctionConstantValues alloc] init];
|
||||
[constants setConstantValue:&addresses type:MTLDataTypeBool atIndex:906];
|
||||
if (kernel == QWEN4_K_MOE_DOWN_MXFP4_PF)
|
||||
[constants setConstantValue:&pair_rows type:MTLDataTypeBool atIndex:907];
|
||||
for (uint32_t i = 0; specialize && i < 4u; i++)
|
||||
[constants setConstantValue:&values[i] type:MTLDataTypeUInt atIndex:901u+i];
|
||||
NSError *error = nil;
|
||||
|
||||
+95
-5
@@ -3471,8 +3471,8 @@ kernel void kernel_qwen4_moe_down_q2k(
|
||||
* accumulation chain. The shipped qwen4_row_dot loop for type 39 compiles
|
||||
* to s = t0*y0 + t1*y1 + t2*y16 + t3*y17 (left to right), acc += s*d; that
|
||||
* order is spelled out with reassociation pinned off (tests/test_qwen4_kernels.c
|
||||
* pins the rows against kernel_qwen4_moe_down). The shared Q8 slot keeps the
|
||||
* generic row dot. */
|
||||
* pins the rows against kernel_qwen4_moe_down). Shared slots retain their
|
||||
* original quantization and per-row accumulation order. */
|
||||
#define QWEN4_MXFP4_PF_ACC_TO(ACC_, E_, Q0_, Q1_, Y0_, Y1_, Y16_, Y17_) do { \
|
||||
const float d_ = ds4_metal_e8m0_to_f32(E_); \
|
||||
const float t0_ = ds4_metal_mxfp4_values[(Q0_) & 0xfu], t1_ = ds4_metal_mxfp4_values[(Q1_) & 0xfu]; \
|
||||
@@ -3486,6 +3486,8 @@ kernel void kernel_qwen4_moe_down_q2k(
|
||||
#define QWEN4_MXFP4_PF_ACC(E_, Q0_, Q1_, Y0_, Y1_, Y16_, Y17_) \
|
||||
QWEN4_MXFP4_PF_ACC_TO(acc, E_, Q0_, Q1_, Y0_, Y1_, Y16_, Y17_)
|
||||
|
||||
constant bool qwen4_mxfp4_pair_rows [[function_constant(907)]];
|
||||
|
||||
kernel void kernel_qwen4_moe_down_mxfp4_pf(
|
||||
constant ds4_metal_args_qwen4_moe & args,
|
||||
device const char *down_base,
|
||||
@@ -3507,10 +3509,21 @@ kernel void kernel_qwen4_moe_down_mxfp4_pf(
|
||||
if (row0 >= args.out_rows || slot >= n_out || tok >= args.n_tokens) return;
|
||||
const uint64_t pair = (uint64_t)tok * n_out + slot;
|
||||
device const float *m = mid + pair * args.in_dim;
|
||||
const bool pair_rows = is_function_constant_defined(qwen4_mxfp4_pair_rows) && qwen4_mxfp4_pair_rows;
|
||||
if (slot == args.n_slots) {
|
||||
for (uint r = row0; r < row0 + nr && r < args.out_rows; r++) {
|
||||
const float v = qwen4_row_dot(sh_down + (uint64_t)r * args.shared_row_bytes, m, st, dim, tiisg);
|
||||
if (tiisg == 0) part[pair * args.out_rows + r] = v;
|
||||
if (pair_rows && nr == 2u && st == 8u && row0 + 1u < args.out_rows) {
|
||||
const uint64_t off = (uint64_t)row0 * args.shared_row_bytes;
|
||||
const float2 v = qwen4_q8_pair_dot(sh_down + off,
|
||||
sh_down + off + args.shared_row_bytes, m, dim, tiisg);
|
||||
if (tiisg == 0) {
|
||||
part[pair * args.out_rows + row0] = v.x;
|
||||
part[pair * args.out_rows + row0 + 1u] = v.y;
|
||||
}
|
||||
} else {
|
||||
for (uint r = row0; r < row0 + nr && r < args.out_rows; r++) {
|
||||
const float v = qwen4_row_dot(sh_down + (uint64_t)r * args.shared_row_bytes, m, st, dim, tiisg);
|
||||
if (tiisg == 0) part[pair * args.out_rows + r] = v;
|
||||
}
|
||||
}
|
||||
return;
|
||||
}
|
||||
@@ -3518,6 +3531,65 @@ kernel void kernel_qwen4_moe_down_mxfp4_pf(
|
||||
(uint)selected[(uint64_t)tok * args.n_slots + slot], args.expert_bytes);
|
||||
const uint ix = tiisg / 8, it = tiisg % 8;
|
||||
const uint nb = dim / 32;
|
||||
if (pair_rows && nr == 2u) {
|
||||
#pragma clang fp reassociate(off)
|
||||
#pragma clang fp contract(off)
|
||||
/* Batched rows benefit from interleaving independent chains. The
|
||||
* function constant removes this register footprint from single-token
|
||||
* decode. Activation loads are shared; block/FMA order is unchanged. */
|
||||
const bool two = row0 + 1u < args.out_rows;
|
||||
device const uchar *rowa = (device const uchar *)(db + (uint64_t)row0 * args.row_bytes);
|
||||
device const uchar *rowb = rowa + (two ? args.row_bytes : 0u);
|
||||
float acca = 0.0f, accb = 0.0f;
|
||||
uint ib = ix;
|
||||
for (; ib + 12u < nb; ib += 16u) {
|
||||
device const uchar *a0 = rowa + (uint64_t)ib * 17u;
|
||||
device const uchar *a1 = a0 + 4u * 17u, *a2 = a0 + 8u * 17u, *a3 = a0 + 12u * 17u;
|
||||
device const uchar *c0 = rowb + (uint64_t)ib * 17u;
|
||||
device const uchar *c1 = c0 + 4u * 17u, *c2 = c0 + 8u * 17u, *c3 = c0 + 12u * 17u;
|
||||
const uchar ea0 = a0[0], ea1 = a1[0], ea2 = a2[0], ea3 = a3[0];
|
||||
const uint pa0 = a0[1 + it * 2u], qa0 = a0[2 + it * 2u];
|
||||
const uint pa1 = a1[1 + it * 2u], qa1 = a1[2 + it * 2u];
|
||||
const uint pa2 = a2[1 + it * 2u], qa2 = a2[2 + it * 2u];
|
||||
const uint pa3 = a3[1 + it * 2u], qa3 = a3[2 + it * 2u];
|
||||
const uchar eb0 = c0[0], eb1 = c1[0], eb2 = c2[0], eb3 = c3[0];
|
||||
const uint pb0 = c0[1 + it * 2u], qb0 = c0[2 + it * 2u];
|
||||
const uint pb1 = c1[1 + it * 2u], qb1 = c1[2 + it * 2u];
|
||||
const uint pb2 = c2[1 + it * 2u], qb2 = c2[2 + it * 2u];
|
||||
const uint pb3 = c3[1 + it * 2u], qb3 = c3[2 + it * 2u];
|
||||
device const float *y0 = m + ib * 32u + it * 2u;
|
||||
device const float *y1 = y0 + 128u, *y2 = y0 + 256u, *y3 = y0 + 384u;
|
||||
const float y0a = y0[0], y0b = y0[1], y0c = y0[16], y0d = y0[17];
|
||||
const float y1a = y1[0], y1b = y1[1], y1c = y1[16], y1d = y1[17];
|
||||
const float y2a = y2[0], y2b = y2[1], y2c = y2[16], y2d = y2[17];
|
||||
const float y3a = y3[0], y3b = y3[1], y3c = y3[16], y3d = y3[17];
|
||||
QWEN4_MXFP4_PF_ACC_TO(acca, ea0, pa0, qa0, y0a, y0b, y0c, y0d);
|
||||
QWEN4_MXFP4_PF_ACC_TO(acca, ea1, pa1, qa1, y1a, y1b, y1c, y1d);
|
||||
QWEN4_MXFP4_PF_ACC_TO(acca, ea2, pa2, qa2, y2a, y2b, y2c, y2d);
|
||||
QWEN4_MXFP4_PF_ACC_TO(acca, ea3, pa3, qa3, y3a, y3b, y3c, y3d);
|
||||
QWEN4_MXFP4_PF_ACC_TO(accb, eb0, pb0, qb0, y0a, y0b, y0c, y0d);
|
||||
QWEN4_MXFP4_PF_ACC_TO(accb, eb1, pb1, qb1, y1a, y1b, y1c, y1d);
|
||||
QWEN4_MXFP4_PF_ACC_TO(accb, eb2, pb2, qb2, y2a, y2b, y2c, y2d);
|
||||
QWEN4_MXFP4_PF_ACC_TO(accb, eb3, pb3, qb3, y3a, y3b, y3c, y3d);
|
||||
}
|
||||
for (; ib < nb; ib += 4u) {
|
||||
device const uchar *a0 = rowa + (uint64_t)ib * 17u;
|
||||
device const uchar *c0 = rowb + (uint64_t)ib * 17u;
|
||||
const uchar ea0 = a0[0], eb0 = c0[0];
|
||||
const uint pa0 = a0[1 + it * 2u], qa0 = a0[2 + it * 2u];
|
||||
const uint pb0 = c0[1 + it * 2u], qb0 = c0[2 + it * 2u];
|
||||
device const float *y0 = m + ib * 32u + it * 2u;
|
||||
const float y0a = y0[0], y0b = y0[1], y0c = y0[16], y0d = y0[17];
|
||||
QWEN4_MXFP4_PF_ACC_TO(acca, ea0, pa0, qa0, y0a, y0b, y0c, y0d);
|
||||
QWEN4_MXFP4_PF_ACC_TO(accb, eb0, pb0, qb0, y0a, y0b, y0c, y0d);
|
||||
}
|
||||
const float va = simd_sum(acca), vb = simd_sum(accb);
|
||||
if (tiisg == 0) {
|
||||
part[pair * args.out_rows + row0] = va;
|
||||
if (two) part[pair * args.out_rows + row0 + 1u] = vb;
|
||||
}
|
||||
return;
|
||||
}
|
||||
for (uint r = row0; r < row0 + nr && r < args.out_rows; r++) {
|
||||
#pragma clang fp reassociate(off)
|
||||
#pragma clang fp contract(off)
|
||||
@@ -3965,6 +4037,24 @@ static inline void qwen4_mm_stage16(device const char *row, uint b, uint q0, uin
|
||||
}
|
||||
return;
|
||||
}
|
||||
if (type == 39) {
|
||||
device const uchar *blk = (device const uchar *)(row + (uint64_t)b * 17);
|
||||
const float d = ds4_metal_e8m0_to_f32(blk[0]);
|
||||
const bool hi = q0 >= 2;
|
||||
/* Seventeen-byte blocks have no wider alignment guarantee. Read the
|
||||
* sixteen packed code bytes once, retaining stage8's multiplication
|
||||
* and final storage conversion for each low/high nibble. */
|
||||
#pragma unroll
|
||||
for (uint j = 0; j < 4; j++) {
|
||||
const packed_uchar4 codes = *(device const packed_uchar4 *)(blk + 1 + j * 4);
|
||||
#pragma unroll
|
||||
for (uint i = 0; i < 4; i++) {
|
||||
const uint byte = codes[i];
|
||||
dst[j * 4 + i] = (D)(d * ds4_metal_mxfp4_values[hi ? (byte >> 4) : (byte & 0xfu)]);
|
||||
}
|
||||
}
|
||||
return;
|
||||
}
|
||||
qwen4_mm_stage8(row, b, q0, type, dst);
|
||||
qwen4_mm_stage8(row, b, q0 + 1, type, dst + 8);
|
||||
}
|
||||
|
||||
@@ -0,0 +1,171 @@
|
||||
/* MXFP4 prefill staging oracle: production stage16 versus two unchanged
|
||||
* stage8 calls. Covers every exponent byte and nibble code at every block
|
||||
* alignment, for half/float storage and generic/specialized dispatch.
|
||||
* Build: clang -O2 -Wall -Wextra -fobjc-arc -framework Foundation -framework Metal \
|
||||
* tests/test_metal_qwen4_mxfp4_stage.m -o /tmp/test_metal_qwen4_mxfp4_stage
|
||||
* Run from the repository root. --source FILE selects a candidate Metal
|
||||
* source; --compile-only creates pipelines without GPU dispatches.
|
||||
*/
|
||||
#import <Foundation/Foundation.h>
|
||||
#import <Metal/Metal.h>
|
||||
#include <stdbool.h>
|
||||
#include <stdint.h>
|
||||
#include <stdio.h>
|
||||
#include <stdlib.h>
|
||||
#include <string.h>
|
||||
|
||||
static const NSUInteger kGuard = 256;
|
||||
static const unsigned kBlocks = 256 * 16;
|
||||
static uint64_t compared, nan_payload_differences;
|
||||
|
||||
static void fail(NSString *s) {
|
||||
fprintf(stderr, "FAIL MXFP4 staging: %s\n", s.UTF8String); exit(1);
|
||||
}
|
||||
static NSString *read_source(NSString *path) {
|
||||
NSError *error = nil;
|
||||
NSString *s = [NSString stringWithContentsOfFile:path encoding:NSUTF8StringEncoding error:&error];
|
||||
if (!s) fail([NSString stringWithFormat:@"read %@: %@", path, error]);
|
||||
return s;
|
||||
}
|
||||
static NSString *between(NSString *s, NSString *begin, NSString *end) {
|
||||
if ([s componentsSeparatedByString:begin].count != 2 ||
|
||||
[s componentsSeparatedByString:end].count != 2) fail(@"source anchors missing/nonunique");
|
||||
NSRange a = [s rangeOfString:begin], b = [s rangeOfString:end];
|
||||
if (b.location <= a.location) fail(@"invalid source range");
|
||||
return [s substringWithRange:NSMakeRange(a.location, b.location - a.location)];
|
||||
}
|
||||
static NSString *kernel_name(bool half, bool specialized, bool candidate) {
|
||||
return [NSString stringWithFormat:@"stage_%@_%u_%u", half ? @"half" : @"float", specialized, candidate];
|
||||
}
|
||||
static NSString *production_source(NSString *repo, NSString *path) {
|
||||
NSString *q = read_source(path), *m = read_source([repo stringByAppendingPathComponent:@"metal/moe.metal"]);
|
||||
NSMutableString *s = [NSMutableString stringWithString:@"#include <metal_stdlib>\nusing namespace metal;\n"];
|
||||
[s appendString:between(m, @"static constant float ds4_metal_mxfp4_values[16]", @"// BEGIN GENERATED MXFP4 HALF LUT")];
|
||||
[s appendString:between(m, @"static constant uchar ds4_metal_ksigns_iq2xs[128]", @"#define kmask_iq2xs")];
|
||||
[s appendString:between(q, @"/* dequantize 8 consecutive values", @"/* Raw words of one 32-block")];
|
||||
for (unsigned h = 0; h < 2; h++) for (unsigned sp = 0; sp < 2; sp++) for (unsigned arm = 0; arm < 2; arm++) {
|
||||
NSString *type = h ? @"half" : @"float";
|
||||
[s appendFormat:@"\nkernel void %@(device const char *input [[buffer(0)]], "
|
||||
"device %@ *output [[buffer(1)]], constant uint &runtime_type [[buffer(2)]], "
|
||||
"uint gid [[thread_position_in_grid]], ushort tid [[thread_index_in_threadgroup]]) {\n"
|
||||
" threadgroup %@ scratch[128 * 16];\n"
|
||||
" threadgroup %@ *dst = scratch + uint(tid) * 16;\n"
|
||||
" const uint b = gid / 2, q0 = (gid & 1u) * 2;\n"
|
||||
" const uint type = %@;\n", kernel_name(h, sp, arm), type, type, type, sp ? @"39u" : @"runtime_type"];
|
||||
if (arm) [s appendString:@" qwen4_mm_stage16(input + 1, b, q0, type, dst);\n"];
|
||||
else [s appendString:@" qwen4_mm_stage8(input + 1, b, q0, type, dst);\n"
|
||||
" qwen4_mm_stage8(input + 1, b, q0 + 1, type, dst + 8);\n"];
|
||||
[s appendString:@" threadgroup_barrier(mem_flags::mem_threadgroup);\n"
|
||||
" for (uint i = 0; i < 16; i++) output[gid * 16 + i] = dst[i];\n}\n"];
|
||||
}
|
||||
return s;
|
||||
}
|
||||
static id<MTLBuffer> buffer(id<MTLDevice> device, size_t bytes) {
|
||||
id<MTLBuffer> b = [device newBufferWithLength:bytes + 2 * kGuard options:MTLResourceStorageModeShared];
|
||||
if (!b) fail(@"buffer allocation");
|
||||
memset(b.contents, 0xa5, b.length); return b;
|
||||
}
|
||||
static void *payload(id<MTLBuffer> b) { return (char *)b.contents + kGuard; }
|
||||
static void check_guard(id<MTLBuffer> b) {
|
||||
const uint8_t *p = b.contents;
|
||||
for (NSUInteger i = 0; i < kGuard; i++)
|
||||
if (p[i] != 0xa5 || p[b.length - 1 - i] != 0xa5) fail(@"output guard overwritten");
|
||||
}
|
||||
static NSArray<id<MTLComputePipelineState>> *pipelines(id<MTLDevice> device, NSString *source, bool safe) {
|
||||
MTLCompileOptions *options = [MTLCompileOptions new];
|
||||
if (safe) {
|
||||
if (@available(macOS 15.0, *)) options.mathMode = MTLMathModeSafe;
|
||||
else {
|
||||
#pragma clang diagnostic push
|
||||
#pragma clang diagnostic ignored "-Wdeprecated-declarations"
|
||||
options.fastMathEnabled = NO;
|
||||
#pragma clang diagnostic pop
|
||||
}
|
||||
}
|
||||
NSError *error = nil;
|
||||
id<MTLLibrary> lib = [device newLibraryWithSource:source options:options error:&error];
|
||||
if (!lib) fail(error.description);
|
||||
NSMutableArray *ps = [NSMutableArray new];
|
||||
for (unsigned h = 0; h < 2; h++) for (unsigned sp = 0; sp < 2; sp++) for (unsigned arm = 0; arm < 2; arm++) {
|
||||
id<MTLFunction> f = [lib newFunctionWithName:kernel_name(h, sp, arm)];
|
||||
if (!f) fail(@"kernel unavailable");
|
||||
id<MTLComputePipelineState> p = [device newComputePipelineStateWithFunction:f error:&error];
|
||||
if (!p) fail(error.description);
|
||||
if (p.maxTotalThreadsPerThreadgroup < 128) fail(@"128-thread staging unavailable");
|
||||
[ps addObject:p];
|
||||
}
|
||||
return ps;
|
||||
}
|
||||
static void dispatch_pair(id<MTLCommandQueue> queue, NSArray *ps, unsigned first,
|
||||
id<MTLBuffer> input, id<MTLBuffer> ref, id<MTLBuffer> got) {
|
||||
id<MTLCommandBuffer> cb = [queue commandBuffer];
|
||||
id<MTLComputeCommandEncoder> en = [cb computeCommandEncoder];
|
||||
if (!cb || !en) fail(@"command allocation");
|
||||
const uint32_t type = 39;
|
||||
for (unsigned arm = 0; arm < 2; arm++) {
|
||||
[en setComputePipelineState:ps[first + arm]];
|
||||
[en setBuffer:input offset:kGuard atIndex:0];
|
||||
[en setBuffer:arm ? got : ref offset:kGuard atIndex:1];
|
||||
[en setBytes:&type length:sizeof(type) atIndex:2];
|
||||
[en dispatchThreadgroups:MTLSizeMake(kBlocks * 2 / 128, 1, 1)
|
||||
threadsPerThreadgroup:MTLSizeMake(128, 1, 1)];
|
||||
}
|
||||
[en endEncoding]; [cb commit]; [cb waitUntilCompleted];
|
||||
if (cb.status != MTLCommandBufferStatusCompleted) fail(cb.error.description);
|
||||
}
|
||||
static void compare(id<MTLBuffer> ref, id<MTLBuffer> got, bool half) {
|
||||
for (size_t i = 0; i < (size_t)kBlocks * 32; i++) {
|
||||
const uint32_t r = half ? ((uint16_t *)payload(ref))[i] : ((uint32_t *)payload(ref))[i];
|
||||
const uint32_t g = half ? ((uint16_t *)payload(got))[i] : ((uint32_t *)payload(got))[i];
|
||||
const uint32_t signless = half ? 0x7fffu : 0x7fffffffu, inf = half ? 0x7c00u : 0x7f800000u;
|
||||
if (r == (half ? 0xa5a5u : 0xa5a5a5a5u) || g == (half ? 0xa5a5u : 0xa5a5a5a5u)) fail(@"output untouched");
|
||||
if ((r & signless) > inf && (g & signless) > inf) nan_payload_differences += r != g;
|
||||
else if (r != g) fail([NSString stringWithFormat:@"%@ block=%zu value=%zu: %08x != %08x",
|
||||
half ? @"half" : @"float", i / 32, i % 32, r, g]);
|
||||
compared++;
|
||||
}
|
||||
check_guard(ref); check_guard(got);
|
||||
}
|
||||
int main(int argc, const char **argv) { @autoreleasepool {
|
||||
NSString *repo = [[NSFileManager defaultManager] currentDirectoryPath], *path = nil;
|
||||
bool compileOnly = false;
|
||||
for (int i = 1; i < argc; i++) {
|
||||
if (!strcmp(argv[i], "--source") && i + 1 < argc) path = @(argv[++i]);
|
||||
else if (!strcmp(argv[i], "--repo") && i + 1 < argc) repo = @(argv[++i]);
|
||||
else if (!strcmp(argv[i], "--compile-only")) compileOnly = true;
|
||||
else fail(@"usage: [--repo PATH] [--source FILE] [--compile-only]");
|
||||
}
|
||||
path = path ?: [repo stringByAppendingPathComponent:@"metal/qwen4.metal"];
|
||||
NSString *source = production_source(repo, path);
|
||||
id<MTLDevice> device = MTLCreateSystemDefaultDevice();
|
||||
if (!device) fail(@"Metal device unavailable");
|
||||
for (unsigned safe = 0; safe < 2; safe++) @autoreleasepool {
|
||||
NSArray *ps = pipelines(device, source, safe);
|
||||
if (compileOnly) continue;
|
||||
id<MTLCommandQueue> queue = [device newCommandQueue];
|
||||
if (!queue) fail(@"command queue unavailable");
|
||||
id<MTLBuffer> input = buffer(device, (size_t)kBlocks * 17 + 1);
|
||||
uint8_t *p = (uint8_t *)payload(input) + 1;
|
||||
/* The odd base and 17-byte stride exercise all sixteen alignments.
|
||||
* Each exponent sees every low/high code at every payload position. */
|
||||
for (unsigned b = 0; b < kBlocks; b++) {
|
||||
const unsigned pattern = b % 16;
|
||||
p[b * 17] = b / 16;
|
||||
for (unsigned j = 0; j < 16; j++)
|
||||
p[b * 17 + 1 + j] = ((pattern + j) & 15u) | (((pattern * 7u + j * 3u) & 15u) << 4);
|
||||
}
|
||||
NSData *before = [NSData dataWithBytes:input.contents length:input.length];
|
||||
for (unsigned h = 0; h < 2; h++) for (unsigned sp = 0; sp < 2; sp++) {
|
||||
const size_t bytes = (size_t)kBlocks * 32 * (h ? 2 : 4);
|
||||
id<MTLBuffer> ref = buffer(device, bytes), got = buffer(device, bytes);
|
||||
dispatch_pair(queue, ps, h * 4 + sp * 2, input, ref, got); compare(ref, got, h);
|
||||
if (memcmp(input.contents, before.bytes, input.length)) fail(@"input or its guards modified");
|
||||
printf("PASS MXFP4 stage8/stage16 safe=%u storage=%s specialized=%u: %u values\n",
|
||||
safe, h ? "half" : "float", sp, kBlocks * 32);
|
||||
}
|
||||
}
|
||||
printf("PASS MXFP4 staging: %llu comparisons, %llu NaN payload differences%s\n",
|
||||
(unsigned long long)compared, (unsigned long long)nan_payload_differences,
|
||||
compileOnly ? " (compile only)" : "");
|
||||
return 0;
|
||||
} }
|
||||
@@ -1967,6 +1967,86 @@ static void test_moe(arena_t *a, uint32_t NE, uint32_t slots, uint32_t E, uint32
|
||||
test_moe_types(a, NE, slots, E, F, T, wtype, wtype ? 8u : 0u);
|
||||
}
|
||||
|
||||
/* A down-only fixture allows odd output rows, unlike the complete Q4 MoE
|
||||
* fixture whose gate input must be a multiple of 256. Exercise both row
|
||||
* reuse and its single-row tail with guarded, unaligned output views. */
|
||||
static void test_mxfp4_down_exact(arena_t *a) {
|
||||
const char *env[] = {"DS4_QWEN4_MOE_MV_SPECIALIZE", "DS4_QWEN4_MOE_MV_NR",
|
||||
"DS4_QWEN4_MOE_MV_NSG", "DS4_QWEN4_MOE_DOWN_PREFETCH"};
|
||||
char *saved[4];
|
||||
for (uint32_t i = 0; i < 4u; i++) {
|
||||
const char *value = getenv(env[i]);
|
||||
saved[i] = value ? strdup(value) : NULL;
|
||||
require_ok(!value || saved[i], "save MXFP4 down environment");
|
||||
require_ok(unsetenv(env[i]) == 0, "clear MXFP4 down environment");
|
||||
}
|
||||
const uint32_t widths[] = {640u, 640u, 672u, 32u};
|
||||
const uint32_t rows[] = {2560u, 2561u, 257u, 1u};
|
||||
const uint32_t shared_types[] = {UINT32_MAX, 8u, 39u};
|
||||
const uint32_t NE = 4u, slots = 3u, guard = 17u;
|
||||
const float sentinel = 127.25f;
|
||||
for (uint32_t shape = 0; shape < 4u; shape++) {
|
||||
const uint32_t F = widths[shape], E = rows[shape];
|
||||
double *shadow;
|
||||
const uint64_t down = arena_mxfp4(a, (uint64_t)NE * E, F, &shadow); free(shadow);
|
||||
const uint64_t sh_q8 = arena_q8_0(a, E, F, &shadow, 0.05f); free(shadow);
|
||||
const uint64_t sh_mx = arena_mxfp4(a, E, F, &shadow); free(shadow);
|
||||
int32_t selected[9];
|
||||
for (uint32_t t = 0; t < 3u; t++) for (uint32_t s = 0; s < slots; s++)
|
||||
selected[t * slots + s] = (int32_t)((t + s * 3u) % NE);
|
||||
ds4_gpu_tensor *gsel = ds4_gpu_tensor_alloc(sizeof(selected));
|
||||
require_ok(gsel && ds4_gpu_tensor_write(gsel, 0, selected, sizeof(selected)), "MXFP4 selected IDs");
|
||||
for (uint32_t shared = 0; shared < 3u; shared++) for (uint32_t T = 1; T <= 3u; T++) {
|
||||
const uint32_t st = shared_types[shared], n_out = slots + (shared != 0u);
|
||||
const uint64_t nx = (uint64_t)T * n_out * F, n = (uint64_t)T * n_out * E;
|
||||
const uint64_t total = n + 2u * guard, sh = st == 8u ? sh_q8 : sh_mx;
|
||||
float *x = rand_vec(nx, 1.0f);
|
||||
for (uint64_t i = 0; i < nx; i++) x[i] *= i % 7u == 0u ? 8.0f : i % 7u == 1u ? 0.125f : 1.0f;
|
||||
ds4_gpu_tensor *gx = upload(x, nx); free(x);
|
||||
ds4_gpu_tensor *storage = upload(NULL, total);
|
||||
ds4_gpu_tensor *out = ds4_gpu_tensor_view(storage, guard * sizeof(float), n * sizeof(float));
|
||||
require_ok(out != NULL, "MXFP4 guarded output view");
|
||||
float *reference = NULL;
|
||||
char label[128];
|
||||
snprintf(label, sizeof(label), "MXFP4 down E=%u F=%u T=%u shared=%u", E, F, T, st);
|
||||
for (uint32_t mode = 0; mode < 5u; mode++) {
|
||||
/* Compare NR2 reuse and specialized NR1/2/4 fallbacks with
|
||||
* the original plain, unspecialized down kernel. */
|
||||
const char *nr[] = {"2", "2", "2", "1", "4"};
|
||||
require_ok(setenv(env[0], mode >= 2u ? "1" : "0", 1) == 0 &&
|
||||
setenv(env[1], nr[mode], 1) == 0 && setenv(env[2], "4", 1) == 0 &&
|
||||
setenv(env[3], mode ? "1" : "0", 1) == 0, "select MXFP4 down oracle or prefetch");
|
||||
require_ok(ds4_gpu_tensor_fill_f32(storage, sentinel, total) &&
|
||||
ds4_gpu_begin_commands() &&
|
||||
ds4_gpu_qwen4_moe_down_tensor(out, gx, gsel, a->base, a->size,
|
||||
down, 39u, NE, T, slots, F, E, sh, st) && ds4_gpu_end_commands(),
|
||||
"MXFP4 down guarded dispatch");
|
||||
float *got = download(storage, total);
|
||||
for (uint32_t i = 0; i < guard; i++)
|
||||
require_ok(got[i] == sentinel && got[guard + n + i] == sentinel, "MXFP4 down guards");
|
||||
if (!mode) {
|
||||
reference = got;
|
||||
check_exact_f32(label, reference, reference, total);
|
||||
for (uint64_t i = guard; i < guard + n; i++)
|
||||
require_ok(got[i] != sentinel, "MXFP4 oracle wrote every output");
|
||||
} else {
|
||||
check_exact_f32(label, got, reference, total);
|
||||
free(got);
|
||||
}
|
||||
}
|
||||
free(reference);
|
||||
ds4_gpu_tensor_free(out); ds4_gpu_tensor_free(storage); ds4_gpu_tensor_free(gx);
|
||||
}
|
||||
ds4_gpu_tensor_free(gsel);
|
||||
}
|
||||
for (uint32_t i = 0; i < 4u; i++) {
|
||||
require_ok(saved[i] ? setenv(env[i], saved[i], 1) == 0 : unsetenv(env[i]) == 0,
|
||||
"restore MXFP4 down environment");
|
||||
free(saved[i]);
|
||||
}
|
||||
printf("MXFP4 down: exact NR1/2/4, odd rows, block tails and shared slots; guards intact\n");
|
||||
}
|
||||
|
||||
static void check_exact_f32(const char *what, const float *got, const float *ref, uint64_t n) {
|
||||
for (uint64_t i = 0; i < n; i++) {
|
||||
uint32_t gb, rb;
|
||||
@@ -4234,6 +4314,7 @@ int main(void) {
|
||||
#endif
|
||||
if (getenv("DS4_TEST_QWEN4_DECODE_FUSIONS")) { test_decode_fusions(&arena); return 0; }
|
||||
if (getenv("DS4_TEST_QWEN4_MV_EXACT")) {
|
||||
test_mxfp4_down_exact(&arena);
|
||||
test_moe_types(&arena, 8, 6, 2560, 640, 1, 16u, 10u);
|
||||
test_moe_types(&arena, 8, 6, 2560, 640, 2, 16u, 10u);
|
||||
test_moe_types(&arena, 8, 6, 256, 256, 9, 16u, 10u);
|
||||
|
||||
Reference in new issue
Block a user