infini::ops::FlashAttnWithKvcache Class Reference

#include <flash_attn_with_kvcache.h>

Inheritance diagram for infini::ops::FlashAttnWithKvcache:

Public Member Functions
	FlashAttnWithKvcache (const Tensor q, Tensor k_cache, Tensor v_cache, Tensor out)
	FlashAttnWithKvcache (const Tensor q, Tensor k_cache, Tensor v_cache, const std::optional< Tensor > k, const std::optional< Tensor > v, const std::optional< Tensor > rotary_cos, const std::optional< Tensor > rotary_sin, const int64_t cache_seqlens, const std::optional< Tensor > cache_batch_idx, const std::optional< Tensor > cache_leftpad, const std::optional< Tensor > block_table, const std::optional< Tensor > alibi_slopes, const std::optional< double > softmax_scale, const bool causal, const std::vector< int64_t > window_size, const double softcap, const bool rotary_interleaved, const int64_t num_splits, const bool return_softmax_lse, Tensor out, std::optional< Tensor > softmax_lse)
	FlashAttnWithKvcache (const Tensor q, Tensor k_cache, Tensor v_cache, const std::optional< Tensor > k, const std::optional< Tensor > v, const std::optional< Tensor > rotary_cos, const std::optional< Tensor > rotary_sin, const std::optional< Tensor > cache_seqlens, const std::optional< Tensor > cache_batch_idx, const std::optional< Tensor > cache_leftpad, const std::optional< Tensor > block_table, const std::optional< Tensor > alibi_slopes, const std::optional< double > softmax_scale, const bool causal, const std::vector< int64_t > window_size, const double softcap, const bool rotary_interleaved, const int64_t num_splits, const bool return_softmax_lse, Tensor out, std::optional< Tensor > softmax_lse)
void	operator() (const Tensor q, Tensor k_cache, Tensor v_cache, Tensor out) const
virtual void	operator() (const Tensor q, Tensor k_cache, Tensor v_cache, const std::optional< Tensor > k, const std::optional< Tensor > v, const std::optional< Tensor > rotary_cos, const std::optional< Tensor > rotary_sin, const int64_t cache_seqlens, const std::optional< Tensor > cache_batch_idx, const std::optional< Tensor > cache_leftpad, const std::optional< Tensor > block_table, const std::optional< Tensor > alibi_slopes, const std::optional< double > softmax_scale, const bool causal, const std::vector< int64_t > window_size, const double softcap, const bool rotary_interleaved, const int64_t num_splits, const bool return_softmax_lse, Tensor out, std::optional< Tensor > softmax_lse) const =0
virtual void	operator() (const Tensor q, Tensor k_cache, Tensor v_cache, const std::optional< Tensor > k, const std::optional< Tensor > v, const std::optional< Tensor > rotary_cos, const std::optional< Tensor > rotary_sin, const std::optional< Tensor > cache_seqlens, const std::optional< Tensor > cache_batch_idx, const std::optional< Tensor > cache_leftpad, const std::optional< Tensor > block_table, const std::optional< Tensor > alibi_slopes, const std::optional< double > softmax_scale, const bool causal, const std::vector< int64_t > window_size, const double softcap, const bool rotary_interleaved, const int64_t num_splits, const bool return_softmax_lse, Tensor out, std::optional< Tensor > softmax_lse) const =0
Public Member Functions inherited from infini::ops::Operator< FlashAttnWithKvcache >
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 Member Functions
void	ValidateIndexVector (const std::optional< Tensor > &tensor) const

