diff --git a/CHANGELOG.md b/CHANGELOG.md index 3265c94..371ade3 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'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 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 c31651a..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 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 maxs[0] = max_score; } for (int i = 0; i < v_per_thread; i++) { - out[i] = static_cast(o[i]); + out[i] = static_cast(o[i]); } } -template +template [[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)]], diff --git a/src/kquant_sdpa.cpp b/src/kquant_sdpa.cpp index ae5e7ff..39d8007 100644 --- a/src/kquant_sdpa.cpp +++ b/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())); diff --git a/tests/test_sdpa.py b/tests/test_sdpa.py index cb13dd6..5f76fff 100644 --- a/tests/test_sdpa.py +++ b/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).