GmmDeqSwigluQuantGmmDeqOperation

Applicable Products

Hardware Model	Supported or Not
Atlas 350 accelerator card	x
Atlas A3 inference products/Atlas A3 training products	√
Atlas A2 training products/Atlas A2 inference products	√
Atlas training products	x
Atlas inference products	x
Atlas 200I/500 A2 inference products	x

Description

The GmmDeqSwigluQuantGmmDeq operator is a fusion operator of gate up and gate down layers on the entire network.

Figure 1 Computation process

Definition

struct GmmDeqSwigluQuantGmmDeqParam {
    enum OutputType {
        OUTPUT_FLOAT16 = 0,
        OUTPUT_BFLOAT16,
        OUTPUT_INVALID
    };
    enum GroupListType {
        GROUP_LIST_CUMSUM = 0,
        GROUP_LIST_SINGLE,
        GROUP_LIST_INVALID
    };
    enum WeightUpPermuteType {
        PERMUTE_N256 = 0,
        PERMUTE_N128,
        PERMUTE_INVALID
    };
    OutputType outputType = OUTPUT_FLOAT16;
    GroupListType groupListType = GROUP_LIST_CUMSUM;
    WeightUpPermuteType weightUpPermuteType = PERMUTE_N256;
    bool transposeWeightUp = false;
    bool transposeWeightDown = true;
    uint8_t rsv[42] = {0};
};

Parameters

Member	Type	Default Value	Value Range	Mandatory or Not	Description
outputType	OutputType	OUTPUT_FLOAT16	OUTPUT_FLOAT16	Yes	Controls the output type. OUTPUT_FLOAT16: default value. The output data type is float16. Currently, only the default value is supported. OUTPUT_BFLOAT16: The output data type is bf16. OUTPUT_INVALID: invalid value.
groupListType	GroupListType	GROUP_LIST_CUMSUM	GROUP_LIST_CUMSUM	Yes	Controls the groupList type. GROUP_LIST_CUMSUM: default value, which is the accumulated sum. Currently, only the default value is supported. GROUP_LIST_SINGLE: single-group value mode. GROUP_LIST_INVALID: invalid value.
weightUpPermuteType	WeightUpPermuteType	PERMUTE_N256	PERMUTE_N256 PERMUTE_N128	Yes	Controls the permutation mode. Currently, only 256 and 128 are supported. PERMUTE_N256: basic block width for permutation on the N axis. PERMUTE_N128: basic block width for permutation on the N axis. PERMUTE_INVALID: invalid value.
transposeWeightUp	bool	false	false	Yes	Controls whether to transpose the weight of the previous GroupedMatmul. Currently, only non-transposition is supported.
transposeWeightDown	bool	true	true	Yes	Controls whether to transpose the weight of the next GroupedMatmul. Currently, only transposition is supported.
rsv[42]	uint8_t	{0}	[0]	No	Reserved

Input

Parameter	Dimension	Data Type	Format	Description
x1	[m, 7168]	int8	ND	Left matrix of the first gmm.
permuteWeight1	[G, 7168, 4096]	int8	NZ	Right matrix of the first gmm after permutation.
permuteScale1	[G, 4096]	float32	ND	Per-channel dequantization scale of the first gmm after permutation.
perTokenScale1	[m]	float32	ND	Per-token dequantization scale of the first gmm.
groupList	[G]	int64	ND	Group list of the grouped matmul, in prefix sum mode.
weight2	[G, 7168, 2048]	int8	NZ	Right matrix of the second gmm, NZ layout after transposition.
scale2	[G, 7168]	float32	ND	Per-channel dequantization scale of the second gmm.

Permutation logic:

def permute_weight(w: torch.Tensor, tile_n=256):
    *dims, n = w.shape
    order = list(range(len(dims))) + [-2, -3, -1]
    return w.reshape(*dims, 2, n // tile_n, tile_n // 2).permute(order).reshape(*dims, n).contiguous()

The value of weightUpPermuteType must correspond to the tile_n parameter in the permutation logic. Currently, PERMUTE_N256 corresponds to tile_n=256, and PERMUTE_N128 corresponds to tile_n=128.

Output

Parameter	Dimension	Data Type	Format	Description
output	[m, 7168]	float16	ND	Result of the second gmm after dequantization.

Constraints

m indicates the total number of tokens, which cannot exceed 128,000.
G indicates the number of experts, that is, the length of groupList. The value cannot exceed 32.
In CUMSUM mode, groupList is a monotonically non-decreasing list with a length of G. The value of the ith bit indicates the sum of the lengths of the ith group and all groups before the ith group on the m axis. Ensure that the maximum value in the list does not exceed m.

Parent topic: atb/infer_op_params.h