Skip to content

Fix int16 overflow in SDPA NAX mask indexing for KV sequences > 32K #3361

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
angeloskath merged 5 commits into from
Apr 10, 2026
Merged

Fix int16 overflow in SDPA NAX mask indexing for KV sequences > 32K #3361

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
ccfad15
Fix int16 overflow in SDPA NAX mask indexing for KV sequences > 32K
Clydingus Apr 3, 2026
fe183de
Add SDPA test for overflow in mask col_pos of KV sequences > 32K
Clydingus Apr 4, 2026
6f14a8c
Add unittest skipif not on gpu
Clydingus Apr 4, 2026
ed93902
Improve the mask check for NAX attention
angeloskath Apr 9, 2026
6f95bb0
Fix condition
angeloskath Apr 9, 2026
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: 6 additions & 0 deletions benchmarks/python/sdpa_bench.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -176,20 +176,26 @@ def get_gflop_count(B, M, N, K):
( 1, 1024, 1024, 64, 32, 8),
( 1, 2048, 2048, 64, 32, 8),
( 1, 4096, 4096, 64, 32, 8),
( 1, 4096, 5000, 64, 32, 8),
( 1, 2048, 32121, 64, 32, 8),
)

shapes_80 = (
# ( B, qsl, ksl, head_dim, n_qh, n_kvh)
( 1, 1024, 1024, 80, 32, 8),
( 1, 2048, 2048, 80, 32, 8),
( 1, 4096, 4096, 80, 32, 8),
( 1, 4096, 5000, 80, 32, 8),
( 1, 2048, 32121, 80, 32, 8),
)

shapes_128 = (
# ( B, qsl, ksl, head_dim, n_qh, n_kvh)
( 1, 1024, 1024, 128, 32, 8),
( 1, 2048, 2048, 128, 32, 8),
( 1, 4096, 4096, 128, 32, 8),
( 1, 4096, 5000, 128, 32, 8),
( 1, 2048, 32121, 128, 32, 8),
)
# fmt: on

