[object Object][object Object][object Object]undefined[object Object]

• API功能：适配增量&全量推理场景的FlashAttention算子，既可以支持全量计算场景（PromptFlashAttention），也可支持增量计算场景（IncreFlashAttention）。当不涉及system prefix、左padding、kv量化参数合一、pertensor全量化的场景，推荐使用本接口，否则使用老接口[object Object]。
• 计算公式：$$Attention(Q,K,V)=Softmax(\frac{QK^T}{\sqrt{d}})V$$ 其中$Q$和$K^T$的乘积代表输入$x$的注意力，$d$表示隐藏层最小的单元尺寸。为避免该值变得过大，通常除以$d$的开根号进行缩放，并对每行进行softmax归一化，与$V$相乘后得到一个$n*d$的矩阵，$n$为输出矩阵的行数。

[object Object][object Object]

[object Object]

• query（[object Object]）：必选参数，表示attention结构的Query输入，对应公式中的[object Object]。不支持非连续的Tensor，数据格式支持ND。

  • [object Object]Atlas A2 训练系列产品/Atlas A2 推理系列产品[object Object]：数据类型支持[object Object]、[object Object]。
  • [object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]：数据类型支持[object Object]、[object Object]。

• key（[object Object]）：必选参数，表示attention结构的Key输入，对应公式中的[object Object]。不支持非连续的Tensor，数据格式支持ND。

  • [object Object]Atlas A2 训练系列产品/Atlas A2 推理系列产品[object Object]：数据类型支持[object Object]、[object Object]、[object Object]、[object Object]（[object Object]）。
  • [object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]：数据类型支持[object Object]、[object Object]、[object Object]、[object Object]（[object Object]）。

• value（[object Object]）：必选参数，表示attention结构的Value输入，对应公式中的[object Object]。不支持非连续的Tensor，数据格式支持ND。

  • [object Object]Atlas A2 训练系列产品/Atlas A2 推理系列产品[object Object]：数据类型支持[object Object]、[object Object]、[object Object]、[object Object]（[object Object]）。
  • [object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]：数据类型支持[object Object]、[object Object]、[object Object]、[object Object]（[object Object]）。

• [object Object]*[object Object]：必选参数，代表其之前的变量是位置相关的，必须按照顺序输入；之后的变量是可选参数，位置无关，需要使用键值对赋值，不赋值会使用默认值。

• query_rope（[object Object]）：可选参数，表示MLA（Multi-head Latent Attention）结构中[object Object]的rope信息，数据类型支持[object Object]、[object Object]，不支持非连续的Tensor，数据格式支持ND。

• key_rope（[object Object]）：可选参数，表示MLA（Multi-head Latent Attention）结构中的[object Object]的rope信息，数据类型支持[object Object]、[object Object]，不支持非连续的Tensor，数据格式支持ND。

• pse_shift（[object Object]）：可选参数，表示attention结构内部的位置编码参数，数据类型支持[object Object]、[object Object]，数据类型与[object Object]数据类型需满足类型推导规则。不支持非连续的Tensor，数据格式支持ND。如不使用该功能可传入None。

  • Q_S不为1，当[object Object]为[object Object]类型时，要求[object Object]为float16或int8类型；当[object Object]为[object Object]类型时，要求[object Object]为[object Object]类型。输入shape类型需为(B, Q_N, Q_S, KV_S)或(1, Q_N, Q_S, KV_S)。对于[object Object]的KV_S为非32对齐的场景，建议padding到32字节来提高性能，多余部分的填充值不做要求。
  • Q_S为1，当[object Object]为[object Object]类型时，要求[object Object]为[object Object]类型；当[object Object]为[object Object]类型时，要求[object Object]为[object Object]类型。输入shape类型需为(B, Q_N, 1, KV_S)或(1, Q_N, 1, KV_S)。对于[object Object]的KV_S为非32对齐的场景，建议padding到32字节来提高性能，多余部分的填充值不做要求。

• atten_mask（[object Object]）：可选参数，对QK结果进行mask，用来指示是否计算Token间的相关性。数据类型支持[object Object]、[object Object]和[object Object]。不支持非连续的Tensor，数据格式支持ND。如不使用该功能可传入None。

  • Q_S不为1时，建议shape输入(Q_S, KV_S)、(B, Q_S, KV_S)、(1, Q_S, KV_S)、(B, 1, Q_S, KV_S)、(1, 1, Q_S, KV_S)。
  • Q_S为1时，建议shape输入(B, KV_S)、(B, 1, KV_S)、(B, 1, 1, KV_S)。

  如果Q_S、KV_S非16或32对齐，可以向上取到对齐的S。综合约束请见。