Protected Attributes
Tensor::Shape	q_shape_
Tensor::Shape	k_cache_shape_
Tensor::Shape	v_cache_shape_
Tensor::Shape	k_shape_
Tensor::Shape	v_shape_
Tensor::Shape	rotary_cos_shape_
Tensor::Shape	rotary_sin_shape_
Tensor::Shape	cache_seqlens_shape_
Tensor::Shape	cache_batch_idx_shape_
Tensor::Shape	cache_leftpad_shape_
Tensor::Shape	block_table_shape_
Tensor::Shape	alibi_slopes_shape_
Tensor::Shape	out_shape_
Tensor::Shape	softmax_lse_shape_
Tensor::Strides	q_strides_
Tensor::Strides	k_cache_strides_
Tensor::Strides	v_cache_strides_
Tensor::Strides	k_strides_
Tensor::Strides	v_strides_
Tensor::Strides	rotary_cos_strides_
Tensor::Strides	rotary_sin_strides_
Tensor::Strides	cache_seqlens_strides_
Tensor::Strides	cache_batch_idx_strides_
Tensor::Strides	cache_leftpad_strides_
Tensor::Strides	block_table_strides_
Tensor::Strides	alibi_slopes_strides_
Tensor::Strides	out_strides_
Tensor::Strides	softmax_lse_strides_
DataType	q_dtype_
DataType	k_cache_dtype_
DataType	v_cache_dtype_
DataType	k_dtype_
DataType	v_dtype_
DataType	rotary_cos_dtype_
DataType	rotary_sin_dtype_
DataType	cache_seqlens_dtype_
DataType	cache_batch_idx_dtype_
DataType	cache_leftpad_dtype_
DataType	block_table_dtype_
DataType	alibi_slopes_dtype_
DataType	out_dtype_
DataType	softmax_lse_dtype_
bool	has_k_ {false}
bool	has_v_ {false}
bool	has_rotary_cos_ {false}
bool	has_rotary_sin_ {false}
bool	has_cache_batch_idx_ {false}
bool	has_block_table_ {false}
bool	has_alibi_slopes_ {false}
bool	has_softmax_lse_ {false}
Tensor::Size	batch_size_ {0}
Tensor::Size	head_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< FlashAttnWithKvcache >
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< FlashAttnWithKvcache >
static constexpr Device::Type	device_type_
static constexpr std::size_t	implementation_index_

Constructor & Destructor Documentation

FlashAttnWithKvcache() [1/3]

infini::ops::FlashAttnWithKvcache::FlashAttnWithKvcache ( const Tensor  q, Tensor  k_cache, Tensor  v_cache, Tensor  out  )

inline

FlashAttnWithKvcache() [2/3]

infini::ops::FlashAttnWithKvcache::FlashAttnWithKvcache ( const Tensor  q, Tensor  k_cache, Tensor  v_cache, const std::optional< Tensor >  k, const std::optional< Tensor >  v, const std::optional< Tensor >  rotary_cos, const std::optional< Tensor >  rotary_sin, const int64_t  cache_seqlens, const std::optional< Tensor >  cache_batch_idx, const std::optional< Tensor >  cache_leftpad, const std::optional< Tensor >  block_table, const std::optional< Tensor >  alibi_slopes, const std::optional< double >  softmax_scale, const bool  causal, const std::vector< int64_t >  window_size, const double  softcap, const bool  rotary_interleaved, const int64_t  num_splits, const bool  return_softmax_lse, Tensor  out, std::optional< Tensor >  softmax_lse  )

inline

FlashAttnWithKvcache() [3/3]

infini::ops::FlashAttnWithKvcache::FlashAttnWithKvcache ( const Tensor  q, Tensor  k_cache, Tensor  v_cache, const std::optional< Tensor >  k, const std::optional< Tensor >  v, const std::optional< Tensor >  rotary_cos, const std::optional< Tensor >  rotary_sin, const std::optional< Tensor >  cache_seqlens, const std::optional< Tensor >  cache_batch_idx, const std::optional< Tensor >  cache_leftpad, const std::optional< Tensor >  block_table, const std::optional< Tensor >  alibi_slopes, const std::optional< double >  softmax_scale, const bool  causal, const std::vector< int64_t >  window_size, const double  softcap, const bool  rotary_interleaved, const int64_t  num_splits, const bool  return_softmax_lse, Tensor  out, std::optional< Tensor >  softmax_lse  )

inline

Member Function Documentation

operator()() [1/3]

virtual void infini::ops::FlashAttnWithKvcache::operator() ( const Tensor  q, Tensor  k_cache, Tensor  v_cache, const std::optional< Tensor >  k, const std::optional< Tensor >  v, const std::optional< Tensor >  rotary_cos, const std::optional< Tensor >  rotary_sin, const int64_t  cache_seqlens, const std::optional< Tensor >  cache_batch_idx, const std::optional< Tensor >  cache_leftpad, const std::optional< Tensor >  block_table, const std::optional< Tensor >  alibi_slopes, const std::optional< double >  softmax_scale, const bool  causal, const std::vector< int64_t >  window_size, const double  softcap, const bool  rotary_interleaved, const int64_t  num_splits, const bool  return_softmax_lse, Tensor  out, std::optional< Tensor >  softmax_lse  ) const

pure virtual

operator()() [2/3]

virtual void infini::ops::FlashAttnWithKvcache::operator() ( const Tensor  q, Tensor  k_cache, Tensor  v_cache, const std::optional< Tensor >  k, const std::optional< Tensor >  v, const std::optional< Tensor >  rotary_cos, const std::optional< Tensor >  rotary_sin, const std::optional< Tensor >  cache_seqlens, const std::optional< Tensor >  cache_batch_idx, const std::optional< Tensor >  cache_leftpad, const std::optional< Tensor >  block_table, const std::optional< Tensor >  alibi_slopes, const std::optional< double >  softmax_scale, const bool  causal, const std::vector< int64_t >  window_size, const double  softcap, const bool  rotary_interleaved, const int64_t  num_splits, const bool  return_softmax_lse, Tensor  out, std::optional< Tensor >  softmax_lse  ) const

pure virtual

operator()() [3/3]

void infini::ops::FlashAttnWithKvcache::operator() ( const Tensor  q, Tensor  k_cache, Tensor  v_cache, Tensor  out  ) const

inline

ValidateIndexVector()

void infini::ops::FlashAttnWithKvcache::ValidateIndexVector ( const std::optional< Tensor > &  tensor ) const

inlineprotected

Member Data Documentation

alibi_slopes_dtype_

DataType infini::ops::FlashAttnWithKvcache::alibi_slopes_dtype_

protected

alibi_slopes_shape_

Tensor::Shape infini::ops::FlashAttnWithKvcache::alibi_slopes_shape_

protected

alibi_slopes_strides_

Tensor::Strides infini::ops::FlashAttnWithKvcache::alibi_slopes_strides_

protected

batch_size_

Tensor::Size infini::ops::FlashAttnWithKvcache::batch_size_ {0}

protected

block_table_dtype_

DataType infini::ops::FlashAttnWithKvcache::block_table_dtype_

protected

block_table_shape_

Tensor::Shape infini::ops::FlashAttnWithKvcache::block_table_shape_

protected

block_table_strides_

Tensor::Strides infini::ops::FlashAttnWithKvcache::block_table_strides_

protected

cache_batch_idx_dtype_

DataType infini::ops::FlashAttnWithKvcache::cache_batch_idx_dtype_

protected

cache_batch_idx_shape_

Tensor::Shape infini::ops::FlashAttnWithKvcache::cache_batch_idx_shape_

protected

cache_batch_idx_strides_

Tensor::Strides infini::ops::FlashAttnWithKvcache::cache_batch_idx_strides_

protected

cache_leftpad_dtype_

DataType infini::ops::FlashAttnWithKvcache::cache_leftpad_dtype_

protected

cache_leftpad_shape_

Tensor::Shape infini::ops::FlashAttnWithKvcache::cache_leftpad_shape_

protected

cache_leftpad_strides_

Tensor::Strides infini::ops::FlashAttnWithKvcache::cache_leftpad_strides_

protected

cache_seqlens_dtype_

DataType infini::ops::FlashAttnWithKvcache::cache_seqlens_dtype_

protected

cache_seqlens_shape_

Tensor::Shape infini::ops::FlashAttnWithKvcache::cache_seqlens_shape_

protected

cache_seqlens_strides_

Tensor::Strides infini::ops::FlashAttnWithKvcache::cache_seqlens_strides_

protected

device_index_

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

protected

has_alibi_slopes_

