Adapts to the
[object Object]operator in the decode ([object Object]) and prefill ([object Object]) inference scenarios.Compared with FusedInferAttentionScoreV4, this API adds the qStartIdxOptional, kvStartIdxOptional, and pseType parameters.
NOTE
KV cache specific to the decode scenario: KV cache is a common technology for optimizing the inference performance of foundation models. During sampling, the transformer model uses the given prompt/context as the initial input for inference (parallel processing supported), and then generates additional tokens one by one to improve the generated sequence (reflecting the auto-regressive property of the model). The transformer performs the self-attention operation during sampling. Therefore, KV vectors need to be extracted for each item (regardless of the prompt/context or generated token) in the current sequence. These vectors are stored in a matrix called KV cache.
Formulas:
Self-attention constructs an attention model by leveraging the relationships within the input samples. The principle assumes there is an input sample sequence of length , where each element of is a -dimensional vector. Each -dimensional vector can be regarded as a token embedding. Such a sequence is transformed by three weight matrices to obtain three matrices.
The computation formula for self-attention is generally defined as follows, where , , and are key attribute elements of the input sample, obtained through spatial transformation and unified into a single feature space. "Attention" in the formula and operator name is an abbreviation for "self-attention."
The
[object Object]function in this operator employs the[object Object]function. The self-attention calculation formula is as follows:The product of and represents the attention to the input . To prevent this value from becoming excessively large, it is typically scaled by dividing by the square root of , followed by row-wise softmax normalization. The result is then multiplied by to produce an matrix.
Note:
[object Object]
Each operator has calls. First, [object Object] is called to obtain the input parameters and compute the required workspace size based on the process. Then, [object Object] is called to perform computation.
Parameters:
[object Object]Returns
[object Object]: status code. For details, see .The first-phase API implements input parameter validation. The following error codes may be returned.
[object Object]
Deterministic computation:
[object Object]defaults to a deterministic implementation.
Common constraints
- Processing when the input parameter is empty:
- An empty tensor indicates that the shapeSize of the required input and output is 0. In the empty tensor scenario, if attentionOut is empty, the return value is empty. Otherwise, the return value is all 0s. If lse is empty, the return value is empty. If lse is not empty, the return value is all inf. When the tensor is not empty, the input is intercepted normally.
- The shapeSize of all tensors in query and attentionOut is 0, which is an empty tensor.
- The shapeSize of all tensors in query and attentionOut is not 0. If lse is not empty and the shapeSize of all tensors in key and value is 0, the tensor is an empty tensor.
- If both attentionOut and lse are empty, the tensor is an empty tensor.
- If the tensor is empty, the verification process is skipped. Otherwise, the normal verification process is performed.
- The restrictions in the BNSD_BSND, BSH_BNSD, BSND_BNSD, BSH_NBSD, BSND_NBSD, and BNSD_NBSD scenarios are as follows:
- When dimension
[object Object]of[object Object]is[object Object]:- Only BSH_NBSD, BSND_NBSD, and BNSD_NBSD are supported.
- Only the decode mla scenario is supported. The queryRope and keyRope cannot be empty, and the d value of queryRope and keyRope is 64.
- When dimension
[object Object]of[object Object]is not[object Object]:- Only BNSD_BSND, BSH_BNSD, and BSND_BNSD are supported.
- The prefill mla or gqa non-quantization scenario is supported. In the prefill mla scenario, either of the following conditions must be met:
- The d value of the query, key, and value is 128. The queryRope and keyRope cannot be empty, and the d value of queryRope and keyRope is 64.
- The d value of the query and key is 192, the d value of the value is 128, and the queryRope and keyRope are empty.
- In the gqa non-quantization scenario, BSH_BNSD and BSND_BNSD support only D = 64 or D = 128. BNSD_BSND supports only D = 16 alignment (32 alignment when the output dtype is int8).
- In the BSH_BNSD and BSND_BNSD scenarios, left padding, tensorlist, pse, and prefix are not supported.
- BSH_BNSD and BSND_BNSD do not support fake-quantization. BNSD_BSND supports fake-quantization.
- In the fake-quantization scenario, BNSD_BSND does not support QS = 1.
- When dimension
- The restrictions on the query, key, and value inputs in the TND, NTD, TND_NTD, and NTD_TND scenarios are as follows:
- When dimension
[object Object]of[object Object]is[object Object]:- Only TND and TND_NTD are supported.
- Only the decode mla scenario is supported. The queryRope and keyRope cannot be empty, and the d value of queryRope and keyRope is 64.
- Left padding, tensorlist, pseType=0, prefix, and fake-quantization are not supported.
- When dimension
[object Object]of[object Object]is not[object Object]:- Only TND, NTD, and NTD_TND are supported.
- The prefill mla or gqa non-quantization scenario is supported. In the prefill mla scenario, either of the following conditions must be met:
- The d value of the query, key, and value is 128. The queryRope and keyRope cannot be empty, and the d value of queryRope and keyRope is 64.
- The d value of the query and key is 192, the d value of the value is 128, and the queryRope and keyRope are empty.
- In the GQA non-quantization scenario, NTD and NTD_TND support only D = 64 or D = 128.
- Left padding, tensorlist, pseType = 0, prefix, and fake-quantization are not supported.
- When dimension
- Processing when the input parameter is empty:
[object Object][object Object]General scenarios
[object Object][object Object][object Object]PseShift
[object Object][object Object][object Object]Mask
[object Object][object Object][object Object]ActualSeqLen
[object Object][object Object][object Object]Fake-quantization parameters
[object Object][object Object][object Object]PagedAttention
[object Object][object Object][object Object]Constraints on the number of input parameters and input and output related to INT8/FP8 quantization:
[object Object][object Object][object Object]Left padding
[object Object][object Object][object Object]Prefix
[object Object][object Object][object Object]MLA ([object Object]queryRope[object Object] and [object Object]keyRope[object Object] not null)
[object Object]
The following example is for reference only. For details, see .