• actual_seq_qlen（[object Object]）：可选参数，表示不同Batch中[object Object]的有效seqlen，数据类型支持[object Object]。如果不指定seqlen可传入None，表示和[object Object]的shape的S长度相同。 [object Object]Atlas A2 训练系列产品/Atlas A2 推理系列产品[object Object]、[object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]：该入参中每个Batch的有效seqlen不超过[object Object]中对应batch的seqlen。当seqlen传入长度为1时，每个Batch使用相同seqlen；当seqlen传入长度>=Batch时，取seqlen的前Batch个数；其他长度不支持。当[object Object]的input_layout为TND时，该入参必须传入，且以该入参元素的数量作为Batch值。该入参中每个元素的值表示当前Batch与之前所有Batch的seqlen和，因此后一个元素的值必须>=前一个元素的值，且不能出现负值。

• actual_seq_kvlen（[object Object]）：可选参数，表示不同Batch中[object Object]/[object Object]的有效seqlenKv，数据类型支持[object Object]。如果不指定None，表示和key/value的shape的S长度相同。不同O_S值有不同的约束，具体参见。

• block_table（[object Object]）：可选参数，表示PageAttention中KV存储使用的block映射表，数据类型支持[object Object]。数据格式支持ND。如不使用该功能可传入None。

• dequant_scale_query（[object Object]）：可选参数，表示[object Object]的反量化参数，仅支持pertoken叠加perhead。数据类型支持[object Object]。数据格式支持ND，如不使用该功能可传入None，综合约束请见。

• dequant_scale_key（[object Object]）：可选参数，kv伪量化参数分离时表示[object Object]的反量化因子。数据格式支持ND。通常支持perchannel、pertensor、pertoken、pertensor叠加perhead、pertoken叠加perhead、pertoken叠加使用page attention模式管理scale、pertoken叠加per head并使用page attention模式管理scale。如不使用该功能可传入None，综合约束请见。

  • [object Object]Atlas A2 训练系列产品/Atlas A2 推理系列产品[object Object]：数据类型支持[object Object]、[object Object]、[object Object]。
  • [object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]：数据类型支持[object Object]、[object Object]、[object Object]。

• dequant_offset_key（[object Object]）：可选参数，kv伪量化参数分离时表示[object Object]的反量化偏移。数据类型支持[object Object]、[object Object]、[object Object]。数据格式支持ND。支持perchannel、pertensor、pertoken、pertensor叠加perhead、pertoken叠加perhead、pertoken叠加使用page attention模式管理offset、pertoken叠加perhead并使用page attention模式管理offset。如不使用该功能可传入None，综合约束请见。

• dequant_scale_value（[object Object]）：可选参数，kv伪量化参数分离时表示[object Object]的反量化因子。数据类型支持[object Object]、[object Object]、[object Object]。数据格式支持ND。支持perchannel、pertensor、pertoken、pertensor叠加perhead、pertoken叠加perhead、pertoken叠加使用page attention模式管理scale、pertoken叠加perhead并使用page attention模式管理scale。如不使用该功能可传入None，综合约束请见。

  • [object Object]Atlas A2 训练系列产品/Atlas A2 推理系列产品[object Object]：数据类型支持[object Object]、[object Object]、[object Object]。
  • [object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]：数据类型支持[object Object]、[object Object]、[object Object]。

• dequant_offset_value（[object Object]）：可选参数，kv伪量化参数分离时表示[object Object]的反量化偏移。数据类型支持[object Object]、[object Object]、[object Object]。数据格式支持ND。支持perchannel、pertensor、pertoken、pertensor叠加perhead、pertoken叠加perhead、pertoken叠加使用page attention模式管理offset、pertoken叠加perhead并使用page attention模式管理offset。如不使用该功能可传入None，综合约束请见。

• dequant_scale_key_rope（[object Object]）：可选参数，预留参数，暂未使用，使用默认值即可。

• quant_scale_out（[object Object]）：可选参数，表示输出的量化因子。数据类型支持[object Object]、[object Object]。数据格式支持ND。支持pertensor、perchannel。当输入为[object Object]时，同时支持[object Object]、[object Object]，否则仅支持[object Object]。perchannel格式，当输出layout为BSH时，要求[object Object]所有维度的乘积等于H；其他layout要求乘积等于Q_N*D（建议输出layout为BSH时，quant_scale_out shape传入(1, 1, H)或(H,)；输出为BNSD时，建议传入(1, Q_N, 1, D)或(Q_N, D)；输出为BSND时，建议传入(1, 1, Q_N, D)或(Q_N, D)）。如不使用该功能可传入None，综合约束请见。

• quant_offset_out（[object Object]）：可选参数，表示输出的量化偏移。数据类型支持[object Object]、[object Object]。数据格式支持ND。支持pertensor、perchannel。若传入[object Object]，需保证其类型和shape信息与[object Object]一致。如不使用该功能可传入None，综合约束请见。

• learnable_sink（[object Object]）：可选参数，表示通过可学习的“Sink Token”起到吸收Attention Score的作用，数据类型支持[object Object]，数据格式支持ND，shape输入为(Q_N,)。默认值为None，综合约束请见。

• num_query_heads（[object Object]）：可选参数，代表query的head个数，数据类型支持int64，在BNSD场景下，需要与shape中的[object Object]的N轴shape值相同，否则执行异常。

• num_key_value_heads（[object Object]）：可选参数，代表[object Object]、[object Object]中head个数，用于支持GQA（Grouped-Query Attention，分组查询注意力）场景，数据类型支持[object Object]。默认值为0，表示[object Object]/[object Object]/[object Object]的head个数相等，需要满足[object Object]整除[object Object]，[object Object]与[object Object]的比值不能大于64。在BSND、BNSD、BNSD_BSND（仅支持Q_S大于1）场景下，还需要与shape中的[object Object]/[object Object]的N轴shape值相同，否则执行异常。

• softmax_scale（[object Object]）：可选参数，公式中d开根号的倒数，代表缩放系数，作为计算流中Muls的scalar值，数据类型支持[object Object]。数据类型与[object Object]数据类型需满足数据类型推导规则。默认值为1.0。

• pre_tokens（[object Object]）：可选参数，用于稀疏计算，表示attention需要和前几个Token计算关联。数据类型支持[object Object]。默认值为2147483647，Q_S为1时该参数无效。

• next_tokens（[object Object]）：可选参数，用于稀疏计算，表示attention需要和后几个Token计算关联。数据类型支持[object Object]。默认值为2147483647，Q_S为1时该参数无效。

• input_layout（[object Object]）：可选参数，用于标识输入[object Object]、[object Object]、[object Object]的数据排布格式，默认值为"BSH"。

  [object Object]

  • [object Object]Atlas A2 训练系列产品/Atlas A2 推理系列产品[object Object]、[object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]：支持BSH、BSND、BNSD、BNSD_BSND（输入为BNSD时，输出格式为BSND，仅支持Q_S大于1）、BSH_NBSD、BSND_NBSD、BNSD_NBSD（输出格式为NBSD时，仅支持Q_S大于1且小于等于16）、TND、TND_NTD、NTD_TND（TND相关场景综合约束请见）。

  其中BNSD_BSND含义指当输入为BNSD，输出格式为BSND，仅支持Q_S大于1。

• sparse_mode（[object Object]）：可选参数，表示sparse的模式。数据类型支持[object Object]。Q_S为1且不带rope输入时该参数无效。input_layout为TND、TND_NTD、NTD_TND时，综合约束请见。

  • sparse_mode为0时，代表defaultMask模式，如果atten_mask未传入则不做mask操作，忽略pre_tokens和next_tokens（内部赋值为INT_MAX）；如果传入，则需要传入完整的atten_mask矩阵（S1*S2），表示pre_tokens和next_tokens之间的部分需要计算。
  • sparse_mode为1时，代表allMask，必须传入完整的atten_mask矩阵（S1*S2）。
  • sparse_mode为2时，代表leftUpCausal模式的mask，需要传入优化后的atten_mask矩阵（2048*2048）。
  • sparse_mode为3时，代表rightDownCausal模式的mask，对应以右顶点为划分的下三角场景，需要传入优化后的atten_mask矩阵（2048*2048）。
  • sparse_mode为4时，代表band模式的mask，需要传入优化后的atten_mask矩阵（2048*2048）。
  • sparse_mode为5、6、7、8时，分别代表prefix、global、dilated、block_local，均暂不支持。默认值为0。综合约束请见。

• block_size（[object Object]）：可选参数，表示PageAttention中KV存储每个block中最大的token个数，默认为0，数据类型支持[object Object]。

• query_quant_mode（[object Object]）：可选参数， 表示query的伪量化方式。仅支持传入3，代表模式3：pertoken叠加perhead模式。

• key_quant_mode（[object Object]）：可选参数，表示key的伪量化方式，默认值为0。取值除了[object Object]为0且[object Object]为1的场景外，其他场景取值需要与[object Object]一致。综合约束请见。

  [object Object]Atlas A2 训练系列产品/Atlas A2 推理系列产品[object Object]、[object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]：当Q_S>=2时，仅支持传入值为0、1；当Q_S=1时，支持取值0、1、2、3、4、5。

  • key_quant_mode为0时，代表perchannel模式（perchannel包含pertensor）。
  • key_quant_mode为1时，代表pertoken模式。
  • key_quant_mode为2时，代表pertensor叠加perhead模式。
  • key_quant_mode为3时，代表pertoken叠加perhead模式。
  • key_quant_mode为4时，代表pertoken叠加使用page attention模式管理scale/offset模式。
  • key_quant_mode为5时，代表pertoken叠加perhead并使用page attention模式管理scale/offset模式。

• value_quant_mode（[object Object]）：可选参数，表示[object Object]的伪量化方式，模式编号与[object Object]一致，默认值为0。取值除了[object Object]为0且[object Object]为1的场景外，其他场景取值需要与[object Object]一致。综合约束请见。

  [object Object]Atlas A2 训练系列产品/Atlas A2 推理系列产品[object Object]、[object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]：当Q_S>=2时，仅支持传入值为0、1；当Q_S=1时，支持取值0、1、2、3、4、5。