bool infini::ops::FlashAttnWithKvcache::has_alibi_slopes_ {false}

protected

has_block_table_

bool infini::ops::FlashAttnWithKvcache::has_block_table_ {false}

protected

has_cache_batch_idx_

bool infini::ops::FlashAttnWithKvcache::has_cache_batch_idx_ {false}

protected

has_k_

bool infini::ops::FlashAttnWithKvcache::has_k_ {false}

protected

has_rotary_cos_

bool infini::ops::FlashAttnWithKvcache::has_rotary_cos_ {false}

protected

has_rotary_sin_

bool infini::ops::FlashAttnWithKvcache::has_rotary_sin_ {false}

protected

has_softmax_lse_

bool infini::ops::FlashAttnWithKvcache::has_softmax_lse_ {false}

protected

has_v_

bool infini::ops::FlashAttnWithKvcache::has_v_ {false}

protected

head_size_

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

protected

k_cache_dtype_

DataType infini::ops::FlashAttnWithKvcache::k_cache_dtype_

protected

k_cache_shape_

Tensor::Shape infini::ops::FlashAttnWithKvcache::k_cache_shape_

protected

k_cache_strides_

Tensor::Strides infini::ops::FlashAttnWithKvcache::k_cache_strides_

protected

k_dtype_

DataType infini::ops::FlashAttnWithKvcache::k_dtype_

protected

k_shape_

Tensor::Shape infini::ops::FlashAttnWithKvcache::k_shape_

protected

k_strides_

Tensor::Strides infini::ops::FlashAttnWithKvcache::k_strides_

protected

out_dtype_

DataType infini::ops::FlashAttnWithKvcache::out_dtype_

protected

out_shape_

Tensor::Shape infini::ops::FlashAttnWithKvcache::out_shape_

protected

out_strides_

Tensor::Strides infini::ops::FlashAttnWithKvcache::out_strides_

protected

q_dtype_

DataType infini::ops::FlashAttnWithKvcache::q_dtype_

protected

q_shape_

Tensor::Shape infini::ops::FlashAttnWithKvcache::q_shape_

protected

q_strides_

Tensor::Strides infini::ops::FlashAttnWithKvcache::q_strides_

protected

rotary_cos_dtype_

DataType infini::ops::FlashAttnWithKvcache::rotary_cos_dtype_

protected

rotary_cos_shape_

Tensor::Shape infini::ops::FlashAttnWithKvcache::rotary_cos_shape_

protected

rotary_cos_strides_

Tensor::Strides infini::ops::FlashAttnWithKvcache::rotary_cos_strides_

protected

rotary_sin_dtype_

DataType infini::ops::FlashAttnWithKvcache::rotary_sin_dtype_

protected

rotary_sin_shape_

Tensor::Shape infini::ops::FlashAttnWithKvcache::rotary_sin_shape_

protected

rotary_sin_strides_

Tensor::Strides infini::ops::FlashAttnWithKvcache::rotary_sin_strides_

protected

softmax_lse_dtype_

DataType infini::ops::FlashAttnWithKvcache::softmax_lse_dtype_

protected

softmax_lse_shape_

Tensor::Shape infini::ops::FlashAttnWithKvcache::softmax_lse_shape_

protected

softmax_lse_strides_

Tensor::Strides infini::ops::FlashAttnWithKvcache::softmax_lse_strides_

protected

v_cache_dtype_

DataType infini::ops::FlashAttnWithKvcache::v_cache_dtype_

protected

v_cache_shape_

Tensor::Shape infini::ops::FlashAttnWithKvcache::v_cache_shape_

protected

v_cache_strides_

Tensor::Strides infini::ops::FlashAttnWithKvcache::v_cache_strides_

protected

v_dtype_

DataType infini::ops::FlashAttnWithKvcache::v_dtype_

protected

v_shape_

Tensor::Shape infini::ops::FlashAttnWithKvcache::v_shape_

protected

v_strides_

Tensor::Strides infini::ops::FlashAttnWithKvcache::v_strides_

protected

---

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

• flash_attn_with_kvcache.h