Constraints

  • When the input data layout format is TYPE_BSND:
    • If PA_ENCODER is not used, query, key, and value on Atlas A2 training products/Atlas A2 inference products and Atlas A3 inference products/Atlas A3 training products can be provided either as a 2D tensor of shape [nTokens, hiddenSize] or as a 4D tensor of shape [batch, seqLen, headNum, headSize].
    • If PA_ENCODER is used:
      • query, key, and value on Atlas A2 training products/Atlas A2 inference products and Atlas A3 inference products/Atlas A3 training products can be provided either as a 2D tensor of shape [nTokens, hiddenSize] or as a 3D tensor of shape [nTokens, headNum, headSize].
      • query, key, and value on Atlas 350 accelerator cards can only be provided as a 3D tensor of shape [nTokens, headNum, headSize].
  • Atlas inference products:
    • 0 < batch <= 2000.
    • The mask size must be the actual value of maxSeqLen.
    • qSeqLen must be equal to kvSeqLen.
    • When maskType is set to MASK_TYPE_NORM_COMPRESS and the mask dimension is [1,128,2048,16] (2048-long window), the preceding constraints have the following exceptions:
      • The constraint that the mask size must be the actual value of maxSeqLen does not apply: In this scenario, a fixed 2048 x 2048 shared compressed mask canvas is used, and the mask size is 2048 x 2048, which is irrelevant to the actual value of maxSeqLen.
      • The mask data type can only be float16.
      • Quantization is not supported (quantType must be TYPE_QUANT_UNDEFINED).
      • scaleType must be SCALE_TYPE_TOR (LOGN is not supported).
      • isTriuMask is forcibly set to 1 in the source code (at the constructor layer, when maskType is MASK_TYPE_NORM_COMPRESS and 310P is used). It is not necessary to pass this parameter externally (a triangular mask, where the lower triangle and diagonal are 0 (visible), and the upper triangle is a large negative value (invisible); -60000 or -10000 is recommended).
      • maskStride = 2048.
      • The mask must be input in NZ format ([1,128,2048,16]). The ND format [2048,2048] must be converted to NZ on the user side.
  • In the following cases, headSize is equal to vHeadSize, and the value range is (0, 256]. In other cases, headSize can be different from vHeadSize, and the value range is (0, 576].
    • Quantization is enabled.
    • The attention uses logN scaling.
    • SWA is enabled.
    • inputLayout is set to TYPE_BNSD.
    • Atlas inference products: headSize must be a multiple of 16.
    • Atlas training products: headSize must be a multiple of 16.
  • nTokens in the tensor dimension:
    • Atlas A2 training products/Atlas A2 inference products and Atlas A3 inference products/Atlas A3 training products: sum of seqLen in each batch.
    • Atlas inference products: If PA_ENCODER is used, the value is the sum of seqLen of all batches rounded up to an integer multiple of 16.
    • In other cases, the value is the sum of seqLen of all batches rounded up to an integer multiple of 16.
  • How to select a mask:
    • In the full phase, set maskType to MASK_TYPE_NORM or MASK_TYPE_ALIBI based on the configuration of the native model.
    • In the incremental phase, if the original model is not alibi and the mask is not padded, maskType can be set to UNDEFINED, that is, the mask is not passed. In other cases, set maskType to MASK_TYPE_NORM or MASK_TYPE_ALIBI based on the configuration of the native model.
    • For long sequences, you can consider compressing the mask. For details about how to construct a compressed mask, see mask.
  • If PA_ENCODER is used:
    • qScale does not take effect.
    • The number of dimensions of query, key, and value must be the same.
  • To use the GQA mode, headNum must be greater than kvHeadNum.
  • The Atlas training products support only PA_ENCODER and basic functions and do not support functions listed in Functions.
  • hiddenSize must be a multiple of 16 when calcType is set to ENCODER or DECODER in the non-bypass scenario.
  • If headSize of query, key, and value is not a multiple of 64, the operator performance deteriorates. To deliver the optimal performance, you are advised to set headSize of query, key, and value to a multiple of 64.
  • For CPU tensors, use hostData to set the data address; no need to set deviceData. For NPU tensors, use deviceData to set the data address; no need to set hostData. However, if the calculation type is ENCODER or DECODER with non-KV-bypass configured, when both deviceData and hostData are passed for the tokenOffset and Seqlen tensors, the data in deviceData is used for kvCache calculation.
  • headNum%kvHeadNum == 0.
  • Atlas 350 accelerator cards: Only MASK_TYPE_ALIBI, MASK_TYPE_NORM, and MASK_TYPE_NORM_COMPRESS are supported.
  • When headSize is not equal to headSizeV:
    • maskType cannot be MASK_TYPE_ALIBI.
    • headSize = 192, headSizeV = 128.