infini::ops::RotaryEmbedding Class Reference

#include <rotary_embedding.h>

Inheritance diagram for infini::ops::RotaryEmbedding:

Public Member Functions
	RotaryEmbedding (const Tensor positions, Tensor query, std::optional< Tensor > key, const Tensor cos_sin_cache, int64_t head_size, bool is_neox, int64_t rope_dim_offset=0, bool inverse=false)
virtual void	operator() (const Tensor positions, Tensor query, std::optional< Tensor > key, const Tensor cos_sin_cache, int64_t head_size, bool is_neox, int64_t rope_dim_offset=0, bool inverse=false) const =0
Public Member Functions inherited from infini::ops::Operator< RotaryEmbedding >
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)

Protected Attributes
Tensor::Shape	positions_shape_
Tensor::Shape	query_shape_
Tensor::Shape	key_shape_
Tensor::Shape	cos_sin_cache_shape_
Tensor::Strides	positions_strides_
Tensor::Strides	query_strides_
Tensor::Strides	key_strides_
Tensor::Strides	cos_sin_cache_strides_
DataType	positions_type_
DataType	query_type_
DataType	cos_sin_cache_type_
Tensor::Size	num_tokens_ {0}
Tensor::Size	positions_ndim_ {0}
int64_t	head_size_ {0}
int64_t	rot_dim_ {0}
int64_t	rope_dim_offset_ {0}
bool	is_neox_ {false}
bool	inverse_ {false}
Tensor::Size	query_hidden_size_ {0}
Tensor::Size	key_hidden_size_ {0}
Tensor::Size	num_heads_ {0}
Tensor::Size	num_kv_heads_ {0}
int64_t	query_token_stride_ {0}
int64_t	key_token_stride_ {0}
int64_t	query_head_stride_ {0}
int64_t	key_head_stride_ {0}
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 Public Member Functions inherited from infini::ops::Operator< RotaryEmbedding >
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)
Static Protected Attributes inherited from infini::ops::Operator< RotaryEmbedding >
static constexpr Device::Type	device_type_
static constexpr std::size_t	implementation_index_

Constructor & Destructor Documentation

RotaryEmbedding()

infini::ops::RotaryEmbedding::RotaryEmbedding ( const Tensor  positions, Tensor  query, std::optional< Tensor >  key, const Tensor  cos_sin_cache, int64_t  head_size, bool  is_neox, int64_t  rope_dim_offset = 0, bool  inverse = false  )

inline

Member Function Documentation

operator()()

virtual void infini::ops::RotaryEmbedding::operator() ( const Tensor  positions, Tensor  query, std::optional< Tensor >  key, const Tensor  cos_sin_cache, int64_t  head_size, bool  is_neox, int64_t  rope_dim_offset = 0, bool  inverse = false  ) const

pure virtual

Member Data Documentation

cos_sin_cache_shape_

Tensor::Shape infini::ops::RotaryEmbedding::cos_sin_cache_shape_

protected

cos_sin_cache_strides_

Tensor::Strides infini::ops::RotaryEmbedding::cos_sin_cache_strides_

protected

cos_sin_cache_type_

DataType infini::ops::RotaryEmbedding::cos_sin_cache_type_

protected

device_index_

int infini::ops::RotaryEmbedding::device_index_ {0}

protected

head_size_

int64_t infini::ops::RotaryEmbedding::head_size_ {0}

protected

inverse_

bool infini::ops::RotaryEmbedding::inverse_ {false}

protected

is_neox_

bool infini::ops::RotaryEmbedding::is_neox_ {false}

protected

key_head_stride_

int64_t infini::ops::RotaryEmbedding::key_head_stride_ {0}

protected

key_hidden_size_

Tensor::Size infini::ops::RotaryEmbedding::key_hidden_size_ {0}

protected

key_shape_

Tensor::Shape infini::ops::RotaryEmbedding::key_shape_

protected

key_strides_

Tensor::Strides infini::ops::RotaryEmbedding::key_strides_

protected

key_token_stride_

int64_t infini::ops::RotaryEmbedding::key_token_stride_ {0}

protected

num_heads_

Tensor::Size infini::ops::RotaryEmbedding::num_heads_ {0}

protected

num_kv_heads_

Tensor::Size infini::ops::RotaryEmbedding::num_kv_heads_ {0}

protected

num_tokens_

Tensor::Size infini::ops::RotaryEmbedding::num_tokens_ {0}

protected

positions_ndim_

Tensor::Size infini::ops::RotaryEmbedding::positions_ndim_ {0}

protected

positions_shape_

Tensor::Shape infini::ops::RotaryEmbedding::positions_shape_

protected

positions_strides_

Tensor::Strides infini::ops::RotaryEmbedding::positions_strides_

protected

positions_type_

DataType infini::ops::RotaryEmbedding::positions_type_

protected

query_head_stride_

int64_t infini::ops::RotaryEmbedding::query_head_stride_ {0}

protected

query_hidden_size_

Tensor::Size infini::ops::RotaryEmbedding::query_hidden_size_ {0}

protected

query_shape_

Tensor::Shape infini::ops::RotaryEmbedding::query_shape_

protected

query_strides_

Tensor::Strides infini::ops::RotaryEmbedding::query_strides_

protected

query_token_stride_

int64_t infini::ops::RotaryEmbedding::query_token_stride_ {0}

protected

query_type_

DataType infini::ops::RotaryEmbedding::query_type_

protected

rope_dim_offset_

int64_t infini::ops::RotaryEmbedding::rope_dim_offset_ {0}

protected

rot_dim_

int64_t infini::ops::RotaryEmbedding::rot_dim_ {0}

protected

---

The documentation for this class was generated from the following file:

• rotary_embedding.h