InfiniOps
Operator Library for Accelerators
Loading...
Searching...
No Matches
Public Member Functions | Static Public Member Functions | Protected Attributes | List of all members
infini::ops::ScaledDotProductAttention Class Referenceabstract

#include <scaled_dot_product_attention.h>

Inheritance diagram for infini::ops::ScaledDotProductAttention:
infini::ops::Operator< ScaledDotProductAttention > infini::ops::OperatorBase infini::ops::OperatorBase

Public Member Functions

 ScaledDotProductAttention (const Tensor query, const Tensor key, const Tensor value, const std::optional< Tensor > attn_mask, double dropout_p, bool is_causal, const std::optional< double > scale, bool enable_gqa, Tensor out)
 
 ScaledDotProductAttention (const Tensor query, const Tensor key, const Tensor value, Tensor out)
 
virtual void operator() (const Tensor query, const Tensor key, const Tensor value, const std::optional< Tensor > attn_mask, double dropout_p, bool is_causal, const std::optional< double > scale, bool enable_gqa, Tensor out) const =0
 
void operator() (const Tensor query, const Tensor key, const Tensor value, Tensor out) const
 
- Public Member Functions inherited from infini::ops::Operator< ScaledDotProductAttention >
void operator() (const Handle &handle, const Args &... args)
 
void operator() (const Args &... args) const
 
void operator() (const Handle &handle, const Args &... args)
 
void operator() (const Args &... args) const
 
- Public Member Functions inherited from infini::ops::OperatorBase
virtual ~OperatorBase ()=default
 
virtual std::size_t workspace_size_in_bytes () const
 
void set_handle (const Handle &handle)
 
void set_config (const Config &config)
 
void set_stream (void *stream)
 
void set_workspace (void *workspace)
 
void set_workspace_size_in_bytes (std::size_t workspace_size_in_bytes)
 
virtual ~OperatorBase ()=default
 
virtual std::size_t workspace_size_in_bytes () const
 
void set_handle (const Handle &handle)
 
void set_config (const Config &config)
 
void set_stream (void *stream)
 
void set_workspace (void *workspace)
 
void set_workspace_size_in_bytes (std::size_t workspace_size_in_bytes)
 

Static Public Member Functions

template<typename TensorLike >
static auto MakeReturnValue (const TensorLike &query, const TensorLike &key, const TensorLike &value, const std::optional< TensorLike > attn_mask=std::nullopt, double dropout_p=0.0, bool is_causal=false, const std::optional< double > scale=std::nullopt, bool enable_gqa=false)
 
- Static Public Member Functions inherited from infini::ops::Operator< ScaledDotProductAttention >
static void clear_cache ()
 
static void clear_cache ()
 
static std::unique_ptr< Operator > Make (const Config &config, const Tensor tensor, Args &&... args)
 
static std::unique_ptr< Operator > Make (const Tensor tensor, Args &&... args)
 
static std::unique_ptr< Operator > Make (const Config &config, const std::vector< Tensor > tensors, Args &&... args)
 
static std::unique_ptr< Operator > Make (const std::vector< Tensor > tensors, Args &&... args)
 
static std::unique_ptr< Operator > Make (const Config &config, const Tensor tensor, Args &&... args)
 
static std::unique_ptr< Operator > Make (const Tensor tensor, Args &&... args)
 
static std::unique_ptr< Operator > Make (const Config &config, const std::vector< Tensor > tensors, Args &&... args)
 
static std::unique_ptr< Operator > Make (const std::vector< Tensor > tensors, Args &&... args)
 
static void Call (const Handle &handle, const Config &config, const Args &... args)
 
static void Call (const Tensor tensor, const Args &... args)
 
static auto Call (const TensorLike &tensor, const Args &... args)
 
static void Call (const Handle &handle, const Config &config, const Args &... args)
 
static void Call (const Tensor tensor, const Args &... args)
 
static auto Call (const TensorLike &tensor, const Args &... args)
 
static std::vector< std::size_t > active_implementation_indices (Device::Type dev_type)
 
static std::vector< std::size_t > active_implementation_indices (Device::Type dev_type)
 

Protected Attributes

Tensor::Shape query_shape_
 
Tensor::Shape key_shape_
 
Tensor::Shape value_shape_
 
Tensor::Shape attn_mask_shape_
 
Tensor::Shape out_shape_
 
Tensor::Strides query_strides_
 
Tensor::Strides key_strides_
 
Tensor::Strides value_strides_
 
