Enum AttentionKernel 

#[non_exhaustive]
pub enum AttentionKernel {
    Vector {
        passes: u8,
        blocks: i32,
        dispatches: i32,
    },
    Tiled,
    Composed,
    ComposedBlocks {
        rows: i32,
    },
}

The kernel family an attention call runs (see plan_attention).

Variants (Non-exhaustive)

Non-exhaustive enums could have additional variants added in future. Therefore, when matching against variants of non-exhaustive enums, an extra wildcard arm must be added to account for any future variants.

Vector

At most 8 queries: MLX’s vector kernels. passes is 1 or 2; blocks the key partitions of the first of two passes (1 for one pass; the largest over dispatches); dispatches is more than 1 when query rows need different one-query configurations (the serial rule).

Fields

passes: u8

Passes over the keys (1 or 2).

blocks: i32

Key partitions of the first pass.

dispatches: i32

Dispatches (groups of query rows with one configuration).

Tiled

MLX’s tiled (steel) attention kernel: no scores are materialised.

Composed

The composition: softmax(scale · q kᵀ + mask) v from matrix products.

ComposedBlocks

The composition per block of rows queries, one block after another.

Fields

rows: i32

Queries per block (the last block may have fewer).

Trait Implementations

impl Clone for AttentionKernel

fn clone(&self) -> Self

Returns a duplicate of the value.

fn clone_from(&mut self, source: &Self)

Performs copy-assignment from source.

impl Copy for AttentionKernel

impl Debug for AttentionKernel

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.

impl Eq for AttentionKernel

impl PartialEq for AttentionKernel

fn eq(&self, other: &Self) -> bool

Equality operator ==.

fn ne(&self, other: &Rhs) -> bool

Inequality operator !=.

impl StructuralPartialEq for AttentionKernel