spec_decode_verify transforms and qmv_quad gs 128 coverage (gemma4-qat#323, #324)

mlx/fast.cpp

@@ -1268,7 +1268,8 @@ static std::vector<array> spec_decode_verify_composition(
           /*axis=*/1,
           /*keepdims=*/false,
           s); // [B]
-  auto n_acc2 = reshape(n_acc, Shape{B, 1}, s);
+  // An index, not a value: transforms see a zero tangent instead of throwing.
+  auto n_acc2 = stop_gradient(reshape(n_acc, Shape{B, 1}, s), s);
   auto corrected = take_along_axis(t, n_acc2, /*axis=*/1, s); // [B, 1]
   auto j1 = broadcast_to(
       reshape(arange(K + 1, int32, s), Shape{1, K + 1}, s), Shape{B, K + 1}, s);

python/tests/test_fast.py

@@ -1465,6 +1465,23 @@ def test_spec_decode_verify_empty_draft(self):
                     self.assertEqual(n.tolist(), [0] * B)
                     self.assertEqual(c.tolist(), [[t] for t in bonus])
 
+    def test_spec_decode_verify_transforms(self):
+        # The outputs are integer token ids, so jvp and vjp give zero
+        # (co)tangents rather than throwing on the accept-length gather.
+        draft = mx.array([[1, 2, 3]], dtype=mx.int32)
+        lg = mx.random.normal((1, 4, 8), key=mx.random.key(0))
+
+        def f(x):
+            return mx.fast.spec_decode_verify(draft, x)
+
+        _, tangents = mx.jvp(f, [lg], [mx.ones_like(lg)])
+        for t in tangents:
+            self.assertTrue(mx.all(t == 0).item())
+        cotangents = [mx.ones((1,), mx.int32), mx.ones((1, 4), mx.int32)]
+        _, (g,) = mx.vjp(f, [lg], cotangents)
+        self.assertEqual(g.shape, lg.shape)
+        self.assertTrue(mx.all(g == 0).item())
+
     def test_spec_decode_verify_empty_batch(self):
         # B = 0: empty, correctly shaped int32 results on every backend.
         for K in (3, 0):

python/tests/test_quantized.py

@@ -911,10 +911,14 @@ def test_qmv_quad_ragged_n(self):
         # quad rows past N must read no scale or bias (gemma4-qat#302). N = 65
         # leaves one valid row in tile 2, N = 72 leaves eight. Explicit and
         # implied (0-d factor) biases are both covered; affine_sym never reaches
-        # qmv_quad.
+        # qmv_quad. Group size 128 runs at K = 128 only.
         key = mx.random.key(7)
         cases = [
-            (bits, gs, K) for bits in (1, 2, 4, 8) for gs in (32, 64) for K in (64, 128)
+            (bits, gs, K)
+            for bits in (1, 2, 4, 8)
+            for gs in (32, 64, 128)
+            for K in (64, 128)
+            if gs <= K
         ]
         for bits, gs, K in cases:
             if K == 64 and bits == 1: