infini::ops::ReshapeAndCache Class Reference

#include <reshape_and_cache.h>

Inheritance diagram for infini::ops::ReshapeAndCache:

Public Member Functions
	ReshapeAndCache (const Tensor key, const Tensor value, Tensor key_cache, Tensor value_cache, const Tensor slot_mapping, const Tensor k_scale, const Tensor v_scale, const std::string kv_cache_dtype)
virtual void	operator() (const Tensor key, const Tensor value, Tensor key_cache, Tensor value_cache, const Tensor slot_mapping, const Tensor k_scale, const Tensor v_scale, const std::string kv_cache_dtype) const =0
Public Member Functions inherited from infini::ops::Operator< ReshapeAndCache >
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	key_shape_
Tensor::Shape	value_shape_
Tensor::Shape	key_cache_shape_
Tensor::Shape	value_cache_shape_
Tensor::Shape	slot_mapping_shape_
Tensor::Shape	k_scale_shape_
Tensor::Shape	v_scale_shape_
Tensor::Strides	key_strides_
Tensor::Strides	value_strides_
Tensor::Strides	key_cache_strides_
Tensor::Strides	value_cache_strides_
Tensor::Strides	slot_mapping_strides_
Tensor::Strides	k_scale_strides_
Tensor::Strides	v_scale_strides_
DataType	key_type_
DataType	key_cache_type_
DataType	value_cache_type_
DataType	k_scale_type_
DataType	v_scale_type_
std::string	kv_cache_dtype_
Tensor::Size	num_tokens_ {0}
Tensor::Size	num_heads_ {0}
Tensor::Size	head_size_ {0}
Tensor::Size	block_size_ {0}
Tensor::Size	x_ {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< ReshapeAndCache >
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< ReshapeAndCache >
static constexpr Device::Type	device_type_
static constexpr std::size_t	implementation_index_

Constructor & Destructor Documentation

ReshapeAndCache()

infini::ops::ReshapeAndCache::ReshapeAndCache ( const Tensor  key, const Tensor  value, Tensor  key_cache, Tensor  value_cache, const Tensor  slot_mapping, const Tensor  k_scale, const Tensor  v_scale, const std::string  kv_cache_dtype  )

inline

Member Function Documentation

operator()()

virtual void infini::ops::ReshapeAndCache::operator() ( const Tensor  key, const Tensor  value, Tensor  key_cache, Tensor  value_cache, const Tensor  slot_mapping, const Tensor  k_scale, const Tensor  v_scale, const std::string  kv_cache_dtype  ) const

pure virtual

Member Data Documentation

block_size_

Tensor::Size infini::ops::ReshapeAndCache::block_size_ {0}

protected

device_index_

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

protected

head_size_

Tensor::Size infini::ops::ReshapeAndCache::head_size_ {0}

protected

k_scale_shape_

Tensor::Shape infini::ops::ReshapeAndCache::k_scale_shape_

protected

k_scale_strides_

Tensor::Strides infini::ops::ReshapeAndCache::k_scale_strides_

protected

k_scale_type_

DataType infini::ops::ReshapeAndCache::k_scale_type_

protected

key_cache_shape_

Tensor::Shape infini::ops::ReshapeAndCache::key_cache_shape_

protected

key_cache_strides_

Tensor::Strides infini::ops::ReshapeAndCache::key_cache_strides_

protected

key_cache_type_

DataType infini::ops::ReshapeAndCache::key_cache_type_

protected

key_shape_

Tensor::Shape infini::ops::ReshapeAndCache::key_shape_

protected

key_strides_

Tensor::Strides infini::ops::ReshapeAndCache::key_strides_

protected

key_type_

DataType infini::ops::ReshapeAndCache::key_type_

protected

kv_cache_dtype_

std::string infini::ops::ReshapeAndCache::kv_cache_dtype_

protected

num_heads_

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

protected

num_tokens_

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

protected

slot_mapping_shape_

Tensor::Shape infini::ops::ReshapeAndCache::slot_mapping_shape_

protected

slot_mapping_strides_

Tensor::Strides infini::ops::ReshapeAndCache::slot_mapping_strides_

protected

v_scale_shape_

Tensor::Shape infini::ops::ReshapeAndCache::v_scale_shape_

protected

v_scale_strides_

Tensor::Strides infini::ops::ReshapeAndCache::v_scale_strides_

protected

v_scale_type_

DataType infini::ops::ReshapeAndCache::v_scale_type_

protected

value_cache_shape_

Tensor::Shape infini::ops::ReshapeAndCache::value_cache_shape_

protected

value_cache_strides_

Tensor::Strides infini::ops::ReshapeAndCache::value_cache_strides_

protected

value_cache_type_

DataType infini::ops::ReshapeAndCache::value_cache_type_

protected

value_shape_

Tensor::Shape infini::ops::ReshapeAndCache::value_shape_

protected

value_strides_

Tensor::Strides infini::ops::ReshapeAndCache::value_strides_

protected

x_

Tensor::Size infini::ops::ReshapeAndCache::x_ {0}

protected

---

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

• reshape_and_cache.h