• inner_precise（[object Object]）：可选参数，数据类型支持[object Object]，支持4种模式：0、1、2、3。一共两位bit位，第0位（bit0）表示高精度或者高性能选择，第1位（bit1）表示是否做行无效修正。当Q_S>1时，sparse_mode为0或1，并传入用户自定义mask的情况下，建议开启行无效；Q_S为1时该参数仅支持取0和1。综合约束请见。

  • inner_precise为0时，代表开启高精度模式，且不做行无效修正。
  • inner_precise为1时，代表高性能模式，且不做行无效修正。
  • inner_precise为2时，代表开启高精度模式，且做行无效修正。
  • inner_precise为3时，代表高性能模式，且做行无效修正。

  [object Object]

• return_softmax_lse（[object Object]）：可选参数，表示是否输出[object Object]，支持S轴外切（增加输出）。true表示输出，false表示不输出；默认值为false。

• query_dtype（[object Object]）：可选参数，表示[object Object]的数据类型，预留参数，暂未使用，使用默认值即可。

• key_dtype（[object Object]）：可选参数，表示[object Object]的数据类型，预留参数，暂未使用，使用默认值即可。

• value_dtype（[object Object]）：可选参数，表示[object Object]的数据类型，预留参数，暂未使用，使用默认值即可。

• query_rope_dtype（[object Object]）：可选参数，表示[object Object]的数据类型，预留参数，暂未使用，使用默认值即可。

• key_rope_dtype（[object Object]）：可选参数，表示[object Object]的数据类型，预留参数，暂未使用，使用默认值即可。

• key_shared_prefix_dtype（[object Object]）：可选参数，表示key_shared_prefix的数据类型，预留参数，暂未使用，使用默认值即可。

• value_shared_prefix_dtype（[object Object]）：可选参数，表示value_shared_prefix的数据类型，预留参数，暂未使用，使用默认值即可。

• dequant_scale_query_dtype（[object Object]）：可选参数，表示[object Object]的数据类型，预留参数，暂未使用，使用默认值即可。

• dequant_scale_key_dtype（[object Object]）：可选参数，表示[object Object]的数据类型，预留参数，暂未使用，使用默认值即可。

• dequant_scale_value_dtype（[object Object]）：可选参数，表示[object Object]的数据类型，预留参数，暂未使用，使用默认值即可。

• dequant_scale_key_rope_dtype（[object Object]）：可选参数，表示[object Object]的数据类型，预留参数，暂未使用，使用默认值即可。

[object Object]

• attention_out（[object Object]）：公式中的输出，数据类型支持[object Object]、[object Object]、[object Object]。数据格式支持ND。限制：该入参的D维度与[object Object]的D保持一致，其余维度需要与入参[object Object]的shape保持一致。
• softmax_lse（[object Object]）：ring attention算法对query乘key的结果先取max得到softmax_max，query乘key的结果减去softmax_max，再取exp，最后取sum，得到softmax_sum，最后对softmax_sum取log，再加上softmax_max得到的结果。数据类型支持[object Object]，当[object Object]为True时，一般情况下输出shape为(B, Q_N, Q_S, 1)，若input_layout为TND/NTD_TND时，输出shape为(T,Q_N,1)；当[object Object]为False时，输出shape为[1]的值为0的Tensor。

[object Object]

• 该接口支持推理场景下使用。

• 该接口支持图模式。

• 该接口与PyTorch配合使用时，需要保证CANN相关包与PyTorch相关包的版本匹配。

• 入参为空的处理：算子内部需要判断参数query是否为空，如果是空则直接返回。参数query不为空Tensor，参数key、value为空tensor（即S2为0），则填充全零的对应shape的输出（填充attention_out）。attention_out为空Tensor时，框架会处理。

• 参数key、value中对应tensor的shape需要完全一致；非连续场景下key、value的tensorlist中的batch只能为1，个数等于query的B，N和D需要相等。

• int8量化相关入参数量与输出数据格式的综合限制：

  • 输出为int8的场景：入参quant_scale_out需要存在，quant_offset_out可选，不传时默认为0。

    • [object Object]Atlas A2 训练系列产品/Atlas A2 推理系列产品[object Object]：输入为int8。
    • [object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]：输入为int8。

  • 输出为float16的场景：若存在入参quant_offset_out或quant_scale_out（即不为None），则报错并返回。

    • [object Object]Atlas A2 训练系列产品/Atlas A2 推理系列产品[object Object]：输入为int8。
    • [object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]：输入为int8。

  • 入参quant_offset_out和quant_scale_out支持pertensor或perchannel格式，数据类型支持float32、bfloat16。