Expand Down
84 changes: 56 additions & 28 deletions mlx/backend/metal/kernels/steel/attn/kernels/steel_attention_nax.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -286,17 +286,15 @@ template <
for (short iq = 0; iq < TQ; iq++) {
STEEL_PRAGMA_UNROLL
for (short ik = 0; ik < TK; ik++) {
const short row_pos = base_row + iq * kU;
const short col_pos = base_col + ik * kU;

thread auto& fg = Stile.frag_at(iq, ik);

STEEL_PRAGMA_UNROLL
for (short ii = 0; ii < stile_t::kFragThrRows; ii++) {
STEEL_PRAGMA_UNROLL
for (short jj = 0; jj < stile_t::kFragThrCols; jj++) {
const auto r = row_pos + ii * stile_t::kFragRowsJump + sm;
const auto c = col_pos + jj + sn;
const auto r =
base_row + iq * kU + ii * stile_t::kFragRowsJump + sm;
const auto c = base_col + ik * kU + jj + sn;
const auto loc = ii * stile_t::kFragThrCols + jj;
fg[loc] = (r < c) ? neg_inf : fg[loc];
}
Expand All @@ -317,32 +315,62 @@ template <
using mtile_t = NAXTile<melem_t, TQ, TK>;
using mfrag_t = typename mtile_t::frag_type;

STEEL_PRAGMA_UNROLL
for (short iq = 0; iq < TQ; iq++) {
STEEL_PRAGMA_UNROLL
for (short ik = 0; ik < TK; ik++) {
const short row_pos = base_row + iq * kU;
const short col_pos = base_col + ik * kU;

mfrag_t mfrag;
mtile_t::NAXFrag_t::load_safe(
mfrag,
mask,
int64_t(mask_params->M_strides[2]),
Int<1>{},
params->qL,
params->kL,
row_pos,
col_pos);
if (base_row + BQ <= params->qL && base_col + BK <= params->kL) {
for (short iq = 0; iq < TQ; iq++) {
STEEL_PRAGMA_UNROLL
for (short ik = 0; ik < TK; ik++) {
const int row_pos = base_row + iq * kU;
const int col_pos = base_col + ik * kU;

thread auto& fg = Stile.frag_at(iq, ik);
mfrag_t mfrag;
mtile_t::NAXFrag_t::load(
mfrag,
mask,
int64_t(mask_params->M_strides[2]),
Int<1>{},
row_pos,
col_pos);

thread auto& fg = Stile.frag_at(iq, ik);

STEEL_PRAGMA_UNROLL
for (short jj = 0; jj < mtile_t::kElemsPerFrag; jj++) {
if constexpr (is_bool) {
fg[jj] = mfrag[jj] ? fg[jj] : neg_inf;
} else {
fg[jj] += M_LOG2E_F * AccumType(mfrag[jj]);
}
}
}
}
} else {
STEEL_PRAGMA_UNROLL
for (short iq = 0; iq < TQ; iq++) {
STEEL_PRAGMA_UNROLL
for (short jj = 0; jj < mtile_t::kElemsPerFrag; jj++) {
if constexpr (is_bool) {
fg[jj] = mfrag[jj] ? fg[jj] : neg_inf;
} else {
fg[jj] += M_LOG2E_F * AccumType(mfrag[jj]);
for (short ik = 0; ik < TK; ik++) {
const int row_pos = base_row + iq * kU;
const int col_pos = base_col + ik * kU;

mfrag_t mfrag;
mtile_t::NAXFrag_t::load_safe(
mfrag,
mask,
int64_t(mask_params->M_strides[2]),
Int<1>{},
params->qL,
params->kL,
row_pos,
col_pos);

thread auto& fg = Stile.frag_at(iq, ik);

STEEL_PRAGMA_UNROLL
for (short jj = 0; jj < mtile_t::kElemsPerFrag; jj++) {
if constexpr (is_bool) {
fg[jj] = mfrag[jj] ? fg[jj] : neg_inf;
} else {
fg[jj] += M_LOG2E_F * AccumType(mfrag[jj]);
}
}
}
}
Expand Down
38 changes: 38 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 @@ -442,6 +442,44 @@ def test_sdpa(self):
diff = mx.abs(out_fst - out_ref) - atol * mx.abs(out_ref)
self.assertLessEqual(mx.max(diff).item(), atol)

@unittest.skipIf(not mx.is_available(mx.gpu), "too slow on CPU")
def test_sdpa_long_masked_sequence(self):
# Test for int16 overflow in steel_attention_nax.h mask
# indexing (col_pos declared as short, overflows when kL > 32767).
D = 64
dtype = mx.float16
atol = 1e-3 # Slightly looser than test_sdpa due to long masked sequences

for kL, active in [
(8192, 1024),
(36864, 1024),
(49152, 1024),
(66048, 1024),
]:
with self.subTest(kL=kL, active=active):
mx.random.seed(0)
qH, kH, qL = 32, 16, 512
scale = 1.0 / math.sqrt(D)

q = mx.random.normal(shape=(1, qH, qL, D)).astype(dtype)
k = mx.random.normal(shape=(1, kH, kL, D)).astype(dtype)
v = mx.random.normal(shape=(1, kH, kL, D)).astype(dtype)

# Additive mask: -1e4 for inactive, 0 for last `active` positions
mask = mx.full((1, 1, 1, kL), -1e4, dtype=dtype)
mask[..., kL - active :] = 0.0

out = mx.fast.scaled_dot_product_attention(
q, k, v, scale=scale, mask=mask
)
ref = mlx_ref_attn(q, k, v, scale=scale, mask=mask)

self.assertFalse(mx.isnan(out).any().item())
self.assertListEqual(list(out.shape), list(ref.shape))

diff = mx.abs(out - ref) - atol * mx.abs(ref)
self.assertLessEqual(mx.max(diff).item(), atol)

def test_sdpa_broadcast_mask(self):
mask = mx.array(True)
D = 64
Expand Down
Loading