| cu_seqlens_q |
可选输入 |
表示不同Batch中Query的有效Sequence Length，shape为(B+1, )，仅layout_q为TND场景需传入。 |
INT32 |
ND |
| cu_seqlens_ori_kv |
可选输入 |
表示不同Batch中ori_kv的有效Sequence Length，shape为(B+1, )，仅layout_kv为TND场景需传入。 |
INT32 |
ND |
| cu_seqlens_cmp_kv |
可选输入 |
表示不同Batch中cmp_kv的有效Sequence Length，shape为(B+1, )，仅layout_kv为TND场景需传入。 |
INT32 |
ND |
| seqused_q |
可选输入 |
表示不同Batch中Query实际参与运算的Sequence Length，shape为(B, )。 |
INT32 |
ND |
| seqused_ori_kv |
可选输入 |
表示不同Batch中ori_kv实际参与运算的Sequence Length，shape为(B, )。 |
INT32 |
ND |
| seqused_cmp_kv |
可选输入 |
表示不同Batch中cmp_kv实际参与运算的Sequence Length，shape为(B, )。 |
INT32 |
ND |
| cmp_residual_kv |
可选输入 |
表示不同Batch中cmp_kv压缩后Sequence Length的余数，配合cmp_ratio实现cmp_kv部分的mask和负载计算。cmp_mask_mode=3且cmp_ratio≠1时必须传入，shape为(B, )。 |
INT32 |
ND |
| ori_topk_length |
可选输入 |
预留参数，当前不生效。 |
INT32 |
ND |
| cmp_topk_length |
可选输入 |
预留参数，当前不生效。 |
INT32 |
ND |
| num_heads_q |
属性 |
表示Query的head个数，当前仅支持2/4/8/16/32/64/128。 |
INT32 |
- |
| num_heads_kv |
属性 |
表示Key和Value对应的多头数，当前仅支持1。 |
INT32 |
- |
| head_dim |
属性 |
表示注意力头的维度，当前仅支持512。 |
INT32 |
- |
| quant_mode |
属性 |
表示量化模式，1表示K、V nope为per-token-group量化，scale类型为bfloat16，2表示K、V nope为per-token-group量化，scale类型为float8_e8m0。 |
INT32 |
- |
| batch_size |
可选属性 |
表示Batch数量，默认值为0。 |
INT32 |
- |
| max_seqlen_q |
可选属性 |
表示Query的最长Sequence Length，默认值为0。 |
INT32 |
- |
| max_seqlen_ori_kv |
可选属性 |
表示ori_kv的最长Sequence Length，默认值为0。 |
INT32 |
- |
| max_seqlen_cmp_kv |
可选属性 |
表示cmp_kv的最长Sequence Length，默认值为0。 |
INT32 |
- |
| ori_topk |
可选属性 |
预留参数，当前不生效，表示ori_kv中筛选出的关键稀疏token的个数，0表示非稀疏场景，默认值为0，当前仅支持0。 |
INT32 |
- |
| cmp_topk |
可选属性 |
表示cmp_kv中筛选出的关键稀疏token的个数，0表示非稀疏场景，默认值为0，当前仅支持512/1024。 |
INT32 |
- |
| rope_head_dim |
可选属性 |
表示rope头的维度，默认值为64，当前仅支持64。 |
INT32 |
- |
| cmp_ratio |
可选属性 |
表示对cmp_kv的压缩率，默认值为1，当前仅支持1/4/128。 |
INT32 |
- |
| ori_mask_mode |
可选属性 |
表示q和ori_kv计算的mask模式，默认值为0，当前仅支持4，表示sliding window模式。 |
INT32 |
- |
| cmp_mask_mode |
可选属性 |
表示q和cmp_kv计算的mask模式，默认值为0，当前仅支持3，表示rightDownCausal模式。 |
INT32 |
- |
| ori_win_left |
可选属性 |
表示q和ori_kv计算中q对过去token计算的数量，-1表示无穷大，默认值为-1，当前仅支持127。 |
INT32 |
- |
| ori_win_right |
可选属性 |
表示q和ori_kv计算中q对未来token计算的数量，-1表示无穷大，默认值为-1，当前仅支持0。 |
INT32 |
- |
| layout_q |
可选属性 |
表示Query的排列格式，支持BSND、TND，默认值为BSND。 |
STRING |
- |
| layout_kv |
可选属性 |
表示Key的排列格式，支持BSND、TND、PA_BBND，默认值为BSND。 |
STRING |
- |
| has_ori_kv |
可选属性 |
用于标识是否含有ori_kv，默认值为true。 |
BOOL |
- |
| has_cmp_kv |
可选属性 |
用于标识是否含有cmp_kv，默认值为true。 |
BOOL |
- |
| metadata |
输出 |
表示负载均衡结果输出，shape固定为[1024]。 |
INT32 |
ND |