• query_rope和key_rope输入时即为MLA场景，参数约束如下：

  • query_rope的数据类型、数据格式与query一致。

  • key_rope的数据类型、数据格式与key一致。

  • query_rope和key_rope要求同时配置或同时不配置，不支持只配置其中一个。

  • 当query_rope和key_rope非空时，支持如下特性：

    • [object Object]Atlas A2 训练系列产品/Atlas A2 推理系列产品[object Object]、[object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]：query的D只支持512、128；
    • 当query的D等于512时：
      • sparse：Q_S等于1时支持sparse=0且不传mask或sparse=4且传入mask，Q_S大于1时支持sparse=3或sparse=4且传入mask；
      • [object Object]Atlas A2 训练系列产品/Atlas A2 推理系列产品[object Object]、[object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]约束如下：
      • sparse不为4时：
        • query_rope配置时要求query的S为1-16、N为1/2/4/8/16/32/64/128，query_rope shape中D为64，其余维度与query一致；
        • key_rope配置时要求key的N为1、D为512，key_rope shape中D为64，其余维度与key一致；
        • 支持key、value、key_rope的input_layout格式为ND或NZ。当input_layout为NZ时，数据类型为float16或bfloat16时，输入参数key和value的格式为[blockNum, KV_N, D/16, blockSize, 16]；当数据类型为int8时，输入参数key和value的格式为[blockNum, KV_N, D/32, blockSize, 32]；
        • input_layout形状支持BSH、BSND、BNSD、BNSD_NBSD、BSND_NBSD、BSH_NBSD、TND、TND_NTD，当数据格式为NZ时input_layout不支持BNSD、BNSD_NBSD。
        • 该场景下，必须开启PageAttention，此时block_size支持16、128，其中数据格式为NZ时block_size不支持配置16。
        • 不支持开启softmax_lse、左padding、tensorlist、pse、prefix、伪量化、后量化、空Tensor。
        • 支持全量化场景，即输入query/key/value全为int8，query_rope和key_rope为bfloat16，输出为bfloat16的场景：
          • 入参dequant_scale_query、dequant_scale_key、dequant_scale_value需要同时存在，且其数据类型仅支持FP32。
          • 不支持传入quant_scale_out、quant_offset_out、dequant_offset_key、dequant_offset_value，否则报错并返回。
          • query_quant_mode仅支持pertoken叠加perhead模式，key_quant_mode和value_quant_mode仅支持pertensor模式。
          • 支持key、value、key_rope的input_layout格式为NZ。
      • sparse为4时：
        • query的每batch的S不大于key的每batch的S、N为128，query_rope shape中D为64，其余维度与query一致；
        • key的S不大于131088，N为1、D为512，key_rope shape中D为64，其余维度与key一致；
        • key、value、key_rope仅支持ND输入；
        • input_layout形状仅支持BSND。
        • 该场景下，支持开启PageAttention，此时block_size支持64、128。
        • 不支持开启SoftMaxLse、tensorlist、pse、伪量化、后量化、全量化场景。
    • 当query的D等于128时：

      • [object Object]Atlas A2 训练系列产品/Atlas A2 推理系列产品[object Object]、[object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]约束如下：

        • input_layout：TND、NTD_TND。

        • query_rope配置时要求query_rope的shape中d为64，其余维度与query一致。

        • key_rope配置时要求key_rope的shape中d为64，其余维度与key一致。

        • 不支持左padding、tensorlist、pse、page attention、prefix、伪量化、全量化、后量化、空Tensor。

      • 其余约束同TND、NTD_TND场景下的综合限制保持一致。

  • TND、TND_NTD、NTD_TND场景下query、key、value输入的综合限制：

    • T小于等于1M；

    • sparse模式仅支持sparse_mode=0且不传mask，或sparse_mode=3且传入mask；

    • actual_seq_qlen和actual_seq_kvlen必须传入，且以该入参元素数量作为Batch值（注意入参元素数量要小于等于4096）。该入参中每个元素的值表示当前Batch与之前所有Batch的Sequence Length和，因此后一个元素的值必须大于等于前一个元素的值；

    • 当query的D等于512时：

      • 支持TND、TND_NTD；
      • 必须开启page attention，此时actual_seq_kvlen长度等于key/value的batch值，代表每个batch的实际长度，值不大于KV_S；
      • 支持query每个batch的s为1-16；
      • 要求query的n为1/2/4/8/16/32/64/128，key、value的n为1；
      • 要求query_rope和key_rope不等于空，query_rope和key_rope的d为64；
      • 不支持开启softmax_lse、左padding、tensorlist、pse、prefix、伪量化、全量化、后量化、空Tensor。

    • 当query的D不等于512时：

      • 当query_rope和key_rope为空时：TND场景，要求Q_D、K_D、V_D等于128，或者Q_D、K_D等于192，V_D等于128/192；NTD_TND场景，要求Q_D、K_D等于128/192，V_D等于128。当query_rope和key_rope不为空时，要求Q_D、K_D、V_D等于128；GQA和PA场景不支持V_D等于192;
      • 支持TND、NTD_TND；
      • 数据类型仅支持bfloat16；
      • 当sparse_mode=3时，要求每个batch单独的actual_seq_qlen<actual_seq_kvlen；
      • sparse模式支持sparse_mode=4且传入mask；当sparse_mode=4时，要求preTokens >= -actual_seq_qlen、nextTokens >= -actual_seq_kvlen、preTokens + nextTokens >= 0；
      • page attention场景下仅支持blocksize 128,512或1024;
      • 不支持左padding、tensorlist、pse、prefix、伪量化、全量化、后量化、空Tensor；
      • 不支持图模式配置Tiling调度优化（tiling_schedule_optimize=True）、reduce-overhead执行模式（config.mode="reduce-overhead"）。
      • actual_seq_qlen和actual_seq_kvlen的元素个数不大于4096。

