Use lower-right causal mask alignment consistently (#2967)

Co-authored-by: Awni Hannun <awni.hannun@gmail.com>
2026-01-29 03:15:14 +02:00
parent c86a9bced1
commit 0c6a895ed7
5 changed files with 29 additions and 21 deletions
@@ -272,6 +272,8 @@ void init_fast(nb::module_& parent_module) {
               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``.
+               The ``"causal"`` mask uses lower-right alignment where the
+               last query aligns with the last key.
            sinks (array, optional): An optional array of attention sinks.
               Default: ``None``.