#include <moe_align_block_size.h>

Inheritance diagram for infini::ops::MoeAlignBlockSize:

Public Member Functions
	MoeAlignBlockSize (const Tensor topk_ids, const std::optional< Tensor > expert_map, const int64_t num_experts, const int64_t block_size, Tensor sorted_token_ids, Tensor experts_ids, Tensor num_tokens_post_pad)

void	operator() (const Tensor topk_ids, const std::optional< Tensor > expert_map, const int64_t num_experts, const int64_t block_size, Tensor sorted_token_ids, Tensor experts_ids, Tensor num_tokens_post_pad) const

Public Member Functions inherited from infini::ops::Operator< MoeAlignBlockSize >
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	ValidateInvocation (const Tensor topk_ids, const std::optional< Tensor > maybe_expert_map, const int64_t num_experts, const int64_t block_size, Tensor sorted_token_ids, Tensor experts_ids, Tensor num_tokens_post_pad) const

bool	CallMetadataMatches (const Tensor topk_ids, const std::optional< Tensor > maybe_expert_map, const Tensor sorted_token_ids, const Tensor experts_ids, const Tensor num_tokens_post_pad) const

void	Validate (const Tensor topk_ids, const std::optional< Tensor > maybe_expert_map, Tensor sorted_token_ids, Tensor experts_ids, Tensor num_tokens_post_pad) const

virtual void	Run (const Tensor topk_ids, const std::optional< Tensor > maybe_expert_map, const int64_t num_experts, const int64_t block_size, Tensor sorted_token_ids, Tensor experts_ids, Tensor num_tokens_post_pad) const =0

Protected Attributes
Tensor	topk_ids_metadata_

std::optional< Tensor >	expert_map_metadata_

Tensor	sorted_token_ids_metadata_

Tensor	experts_ids_metadata_

Tensor	num_tokens_post_pad_metadata_

Tensor::Size	numel_ {0}

int64_t	num_experts_ {0}

int64_t	block_size_ {0}

Tensor::Size	sorted_token_ids_size_ {0}

Tensor::Size	experts_ids_size_ {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< MoeAlignBlockSize >
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< MoeAlignBlockSize >
static constexpr Device::Type	device_type_

static constexpr std::size_t	implementation_index_

Constructor & Destructor Documentation

◆ MoeAlignBlockSize()

infini::ops::MoeAlignBlockSize::MoeAlignBlockSize	(	const Tensor	topk_ids,
		const std::optional< Tensor >	expert_map,
		const int64_t	num_experts,
		const int64_t	block_size,
		Tensor	sorted_token_ids,
		Tensor	experts_ids,
		Tensor	num_tokens_post_pad
	)

inline

Member Function Documentation

◆ CallMetadataMatches()

bool infini::ops::MoeAlignBlockSize::CallMetadataMatches	(	const Tensor	topk_ids,
		const std::optional< Tensor >	maybe_expert_map,
		const Tensor	sorted_token_ids,
		const Tensor	experts_ids,
		const Tensor	num_tokens_post_pad
	)		const

inlineprotected

◆ operator()()

void infini::ops::MoeAlignBlockSize::operator()	(	const Tensor	topk_ids,
		const std::optional< Tensor >	expert_map,
		const int64_t	num_experts,
		const int64_t	block_size,
		Tensor	sorted_token_ids,
		Tensor	experts_ids,
		Tensor	num_tokens_post_pad
	)		const

inline

◆ Run()

virtual void infini::ops::MoeAlignBlockSize::Run	(	const Tensor	topk_ids,
		const std::optional< Tensor >	maybe_expert_map,
		const int64_t	num_experts,
		const int64_t	block_size,
		Tensor	sorted_token_ids,
		Tensor	experts_ids,
		Tensor	num_tokens_post_pad
	)		const

protectedpure virtual

◆ Validate()

void infini::ops::MoeAlignBlockSize::Validate	(	const Tensor	topk_ids,
		const std::optional< Tensor >	maybe_expert_map,
		Tensor	sorted_token_ids,
		Tensor	experts_ids,
		Tensor	num_tokens_post_pad
	)		const

inlineprotected

◆ ValidateInvocation()

void infini::ops::MoeAlignBlockSize::ValidateInvocation	(	const Tensor	topk_ids,
		const std::optional< Tensor >	maybe_expert_map,
		const int64_t	num_experts,
		const int64_t	block_size,
		Tensor	sorted_token_ids,
		Tensor	experts_ids,
		Tensor	num_tokens_post_pad
	)		const

inlineprotected

Member Data Documentation

◆ block_size_

int64_t infini::ops::MoeAlignBlockSize::block_size_ {0}

protected

◆ expert_map_metadata_

std::optional<Tensor> infini::ops::MoeAlignBlockSize::expert_map_metadata_

protected

◆ experts_ids_metadata_

Tensor infini::ops::MoeAlignBlockSize::experts_ids_metadata_

protected

◆ experts_ids_size_

Tensor::Size infini::ops::MoeAlignBlockSize::experts_ids_size_ {0}

protected

◆ num_experts_

int64_t infini::ops::MoeAlignBlockSize::num_experts_ {0}

protected

◆ num_tokens_post_pad_metadata_

Tensor infini::ops::MoeAlignBlockSize::num_tokens_post_pad_metadata_

protected

◆ numel_

Tensor::Size infini::ops::MoeAlignBlockSize::numel_ {0}

protected

◆ sorted_token_ids_metadata_

Tensor infini::ops::MoeAlignBlockSize::sorted_token_ids_metadata_

protected

◆ sorted_token_ids_size_

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

protected

◆ topk_ids_metadata_

Tensor infini::ops::MoeAlignBlockSize::topk_ids_metadata_

protected

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

moe_align_block_size.h

Public Member Functions

Protected Member Functions

Protected Attributes

Additional Inherited Members

Constructor & Destructor Documentation

◆ MoeAlignBlockSize()

Member Function Documentation

◆ CallMetadataMatches()

◆ operator()()

◆ Run()

◆ Validate()

◆ ValidateInvocation()

Member Data Documentation

◆ block_size_

◆ expert_map_metadata_

◆ experts_ids_metadata_

◆ experts_ids_size_

◆ num_experts_

◆ num_tokens_post_pad_metadata_

◆ numel_

◆ sorted_token_ids_metadata_

◆ sorted_token_ids_size_

◆ topk_ids_metadata_