zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-10-19 00:04:41 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Support fused masking in Attention (#1924)

Browse Source 
* Update API to allow mask='causal' in fast::sdpa

* Add fallback

* Update steel::AttnParams

* Fix typo

* WIP, basic causal

* Update tests

* Update benchmarking

* Update masking loop limits

* Add bool masking and update tests

* Update additive mask

* Update benchmarks

* Update benchmarks

* Update tests

* Update for bfloat error

* Update early exit

* Add random seed to tests
This commit is contained in:
Jagrit Digani
2025-03-20 11:01:32 -07:00
committed by GitHub
parent 3c164fca8c
commit 9adcd1a650
11 changed files with 504 additions and 148 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
python/src/fast.cpp
View File

@@ -134,7 +134,7 @@ void init_fast(nb::module_& parent_module) {
"memory_efficient_threshold"_a = nb::none(),
"stream"_a = nb::none(),
nb::sig(
"def scaled_dot_product_attention(q: array, k: array, v: array, *, scale: float, mask: Optional[array] = None, stream: Union[None, Stream, Device] = None) -> array"),
"def scaled_dot_product_attention(q: array, k: array, v: array, *, scale: float, mask: Union[None, str, array] = None, stream: Union[None, Stream, Device] = None) -> array"),
R"pbdoc(
A fast implementation of multi-head attention: ``O = softmax(Q @ K.T, dim=-1) @ V``.
@@ -164,11 +164,11 @@ void init_fast(nb::module_& parent_module) {
k (array): Keys with shape ``[B, N_kv, T_kv, D]``.
v (array): Values with shape ``[B, N_kv, T_kv, D]``.
scale (float): Scale for queries (typically ``1.0 / sqrt(q.shape(-1)``)
mask (array, optional): A boolean or additive mask to apply to the
query-key scores. The mask can have at most 4 dimensions and must
be broadcast-compatible with the shape ``[B, N, T_q, T_kv]``. If an
additive mask is given its type must promote to the promoted
type of ``q``, ``k``, and ``v``.
mask (Union[None, str, array], optional): A causal, boolean or additive
mask to apply to the query-key scores. The mask can have at most 4
dimensions and must be broadcast-compatible with the shape
``[B, N, T_q, T_kv]``. If an additive mask is given its type must
promote to the promoted type of ``q``, ``k``, and ``v``.
Returns:
array: The output array.
)pbdoc");