昇腾社区首页
开发者
下载
支持
积分兑换 NEW
文档中心
sparse_flash_mla_grad
[object Object][object Object]
  • [object Object]Ascend 950PR/Ascend 950DT[object Object]：支持
  • [object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]：支持
  • [object Object]Atlas A2 训练系列产品/Atlas A2 推理系列产品[object Object]：支持
  • [object Object]Atlas 200I/500 A2 推理产品[object Object]：不支持
  • [object Object]Atlas 推理系列产品[object Object]：不支持
  • [object Object]Atlas 训练系列产品[object Object]：不支持
[object Object]

  • API功能：

    • [object Object]：计算[object Object]训练场景下注意力的反向输出，支持Sliding Window Attention、Compressed Attention以及Sparse Compressed Attention。
    • [object Object]：接口用于生成一个任务列表，包含每个AIcore的Attention计算任务的起止点的Batch、Head、以及Q和K的分块的索引，供后续[object Object]算子使用。

  • 计算公式：

    阶段一：根据不同cmp_ratio场景，对输入ori_kv与cmp_kv进行选择

    • 当cmp_ratio = 1 (SWA)：
    selectedKv = orikvselectedKv\text{ }=\text{ }orikv
    • 当cmp_ratio = 4 (SCFA)：
    selectedKv =concat(oriKv, Gather(cmpkv,topkIndices[i])), 0 <=i< selectBlockCountselectedKv\text{ }=concat(oriKv, \text{ }Gather \left( cmpkv,topkIndices \left[ i \left] \left)) ,\text{ }0\text{ } < =i < \text{ }selectBlockCount\right. \right. \right. \right.
    • else (CFA):
    selectedKv =concat(oriKv, cmpkv)selectedKv\text{ }=concat(oriKv, \text{ }cmpkv)

    阶段二：计算P、dP、dS

    P=SimpleSoftmax(Mask(Q @ selectedKvT⋅ scale),lse)P = SimpleSoftmax(Mask(Q \text{ }@\text{ } selectedKv^{{T}} \cdot \text{ } scale), lse) dP=dO @ selectedKvTdP = dO \text{ }@\text{ } selectedKv^{{T}} dS=P×(dP − SoftmaxGrad(dO,O))dS = P \times (dP\text{ } -\text{ } SoftmaxGrad(dO, O))

    阶段三：计算dQ, dKV, dSinks

    dQ=dS @ selectedKv ⋅ scaledQ = dS \text{ } @ \text{ } selectedKv \text{ } \cdot \text{ } scale dKV=dST @ Q ⋅ scale+PT@ dOdKV = dS^{{T}} \text{ } @ \text{ } Q \text{ } \cdot \text{ } scale + P^{{T}} @ \text{ } dO dSinks=ReduceSum(−P × dP × SimpleSoftmax(sinks,lse),dim=−1)dSinks = ReduceSum(-P \text{ }\times\text{ } dP \text{ }\times\text{ } SimpleSoftmax(sinks, lse), dim=-1)
[object Object]
[object Object]
[object Object]
[object Object][object Object]

  • q（[object Object]）：必选参数，对应公式中的QQ。[object Object]="BSND" 时 shape 为 [object Object]；[object Object]="TND" 时 shape 为 [object Object]。B：支持泛化；S1：支持泛化；D：512；T1：B × S1。支持非连续，数据格式支持ND，数据类型支持[object Object]和[object Object]。

  • dout（[object Object]）：必选参数，注意力正向输出矩阵的梯度，对应公式中的dOdO。数据类型和 shape 均与 q 保持一致。支持非连续，数据格式支持ND，数据类型支持[object Object]和[object Object]。

  • attn_out（[object Object]）：必选参数，注意力正向输出矩阵，对应公式中的OO。数据类型和 shape 均与 q 保持一致。支持非连续，数据格式支持ND，数据类型支持[object Object]和[object Object]。

  • softmax_lse（[object Object]）：必选参数，注意力正向计算的输出 lse。[object Object]="BSND" 时 shape 为 [object Object]；[object Object]="TND" 时 shape 为 [object Object]。B：与 q 的 B 保持一致；N2：1；S1：与 q 的 S1 保持一致；G：N1/N2；T1：B × S1。支持非连续，数据格式支持ND，数据类型支持[object Object]。

  • ori_kv（[object Object]）：可选参数，对应公式中的oriKvoriKv。[object Object]="BSND" 时 shape 为 [object Object]；[object Object]="TND" 时 shape 为 [object Object]。B：与 q 的 B 保持一致；S2：支持泛化；N2：1；D：512；T2：B × S2。支持非连续，数据格式支持ND，数据类型支持[object Object]和[object Object]。

  • cmp_kv（[object Object]）：可选参数，对应公式中的cmpkvcmpkv。[object Object]="BSND" 时 shape 为 [object Object]；[object Object]="TND" 时 shape 为 [object Object]。B：与 q 的 B 保持一致；S3：支持泛化；N2：1；D：512；T3：B × S3。传 None 时按 SWA 场景计算。支持非连续，数据格式支持ND，数据类型支持[object Object]和[object Object]。

  • ori_sparse_indices（[object Object]）：可选参数，对应 oriKv 部分的 topk 索引。[object Object]="BSND" 时 shape 为 [object Object]；[object Object]="TND" 时 shape 为 [object Object]。B：与 q 的 B 保持一致；S1：与 q 的 S1 保持一致；N2：1；K1：支持泛化；T1：B × S1。支持非连续，数据格式支持ND，数据类型支持[object Object]。

  • cmp_sparse_indices（[object Object]）：可选参数，对应公式中的topkIndicestopkIndices。[object Object]="BSND" 时 shape 为 [object Object]；[object Object]="TND" 时 shape 为 [object Object]。B：与 q 的 B 保持一致；S1：与 q 的 S1 保持一致；N2：1；K2：支持泛化；T1：B × S1。若 cmp_kv 不为 None，此时 cmp_sparse_indices 不为 None 时按 SCFA 场景计算，为 None 时按 CFA 场景计算；若 cmp_kv 为 None，则 cmp_sparse_indices 只能为 None，此时按 SWA 场景计算。支持非连续，数据格式支持ND，数据类型支持[object Object]。

  • cu_seqlens_q（[object Object]）：可选参数，每个 Batch 中 q 的有效 token 数的累加和形式。[object Object]="TND" 时必传。shape 为 [object Object]，累加和与 T1 保持一致。支持非连续，数据格式支持ND，数据类型支持[object Object]。

  • cu_seqlens_ori_kv（[object Object]）：可选参数，每个 Batch 中 ori_kv 的有效 token 数的累加和形式。[object Object]="TND" 时必传。shape 为 [object Object]，累加和与 T2 保持一致。支持非连续，数据格式支持ND，数据类型支持[object Object]。

  • cu_seqlens_cmp_kv（[object Object]）：可选参数，每个 Batch 中 cmp_kv 的有效 token 数的累加和形式。[object Object]="TND" 时必传。shape 为 [object Object]，累加和与 T3 保持一致。支持非连续，数据格式支持ND，数据类型支持[object Object]。

  • seqused_q（[object Object]）：可选参数，表示不同 batch 中 query 实际参与运算的 token 数。shape 为 [object Object]。支持非连续，数据格式支持ND，数据类型支持[object Object]。

  • seqused_ori_kv（[object Object]）：可选参数，表示不同 batch 中 ori_kv 实际参与运算的 token 数。shape 为 [object Object]。支持非连续，数据格式支持ND，数据类型支持[object Object]。

  • seqused_cmp_kv（[object Object]）：可选参数，表示不同 batch 中 cmp_kv 实际参与运算的 token 数。shape 为 [object Object]。支持非连续，数据格式支持ND，数据类型支持[object Object]。

  • cmp_residual_kv（[object Object]）：可选参数，表示每个 batch S2 // cmpRatio 后的余数。shape 为 [object Object]。当 cmp_kv 不为空且 cmp_mask_mode=3 时必须传入。支持非连续，数据格式支持ND，数据类型支持[object Object]。

  • ori_topk_length（[object Object]）：可选参数，表示每行 query 对应的 ori_kv 实际可选的 topk 长度。shape 为 [object Object]（BSND）或 [object Object]（TND）。当 ori_mask_mode=0 且 ori_sparse_indices 不为 None 时必须传入且必须为准确值。支持非连续，数据格式支持ND，数据类型支持[object Object]。

  • cmp_topk_length（[object Object]）：可选参数，表示每行 query 对应的 cmp_kv 实际可选的 topk 长度。shape 为 [object Object]（BSND）或 [object Object]（TND）。当 cmp_mask_mode=0 且 cmp_sparse_indices 不为 None 时必须传入且必须为准确值。支持非连续，数据格式支持ND，数据类型支持[object Object]。

  • sinks（[object Object]）：可选参数，注意力下沉 tensor。shape 为 [object Object]。支持非连续，数据格式支持ND，数据类型支持[object Object]。

  • metadata（[object Object]）：可选参数，表示 tiling 下沉的 aicpu 算子输出结果。支持非连续，数据格式支持ND，数据类型支持[object Object]。

  • softmax_scale（[object Object]）：可选参数，代表缩放系数。数据类型支持[object Object]，默认值：1.0 / sqrt(D)。

  • cmp_ratio（[object Object]）：可选参数，代表压缩率，取值范围 1~128。数据类型支持[object Object]，默认值：1。

  • ori_mask_mode（[object Object]）：可选参数，q 和 ori_kv 计算的 mask 模式。模式 0 为不做 mask 操作；模式 3 为 rightDownCausal；模式 4 为 band（滑窗，起点右下角）。数据类型支持[object Object]。

  • cmp_mask_mode（[object Object]）：可选参数，q 和 cmp_kv 计算的 mask 模式。模式 0 为不做 mask 操作；模式 3 为 rightDownCausal。数据类型支持[object Object]。

  • ori_win_left（[object Object]）：可选参数，q 和 ori_kv 计算中 q 对过去 token 计算的数量。当前仅支持取值 127。数据类型支持[object Object]。

  • ori_win_right（[object Object]）：可选参数，q 和 ori_kv 计算中 q 对未来 token 计算的数量。当前仅支持取值 0。数据类型支持[object Object]。

  • layout_q（[object Object]）：可选参数，q 的数据排布格式。支持 "BSND"、"TND"。数据类型支持[object Object]。

  • layout_kv（[object Object]）：可选参数，ori_kv、cmp_kv 的数据排布格式。支持 "BSND"、"TND"，当前必须与 layout_q 保持一致。数据类型支持[object Object]。

  • [object Object]Atlas A2 训练系列产品/Atlas A2 推理系列产品[object Object]、[object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]：暂不支持 sequsedQOptional、sequsedOriKvOptional、sequsedCmpKvOptional、oriTopkLengthOptional、cmpTopkLengthOptional、metadataOptional 字段。

