InfiniOps
Operator Library for Accelerators
Loading...
Searching...
No Matches
infini::ops::KvCachingInfinilm Class Referenceabstract

#include <kv_caching_infinilm.h>

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

Public Member Functions

 KvCachingInfinilm (const Tensor k, const Tensor v, const Tensor past_kv_lengths, Tensor k_cache, Tensor v_cache)
 
virtual void operator() (const Tensor k, const Tensor v, const Tensor past_kv_lengths, Tensor k_cache, Tensor v_cache) const =0
 
- Public Member Functions inherited from infini::ops::Operator< KvCachingInfinilm >
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 k_cache_shape_
 
Tensor::Strides k_cache_strides_
 
Tensor::Shape v_cache_shape_
 
Tensor::Strides v_cache_strides_
 
Tensor::Shape k_shape_
 
Tensor::Strides k_strides_
 
Tensor::Shape v_shape_
 
Tensor::Strides v_strides_
 
Tensor::Shape past_kv_lengths_shape_
 
DataType data_type_
 
DataType past_kv_lengths_type_
 
Tensor::Size batch_size_ {0}
 
Tensor::Size num_kv_heads_ {0}
 
Tensor::Size max_seq_len_ {0}
 
Tensor::Size seq_len_ {0}
 
Tensor::Size hidden_size_ {0}
 
Tensor::Size output_size_ {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< KvCachingInfinilm >
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< KvCachingInfinilm >
static constexpr Device::Type device_type_
 
static constexpr std::size_t implementation_index_
 

Detailed Description

Deprecated:
Migrate to an open-source-aligned operator when available. This interface will be removed in a future release.

Constructor & Destructor Documentation

◆ KvCachingInfinilm()

infini::ops::KvCachingInfinilm::KvCachingInfinilm ( const Tensor  k,
const Tensor  v,
const Tensor  past_kv_lengths,
Tensor  k_cache,
Tensor  v_cache 
)
inline

Member Function Documentation

◆ operator()()

virtual void infini::ops::KvCachingInfinilm::operator() ( const Tensor  k,
const Tensor  v,
const Tensor  past_kv_lengths,
Tensor  k_cache,
Tensor  v_cache 
) const
pure virtual

Member Data Documentation

◆ batch_size_

Tensor::Size infini::ops::KvCachingInfinilm::batch_size_ {0}
protected

◆ data_type_

DataType infini::ops::KvCachingInfinilm::data_type_
protected

◆ device_index_

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

◆ hidden_size_

Tensor::Size infini::ops::KvCachingInfinilm::hidden_size_ {0}
protected

◆ k_cache_shape_

Tensor::Shape infini::ops::KvCachingInfinilm::k_cache_shape_
protected

◆ k_cache_strides_

Tensor::Strides infini::ops::KvCachingInfinilm::k_cache_strides_
protected

◆ k_shape_

Tensor::Shape infini::ops::KvCachingInfinilm::k_shape_
protected

◆ k_strides_

Tensor::Strides infini::ops::KvCachingInfinilm::k_strides_
protected

◆ max_seq_len_

Tensor::Size infini::ops::KvCachingInfinilm::max_seq_len_ {0}
protected

◆ num_kv_heads_

Tensor::Size infini::ops::KvCachingInfinilm::num_kv_heads_ {0}
protected

◆ output_size_

Tensor::Size infini::ops::KvCachingInfinilm::output_size_ {0}
protected

◆ past_kv_lengths_shape_

Tensor::Shape infini::ops::KvCachingInfinilm::past_kv_lengths_shape_
protected

◆ past_kv_lengths_type_

DataType infini::ops::KvCachingInfinilm::past_kv_lengths_type_
protected

◆ seq_len_

Tensor::Size infini::ops::KvCachingInfinilm::seq_len_ {0}
protected

◆ v_cache_shape_

Tensor::Shape infini::ops::KvCachingInfinilm::v_cache_shape_
protected

◆ v_cache_strides_

Tensor::Strides infini::ops::KvCachingInfinilm::v_cache_strides_
protected

◆ v_shape_

Tensor::Shape infini::ops::KvCachingInfinilm::v_shape_
protected

◆ v_strides_

Tensor::Strides infini::ops::KvCachingInfinilm::v_strides_
protected

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