L-electron-Rare/mlx
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

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

Browse Source 
Co-authored-by: Angelos Katharopoulos <a_katharopoulos@apple.com>
This commit is contained in:
Clydingus
2026-04-10 15:01:47 +08:00
committed by GitHub
parent d6d9b24801
commit a33b791615
3 changed files with 102 additions and 30 deletions
Download Patch File Download Diff File Expand all files Collapse all files
benchmarks/python/sdpa_bench.py
+6
View File
@@ -176,6 +176,8 @@ if __name__ == "__main__":
( 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 = (
@@ -183,6 +185,8 @@ if __name__ == "__main__":
( 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 = (
@@ -190,6 +194,8 @@ if __name__ == "__main__":
( 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
mlx/backend/metal/kernels/steel/attn/kernels/steel_attention_nax.h
+58 -30
View File
@@ -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];
}
@@ -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);
thread auto& fg = Stile.frag_at(iq, ik);
if (base_row + BQ <= params->qL && base_col + BK <= params->kL) {
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(
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 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]);
}
}
}
}
python/tests/test_fast_sdpa.py
+38
View File
@@ -442,6 +442,44 @@ class TestFastSDPA(mlx_tests.MLXTestCase):
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