From 144ca794a9a80e00de8617528cd46fcea64c9128 Mon Sep 17 00:00:00 2001 From: Asher Feldman <59994+asher@users.noreply.github.com> Date: Sat, 15 Aug 2026 14:43:12 -0700 Subject: [PATCH 1/2] fix(sdpa): float32 pass-1 partials in sdpa_vector --- CHANGELOG.md | 5 +++++ metal/mlx/backend/metal/kernels/kq_sdpa.h | 6 +++--- src/kquant_sdpa.cpp | 5 ++++- tests/test_sdpa.py | 18 ++++++++++++++++++ 4 files changed, 30 insertions(+), 4 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 3265c94..f5f8417 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,11 @@ adhere to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). ## [Unreleased] +### Fixed +- sdpa_vector stored its pass-1 partials in the activation dtype; float16 + activations could overflow them to inf on long flat attention with large + V outliers. Partials are float32 now, matching the GQA decode kernels. + ## [0.3.13] ### Changed diff --git a/metal/mlx/backend/metal/kernels/kq_sdpa.h b/metal/mlx/backend/metal/kernels/kq_sdpa.h index c31651a..e14d9ce 100644 --- a/metal/mlx/backend/metal/kernels/kq_sdpa.h +++ b/metal/mlx/backend/metal/kernels/kq_sdpa.h @@ -37,7 +37,7 @@ template const device T* queries [[buffer(0)]], const device T* keys [[buffer(1)]], const device T* values [[buffer(2)]], - device T* out [[buffer(3)]], + device float* out [[buffer(3)]], device float* sums [[buffer(4)]], device float* maxs [[buffer(5)]], const constant int& N [[buffer(6)]], @@ -119,13 +119,13 @@ template maxs[0] = max_score; } for (int i = 0; i < v_per_thread; i++) { - out[i] = static_cast(o[i]); + out[i] = o[i]; } } template [[kernel]] void kq_sdpa_vector_2pass_2( - const device T* partials [[buffer(0)]], + const device float* partials [[buffer(0)]], const device float* sums [[buffer(1)]], const device float* maxs [[buffer(2)]], device T* out [[buffer(3)]], diff --git a/src/kquant_sdpa.cpp b/src/kquant_sdpa.cpp index ae5e7ff..76bffc4 100644 --- a/src/kquant_sdpa.cpp +++ b/src/kquant_sdpa.cpp @@ -125,7 +125,10 @@ 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, {}); + // f32 partials: un-normalized online-softmax accumulator state is + // unbounded by the model, so a 16-bit store can overflow (fp16) or lose + // mantissa (bf16). The gqa/cascade/paged siblings already use f32. + array partials(part_shape, mx::float32, nullptr, {}); array sums(red_shape, mx::float32, nullptr, {}); array maxs(red_shape, mx::float32, nullptr, {}); partials.set_data(mx::allocator::malloc(partials.nbytes())); diff --git a/tests/test_sdpa.py b/tests/test_sdpa.py index cb13dd6..1e33f60 100644 --- a/tests/test_sdpa.py +++ b/tests/test_sdpa.py @@ -125,6 +125,24 @@ 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}" + + 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). From 606eaa5d5ef1d5795d2ae21e95d0155aed7d6e59 Mon Sep 17 00:00:00 2001 From: Asher Feldman <59994+asher@users.noreply.github.com> Date: Sun, 16 Aug 2026 00:08:38 -0700 Subject: [PATCH 2/2] fix(sdpa): dtype-gate f32 pass-1 partials, float16 inputs only --- CHANGELOG.md | 6 +++--- metal/kq_sdpa.metal | 13 ++++++++----- metal/mlx/backend/metal/kernels/kq_sdpa.h | 13 ++++++++----- src/kquant_sdpa.cpp | 10 ++++++---- tests/test_sdpa.py | 18 ++++++++++++++++++ 5 files changed, 43 insertions(+), 17 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index f5f8417..371ade3 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,9 +7,9 @@ adhere to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). ## [Unreleased] ### Fixed -- sdpa_vector stored its pass-1 partials in the activation dtype; float16 - activations could overflow them to inf on long flat attention with large - V outliers. Partials are float32 now, matching the GQA decode kernels. +- 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] diff --git a/metal/kq_sdpa.metal b/metal/kq_sdpa.metal index 72bea44..2022859 100644 --- a/metal/kq_sdpa.metal +++ b/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( \ diff --git a/metal/mlx/backend/metal/kernels/kq_sdpa.h b/metal/mlx/backend/metal/kernels/kq_sdpa.h index e14d9ce..1aee1b0 100644 --- a/metal/mlx/backend/metal/kernels/kq_sdpa.h +++ b/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 +// 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 [[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 float* 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 maxs[0] = max_score; } for (int i = 0; i < v_per_thread; i++) { - out[i] = o[i]; + out[i] = static_cast(o[i]); } } -template +template [[kernel]] void kq_sdpa_vector_2pass_2( - const device float* 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)]], diff --git a/src/kquant_sdpa.cpp b/src/kquant_sdpa.cpp index 76bffc4..39d8007 100644 --- a/src/kquant_sdpa.cpp +++ b/src/kquant_sdpa.cpp @@ -125,10 +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}; - // f32 partials: un-normalized online-softmax accumulator state is - // unbounded by the model, so a 16-bit store can overflow (fp16) or lose - // mantissa (bf16). The gqa/cascade/paged siblings already use f32. - array partials(part_shape, mx::float32, 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())); diff --git a/tests/test_sdpa.py b/tests/test_sdpa.py index 1e33f60..5f76fff 100644 --- a/tests/test_sdpa.py +++ b/tests/test_sdpa.py @@ -143,6 +143,24 @@ def test_sdpa_vector_f32_partials_outlier_v(D): 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).