Struct rai_core::primitives::FlashAttention
source · 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§
source§impl FlashAttention
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§
source§impl Clone for FlashAttention
impl Clone for FlashAttention
source§fn clone(&self) -> FlashAttention
fn clone(&self) -> FlashAttention
Returns a copy of the value. Read more
1.0.0 · source§fn clone_from(&mut self, source: &Self)
fn clone_from(&mut self, source: &Self)
Performs copy-assignment from
source
. Read moresource§impl Debug for FlashAttention
impl Debug for FlashAttention
source§impl PartialEq for FlashAttention
impl PartialEq for FlashAttention
source§fn eq(&self, other: &FlashAttention) -> bool
fn eq(&self, other: &FlashAttention) -> bool
This method tests for
self
and other
values to be equal, and is used
by ==
.source§impl Primitive for FlashAttention
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
Auto Trait Implementations§
impl Freeze for FlashAttention
impl !RefUnwindSafe for FlashAttention
impl !Send for FlashAttention
impl !Sync for FlashAttention
impl Unpin for FlashAttention
impl !UnwindSafe for FlashAttention
Blanket Implementations§
source§impl<T, M> ApplyModule<M> for Twhere
M: Module<Input = T>,
impl<T, M> ApplyModule<M> for Twhere
M: Module<Input = T>,
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
Mutably borrows from an owned value. Read more