• GQA伪量化场景下KV为NZ格式时的参数约束如下：

  • 支持perchannel和pertoken模式，query数据类型固定为bfloat16，key&value固定为int8；query&key&value的D仅支持128；query Sequence Length仅支持1-16；
  • input_layout仅支持BSH、BSND、BNSD；
  • 仅支持page_attention场景，blockSize仅支持128或512；
  • key&value仅支持NZ输入，输入格式为[blockNum, KV_N, D/32, blockSize, 32]；
  • dequant_scale_key和dequant_scale_value的dtype：perchannel模式下，仅支持bfloat16类型；pertoken模式下，仅支持float32类型；
  • dequant_scale_key和dequant_scale_value的shape：perchannel模式下，当layout为BSH时，必须传入[H]；layout为BNSD时，必须传入[KV_N,1,D]；输出为BSND时，必须传入[KV_N, D]；pertoken模式下，必须传入[B,KV_S]，S需要大于等于blockTable的第二维*blockSize；
  • 仅支持KV分离；
  • 仅支持高性能模式；
  • 当MTP等于0时，支持sparse_mode=0且不传mask；当MTP大于0、小于16时，支持sparse_mode=3且传入优化后的atten_mask矩阵，atten_mask矩阵shape必须传入（2048*2048）；
  • 不支持配置dequant_offset_key和dequant_offset_value;
  • 不支持配置query_rope和key_rope；
  • 不支持左padding、tensorlist、pse、prefix、后量化；
  • num_query_heads与num_key_value_heads支持组合有(10, 1)、(64, 8)、(80, 8)、(128, 16)。

• learnable_sink的参数约束如下：

  • 仅支持TND、NTD_TND；
  • 仅支持value的d小于等于128；

