diff --git a/codegen/templates/_op_plugin_docs.py b/codegen/templates/_op_plugin_docs.py index 3adf4be2702b43e37b6bd6b4576cdd95a4540ed6..77340df2924997a6242ac9dba0796c122b4b9405 100644 --- a/codegen/templates/_op_plugin_docs.py +++ b/codegen/templates/_op_plugin_docs.py @@ -6576,19 +6576,21 @@ actual_seq_lengths_kv: 该参数应为非负数, 在input_layout不同时, 其 page attention场景: 使能必要条件是block_table存在且有效, 同时key、value是按照block_table中的索引在一片连续内存中排布, 在该场景下key、value的input_layout参数无效. Atlas A2 训练系列产品/Atlas 800I A2 推理产品: 支持key、value数据类型为float16、bfloat16、int8. -Atlas A3 训练系列产品: 支持key、value数据类型为float16、bfloat16、int8. -该场景下, block_size是用户自定义的参数, 该参数的取值会影响page attention的性能. key、value输入类型为float16、bfloat16时需要16对齐, key、value输入类型为int8时需要32对齐, 推荐使用128. 通常情况下, page attention可以提高吞吐量, 但会带来性能上的下降. -参数key、value各自对应tensor的shape所有维度相乘不能超过int32的表示范围. -Atlas A2 训练系列产品/Atlas 800I A2 推理产品: 不支持Q为bfloat16、float16、key、value为int4(int32)的场景. -Atlas A3 训练系列产品: 不支持Q为bfloat16、float16、key、value为int4(int32)的场景. +该场景下, block_size是用户自定义的参数, 该参数的取值会影响page attention的性能,在使能page attention场景下，blocksize需要传入非0值，且blocksize最大不超过512. 通常情况下, page attention可以提高吞吐量, 但会带来性能上的下降. +Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件：key、value输入类型为float16、bfloat16时需要16对齐，key、value输入类型为int8时需要32对齐，推荐使用128。 page attention场景下, blockTable必须为二维, 第一维长度需等于B, 第二维长度不能小于maxBlockNumPerSeq(maxBlockNumPerSeq为不同batch中最大actual_seq_lengths_kv对应的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_lengths_kv和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不支持左padding场景. -Atlas A2 训练系列产品/Atlas 800I A2 推理产品: 不支持Q为BF16/FP16且KV为INT4(INT32)的场景. -Atlas A3 训练系列产品: 不支持Q为BF16/FP16且KV为INT4(INT32)的场景. -page attention场景的参数key、value各自对应tensor的shape所有维度相乘不能超过int32的表示范围. +page attention场景下，必须传入actualSeqLengthsKv. +page attentionbu不支持tensorlist场景，不支持左padding场景. +Atlas A2 训练系列产品/Atlas 800I A2 推理产品: 不支持Q为bfloat16、float16、key、value为int4(int32)的场景. +page attention的使能场景下，以下场景输入S需要大于等于blockTable的第二维 * blockSize: +使能Attention mask: 如mask shape为(B, 1, 1, S). +使能pseShift: 如pseShift shape为(B, N, 1, S). +使能伪量化per-token模式：输入参数antiquantScale和antiquantOffset的shape均为(2, B, S). +使能per-token叠加per-head模式：两个参数的shape均为(B, N, S)，数据类型固定为float32，当key、value数据类型为int8、int4(int32)时支持. +使能per-token-group模式：antiquantScale的shape为(1, B, N, S, D/32), 数据类型固定为FLOAT8_E8M0，不支持带antiquantOffset。当key、value数据类型为FLOAT4_E1M2、FLOAT4_E2M1时支持. kv左padding场景: kvCache的搬运起点计算公式为: Smax-kv_padding_size-actual_seq_lengths. kvCache的搬运终点计算公式为: Smax-kv_padding_size. 其中kvCache的搬运起点或终点小于0时, 返回数据结果为全0. kv_padding_size小于0时将被置为0. diff --git a/docs/context/torch_npu-npu_fused_infer_attention_score.md b/docs/context/torch_npu-npu_fused_infer_attention_score.md index 30a657f653891b065ef224951fc2023524c08e8b..94ea6cc3f45c7fa74204482b6c5c96829de50165 100644 --- a/docs/context/torch_npu-npu_fused_infer_attention_score.md +++ b/docs/context/torch_npu-npu_fused_infer_attention_score.md @@ -383,19 +383,23 @@ torch_npu.npu_fused_infer_attention_score(query, key, value, *, pse_shift=None, - page attention场景： - 使能必要条件是`block_table`存在且有效，同时`key`、`value`是按照`block_table`中的索引在一片连续内存中排布，在该场景下`key`、`value`的`input_layout`参数无效。 - - Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件、Atlas A3 训练系列产品/Atlas A3 推理系列产品： - - 支持`key`、`value`数据类型为`float16`、`bfloat16`、`int8`。 - - 不支持Q为`bfloat16`、`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场景下，`block_table`必须为二维，第一维长度需等于B，第二维长度不能小于maxBlockNumPerSeq（maxBlockNumPerSeq为不同batch中最大`actual_seq_lengths_kv`对应的block数量）。 + - Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件：支持`key`、`value`数据类型为`float16`、`bfloat16`、`int8`。 + - 该场景下，`block_size`是用户自定义的参数，该参数的取值会影响page attention的性能，在使能page attention场景下，blocksize需要传入非0值，且blocksize最大不超过512。通常情况下，page attention可以提高吞吐量，但会带来性能上的下降。 + - `key`、`value`输入类型为`float16`/`bfloat16`时需要16对齐，`key`、`value`输入类型为`int8`时需要32对齐，推荐使用128。 - 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_lengths_kv`和`block_size`计算的每个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不支持左padding场景。 - - page attention场景的参数`key`、`value`各自对应tensor的shape所有维度相乘不能超过`int32`的表示范围。 - + - page attentionbu不支持tensorlist场景，不支持左padding场景。 + - 不支持Q为`bfloat16`/`float16`、KV为`int4`（`int32`）的场景。 + - page attention场景下，必须传入`actualSeqLengthsKv`。 + - page attention场景下，`block_table`必须为二维，第一维长度需等于B，第二维长度不能小于maxBlockNumPerSeq（maxBlockNumPerSeq为不同batch中最大`actual_seq_lengths_kv`对应的block数量）。 + - page attention的使能场景下，以下场景输入S需要大于等于blockTable的第二维 * blockSize。 + - 使能Attention mask，如`mask` shape为 \(B, 1, 1, S\)。 + - 使能`pseShift`，如`pseShift` shape为\(B, N, 1, S\)。 + - 使能伪量化`per-token`模式：输入参数`antiquantScale`和`antiquantOffset`的shape均为\(2, B, S\)。 + - 使能`per-token`叠加`per-head`模式：两个参数的shape均为\(B, N, S\)，数据类型固定为`float32`，当`key`、`value`数据类型为`int8`、`int4`\(`int32`\)时支持。 + - 使能`per-token-group`模式：`antiquantScale`的shape为\(1, B, N, S, D/32\), 数据类型固定为FLOAT8_E8M0，不支持带`antiquantOffset`。当`key`、`value`数据类型为FLOAT4_E1M2、FLOAT4_E2M1时支持。 + - kv左padding场景： - kvCache的搬运起点计算公式为：Smax-kv\_padding\_size-actual\_seq\_lengths。kvCache的搬运终点计算公式为：Smax-kv\_padding\_size。其中kvCache的搬运起点或终点小于0时，返回数据结果为全0。 - `kv_padding_size`小于0时将被置为0。