diff --git a/benchmarks/python/sdpa_bench.py b/benchmarks/python/sdpa_bench.py index 5eb789de..4e1c0234 100644 --- a/benchmarks/python/sdpa_bench.py +++ b/benchmarks/python/sdpa_bench.py @@ -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 diff --git a/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention_nax.h b/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention_nax.h index 084ced69..adc9a427 100644 --- a/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention_nax.h +++ b/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention_nax.h @@ -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; 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]); + } } } } diff --git a/python/tests/test_fast_sdpa.py b/python/tests/test_fast_sdpa.py index 7606373c..115fd5f8 100644 --- a/python/tests/test_fast_sdpa.py +++ b/python/tests/test_fast_sdpa.py @@ -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