• 当Q_S大于1时：

  • query、key、value输入，功能使用限制如下：

    • 支持B轴小于等于65536，D轴32byte不对齐时仅支持到128。

    • 支持N轴小于等于256，支持D轴小于等于512；input_layout为BSH或者BSND时，建议N*D小于65535。

    • S支持小于等于20971520（20M）。部分长序列场景下，如果计算量过大可能会导致PFA算子执行超时（aicore error类型报错，errorStr为timeout or trap error），此场景下建议做S切分处理（注：这里计算量会受B、S、N、D等的影响，值越大计算量越大），典型的会超时的长序列（即B、S、N、D的乘积较大）场景包括但不限于：

      • B=1，Q_N=20，Q_S=2097152，D=256，KV_N=1，KV_S=2097152。
      • B=1，Q_N=2，Q_S=20971520，D=256，KV_N=2，KV_S=20971520。
      • B=20，Q_N=1，Q_S=2097152，D=256，KV_N=1，KV_S=2097152。
      • B=1，Q_N=10，Q_S=2097152，D=512，KV_N=1，KV_S=2097152。

    • query、key、value输入类型包含int8时，D轴需要32对齐；输入类型全为[object Object]、[object Object]时，D轴需16对齐。

    • D轴限制：

      • [object Object]Atlas A2 训练系列产品/Atlas A2 推理系列产品[object Object]、[object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]：query、key、value输入类型包含[object Object]时，D轴需要32对齐；query、key、value或attentionOut类型包含[object Object]时，D轴需要64对齐；输入类型全为[object Object]、[object Object]时，D轴需16对齐。

  • actual_seq_qlen：

    [object Object]Atlas A2 训练系列产品/Atlas A2 推理系列产品[object Object]、[object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]：该入参中每个batch的有效Sequence Length应该不大于query中对应batch的Sequence Length。seqlen的传入长度为1时，每个Batch使用相同seqlen；传入长度大于等于Batch时取seqlen的前Batch个数。其他长度不支持。当query的input_layout为TND/NTD_TND时，综合约束请见。

  • actual_seq_kvlen：

    [object Object]Atlas A2 训练系列产品/Atlas A2 推理系列产品[object Object]、[object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]：该入参中每个batch的有效Sequence Length应该不大于key/value中对应batch的Sequence Length。seqlenKv的传入长度为1时，每个Batch使用相同seqlenKv；传入长度大于等于Batch时取seqlenKv的前Batch个数。其他长度不支持。当key/value的input_layout为TND/NTD_TND时，综合约束请见。

  • 参数sparse_mode当前仅支持值为0、1、2、3、4的场景，取其它值时会报错。

    • sparse_mode=0时，atten_mask如果为None，则忽略入参pre_tokens、next_tokens（内部赋值为INT_MAX）。
    • sparse_mode=2、3、4时，atten_mask的shape需要为(S, S)或(1, S, S)或(1, 1, S, S)，其中S的值需要固定为2048，且需要用户保证传入的atten_mask为下三角，不传入atten_mask或者传入的shape不正确报错。
    • sparse_mode=1、2、3的场景忽略入参pre_tokens、next_tokens并按照相关规则赋值。

  • page attention场景：

    • page attention的使能必要条件是block_table存在且有效，同时key、value是按照block_table中的索引在一片连续内存中排布，支持key、value数据类型为[object Object]、[object Object]。在该场景下key、value的input_layout参数无效。block_table中填充的是blockid，当前不会对blockid的合法性进行校验，需用户自行保证。

    • block_size是用户自定义的参数，该参数的取值会影响page attention的性能，在使能page attention场景下，block_size最小为128，最大为512，且要求是128的倍数。通常情况下，page attention可以提高吞吐量，但会带来性能上的下降。

    • page attention场景下，当输入kv cache排布格式为（blocknum, blocksize, H），且KV_N*D超过65535时，受硬件指令约束，会被拦截报错。可通过使能GQA（减小KV_N）或调整kv cache排布格式为（blocknum, KV_N, blocksize, D）解决。当query的input_layout为BNSD、TND时，kv cache排布支持（blocknum, blocksize, H）和（blocknum, KV_N, blocksize, D）两种格式，当query的input_layout为BSH、BSND时，kv cache排布只支持（blocknum, blocksize, H）一种格式。blocknum不能小于根据actual_seq_kvlen和blockSize计算的每个batch的block数量之和。且key和value的shape需保证一致。

    • page attention不支持伪量化场景，不支持tensorlist场景。

    • page attention场景下，必须传入actual_seq_kvlen。

    • page attention场景下，block_table必须为二维，第一维长度需等于B，第二维长度不能小于maxBlockNumPerSeq（maxBlockNumPerSeq为不同batch中最大actual_seq_kvlen对应的block数量）。

    • page attention场景下，支持两种格式和float32/bfloat16，不支持输入query为int8的场景。

    • page attention使能场景下，以下场景输入需满足KV_S>=maxBlockNumPerSeq*blockSize：

      • 传入atten_mask时，如mask shape为（B, 1, Q_S, KV_S）。
      • 传入pse_shift时，如pse_shift shape为（B, Q_N, Q_S, KV_S）。

  • 入参quant_scale_out和quant_offset_out支持pertensor、perchannel量化，支持float32、bfloat16类型。若传入quant_offset_out，需保证其类型和shape信息与quant_scale_out一致。当输入为bfloat16时，同时支持float32和bfloat16，否则仅支持float32。perchannel场景下，当输出layout为BSH时，要求quant_scale_out所有维度的乘积等于H；其他layout要求乘积等于Q_N*D。当输出layout为BSH时，quant_scale_out shape建议传入(1, 1, H)或(H,)；当输出layout为BNSD时，建议传入(1, Q_N, 1, D)或(Q_N, D)；当输出为BSND时，建议传入(1, 1, Q_N, D)或(Q_N, D)。

  • 输出为int8，quant_scale_out和quant_offset_out为perchannel时，暂不支持Ring Attention或者D非32Byte对齐的场景。

  • 输出为int8时，暂不支持sparse为band且preTokens/nextTokens为负数。

  • pse_shift功能使用限制如下：

    • 支持query数据类型为float16、bfloat16、int8场景下使用该功能。
    • query、key、value数据类型为float16且pse_shift存在时，强制走高精度模式，对应的限制继承自高精度模式的限制。
    • Q_S需大于等于query的S长度，KV_S需大于等于key的S长度。

  • 输出为int8，入参quant_offset_out传入非None和非空tensor值，并且sparse_mode、pre_tokens和next_tokens满足以下条件，矩阵会存在某几行不参与计算的情况，导致计算结果误差，该场景会拦截：

    • sparse_mode=0，atten_mask如果非None，每个batch actual_seq_qlen-actual_seq_kvlen-pre_tokens>0或next_tokens<0时，满足拦截条件。
    • sparse_mode=1或2，不会出现满足拦截条件的情况。
    • sparse_mode=3，每个batch actual_seq_kvlen-actual_seq_qlen<0，满足拦截条件。
    • sparse_mode=4，pre_tokens<0或每个batch next_tokens+actual_seq_kvlen-actual_seq_qlen<0时，满足拦截条件。

  • kv伪量化参数分离：

    • 当伪量化参数和KV分离量化参数同时传入时，以KV分离量化参数为准。

    • key_quant_mode和value_quant_mode取值需要保持一致。

    • dequant_scale_key和dequant_scale_value要么都为空，要么都不为空；dequant_offset_key和dequant_offset_value要么都为空，要么都不为空。

    • dequant_scale_key和dequant_scale_value都不为空时，其shape需要保持一致；dequant_offset_key和dequant_offset_value都不为空时，其shape需要保持一致。

    • [object Object]Atlas A2 训练系列产品/Atlas A2 推理系列产品[object Object]、[object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]：

      • 仅支持pertoken和perchannel模式，pertoken模式下要求两个参数的shape均为(B, KV_S)，数据类型固定为float32；perchannel模式下要求两个参数的shape为（KV_N, D），(KV_N, D)，(H)，数据类型固定为bfloat16,H为KV_N*D。
      • dequant_scale_key与dequant_scale_value非空场景，要求query的s小于等于16；要求query的dtype为bfloat16，key、value的dtype为int8，输出的dtype为bfloat16；不支持tensorlist、page attention特性。

    • 管理scale/offset的量化模式如下：

      [object Object]

      [object Object][object Object]

      [object Object]