[object Object][object Object]undefined
  • [object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]、[object Object]Atlas A2 训练系列产品/Atlas A2 推理系列产品[object Object]：不支持seqused_q、seqused_ori_kv、seqused_cmp_kv、ori_topk_length、cmp_topk_length，ori_mask_mode仅支持4，cmp_mask_mode仅支持3，ori_win_left仅支持127，ori_win_right仅支持0。
[object Object][object Object]

  • dq（[object Object]）：对应公式中的dQdQ，支持非连续，数据格式支持ND，数据类型、shape与输入q保持一致。

  • dori_kv（[object Object]）：可选输出，表示输入ori_kv的梯度，支持非连续，数据格式支持ND，数据类型、shape与输入ori_kv保持一致。

  • dcmp_kv（[object Object]）：可选输出，表示输入cmp_kv的梯度，支持非连续，数据格式支持ND，数据类型、shape与输入cmp_kv保持一致；当cmp_kv为None时，dcmp_kv也为None。

  • dsinks（[object Object]）：可选输出，表示输入sinks的梯度，支持非连续，数据格式支持ND，数据类型支持[object Object]，shape与输入sinks保持一致。

  • ori_softmax_l1norm（[object Object]）：可选输出，表示q与ori_kv计算得出的softmax的L1Norm结果，公式为reduceG(softmax)/G；数据类型为[object Object]。[object Object]为BSND时shape为[object Object]，当[object Object]为TND时shape为[object Object]。当 ori_sparse_indices 不为 None 时该输出不为空，其他场景下输出为 None。

  • cmp_softmax_l1norm（[object Object]）：可选输出，表示q与cmp_kv计算得出的softmax的L1Norm结果，公式为reduceG(softmax)/G；当 cmp_sparse_indices 不为 None 时该输出不为空，其他场景下输出为 None。

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

  • 该接口支持训练场景下使用。

  • 该接口支持单算子模式和aclgraph模式。

  • 参数q、dout、attn_out、ori_kv、cmp_kv的数据类型必须保持一致。

  • 入参为空的场景处理：q 为空 Tensor 时直接返回。

  • 各个场景关于cmp_kv、cmp_sparse_indices的使用说明如下：

    • SWA场景：要求cmp_kv == None && cmp_sparse_indices == None
    • SCFA场景：要求cmp_kv != None && cmp_sparse_indices != None
    • CFA场景：要求cmp_kv != None && cmp_sparse_indices == None

  • 确定性计算：Ascend 950PR/Ascend 950DT 默认非确定性计算，支持通过 [object Object] 开启确定性计算；Atlas A2/A3不支持开启确定性计算。

  • Mask 模式支持：

    [object Object]undefined

  • 规格约束：

    [object Object]undefined
    • ori_kv/cmp_kv 传 None 的支持情况:
      • [object Object]Atlas A2 训练系列产品/Atlas A2 推理系列产品[object Object]、[object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]：不支持。
      • [object Object]Ascend 950PR/Ascend 950DT[object Object]：支持。
    • ori_sparse_indices 的支持情况:
      • [object Object]Atlas A2 训练系列产品/Atlas A2 推理系列产品[object Object]、[object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]：不支持。
      • [object Object]Ascend 950PR/Ascend 950DT[object Object]：支持。
    • seqused 的支持情况:
      • [object Object]Atlas A2 训练系列产品/Atlas A2 推理系列产品[object Object]、[object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]：不支持（需传 None）。
      • [object Object]Ascend 950PR/Ascend 950DT[object Object]：支持。
    • ori_topk_length / cmp_topk_length 的支持情况:
      • [object Object]Atlas A2 训练系列产品/Atlas A2 推理系列产品[object Object]、[object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]：不支持（需传 None）。
      • [object Object]Ascend 950PR/Ascend 950DT[object Object]：支持。
    • sinks 传 None 的支持情况:
      • [object Object]Atlas A2 训练系列产品/Atlas A2 推理系列产品[object Object]、[object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]：不支持。
      • [object Object]Ascend 950PR/Ascend 950DT[object Object]：支持。
    • metadata 的支持情况:
      • [object Object]Atlas A2 训练系列产品/Atlas A2 推理系列产品[object Object]、[object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]：仅支持传 None。
      • [object Object]Ascend 950PR/Ascend 950DT[object Object]：必须传。
    • ori_mask_mode 的支持情况:
      • [object Object]Atlas A2 训练系列产品/Atlas A2 推理系列产品[object Object]、[object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]：仅支持模式 4。
      • [object Object]Ascend 950PR/Ascend 950DT[object Object]：支持模式 0、3、4。
    • cmp_mask_mode 的支持情况:
      • [object Object]Atlas A2 训练系列产品/Atlas A2 推理系列产品[object Object]、[object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]：仅支持模式 3。
      • [object Object]Ascend 950PR/Ascend 950DT[object Object]：支持模式 0、3。
[object Object]
  • 默认支持确定性计算
[object Object]

  • 单算子模式调用

    [object Object]
我们使用cookie来确保您的高速浏览体验。继续浏览本站，即表示您同意我们使用cookie。查看详情