maskShape Constraints

Atlas A2 training products/Atlas A2 inference products PA_ENCODER (ND)

MaskType

dtype

Shape

Extra

Remarks

UNDEFINED

-

No input

-

-

NORM

Same as query

[maxSeqLen,maxSeqLen] / [batch,maxSeqLen,maxSeqLen]

-

In 3D cases, batch is the same as seqLen.

ALIBI

Same as query

1. [headNum, maxSeqLen, maxSeqLen]

2. [batch, headNum, maxSeqLen, maxSeqLen]

-

The values respectively indicate:

1. The same mask for each batch;

2. Different masks for each batch

NORM_COMPRESS

Same as query

[128,128]

-

Quantization is not supported.

ALIBI_COMPRESS/SQRT

Same as query

[256,256] or [headNum,seqlen,128]

slopes [headNum]

fp16 masks must be HIGH_PRECISION.

ALIBI_LEFT_ALIGN

Same as query

[256,256] only

slopes [headNum]

-

SWA_NORM

Same as query

[S,S]

-

windowSize>0

SWA_COMPRESS

Same as query

[512,512]

-

windowSize>0

Atlas A2 training products/Atlas A2 inference products PREFIX_ENCODER

MaskType

Input Count

mask shape

slopes

ALIBI_COMPRESS/SQRT

8

Same as PA ALIBI compress

√

NORM_COMPRESS

7

Same as PA [128,128]

×

CAUSAL_MASK

6

None (internally generated)

×

Atlas 350 accelerator cards PA_ENCODER (ND)

MaskType

dtype

Shape

NORM

int8

[maxSeqLen,maxSeqLen] / [batch,maxSeqLen,maxSeqLen]

NORM_COMPRESS

int8

[2048,2048]

ALIBI

int8

[batch,headNum,maxSeqLen,maxSeqLen]

Atlas training products, Atlas inference products PA (NZ)

MaskType

Format

Shape

Remarks

NORM_COMPRESS

Same as query

[1,8,128,16]

legacy

NORM_COMPRESS long

Same as query

[1,128,2048,16]

Atlas training products only + fp16

SWA_COMPRESS

Same as query

[1,32,512,16]