BNSD Dimension Input

Description

Typically, the dimensions of q, k, and v passed to the SelfAttention operator are [batch, seqLen, headNum, head_dim], that is, [b, s, n, d], or the variants of [b, s, n, d] after axis combination. In certain cases, passing [b, n, s, d] delivers better performance.

How to Enable

Set inputLayout to TYPE_BNSD.

The input parameters are as follows:

• When calcType is UNDEFINED, ENCODER, or DECODER on Atlas A2 training products/Atlas A2 inference products and Atlas A3 inference products/Atlas A3 training products:

  Input Tensor	Dimension	Data Type	Format	Description
  query	[batch, headNum, seqLen, headSize]	float16/bf16	ND	Query matrix.
  cacheK	[layer, batch, headNum, seqLen, headSize]	float16/bf16	ND	NPU: Stores all previous keys. During execution, keys are updated to cacheK. CPU: The input is the prepared cacheK. The input is divided into batch tensors as the input std::vector<tensor>. In this case, the layer dimension must be 1.
  cacheV	[layer, batch, headNum, seqLen, headSize]	float16/bf16	ND	NPU: Stores all previous values. During execution, values are updated to cacheV. CPU: The input is the prepared cacheV. The input is divided into batch tensors as the input std::vector<tensor>. In this case, the layer dimension must be 1.

• If calcType is UNDEFINED, ENCODER, or DECODER on Atlas inference products:

  Input Tensor	Dimension	Data Type	Format	Description
  query	[batch, headNum, seqLen, headSize]	float16/bf16	ND	Query matrix.
  cacheK	[layer, batch*headNum, headSize / 16, kvMaxSeq, 16]	float16/bf16	NZ	NPU: Stores all previous keys. During execution, keys are updated to cacheK. CPU: The input is the prepared cacheK. The input is divided into batch tensors as the input std::vector<tensor>. In this case, the layer dimension must be 1.
  cacheV	[layer, batch*headNum, headSize / 16, kvMaxSeq, 16]	float16/bf16	NZ	NPU: Stores all previous values. During execution, values are updated to cacheV. CPU: The input is the prepared cacheV. The input is divided into batch tensors as the input std::vector<tensor>. In this case, the layer dimension must be 1.

• When calcType is set to PA_ENCODER:

  Parameter	Dimension	Data Type	Format	Description
  query	[batch, headNum, qSeqLen, headSize]	bf16	ND	Query matrix.
  cacheK	[batch, kvHeadNum, kvSeqLen, headSize]	bf16	ND	Key matrix.
  cacheV	[batch, kvHeadNum, kvSeqLen, headSize]	bf16	ND	Value matrix.
  SeqLen	[batch] / [2, batch]	int32/uint32	ND	If the shape is [batch], it indicates the sequence length of each batch. The sequence length of each batch is the same for query, cacheK, and cacheV. If the shape is [2, batch], SeqLen[0] indicates the sequence length of query, and SeqLen[1] indicates the sequence length of cacheK and cacheV.
  attnOut	[batch, headNum, qSeqLen, headSize]	bf16	ND	Output tensor.

Constraints

• BNSD is available only when kv-bypass is enabled, that is, kvcacheCfg is set to K_BYPASS_V_BYPASS or calcType is set to PA_ENCODER.
• When the BNSD dimension is used for input and calcType is not PA_ENCODER, maskType cannot be MASK_TYPE_UNDEFINED. When calcType is set to PA_ENCODER, maskType can only be MASK_TYPE_UNDEFINED.
• dimNum of seqlen can be 2 only when calcType is set to PA_ENCODER and BNSD is enabled.
• When calcType is set to PA_ENCODER, the function is not supported by Atlas inference products.
• The values of headSize for query, key, and value must be the same and less than or equal to 256.
• When calcType is set to PA_ENCODER, BNSD dimension input is used, and high precision is not supported.