Description: Performs selected attention computation during the Native Sparse Attention (NSA) inference.
Formula:
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 calculation of selected attention is formed by combining the topk index data obtaining and attention calculation, and the paged attention obtains the kvCache. First, is obtained from and is obtained from by using . The self-attention computation 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.
Each operator has calls. First, [object Object] is called to obtain the workspace size required for computation and the executor that contains the operator computation 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 verification. The following errors may be thrown.
[object Object]
- Deterministic computation:
[object Object]defaults to a deterministic implementation.
- The B axis must be less than or equal to 3072.
- Only paged attention is supported.
- In multi-token inference scenarios, the S axis of
[object Object]can be at most 4 and for each batch,[object Object]must be less than or equal to[object Object].
The following example is for reference only. For details, see .