infini::ops::FlashAttnVarlenFunc Class Reference

#include <flash_attn_varlen_func.h>

Inheritance diagram for infini::ops::FlashAttnVarlenFunc:

Public Member Functions
	FlashAttnVarlenFunc (const Tensor q, const Tensor k, const Tensor v, const Tensor cu_seqlens_q, const Tensor cu_seqlens_k, const int64_t max_seqlen_q, const int64_t max_seqlen_k, Tensor out)
	FlashAttnVarlenFunc (const Tensor q, const Tensor k, const Tensor v, const Tensor cu_seqlens_q, const Tensor cu_seqlens_k, const std::optional< Tensor > alibi_slopes, const std::optional< Tensor > block_table, const int64_t max_seqlen_q, const int64_t max_seqlen_k, const double dropout_p, const std::optional< double > softmax_scale, const bool causal, const std::vector< int64_t > window_size, const double softcap, const bool deterministic, const bool return_attn_probs, Tensor out, std::optional< Tensor > softmax_lse, std::optional< Tensor > s_dmask)
void	operator() (const Tensor q, const Tensor k, const Tensor v, const Tensor cu_seqlens_q, const Tensor cu_seqlens_k, const int64_t max_seqlen_q, const int64_t max_seqlen_k, Tensor out) const
virtual void	operator() (const Tensor q, const Tensor k, const Tensor v, const Tensor cu_seqlens_q, const Tensor cu_seqlens_k, const std::optional< Tensor > alibi_slopes, const std::optional< Tensor > block_table, const int64_t max_seqlen_q, const int64_t max_seqlen_k, const double dropout_p, const std::optional< double > softmax_scale, const bool causal, const std::vector< int64_t > window_size, const double softcap, const bool deterministic, const bool return_attn_probs, Tensor out, std::optional< Tensor > softmax_lse, std::optional< Tensor > s_dmask) const =0
Public Member Functions inherited from infini::ops::Operator< FlashAttnVarlenFunc >
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	q_shape_
Tensor::Shape	k_shape_
Tensor::Shape	v_shape_
Tensor::Shape	cu_seqlens_q_shape_
Tensor::Shape	cu_seqlens_k_shape_
Tensor::Shape	alibi_slopes_shape_
Tensor::Shape	block_table_shape_
Tensor::Shape	out_shape_
Tensor::Shape	softmax_lse_shape_
Tensor::Shape	s_dmask_shape_
Tensor::Strides	q_strides_
Tensor::Strides	k_strides_
Tensor::Strides	v_strides_
Tensor::Strides	cu_seqlens_q_strides_
Tensor::Strides	cu_seqlens_k_strides_
Tensor::Strides	alibi_slopes_strides_
Tensor::Strides	block_table_strides_
Tensor::Strides	out_strides_
Tensor::Strides	softmax_lse_strides_
Tensor::Strides	s_dmask_strides_
DataType	q_dtype_
DataType	k_dtype_
DataType	v_dtype_
DataType	cu_seqlens_q_dtype_
DataType	cu_seqlens_k_dtype_
DataType	alibi_slopes_dtype_
DataType	block_table_dtype_
DataType	out_dtype_
DataType	softmax_lse_dtype_
DataType	s_dmask_dtype_
bool	has_auxiliary_outputs_ {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 Public Member Functions inherited from infini::ops::Operator< FlashAttnVarlenFunc >
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< FlashAttnVarlenFunc >
static constexpr Device::Type	device_type_
static constexpr std::size_t	implementation_index_

Constructor & Destructor Documentation

FlashAttnVarlenFunc() [1/2]

infini::ops::FlashAttnVarlenFunc::FlashAttnVarlenFunc ( const Tensor  q, const Tensor  k, const Tensor  v, const Tensor  cu_seqlens_q, const Tensor  cu_seqlens_k, const int64_t  max_seqlen_q, const int64_t  max_seqlen_k, Tensor  out  )

inline

FlashAttnVarlenFunc() [2/2]

infini::ops::FlashAttnVarlenFunc::FlashAttnVarlenFunc ( const Tensor  q, const Tensor  k, const Tensor  v, const Tensor  cu_seqlens_q, const Tensor  cu_seqlens_k, const std::optional< Tensor >  alibi_slopes, const std::optional< Tensor >  block_table, const int64_t  max_seqlen_q, const int64_t  max_seqlen_k, const double  dropout_p, const std::optional< double >  softmax_scale, const bool  causal, const std::vector< int64_t >  window_size, const double  softcap, const bool  deterministic, const bool  return_attn_probs, Tensor  out, std::optional< Tensor >  softmax_lse, std::optional< Tensor >  s_dmask  )

inline

Member Function Documentation

operator()() [1/2]

void infini::ops::FlashAttnVarlenFunc::operator() ( const Tensor  q, const Tensor  k, const Tensor  v, const Tensor  cu_seqlens_q, const Tensor  cu_seqlens_k, const int64_t  max_seqlen_q, const int64_t  max_seqlen_k, Tensor  out  ) const

inline

operator()() [2/2]

virtual void infini::ops::FlashAttnVarlenFunc::operator() ( const Tensor  q, const Tensor  k, const Tensor  v, const Tensor  cu_seqlens_q, const Tensor  cu_seqlens_k, const std::optional< Tensor >  alibi_slopes, const std::optional< Tensor >  block_table, const int64_t  max_seqlen_q, const int64_t  max_seqlen_k, const double  dropout_p, const std::optional< double >  softmax_scale, const bool  causal, const std::vector< int64_t >  window_size, const double  softcap, const bool  deterministic, const bool  return_attn_probs, Tensor  out, std::optional< Tensor >  softmax_lse, std::optional< Tensor >  s_dmask  ) const

pure virtual

Member Data Documentation

alibi_slopes_dtype_

DataType infini::ops::FlashAttnVarlenFunc::alibi_slopes_dtype_

protected

alibi_slopes_shape_

Tensor::Shape infini::ops::FlashAttnVarlenFunc::alibi_slopes_shape_

protected

alibi_slopes_strides_

Tensor::Strides infini::ops::FlashAttnVarlenFunc::alibi_slopes_strides_

protected

block_table_dtype_

DataType infini::ops::FlashAttnVarlenFunc::block_table_dtype_

protected

block_table_shape_

Tensor::Shape infini::ops::FlashAttnVarlenFunc::block_table_shape_

protected

block_table_strides_

Tensor::Strides infini::ops::FlashAttnVarlenFunc::block_table_strides_

protected

cu_seqlens_k_dtype_

DataType infini::ops::FlashAttnVarlenFunc::cu_seqlens_k_dtype_

protected

cu_seqlens_k_shape_

Tensor::Shape infini::ops::FlashAttnVarlenFunc::cu_seqlens_k_shape_

protected

cu_seqlens_k_strides_

Tensor::Strides infini::ops::FlashAttnVarlenFunc::cu_seqlens_k_strides_

protected

cu_seqlens_q_dtype_

DataType infini::ops::FlashAttnVarlenFunc::cu_seqlens_q_dtype_

protected

cu_seqlens_q_shape_

Tensor::Shape infini::ops::FlashAttnVarlenFunc::cu_seqlens_q_shape_

protected

cu_seqlens_q_strides_

Tensor::Strides infini::ops::FlashAttnVarlenFunc::cu_seqlens_q_strides_

protected

device_index_

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

protected

has_auxiliary_outputs_

bool infini::ops::FlashAttnVarlenFunc::has_auxiliary_outputs_ {false}

protected

k_dtype_

DataType infini::ops::FlashAttnVarlenFunc::k_dtype_

protected

k_shape_

Tensor::Shape infini::ops::FlashAttnVarlenFunc::k_shape_

protected

k_strides_

Tensor::Strides infini::ops::FlashAttnVarlenFunc::k_strides_

protected

out_dtype_

DataType infini::ops::FlashAttnVarlenFunc::out_dtype_

protected

out_shape_

Tensor::Shape infini::ops::FlashAttnVarlenFunc::out_shape_

protected

out_strides_

Tensor::Strides infini::ops::FlashAttnVarlenFunc::out_strides_

protected

q_dtype_

DataType infini::ops::FlashAttnVarlenFunc::q_dtype_

protected

q_shape_

Tensor::Shape infini::ops::FlashAttnVarlenFunc::q_shape_

protected

q_strides_

Tensor::Strides infini::ops::FlashAttnVarlenFunc::q_strides_

protected

s_dmask_dtype_

DataType infini::ops::FlashAttnVarlenFunc::s_dmask_dtype_

protected

s_dmask_shape_

Tensor::Shape infini::ops::FlashAttnVarlenFunc::s_dmask_shape_

protected

s_dmask_strides_

Tensor::Strides infini::ops::FlashAttnVarlenFunc::s_dmask_strides_

protected

softmax_lse_dtype_

DataType infini::ops::FlashAttnVarlenFunc::softmax_lse_dtype_

protected

softmax_lse_shape_

Tensor::Shape infini::ops::FlashAttnVarlenFunc::softmax_lse_shape_

protected

softmax_lse_strides_

Tensor::Strides infini::ops::FlashAttnVarlenFunc::softmax_lse_strides_

protected

v_dtype_

DataType infini::ops::FlashAttnVarlenFunc::v_dtype_

protected

v_shape_

Tensor::Shape infini::ops::FlashAttnVarlenFunc::v_shape_

protected

v_strides_

Tensor::Strides infini::ops::FlashAttnVarlenFunc::v_strides_

protected

---

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

• flash_attn_varlen_func.h