MmDeqSwigluQuantMmDeqOperation

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 MmDeqSwigluQuantMmDeq operator is a fusion operator of gate up and gate down layers on the entire network.

Figure 1 Computation process

Definition

struct MmDeqSwigluQuantMmDeqParam {
    enum OutputType {
        OUTPUT_FLOAT16 = 0,
        OUTPUT_BFLOAT16,
        OUTPUT_INVALID
    };

    enum WeightUpPermuteType {
        PERMUTE_N256 = 0,
        PERMUTE_N128,
        PERMUTE_INVALID
    };
    OutputType outputType = OUTPUT_FLOAT16;
    WeightUpPermuteType weightUpPermuteType = PERMUTE_N256;
    bool transposeWeightUp = false;
    bool transposeWeightDown = true;
    uint8_t rsv[46] = {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.
weightUpPermuteType	WeightUpPermuteType	PERMUTE_N256	PERMUTE_N256 PERMUTE_N128	Yes	Controls the permutation mode. Currently, only PERMUTE_N256 and PERMUTE_N128 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[46]	uint8_t	{0}	[0]	No	Reserved parameter.

Input

Parameter	Dimension	Data Type	Format	Description
x1	[m, 7168]	int8	ND	Left matrix of the first mm.
permuteWeight1	[7168, 4096]	int8	NZ	Right matrix of the first mm after permutation.
permuteScale1	[4096]	float32	ND	Per-channel dequantization scale of the first mm after permutation.
perTokenScale1	[m]	float32	ND	Per-token dequantization scale of the first mm.
weight2	[7168, 2048]	int8	NZ	Right matrix of the second mm, NZ layout after transposition.
scale2	[7168]	float32	ND	Per-channel dequantization scale of the second mm.

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 mm after dequantization.

Constraints

m indicates the total number of tokens, which cannot exceed 128,000.

Parent topic: atb/infer_op_params.h