From ffd88593251e5f357b799cc4eaaa41f4c486649f Mon Sep 17 00:00:00 2001 From: l00949320 Date: Fri, 12 Sep 2025 11:33:07 +0800 Subject: [PATCH 1/5] Fix torch_npu_npu.npu_fused_infer_attention_score doc --- codegen/templates/_op_plugin_docs.py | 10 +++++++++- .../torch_npu-npu_fused_infer_attention_score.md | 13 +++++++++++-- 2 files changed, 20 insertions(+), 3 deletions(-) diff --git a/codegen/templates/_op_plugin_docs.py b/codegen/templates/_op_plugin_docs.py index 3adf4be27..ec50ce6f7 100644 --- a/codegen/templates/_op_plugin_docs.py +++ b/codegen/templates/_op_plugin_docs.py @@ -6578,6 +6578,7 @@ page attention场景: 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可以提高吞吐量, 但会带来性能上的下降. +Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件：key、value输入类型为float16/bfloat16时需要16对齐；key、value 输入类型为int8时需要32对齐，推荐使用128. 参数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)的场景. @@ -6585,10 +6586,17 @@ page attention场景下, blockTable必须为二维, 第一维长度需等于B, 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场景. +page attention场景下，必须传入actualSeqLengthsKv. +page attentionbu不支持tensorlist场景，不支持左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的使能场景下，以下场景输入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 30a657f65..6c99eeb1f 100644 --- a/docs/context/torch_npu-npu_fused_infer_attention_score.md +++ b/docs/context/torch_npu-npu_fused_infer_attention_score.md @@ -388,14 +388,23 @@ torch_npu.npu_fused_infer_attention_score(query, key, value, *, pse_shift=None, - 不支持Q为`bfloat16`、`float16`、`key`、`value`为`int4`（`int32`）的场景。 - 该场景下，`block_size`是用户自定义的参数，该参数的取值会影响page attention的性能。`key`、`value`输入类型为`float16`、`bfloat16`时需要16对齐，`key`、`value`输入类型为`int8`时需要32对齐，推荐使用128。通常情况下，page attention可以提高吞吐量，但会带来性能上的下降。 + - Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件：`key`、`value`输入类型为`float16`/`bfloat16`时需要16对齐；`key`、`value` 输入类型为`int8`时需要32对齐，推荐使用128。 - 参数`key`、`value`各自对应tensor的shape所有维度相乘不能超过`int32`的表示范围。 - page attention场景下，`block_table`必须为二维，第一维长度需等于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`和`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 attentionbu不支持tensorlist场景，不支持左padding场景。 + - Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件：不支持Q为`bf16`/`fp16`、KV为`int4`（`int32`）的场景。 + - page attention场景下，必须传入`actualSeqLengthsKv`。 - page attention场景的参数`key`、`value`各自对应tensor的shape所有维度相乘不能超过`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。 -- Gitee From 163af47c7f29b857a95b0b4f4bb68830c45df4b1 Mon Sep 17 00:00:00 2001 From: l00949320 Date: Mon, 15 Sep 2025 09:10:25 +0800 Subject: [PATCH 2/5] =?UTF-8?q?=E5=88=A0=E9=99=A4=E9=87=8D=E5=A4=8D?= =?UTF-8?q?=E6=8F=8F=E8=BF=B0?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- docs/context/torch_npu-npu_fused_infer_attention_score.md | 1 - 1 file changed, 1 deletion(-) 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 6c99eeb1f..0fcf6ccc7 100644 --- a/docs/context/torch_npu-npu_fused_infer_attention_score.md +++ b/docs/context/torch_npu-npu_fused_infer_attention_score.md @@ -388,7 +388,6 @@ torch_npu.npu_fused_infer_attention_score(query, key, value, *, pse_shift=None, - 不支持Q为`bfloat16`、`float16`、`key`、`value`为`int4`（`int32`）的场景。 - 该场景下，`block_size`是用户自定义的参数，该参数的取值会影响page attention的性能。`key`、`value`输入类型为`float16`、`bfloat16`时需要16对齐，`key`、`value`输入类型为`int8`时需要32对齐，推荐使用128。通常情况下，page attention可以提高吞吐量，但会带来性能上的下降。 - - Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件：`key`、`value`输入类型为`float16`/`bfloat16`时需要16对齐；`key`、`value` 输入类型为`int8`时需要32对齐，推荐使用128。 - 参数`key`、`value`各自对应tensor的shape所有维度相乘不能超过`int32`的表示范围。 - page attention场景下，`block_table`必须为二维，第一维长度需等于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`和`block_size`计算的每个batch的block数量之和。且`key`和`value`的shape需保证一致。 -- Gitee From 7b20406ad7d52571a5629c47a04a37849f9ff711 Mon Sep 17 00:00:00 2001 From: l00949320 Date: Mon, 15 Sep 2025 09:33:42 +0800 Subject: [PATCH 3/5] =?UTF-8?q?=E5=88=A0=E9=99=A4=E6=8E=A5=E5=8F=A3?= =?UTF-8?q?=E4=B8=AD=E5=8E=9F=E6=9D=A5=E7=9A=84=E9=87=8D=E5=A4=8D=E6=8F=8F?= =?UTF-8?q?=E8=BF=B0=EF=BC=8C=E5=92=8CAclnn=E5=AF=B9=E9=BD=90?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- codegen/templates/_op_plugin_docs.py | 4 ---- docs/context/torch_npu-npu_fused_infer_attention_score.md | 2 -- 2 files changed, 6 deletions(-) diff --git a/codegen/templates/_op_plugin_docs.py b/codegen/templates/_op_plugin_docs.py index ec50ce6f7..396e6c4fe 100644 --- a/codegen/templates/_op_plugin_docs.py +++ b/codegen/templates/_op_plugin_docs.py @@ -6578,7 +6578,6 @@ page attention场景: 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可以提高吞吐量, 但会带来性能上的下降. -Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件：key、value输入类型为float16/bfloat16时需要16对齐；key、value 输入类型为int8时需要32对齐，推荐使用128. 参数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)的场景. @@ -6588,9 +6587,6 @@ page attention场景下, kv cache排布为(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场景下，必须传入actualSeqLengthsKv. page attentionbu不支持tensorlist场景，不支持左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的使能场景下，以下场景输入S需要大于等于blockTable的第二维 * blockSize: 使能Attention mask: 如mask shape为(B, 1, 1, S). 使能pseShift: 如pseShift shape为(B, N, 1, S). 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 0fcf6ccc7..4aadde699 100644 --- a/docs/context/torch_npu-npu_fused_infer_attention_score.md +++ b/docs/context/torch_npu-npu_fused_infer_attention_score.md @@ -394,9 +394,7 @@ torch_npu.npu_fused_infer_attention_score(query, key, value, *, pse_shift=None, - 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 attentionbu不支持tensorlist场景，不支持左padding场景。 - - Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件：不支持Q为`bf16`/`fp16`、KV为`int4`（`int32`）的场景。 - page attention场景下，必须传入`actualSeqLengthsKv`。 - - page attention场景的参数`key`、`value`各自对应tensor的shape所有维度相乘不能超过`int32`的表示范围。 - page attention的使能场景下，以下场景输入S需要大于等于blockTable的第二维 * blockSize。 - 使能Attention mask，如`mask` shape为 \(B, 1, 1, S\)。 - 使能`pseShift`，如`pseShift` shape为\(B, N, 1, S\)。 -- Gitee From e85cb1aed3e92f3e18e859debe877e9b2f990fe0 Mon Sep 17 00:00:00 2001 From: l00949320 Date: Mon, 15 Sep 2025 11:49:22 +0800 Subject: [PATCH 4/5] =?UTF-8?q?=E8=B7=9Faclnn=E5=AF=B9=E9=BD=90?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- codegen/templates/_op_plugin_docs.py | 8 +++----- .../torch_npu-npu_fused_infer_attention_score.md | 12 +++++------- 2 files changed, 8 insertions(+), 12 deletions(-) diff --git a/codegen/templates/_op_plugin_docs.py b/codegen/templates/_op_plugin_docs.py index 396e6c4fe..77340df29 100644 --- a/codegen/templates/_op_plugin_docs.py +++ b/codegen/templates/_op_plugin_docs.py @@ -6576,17 +6576,15 @@ 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场景下，必须传入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). 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 4aadde699..c5f8f1a92 100644 --- a/docs/context/torch_npu-npu_fused_infer_attention_score.md +++ b/docs/context/torch_npu-npu_fused_infer_attention_score.md @@ -383,18 +383,16 @@ 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可以提高吞吐量，但会带来性能上的下降。 + - Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件：`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 attentionbu不支持tensorlist场景，不支持左padding场景。 + - Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件：不支持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\)。 -- Gitee From d88f7203fc1ffff1853b261572358d97134941dc Mon Sep 17 00:00:00 2001 From: l00949320 Date: Tue, 23 Sep 2025 10:38:02 +0800 Subject: [PATCH 5/5] =?UTF-8?q?=E5=88=A0=E9=99=A4=E9=87=8D=E5=A4=8D?= =?UTF-8?q?=E4=BA=A7=E5=93=81=E5=88=97=E4=B8=BE?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- docs/context/torch_npu-npu_fused_infer_attention_score.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) 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 c5f8f1a92..94ea6cc3f 100644 --- a/docs/context/torch_npu-npu_fused_infer_attention_score.md +++ b/docs/context/torch_npu-npu_fused_infer_attention_score.md @@ -385,12 +385,12 @@ torch_npu.npu_fused_infer_attention_score(query, key, value, *, pse_shift=None, - 使能必要条件是`block_table`存在且有效，同时`key`、`value`是按照`block_table`中的索引在一片连续内存中排布，在该场景下`key`、`value`的`input_layout`参数无效。 - 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可以提高吞吐量，但会带来性能上的下降。 - - Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件：`key`、`value`输入类型为`float16`/`bfloat16`时需要16对齐，`key`、`value`输入类型为`int8`时需要32对齐，推荐使用128。 + - `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 attentionbu不支持tensorlist场景，不支持左padding场景。 - - Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件：不支持Q为`bfloat16`/`float16`、KV为`int4`（`int32`）的场景。 + - 不支持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。 -- Gitee