#[non_exhaustive]pub struct AttentionOptions<'a> {
pub mask: AttentionMask<'a>,
pub sinks: Option<&'a Array>,
pub window: Option<i32>,
pub max_score_bytes: Option<usize>,
}Expand description
The options of scaled_dot_product_attention and plan_attention: the mask,
attention sinks, a sliding window and a cap on score memory. Build them with
AttentionOptions::causal or Default and the chained setters:
use nuvec_core::fast::AttentionOptions;
let opts = AttentionOptions::causal().window(4096).max_score_bytes(512 << 20);
assert_eq!(opts.window, Some(4096));This is the one options type of attention: new features (window and head-size-512 kernels, backward routing) extend it, never a second entry point.
Fields (Non-exhaustive)§
This struct is marked as non-exhaustive
Struct { .. } syntax; cannot be matched against without a wildcard ..; and struct update syntax will not work.mask: AttentionMask<'a>The mask: none (the default), causal, or an array.
sinks: Option<&'a Array>Attention sinks ([Hq]): one logit per query head that takes part in the softmax
but attends to nothing.
window: Option<i32>A sliding window: query i attends only to the window keys up to and including
its position (as mlx-lm’s create_causal_mask(..., window_size)). Needs the causal
mask.
max_score_bytes: Option<usize>A cap, in bytes, on every score-sized intermediate (scores, masked scores, the
float32 sink-softmax intermediates, materialised masks); see
scaled_dot_product_attention.
Implementations§
Source§impl<'a> AttentionOptions<'a>
impl<'a> AttentionOptions<'a>
Sourcepub fn causal() -> Self
pub fn causal() -> Self
Causal masking: query i attends to keys 0..=i + Lk - Lq (the queries are the last
Lq positions).
Sourcepub fn mask(self, mask: &'a Array) -> Self
pub fn mask(self, mask: &'a Array) -> Self
A boolean (true attends) or additive mask, broadcastable to [B, Hq, Lq, Lk].
Sourcepub fn sinks(self, sinks: impl Into<Option<&'a Array>>) -> Self
pub fn sinks(self, sinks: impl Into<Option<&'a Array>>) -> Self
Attention sinks, [Hq] (None for none).
Sourcepub fn window(self, keys: i32) -> Self
pub fn window(self, keys: i32) -> Self
A sliding window of keys keys (needs AttentionOptions::causal): query i
attends to keys i + Lk − Lq − keys < j ≤ i + Lk − Lq. The vector and tiled kernels
skip the keys before the window, so a full cache needs no slicing; a cache that keeps
only the last W − 1 + L positions (nuvec-lm’s sliding-window KvCache) passes
just those, with the same result:
use nuvec_core::{DType, fast::{self, AttentionOptions}, random};
let (w, l) = (4, 2);
let q = random::normal([1, 2, l, 8], DType::F32, 0.0, 1.0, None);
let k = random::normal([1, 2, 10, 8], DType::F32, 0.0, 1.0, None);
let opts = AttentionOptions::causal().window(w);
let full = fast::scaled_dot_product_attention(&q, &k, &k, 0.35, &opts);
let tail = k.i((.., .., 10 - (w - 1 + l)..));
let sliced = fast::scaled_dot_product_attention(&q, &tail, &tail, 0.35, &opts);
let diff = full.subtract(&sliced).abs().max(..).item::<f32>().unwrap();
assert!(diff < 1e-6);Sourcepub fn max_score_bytes(self, bytes: usize) -> Self
pub fn max_score_bytes(self, bytes: usize) -> Self
Caps every score-sized intermediate at bytes (positive).
Trait Implementations§
Source§impl<'a> Clone for AttentionOptions<'a>
impl<'a> Clone for AttentionOptions<'a>
impl<'a> Copy for AttentionOptions<'a>
Source§impl<'a> Debug for AttentionOptions<'a>
impl<'a> Debug for AttentionOptions<'a>
Source§impl<'a> Default for AttentionOptions<'a>
impl<'a> Default for AttentionOptions<'a>
Source§impl<'a> From<AttentionMask<'a>> for AttentionOptions<'a>
impl<'a> From<AttentionMask<'a>> for AttentionOptions<'a>
Source§fn from(mask: AttentionMask<'a>) -> Self
fn from(mask: AttentionMask<'a>) -> Self
Options with this mask and nothing else.
Auto Trait Implementations§
impl<'a> !RefUnwindSafe for AttentionOptions<'a>
impl<'a> !UnwindSafe for AttentionOptions<'a>
impl<'a> Freeze for AttentionOptions<'a>
impl<'a> Send for AttentionOptions<'a>
impl<'a> Sync for AttentionOptions<'a>
impl<'a> Unpin for AttentionOptions<'a>
impl<'a> UnsafeUnpin for AttentionOptions<'a>
Blanket Implementations§
impl<T> AutoreleaseSafe for Twhere
T: ?Sized,
Source§impl<T> BorrowMut<T> for Twhere
T: ?Sized,
impl<T> BorrowMut<T> for Twhere
T: ?Sized,
Source§fn borrow_mut(&mut self) -> &mut T
fn borrow_mut(&mut self) -> &mut T
impl<ST, DT> CastableFrom<ST, Initialized, Initialized> for DT
impl<ST, DT> CastableFrom<ST, Uninit, Uninit> for DT
Source§impl<T> CloneToUninit for Twhere
T: Clone,
impl<T> CloneToUninit for Twhere
T: Clone,
Source§impl<T> IntoEither for T
impl<T> IntoEither for T
Source§fn into_either(self, into_left: bool) -> Either<Self, Self> ⓘ
fn into_either(self, into_left: bool) -> Either<Self, Self> ⓘ
self into a Left variant of Either<Self, Self>
if into_left is true.
Converts self into a Right variant of Either<Self, Self>
otherwise. Read moreSource§fn into_either_with<F>(self, into_left: F) -> Either<Self, Self> ⓘ
fn into_either_with<F>(self, into_left: F) -> Either<Self, Self> ⓘ
self into a Left variant of Either<Self, Self>
if into_left(&self) returns true.
Converts self into a Right variant of Either<Self, Self>
otherwise. Read more