Struct rai_core::primitives::FlashAttention

pub struct FlashAttention {
    pub softmax_scale: f32,
    pub window_size_left: Option<usize>,
    pub window_size_right: Option<usize>,
    pub alibi_slopes: Option<Tensor>,
}

Fields

softmax_scale: f32

window_size_left: Option<usize>

window_size_right: Option<usize>

alibi_slopes: Option<Tensor>

Implementations

impl FlashAttention

pub fn new( softmax_scale: f32, window_size_left: Option<usize>, window_size_right: Option<usize>, alibi_slopes: Option<Tensor> ) -> Self

pub fn softmax_scale(&self) -> f32

pub fn window_size_left(&self) -> Option<usize>

pub fn window_size_right(&self) -> Option<usize>

pub fn alibi_slopes(&self) -> Option<&Tensor>

Trait Implementations

impl Clone for FlashAttention

fn clone(&self) -> FlashAttention

Returns a copy of the value.

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

Performs copy-assignment from source.

impl Debug for FlashAttention

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

Formats the value using the given formatter.

impl PartialEq for FlashAttention

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

This method tests for self and other values to be equal, and is used by ==.

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

This method tests for !=. The default implementation is almost always sufficient, and should not be overridden without very good reason.

impl Primitive for FlashAttention

fn clone_boxed(&self) -> Box<dyn Primitive>

fn as_any(&self) -> &dyn Any

fn dot_label(&self) -> String

fn jvp( &self, _output: &Tensor, _primals: &[Tensor], tangents: &[Tensor] ) -> Tensor

fn vjp( &self, _output: &Tensor, _primals: &[Tensor], cotangent: &Tensor ) -> Vec<Tensor>

impl StructuralPartialEq for FlashAttention