Tensor::Strides attn_mask_strides_
 
Tensor::Strides out_strides_
 
DataType query_type_
 
DataType attn_mask_type_
 
double dropout_p_ {0.0}
 
bool is_causal_ {false}
 
std::optional< double > scale_
 
bool enable_gqa_ {false}
 
int device_index_ {0}
 
- Protected Attributes inherited from infini::ops::OperatorBase
std::unique_ptr< Handle > handle_ptr_
 
std::unique_ptr< Config > config_ptr_
 
void * stream_ {nullptr}
 
void * workspace_ {nullptr}
 
std::size_t workspace_size_in_bytes_ {0}
 

Additional Inherited Members

- Static Protected Attributes inherited from infini::ops::Operator< ScaledDotProductAttention >
static constexpr Device::Type device_type_
 
static constexpr std::size_t implementation_index_
 

Constructor & Destructor Documentation

◆ ScaledDotProductAttention() [1/2]

infini::ops::ScaledDotProductAttention::ScaledDotProductAttention ( const Tensor  query,
const Tensor  key,
const Tensor  value,
const std::optional< Tensor >  attn_mask,
double  dropout_p,
bool  is_causal,
const std::optional< double >  scale,
bool  enable_gqa,
Tensor  out 
)
inline

◆ ScaledDotProductAttention() [2/2]

infini::ops::ScaledDotProductAttention::ScaledDotProductAttention ( const Tensor  query,
const Tensor  key,
const Tensor  value,
Tensor  out 
)
inline

Member Function Documentation

◆ MakeReturnValue()

template<typename TensorLike >
static auto infini::ops::ScaledDotProductAttention::MakeReturnValue ( const TensorLike &  query,
const TensorLike &  key,
const TensorLike &  value,
const std::optional< TensorLike >  attn_mask = std::nullopt,
double  dropout_p = 0.0,
bool  is_causal = false,
const std::optional< double >  scale = std::nullopt,
bool  enable_gqa = false 
)
inlinestatic

◆ operator()() [1/2]

virtual void infini::ops::ScaledDotProductAttention::operator() ( const Tensor  query,
const Tensor  key,
const Tensor  value,
const std::optional< Tensor >  attn_mask,
double  dropout_p,
bool  is_causal,
const std::optional< double >  scale,
bool  enable_gqa,
Tensor  out 
) const
pure virtual

◆ operator()() [2/2]

void infini::ops::ScaledDotProductAttention::operator() ( const Tensor  query,
const Tensor  key,
const Tensor  value,
Tensor  out 
) const
inline

Member Data Documentation

◆ attn_mask_shape_

Tensor::Shape infini::ops::ScaledDotProductAttention::attn_mask_shape_
protected

◆ attn_mask_strides_

Tensor::Strides infini::ops::ScaledDotProductAttention::attn_mask_strides_
protected

◆ attn_mask_type_

DataType infini::ops::ScaledDotProductAttention::attn_mask_type_
protected

◆ device_index_

int infini::ops::ScaledDotProductAttention::device_index_ {0}
protected

◆ dropout_p_

double infini::ops::ScaledDotProductAttention::dropout_p_ {0.0}
protected

◆ enable_gqa_

bool infini::ops::ScaledDotProductAttention::enable_gqa_ {false}
protected

◆ is_causal_

bool infini::ops::ScaledDotProductAttention::is_causal_ {false}
protected

◆ key_shape_

Tensor::Shape infini::ops::ScaledDotProductAttention::key_shape_
protected

◆ key_strides_

Tensor::Strides infini::ops::ScaledDotProductAttention::key_strides_
protected

◆ out_shape_

Tensor::Shape infini::ops::ScaledDotProductAttention::out_shape_
protected

◆ out_strides_

Tensor::Strides infini::ops::ScaledDotProductAttention::out_strides_
protected

◆ query_shape_

Tensor::Shape infini::ops::ScaledDotProductAttention::query_shape_
protected

◆ query_strides_

Tensor::Strides infini::ops::ScaledDotProductAttention::query_strides_
protected

◆ query_type_

DataType infini::ops::ScaledDotProductAttention::query_type_
protected

◆ scale_

std::optional<double> infini::ops::ScaledDotProductAttention::scale_
protected

◆ value_shape_

Tensor::Shape infini::ops::ScaledDotProductAttention::value_shape_
protected

◆ value_strides_

Tensor::Strides infini::ops::ScaledDotProductAttention::value_strides_
protected
The documentation for this class was generated from the following file: