---
title: PrefixEncoder
description: "支持通过C接口直调接入PyTorch，在整网中进行亲和算子替换。"
url: https://www.hiascend.com/document/detail/zh/canncommercial/latest/API/ascendtb/ascendtb_01_0320.html
sourcePath: /source/zh/canncommercial/900/API/ascendtb/ascendtb_01_0320.html
indexId: e09e8d3361def25f124265fe9b721b8df1bcc8fe8b07c732c7e843ccc26b5ddc64
---
# 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 | 张量 | [nTokens, qHiddenSize] [nTokens, headNum, headSize] | float16/bf16 | ND | \- | 是 | query矩阵。 |
| key | 张量 | [numBlocks, blockSize, kvHiddenSize] [numBlocks, blockSize, headNum, headSize] | float16/bf16 | ND | \- | 是 | key矩阵。 |
| value | 张量 | [numBlocks, blockSize, kvHiddenSize] [numBlocks, blockSize, headNum, headSize] | float16/bf16 | ND | \- | 是 | value矩阵。 |
| blockTables | 张量 | [batch, maxBlockNum] | int32 | ND | \- | 是 | 每个query的kvcache的block table，第一维是token索引，第二维表示block索引。 |
| mask | 张量 | [256, 256] [headNum, seqLen, 128] | 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左对齐，只支持 Atlas A2 训练系列产品 / Atlas A2 推理系列产品 和 Atlas A3 推理系列产品 / Atlas A3 训练系列产品 。 9：MASK\_TYPE\_CAUSAL\_MASK, mask内部生成。 |
| headNum | 标量 | \- | int32\_t | \- | 0 | 是 | query头大小，需大于0。 |
| kvHeadNum | 标量 | \- | int32\_t | \- | 0 | 是 | kv头数量，该值需要用户根据使用的模型实际情况传入。 |
| qkScale | 标量 | \- | float | \- | 1 | 是 | Q\*K^T后乘以的缩放系数。 |
| attnOut | 张量 | [nTokens, qHiddenSize] [nTokens, headNum, headSize] | float16/bf16 | ND | \- | 是 | attention输出。 |


#### 原始接口

请参见calcType为PREFIX_ENCODER时输入输出。