Struct MultiHeadAttnFwdRequest 

pub struct MultiHeadAttnFwdRequest<T: CudnnSupported> {
    pub q: GpuRef<T>,
    pub k: GpuRef<T>,
    pub v: GpuRef<T>,
    pub o: GpuRef<T>,
    pub stats: Option<GpuRef<T>>,
    pub bias: Option<GpuRef<T>>,
    pub layout: TensorLayout,
    pub params: AttentionParams,
    pub reply: Sender<Result<(), GpuError>>,
    pub _ty: PhantomData<T>,
}

MHA forward request.

Fields

q: GpuRef<T>

k: GpuRef<T>

v: GpuRef<T>

o: GpuRef<T>

stats: Option<GpuRef<T>>

Optional saved softmax-stats for backward.

bias: Option<GpuRef<T>>

Optional bias added to attention scores.

layout: TensorLayout

params: AttentionParams

reply: Sender<Result<(), GpuError>>

_ty: PhantomData<T>

Implementations

impl<T: CudnnSupported> MultiHeadAttnFwdRequest<T>

pub fn graph_spec(&self) -> OperationGraphSpec

Trait Implementations

impl<T: CudnnSupported> CudnnDispatch for MultiHeadAttnFwdRequest<T>

fn dtype_name(&self) -> &'static str

fn op_kind(&self) -> &'static str

fn dispatch(self: Box<Self>, _ctx: &CudnnDispatchCtx<'_>)