Skip to content

Preserve FP32 partial sums in two-pass vector SDPA #16

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Gajesh2007 merged 1 commit into from
Sep 6, 2026
Merged

Preserve FP32 partial sums in two-pass vector SDPA #16

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 3 additions & 3 deletions mlx/backend/metal/kernels/scaled_dot_product_attention.metal
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@ using namespace metal;
// SDPA vector instantiations
#define instantiate_sdpa_vector_aggregation(type, value_dim) \
instantiate_kernel( \
"sdpa_vector_2pass_2_" #type "_" #value_dim, \
"sdpa_vector_2pass_fp32partials_2_" #type "_" #value_dim, \
sdpa_vector_2pass_2, \
type, \
value_dim)
Expand All @@ -22,15 +22,15 @@ using namespace metal;
qk_dim, \
value_dim) \
instantiate_kernel( \
"sdpa_vector_2pass_1_" #type "_" #qk_dim "_" #value_dim, \
"sdpa_vector_2pass_fp32partials_1_" #type "_" #qk_dim "_" #value_dim, \
sdpa_vector_2pass_1, \
type, \
qk_dim, \
value_dim)

#define instantiate_sdpa_vector_gqa(type, qk_dim, value_dim, hpt) \
instantiate_kernel( \
"sdpa_vector_2pass_1_gqa_" #type "_" #qk_dim "_" #value_dim, \
"sdpa_vector_2pass_fp32partials_1_gqa_" #type "_" #qk_dim "_" #value_dim, \
sdpa_vector_2pass_1_gqa, \
type, \
qk_dim, \
Expand Down
12 changes: 6 additions & 6 deletions mlx/backend/metal/kernels/sdpa_vector.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -181,7 +181,7 @@ template <typename T, int D, int V = D>
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(7)]],
Expand Down Expand Up @@ -313,7 +313,7 @@ template <typename T, int D, int V = D>
}

for (int i = 0; i < v_per_thread; i++) {
out[i] = static_cast<T>(o[i]);
out[i] = o[i];
}
}

Expand All @@ -327,7 +327,7 @@ template <typename T, int D, int V, int G, int HPT>
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(7)]],
Expand Down Expand Up @@ -449,10 +449,10 @@ template <typename T, int D, int V, int G, int HPT>
}

const int o_offset = base_head + g;
device T* op =
device float* op =
out + o_offset * blocks * V + block_idx * V + simd_lid * v_per_thread;
for (int i = 0; i < v_per_thread; i++) {
op[i] = static_cast<T>(acc[i]);
op[i] = acc[i];
}
if (simd_lid == 0) {
sums[o_offset * blocks + block_idx] = denom;
Expand All @@ -462,7 +462,7 @@ template <typename T, int D, int V, int G, int HPT>

template <typename T, int D>
[[kernel]] void 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)]],
Expand Down
7 changes: 4 additions & 3 deletions mlx/backend/metal/scaled_dot_product_attention.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -465,7 +465,7 @@ void sdpa_vector_2pass(
// Set the kernel name
std::string kname;
kname.reserve(64);
kname += "sdpa_vector_2pass_1";
kname += "sdpa_vector_2pass_fp32partials_1";
if (!mask && !sinks && q.shape(2) == 1 && q.shape(1) == 8 * k.shape(1) &&
q.shape(-1) == v.shape(-1) && (q.shape(-1) == 64 || q.shape(-1) == 128) &&
k.shape(2) >= 8192) {
Expand Down Expand Up @@ -535,7 +535,8 @@ void sdpa_vector_2pass(
intermediate_shape.end(), out.shape().begin(), out.shape().end() - 1);
intermediate_shape.push_back(blocks);
intermediate_shape.push_back(out.shape().back());
array intermediate(intermediate_shape, q.dtype(), nullptr, {});
// Keep unnormalized partials in float until the final reduction.
array intermediate(intermediate_shape, float32, nullptr, {});
intermediate_shape.pop_back();
array sums(intermediate_shape, float32, nullptr, {});
array maxs(std::move(intermediate_shape), float32, nullptr, {});
Expand Down Expand Up @@ -607,7 +608,7 @@ void sdpa_vector_2pass(

// Final pass
kname.clear();
kname = "sdpa_vector_2pass_2_";
kname = "sdpa_vector_2pass_fp32partials_2_";
kname += get_type_string(q.dtype());
kname += "_";
kname += std::to_string(v.shape(-1));
Expand Down
47 changes: 47 additions & 0 deletions python/tests/test_fast_sdpa.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@
import os
import unittest
from itertools import product
from unittest.mock import patch

import mlx.core as mx
import mlx_tests
Expand Down Expand Up @@ -356,6 +357,52 @@ def test_sdpa_vector_gqa_long(self):
out = mx.fast.scaled_dot_product_attention(q, k, v, scale=scale)
self.assertTrue(mx.allclose(ref, out, atol=1e-4, rtol=1e-4))

@unittest.skipIf(not mx.is_available(mx.gpu), "GPU kernel path only")
def test_sdpa_two_pass_partial_cancellation(self):
# Uniform attention has the exact output 1 / L. Casting an
# unnormalized block sum to the input dtype loses the residual 1.
L = 8192
for dtype, D, blocks, masked in product(
(mx.bfloat16, mx.float16, mx.float32),
(64, 128, 256),
(32, 128),
(False, True),
):
with self.subTest(dtype=dtype, D=D, blocks=blocks, masked=masked):
amplitude = 2048 if dtype == mx.float16 else 256
values = np.zeros((1, 1, L, D), dtype=np.float32)
values[:, :, 0, :] = amplitude
# D64/D128 without a mask uses contiguous GQA chunks.
gqa = D in (64, 128) and not masked
values[:, :, 1 if gqa else blocks, :] = 1
values[:, :, L - 1 if gqa else 1, :] = -amplitude
q = mx.zeros((1, 8, 1, D), dtype=dtype)
k = mx.zeros((1, 1, L, D), dtype=dtype)
v = mx.array(values, dtype=dtype)
mask = mx.ones((L,), dtype=mx.bool_) if masked else None
with patch.dict(os.environ, MLX_SDPA_BLOCKS=str(blocks)):
out = mx.fast.scaled_dot_product_attention(
q, k, v, scale=1, mask=mask
)
mx.eval(out)
self.assertEqual(out.dtype, dtype)
self.assertTrue(mx.array_equal(out, mx.full(out.shape, 1 / L, dtype)))

@unittest.skipIf(not mx.is_available(mx.gpu), "GPU kernel path only")
def test_sdpa_two_pass_partial_overflow(self):
# The final mean is representable in FP16, but its partial sums are not.
L = 8192
for D, blocks in product((64, 128, 256), (32, 128)):
with self.subTest(D=D, blocks=blocks):
q = mx.zeros((1, 8, 1, D), dtype=mx.float16)
k = mx.zeros((1, 1, L, D), dtype=mx.float16)
v = mx.full((1, 1, L, D), 32768, dtype=mx.float16)
with patch.dict(os.environ, MLX_SDPA_BLOCKS=str(blocks)):
out = mx.fast.scaled_dot_product_attention(q, k, v, scale=1)
mx.eval(out)
self.assertTrue(mx.all(mx.isfinite(out)))
self.assertTrue(mx.array_equal(out, mx.full(out.shape, 32768, mx.float16)))

def test_sdpa_fully_masked(self):
Lkv = 8
mask = mx.array(False)
Expand Down