Skip to content

Improve metal memory usage for SDPA D256 #4505

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

Open
dhiltgen wants to merge 1 commit into
base: main
Choose a base branch
from
+165 −38
Open

Improve metal memory usage for SDPA D256 #4505

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
111 changes: 89 additions & 22 deletions mlx/backend/metal/kernels/steel/attn/kernels/steel_attention.h
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
// Copyright © 2024-25 Apple Inc.
// Copyright © 2024-26 Apple Inc.

#include "mlx/backend/metal/kernels/steel/attn/attn.h"

Expand Down Expand Up @@ -119,8 +119,10 @@ template <
constexpr short tgp_mem_0 = (BK + padK) * (BD);
constexpr short tgp_mem_1 = BK * (BD + padV);
constexpr short tgp_mem_s = tgp_mem_0 > tgp_mem_1 ? tgp_mem_0 : tgp_mem_1;
constexpr bool q_in_registers = WN == 2 && (BK == 32 || sizeof(T) == 4);
constexpr short q_tgp_mem_s = q_in_registers ? 1 : BQ * (BD + padQ);

threadgroup T Q_smem[BQ * (BD + padQ)];
threadgroup T Q_smem[q_tgp_mem_s];
threadgroup T KV_smem[tgp_mem_s];

threadgroup T* Qs = Q_smem;
Expand Down Expand Up @@ -168,8 +170,13 @@ template <
// Prepare MMA tiles
constexpr short kFragSize = 8; // MMAFrag size
using MMAFrag_acc_t = BaseMMAFrag<AccumType, kFragSize, kFragSize>;
using MMAInType = metal::conditional_t<WN == 2, T, AccumType>;
using MMAFrag_in_t = BaseMMAFrag<MMAInType, kFragSize, kFragSize>;

constexpr int kNWarps = WM * WN;
static_assert(WN == 1 || WN == 2, "WN must be 1 or 2");
static_assert(BD % WN == 0, "The head dim must split evenly across WN");

constexpr int kNWarps = WM;
static_assert(
BQ >= (kNWarps * kFragSize) && BQ % (kNWarps * kFragSize) == 0,
"Each simdgroup must host atleast 1 simdgroup matrix along Q sequence.");
Expand All @@ -180,37 +187,63 @@ template <
constexpr int TK = BK / kFragSize;
// HeadDim frags (all warps load the same frags)
constexpr int TD = BD / kFragSize;
constexpr int TDh = TD / WN;
constexpr int BDh = BD / WN;

static_assert(TQ == 1, "Check TQ");

MMATile<AccumType, TQ, 1, MMAFrag_acc_t> Qtile;
MMATile<AccumType, 1, TK, MMAFrag_acc_t> Ktile;
MMATile<MMAInType, TQ, 1, MMAFrag_in_t> Qtile;
MMATile<MMAInType, 1, TK, MMAFrag_in_t> Ktile;
MMATile<AccumType, TQ, TK, MMAFrag_acc_t> Stile;
MMATile<AccumType, 1, 1, MMAFrag_acc_t> Vtile;
MMATile<AccumType, TQ, TD, MMAFrag_acc_t> Otile;
MMATile<MMAInType, 1, 1, MMAFrag_in_t> Vtile;
MMATile<AccumType, TQ, TDh, MMAFrag_acc_t> Otile;
MMATile<MMAInType, 1, 1, MMAFrag_in_t> Qtiles[q_in_registers ? TDh : 1];

Otile.clear();

// Prepare mma tile offsets
const short2 simd_coord = MMAFrag_acc_t::get_coord(simd_lane_id);
const short sm = simd_coord.y;
const short sn = simd_coord.x;
const short tm = kFragSize * TQ * simd_group_id;
const short row_group = simd_group_id / WN;
const short d_half = simd_group_id % WN;
const short tm = kFragSize * TQ * row_group;

const short Qs_offset = (tm + sm) * LDQ_tgp + sn;
const short Ks_offset = sm * LDK_tgp + sn;
const short Vs_offset = sm * LDV_tgp + sn;
const short Qs_offset = (tm + sm) * LDQ_tgp + d_half * BDh + sn;
const short Ks_offset = d_half * BDh * LDK_tgp + sm * LDK_tgp + sn;
const short Vs_offset = sm * LDV_tgp + d_half * BDh + sn;

