| num_heads_q |
属性 |
Q的多头数，目前仅支持64。 |
INT32 |
- |
| num_heads_kv |
属性 |
K和V的多头数，目前仅支持1。 |
INT32 |
- |
| head_dim |
属性 |
注意力头的维度，目前仅支持512。 |
INT32 |
- |
| cu_seqlens_q |
可选输入 |
当layout_query为TND时，表示不同Batch中q的有效token数，维度为B+1，大小为参数中每个元素的值表示当前batch与之前所有batch的token数总和，即前缀和。 |
INT32 |
ND |
| cu_seqlens_ori_kv |
可选输入 |
当layout_kv为TND时，表示不同Batch中ori_kv的有效token数，维度为B+1，大小为参数中每个元素的值表示当前batch与之前所有batch的token数总和，即前缀和。目前layout_kv仅支持PA_ND，故设置此参数无效。 |
INT32 |
ND |
| cu_seqlens_cmp_kv |
可选输入 |
当layout_kv为TND时，表示不同Batch中cmp_kv的有效token数，维度为B+1，大小为参数中每个元素的值表示当前batch与之前所有batch的token数总和，即前缀和。目前layout_kv仅支持PA_ND，故设置此参数无效。 |
INT32 |
ND |
| seqused_q |
可选输入 |
表示不同Batch中q实际参与运算的token数，维度为B。目前暂不支持指定该参数。 |
INT32 |
ND |
| seqused_kv |
可选输入 |
表示不同Batch中ori_kv实际参与运算的token数，维度为B。 |
INT32 |
ND |
| batch_size |
可选属性 |
输入样本批量大小。 |
INT32 |
- |
| max_seqlen_q |
可选属性 |
表示所有batch中`q`的最大有效token数。 |
INT32 |
- |
| max_seqlen_kv |
可选属性 |
表示所有batch中`ori_kv`的最大有效token数。 |
INT32 |
- |
| ori_topk |
可选属性 |
表示通过QLI算法从`ori_kv`中筛选出的关键稀疏token的个数。目前暂不支持指定该参数。 |
INT32 |
- |
| cmp_topk |
可选属性 |
表示通过QLI算法从`cmp_kv`中筛选出的关键稀疏token的个数，目前仅支持512。 |
INT32 |
- |
| cmp_ratio |
可选属性 |
表示对`ori_kv`的压缩率，数据范围支持4/128， |
INT32 |
- |
| ori_mask_mode |
可选属性 |
表示q和ori_kv计算的mask模式，仅支持输入默认值4，代表band模式的mask。 |
INT32 |
- |
| cmp_mask_mode |
可选属性 |
表示q和cmp_kv计算的mask模式，仅支持输入默认值3，代表rightDownCausal模式的mask。 |
INT32 |
- |
| ori_win_left |
可选属性 |
表示q和ori_kv计算中q对过去token计算的数量，仅支持默认值127。 |
INT32 |
- |
| ori_win_right |
可选属性 |
表示q和ori_kv计算中q对未来token计算的数量，仅支持默认值0。 |
INT32 |
- |
| layout_q |
可选属性 |
用于标识输入q的数据排布格式,默认值为BSND，目前支持传入BSND和TND。 |
STRING |
| layout_kv |
可选属性 |
用于标识输入ori_kv和cmp_kv的数据排布格式，目前仅支持传入默认值PA_ND。 |
STRING |
- |
| has_ori_kv |
可选属性 |
是否传入ori_kv。 |
BOOL |
- |
| has_cmp_kv |
可选属性 |
是否传入cmp_kv。 |
BOOL |
- |
| device |
可选属性 |
用于获取设备信息。 |
STRING |
- |
| metadata |
输出 |
每个cube核上FlashAttention计算任务的Batch、Head、以及 Q 和 K 的分块的索引，以及每个vector核上FlashDecode的规约任务索引。 |
INT32 |
- |
