PrefixEncoder
支持通过C接口直调接入PyTorch，在整网中进行亲和算子替换。
定义
atb::Status AtbSelfAttentionPrefixEncoderGetWorkspaceSize(const aclTensor *query, const aclTensor *key,
const aclTensor *value, const aclTensor *blockTables,
const aclTensor *mask, const aclTensor *seqLen,
const aclTensor *kvSeqLen, const aclTensor *slopes,
int maskType, int32_t headNum, int32_t kvHeadNum,
float qkScale, aclTensor *attnOut, uint64_t *workspaceSize,
atb::Operation **op, atb::Context *context);
atb::Status AtbSelfAttentionPrefixEncoder(void *workspace, uint64_t workspaceSize, atb::Operation *op,
atb::Context *context);
AtbSelfAttentionPrefixEncoderGetWorkspaceSize成员
参数 |
标量/张量 |
维度 |
数据类型 |
格式 |
默认值 |
是否必选 |
描述 |
|---|---|---|---|---|---|---|---|
query |
张量 |
|
float16/bf16 |
ND |
- |
是 |
query矩阵。 |
key |
张量 |
|
float16/bf16 |
ND |
- |
是 |
key矩阵。 |
value |
张量 |
|
float16/bf16 |
ND |
- |
是 |
value矩阵。 |
blockTables |
张量 |
[batch, maxBlockNum] |
int32 |
ND |
- |
是 |
每个query的kvcache的block table，第一维是token索引，第二维表示block索引。 |
mask |
张量 |
|
float16/bf16 |
ND |
- |
否 |
alibiMask。 maskType不为MASK_TYPE_ALIBI_COMPRESS或MASK_TYPE_ALIBI_COMPRESS_SQRT时，需要置为nullptr。 |
seqLen |
张量 |
[batch] |
int32/uint32 |
ND |
- |
是 |
query对应的每个batch的序列长度 Host侧。 |
kvSeqLen |
张量 |
[batch] |
int32/uint32 |
ND |
- |
是 |
key、value对应的每个batch的序列长度。Host侧tensor。 |
slopes |
张量 |
[head] |
fp32 |
ND |
- |
否 |
即alibi coefficient，当使用alibiMask时传入tensor时，为alibiMask的每个head的系数。 maskType为MASK_TYPE_CAUSAL_MASK时，需要置为nullptr。 |
maskType |
标量 |
- |
int |
- |
0 |
是 |
mask类型。 仅支持 3：压缩mask MASK_TYPE_ALIBI_COMPRESS，alibi压缩mask。 4：MASK_TYPE_ALIBI_COMPRESS_SQRT，alibi压缩开平方mask。 5： MASK_TYPE_ALIBI_COMPRESS_LEFT_ALIGN，alibi压缩mask左对齐，只支持 9：MASK_TYPE_CAUSAL_MASK, mask内部生成。 |
headNum |
标量 |
- |
int32_t |
- |
0 |
是 |
query头大小，需大于0。 |
kvHeadNum |
标量 |
- |
int32_t |
- |
0 |
是 |
kv头数量，该值需要用户根据使用的模型实际情况传入。 |
qkScale |
标量 |
- |
float |
- |
1 |
是 |
Q*K^T后乘以的缩放系数。 |
attnOut |
张量 |
|
float16/bf16 |
ND |
- |
是 |
attention输出。 |