• 当Q_S等于1时：

  • query、key、value输入，功能使用限制如下：

    • 支持B轴小于等于65536，支持N轴小于等于256，支持S轴小于等于262144，支持D轴小于等于512。
    • query、key、value输入类型均为int8的场景暂不支持。
    • 在int4（int32）伪量化场景下，PyTorch入图调用仅支持KV int4拼接成int32输入（建议通过dynamicQuant生成int4格式的数据，因为dynamicQuant就是一个int32包括8个int4）。
    • 在int4（int32）伪量化场景下，若KV int4拼接成int32输入，那么KV的N、D或者H是实际值的八分之一。并且，int4伪量化仅支持D 64对齐（int32支持D 8对齐）。

  • actual_seq_qlen：

    • [object Object]Atlas A2 训练系列产品/Atlas A2 推理系列产品[object Object]、[object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]：当query的input_layout不为TND时，Q_S为1时该参数无效。当query的input_layout为TND/TND_NTD时，综合约束请见。

  • actual_seq_kvlen：

    • [object Object]Atlas A2 训练系列产品/Atlas A2 推理系列产品[object Object]、[object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]：该入参中每个batch的有效Sequence Length应该不大于key/value中对应batch的Sequence Length。seqlenKv的传入长度为1时，每个Batch使用相同seqlenKv；传入长度大于等于Batch时取seqlenKv的前Batch个数。其他长度不支持。当key/value的input_layout为TND/TND_NTD时，综合约束请见。

  • page attention场景：

    • 使能必要条件是block_table存在且有效，同时key、value是按照block_table中的索引在一片连续内存中排布，在该场景下key、value的input_layout参数无效。

    • [object Object]Atlas A2 训练系列产品/Atlas A2 推理系列产品[object Object]、[object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]：

      • 支持key、value数据类型为float16、bfloat16、int8。
      • 不支持Q为[object Object]、float16、key、value为int4（int32）的场景。

    • 该场景下，block_size是用户自定义的参数，该参数的取值会影响page attention的性能。key、value输入类型为float16、bfloat16时需要16对齐，key、value输入类型为int8时需要32对齐，推荐使用128。通常情况下，page attention可以提高吞吐量，但会带来性能上的下降。

    • 参数key、value各自对应tensor的shape所有维度相乘不能超过int32的表示范围。

    • page attention场景下，blockTable必须为二维，第一维长度需等于B，第二维长度不能小于maxBlockNumPerSeq（maxBlockNumPerSeq为不同batch中最大actual_seq_kvlen对应的block数量）。

    • page attention场景下，当query的input_layout为BNSD、TND时，kv cache排布支持（blocknum, blocksize, H）和（blocknum, KV_N, blocksize, D）两种格式，当query的input_layout为BSH、BSND时，kv cache排布只支持（blocknum, blocksize, H）一种格式。blocknum不能小于根据actual_seq_kvlen和blockSize计算的每个batch的block数量之和。且key和value的shape需保证一致。

    • page attention场景下，kv cache排布为（blocknum, KV_N, blocksize, D）时性能通常优于kv cache排布为（blocknum, blocksize, H）时的性能，建议优先选择（blocknum, KV_N, blocksize, D）格式。

    • page attention使能场景下，当输入kv cache排布格式为（blocknum, blocksize, H），且numKvHeads * headDim 超过64k时，受硬件指令约束，会被拦截报错。可通过使能GQA（减小 numKvHeads）或调整kv cache排布格式为（blocknum, numKvHeads, blocksize, D）解决。

    • page attention场景的参数key、value各自对应tensor的shape所有维度相乘不能超过int32的表示范围。

  • kv伪量化参数分离：

    • 除了key_quant_mode为0并且value_quant_mode为1的场景外，key_quant_mode和value_quant_mode取值需要保持一致。

    • dequant_scale_key和dequant_scale_value要么都为空，要么都不为空；dequant_offset_key和dequant_offset_value要么都为空，要么都不为空。

    • dequant_scale_key和dequant_scale_value都不为空时，除了key_quant_mode为0并且value_quant_mode为1的场景外，其shape需要保持一致；dequant_offset_key和dequant_offset_value都不为空时，除了key_quant_mode为0并且value_quant_mode为1的场景外，其shape需要保持一致。

    • int4（int32）伪量化场景不支持后量化。

    • 管理scale/offset的量化模式如下：

      [object Object]

      [object Object][object Object]

      [object Object]
    [object Object]

  • pse_shift功能使用限制如下：

    • pse_shift数据类型需与query数据类型保持一致。
    • 仅支持D轴对齐，即D轴可以被16整除。

[object Object]

• 单算子模式调用

  [object Object]

• 图模式调用

  [object Object]