constexpr short Qs_tile_stride = kFragSize;
constexpr short Ks_tile_stride = kFragSize * LDK_tgp;

threadgroup_barrier(mem_flags::mem_threadgroup);

// Load Q blocks
if (!align_Q && int(tid.x) == (params->NQ_aligned)) {
loader_q.load_safe(short2(BD, params->qL_rem));
if constexpr (!q_in_registers) {
if (!align_Q && int(tid.x) == (params->NQ_aligned)) {
loader_q.load_safe(short2(BD, params->qL_rem));
} else {
loader_q.load_unsafe();
}
} else {
loader_q.load_unsafe();
STEEL_PRAGMA_UNROLL
for (short dd = 0; dd < TDh; dd++) {
const device T* q_src = Q + (tm + sm) * params->Q_strides[2] +
d_half * BDh + dd * kFragSize + sn;
if (!align_Q && int(tid.x) == (params->NQ_aligned)) {
MMAFrag_in_t::load_safe(
Qtiles[dd].frag_at(0, 0),
q_src,
params->Q_strides[2],
Int<1>{},
params->qL_rem - (tm + sm),
BDh - (dd * kFragSize + sn),
Int<0>{},
Int<0>{});
} else {
MMAFrag_in_t::load(
Qtiles[dd].frag_at(0, 0), q_src, params->Q_strides[2], Int<1>{});
}
}
}

// Init row reduction variables
Expand Down Expand Up @@ -262,17 +295,51 @@ template <
threadgroup_barrier(mem_flags::mem_threadgroup);

STEEL_PRAGMA_UNROLL
for (short dd = 0; dd < TD; dd++) {
for (short dd = 0; dd < TDh; dd++) {
simdgroup_barrier(mem_flags::mem_none);

Qtile.template load<T, 1, 1, LDQ_tgp, 1>(
&Qs[Qs_offset + dd * Qs_tile_stride]);
if constexpr (!q_in_registers) {
Qtile.template load<T, 1, 1, LDQ_tgp, 1>(
&Qs[Qs_offset + dd * Qs_tile_stride]);
}
Ktile.template load<T, 1, 1, LDK_tgp, 1>(
&Ks[Ks_offset + dd * Ks_tile_stride]);

simdgroup_barrier(mem_flags::mem_none);

tile_matmad(Stile, Qtile, Ktile, Stile);
if constexpr (!q_in_registers) {
tile_matmad(Stile, Qtile, Ktile, Stile);
} else {
tile_matmad(Stile, Qtiles[dd], Ktile, Stile);
}
}

if constexpr (WN == 2) {
constexpr int s_xchg_stride = decltype(Stile)::kElemsPerTile * 32;
constexpr int s_xchg_size = WM * WN * s_xchg_stride;
static_assert(
s_xchg_size * sizeof(AccumType) <= tgp_mem_s * sizeof(T),
"The score exchange must fit in KV_smem");

threadgroup_barrier(mem_flags::mem_threadgroup);
threadgroup AccumType* s_xchg =
reinterpret_cast<threadgroup AccumType*>(KV_smem);
threadgroup AccumType* slot =
s_xchg + (row_group * WN + d_half) * s_xchg_stride;
const short base = short(simd_lane_id) * decltype(Stile)::kElemsPerTile;

STEEL_PRAGMA_UNROLL
for (short ii = 0; ii < decltype(Stile)::kElemsPerTile; ii++) {
slot[base + ii] = Stile.elems()[ii];
}
threadgroup_barrier(mem_flags::mem_threadgroup);

const threadgroup AccumType* peer =
s_xchg + (row_group * WN + 1 - d_half) * s_xchg_stride;
STEEL_PRAGMA_UNROLL
for (short ii = 0; ii < decltype(Stile)::kElemsPerTile; ii++) {
Stile.elems()[ii] += peer[base + ii];
}
}

// Apply scale in float32
Expand Down Expand Up @@ -425,7 +492,7 @@ template <
STEEL_PRAGMA_UNROLL
for (short iq = 0; iq < TQ; iq++) {
STEEL_PRAGMA_UNROLL
for (short id = 0; id < TD; id++) {
for (short id = 0; id < TDh; id++) {
STEEL_PRAGMA_UNROLL
for (short ik = 0; ik < TK; ik++) {
if constexpr (BD >= 128) {
Expand Down Expand Up @@ -461,10 +528,10 @@ template <
threadgroup_barrier(mem_flags::mem_none);

// Store results
O += (tm + sm) * params->O_strides[2] + sn;
O += (tm + sm) * params->O_strides[2] + d_half * BDh + sn;

if (!align_Q && int(tid.x) == (params->NQ_aligned)) {
auto dst_tile_dims = short2(BD - sn, params->qL_rem - (tm + sm));
auto dst_tile_dims = short2(BDh - sn, params->qL_rem - (tm + sm));

if (dst_tile_dims.x <= 0 || dst_tile_dims.y <= 0)
return;
Expand Down
22 changes: 13 additions & 9 deletions mlx/backend/metal/kernels/steel/attn/kernels/steel_attention.metal
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
// Copyright © 2024-25 Apple Inc.
// Copyright © 2024-26 Apple Inc.

// clang-format off
#include "mlx/backend/metal/kernels/utils.h"
Expand All @@ -11,21 +11,25 @@
"_wm" #wm "_wn" #wn "_mask" #mname, \
attention, dtype, bq, bk, bd, wm, wn, mtype, float)

#define instantiate_attn_shapes_helper(iname, itype, mname, mtype) \
instantiate_attn(iname, itype, 32, 16, 256, 4, 1, mname, mtype) \
#define instantiate_attn_shapes_helper(iname, itype, mname, mtype, bk256) \
instantiate_attn(iname, itype, 32, bk256, 256, 4, 2, mname, mtype) \
instantiate_attn(iname, itype, 32, 16, 192, 4, 1, mname, mtype) \
instantiate_attn(iname, itype, 32, 16, 128, 4, 1, mname, mtype) \
instantiate_attn(iname, itype, 32, 32, 96, 4, 1, mname, mtype) \
instantiate_attn(iname, itype, 32, 32, 80, 4, 1, mname, mtype) \
instantiate_attn(iname, itype, 32, 32, 72, 4, 1, mname, mtype) \
instantiate_attn(iname, itype, 32, 32, 64, 4, 1, mname, mtype)

#define instantiate_attn_mask_helper(iname, itype) \
instantiate_attn_shapes_helper(iname, itype, iname, itype) \
instantiate_attn_shapes_helper(iname, itype, bool_, bool)
#define instantiate_attn_mask_helper(iname, itype, bk256) \
instantiate_attn_shapes_helper(iname, itype, iname, itype, bk256) \
instantiate_attn_shapes_helper(iname, itype, bool_, bool, bk256)

instantiate_attn_mask_helper(float16, half);
instantiate_attn_mask_helper(bfloat16, bfloat16_t);
instantiate_attn_mask_helper(float16, half, 32);
instantiate_attn_mask_helper(bfloat16, bfloat16_t, 32);
instantiate_attn(float16, half, 32, 16, 256, 4, 2, float16, half);
instantiate_attn(float16, half, 32, 16, 256, 4, 2, bool_, bool);
instantiate_attn(bfloat16, bfloat16_t, 32, 16, 256, 4, 2, bfloat16, bfloat16_t);
instantiate_attn(bfloat16, bfloat16_t, 32, 16, 256, 4, 2, bool_, bool);

instantiate_attn_mask_helper(float32, float);
instantiate_attn_mask_helper(float32, float, 16);
// clang-format on
15 changes: 12 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
@@ -1,4 +1,4 @@
// Copyright © 2024 Apple Inc.
// Copyright © 2024-26 Apple Inc.
#include <sstream>

#include "mlx/backend/common/compiled.h"
Expand Down Expand Up @@ -304,11 +304,14 @@ void sdpa_full_self_attention_metal(
using namespace mlx::steel;

int wm = 4;
int wn = 1;

int bd = q.shape(-1);
char devc = d.get_architecture().back();
int wn = bd == 256 ? 2 : 1;
int bq = 32;
int bk = bd < 128 ? 32 : 16;
int bk = bd == 256 && q.dtype() != float32 && devc == 'd'
? 32
: (bd < 128 ? 32 : 16);

const bool align_Q = (qL % bq) == 0;
const bool align_K = (kL % bk) == 0;
Expand Down Expand Up @@ -862,6 +865,12 @@ bool ScaledDotProductAttention::use_fallback(
return false;
}

if (!metal::is_nax_available() && q.dtype() != float32 && do_causal &&
query_head_dim == 256 && query_sequence_length >= 2048 &&
query_sequence_length == k.shape(2)) {
return false;
}

// Unfused path is faster for following shapes.
if (query_sequence_length > 8) {
return query_head_dim == 192 || query_head_dim == 256;
Expand Down
55 changes: 51 additions & 4 deletions python/tests/test_fast_sdpa.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -929,15 +929,35 @@ def test_grad(slow, fast, args):
).sum()
test_grad(loss_slow, loss_fast, [q, k, v])

@unittest.skipIf(not mx.metal.is_available(), "Metal kernel path only")
def test_sdpa_d256_default_metal(self):
if mx.default_device() != mx.gpu:
self.skipTest("requires GPU")
mx.random.seed(0)

for dtype in (mx.float16, mx.bfloat16):
with self.subTest(dtype=dtype):
q = (0.5 * mx.random.normal((1, 2, 2048, 256))).astype(dtype)
k = (0.5 * mx.random.normal((1, 1, 2048, 256))).astype(dtype)
v = (0.5 * mx.random.normal((1, 1, 2048, 256))).astype(dtype)
scale = 256**-0.5
ref = mlx_ref_attn(q, k, v, scale=scale, mask="causal")
out = mx.fast.scaled_dot_product_attention(
q, k, v, scale=scale, mask="causal"
)
tol = 1e-2 if dtype == mx.bfloat16 else 1e-3
self.assertTrue(mx.allclose(ref, out, atol=tol, rtol=tol))

@unittest.skipIf(not mx.metal.is_available(), "Metal kernel path only")
def test_sdpa_force_fused_metal(self):
if mx.default_device() != mx.gpu:
self.skipTest("requires GPU")
mx.random.seed(0)

def make_qkv(qL, kL, D, qH=8, kH=8):
q = mx.random.normal((1, qH, qL, D), mx.float16)
k = mx.random.normal((1, kH, kL, D), mx.float16)
v = mx.random.normal((1, kH, kL, D), mx.float16)
def make_qkv(qL, kL, D, qH=8, kH=8, dtype=mx.float16):
q = mx.random.normal((1, qH, qL, D), dtype)
k = mx.random.normal((1, kH, kL, D), dtype)
v = mx.random.normal((1, kH, kL, D), dtype)
return q, k, v

# Full attention kernel.
Expand All @@ -951,6 +971,33 @@ def make_qkv(qL, kL, D, qH=8, kH=8):
)
self.assertTrue(mx.allclose(ref, out, atol=1e-3, rtol=1e-3))

for dtype, mask_kind in product(
(mx.bfloat16, mx.float32), (None, "causal", "bool", "additive")
):
with self.subTest(head_dim=256, dtype=dtype, mask=mask_kind):
q, k, v = make_qkv(9, 31, 256, 8, 4, dtype)
if mask_kind == "bool":
mask = mx.random.uniform(shape=(1, 8, 9, 31)) > 0.2
elif mask_kind == "additive":
mask = mx.random.normal((1, 8, 9, 31), dtype)
else:
mask = mask_kind
scale = 256**-0.5
ref = mlx_ref_attn(q, k, v, scale=scale, mask=mask)
out = mx.fast.scaled_dot_product_attention(
q, k, v, scale=scale, mask=mask, force_fused=True
)
tol = 1e-2 if dtype == mx.bfloat16 else 1e-3
self.assertTrue(mx.allclose(ref, out, atol=tol, rtol=tol))

q, k, _ = make_qkv(9, 31, 256, 8, 4, mx.bfloat16)
v_half = mx.random.normal((1, 4, 31, 128), mx.bfloat16)
v = mx.concatenate((v_half, v_half), axis=-1)
out = mx.fast.scaled_dot_product_attention(
q, k, v, scale=256**-0.5, force_fused=True
)
self.assertTrue(mx.array_equal(out[..., :128], out[..., 128:]))

# Vector attention kernel.
for D in (192, 256, 512):
with self.subTest(head_dim=D):
Expand Down