开发者
下载
[object Object]

[object Object][object Object]undefined
[object Object]
  • 接口功能:该接口为AI CPU算子接口,是aclnnMixedQuantSparseFlashMla算子的前置算子接口。根据aclnnMixedQuantSparseFlashMla算子接口的输入信息,计算并输出负载均衡结果。输出结果可以作为aclnnMixedQuantSparseFlashMla算子接口的输入,减少aclnnMixedQuantSparseFlashMla算子接口的执行耗时。

    该算子不建议单独使用,建议与aclnnMixedQuantSparseFlashMla算子配合使用,形成完整的工作流。

    1. 接受aclnnMixedQuantSparseFlashMla算子接口输入数据shape信息,包含batchSize、qSeqlen、kSeqlen、mask。通过对输入分块并模拟计算耗时,均匀分配分块到可用核上,以降低aclnnMixedQuantSparseFlashMla算子的整体计算耗时,并提高硬件利用率。
    2. 分配结果输出后,后续作为输入供aclnnMixedQuantSparseFlashMla算子使用。
    3. 分配结果包含每个AIC核基本块的起始点和终止点,已经每个AIV核的FD任务信息。详细内容可以参考
[object Object]

每个算子分为,必须先调用"aclnnMixedQuantSparseFlashMlaMetadataGetWorkspaceSize"获取workspace大小,在调用"aclnnMixedQuantSparseFlashMlaMetadata"执行计算

[object Object]
[object Object]
[object Object]
  • 参数说明

    [object Object]
  • 返回值:

    返回aclnnStatus状态码,具体参见

    第一段接口完成入参校验,出现以下场景时报错:

    [object Object]
[object Object]
  • 参数说明:

    [object Object]
  • 返回值:

    返回aclnnStatus状态码,具体参见

[object Object]
  • aclnnMixedQuantSparseFlashMlaMetadata默认确定性实现。
  • B(Batch)表示输入样本批量大小。
  • Batch取值规则
    • 优先获取sequsedQOptional中的Batch信息。
    • 如果未传入sequsedQOptional,且layoutQOptional为TND和传入了cuSeqlensQOptional,则获取cuSeqlensQOptional中的Batch信息。
    • 除上所述,使用batchSize。
  • Query Sequence Length取值规则
    • 优先获取sequsedQOptional中的Sequence Length信息。
    • 如果未传入sequsedQOptional,且layoutQOptional为TND和传入了cuSeqlensQOptional,则获取cuSeqlensQOptional中的Sequence Length信息。
    • 除上所述,使用maxSeqlenQ。
    • ori_kv、cmp_kv Sequence Length与Query的获取规则一致。
  • BSND场景
    • 当传入的layoutQOptional为"BSND"时,在未传入sequsedQOptional的情况下,必传maxSeqlenQ参数。
    • 当传入的layoutKvOptional为"BSND"时,若hasOriKv为true,在未传入sequsedOriKvOptional的情况下,必传maxSeqlenOriKv参数;若hasCmpKv为true,在未传入sequsedCmpKvOptional的情况下,必传maxSeqlenCmpKv参数。
  • TND场景
    • 当传入的layoutQOptional为"TND"时,必传cuSeqlensQOptional参数。
    • 当传入的layoutKvOptional为"TND"时,若hasOriKv为true,必传cuSeqlensOriKvOptional;若hasCmpKv为true,必传cuSeqlensCmpKvOptional参数。
[object Object]

示例代码如下,仅供参考,具体编译和执行过程请参考

[object Object]