fix(sdpa_vector): float16 partials could overflow to inf on long flat attention

CHANGELOG.md

@@ -6,6 +6,11 @@ adhere to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
 ## [Unreleased]
 
+### Fixed
+- sdpa_vector's float16 pass-1 partials could overflow to inf under long
+  flat attention with large V outliers; they now store as float32
+  (bfloat16 partials are unchanged).
+
 ## [0.3.13]
 
 ### Changed

metal/kq_sdpa.metal

@@ -5,22 +5,25 @@
 #include "mlx/backend/metal/kernels/steel/attn/mma.h"
 #include "mlx/backend/metal/kernels/kq_sdpa.h"
 
-#define instantiate_kq_sdpa(type, D)                                  \
+#define instantiate_kq_sdpa(type, ptype, D)                           \
   instantiate_kernel(                                                 \
       "kq_sdpa_vector_2pass_1_" #type "_" #D,                         \
       kq_sdpa_vector_2pass_1,                                         \
       type,                                                           \
+      ptype,                                                          \
       D)                                                              \
   instantiate_kernel(                                                 \
       "kq_sdpa_vector_2pass_2_" #type "_" #D,                         \
       kq_sdpa_vector_2pass_2,                                         \
       type,                                                           \
+      ptype,                                                          \
       D)
 
-instantiate_kq_sdpa(bfloat16_t, 256)
-instantiate_kq_sdpa(bfloat16_t, 512)
-instantiate_kq_sdpa(float16_t, 256)
-instantiate_kq_sdpa(float16_t, 512)
+// Partial store type per input dtype (see the PT note in kq_sdpa.h).
+instantiate_kq_sdpa(bfloat16_t, bfloat16_t, 256)
+instantiate_kq_sdpa(bfloat16_t, bfloat16_t, 512)
+instantiate_kq_sdpa(float16_t, float, 256)
+instantiate_kq_sdpa(float16_t, float, 512)
 
 #define instantiate_kq_sdpa_gqa(type, D, C)                           \
   instantiate_kernel(                                                 \

metal/mlx/backend/metal/kernels/kq_sdpa.h

@@ -32,12 +32,15 @@ constant bool gqa_cascade [[function_constant(7)]];
 // [0, N) axis. Compiled out when false.
 constant bool gqa_paged [[function_constant(8)]];
 
-template <typename T, int D, int V = D>
+// PT = pass-1 partial store type. float16 inputs use float: the
+// un-normalized accumulator state can exceed the fp16 ceiling. bfloat16
+// keeps 16-bit stores (range-safe) and the pre-fix bandwidth.
+template <typename T, typename PT, int D, int V = D>
 [[kernel]] void kq_sdpa_vector_2pass_1(
     const device T* queries [[buffer(0)]],
     const device T* keys [[buffer(1)]],
     const device T* values [[buffer(2)]],
-    device T* out [[buffer(3)]],
+    device PT* out [[buffer(3)]],
     device float* sums [[buffer(4)]],
     device float* maxs [[buffer(5)]],
     const constant int& N [[buffer(6)]],
@@ -119,13 +122,13 @@ template <typename T, int D, int V = D>
     maxs[0] = max_score;
   }
   for (int i = 0; i < v_per_thread; i++) {
-    out[i] = static_cast<T>(o[i]);
+    out[i] = static_cast<PT>(o[i]);
   }
 }
 
-template <typename T, int D>
+template <typename T, typename PT, int D>
 [[kernel]] void kq_sdpa_vector_2pass_2(
-    const device T* partials [[buffer(0)]],
+    const device PT* partials [[buffer(0)]],
     const device float* sums [[buffer(1)]],
     const device float* maxs [[buffer(2)]],
     device T* out [[buffer(3)]],

src/kquant_sdpa.cpp

@@ -125,7 +125,12 @@ void KQuantSDPA::eval_gpu(
   // Per-block partials + running max/sum, reduced by pass 2.
   mx::Shape part_shape = {B, n_q_heads, qL, blocks, D};
   mx::Shape red_shape = {B, n_q_heads, qL, blocks};
-  array partials(part_shape, q.dtype(), nullptr, {});
+  // Un-normalized online-softmax accumulator state is unbounded by the
+  // model, so a float16 store can overflow: float16 inputs get f32
+  // partials. bfloat16 has the range and keeps 16-bit stores (pre-fix
+  // bandwidth). Must match the PT instantiation map in kq_sdpa.metal.
+  auto part_dtype = q.dtype() == mx::float16 ? mx::float32 : q.dtype();
+  array partials(part_shape, part_dtype, nullptr, {});
   array sums(red_shape, mx::float32, nullptr, {});
   array maxs(red_shape, mx::float32, nullptr, {});
   partials.set_data(mx::allocator::malloc(partials.nbytes()));

tests/test_sdpa.py

@@ -125,6 +125,42 @@ def test_sdpa_vector_gqa(Hq, Hkv):
     _check(512, qL=4, kL=2048, dtype=mx.bfloat16, Hq=Hq, Hkv=Hkv)
 
 
+@pytest.mark.parametrize("D", [256, 512])
+def test_sdpa_vector_f32_partials_outlier_v(D):
+    # Un-normalized pass-1 partials scale with keys-per-block times |v|.
+    # A float16 partial store overflows 65504 under flat attention with
+    # outlier V channels; f32 partials must stay finite and rounding-level.
+    scale = 1.0 / (D**0.5)
+    q, k, v = _make(1, 8, 2, 1, 16384, D, mx.float16, seed=3, strided=False)
+    k = (0.05 * k.astype(mx.float32)).astype(mx.float16)  # flatten attention
+    v[:, :, :, ::64] = 2048.0
+    mx.eval(k, v)
+    got = kq.sdpa_vector(q, k, v, scale, causal=False)
+    ref = _ref_sdpa(q, k, v, scale, causal=False)
+    _eval_or_skip(got, ref)
+    assert bool(mx.all(mx.isfinite(got.astype(mx.float32))).item())
+    rel = _rel(got, ref)
+    assert rel < REL_BOUND[mx.float16], f"D={D} rel {rel:.3e}"
+
+
+@pytest.mark.parametrize("D", [256, 512])
+def test_sdpa_vector_bf16_partials_outlier_v(D):
+    # bfloat16 keeps 16-bit pass-1 partials (range covers the outlier
+    # magnitudes that overflow fp16); same stress must stay finite and
+    # within the bf16 rounding bound.
+    scale = 1.0 / (D**0.5)
+    q, k, v = _make(1, 8, 2, 1, 16384, D, mx.bfloat16, seed=3, strided=False)
+    k = (0.05 * k.astype(mx.float32)).astype(mx.bfloat16)
+    v[:, :, :, ::64] = 2048.0
+    mx.eval(k, v)
+    got = kq.sdpa_vector(q, k, v, scale, causal=False)
+    ref = _ref_sdpa(q, k, v, scale, causal=False)
+    _eval_or_skip(got, ref)
+    assert bool(mx.all(mx.isfinite(got.astype(mx.float32))).item())
+    rel = _rel(got, ref)
+    assert rel < REL_BOUND[mx.bfloat16], f"D={D} rel {rel:.3e}"
+
+
 def _ref_sdpa_sinks(q, k, v, scale, sinks):
     """f32 reference with per-q-head sink logits: an extra softmax column
     with no value row (raises the max / adds to the denominator only).