BNSD Dimension Input

Description

Generally, 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 some scenarios, the performance is better when [b, n, s, d] is passed.

How to Enable

Set inputLayout to TYPE_BNSD.

The input parameters are as follows:

• When calctype is UNDEFINED, ENCODER, or DECODER, on the 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 the 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 the kv-bypass function 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 the 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, the BNSD dimension input function is used, and high precision is not supported.