MLA Prefill
The PyTorch can be directly connected through the C interface to replace affinity operators on the entire network.
Definition
atb::Status AtbMLAPreFillGetWorkspaceSize(const aclTensor *q, const aclTensor *qRope, const aclTensor *k,
const aclTensor *kRope, const aclTensor *v, const aclTensor *qSeqLen, const aclTensor *kvSeqLen,
const aclTensor *mask, int32_t headNum, float qkScale, int32_t kvHeadNum,
int maskType, uint8_t cacheMode, aclTensor *attenOut,
uint64_t *workspaceSize, atb::Operation **op, atb::Context *context);
atb::Status AtbMLAPreFill(void* workspace, uint64_t workspaceSize, atb::Operation *op, atb::Context *context);
AtbMLAPreFillGetWorkspaceSize Members
Parameter |
Scalar/Tensor |
Dimension |
Data Type |
Format |
Default Value |
Mandatory or Not |
Description |
|---|---|---|---|---|---|---|---|
Q |
Tensor |
[nTokens, qhead * embeddimV] [nTokens, qhead, embeddimV] |
float16/bf16 |
ND |
- |
Yes |
Query. |
QRope |
Tensor |
[nTokens, qhead * 64] [nTokens, qhead, 64] |
float16/bf16 |
ND |
- |
Yes |
Query with rotational positional encoding. |
K |
Tensor |
[batch, max_seq, kv_head * embeddimV] [batch*seq, kv_head, embeddimV] |
float16/bf16 |
ND |
- |
Yes |
Key. |
kRope |
Tensor |
[batch, max_seq, kv_head * 64] [batch*seq, kv_head, 64] |
float16/bf16 |
ND |
- |
Yes |
Key with rotational positional encoding. |
V |
Tensor |
[batch, max_seq, kv_head * embeddimV] [batch*seq, kv_head, embeddimV] |
float16/bf16 |
ND |
- |
Yes |
Value. |
qSeqLen |
Tensor |
[batch] |
int32 |
ND |
- |
Yes |
Input tensor. |
kvSeqLen |
Tensor |
[batch] |
int32 |
ND |
- |
Yes |
Input tensor. |
Mask |
Tensor |
None [512, 512] |
float16/bf16 |
ND |
- |
No |
Attention mask. For UNDEFINED or MASK_TYPE_CAUSAL_MASK, the value can be nullptr. For MASK_TYPE_MASK_FREE, the value is an input tensor. |
headNum |
Scalar |
- |
int32_t |
- |
0 |
Yes |
Number of query headers. |
qkScale |
Scalar |
- |
float |
- |
1.0 |
Yes |
Scaling coefficient post-multiplied by Q*K^T. |
kvHeadNum |
Scalar |
- |
int32_t |
- |
0 |
Yes |
Number of KV heads. |
maskType |
Scalar |
- |
int |
- |
0 |
No |
Mask type. UNDEFINED (default): all-0 mask. MASK_TYPE_SPEC: mask when qseqlen > 1. MASK_TYPE_MASK_FREE: maskfree function. |
calcType |
Scalar |
- |
int |
- |
0 |
No |
Computation type.
|
cacheMode |
Scalar |
- |
int |
- |
0 |
Yes |
|
attenOut |
Tensor |
[nTokens, qhead * embeddimV] [ntokens, kvhead, embedv] |
float16/bf16 |
ND |
- |
Yes |
Attention output. |
Original APIs
For details, see Input and Output When calcType Is CALC_TYPE_PREFILL.