From 5dec79703e8a99b485ea27cdc6c2b67054ceee51 Mon Sep 17 00:00:00 2001 From: Glenn Neuber Date: Fri, 18 Sep 2026 20:43:19 +1000 Subject: [PATCH] =?UTF-8?q?test(mlx):=20guard=20MLX=20#3912=20=E2=80=94=20?= =?UTF-8?q?fp=5Fqmm=5Ft=20misreads=20past=20K=20when=20K=20mod=2032=20=3D?= =?UTF-8?q?=3D=2016?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Upstream MLX #3912 fixed a Metal kernel bug in fp_qmm_t that reads past the K dimension when K mod 32 == 16. Our v0.34.1 fold crossed that fix when the MLX pin moved ce916dbb -> d9add9d1, which is why gemma4 26b/31b vision embeddings changed across #302 while 12b did not. K=4304 is the vision tower's mlp.down_proj contraction dimension in the gemma4 26b/31b nvfp4 checkpoints: stored packed as [1152, 538] at 8 values per word, cross-checked by the group-16 scale tensor at [1152, 269]. 12b carries no nvfp4 vision down_proj at all — its text ones are K = 15360 and 65536, both mod 32 == 0 — so it was structurally immune. K=4288 is the nearest aligned K and is the control. Verified in both directions rather than only passing: MLX ce916dbb (pre-fix) K=4288 max 0.000152588 0/294912 over 1.0 K=4304 max 24.0494 232722/294912 -> FAIL MLX d9add9d1 (post-fix) K=4288 max 0.000152588 0/294912 K=4304 max 0.137939 0/294912 -> PASS Two things the comment block records because they are what made this expensive to find. First, the production symptom was not a crash or an obviously wrong answer: one fine-text OCR tier moved from 4 to 3 on one model in one think mode, and localising it from there took a day. Second, M must exceed get_qmv_batch_limit (13 for this K and N) or the dispatch goes to the matrix-vector kernel and fp_qmm_t is never reached — a first attempt at M=8 came back clean on the buggy library and read as a refutation of a correct diagnosis. Shrinking M silently voids this test. Co-Authored-By: Claude Opus 5 --- x/mlxrunner/mlx/fp_qmm_t_kmod32_test.go | 102 ++++++++++++++++++++++++ 1 file changed, 102 insertions(+) create mode 100644 x/mlxrunner/mlx/fp_qmm_t_kmod32_test.go diff --git a/x/mlxrunner/mlx/fp_qmm_t_kmod32_test.go b/x/mlxrunner/mlx/fp_qmm_t_kmod32_test.go new file mode 100644 index 00000000000..a2cc9275537 --- /dev/null +++ b/x/mlxrunner/mlx/fp_qmm_t_kmod32_test.go @@ -0,0 +1,102 @@ +package mlx + +import ( + "math" + "math/rand" + "testing" + + "github.com/ollama/ollama/x/internal/mlxthreadtest" +) + +// TestFpQmmTKMod32 guards upstream MLX #3912: fp_qmm_t misread past the K +// dimension when K mod 32 == 16, silently corrupting most of the output. +// +// WHY THIS K. 4304 is the vision tower's mlp.down_proj contraction dimension in +// the gemma4 26b/31b nvfp4 checkpoints — stored packed as [1152, 538] at 8 +// values per word, cross-checked by the group-16 scale tensor at [1152, 269]. +// 12b carries no nvfp4 vision down_proj at all (its text ones are K = 15360 and +// 65536, both mod 32 == 0), which is why it was unaffected. 4288 is the nearest +// mod-32-aligned K and is the control: it must stay clean on every build, so a +// failure at 4288 means something other than this bug. +// +// WHAT IT LOOKED LIKE IN PRODUCTION. Not a crash and not a wrong answer that +// anyone would notice — a single fine-text OCR tier moving from 4 to 3 on one +// model in one think mode. It cost a day to localise from that symptom. Measured +// here at MLX ce916dbb: 232722 of 294912 output elements wrong by more than 1.0, +// max error 24.05. At d9add9d1 (post-fix): zero, max error 0.138. +// +// M MUST EXCEED get_qmv_batch_limit. quantized.cpp dispatches to the matrix- +// VECTOR kernel when M is at or below that limit, which for K=4304, N=1152 is +// 13. A first attempt at M=8 never reached fp_qmm_t, came back clean on the +// buggy library, and read as a refutation of a correct diagnosis. If this test +// is ever "simplified" by shrinking M, it stops testing anything. +func TestFpQmmTKMod32(t *testing.T) { + withMLXThread(t, func(t *mlxthreadtest.T) { + const ( + N = 1152 // gemma4 vision tower down_proj output dim + M = 256 // must exceed get_qmv_batch_limit (13 here); see above + group = 16 + bits = 4 + // Post-fix the aligned K lands at 1.5e-4 and the trigger K at 0.14, + // both pure quantisation error. The bug produced 24.05. Anything + // above 1.0 is the kernel, not the quantiser. + tolerance = 1.0 + ) + + for _, tc := range []struct { + name string + K int + }{ + {"aligned_control", 4288}, // K mod 32 == 0 + {"gemma4_vision_down_proj", 4304}, // K mod 32 == 16, the trigger + } { + func() { + rng := rand.New(rand.NewSource(42)) + wv := make([]float32, N*tc.K) + for i := range wv { + wv[i] = float32(rng.NormFloat64()) + } + xv := make([]float32, M*tc.K) + for i := range xv { + xv[i] = float32(rng.NormFloat64()) + } + + w := FromValues(wv, N, tc.K) + x := FromValues(xv, M, tc.K) + wq, scales, biases := Quantize(w, group, bits, "nvfp4") + + // Reference deliberately avoids fp_qmm_t: dequantize, then a + // dense matmul. Any disagreement beyond quantisation error is + // the quantized kernel. + deq := Dequantize(wq, scales, biases, group, bits, "nvfp4", nil) + want := Matmul(x, Transpose(deq, 1, 0)) + got := QuantizedMatmul(x, wq, scales, biases, true, group, bits, "nvfp4", nil) + Eval(want, got) + + a, b := want.Floats(), got.Floats() + if len(a) != len(b) { + t.Fatalf("shape mismatch: reference %d elements, kernel %d", len(a), len(b)) + } + var maxAbs float64 + var bad int + for i := range a { + d := math.Abs(float64(a[i]) - float64(b[i])) + if d > maxAbs { + maxAbs = d + } + if d > tolerance { + bad++ + } + } + t.Logf("%s: K=%d (K%%32=%d) max|qmm-reference|=%.6g, %d/%d elements over %.1f", + tc.name, tc.K, tc.K%32, maxAbs, bad, len(a), tolerance) + if bad > 0 { + t.Errorf("fp_qmm_t disagrees with dequantize+matmul at K=%d (K%%32=%d): "+ + "%d of %d elements off by more than %.1f, max %.6g. "+ + "This is the MLX #3912 signature; check the MLX pin (%s).", + tc.K, tc.K%32, bad, len(a), tolerance, maxAbs, Version()) + } + }() + } + }) +}