FusedAddTopkDivOperation

The PyTorch can be directly connected through the C interface to replace affinity operators on the entire network.

Definition

atb::Status AtbFusedAddTopkDivGetWorkspaceSize(const aclTensor *x, const aclTensor *addNum, const aclTensor *mappingNum,
                                               const aclTensor *mappingTable, uint32_t groupNum, uint32_t groupTopk,
                                               uint32_t n, uint32_t k, int activationType, bool isNorm, float scale,
                                               bool enableExpertMapping, aclTensor *y, aclTensor *indices,
                                               uint64_t *workspaceSize, atb::Operation **op, atb::Context *context);
atb::Status AtbFusedAddTopkDiv(void *workspace, uint64_t workspaceSize, atb::Operation *op, atb::Context *context);

AtbFusedAddTopkDivGetWorkspaceSize Members

Parameter

Scalar/Tensor

Dimension

Data Type

Format

Default Value

Mandatory or Not

Description

x

Tensor

[batch_size, expert_num]

float16/float32/bf16

ND

-

Yes

Input tensor.

addNum

Tensor

[expert_num]

float16/float32/bf16

ND

-

Yes

Input tensor, which is used to add to x. The data type and format are the same as those of x.

mappingNum

Tensor

[expert_num]

int32

ND

-

No

This parameter is not used when enableExpertMapping is set to false. When enableExpertMapping is set to true, an input tensor is used, indicating the number of logical experts to which each physical expert is actually mapped.

mappingTable

Tensor

[expert_num, max_redundant_expert_num]

int32

ND

-

No

This parameter is not used when enableExpertMapping is set to false. When enableExpertMapping is set to true, an input tensor for the mapping table between physical experts and logical experts is used.

groupNum

Scalar

-

uint32_t

-

1

Yes

Number of groups.

groupTopk

Scalar

-

uint32_t

-

1

Yes

Select k groups.

n

Scalar

-

uint32_t

-

1

Yes

Select n maximum values in a group and sum them up.

k

Scalar

-

uint32_t

-

1

Yes

Select the top k values.

activationType

Scalar

-

int

-

8

Yes

Activation type.

Only 8: ACTIVATION_SIGMOID is supported.

isNorm

Scalar

-

bool

-

true

Yes

Normalized or not.

scale

Scalar

-

float

-

1.0

Yes

Multiplication coefficient after normalization.

enableExpertMapping

Scalar

-

bool

-

false

Yes

Whether to enable the mapping from physical experts to logical experts. If this parameter is set to false, two tensors are input. If this parameter is set to true, four tensors are input.

y

Tensor

[batch_size, k]

float32

ND

-

Yes

Output tensor.

indices

Tensor

[batch_size, k]

int32

ND

-

Yes

Output tensor.

Original APIs

For details, see FusedAddTopkDivOperation.