#include <moe_wna16_gemm.h>

Inheritance diagram for infini::ops::MoeWna16Gemm:

Public Member Functions
	MoeWna16Gemm (const Tensor input, const Tensor b_qweight, const Tensor b_scales, std::optional< Tensor > b_qzeros, std::optional< Tensor > topk_weights, const Tensor sorted_token_ids, const Tensor expert_ids, const Tensor num_tokens_post_pad, const int64_t top_k, const int64_t block_size_m, const int64_t block_size_n, const int64_t block_size_k, const int64_t bit, Tensor output)

virtual void	operator() (const Tensor input, const Tensor b_qweight, const Tensor b_scales, std::optional< Tensor > b_qzeros, std::optional< Tensor > topk_weights, const Tensor sorted_token_ids, const Tensor expert_ids, const Tensor num_tokens_post_pad, const int64_t top_k, const int64_t block_size_m, const int64_t block_size_n, const int64_t block_size_k, const int64_t bit, Tensor output) const =0

Public Member Functions inherited from infini::ops::Operator< MoeWna16Gemm >
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	ValidateCallMetadata (const Tensor input, const Tensor b_qweight, const Tensor b_scales, std::optional< Tensor > b_qzeros, std::optional< Tensor > topk_weights, const Tensor sorted_token_ids, const Tensor expert_ids, const Tensor num_tokens_post_pad, const int64_t top_k, const int64_t block_size_m, const int64_t block_size_n, const int64_t block_size_k, const int64_t bit, const Tensor output) const

Protected Attributes
Tensor::Size	num_experts_ {0}

Tensor::Size	m_ {0}

Tensor::Size	n_ {0}

Tensor::Size	k_ {0}

Tensor::Size	num_groups_ {0}

Tensor::Size	group_size_ {0}

Tensor::Size	sorted_token_ids_size_ {0}

Tensor::Size	expert_ids_size_ {0}

int64_t	top_k_ {0}

int64_t	block_size_m_ {0}

int64_t	block_size_n_ {0}

int64_t	block_size_k_ {0}

int64_t	bit_ {0}

DataType	dtype_

bool	has_qzeros_ {false}

bool	has_topk_weights_ {false}

Device::Type	device_type_

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< MoeWna16Gemm >
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< MoeWna16Gemm >
static constexpr Device::Type	device_type_

static constexpr std::size_t	implementation_index_

Constructor & Destructor Documentation

◆ MoeWna16Gemm()

infini::ops::MoeWna16Gemm::MoeWna16Gemm	(	const Tensor	input,
		const Tensor	b_qweight,
		const Tensor	b_scales,
		std::optional< Tensor >	b_qzeros,
		std::optional< Tensor >	topk_weights,
		const Tensor	sorted_token_ids,
		const Tensor	expert_ids,
		const Tensor	num_tokens_post_pad,
		const int64_t	top_k,
		const int64_t	block_size_m,
		const int64_t	block_size_n,
		const int64_t	block_size_k,
		const int64_t	bit,
		Tensor	output
	)

inline

Member Function Documentation

◆ operator()()

virtual void infini::ops::MoeWna16Gemm::operator()	(	const Tensor	input,
		const Tensor	b_qweight,
		const Tensor	b_scales,
		std::optional< Tensor >	b_qzeros,
		std::optional< Tensor >	topk_weights,
		const Tensor	sorted_token_ids,
		const Tensor	expert_ids,
		const Tensor	num_tokens_post_pad,
		const int64_t	top_k,
		const int64_t	block_size_m,
		const int64_t	block_size_n,
		const int64_t	block_size_k,
		const int64_t	bit,
		Tensor	output
	)		const

pure virtual

◆ ValidateCallMetadata()

void infini::ops::MoeWna16Gemm::ValidateCallMetadata	(	const Tensor	input,
		const Tensor	b_qweight,
		const Tensor	b_scales,
		std::optional< Tensor >	b_qzeros,
		std::optional< Tensor >	topk_weights,
		const Tensor	sorted_token_ids,
		const Tensor	expert_ids,
		const Tensor	num_tokens_post_pad,
		const int64_t	top_k,
		const int64_t	block_size_m,
		const int64_t	block_size_n,
		const int64_t	block_size_k,
		const int64_t	bit,
		const Tensor	output
	)		const

inlineprotected

Member Data Documentation

◆ bit_

int64_t infini::ops::MoeWna16Gemm::bit_ {0}

protected

◆ block_size_k_

int64_t infini::ops::MoeWna16Gemm::block_size_k_ {0}

protected

◆ block_size_m_

int64_t infini::ops::MoeWna16Gemm::block_size_m_ {0}

protected

◆ block_size_n_

int64_t infini::ops::MoeWna16Gemm::block_size_n_ {0}

protected

◆ device_index_

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

protected

◆ device_type_

Device::Type infini::ops::MoeWna16Gemm::device_type_

protected

◆ dtype_

DataType infini::ops::MoeWna16Gemm::dtype_

protected

◆ expert_ids_size_

Tensor::Size infini::ops::MoeWna16Gemm::expert_ids_size_ {0}

protected

◆ group_size_

Tensor::Size infini::ops::MoeWna16Gemm::group_size_ {0}

protected

◆ has_qzeros_

bool infini::ops::MoeWna16Gemm::has_qzeros_ {false}

protected

◆ has_topk_weights_

bool infini::ops::MoeWna16Gemm::has_topk_weights_ {false}

protected

◆ k_

Tensor::Size infini::ops::MoeWna16Gemm::k_ {0}

protected

◆ m_

Tensor::Size infini::ops::MoeWna16Gemm::m_ {0}

protected

◆ n_

Tensor::Size infini::ops::MoeWna16Gemm::n_ {0}

protected

◆ num_experts_

Tensor::Size infini::ops::MoeWna16Gemm::num_experts_ {0}

protected

◆ num_groups_

Tensor::Size infini::ops::MoeWna16Gemm::num_groups_ {0}

protected

◆ sorted_token_ids_size_

Tensor::Size infini::ops::MoeWna16Gemm::sorted_token_ids_size_ {0}

protected

◆ top_k_

int64_t infini::ops::MoeWna16Gemm::top_k_ {0}

protected

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

moe_wna16_gemm.h

Public Member Functions

Protected Member Functions

Protected Attributes

Additional Inherited Members

Constructor & Destructor Documentation

◆ MoeWna16Gemm()

Member Function Documentation

◆ operator()()

◆ ValidateCallMetadata()

Member Data Documentation

◆ bit_

◆ block_size_k_

◆ block_size_m_

◆ block_size_n_

◆ device_index_

◆ device_type_

◆ dtype_

◆ expert_ids_size_

◆ group_size_

◆ has_qzeros_

◆ has_topk_weights_

◆ k_

◆ m_

◆ n_

◆ num_experts_

◆ num_groups_

◆ sorted_token_ids_size_

◆ top_k_