开发者
下载
[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]
  • 接口功能:

    [object Object]接口用于生成一个任务列表,包含每个AIcore的Attention计算任务的起止点的Batch、Head、以及Q和K的分块的索引,供后续[object Object]算子使用。

    [object Object]接口基于一系列操作得到每一个token对应的top-k个位置。主要计算过程为:

    1. 将某个token对应的输入参数[object Object]QindexRg×dQ_{index}\in\R^{g\times d})乘以给定上下文[object Object]KindexRSk×dK_{index}\in\R^{S_{k}\times d}),得到相关性。
    2. 通过激活函数ReLUReLU过滤无效负相关信号后,得到当前Token与所有前序Token的相关性分数向量。
    3. 将其与权重系数[object Object]WW)相乘后,沿g的方向,选取前TopkTop-k个索引值得到输出sparseIndicessparseIndices,并输出对应的sparseValuessparseValues,作为Attention的输入。
  • 计算公式:

    Topk{[1]1×g@[(W@[1]1×Sk)ReLU(Qindex@KTindex)]}Top-k \left\{ \left[ 1 \left] \mathop{{}}\nolimits_{{1 \times \text{ }g}}\text{@} \left[ \left( W\text{@} \left[ 1 \left] \mathop{{}}\nolimits_{{1\text{ } \times \text{ }S\mathop{{}}\nolimits_{{k}}}} \left) \text{ } \odot \text{ }ReLU \left( Q\mathop{{}}\nolimits_{{index}}\text{@}K\mathop{{}}\nolimits_{{T}}^{{index}} \left) \left] \right\} \right. \right. \right. \right. \right. \right. \right. \right. \right. \right.
[object Object]

调用lightning_indexer接口之前,先调用前置接口lightning_indexer_metadata,完成lightning_indexer负载均衡的计算。

[object Object]
[object Object]
[object Object]
[object Object]
[object Object][object Object]undefined
[object Object][object Object]undefined
  • [object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]、[object Object]Atlas A2 训练系列产品/Atlas A2 推理系列产品[object Object]:
    • topk取值范围当前仅支持[1, 2048],以及3072、4096、5120、6144、7168、8192。
    • 当前不支持sequsedQOptional、outputIdxOffsetOptional、maxSeqlenQ功能,不建议传入这些参数。
    • 当layout_k为PA_BBND时,必须传入sequsedKOptional;当layout_k不为PA_BBND时,不支持sequsedKOptional功能,不建议传入该参数。
[object Object][object Object][object Object]undefined
[object Object][object Object]undefined
[object Object]
  • 该接口支持推理场景下使用。
  • 该接口支持单算子模式和aclgraph图模式调用。
  • lightning_indexer_metadata接口需与lightning_indexer算子配套使用。
  • B(Batch)表示输入样本批量大小。
  • 参数cu_seqlens_q、cu_seqlens_k要求其值为当前Batch与前序Batch有效token数的累加值,后一个元素的值必须大于等于前一个元素的值。
  • 参数seqused_q、seqused_k要求其值表示每个Batch中的有效token数。
  • 参数cmp_residual_k需满足cmp_residual_k[i] < cmp_ratio。
  • mask_mode所表示的mask模式的详细介绍见
  • 参数q的N支持1~64,k的N支持1,headdim支持128。
  • pa_kv_cache支持0轴非连续;pa_block_size支持1~1024,满足block大小32B对齐。
  • 参数q、k的数据类型应保持一致。
  • sparse_indices无效部分填-1;sparse_values无效部分填-inf。
  • [object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]、[object Object]Atlas A2 训练系列产品/Atlas A2 推理系列产品[object Object]:
    • topk取值范围当前仅支持[1, 2048],以及3072、4096、5120、6144、7168、8192。
    • 当前不支持sequsedQOptional、outputIdxOffsetOptional、maxSeqlenQ功能,不建议传入这些参数。
    • 当layout_k为PA_BBND时,必须传入sequsedKOptional;当layout_k不为PA_BBND时,不支持sequsedKOptional功能,不建议传入该参数。
  • [object Object]Ascend 950PR/Ascend 950DT[object Object]:
    • 当layout_q为BSND时,不支持传入cuSeqlensQOptional;当layout_k为BSND或PA_BBND时,不支持传入cuSeqlensKOptional。
    • 当传入outputIdxOffsetOptional时,只支持大于0的索引偏移值;且应满足约束:加上传入的索引偏移值后,得到的sparseIndice值不超过INT32的最大值。
    • 当layout_q为TND时,必须传入cuSeqlensQOptional,如果也传入sequsedQOptional,应保证由sequsedQOptional传入的各个batch的query长度不超过根据cuSeqlensQOptional计算出的各个batch的q序列长度。当某个batch由sequsedQOptional传入的q序列长度seqlen1小于由cuSeqlensQOptional计算出的query长度seqlen2时,会启用TND Padding功能,将该batch的从seqlen1 + 1到seqlen2的query输出的sparseIndices和sparseValues全部置为无效值。
    • 当传入的cmpRatio > 1且maskMode = 3时,必须传入cmpResidualKOptional,其余情况不传入。
[object Object]
  • 默认支持确定性计算。
[object Object]
  • 单算子模式调用:

    [object Object]
  • aclgraph图模式调用:

    [object Object]