diff --git a/xllm_ops/attention/quant_lightning_indexer_v2/CMakeLists.txt b/xllm_ops/attention/quant_lightning_indexer_v2/CMakeLists.txt new file mode 100644 index 0000000..a3ab34a --- /dev/null +++ b/xllm_ops/attention/quant_lightning_indexer_v2/CMakeLists.txt @@ -0,0 +1,19 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +file(GLOB CURRENT_DIRS RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*) +if(NOT ENABLE_TEST AND NOT BENCHMARK) + list(REMOVE_ITEM CURRENT_DIRS tests) +endif() +foreach(SUB_DIR ${CURRENT_DIRS}) + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/${SUB_DIR}/CMakeLists.txt") + add_subdirectory(${SUB_DIR}) + endif() +endforeach() \ No newline at end of file diff --git a/xllm_ops/attention/quant_lightning_indexer_v2/README.md b/xllm_ops/attention/quant_lightning_indexer_v2/README.md new file mode 100644 index 0000000..00a3964 --- /dev/null +++ b/xllm_ops/attention/quant_lightning_indexer_v2/README.md @@ -0,0 +1,78 @@ +# QuantLightningIndexerV2 + +## 产品支持情况 + +| 产品 | 是否支持 | +| ------------------------------------------------------------ | :------: | +|Ascend 950PR/Ascend 950DT| √ | +|Atlas A3 训练系列产品/Atlas A3 推理系列产品| √ | +|Atlas A2 训练系列产品/Atlas A2 推理系列产品| √ | +|Atlas 200I/500 A2 推理产品| × | +|Atlas 推理系列产品| × | +|Atlas 训练系列产品| × | + +## 功能说明 + +- API功能：QuantLightningIndexerV2是推理场景下，稀疏attention前处理的计算，选出关键的稀疏token，并对输入query和key进行量化实现存8算8，获取最大收益。 + +- 计算公式： + $$out = \text{Top-}k\left\{[1]_{1\times g}@\left[(W@[1]_{1\times S_{k}})\odot\text{ReLU}\left(\left(Scale_Q@Scale_K^T\right)\odot\left(Q_{index}^{Quant}@{\left(K_{index}^{Quant}\right)}^T\right)\right)\right]\right\}$$ + 主要计算过程为： + 1. 将某个token对应的输入参数`query`（$Q_{index}^{Quant}\in\R^{g\times d}$）乘以给定上下文`key`（$K_{index}^{Quant}\in\R^{S_{k}\times d}$），得到相关性。 + 2. 相关性结果与`query`和`key`对应的反量化系数`query_dequant_scale`（$Scale_Q$）和`key_dequant_scale`（$Scale_K^T$）相乘，通过激活函数$ReLU$过滤无效负相关信号后，得到当前Token与所有前序Token的相关性分数向量。 + 3. 将其与权重系数`weights`（$W$）相乘后，沿g的方向，选取前$Top-k$个索引值得到输出$out$，作为Attention的输入。 + +## 参数说明 + +| 参数名 | 输入/输出/属性 | 描述 | 数据类型 | 数据格式 | +|----------------------------|-----------|----------------------------------------------------------------------|----------------|------------| +| query | 输入 | 公式中的$Q_{index}^{Quant}\in\R^{g\times d}$，表示输入Index Query，不支持非连续。| INT8、FLOAT8_e4m3fn、HIFLOAT8、FLOAT4_e2m1 | ND | +| key | 输入 | 公式中的$K_{index}^{Quant}\in\R^{S_{k}\times d}$，表示压缩后的输入Index Key，支持0轴非连续。| INT8、FLOAT8_e4m3fn、HIFLOAT8、FLOAT4_e2m1 | ND | +| weights | 输入 | 公式中的$W$，表示权重系数，不支持非连续。 | FLOAT16、FLOAT32 | ND | +| query_dequant_scale | 输入 | 公式中的$Scale_Q$，表示Index Query的反量化系数，不支持非连续。`quant_mode`为3/5时，shape为将`query`的D轴替换为(D/64, 2)；`quant_mode`为4时，shape为(1,)；其他场景shape与weights一致 | FLOAT16、FLOAT32、FLOAT8_e8m0 | ND | +| key_dequant_scale | 输入 | 公式中的$Scale_K$，表示Index Key的反量化系数，支持0轴非连续。`quant_mode`为3/5时，shape为将`key`的D轴替换为(D/64, 2)；`quant_mode`为4时，shape为(1,)；其他场景shape为移除`key`的D轴 | FLOAT16、FLOAT32、FLOAT8_e8m0 | ND | +| cu_seqlens_q | 可选输入 | layout_q为TND时必须传入，表示每个Batch中`query`的有效token数前缀和。；layout_q为BSND时不能传入 | INT32 | ND | +| cu_seqlens_k | 可选输入 | layout_k为TND时必须传入，表示每个Batch中`key`的有效token数前缀和；layout_k为PA_BSND或BSND时不能传入 | INT32 | ND | +| seqused_q | 可选输入 | layout_q为BSND时可选传入，表示每个Batch中`query`的有效token数 | INT32 | ND | +| seqused_k | 可选输入 | layout_k为PA_BSND或BSND时使用，表示每个Batch中`key`的有效token数。| INT32 | ND | +| cmp_residual_k | 可选输入 | 压缩场景下Key的残余长度，需满足0 \<= cmp_residual_k\[i\] \< cmp_ratio。| INT32 | ND | +| block_table | 可选输入 | 表示PageAttention中KV存储使用的block映射表。 | INT32 | ND | +| output_idx_offset | 可选输入 | 输出索引的偏移量 | INT32 | ND | +| metadata | 可选输入 | QuantLightningIndexerV2Metadata算子传入的分核信息，包含使用核数、分块大小以及每个核处理数据的起始点等内容。 | INT32 | ND | +| quant_mode | 属性 | 用于标识输入的量化模式。 | INT32 | - | +| max_seqlen_q | 可选属性| Query的最大序列长度，默认值-1表示任意可能长度 | INT32 | - | +| layout_q | 可选属性| 用于标识输入`query`的数据排布格式，默认值"BSND"。 | STRING | - | +| layout_k | 可选属性 | 用于标识输入`key`的数据排布格式，默认值"BSND"。| STRING | - | +| topk | 属性 | 代表topK阶段需要保留的索引数量，默认值2048。 | INT32 | - | +| mask_mode | 可选属性 | 表示mask的模式，默认值0。 | INT32 | - | +| cmp_ratio | 可选属性 | 用于稀疏计算，表示key的压缩倍数，默认值1。 | INT32 | - | +| return_value | 可选属性 | 表示是否输出`sparse_values`，默认值0。 | INT32 | - | +| sparse_indices | 输出 | 公式中的输出Out，参与稀疏attention计算的token索引值。 | INT32 | ND | +| sparse_values | 输出 | 公式中的Indices输出对应的value值。`return_value`为1时shape与`sparse_indices`一致，`return_value`为0时shape为(0,) | BFLOAT16 | ND | + +## 约束说明 + +- Ascend 950PR/Ascend 950DT： + - `query`、`key`在`quant_mode`为1/3时支持FLOAT8_e4m3fn，`quant_mode`为2时支持INT8，`quant_mode`为4时支持HIFLOAT8，`quant_mode`为5时支持FLOAT4_e2m1。 + - `query_dequant_scale`和`key_dequant_scale`在`quant_mode`为1/4时支持FLOAT32，`quant_mode`为2时支持FLOAT16，`quant_mode`为3/5时支持FLOAT8_e8m0。 + - `weights`在`quant_mode`为2时支持FLOAT16，`quant_mode`为1/3/4/5时支持FLOAT32。 + - `quant_mode`为3/5时，`query_dequant_scale`和`key_dequant_scale`的维数分别比`query`和`key`多1，前缀维度保持一致，末两维为(D/64,2)；D必须为64的倍数，每个scale对应连续32个D轴逻辑元素。 + - `query`的N支持[1, 64]，`key`的N仅支持1。 + - `topk`支持[1, 8192]。 + - 当传入的参数layout_query为TND时，必须传入cu_seqlens_q，如果也传入seqused_q，应保证由seqused_q传入的各个batch的query长度不超过根据cu_seqlens_q计算出的各个batch的q序列长度。当某个batch由seqused_q传入的q序列长度seqlen1小于由cu_seqlens_q计算出的query长度seqlen2时，会启用TND Padding功能，将该batch的seqlen2与seqlen1差值部分的query输出的sparse_indices和sparse_values全部置为无效值。部分长序列场景下，如果需要填充的无效数据过多，由于硬件限制可能会导致aicore执行超时，可以通过(seqlen2 - seqlen1) * topk来计算需要填充的数据量，建议将这个数据量控制在4亿以内。 + +- Atlas A3 训练系列产品/Atlas A3 推理系列产品、Atlas A2 训练系列产品/Atlas A2 推理系列产品： + - `quant_mode`仅支持2。 + - `query`、`key`支持INT8，不支持FLOAT8_e4m3fn、HIFLOAT8和FLOAT4_e2m1。 + - `query_dequant_scale`和`key_dequant_scale`支持FLOAT16，不支持FLOAT32和FLOAT8_e8m0。 + - `weights`支持FLOAT16，不支持FLOAT32。 + - 不支持`output_idx_offset`和`return_value`。 + - `query`的N仅支持64，`key`的N仅支持1。 + - `topk`支持[1, 2048]。 + +## 调用示例 + +| 调用方式 | 调用样例 | 说明 | +|----------|----------|------| +| PyTorch API | - | 通过[torch.ops.cann_ops_transformer.quant_lightning_indexer](../../torch_extension/cann_ops_transformer/docs/zh/quant_lightning_indexer.md)接口调用QuantLightningIndexerV2算子。 | +| aclnn API | [test_aclnn_quant_lightning_indexer_v2](examples/test_aclnn_quant_lightning_indexer_v2.cpp) | 通过[aclnnQuantLightningIndexerV2](docs/aclnnQuantLightningIndexerV2.md)两段式接口调用QuantLightningIndexerV2算子。 | diff --git a/xllm_ops/attention/quant_lightning_indexer_v2/docs/aclnnQuantLightningIndexerV2.md b/xllm_ops/attention/quant_lightning_indexer_v2/docs/aclnnQuantLightningIndexerV2.md new file mode 100644 index 0000000..dbdbb8f --- /dev/null +++ b/xllm_ops/attention/quant_lightning_indexer_v2/docs/aclnnQuantLightningIndexerV2.md @@ -0,0 +1,898 @@ +# aclnnQuantLightningIndexerV2 + +[📄 查看源码](https://gitcode.com/cann/ops-transformer/tree/master/attention/quant_lightning_indexer_v2) + +## 产品支持情况 + + +- Ascend 950PR/Ascend 950DT：支持 + + +- Atlas A3 训练系列产品/Atlas A3 推理系列产品：支持 + + +- Atlas A2 训练系列产品/Atlas A2 推理系列产品：支持 + + +- Atlas 200I/500 A2 推理产品：不支持 + + +- Atlas 推理系列产品：不支持 + + +- Atlas 训练系列产品：不支持 + + +## 功能说明 + +- 接口功能：`QuantLightningIndexerV2`是推理场景下，稀疏attention前处理的计算，选出关键的稀疏token，并对输入query和key进行量化实现存8算8，获取最大收益。 + +- 版本演进：在QuantLightningIndexer的基础上，新增压缩key场景、分核计算metadata、稀疏value输出等能力。 + +- 计算公式： + +$$ +out = \text{Top-}k\left\{[1]_{1\times g}@\left[(W@[1]_{1\times S_{k}})\odot\text{ReLU}\left(\left(Scale_Q@Scale_K^T\right)\odot\left(Q_{index}^{Quant}@{\left(K_{index}^{Quant}\right)}^T\right)\right)\right]\right\} +$$ + +主要计算过程为： + +1. 将某个token对应的输入参数`query`（$Q_{index}^{Quant}\in\R^{g\times d}$）乘以给定上下文`key`（$K_{index}^{Quant}\in\R^{S_{k}\times d}$），得到相关性。 +2. 相关性结果与`query`和`key`对应的反量化系数`query_dequant_scale`（$Scale_Q$）和`key_dequant_scale`（$Scale_K^T$）相乘，通过激活函数$ReLU$过滤无效负相关信号后，得到当前Token与所有前序Token的相关性分数向量。 +3. 将其与权重系数`weights`（$W$）相乘后，沿g的方向，选取前$Top-k$个索引值得到输出$out$，作为Attention的输入。 + +## 函数原型 + +每个算子分为[两段式接口](../../../docs/zh/context/two_phase_api.md)，必须先调用"aclnnQuantLightningIndexerV2GetWorkspaceSize"接口获取计算所需workspace大小以及包含了算子计算流程的执行器，再调用"aclnnQuantLightningIndexerV2"接口执行计算。 + +```Cpp +aclnnStatus aclnnQuantLightningIndexerV2GetWorkspaceSize( + const aclTensor *query, + const aclTensor *key, + const aclTensor *weights, + const aclTensor *queryDequantScale, + const aclTensor *keyDequantScale, + const aclTensor *cuSeqLensQOptional, + const aclTensor *cuSeqLensKOptional, + const aclTensor *sequsedQOptional, + const aclTensor *sequsedKOptional, + const aclTensor *cmpResidualKOptional, + const aclTensor *blockTableOptional, + const aclTensor *outputIdxOffsetOptional, + const aclTensor *metadataOptional, + int64_t topk, + int64_t quantMode, + int64_t maxSeqlenQOptional, + char *layoutQOptional, + char *layoutKOptional, + int64_t maskModeOptional, + int64_t cmpRatioOptional, + int64_t returnValueOptional, + const aclTensor *sparseIndicesOut, + const aclTensor *sparseValuesOut, + uint64_t *workspaceSize, + aclOpExecutor **executor) +``` + +```Cpp +aclnnStatus aclnnQuantLightningIndexerV2( + void *workspace, + uint64_t workspaceSize, + aclOpExecutor *executor, + const aclrtStream stream) +``` + +## aclnnQuantLightningIndexerV2GetWorkspaceSize + +- **参数说明：** + +> [!NOTE] +> +> - query、key、weights参数维度含义：B（Batch Size）表示输入样本批量大小、S（Sequence Length）表示输入样本序列长度、H（Head Size）表示hidden层的大小、N（Head Num）表示多头数、D（Head Dim）表示hidden层最小的单元尺寸，且满足D=H/N、T表示所有Batch输入样本序列长度的累加和。 +> - S1表示query shape中的S，S2表示key shape中的S，T1表示query shape中的T，N1表示query shape中的N，N2表示key shape中的N。 +> - maxBlockNumPerSeq表示每个Batch中最大sequsedK对应的block数量，S2_MAX表示sequsedK中的最大值 + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +
参数名输入/输出描述使用说明数据类型数据格式维度(shape)非连续Tensor
query输入公式中量化后的 Query。不支持空tensor。INT8、FLOAT8_e4m3fn、HIFLOAT8、FLOAT4_e2m1ND +
    +
  • layout_query为BSND时，shape为(B,S1,N1,D)。
    • +
  • layout_query为TND时，shape为(T1,N1,D)。
    • +
+ 		x
key输入公式中量化后的 Key。 +
    +
  • 不支持空tensor。
    • +
  • block_num为PageAttention时block总数，block_size为一个block的token数。
    • +
  • layout_key为PA_BSND时，shape为(block_num, block_size, N2, D)。
    • +
  • layout_key为BSND时，shape为(B, K_S, N2, D)，layout_key为TND时，shape为(K_T, N2, D)。
    • +
+ 		INT8、FLOAT8_e4m3fn、HIFLOAT8、FLOAT4_e2m1ND +
    +
  • layout_key为PA_BSND时，shape为(block_num, block_size, N2, D)。
    • +
+ 		支持0轴非连续
weights输入公式中的权重系数 W。不支持空tensor。FLOAT16、FLOAT32ND +
    +
  • layout_query为BSND时，shape为(B,S1,N1)。
    • +
  • layout_query为TND时，shape为(T1,N1)。
    • +
+ 		x
queryDequantScale输入公式中 Query 的反量化系数。不支持空tensor。FLOAT16、FLOAT32、FLOAT8_e8m0ND +
    +
  • quantMode为3/5时，layout_query为BSND时shape为(B,S1,N1,D/64,2)，layout_query为TND时shape为(T1,N1,D/64,2)。
    • +
  • quantMode为4时，shape为(1,)。
    • +
  • 其他场景shape与weights保持一致。
    • +
+ 		x
keyDequantScale输入公式中 Key 的反量化系数。不支持空tensor。FLOAT16、FLOAT32、FLOAT8_e8m0ND +
    +
  • quantMode为3/5时，layout_key为PA_BSND、BSND、TND对应的shape分别为(block_num,block_size,N2,D/64,2)、(B,K_S,N2,D/64,2)、(K_T,N2,D/64,2)。
    • +
  • quantMode为4时，shape为(1,)。
    • +
  • 其他场景下，layout_key为PA_BSND、BSND、TND对应的shape分别为(block_num,block_size,N2)、(B,K_S,N2)、(K_T,N2)。
    • +
+ 		支持0轴非连续
cuSeqLensQOptional输入每个Batch中，Query的有效token数（TND场景使用cu_seqlens格式）。 +
    +
  • 当layout_query为TND时，该入参必须传入，且以该入参元素的数量作为B值，该入参中每个元素的值表示当前batch与之前所有batch的token数总和，即前缀和。
    • +
+ 		INT32ND(B+1,)x
cuSeqLensKOptional输入每个Batch中，Key的有效token数（TND场景使用cu_seqlens格式）。 +
    +
  • 当layout_key为TND时，该入参必须传入。
    • +
+ 		INT32ND(B+1,)x
sequsedQOptional输入每个Batch中，Query的有效token数（BSND场景使用seqused格式）。该入参中每个Batch的有效token数不超过query中的维度S大小且不小于0。INT32ND(B,)x
sequsedKOptional输入每个Batch中，Key的有效token数（BSND场景使用seqused格式）。 +
    +
  • 该入参中每个Batch的有效token数不超过key中的维度S大小且不小于0。
    • +
  • 当layout_key为PA_BSND时，该入参必须传入。
    • +
+ 		INT32ND(B,)x
cmpResidualKOptional输入压缩场景下Key的残余长度。需满足0 <= cmpResidualKOptional[i] < cmpRatioOptional。INT32ND(B,)x
blockTableOptional输入表示PageAttention中KV存储使用的block映射表。 +
    +
  • 不支持空tensor。
    • +
  • PageAttention场景下，block_table必须为二维，第一维长度需要等于B，第二维长度不能小于maxBlockNumPerSeq。
    • +
+ 		INT32ND(B, S2_MAX/block_size)x
outputIdxOffsetOptional输入输出索引的偏移量。-INT32NDlayout_query为BSND时shape为(B,S1,N2)，layout_query为TND时shape为(T1,N2)。x
metadataOptional输入QuantLightningIndexerV2Metadata算子传入的分核信息。 +
    +
  • 包含使用核数、分块大小以及每个核处理数据的起始点等内容。
    • +
  • shape大小为[1024]，当前不支持传空。
    • +
+ 		INT32ND(1024,)x
topk输入topK阶段需要保留的Key token索引数量。支持[1, 8192]。INT64---
quantMode输入量化模式。 +
    +
  • 支持传入 1（FLOAT8_e4m3fn量化）、2（Per-Token-Head量化）、3（MXFP8量化）、4（HIFLOAT8量化）、5（MXFP4量化）。
    • +
+ 		INT64---
maxSeqlenQOptional输入Query的最大序列长度。-INT64---
layoutQOptional输入用于标识输入Query的数据排布格式。 +
    +
  • 支持BSND、TND。
    • +
+ 		STRING---
layoutKOptional输入用于标识输入Key的数据排布格式。 +
    +
  • 支持 PA_BSND、BSND、TND。
    • +
+ 		STRING---
maskModeOptional输入表示sparse的模式。 +
    +
  • 0代表defaultMask模式。
    • +
  • 3代表rightDownCausal模式的mask，对应以右顶点为划分的下三角场景。
    • +
+ 		INT64---
cmpRatioOptional输入key的压缩倍数。 +
    +
  • 支持 (0, 128] 内的正整数。
    • +
+ 		INT64---
returnValueOptional输入表示是否输出sparseValuesOut。 +
    +
  • 1表示输出，0表示不输出。
    • +
+ 		INT64---
sparseIndicesOut输出公式中的Indices输出。不支持空tensor。INT32ND +
    +
  • layout_query为"BSND"时输出shape为[B, S1, N2, topk]。
    • +
  • layout_query为"TND"时输出shape为[T1, N2, topk]。
    • +
+ 		x
sparseValuesOut输出公式中的Indices输出对应的value值。 +
    +
  • returnValue为1时输出有效值，无效部分填bf16负无穷；returnValue为0时输出shape为(0,)的空tensor。
    • +
+ 		BFLOAT16NDreturnValue为1时shape与sparseIndicesOut保持一致；returnValue为0时shape为(0,)。x
workspaceSize输出返回需要在Device侧申请的workspace大小。-----
executor输出返回op执行器，包含了算子计算流程。-----
+ + +- Ascend 950PR/Ascend 950DT： + - `layout_key` 额外支持 BSND 和 TND；支持 PA_BSND、BSND、TND。 + - `quant_mode` 支持 1（FLOAT8_e4m3fn量化）、2（INT8量化）、3（MXFP8量化）、4（HIFLOAT8量化）和 5（MXFP4量化）。 + - `cmp_ratio` 支持 (0, 128] 内任意正整数。 + - 支持 `return_value`。 + - query 和 key：`quant_mode` 为 1/3 时支持 FLOAT8_e4m3fn，`quant_mode` 为 2 时支持 INT8，`quant_mode` 为 4 时支持 HIFLOAT8，`quant_mode` 为 5 时支持 FLOAT4_e2m1。 + - query_dequant_scale 和 key_dequant_scale：`quant_mode` 为 1/4 时支持 FLOAT32，`quant_mode` 为 2 时支持 FLOAT16，`quant_mode` 为 3/5 时支持 FLOAT8_e8m0。 + - weights：`quant_mode` 为 2 时支持 FLOAT16，`quant_mode` 为 1/3/4/5 时支持 FLOAT32。 + - query Q_N 支持 [1, 64]。 + + +- Atlas A3 训练系列产品/Atlas A3 推理系列产品、Atlas A2 训练系列产品/Atlas A2 推理系列产品： + - `layout_key` 仅支持 PA_BSND。 + - `quant_mode` 仅支持 2（Per-Token-Head量化）。 + - `cmp_ratio` 仅支持 2 的幂次方且范围为 [1, 128]，即 1/2/4/8/16/32/64/128。 + - 不支持 `outputIdxOffsetOptional`。 + - 不支持 `return_value`。 + - query 和 key：支持 INT8，不支持 FLOAT8_e4m3fn、HIFLOAT8 和 FLOAT4_e2m1。 + - query_dequant_scale 和 key_dequant_scale：支持 FLOAT16，不支持 FLOAT32 和 FLOAT8_e8m0。 + - weights：支持 FLOAT16，不支持 FLOAT32。 + - query Q_N 仅支持 64。 + - topk 仅支持 [1, 2048]。 + + +- **返回值：** + + aclnnStatus：返回状态码，具体参见[aclnn返回码](../../../docs/zh/context/aclnn_return_code.md)。 + + 第一段接口会完成入参校验，出现以下场景时报错： + + + + + + + + + + + + + + + + + + + + + + + +
返回值错误码描述
ACLNN_ERR_PARAM_NULLPTR161001如果传入参数是必选输入，输出或者必选属性，且是空指针，则返回161001。
ACLNN_ERR_PARAM_INVALID161002query、key、weights、queryDequantScale、keyDequantScale、cuSeqLensQOptional、cuSeqLensKOptional、sequsedQOptional、sequsedKOptional、cmpResidualKOptional、blockTableOptional、metadataOptional、layoutQOptional、layoutKOptional、topk、quantMode、maskModeOptional、cmpRatioOptional、returnValueOptional、sparseIndicesOut、sparseValuesOut的数据类型和数据格式不在支持的范围内。
+ +## aclnnQuantLightningIndexerV2 + +- **参数说明：** + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +
参数名输入/输出描述
workspace输入在Device侧申请的workspace内存地址。
workspaceSize输入在Device侧申请的workspace大小，由第一段接口aclnnQuantLightningIndexerV2GetWorkspaceSize获取。
executor输入op执行器，包含了算子计算流程。
stream输入指定执行任务的Stream。
+ +- **返回值：** + + aclnnStatus：返回状态码，具体参见[aclnn返回码](../../../docs/zh/context/aclnn_return_code.md)。 + +## 约束说明 + +- headdim 支持 128。 +- block_size 取值为 16 的倍数，最大支持 1024。 +- 当 `layout_key` 不为 PA_BSND 时，`layout_query` 和 `layout_key` 必须一致。 +- 当 `quant_mode` 为 3/5 时，`queryDequantScale` 和 `keyDequantScale` 的维数分别比 `query` 和 `key` 多 1，前缀维度保持一致，末两维为(D/64, 2)；D必须为64的倍数，每个scale对应D轴上连续32个逻辑元素。 +- 当传入的参数layout_query为TND时，必须传入cuSeqlensQOptional，如果也传入sequsedQOptional，应保证由sequsedQOptional传入的各个batch的query长度不超过根据cuSeqlensQOptional计算出的各个batch的q序列长度。当某个batch由sequsedQOptional传入的q序列长度seqlen1小于由cuSeqlensQOptional计算出的query长度seqlen2时，会启用TND Padding功能，将该batch的seqlen2与seqlen1差值部分的query输出的sparseIndices和sparseValues全部置为无效值。部分长序列场景下，如果需要填充的无效数据过多，由于硬件限制可能会导致aicore执行超时，可以通过(seqlen2 - seqlen1) * topk来计算需要填充的数据量，建议将这个数据量控制在4亿以内。 +- **确定性说明：** aclnnQuantLightningIndexerV2 默认确定性实现。 + +## 调用示例 + +示例代码如下，仅供参考，具体编译和执行过程请参考[编译与运行样例](../../../docs/zh/context/compile_and_run_sample.md)。 + +```Cpp +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file test_aclnn_quant_lightning_indexer_v2.cpp + * \brief + */ +#include +#include +#include +#include +#include "securec.h" +#include "acl/acl.h" +#include "aclnnop/aclnn_quant_lightning_indexer_v2.h" +#include "aclnn/opdev/platform.h" + +using namespace std; + +namespace { + +#define CHECK_RET(cond) ((cond) ? true :(false)) + +#define LOG_PRINT(message, ...) \ + do { \ + (void)printf(message, ##__VA_ARGS__); \ + } while (0) + +int64_t GetShapeSize(const std::vector& shape) { + int64_t shapeSize = 1; + for (auto i : shape) { + shapeSize *= i; + } + return shapeSize; +} + +int Init(int32_t deviceId, aclrtStream* stream) { + auto ret = aclInit(nullptr); + if (!CHECK_RET(ret == ACL_SUCCESS)) { + LOG_PRINT("aclInit failed. ERROR: %d\n", ret); + return ret; + } + ret = aclrtSetDevice(deviceId); + if (!CHECK_RET(ret == ACL_SUCCESS)) { + LOG_PRINT("aclrtSetDevice failed. ERROR: %d\n", ret); + return ret; + } + ret = aclrtCreateStream(stream); + if (!CHECK_RET(ret == ACL_SUCCESS)) { + LOG_PRINT("aclrtCreateStream failed. ERROR: %d\n", ret); + return ret; + } + return 0; +} + +template +int CreateAclTensor(const std::vector& hostData, const std::vector& shape, void** deviceAddr, + aclDataType dataType, aclTensor** tensor) { + auto size = GetShapeSize(shape) * sizeof(T); + auto ret = aclrtMalloc(deviceAddr, size, ACL_MEM_MALLOC_HUGE_FIRST); + if (!CHECK_RET(ret == ACL_SUCCESS)) { + LOG_PRINT("aclrtMalloc failed. ERROR: %d\n", ret); + return ret; + } + + ret = aclrtMemcpy(*deviceAddr, size, hostData.data(), size, ACL_MEMCPY_HOST_TO_DEVICE); + if (!CHECK_RET(ret == ACL_SUCCESS)) { + LOG_PRINT("aclrtMemcpy failed. ERROR: %d\n", ret); + return ret; + } + + std::vector strides(shape.size(), 1); + for (int64_t i = shape.size() - 2; i >= 0; i--) { + strides[i] = shape[i + 1] * strides[i + 1]; + } + + *tensor = aclCreateTensor(shape.data(), shape.size(), dataType, strides.data(), 0, aclFormat::ACL_FORMAT_ND, + shape.data(), shape.size(), *deviceAddr); + return 0; +} + +struct TensorResources { + void* queryDeviceAddr = nullptr; + void* keyDeviceAddr = nullptr; + void* weightsDeviceAddr = nullptr; + void* qScaleDeviceAddr = nullptr; + void* kScaleDeviceAddr = nullptr; + void* metadataDeviceAddr = nullptr; + void* sparseIndicesDeviceAddr = nullptr; + void* sparseValuesDeviceAddr = nullptr; + + aclTensor* queryTensor = nullptr; + aclTensor* keyTensor = nullptr; + aclTensor* weightsTensor = nullptr; + aclTensor* qScaleTensor = nullptr; + aclTensor* kScaleTensor = nullptr; + aclTensor* metadataTensor = nullptr; + aclTensor* sparseIndicesTensor = nullptr; + aclTensor* sparseValuesTensor = nullptr; +}; + +int InitializeTensors(TensorResources& resources) { + int64_t B = 2; + int64_t S1 = 4; + int64_t S2 = 8; + int64_t N1 = 64; + int64_t N2 = 1; + int64_t D = 128; + int64_t topk = 512; + + std::vector queryShape = {B, S1, N1, D}; + std::vector keyShape = {B, S2, N2, D}; + std::vector weightsShape = {B, S1, N1}; + std::vector qScaleShape = {B, S1, N1}; + std::vector kScaleShape = {B, S2, N2}; + std::vector metadataShape = {1024}; + std::vector sparseIndicesShape = {B, S1, N2, topk}; + std::vector sparseValuesShape = {B, S1, N2, topk}; + + int64_t queryShapeSize = GetShapeSize(queryShape); + int64_t keyShapeSize = GetShapeSize(keyShape); + int64_t weightsShapeSize = GetShapeSize(weightsShape); + int64_t qScaleShapeSize = GetShapeSize(qScaleShape); + int64_t kScaleShapeSize = GetShapeSize(kScaleShape); + int64_t metadataShapeSize = GetShapeSize(metadataShape); + int64_t sparseIndicesShapeSize = GetShapeSize(sparseIndicesShape); + int64_t sparseValuesShapeSize = GetShapeSize(sparseValuesShape); + + std::vector queryHostData(queryShapeSize, 0x38); + std::vector keyHostData(keyShapeSize, 0x38); + std::vector weightsHostData(weightsShapeSize, 0.01f); + std::vector qScaleHostData(qScaleShapeSize, 1.0f); + std::vector kScaleHostData(kScaleShapeSize, 1.0f); + std::vector metadataHostData(metadataShapeSize, 0); + std::vector sparseIndicesHostData(sparseIndicesShapeSize, 0); + std::vector sparseValuesHostData(sparseValuesShapeSize, 0); + + int ret = CreateAclTensor(queryHostData, queryShape, &resources.queryDeviceAddr, + aclDataType::ACL_FLOAT8_E4M3FN, &resources.queryTensor); + if (!CHECK_RET(ret == ACL_SUCCESS)) { return ret; } + + ret = CreateAclTensor(keyHostData, keyShape, &resources.keyDeviceAddr, + aclDataType::ACL_FLOAT8_E4M3FN, &resources.keyTensor); + if (!CHECK_RET(ret == ACL_SUCCESS)) { return ret; } + + ret = CreateAclTensor(weightsHostData, weightsShape, &resources.weightsDeviceAddr, + aclDataType::ACL_FLOAT, &resources.weightsTensor); + if (!CHECK_RET(ret == ACL_SUCCESS)) { return ret; } + + ret = CreateAclTensor(qScaleHostData, qScaleShape, &resources.qScaleDeviceAddr, + aclDataType::ACL_FLOAT, &resources.qScaleTensor); + if (!CHECK_RET(ret == ACL_SUCCESS)) { return ret; } + + ret = CreateAclTensor(kScaleHostData, kScaleShape, &resources.kScaleDeviceAddr, + aclDataType::ACL_FLOAT, &resources.kScaleTensor); + if (!CHECK_RET(ret == ACL_SUCCESS)) { return ret; } + + ret = CreateAclTensor(metadataHostData, metadataShape, &resources.metadataDeviceAddr, + aclDataType::ACL_INT32, &resources.metadataTensor); + if (!CHECK_RET(ret == ACL_SUCCESS)) { return ret; } + + ret = CreateAclTensor(sparseIndicesHostData, sparseIndicesShape, &resources.sparseIndicesDeviceAddr, + aclDataType::ACL_INT32, &resources.sparseIndicesTensor); + if (!CHECK_RET(ret == ACL_SUCCESS)) { return ret; } + + ret = CreateAclTensor(sparseValuesHostData, sparseValuesShape, &resources.sparseValuesDeviceAddr, + aclDataType::ACL_BF16, &resources.sparseValuesTensor); + if (!CHECK_RET(ret == ACL_SUCCESS)) { return ret; } + + return ACL_SUCCESS; +} + +int ExecuteQuantLightningIndexerV2(TensorResources& resources, aclrtStream stream, + void** workspaceAddr, uint64_t* workspaceSize) { + int64_t topk = 512; + int64_t quantMode = 1; + int64_t maskMode = 0; + int64_t cmpRatio = 1; + int64_t returnValue = 1; + constexpr const char layoutQStr[] = "BSND"; + constexpr const char layoutKStr[] = "BSND"; + constexpr size_t layoutQLen = sizeof(layoutQStr); + constexpr size_t layoutKLen = sizeof(layoutKStr); + char layoutQ[layoutQLen]; + char layoutK[layoutKLen]; + errno_t memcpyRet = memcpy_s(layoutQ, sizeof(layoutQ), layoutQStr, layoutQLen); + if (!CHECK_RET(memcpyRet == 0)) { + LOG_PRINT("memcpy_s layoutQ failed. ERROR: %d\n", memcpyRet); + return -1; + } + memcpyRet = memcpy_s(layoutK, sizeof(layoutK), layoutKStr, layoutKLen); + if (!CHECK_RET(memcpyRet == 0)) { + LOG_PRINT("memcpy_s layoutK failed. ERROR: %d\n", memcpyRet); + return -1; + } + aclOpExecutor* executor; + + int ret = aclnnQuantLightningIndexerV2GetWorkspaceSize( + resources.queryTensor, resources.keyTensor, resources.weightsTensor, + resources.qScaleTensor, resources.kScaleTensor, + nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, + resources.metadataTensor, + topk, quantMode, -1, layoutQ, layoutK, maskMode, cmpRatio, returnValue, + resources.sparseIndicesTensor, resources.sparseValuesTensor, + workspaceSize, &executor); + + if (!CHECK_RET(ret == ACL_SUCCESS)) { + LOG_PRINT("aclnnQuantLightningIndexerV2GetWorkspaceSize failed. ERROR: %d\n", ret); + return ret; + } + + if (*workspaceSize > 0ULL) { + ret = aclrtMalloc(workspaceAddr, *workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST); + if (!CHECK_RET(ret == ACL_SUCCESS)) { + LOG_PRINT("allocate workspace failed. ERROR: %d\n", ret); + return ret; + } + } + + ret = aclnnQuantLightningIndexerV2(*workspaceAddr, *workspaceSize, executor, stream); + if (!CHECK_RET(ret == ACL_SUCCESS)) { + LOG_PRINT("aclnnQuantLightningIndexerV2 failed. ERROR: %d\n", ret); + return ret; + } + + return ACL_SUCCESS; +} + +int PrintOutResult(const std::vector& shape, void* deviceAddr) { + auto size = GetShapeSize(shape); + std::vector resultData(size, 0); + auto ret = aclrtMemcpy(resultData.data(), resultData.size() * sizeof(resultData[0]), + deviceAddr, size * sizeof(resultData[0]), ACL_MEMCPY_DEVICE_TO_HOST); + if (!CHECK_RET(ret == ACL_SUCCESS)) { + LOG_PRINT("copy result from device to host failed. ERROR: %d\n", ret); + return ret; + } + LOG_PRINT("sparse_indices result (first 10 elements):\n"); + for (int64_t i = 0; i < size && i < 10; i++) { + LOG_PRINT(" [%ld] = %d\n", i, resultData[i]); + } + return ACL_SUCCESS; +} + +void CleanupResources(TensorResources& resources, void* workspaceAddr, + aclrtStream stream, int32_t deviceId) { + if (resources.queryTensor) { aclDestroyTensor(resources.queryTensor); } + if (resources.keyTensor) { aclDestroyTensor(resources.keyTensor); } + if (resources.weightsTensor) { aclDestroyTensor(resources.weightsTensor); } + if (resources.qScaleTensor) { aclDestroyTensor(resources.qScaleTensor); } + if (resources.kScaleTensor) { aclDestroyTensor(resources.kScaleTensor); } + if (resources.metadataTensor) { aclDestroyTensor(resources.metadataTensor); } + if (resources.sparseIndicesTensor) { aclDestroyTensor(resources.sparseIndicesTensor); } + if (resources.sparseValuesTensor) { aclDestroyTensor(resources.sparseValuesTensor); } + + if (resources.queryDeviceAddr) { aclrtFree(resources.queryDeviceAddr); } + if (resources.keyDeviceAddr) { aclrtFree(resources.keyDeviceAddr); } + if (resources.weightsDeviceAddr) { aclrtFree(resources.weightsDeviceAddr); } + if (resources.qScaleDeviceAddr) { aclrtFree(resources.qScaleDeviceAddr); } + if (resources.kScaleDeviceAddr) { aclrtFree(resources.kScaleDeviceAddr); } + if (resources.metadataDeviceAddr) { aclrtFree(resources.metadataDeviceAddr); } + if (resources.sparseIndicesDeviceAddr) { aclrtFree(resources.sparseIndicesDeviceAddr); } + if (resources.sparseValuesDeviceAddr) { aclrtFree(resources.sparseValuesDeviceAddr); } + + if (workspaceAddr) { aclrtFree(workspaceAddr); } + if (stream) { aclrtDestroyStream(stream); } + aclrtResetDevice(deviceId); + aclFinalize(); +} + +} // namespace + +int main() { + if (op::GetCurrentPlatformInfo().GetCurNpuArch() != NpuArch::DAV_3510) { + return 0; + } + int32_t deviceId = 0; + aclrtStream stream = nullptr; + TensorResources resources = {}; + void* workspaceAddr = nullptr; + uint64_t workspaceSize = 0; + int64_t B = 2; + int64_t S1 = 4; + int64_t N2 = 1; + int64_t topk = 512; + std::vector sparseIndicesShape = {B, S1, N2, topk}; + int ret = ACL_SUCCESS; + + ret = Init(deviceId, &stream); + if (!CHECK_RET(ret == ACL_SUCCESS)) { + LOG_PRINT("Init acl failed. ERROR: %d\n", ret); + return ret; + } + + ret = InitializeTensors(resources); + if (!CHECK_RET(ret == ACL_SUCCESS)) { + LOG_PRINT("InitializeTensors failed. ERROR: %d\n", ret); + CleanupResources(resources, workspaceAddr, stream, deviceId); + return ret; + } + + ret = ExecuteQuantLightningIndexerV2(resources, stream, &workspaceAddr, &workspaceSize); + if (!CHECK_RET(ret == ACL_SUCCESS)) { + LOG_PRINT("ExecuteQuantLightningIndexerV2 failed. ERROR: %d\n", ret); + CleanupResources(resources, workspaceAddr, stream, deviceId); + return ret; + } + + ret = aclrtSynchronizeStream(stream); + if (!CHECK_RET(ret == ACL_SUCCESS)) { + LOG_PRINT("aclrtSynchronizeStream failed. ERROR: %d\n", ret); + CleanupResources(resources, workspaceAddr, stream, deviceId); + return ret; + } + + PrintOutResult(sparseIndicesShape, resources.sparseIndicesDeviceAddr); + + CleanupResources(resources, workspaceAddr, stream, deviceId); + return 0; +} +``` diff --git a/xllm_ops/attention/quant_lightning_indexer_v2/examples/test_aclnn_quant_lightning_indexer_v2.cpp b/xllm_ops/attention/quant_lightning_indexer_v2/examples/test_aclnn_quant_lightning_indexer_v2.cpp new file mode 100644 index 0000000..e5c52f0 --- /dev/null +++ b/xllm_ops/attention/quant_lightning_indexer_v2/examples/test_aclnn_quant_lightning_indexer_v2.cpp @@ -0,0 +1,455 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file test_aclnn_quant_lightning_indexer_v2.cpp + * \brief + */ +#include +#include +#include +#include +#include "securec.h" +#include "acl/acl.h" +#include "aclnnop/aclnn_quant_lightning_indexer_v2.h" +#include "aclnnop/aclnn_quant_lightning_indexer_v2_metadata.h" +#include "aclnn/opdev/platform.h" + +using namespace std; + +namespace { + +#define CHECK_RET(cond) ((cond) ? true : (false)) + +#define LOG_PRINT(message, ...) \ + do { \ + (void)printf(message, ##__VA_ARGS__); \ + } while (0) + +int64_t GetShapeSize(const std::vector &shape) +{ + int64_t shapeSize = 1; + for (auto i : shape) { + shapeSize *= i; + } + return shapeSize; +} + +int Init(int32_t deviceId, aclrtStream *stream) +{ + auto ret = aclInit(nullptr); + if (!CHECK_RET(ret == ACL_SUCCESS)) { + LOG_PRINT("aclInit failed. ERROR: %d\n", ret); + return ret; + } + ret = aclrtSetDevice(deviceId); + if (!CHECK_RET(ret == ACL_SUCCESS)) { + LOG_PRINT("aclrtSetDevice failed. ERROR: %d\n", ret); + return ret; + } + ret = aclrtCreateStream(stream); + if (!CHECK_RET(ret == ACL_SUCCESS)) { + LOG_PRINT("aclrtCreateStream failed. ERROR: %d\n", ret); + return ret; + } + return 0; +} + +template +int CreateAclTensor(const std::vector &hostData, const std::vector &shape, void **deviceAddr, + aclDataType dataType, aclTensor **tensor) +{ + auto size = GetShapeSize(shape) * sizeof(T); + auto ret = aclrtMalloc(deviceAddr, size, ACL_MEM_MALLOC_HUGE_FIRST); + if (!CHECK_RET(ret == ACL_SUCCESS)) { + LOG_PRINT("aclrtMalloc failed. ERROR: %d\n", ret); + return ret; + } + + ret = aclrtMemcpy(*deviceAddr, size, hostData.data(), size, ACL_MEMCPY_HOST_TO_DEVICE); + if (!CHECK_RET(ret == ACL_SUCCESS)) { + LOG_PRINT("aclrtMemcpy failed. ERROR: %d\n", ret); + return ret; + } + + std::vector strides(shape.size(), 1); + for (int64_t i = shape.size() - 2; i >= 0; i--) { + strides[i] = shape[i + 1] * strides[i + 1]; + } + + *tensor = aclCreateTensor(shape.data(), shape.size(), dataType, strides.data(), 0, aclFormat::ACL_FORMAT_ND, + shape.data(), shape.size(), *deviceAddr); + return 0; +} + +struct TensorResources { + void *queryDeviceAddr = nullptr; + void *keyDeviceAddr = nullptr; + void *weightsDeviceAddr = nullptr; + void *qScaleDeviceAddr = nullptr; + void *kScaleDeviceAddr = nullptr; + void *metadataDeviceAddr = nullptr; + void *sparseIndicesDeviceAddr = nullptr; + void *sparseValuesDeviceAddr = nullptr; + + aclTensor *queryTensor = nullptr; + aclTensor *keyTensor = nullptr; + aclTensor *weightsTensor = nullptr; + aclTensor *qScaleTensor = nullptr; + aclTensor *kScaleTensor = nullptr; + aclTensor *metadataTensor = nullptr; + aclTensor *sparseIndicesTensor = nullptr; + aclTensor *sparseValuesTensor = nullptr; +}; + +int InitializeTensors(TensorResources &resources) +{ + int64_t B = 2; + int64_t S1 = 4; + int64_t S2 = 8; + int64_t N1 = 64; + int64_t N2 = 1; + int64_t D = 128; + int64_t topk = 512; + + std::vector queryShape = {B, S1, N1, D}; + std::vector keyShape = {B, S2, N2, D}; + std::vector weightsShape = {B, S1, N1}; + std::vector qScaleShape = {B, S1, N1}; + std::vector kScaleShape = {B, S2, N2}; + std::vector metadataShape = {1024}; + std::vector sparseIndicesShape = {B, S1, N2, topk}; + std::vector sparseValuesShape = {B, S1, N2, topk}; + + int64_t queryShapeSize = GetShapeSize(queryShape); + int64_t keyShapeSize = GetShapeSize(keyShape); + int64_t weightsShapeSize = GetShapeSize(weightsShape); + int64_t qScaleShapeSize = GetShapeSize(qScaleShape); + int64_t kScaleShapeSize = GetShapeSize(kScaleShape); + int64_t metadataShapeSize = GetShapeSize(metadataShape); + int64_t sparseIndicesShapeSize = GetShapeSize(sparseIndicesShape); + int64_t sparseValuesShapeSize = GetShapeSize(sparseValuesShape); + + std::vector queryHostData(queryShapeSize, 0x38); + std::vector keyHostData(keyShapeSize, 0x38); + std::vector weightsHostData(weightsShapeSize, 0.01f); + std::vector qScaleHostData(qScaleShapeSize, 1.0f); + std::vector kScaleHostData(kScaleShapeSize, 1.0f); + std::vector metadataHostData(metadataShapeSize, 0); + std::vector sparseIndicesHostData(sparseIndicesShapeSize, 0); + std::vector sparseValuesHostData(sparseValuesShapeSize, 0); + + int ret = CreateAclTensor(queryHostData, queryShape, &resources.queryDeviceAddr, aclDataType::ACL_FLOAT8_E4M3FN, + &resources.queryTensor); + if (!CHECK_RET(ret == ACL_SUCCESS)) { + return ret; + } + + ret = CreateAclTensor(keyHostData, keyShape, &resources.keyDeviceAddr, aclDataType::ACL_FLOAT8_E4M3FN, + &resources.keyTensor); + if (!CHECK_RET(ret == ACL_SUCCESS)) { + return ret; + } + + ret = CreateAclTensor(weightsHostData, weightsShape, &resources.weightsDeviceAddr, aclDataType::ACL_FLOAT, + &resources.weightsTensor); + if (!CHECK_RET(ret == ACL_SUCCESS)) { + return ret; + } + + ret = CreateAclTensor(qScaleHostData, qScaleShape, &resources.qScaleDeviceAddr, aclDataType::ACL_FLOAT, + &resources.qScaleTensor); + if (!CHECK_RET(ret == ACL_SUCCESS)) { + return ret; + } + + ret = CreateAclTensor(kScaleHostData, kScaleShape, &resources.kScaleDeviceAddr, aclDataType::ACL_FLOAT, + &resources.kScaleTensor); + if (!CHECK_RET(ret == ACL_SUCCESS)) { + return ret; + } + + ret = CreateAclTensor(metadataHostData, metadataShape, &resources.metadataDeviceAddr, aclDataType::ACL_INT32, + &resources.metadataTensor); + if (!CHECK_RET(ret == ACL_SUCCESS)) { + return ret; + } + + ret = CreateAclTensor(sparseIndicesHostData, sparseIndicesShape, &resources.sparseIndicesDeviceAddr, + aclDataType::ACL_INT32, &resources.sparseIndicesTensor); + if (!CHECK_RET(ret == ACL_SUCCESS)) { + return ret; + } + + ret = CreateAclTensor(sparseValuesHostData, sparseValuesShape, &resources.sparseValuesDeviceAddr, + aclDataType::ACL_BF16, &resources.sparseValuesTensor); + if (!CHECK_RET(ret == ACL_SUCCESS)) { + return ret; + } + + return ACL_SUCCESS; +} + +int GenerateMetadata(TensorResources &resources, aclrtStream stream, int64_t B, int64_t S1, int64_t S2, int64_t N1, + int64_t N2, int64_t D, int64_t topk, int64_t quantMode, int64_t maskMode, int64_t cmpRatio) +{ + constexpr const char layoutQ[] = "BSND"; + constexpr const char layoutK[] = "BSND"; + constexpr size_t layoutLen = sizeof(layoutQ); + char layoutQCopy[layoutLen]; + char layoutKCopy[layoutLen]; + errno_t memcpyRet = memcpy_s(layoutQCopy, sizeof(layoutQCopy), layoutQ, layoutLen); + if (!CHECK_RET(memcpyRet == 0)) { + LOG_PRINT("metadata memcpy_s layoutQ failed. ERROR: %d\n", memcpyRet); + return -1; + } + memcpyRet = memcpy_s(layoutKCopy, sizeof(layoutKCopy), layoutK, layoutLen); + if (!CHECK_RET(memcpyRet == 0)) { + LOG_PRINT("metadata memcpy_s layoutK failed. ERROR: %d\n", memcpyRet); + return -1; + } + + aclOpExecutor *executor; + uint64_t workspaceSize = 0; + int ret = aclnnQuantLightningIndexerV2MetadataGetWorkspaceSize( + nullptr, nullptr, nullptr, nullptr, nullptr, N1, N2, D, topk, quantMode, B, S1, S2, layoutQCopy, layoutKCopy, + maskMode, cmpRatio, resources.metadataTensor, &workspaceSize, &executor); + if (!CHECK_RET(ret == ACL_SUCCESS)) { + LOG_PRINT("aclnnQuantLightningIndexerV2MetadataGetWorkspaceSize failed. ERROR: %d\n", ret); + return ret; + } + + void *metadataWsAddr = nullptr; + if (workspaceSize > 0ULL) { + ret = aclrtMalloc(&metadataWsAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST); + if (!CHECK_RET(ret == ACL_SUCCESS)) { + LOG_PRINT("metadata allocate workspace failed. ERROR: %d\n", ret); + return ret; + } + } + + ret = aclnnQuantLightningIndexerV2Metadata(metadataWsAddr, workspaceSize, executor, stream); + if (!CHECK_RET(ret == ACL_SUCCESS)) { + LOG_PRINT("aclnnQuantLightningIndexerV2Metadata failed. ERROR: %d\n", ret); + if (metadataWsAddr) { + (void)aclrtFree(metadataWsAddr); + } + return ret; + } + + ret = aclrtSynchronizeStream(stream); + if (!CHECK_RET(ret == ACL_SUCCESS)) { + LOG_PRINT("metadata synchronize stream failed. ERROR: %d\n", ret); + if (metadataWsAddr) { + (void)aclrtFree(metadataWsAddr); + } + return ret; + } + + if (metadataWsAddr) { + (void)aclrtFree(metadataWsAddr); + } + return ACL_SUCCESS; +} + +int ExecuteQuantLightningIndexerV2(TensorResources &resources, aclrtStream stream, void **workspaceAddr, + uint64_t *workspaceSize) +{ + int64_t topk = 512; + int64_t quantMode = 1; + int64_t maskMode = 0; + int64_t cmpRatio = 1; + int64_t returnValue = 1; + constexpr const char layoutQStr[] = "BSND"; + constexpr const char layoutKStr[] = "BSND"; + constexpr size_t layoutQLen = sizeof(layoutQStr); + constexpr size_t layoutKLen = sizeof(layoutKStr); + char layoutQ[layoutQLen]; + char layoutK[layoutKLen]; + errno_t memcpyRet = memcpy_s(layoutQ, sizeof(layoutQ), layoutQStr, layoutQLen); + if (!CHECK_RET(memcpyRet == 0)) { + LOG_PRINT("memcpy_s layoutQ failed. ERROR: %d\n", memcpyRet); + return -1; + } + memcpyRet = memcpy_s(layoutK, sizeof(layoutK), layoutKStr, layoutKLen); + if (!CHECK_RET(memcpyRet == 0)) { + LOG_PRINT("memcpy_s layoutK failed. ERROR: %d\n", memcpyRet); + return -1; + } + aclOpExecutor *executor; + + int ret = aclnnQuantLightningIndexerV2GetWorkspaceSize( + resources.queryTensor, resources.keyTensor, resources.weightsTensor, resources.qScaleTensor, + resources.kScaleTensor, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, resources.metadataTensor, + topk, quantMode, -1, layoutQ, layoutK, maskMode, cmpRatio, returnValue, resources.sparseIndicesTensor, + resources.sparseValuesTensor, workspaceSize, &executor); + + if (!CHECK_RET(ret == ACL_SUCCESS)) { + LOG_PRINT("aclnnQuantLightningIndexerV2GetWorkspaceSize failed. ERROR: %d\n", ret); + return ret; + } + + if (*workspaceSize > 0ULL) { + ret = aclrtMalloc(workspaceAddr, *workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST); + if (!CHECK_RET(ret == ACL_SUCCESS)) { + LOG_PRINT("allocate workspace failed. ERROR: %d\n", ret); + return ret; + } + } + + ret = aclnnQuantLightningIndexerV2(*workspaceAddr, *workspaceSize, executor, stream); + if (!CHECK_RET(ret == ACL_SUCCESS)) { + LOG_PRINT("aclnnQuantLightningIndexerV2 failed. ERROR: %d\n", ret); + return ret; + } + + return ACL_SUCCESS; +} + +int PrintOutResult(const std::vector &shape, void *deviceAddr) +{ + auto size = GetShapeSize(shape); + std::vector resultData(size, 0); + auto ret = aclrtMemcpy(resultData.data(), resultData.size() * sizeof(resultData[0]), deviceAddr, + size * sizeof(resultData[0]), ACL_MEMCPY_DEVICE_TO_HOST); + if (!CHECK_RET(ret == ACL_SUCCESS)) { + LOG_PRINT("copy result from device to host failed. ERROR: %d\n", ret); + return ret; + } + LOG_PRINT("sparse_indices result (first 10 elements):\n"); + for (int64_t i = 0; i < size && i < 10; i++) { + LOG_PRINT(" [%ld] = %d\n", i, resultData[i]); + } + return ACL_SUCCESS; +} + +void CleanupResources(TensorResources &resources, void *workspaceAddr, aclrtStream stream, int32_t deviceId) +{ + if (resources.queryTensor) { + aclDestroyTensor(resources.queryTensor); + } + if (resources.keyTensor) { + aclDestroyTensor(resources.keyTensor); + } + if (resources.weightsTensor) { + aclDestroyTensor(resources.weightsTensor); + } + if (resources.qScaleTensor) { + aclDestroyTensor(resources.qScaleTensor); + } + if (resources.kScaleTensor) { + aclDestroyTensor(resources.kScaleTensor); + } + if (resources.metadataTensor) { + aclDestroyTensor(resources.metadataTensor); + } + if (resources.sparseIndicesTensor) { + aclDestroyTensor(resources.sparseIndicesTensor); + } + if (resources.sparseValuesTensor) { + aclDestroyTensor(resources.sparseValuesTensor); + } + + if (resources.queryDeviceAddr) { + aclrtFree(resources.queryDeviceAddr); + } + if (resources.keyDeviceAddr) { + aclrtFree(resources.keyDeviceAddr); + } + if (resources.weightsDeviceAddr) { + aclrtFree(resources.weightsDeviceAddr); + } + if (resources.qScaleDeviceAddr) { + aclrtFree(resources.qScaleDeviceAddr); + } + if (resources.kScaleDeviceAddr) { + aclrtFree(resources.kScaleDeviceAddr); + } + if (resources.metadataDeviceAddr) { + aclrtFree(resources.metadataDeviceAddr); + } + if (resources.sparseIndicesDeviceAddr) { + aclrtFree(resources.sparseIndicesDeviceAddr); + } + if (resources.sparseValuesDeviceAddr) { + aclrtFree(resources.sparseValuesDeviceAddr); + } + + if (workspaceAddr) { + aclrtFree(workspaceAddr); + } + if (stream) { + aclrtDestroyStream(stream); + } + aclrtResetDevice(deviceId); + aclFinalize(); +} + +} // namespace + +int main() +{ + if (op::GetCurrentPlatformInfo().GetCurNpuArch() != NpuArch::DAV_3510) { + return 0; + } + int32_t deviceId = 0; + aclrtStream stream = nullptr; + TensorResources resources = {}; + void *workspaceAddr = nullptr; + uint64_t workspaceSize = 0; + int64_t B = 2; + int64_t S1 = 4; + int64_t S2 = 8; + int64_t N1 = 64; + int64_t N2 = 1; + int64_t D = 128; + int64_t topk = 512; + std::vector sparseIndicesShape = {B, S1, N2, topk}; + int ret = ACL_SUCCESS; + + ret = Init(deviceId, &stream); + if (!CHECK_RET(ret == ACL_SUCCESS)) { + LOG_PRINT("Init acl failed. ERROR: %d\n", ret); + return ret; + } + + ret = InitializeTensors(resources); + if (!CHECK_RET(ret == ACL_SUCCESS)) { + LOG_PRINT("InitializeTensors failed. ERROR: %d\n", ret); + CleanupResources(resources, workspaceAddr, stream, deviceId); + return ret; + } + + ret = GenerateMetadata(resources, stream, B, S1, S2, N1, N2, D, topk, 1, 0, 1); + if (!CHECK_RET(ret == ACL_SUCCESS)) { + LOG_PRINT("GenerateMetadata failed. ERROR: %d\n", ret); + CleanupResources(resources, workspaceAddr, stream, deviceId); + return ret; + } + + ret = ExecuteQuantLightningIndexerV2(resources, stream, &workspaceAddr, &workspaceSize); + if (!CHECK_RET(ret == ACL_SUCCESS)) { + LOG_PRINT("ExecuteQuantLightningIndexerV2 failed. ERROR: %d\n", ret); + CleanupResources(resources, workspaceAddr, stream, deviceId); + return ret; + } + + ret = aclrtSynchronizeStream(stream); + if (!CHECK_RET(ret == ACL_SUCCESS)) { + LOG_PRINT("aclrtSynchronizeStream failed. ERROR: %d\n", ret); + CleanupResources(resources, workspaceAddr, stream, deviceId); + return ret; + } + + PrintOutResult(sparseIndicesShape, resources.sparseIndicesDeviceAddr); + + CleanupResources(resources, workspaceAddr, stream, deviceId); + return 0; +} diff --git a/xllm_ops/attention/quant_lightning_indexer_v2/op_host/CMakeLists.txt b/xllm_ops/attention/quant_lightning_indexer_v2/op_host/CMakeLists.txt new file mode 100644 index 0000000..c05d684 --- /dev/null +++ b/xllm_ops/attention/quant_lightning_indexer_v2/op_host/CMakeLists.txt @@ -0,0 +1,30 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- +add_op_to_compiled_list() + +if (BUILD_OPEN_PROJECT) + set(quant_lightning_indexer_v2_depends attention/lightning_indexer_v2 CACHE INTERNAL "Dependencies for quant_lightning_indexer_v2") + target_sources(op_host_aclnn PRIVATE + quant_lightning_indexer_v2_def.cpp + ) +endif() + +add_ops_compile_options( + OP_NAME QuantLightningIndexerV2 + OPTIONS --cce-auto-sync=off + -Wno-deprecated-declarations + -Werror + -mllvm -cce-vf-remove-membar=false + -mllvm -cce-aicore-hoist-movemask=false +) + +if (NOT BUILD_OPS_RTY_KERNEL) + add_modules_sources(OPTYPE quant_lightning_indexer_v2 ACLNNTYPE aclnn) +endif() \ No newline at end of file diff --git a/xllm_ops/attention/quant_lightning_indexer_v2/op_host/quant_lightning_indexer_v2_def.cpp b/xllm_ops/attention/quant_lightning_indexer_v2/op_host/quant_lightning_indexer_v2_def.cpp new file mode 100644 index 0000000..93d0fd5 --- /dev/null +++ b/xllm_ops/attention/quant_lightning_indexer_v2/op_host/quant_lightning_indexer_v2_def.cpp @@ -0,0 +1,181 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file quant_lightning_indexer_v2_def.cpp + * \brief + */ + +#include "register/op_def_registry.h" + +namespace ops { +class QuantLightningIndexerV2 : public OpDef { +public: + explicit QuantLightningIndexerV2(const char *name) : OpDef(name) + { + this->Input("q").ParamType(REQUIRED).DataType({ge::DT_INT8}).FormatList({ge::FORMAT_ND}).AutoContiguous(); + this->Input("k").ParamType(REQUIRED).DataType({ge::DT_INT8}).FormatList({ge::FORMAT_ND}).IgnoreContiguous(); + this->Input("w").ParamType(REQUIRED).DataType({ge::DT_FLOAT16}).FormatList({ge::FORMAT_ND}).AutoContiguous(); + this->Input("q_descale") + .ParamType(REQUIRED) + .DataTypeList({ge::DT_FLOAT16}) + .FormatList({ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("k_descale") + .ParamType(REQUIRED) + .DataTypeList({ge::DT_FLOAT16}) + .FormatList({ge::FORMAT_ND}) + .IgnoreContiguous(); + this->Input("cu_seqlens_q") + .ParamType(OPTIONAL) + .DataTypeList({ge::DT_INT32}) + .FormatList({ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("cu_seqlens_k") + .ParamType(OPTIONAL) + .DataTypeList({ge::DT_INT32}) + .FormatList({ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("seqused_q") + .ParamType(OPTIONAL) + .DataTypeList({ge::DT_INT32}) + .FormatList({ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("seqused_k") + .ParamType(OPTIONAL) + .DataTypeList({ge::DT_INT32}) + .FormatList({ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("cmp_residual_k") + .ParamType(OPTIONAL) + .DataTypeList({ge::DT_INT32}) + .FormatList({ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("block_table") + .ParamType(OPTIONAL) + .DataTypeList({ge::DT_INT32}) + .FormatList({ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("output_idx_offset") + .ParamType(OPTIONAL) + .DataTypeList({ge::DT_INT32}) + .FormatList({ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("metadata") + .ParamType(OPTIONAL) + .DataTypeList({ge::DT_INT32}) + .FormatList({ge::FORMAT_ND}) + .AutoContiguous(); + this->Output("sparse_indices").ParamType(REQUIRED).DataTypeList({ge::DT_INT32}).FormatList({ge::FORMAT_ND}); + this->Output("sparse_values").ParamType(REQUIRED).DataTypeList({ge::DT_BF16}).FormatList({ge::FORMAT_ND}); + this->Attr("topk").AttrType(REQUIRED).Int(2048); // 2048: 筛选前2048个作为输出index + this->Attr("quant_mode").AttrType(REQUIRED).Int(1); // 1: per-token-head + this->Attr("max_seqlen_q").AttrType(OPTIONAL).Int(-1); // -1: 默认值，表示任意可能长度 + this->Attr("layout_q").AttrType(OPTIONAL).String("BSND"); + this->Attr("layout_k").AttrType(OPTIONAL).String("BSND"); + this->Attr("mask_mode").AttrType(OPTIONAL).Int(0); // 0: 默认值，无mask + this->Attr("cmp_ratio").AttrType(OPTIONAL).Int(1); + this->Attr("return_value").AttrType(OPTIONAL).Int(0); // 0: 默认值 + OpAICoreConfig aicore_config; + aicore_config.DynamicCompileStaticFlag(true) + .DynamicFormatFlag(true) + .DynamicRankSupportFlag(true) + .DynamicShapeSupportFlag(true) + .NeedCheckSupportFlag(false) + .PrecisionReduceFlag(true); + this->AICore().AddConfig("ascend910b", aicore_config); + this->AICore().AddConfig("ascend910_93", aicore_config); + + OpAICoreConfig aicore_config_95; + // fp8/mxfp8/hif8/mxfp4 + aicore_config_95.Input("q") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT8_E4M3FN, ge::DT_FLOAT8_E4M3FN, ge::DT_HIFLOAT8, ge::DT_FLOAT4_E2M1, ge::DT_INT8}) + .FormatList({ge::FORMAT_ND}) + .AutoContiguous(); + aicore_config_95.Input("k") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT8_E4M3FN, ge::DT_FLOAT8_E4M3FN, ge::DT_HIFLOAT8, ge::DT_FLOAT4_E2M1, ge::DT_INT8}) + .FormatList({ge::FORMAT_ND}) + .IgnoreContiguous(); + aicore_config_95.Input("w") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT16}) + .FormatList({ge::FORMAT_ND}) + .AutoContiguous(); + aicore_config_95.Input("q_descale") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT, ge::DT_FLOAT8_E8M0, ge::DT_FLOAT, ge::DT_FLOAT8_E8M0, ge::DT_FLOAT16}) + .FormatList({ge::FORMAT_ND}) + .AutoContiguous(); + aicore_config_95.Input("k_descale") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT, ge::DT_FLOAT8_E8M0, ge::DT_FLOAT, ge::DT_FLOAT8_E8M0, ge::DT_FLOAT16}) + .FormatList({ge::FORMAT_ND}) + .IgnoreContiguous(); + aicore_config_95.Input("cu_seqlens_q") + .ParamType(OPTIONAL) + .DataTypeList({ge::DT_INT32}) + .FormatList({ge::FORMAT_ND}) + .AutoContiguous(); + aicore_config_95.Input("cu_seqlens_k") + .ParamType(OPTIONAL) + .DataTypeList({ge::DT_INT32}) + .FormatList({ge::FORMAT_ND}) + .AutoContiguous(); + aicore_config_95.Input("seqused_q") + .ParamType(OPTIONAL) + .DataTypeList({ge::DT_INT32}) + .FormatList({ge::FORMAT_ND}) + .AutoContiguous(); + aicore_config_95.Input("seqused_k") + .ParamType(OPTIONAL) + .DataTypeList({ge::DT_INT32}) + .FormatList({ge::FORMAT_ND}) + .AutoContiguous(); + aicore_config_95.Input("cmp_residual_k") + .ParamType(OPTIONAL) + .DataTypeList({ge::DT_INT32}) + .FormatList({ge::FORMAT_ND}) + .AutoContiguous(); + aicore_config_95.Input("block_table") + .ParamType(OPTIONAL) + .DataTypeList({ge::DT_INT32}) + .FormatList({ge::FORMAT_ND}) + .AutoContiguous(); + aicore_config_95.Input("output_idx_offset") + .ParamType(OPTIONAL) + .DataTypeList({ge::DT_INT32}) + .FormatList({ge::FORMAT_ND}) + .AutoContiguous(); + aicore_config_95.Input("metadata") + .ParamType(OPTIONAL) + .DataTypeList({ge::DT_INT32}) + .FormatList({ge::FORMAT_ND}) + .AutoContiguous(); + aicore_config_95.Output("sparse_indices") + .ParamType(REQUIRED) + .DataTypeList({ge::DT_INT32}) + .FormatList({ge::FORMAT_ND}); + aicore_config_95.Output("sparse_values") + .ParamType(REQUIRED) + .DataTypeList({ge::DT_BF16}) + .FormatList({ge::FORMAT_ND}); + aicore_config_95.DynamicCompileStaticFlag(true) + .DynamicFormatFlag(true) + .DynamicRankSupportFlag(true) + .DynamicShapeSupportFlag(true) + .NeedCheckSupportFlag(false) + .PrecisionReduceFlag(true); + this->AICore().AddConfig("ascend950", aicore_config_95); + } +}; +OP_ADD(QuantLightningIndexerV2); +} // namespace ops diff --git a/xllm_ops/attention/quant_lightning_indexer_v2/op_host/quant_lightning_indexer_v2_error_log.h b/xllm_ops/attention/quant_lightning_indexer_v2/op_host/quant_lightning_indexer_v2_error_log.h new file mode 100644 index 0000000..e6cb866 --- /dev/null +++ b/xllm_ops/attention/quant_lightning_indexer_v2/op_host/quant_lightning_indexer_v2_error_log.h @@ -0,0 +1,93 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#pragma once + +#include "log/log.h" + +#ifndef OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON +#define OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON(opname, param, actual, reason) \ + OP_LOGE(opname, "Invalid argument for %s, actual: %s, reason: %s", param, actual, reason) +#endif + +#ifndef OP_LOGE_FOR_INVALID_DTYPE +#define OP_LOGE_FOR_INVALID_DTYPE(opname, param, actual, expected) \ + OP_LOGE(opname, "Invalid dtype for %s, actual: %s, expected: %s", param, actual, expected) +#endif + +#ifndef OP_LOGE_FOR_INVALID_DTYPE_WITH_REASON +#define OP_LOGE_FOR_INVALID_DTYPE_WITH_REASON(opname, param, actual, reason) \ + OP_LOGE(opname, "Invalid dtype for %s, actual: %s, reason: %s", param, actual, reason) +#endif + +#ifndef OP_LOGE_FOR_INVALID_DTYPES_WITH_REASON +#define OP_LOGE_FOR_INVALID_DTYPES_WITH_REASON(opname, param, actual, reason) \ + OP_LOGE(opname, "Invalid dtypes for %s, actual: %s, reason: %s", param, actual, reason) +#endif + +#ifndef OP_LOGE_FOR_INVALID_SHAPE +#define OP_LOGE_FOR_INVALID_SHAPE(opname, param, actual, expected) \ + OP_LOGE(opname, "Invalid shape for %s, actual: %s, expected: %s", param, actual, expected) +#endif + +#ifndef OP_LOGE_FOR_INVALID_SHAPE_WITH_REASON +#define OP_LOGE_FOR_INVALID_SHAPE_WITH_REASON(opname, param, actual, reason) \ + OP_LOGE(opname, "Invalid shape for %s, actual: %s, reason: %s", param, actual, reason) +#endif + +#ifndef OP_LOGE_FOR_INVALID_SHAPES_WITH_REASON +#define OP_LOGE_FOR_INVALID_SHAPES_WITH_REASON(opname, param, actual, reason) \ + OP_LOGE(opname, "Invalid shapes for %s, actual: %s, reason: %s", param, actual, reason) +#endif + +#ifndef OP_LOGE_FOR_INVALID_SHAPEDIM +#define OP_LOGE_FOR_INVALID_SHAPEDIM(opname, param, actual, expected) \ + OP_LOGE(opname, "Invalid shape dim for %s, actual: %s, expected: %s", param, actual, expected) +#endif + +#ifndef OP_LOGE_FOR_INVALID_SHAPEDIM_WITH_REASON +#define OP_LOGE_FOR_INVALID_SHAPEDIM_WITH_REASON(opname, param, actual, reason) \ + OP_LOGE(opname, "Invalid shape dim for %s, actual: %s, reason: %s", param, actual, reason) +#endif + +#ifndef OP_LOGE_FOR_INVALID_SHAPEDIMS_WITH_REASON +#define OP_LOGE_FOR_INVALID_SHAPEDIMS_WITH_REASON(opname, param, actual, reason) \ + OP_LOGE(opname, "Invalid shape dims for %s, actual: %s, reason: %s", param, actual, reason) +#endif + +#ifndef OP_LOGE_FOR_INVALID_SHAPESIZE +#define OP_LOGE_FOR_INVALID_SHAPESIZE(opname, param, actual, expected) \ + OP_LOGE(opname, "Invalid shape size for %s, actual: %s, expected: %s", param, actual, expected) +#endif + +#ifndef OP_LOGE_FOR_INVALID_SHAPESIZE_WITH_REASON +#define OP_LOGE_FOR_INVALID_SHAPESIZE_WITH_REASON(opname, param, actual, reason) \ + OP_LOGE(opname, "Invalid shape size for %s, actual: %s, reason: %s", param, actual, reason) +#endif + +#ifndef OP_LOGE_FOR_INVALID_SHAPESIZES_WITH_REASON +#define OP_LOGE_FOR_INVALID_SHAPESIZES_WITH_REASON(opname, param, actual, reason) \ + OP_LOGE(opname, "Invalid shape size for %s, actual: %s, reason: %s", param, actual, reason) +#endif + +#ifndef OP_LOGE_FOR_INVALID_VALUE +#define OP_LOGE_FOR_INVALID_VALUE(opname, param, actual, expected) \ + OP_LOGE(opname, "Invalid value for %s, actual: %s, expected: %s", param, actual, expected) +#endif + +#ifndef OP_LOGE_FOR_INVALID_VALUE_WITH_REASON +#define OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(opname, param, actual, reason) \ + OP_LOGE(opname, "Invalid value for %s, actual: %s, reason: %s", param, actual, reason) +#endif + +#ifndef OP_LOGE_FOR_INVALID_VALUES_WITH_REASON +#define OP_LOGE_FOR_INVALID_VALUES_WITH_REASON(opname, param, actual, reason) \ + OP_LOGE(opname, "Invalid values for %s, actual: %s, reason: %s", param, actual, reason) +#endif \ No newline at end of file diff --git a/xllm_ops/attention/quant_lightning_indexer_v2/op_host/quant_lightning_indexer_v2_infershape.cpp b/xllm_ops/attention/quant_lightning_indexer_v2/op_host/quant_lightning_indexer_v2_infershape.cpp new file mode 100644 index 0000000..0213d71 --- /dev/null +++ b/xllm_ops/attention/quant_lightning_indexer_v2/op_host/quant_lightning_indexer_v2_infershape.cpp @@ -0,0 +1,108 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file quant_lightning_indexer_v2_infershape.cpp + * \brief + */ +#include +#include + +#include "err/ops_err.h" +#include "log/log.h" + +using namespace ge; + +namespace ops { +constexpr uint32_t QUERY_INDEX = 0; +constexpr uint32_t KEY_INDEX = 1; +constexpr uint32_t ATTR_SPARSE_COUNT_INDEX = 0; +constexpr uint32_t ATTR_QUERY_LAYOUT_INDEX = 3; +constexpr uint32_t ATTR_KV_LAYOUT_INDEX = 4; +constexpr uint32_t ATTR_RETURN_VALUE_INDEX = 7; +constexpr uint32_t DIM_NUM_3 = 3; +constexpr uint32_t DIM_NUM_4 = 4; + +static ge::graphStatus InferShapeQuantLightningIndexerV2(gert::InferShapeContext *context) +{ + if (context == nullptr) { + OP_LOGE("QuantLightningIndexerV2", "context is nullptr!"); + return ge::GRAPH_FAILED; + } + const gert::Shape *queryShape = context->GetInputShape(QUERY_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context, queryShape); + const gert::Shape *keyShape = context->GetInputShape(KEY_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context, keyShape); + gert::Shape *sparseIndicesShape = context->GetOutputShape(0); + OP_CHECK_NULL_WITH_CONTEXT(context, sparseIndicesShape); + gert::Shape *sparseValuesShape = context->GetOutputShape(1); + OP_CHECK_NULL_WITH_CONTEXT(context, sparseValuesShape); + + auto attrs = context->GetAttrs(); + OP_CHECK_NULL_WITH_CONTEXT(context, attrs); + const char *inputLayoutQueryPtr = attrs->GetAttrPointer(ATTR_QUERY_LAYOUT_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context, inputLayoutQueryPtr); + const char *inputLayoutKeyPtr = attrs->GetAttrPointer(ATTR_KV_LAYOUT_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context, inputLayoutKeyPtr); + const int64_t *sparse_count = attrs->GetInt(ATTR_SPARSE_COUNT_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context, sparse_count); + + std::string inputLayoutQueryPtrStr = std::string(inputLayoutQueryPtr); + std::string inputLayoutKeyPtrStr = std::string(inputLayoutKeyPtr); + if (inputLayoutQueryPtrStr != "TND" && inputLayoutQueryPtrStr != "BSND") { + OP_LOGE("QuantLightningIndexerV2", "The input layout_q should be BSND or TND, but got %s.", + inputLayoutQueryPtrStr.c_str()); + return GRAPH_FAILED; + } + + int64_t keyHeadNum = (inputLayoutKeyPtrStr == "TND") ? keyShape->GetDim(1) : keyShape->GetDim(2); + if (inputLayoutQueryPtrStr == "BSND") { + sparseIndicesShape->SetDimNum(DIM_NUM_4); + sparseIndicesShape->SetDim(0, queryShape->GetDim(0)); // 0:Dim B + sparseIndicesShape->SetDim(1, queryShape->GetDim(1)); // 1:Dim S + sparseIndicesShape->SetDim(2, keyHeadNum); // 2:Dim N + sparseIndicesShape->SetDim(3, *sparse_count); // 3:Dim K + } else { + sparseIndicesShape->SetDimNum(DIM_NUM_3); + sparseIndicesShape->SetDim(0, queryShape->GetDim(0)); // 0:Dim T + sparseIndicesShape->SetDim(1, keyHeadNum); // 1:output shape's N Dim, 2: key shape's N Dim + sparseIndicesShape->SetDim(2, *sparse_count); // 2:Dim K + } + const int32_t *return_value = attrs->GetAttrPointer(ATTR_RETURN_VALUE_INDEX); + bool returnValueFlag = (return_value != nullptr) ? (*return_value != 0) : false; + if (returnValueFlag) { + *sparseValuesShape = *sparseIndicesShape; + } else { + sparseValuesShape->SetDimNum(1); + sparseValuesShape->SetDim(0, 0); + } + + OP_LOGD(context->GetNodeName(), "QuantLightningIndexerV2 InferShape end."); + return ge::GRAPH_SUCCESS; +} + +static ge::graphStatus InferDataTypeQuantLightningIndexerV2(gert::InferDataTypeContext *context) +{ + if (context == nullptr) { + OP_LOGE("QuantLightningIndexerV2", "InferDataTypeContext context is nullptr!"); + return ge::GRAPH_FAILED; + } + OP_LOGD(context->GetNodeName(), "Enter QuantLightningIndexerV2 InferDataType impl."); + // default index data type is int32 + ge::DataType outputType = ge::DT_INT32; + context->SetOutputDataType(0, outputType); + OP_LOGD(context->GetNodeName(), "QuantLightningIndexerV2 InferDataType end."); + return GRAPH_SUCCESS; +} + +IMPL_OP_INFERSHAPE(QuantLightningIndexerV2) + .InferShape(InferShapeQuantLightningIndexerV2) + .InferDataType(InferDataTypeQuantLightningIndexerV2); +} // namespace ops diff --git a/xllm_ops/attention/quant_lightning_indexer_v2/op_host/quant_lightning_indexer_v2_tiling.cpp b/xllm_ops/attention/quant_lightning_indexer_v2/op_host/quant_lightning_indexer_v2_tiling.cpp new file mode 100644 index 0000000..849db66 --- /dev/null +++ b/xllm_ops/attention/quant_lightning_indexer_v2/op_host/quant_lightning_indexer_v2_tiling.cpp @@ -0,0 +1,1660 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file quant_lightning_indexer_v2_tiling.cpp + * \brief + */ + +#include "quant_lightning_indexer_v2_tiling.h" + +#include "quant_lightning_indexer_v2_error_log.h" +#include "../op_kernel/quant_lightning_indexer_v2_template_tiling_key.h" + +using namespace ge; +using namespace AscendC; +using std::map; +using std::string; +namespace optiling { + +static const std::map DATATYPE_TO_STRING_MAP = { + {ge::DT_UNDEFINED, "DT_UNDEFINED"}, // Used to indicate a DataType field has not been set. + {ge::DT_FLOAT, "DT_FLOAT"}, // float type + {ge::DT_FLOAT16, "DT_FLOAT16"}, // fp16 type + {ge::DT_INT8, "DT_INT8"}, // int8 type + {ge::DT_INT16, "DT_INT16"}, // int16 type + {ge::DT_UINT16, "DT_UINT16"}, // uint16 type + {ge::DT_UINT8, "DT_UINT8"}, // uint8 type + {ge::DT_INT32, "DT_INT32"}, // uint32 type + {ge::DT_INT64, "DT_INT64"}, // int64 type + {ge::DT_UINT32, "DT_UINT32"}, // unsigned int32 + {ge::DT_UINT64, "DT_UINT64"}, // unsigned int64 + {ge::DT_BOOL, "DT_BOOL"}, // bool type + {ge::DT_DOUBLE, "DT_DOUBLE"}, // double type + {ge::DT_DUAL, "DT_DUAL"}, // dual output type + {ge::DT_DUAL_SUB_INT8, "DT_DUAL_SUB_INT8"}, // dual output int8 type + {ge::DT_DUAL_SUB_UINT8, "DT_DUAL_SUB_UINT8"}, // dual output uint8 type + {ge::DT_COMPLEX32, "DT_COMPLEX32"}, // complex32 type + {ge::DT_COMPLEX64, "DT_COMPLEX64"}, // complex64 type + {ge::DT_COMPLEX128, "DT_COMPLEX128"}, // complex128 type + {ge::DT_QINT8, "DT_QINT8"}, // qint8 type + {ge::DT_QINT16, "DT_QINT16"}, // qint16 type + {ge::DT_QINT32, "DT_QINT32"}, // qint32 type + {ge::DT_QUINT8, "DT_QUINT8"}, // quint8 type + {ge::DT_QUINT16, "DT_QUINT16"}, // quint16 type + {ge::DT_RESOURCE, "DT_RESOURCE"}, // resource type + {ge::DT_STRING_REF, "DT_STRING_REF"}, // string ref type + {ge::DT_STRING, "DT_STRING"}, // string type + {ge::DT_VARIANT, "DT_VARIANT"}, // dt_variant type + {ge::DT_BF16, "DT_BFLOAT16"}, // dt_bfloat16 type + {ge::DT_INT4, "DT_INT4"}, // dt_variant type + {ge::DT_UINT1, "DT_UINT1"}, // dt_variant type + {ge::DT_INT2, "DT_INT2"}, // dt_variant type + {ge::DT_UINT2, "DT_UINT2"}, // dt_variant type + {ge::DT_HIFLOAT8, "DT_HIFLOAT8"}, + {ge::DT_FLOAT8_E4M3FN, "DT_FLOAT8_E4M3FN"}, + {ge::DT_FLOAT8_E8M0, "DT_FLOAT8_E8M0"}, + {ge::DT_FLOAT4_E2M1, "DT_FLOAT4_E2M1"}}; + +std::string QLIV2DataTypeToSerialString(ge::DataType type) +{ + const auto it = DATATYPE_TO_STRING_MAP.find(type); + if (it != DATATYPE_TO_STRING_MAP.end()) { + return it->second; + } else { + OP_LOGE("QLIV2DataTypeToSerialString ", "datatype %d not support", type); + return "UNDEFINED"; + } +} + +static std::vector ToVector(const gert::Shape &shape) +{ + size_t shapeSize = shape.GetDimNum(); + std::vector shapeVec(shapeSize, 0); + + for (size_t i = 0; i < shapeSize; i++) { + shapeVec[i] = shape.GetDim(i); + } + return shapeVec; +} + +static std::string ToStringRaw(const gert::Shape &shape) +{ + std::ostringstream oss; + auto v = ToVector(shape); + if (v.size() > 0) { + for (size_t i = 0; i < v.size() - 1; ++i) { + oss << v[i] << ", "; + } + oss << v[v.size() - 1]; + } + return oss.str(); +} + +// --------------------------QLIV2InfoParser类成员函数定义------------------------------------- +ge::graphStatus QLIV2InfoParser::CheckRequiredInOutExistence() const +{ + OP_CHECK_IF(opParamInfo_.query.shape == nullptr, + OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON(opName_, "q", "The shape of q is nullptr"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(opParamInfo_.query.desc == nullptr, + OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON(opName_, "q", "The desc of q is nullptr"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(opParamInfo_.key.shape == nullptr, + OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON(opName_, "k", "The shape of k is nullptr"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(opParamInfo_.key.desc == nullptr, + OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON(opName_, "k", "The desc of k is nullptr"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(opParamInfo_.weights.shape == nullptr, + OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON(opName_, "w", "The shape of w is nullptr"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(opParamInfo_.weights.desc == nullptr, + OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON(opName_, "w", "The desc of w is nullptr"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(opParamInfo_.query_dequant_scale.shape == nullptr, + OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON(opName_, "query_dequant_scale", + "The shape of query_dequant_scale is nullptr"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(opParamInfo_.query_dequant_scale.desc == nullptr, + OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON(opName_, "query_dequant_scale", + "The desc of query_dequant_scale is nullptr"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(opParamInfo_.key_dequant_scale.shape == nullptr, + OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON(opName_, "key_dequant_scale", + "The shape of key_dequant_scale is nullptr"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(opParamInfo_.key_dequant_scale.desc == nullptr, + OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON(opName_, "key_dequant_scale", + "The desc of key_dequant_scale is nullptr"), + return ge::GRAPH_FAILED); + OP_CHECK_IF( + opParamInfo_.attenOut.shape == nullptr, + OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON(opName_, "sparse_indices", "The shape of sparse_indices is nullptr"), + return ge::GRAPH_FAILED); + OP_CHECK_IF( + opParamInfo_.attenOut.desc == nullptr, + OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON(opName_, "sparse_indices", "The desc of sparse_indices is nullptr"), + return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QLIV2InfoParser::CheckRequiredAttrExistence() const +{ + OP_CHECK_IF(opParamInfo_.layOutQuery == nullptr, + OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON(opName_, "layout_q", "Layout_q is nullptr"), + return ge::GRAPH_FAILED); + + OP_CHECK_IF(opParamInfo_.layOutKey == nullptr, + OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON(opName_, "layout_k", "Layout_k is nullptr"), + return ge::GRAPH_FAILED); + + OP_CHECK_IF(opParamInfo_.sparseCount == nullptr, + OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON(opName_, "sparse_count", "Sparse_count is nullptr"), + return ge::GRAPH_FAILED); + + OP_CHECK_IF(opParamInfo_.sparseMode == nullptr, + OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON(opName_, "sparse_mode", "Sparse_mode is nullptr"), + return ge::GRAPH_FAILED); + + OP_CHECK_IF(opParamInfo_.quantMode == nullptr, + OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON(opName_, "query_quant_mode", "Query_quant_mode is nullptr"), + return ge::GRAPH_FAILED); + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QLIV2InfoParser::CheckRequiredParaExistence() const +{ + if (CheckRequiredInOutExistence() != ge::GRAPH_SUCCESS || CheckRequiredAttrExistence() != ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QLIV2InfoParser::GetOpName() +{ + if (context_->GetNodeName() == nullptr) { + OP_LOGE("LightningIndexerV2", "opName got from TilingContext is nullptr"); + return ge::GRAPH_FAILED; + } + opName_ = context_->GetNodeName(); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QLIV2InfoParser::GetNpuInfo() +{ + platformInfo_ = context_->GetPlatformInfo(); + OP_CHECK_IF(platformInfo_ == nullptr, OP_LOGE(opName_, "GetPlatformInfo is nullptr."), return ge::GRAPH_FAILED); + + auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfo_); + uint32_t aivNum = ascendcPlatform.GetCoreNumAiv(); + uint32_t aicNum = ascendcPlatform.GetCoreNumAic(); + OP_CHECK_IF(aicNum == 0 || aivNum == 0, OP_LOGE(opName_, "num of core obtained is 0."), return GRAPH_FAILED); + + socVersion_ = ascendcPlatform.GetSocVersion(); + npuArch_ = ascendcPlatform.GetCurNpuArch(); + if ((npuArch_ != NpuArch::DAV_2201) && (npuArch_ != NpuArch::DAV_3510)) { + OP_LOGE(opName_, "NpuArch[%d] is not support.", static_cast(npuArch_)); + return GRAPH_FAILED; + } + OP_CHECK_IF(context_->GetWorkspaceSizes(1) == nullptr, OP_LOGE(opName_, "workSpaceSize got from ge is nullptr"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(context_->GetRawTilingData() == nullptr, + OP_LOGE(context_->GetNodeName(), "RawTilingData got from GE context is nullptr."), + return ge::GRAPH_FAILED); + + return ge::GRAPH_SUCCESS; +} + +void QLIV2InfoParser::GetOptionalInputParaInfo() +{ + opParamInfo_.cuSeqLensQ.tensor = context_->GetOptionalInputTensor(CU_SEQLENS_Q_INDEX); + opParamInfo_.cuSeqLensQ.desc = context_->GetOptionalInputDesc(CU_SEQLENS_Q_INDEX); + opParamInfo_.cuSeqLensK.tensor = context_->GetOptionalInputTensor(CU_SEQLENS_K_INDEX); + opParamInfo_.cuSeqLensK.desc = context_->GetOptionalInputDesc(CU_SEQLENS_K_INDEX); + opParamInfo_.sequsedQ.tensor = context_->GetOptionalInputTensor(SEQUSED_Q_INDEX); + opParamInfo_.sequsedQ.desc = context_->GetOptionalInputDesc(SEQUSED_Q_INDEX); + opParamInfo_.sequsedK.tensor = context_->GetOptionalInputTensor(SEQUSED_K_INDEX); + opParamInfo_.sequsedK.desc = context_->GetOptionalInputDesc(SEQUSED_K_INDEX); + opParamInfo_.cmpResidualK.tensor = context_->GetOptionalInputTensor(CMP_RESIDUAL_K_INDEX); + opParamInfo_.cmpResidualK.desc = context_->GetOptionalInputDesc(CMP_RESIDUAL_K_INDEX); + opParamInfo_.blockTable.tensor = context_->GetOptionalInputTensor(BLOCK_TABLE_INDEX); + opParamInfo_.blockTable.desc = context_->GetOptionalInputDesc(BLOCK_TABLE_INDEX); + opParamInfo_.outputIdxOffset.tensor = context_->GetOptionalInputTensor(OUTPUT_IDX_OFFSET_INDEX); + opParamInfo_.outputIdxOffset.desc = context_->GetOptionalInputDesc(OUTPUT_IDX_OFFSET_INDEX); + opParamInfo_.metadata.tensor = context_->GetOptionalInputTensor(METADATA_INDEX); + opParamInfo_.metadata.desc = context_->GetOptionalInputDesc(METADATA_INDEX); +} + +void QLIV2InfoParser::GetInputParaInfo() +{ + opParamInfo_.query.desc = context_->GetInputDesc(QUERY_INDEX); + opParamInfo_.query.shape = context_->GetInputShape(QUERY_INDEX); + opParamInfo_.key.desc = context_->GetInputDesc(KEY_INDEX); + opParamInfo_.key.shape = context_->GetInputShape(KEY_INDEX); + opParamInfo_.weights.desc = context_->GetInputDesc(WEIGTHS_INDEX); + opParamInfo_.weights.shape = context_->GetInputShape(WEIGTHS_INDEX); + opParamInfo_.query_dequant_scale.desc = context_->GetInputDesc(QUERY_DEQUANT_SCALE_INDEX); + opParamInfo_.query_dequant_scale.shape = context_->GetInputShape(QUERY_DEQUANT_SCALE_INDEX); + opParamInfo_.key_dequant_scale.desc = context_->GetInputDesc(KEY_DEQUANT_SCALE_INDEX); + opParamInfo_.key_dequant_scale.shape = context_->GetInputShape(KEY_DEQUANT_SCALE_INDEX); + GetOptionalInputParaInfo(); +} + +void QLIV2InfoParser::GetOutputParaInfo() +{ + opParamInfo_.attenOut.desc = context_->GetOutputDesc(SPARSE_INDICES_INDEX); + opParamInfo_.attenOut.shape = context_->GetOutputShape(SPARSE_INDICES_INDEX); + opParamInfo_.sparseValues.desc = context_->GetOutputDesc(SPARSE_VALUES_INDEX); + opParamInfo_.sparseValues.shape = context_->GetOutputShape(SPARSE_VALUES_INDEX); +} + +ge::graphStatus QLIV2InfoParser::GetAttrParaInfo() +{ + auto attrs = context_->GetAttrs(); + OP_CHECK_IF(attrs == nullptr, OP_LOGE(context_->GetNodeName(), "attrs got from ge is nullptr"), + return ge::GRAPH_FAILED); + + OP_LOGI(context_->GetNodeName(), "GetAttrParaInfo start"); + + opParamInfo_.quantMode = attrs->GetAttrPointer(ATTR_QUANT_MODE_INDEX); + opParamInfo_.maxSeqlenQ = attrs->GetAttrPointer(ATTR_MAX_SEQLEN_Q_INDEX); + opParamInfo_.layOutQuery = attrs->GetStr(ATTR_QUERY_LAYOUT_INDEX); + opParamInfo_.layOutKey = attrs->GetStr(ATTR_KEY_LAYOUT_INDEX); + opParamInfo_.sparseCount = attrs->GetAttrPointer(ATTR_TOPK_INDEX); + opParamInfo_.sparseMode = attrs->GetAttrPointer(ATTR_MASK_MODE_INDEX); + opParamInfo_.cmpRatio = attrs->GetAttrPointer(ATTR_CMP_RATIO_INDEX); + opParamInfo_.returnValue = attrs->GetAttrPointer(ATTR_RETURN_VALUE_INDEX); + auto keyStrides = context_->GetDynamicInputStride(KEY_INDEX, 0); + auto keyDequantScaleStrides = context_->GetDynamicInputStride(KEY_DEQUANT_SCALE_INDEX, 0); + if (keyStrides != nullptr && keyStrides->GetDimNum() > 0) { + for (size_t i = 0; i < keyStrides->GetDimNum(); i++) { + keyStridesVec_.push_back(keyStrides->GetStride(i)); + } + } + if (keyDequantScaleStrides != nullptr && keyDequantScaleStrides->GetDimNum() > 0) { + for (size_t i = 0; i < keyDequantScaleStrides->GetDimNum(); i++) { + keyDequantScaleStridesVec_.push_back(keyDequantScaleStrides->GetStride(i)); + } + } + + if (opParamInfo_.layOutQuery != nullptr) { + OP_LOGI(context_->GetNodeName(), "layout_query is:%s", opParamInfo_.layOutQuery); + } + if (opParamInfo_.layOutKey != nullptr) { + OP_LOGI(context_->GetNodeName(), "layout_key is:%s", opParamInfo_.layOutKey); + } + if (opParamInfo_.sparseCount != nullptr) { + OP_LOGI(context_->GetNodeName(), "selscted count is:%d", *opParamInfo_.sparseCount); + } + if (opParamInfo_.sparseMode != nullptr) { + OP_LOGI(context_->GetNodeName(), "sparse mode is:%d", *opParamInfo_.sparseMode); + } + if (opParamInfo_.cmpRatio != nullptr) { + OP_LOGI(context_->GetNodeName(), "cmpRatio is:%d", *opParamInfo_.cmpRatio); + } + if (opParamInfo_.returnValue != nullptr) { + OP_LOGI(context_->GetNodeName(), "returnValue is:%s", *opParamInfo_.returnValue ? "true" : "false"); + } + if (opParamInfo_.maxSeqlenQ != nullptr) { + OP_LOGI(context_->GetNodeName(), "maxSeqlenQ is:%d", *opParamInfo_.maxSeqlenQ); + } + if (opParamInfo_.quantMode != nullptr) { + OP_LOGI(context_->GetNodeName(), "query_quant_mode mode is:%d", *opParamInfo_.quantMode); + } + OP_LOGI(context_->GetNodeName(), "GetAttrParaInfo end"); + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QLIV2InfoParser::CheckAttrParaInfo() +{ + std::string layout_key(opParamInfo_.layOutKey); + std::string layout_query(opParamInfo_.layOutQuery); + + if (npuArch_ == NpuArch::DAV_2201) { + OP_CHECK_IF((std::string(opParamInfo_.layOutKey) != "PA_BBND"), + OP_LOGE(opName_, + "input attr layout_key only supported PA_BBND," + "but now layout_key is %s.", + layout_key.c_str()), + return ge::GRAPH_FAILED); + } else if (npuArch_ == NpuArch::DAV_3510) { + OP_CHECK_IF( + ((std::string(opParamInfo_.layOutKey) != "PA_BBND") && (std::string(opParamInfo_.layOutKey) != "BSND") && + (std::string(opParamInfo_.layOutKey) != "TND")), + OP_LOGE(opName_, "invalid layout_k: %s, Layout_k only supports PA_BBND, BSND or TND", + std::string(opParamInfo_.layOutKey).c_str()), + return ge::GRAPH_FAILED); + } + + if (npuArch_ == NpuArch::DAV_2201) { + OP_CHECK_IF(!((*opParamInfo_.sparseCount > 0) && (*opParamInfo_.sparseCount <= SPARSE_LIMIT)), + OP_LOGE(opName_, "input attr sparse_count must > 0 and <= %d, but now sparse_count is %d", + SPARSE_LIMIT, *opParamInfo_.sparseCount), + return ge::GRAPH_FAILED); + OP_CHECK_IF((*opParamInfo_.cmpRatio <= 0) || (*opParamInfo_.cmpRatio > 128) || + ((*opParamInfo_.cmpRatio & (*opParamInfo_.cmpRatio - 1)) != 0), + OP_LOGE(opName_, + "input attr cmpRatio must > 0 and <= 128 and should be powers of 2," + " but now cmpRatio is %ld.", + *opParamInfo_.cmpRatio), + return ge::GRAPH_FAILED); + } else if (npuArch_ == NpuArch::DAV_3510) { + OP_CHECK_IF( + !((*opParamInfo_.sparseCount > 0) && (*opParamInfo_.sparseCount <= SPARSE_LIMIT_8K)), + OP_LOGE(opName_, "invalid topk: %s, Sparse_count must > 0 and <= %s", + std::to_string(*opParamInfo_.sparseCount).c_str(), std::to_string(SPARSE_LIMIT_8K).c_str()), + return ge::GRAPH_FAILED); + OP_CHECK_IF((*opParamInfo_.cmpRatio <= 0) || (*opParamInfo_.cmpRatio > 128), + OP_LOGE(opName_, "invalid cmp_ratio: %s, Cmp_ratio must > 0 and <= 128", + std::to_string(*opParamInfo_.cmpRatio).c_str()), + return ge::GRAPH_FAILED); + } + + OP_CHECK_IF( + ((std::string(opParamInfo_.layOutQuery) != "BSND") && (std::string(opParamInfo_.layOutQuery) != "TND")), + OP_LOGE(opName_, "invalid layout_q: %s, Layout_q only supports BSND or TND", + std::string(opParamInfo_.layOutQuery).c_str()), + return ge::GRAPH_FAILED); + OP_CHECK_IF( + ((std::string(opParamInfo_.layOutKey) != "PA_BBND") && + (std::string(opParamInfo_.layOutQuery)) != (std::string(opParamInfo_.layOutKey))), + OP_LOGE(opName_, "invalid layout_q and layout_k: %s, Outside of PA, layout_q and layout_k must be the same", + (layout_query + " and " + layout_key).c_str()), + return ge::GRAPH_FAILED); + OP_CHECK_IF( + !((*opParamInfo_.sparseMode == 0) || (*opParamInfo_.sparseMode == SPARSE_MODE_LOWER)), + OP_LOGE(opName_, "invalid sparse_mode: %s, Sparse_mode only supports 0 or 3", + std::to_string(*opParamInfo_.sparseMode).c_str()), + return ge::GRAPH_FAILED); + if (npuArch_ == NpuArch::DAV_2201) { + OP_CHECK_IF(*opParamInfo_.quantMode != 2, OP_LOGE(opName_, "input attr quant_mode only supported 2."), + return ge::GRAPH_FAILED); + } else if (npuArch_ == NpuArch::DAV_3510) { + OP_CHECK_IF((*opParamInfo_.quantMode != QUANT_MODE_FP8) && (*opParamInfo_.quantMode != QUANT_MODE_INT8) && + (*opParamInfo_.quantMode != QUANT_MODE_MXFP8) && + (*opParamInfo_.quantMode != QUANT_MODE_HIFLOAT8) && + (*opParamInfo_.quantMode != QUANT_MODE_MXFP4), + OP_LOGE(opName_, "invalid quant_mode: %s, Quant_mode only supports 1, 2, 3, 4 and 5", + std::to_string(*opParamInfo_.quantMode).c_str()), + return ge::GRAPH_FAILED); + } + + if (npuArch_ == NpuArch::DAV_2201) { + OP_CHECK_IF(*opParamInfo_.returnValue, OP_LOGE(opName_, "input attr returnValue only supported False."), + return ge::GRAPH_FAILED); + } else if (npuArch_ == NpuArch::DAV_3510) { + OP_CHECK_IF((*opParamInfo_.returnValue != 0) && (*opParamInfo_.returnValue != 1), + OP_LOGE(opName_, "invalid return_value: %s, Return_value only supports 0 or 1", + std::to_string(*opParamInfo_.returnValue).c_str()), + return ge::GRAPH_FAILED); + } + OP_CHECK_IF( + (*opParamInfo_.maxSeqlenQ < -1), + OP_LOGE(opName_, "invalid max_seqlen_q: %s, Max_seqlen_q must >= -1", + std::to_string(*opParamInfo_.maxSeqlenQ).c_str()), + return ge::GRAPH_FAILED); + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QLIV2InfoParser::GetOpParaInfo() +{ + GetInputParaInfo(); + GetOutputParaInfo(); + if (ge::GRAPH_SUCCESS != GetAttrParaInfo()) { + return ge::GRAPH_FAILED; + } + if (ge::GRAPH_SUCCESS != CheckAttrParaInfo()) { + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QLIV2InfoParser::GetAndCheckInOutDataType() +{ + inputQType_ = opParamInfo_.query.desc->GetDataType(); + inputKType_ = opParamInfo_.key.desc->GetDataType(); + weightsType_ = opParamInfo_.weights.desc->GetDataType(); + inputQueryScaleType_ = opParamInfo_.query_dequant_scale.desc->GetDataType(); + inputKeyScaleType_ = opParamInfo_.key_dequant_scale.desc->GetDataType(); + outputType_ = opParamInfo_.attenOut.desc->GetDataType(); + valuesOutType_ = opParamInfo_.sparseValues.desc->GetDataType(); + + OP_CHECK_IF(!(inputQType_ == inputKType_), + OP_LOGE_FOR_INVALID_DTYPES_WITH_REASON( + opName_, "q and k", + QLIV2DataTypeToSerialString(inputQType_) + " and " + QLIV2DataTypeToSerialString(inputKType_), + "The dtype of q and k must be same"), + return ge::GRAPH_FAILED); + + OP_CHECK_IF(!(inputQueryScaleType_ == inputKeyScaleType_), + OP_LOGE_FOR_INVALID_DTYPES_WITH_REASON(opName_, "q_descale and k_descale", + QLIV2DataTypeToSerialString(inputQueryScaleType_) + " and " + + QLIV2DataTypeToSerialString(inputKeyScaleType_), + "The dtype of q_descale and k_descale must be same"), + return ge::GRAPH_FAILED); + if (npuArch_ == NpuArch::DAV_2201) { + OP_CHECK_IF(inputQType_ != ge::DT_INT8, + OP_LOGE(opName_, "The data types of the input query and key must be int8."), + return ge::GRAPH_FAILED); + OP_CHECK_IF( + inputQueryScaleType_ != ge::DT_FLOAT16, + OP_LOGE(opName_, "The data types of the input query_dequant_scale and key_dequant_scale must be float16."), + return ge::GRAPH_FAILED); + } else if (npuArch_ == NpuArch::DAV_3510) { + ge::DataType expectQType = ge::DT_FLOAT8_E4M3FN; + ge::DataType expectScaleType = ge::DT_FLOAT; + if (*opParamInfo_.quantMode == QUANT_MODE_MXFP8) { + expectScaleType = ge::DT_FLOAT8_E8M0; + } else if (*opParamInfo_.quantMode == QUANT_MODE_HIFLOAT8) { + expectQType = ge::DT_HIFLOAT8; + } else if (*opParamInfo_.quantMode == QUANT_MODE_MXFP4) { + expectQType = ge::DT_FLOAT4_E2M1; + expectScaleType = ge::DT_FLOAT8_E8M0; + } else if (*opParamInfo_.quantMode == QUANT_MODE_INT8) { + expectQType = ge::DT_INT8; + expectScaleType = ge::DT_FLOAT16; + } + OP_CHECK_IF(inputQType_ != expectQType, + OP_LOGE_FOR_INVALID_DTYPES_WITH_REASON( + opName_, "q and k", + QLIV2DataTypeToSerialString(inputQType_) + " and " + QLIV2DataTypeToSerialString(inputKType_), + "The dtype of q and k must match quant_mode"), + return ge::GRAPH_FAILED); + OP_CHECK_IF( + inputQueryScaleType_ != expectScaleType, + OP_LOGE_FOR_INVALID_DTYPES_WITH_REASON(opName_, "q_descale and k_descale", + QLIV2DataTypeToSerialString(inputQueryScaleType_) + " and " + + QLIV2DataTypeToSerialString(inputKeyScaleType_), + "The dtype of q_descale and k_descale must match quant_mode"), + return ge::GRAPH_FAILED); + } + + if (npuArch_ == NpuArch::DAV_2201) { + OP_CHECK_IF(weightsType_ != ge::DT_FLOAT16, + OP_LOGE(opName_, "The data types of the input weights must be float16."), return ge::GRAPH_FAILED); + } else if (npuArch_ == NpuArch::DAV_3510) { + if (inputQType_ == ge::DT_INT8) { + OP_CHECK_IF(weightsType_ != ge::DT_FLOAT16, + OP_LOGE_FOR_INVALID_DTYPE_WITH_REASON( + opName_, "w", QLIV2DataTypeToSerialString(weightsType_).c_str(), + "When the dtype of query is int8, the dtype of w must be float16"), + return ge::GRAPH_FAILED); + } else { + OP_CHECK_IF(weightsType_ != ge::DT_FLOAT, + OP_LOGE_FOR_INVALID_DTYPE_WITH_REASON( + opName_, "w", QLIV2DataTypeToSerialString(weightsType_).c_str(), + "When the dtype of query is not int8, the dtype of w must be float"), + return ge::GRAPH_FAILED); + } + } + + OP_CHECK_IF(outputType_ != ge::DT_INT32, + OP_LOGE_FOR_INVALID_DTYPE_WITH_REASON(opName_, "sparse_indices", + QLIV2DataTypeToSerialString(outputType_).c_str(), + "The dtype of sparse_indices must be int32"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(valuesOutType_ != ge::DT_BF16, + OP_LOGE_FOR_INVALID_DTYPE_WITH_REASON(opName_, "sparse_values", + QLIV2DataTypeToSerialString(valuesOutType_).c_str(), + "The dtype of sparse_values must be bfloat16"), + return ge::GRAPH_FAILED); + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QLIV2InfoParser::GetQueryKeyAndOutLayout() +{ + // 获取query,key的Layout基准值 + const map layoutQueryMap = {{"BSND", DataLayout::BSND}, {"TND", DataLayout::TND}}; + + std::string layout_query(opParamInfo_.layOutQuery); + auto QLayout_ = layoutQueryMap.find(layout_query); + if (QLayout_ != layoutQueryMap.end()) { + qLayout_ = QLayout_->second; + } + + const map layoutKeyMap = {{"BSND", DataLayout::BSND}, + {"TND", DataLayout::TND}, + {"PA_BSND", DataLayout::PA_BBND}, + {"PA_BBND", DataLayout::PA_BBND}}; + std::string layout_key(opParamInfo_.layOutKey); + auto KLayout = layoutKeyMap.find(layout_key); + if (KLayout != layoutKeyMap.end()) { + kLayout_ = KLayout->second; + } + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QLIV2InfoParser::GetAndCheckOptionalInput() +{ + // =============== K 侧校验 =============== + if (kLayout_ == DataLayout::PA_BBND) { + // PA_BBND: block_table 必传, seqused_k 必传, cu_seqlens_k 不传 + OP_CHECK_IF(opParamInfo_.blockTable.tensor == nullptr, + OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON(opName_, "block_table", + "When layout_k is PA_BBND, block_table must not be null"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(opParamInfo_.sequsedK.tensor == nullptr, + OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON(opName_, "seqused_k", + "When layout_k is PA_BBND, seqused_k must not be null"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(opParamInfo_.cuSeqLensK.tensor != nullptr, + OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON( + opName_, "cu_seqlens_k", "When layout_k is PA_BBND, cu_seqlens_k must not be provided"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(opParamInfo_.blockTable.desc->GetDataType() != ge::DT_INT32, + OP_LOGE_FOR_INVALID_DTYPE_WITH_REASON( + opName_, "block_table", + QLIV2DataTypeToSerialString(opParamInfo_.blockTable.desc->GetDataType()).c_str(), + "The dtype of block_table only supports int32"), + return ge::GRAPH_FAILED); + OP_CHECK_IF( + opParamInfo_.sequsedK.desc->GetDataType() != ge::DT_INT32, + OP_LOGE_FOR_INVALID_DTYPE_WITH_REASON( + opName_, "seqused_k", QLIV2DataTypeToSerialString(opParamInfo_.sequsedK.desc->GetDataType()).c_str(), + "The dtype of seqused_k only supports int32"), + return ge::GRAPH_FAILED); + } else if (kLayout_ == DataLayout::TND) { + // TND: cu_seqlens_k 必传, seqused_k 可选, cu_seqlens_k 不传 + OP_CHECK_IF(opParamInfo_.cuSeqLensK.tensor == nullptr, + OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON(opName_, "cu_seqlens_k", + "When layout_k is TND, cu_seqlens_k must not be null"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(opParamInfo_.cuSeqLensK.desc->GetDataType() != ge::DT_INT32, + OP_LOGE_FOR_INVALID_DTYPE_WITH_REASON( + opName_, "cu_seqlens_k", + QLIV2DataTypeToSerialString(opParamInfo_.cuSeqLensK.desc->GetDataType()).c_str(), + "The dtype of cu_seqlens_k only supports int32"), + return ge::GRAPH_FAILED); + // seqused_k 可选 - 仅校验数据类型 + if (opParamInfo_.sequsedK.tensor != nullptr) { + OP_CHECK_IF(opParamInfo_.sequsedK.desc->GetDataType() != ge::DT_INT32, + OP_LOGE_FOR_INVALID_DTYPE_WITH_REASON( + opName_, "seqused_k", + QLIV2DataTypeToSerialString(opParamInfo_.sequsedK.desc->GetDataType()).c_str(), + "The dtype of seqused_k only supports int32"), + return ge::GRAPH_FAILED); + } + } else { + // BSND: cu_seqlens_k 不传, seqused_k 可选 + OP_CHECK_IF(opParamInfo_.cuSeqLensK.tensor != nullptr, + OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON(opName_, "cu_seqlens_k", + "When layout_k is BSND, cu_seqlens_k must not be null"), + return ge::GRAPH_FAILED); + if (opParamInfo_.sequsedK.tensor != nullptr) { + OP_CHECK_IF(opParamInfo_.sequsedK.desc->GetDataType() != ge::DT_INT32, + OP_LOGE_FOR_INVALID_DTYPE_WITH_REASON( + opName_, "seqused_k", + QLIV2DataTypeToSerialString(opParamInfo_.sequsedK.desc->GetDataType()).c_str(), + "The dtype of seqused_k only supports int32"), + return ge::GRAPH_FAILED); + } + } + + // block_table 非 PA 场景必须为空 + if (kLayout_ != DataLayout::PA_BBND) { + OP_CHECK_IF(opParamInfo_.blockTable.tensor != nullptr, + OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON( + opName_, "block_table", "When layout_k is not PA_BBND, block_table must not be null"), + return ge::GRAPH_FAILED); + } + + // =============== cmpResidualK 校验 =============== + // cmpRatio 不等于 1 且 sparseMode 不等于 0 时 cmpResidualK 必传 + if (opParamInfo_.cmpRatio != nullptr && *opParamInfo_.cmpRatio != 1 && opParamInfo_.sparseMode != nullptr && + *opParamInfo_.sparseMode != 0) { + OP_CHECK_IF(opParamInfo_.cmpResidualK.tensor == nullptr, + OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON( + opName_, "cmp_residual_k", + "Cmp_ratio is not 1 and sparse_mode is not 0, cmp_residual_k must not be null"), + return ge::GRAPH_FAILED); + // cmpResidualK 传入时校验维度 & 数据类型 + if (qLayout_ == DataLayout::BSND) { + OP_CHECK_IF( + opParamInfo_.query.shape->GetStorageShape().GetDim(DIM_IDX_ZERO) != + opParamInfo_.cmpResidualK.tensor->GetStorageShape().GetShapeSize(), + OP_LOGE_FOR_INVALID_SHAPE_WITH_REASON( + opName_, "cmp_residual_k", ToStringRaw(opParamInfo_.cmpResidualK.tensor->GetStorageShape()).c_str(), + "When layout_q is BSND, the shape of cmp_residual_k must be (B,)"), + return ge::GRAPH_FAILED); + } else if (qLayout_ == DataLayout::TND) { + OP_CHECK_IF(opParamInfo_.cmpResidualK.tensor->GetStorageShape().GetShapeSize() != + opParamInfo_.cuSeqLensQ.tensor->GetStorageShape().GetShapeSize() - 1, + OP_LOGE_FOR_INVALID_SHAPESIZE_WITH_REASON( + opName_, "cmp_residual_k", + std::to_string(opParamInfo_.cmpResidualK.tensor->GetStorageShape().GetShapeSize()), + "When layout_q is TND, the shape size of cmp_residual_k " + "must equal the shape size - 1 of cu_seqlens_q"), + return ge::GRAPH_FAILED); + } + } else { + OP_CHECK_IF(opParamInfo_.cmpResidualK.tensor != nullptr, + OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON( + opName_, "cmp_residual_k", "Cmp_ratio is 1 or sparse_mode is 0, cmp_residual_k must be null"), + return ge::GRAPH_FAILED); + } + // cmpResidualK 传入时校验数据类型 + if (opParamInfo_.cmpResidualK.tensor != nullptr) { + OP_CHECK_IF(opParamInfo_.cmpResidualK.desc->GetDataType() != ge::DT_INT32, + OP_LOGE_FOR_INVALID_DTYPE_WITH_REASON( + opName_, "cmp_residual_k", + QLIV2DataTypeToSerialString(opParamInfo_.cmpResidualK.desc->GetDataType()).c_str(), + "The dtype of cmp_residual_k supports int32"), + return ge::GRAPH_FAILED); + } + + // =============== Q 侧校验 =============== + if (qLayout_ == DataLayout::TND) { + // TND: cu_seqlens_q 必传, seqused_q 可选 + OP_CHECK_IF(opParamInfo_.cuSeqLensQ.tensor == nullptr, + OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON(opName_, "cu_seqlens_q", + "When layout_q is TND, cu_seqlens_q must not be null"), + return ge::GRAPH_FAILED); + if (kLayout_ == DataLayout::PA_BBND) { + // k为PA_BBND必传sequsedK, 用sequsedK的维度校验 + OP_CHECK_IF(opParamInfo_.cuSeqLensQ.tensor->GetStorageShape().GetShapeSize() != + opParamInfo_.sequsedK.tensor->GetStorageShape().GetShapeSize() + 1, + OP_LOGE_FOR_INVALID_SHAPESIZE_WITH_REASON( + opName_, "cmp_residual_k", + std::to_string(opParamInfo_.cmpResidualK.tensor->GetStorageShape().GetShapeSize()), + "When layout_q is TND and layout_k is PA_BBND, " + "the shape size of cu_seqlens_q must equal the shape size + 1 of seqused_k"), + return ge::GRAPH_FAILED); + } else if (kLayout_ == DataLayout::TND) { + // q、k都为TND, cuSeqlensQ与cuSeqlensK维度一致校验 + OP_CHECK_IF(opParamInfo_.cuSeqLensQ.tensor->GetStorageShape().GetShapeSize() != + opParamInfo_.cuSeqLensK.tensor->GetStorageShape().GetShapeSize(), + OP_LOGE_FOR_INVALID_SHAPES_WITH_REASON( + opName_, "cu_seqlens_q and cu_seqlens_k", + Ops::Base::ToString(opParamInfo_.cuSeqLensQ.tensor->GetStorageShape()) + " and " + + Ops::Base::ToString(opParamInfo_.cuSeqLensK.tensor->GetStorageShape()), + "When layout_q is TND and layout_k is TND, " + "the shape of cu_seqlens_q must equal the shape of cu_seqlens_k"), + return ge::GRAPH_FAILED); + } + OP_CHECK_IF(opParamInfo_.cuSeqLensQ.desc->GetDataType() != ge::DT_INT32, + OP_LOGE_FOR_INVALID_DTYPE_WITH_REASON( + opName_, "cu_seqlens_q", + QLIV2DataTypeToSerialString(opParamInfo_.cuSeqLensK.desc->GetDataType()).c_str(), + "The dtype of cu_seqlens_q only supports int32"), + return ge::GRAPH_FAILED); + // seqused_q 可选 - 仅校验数据类型 + if (opParamInfo_.sequsedQ.tensor != nullptr) { + OP_CHECK_IF(opParamInfo_.sequsedQ.desc->GetDataType() != ge::DT_INT32, + OP_LOGE_FOR_INVALID_DTYPE_WITH_REASON( + opName_, "seqused_q", + QLIV2DataTypeToSerialString(opParamInfo_.sequsedQ.desc->GetDataType()).c_str(), + "The dtype of seqused_q only supports int32"), + return ge::GRAPH_FAILED); + } + } else { + // BSND: cu_seqlens_q 不传, seqused_q 可选 + OP_CHECK_IF(opParamInfo_.cuSeqLensQ.tensor != nullptr, + OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON( + opName_, "cu_seqlens_q", "When layout_q is BSND, cu_seqlens_q must not be provided"), + return ge::GRAPH_FAILED); + if (opParamInfo_.sequsedQ.tensor != nullptr) { + OP_CHECK_IF(opParamInfo_.sequsedQ.desc->GetDataType() != ge::DT_INT32, + OP_LOGE_FOR_INVALID_DTYPE_WITH_REASON( + opName_, "seqused_q", + QLIV2DataTypeToSerialString(opParamInfo_.sequsedQ.desc->GetDataType()).c_str(), + "The dtype of seqused_q only supports int32"), + return ge::GRAPH_FAILED); + } + } + if (npuArch_ == NpuArch::DAV_3510) { + if (opParamInfo_.outputIdxOffset.tensor != nullptr) { + OP_CHECK_IF(opParamInfo_.outputIdxOffset.desc->GetDataType() != ge::DT_INT32, + OP_LOGE_FOR_INVALID_DTYPE_WITH_REASON( + opName_, "output_idx_offset", + QLIV2DataTypeToSerialString(opParamInfo_.sequsedQ.desc->GetDataType()).c_str(), + "The dtype of output_idx_offset only supports int32"), + return ge::GRAPH_FAILED); + } + } + // metadata 必传 + OP_CHECK_IF(opParamInfo_.metadata.tensor == nullptr, + OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON(opName_, "metadata", "Metadata must not be null"), + return ge::GRAPH_FAILED); + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QLIV2InfoParser::CheckShapeDim() +{ + OP_CHECK_IF((opParamInfo_.blockTable.tensor != nullptr) && + (opParamInfo_.blockTable.tensor->GetStorageShape().GetDimNum() != DIM_NUM_TWO), + OP_LOGE_FOR_INVALID_SHAPEDIM( + opName_, "block_table", + std::to_string(opParamInfo_.blockTable.tensor->GetStorageShape().GetDimNum()).c_str(), "2"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(((kLayout_ == DataLayout::PA_BBND) || (kLayout_ == DataLayout::BSND)) && + (opParamInfo_.key.shape->GetStorageShape().GetDimNum() != DIM_NUM_FOUR), + OP_LOGE_FOR_INVALID_SHAPEDIM( + opName_, "k", std::to_string(opParamInfo_.key.shape->GetStorageShape().GetDimNum()).c_str(), "4"), + return ge::GRAPH_FAILED); + OP_CHECK_IF( + (kLayout_ == DataLayout::TND) && (opParamInfo_.key.shape->GetStorageShape().GetDimNum() != DIM_NUM_THREE), + OP_LOGE_FOR_INVALID_SHAPEDIM( + opName_, "k", std::to_string(opParamInfo_.key.shape->GetStorageShape().GetDimNum()).c_str(), "3"), + return ge::GRAPH_FAILED); + + uint32_t qShapeDim = opParamInfo_.query.shape->GetStorageShape().GetDimNum(); + uint32_t weightsShapeDim = opParamInfo_.weights.shape->GetStorageShape().GetDimNum(); + uint32_t outShapeDim = opParamInfo_.attenOut.shape->GetStorageShape().GetDimNum(); + uint32_t expectShapeDim = DIM_NUM_FOUR; + if (qLayout_ == DataLayout::TND) { + expectShapeDim = DIM_NUM_THREE; + } + OP_CHECK_IF(qShapeDim != expectShapeDim, + OP_LOGE_FOR_INVALID_SHAPEDIM(opName_, "q", std::to_string(qShapeDim).c_str(), + std::to_string(expectShapeDim).c_str()), + return ge::GRAPH_FAILED); + OP_CHECK_IF(outShapeDim != expectShapeDim, + OP_LOGE_FOR_INVALID_SHAPEDIM(opName_, "sparse_indices", std::to_string(outShapeDim).c_str(), + std::to_string(expectShapeDim).c_str()), + return ge::GRAPH_FAILED); + if (opParamInfo_.outputIdxOffset.tensor != nullptr) { + uint32_t outputIdxOffsetShapeDim = opParamInfo_.outputIdxOffset.tensor->GetStorageShape().GetDimNum(); + OP_CHECK_IF( + (outputIdxOffsetShapeDim != expectShapeDim - 1), + OP_LOGE_FOR_INVALID_SHAPEDIM(opName_, "output_idx_offset", std::to_string(outputIdxOffsetShapeDim).c_str(), + std::to_string(expectShapeDim - 1).c_str()), + return ge::GRAPH_FAILED); + } + if (npuArch_ == NpuArch::DAV_3510 && *opParamInfo_.returnValue == 1) { + uint32_t sparseValuesShapeDim = opParamInfo_.sparseValues.shape->GetStorageShape().GetDimNum(); + OP_CHECK_IF(sparseValuesShapeDim != expectShapeDim, + OP_LOGE_FOR_INVALID_SHAPEDIM(opName_, "sparse_values", std::to_string(sparseValuesShapeDim).c_str(), + std::to_string(expectShapeDim).c_str()), + return ge::GRAPH_FAILED); + } + OP_CHECK_IF(!(weightsShapeDim == expectShapeDim - 1), + OP_LOGE_FOR_INVALID_SHAPEDIM(opName_, "w", std::to_string(weightsShapeDim).c_str(), + std::to_string(expectShapeDim - 1).c_str()), + return ge::GRAPH_FAILED); + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QLIV2InfoParser::GetN1Size() +{ + if (qLayout_ == DataLayout::BSND) { + n1Size_ = static_cast(opParamInfo_.query.shape->GetStorageShape().GetDim(DIM_IDX_TWO)); + } else { + // TND + n1Size_ = static_cast(opParamInfo_.query.shape->GetStorageShape().GetDim(DIM_IDX_ONE)); + } + OP_LOGI(context_->GetNodeName(), "n1Size is %d", n1Size_); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QLIV2InfoParser::GetActualSeqLenSize(uint32_t &size, const gert::Tensor *tensor, + const std::string &actualSeqLenName) const +{ + size = static_cast(tensor->GetShapeSize()); + if (size <= 0) { + OP_LOGE_FOR_INVALID_SHAPESIZE_WITH_REASON( + opName_, actualSeqLenName.c_str(), std::to_string(size).c_str(), + "The shape size of " + actualSeqLenName + " should be greater than 0"); + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QLIV2InfoParser::GetAndCheckN2Size() +{ + // PA_BBND + if (kLayout_ == DataLayout::TND) { + n2Size_ = static_cast(opParamInfo_.key.shape->GetStorageShape().GetDim(DIM_IDX_ONE)); + } else { + n2Size_ = static_cast(opParamInfo_.key.shape->GetStorageShape().GetDim(DIM_IDX_TWO)); + } + OP_LOGI(context_->GetNodeName(), "N2 is %d", n2Size_); + OP_CHECK_IF(n2Size_ != 1, + OP_LOGE_FOR_INVALID_SHAPE_WITH_REASON(opName_, "k", + ToStringRaw(opParamInfo_.key.shape->GetStorageShape()).c_str(), + "The head num of k only supports 1"), + return ge::GRAPH_FAILED); + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QLIV2InfoParser::GetGSize() +{ + if (n1Size_ % n2Size_ != 0) { + OP_LOGE_FOR_INVALID_SHAPES_WITH_REASON(opName_, "q and k", + Ops::Base::ToString(opParamInfo_.query.shape->GetStorageShape()) + + " and " + + Ops::Base::ToString(opParamInfo_.key.shape->GetStorageShape()), + "The head num of q can not be a multiple of the head num of k"); + return ge::GRAPH_FAILED; + } + gSize_ = n1Size_ / n2Size_; + + if (npuArch_ == NpuArch::DAV_3510) { + OP_CHECK_IF(gSize_ > G_SIZE_LIMIT, + OP_LOGE_FOR_INVALID_SHAPES_WITH_REASON( + opName_, "q and k", + Ops::Base::ToString(opParamInfo_.query.shape->GetStorageShape()) + " and " + + Ops::Base::ToString(opParamInfo_.key.shape->GetStorageShape()), + "The value of (the head num of q divided by the head num of k) must <= 64"), + return ge::GRAPH_FAILED); + } else { + OP_CHECK_IF(gSize_ != G_SIZE_LIMIT, + OP_LOGE(opName_, "N1 is %u, N2 is %u, N1 divided by N2 must equal 64.", n1Size_, n2Size_), + return ge::GRAPH_FAILED); + } + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QLIV2InfoParser::GetBatchSize() +{ + // 获取B基准值 + // 1、非TND时, 以query的batch_size维度为基准; + // 2、TND时, 以cu_seqlens_q的shape[0]-1为B轴大小 + if (qLayout_ == DataLayout::BSND) { + bSize_ = opParamInfo_.query.shape->GetStorageShape().GetDim(DIM_IDX_ZERO); + OP_LOGI(context_->GetNodeName(), "b: %d, s: %d, n: %d,d :%d", + opParamInfo_.query.shape->GetStorageShape().GetDim(DIM_IDX_ZERO), + opParamInfo_.query.shape->GetStorageShape().GetDim(DIM_IDX_ONE), + opParamInfo_.query.shape->GetStorageShape().GetDim(DIM_IDX_TWO), + opParamInfo_.query.shape->GetStorageShape().GetDim(DIM_IDX_THREE)); + return ge::GRAPH_SUCCESS; + } else { // TND + // cu_seqlens_q shape is [B+1], batch_size = shape[0] - 1 + uint32_t cuSeqLensQSize = 0; + if (GetActualSeqLenSize(cuSeqLensQSize, opParamInfo_.cuSeqLensQ.tensor, "input cu_seqlens_q") != + ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } + OP_CHECK_IF( + cuSeqLensQSize <= 1, + OP_LOGE_FOR_INVALID_SHAPESIZE_WITH_REASON(opName_, "cu_seqlens_q", std::to_string(cuSeqLensQSize).c_str(), + "The shape size of cu_seqlens_q should be greater than 1 (B+1)"), + return ge::GRAPH_FAILED); + bSize_ = cuSeqLensQSize - 1; + + // Validate key side batch size consistency + if (kLayout_ == DataLayout::TND) { + uint32_t cuSeqLensKSize = 0; + if (GetActualSeqLenSize(cuSeqLensKSize, opParamInfo_.cuSeqLensK.tensor, "cu_seqlens_k") != + ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } + OP_CHECK_IF((cuSeqLensKSize - 1) != bSize_, + OP_LOGE_FOR_INVALID_SHAPES_WITH_REASON( + opName_, "cu_seqlens_q and cu_seqlens_k", + Ops::Base::ToString(opParamInfo_.cuSeqLensK.tensor->GetStorageShape()) + " and " + + Ops::Base::ToString(opParamInfo_.cuSeqLensK.tensor->GetStorageShape()), + "The batch sizes derived from cu_seqlens_q and cu_seqlens_k must be same"), + return ge::GRAPH_FAILED); + } + return ge::GRAPH_SUCCESS; + } +} + +ge::graphStatus QLIV2InfoParser::GetHeadDim() +{ + // 以query的D维度为基准 + uint32_t dIndex = DIM_IDX_TWO; + // 根据layout确定D维度在shape中的位置 + switch (qLayout_) { + case DataLayout::TND: + // TND格式: [Total, N, D] -> D是第2维(索引2) + dIndex = DIM_IDX_TWO; + break; + case DataLayout::BSND: + // BSND格式: [Batch, SeqLen, N, D] -> D是第3维(索引3) + dIndex = DIM_IDX_THREE; + break; + default: + OP_LOGE(opName_, "unsupported layout for getting head dim."); + return ge::GRAPH_FAILED; + } + headDim_ = opParamInfo_.query.shape->GetStorageShape().GetDim(dIndex); + OP_CHECK_IF( + headDim_ != HEAD_DIM_LIMIT, + OP_LOGE_FOR_INVALID_SHAPE_WITH_REASON(opName_, "q", ToStringRaw(opParamInfo_.query.shape->GetStorageShape()), + "The head dim of q only supports 128"), + return ge::GRAPH_FAILED); + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QLIV2InfoParser::GetS1Size() +{ + if (qLayout_ == DataLayout::BSND) { + s1Size_ = opParamInfo_.query.shape->GetStorageShape().GetDim(1); + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QLIV2InfoParser::GetAndCheckBlockSize() +{ + blockSize_ = static_cast(opParamInfo_.key.shape->GetStorageShape().GetDim(1)); + OP_LOGI(context_->GetNodeName(), "blockSize_ is %d", blockSize_); + + OP_CHECK_IF( + ((blockSize_ % BLOCK_SIZE_FACTOR != 0) || (blockSize_ == 0) || (blockSize_ > BLOCK_SIZE_LIMIT)), + OP_LOGE_FOR_INVALID_SHAPE_WITH_REASON(opName_, "k", ToStringRaw(opParamInfo_.key.shape->GetStorageShape()), + "The block_size of k must be a multiple of 16 and belong to (0, 1024]"), + return ge::GRAPH_FAILED); + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QLIV2InfoParser::GetS2SizeForPageAttention() +{ + if (GetAndCheckBlockSize() != ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } + + int32_t blockCount_ = static_cast(opParamInfo_.key.shape->GetStorageShape().GetDim(0)); + OP_CHECK_IF( + (blockCount_ == 0), + OP_LOGE_FOR_INVALID_SHAPE_WITH_REASON(opName_, "k", ToStringRaw(opParamInfo_.key.shape->GetStorageShape()), + "The block_count of k cannot be 0"), + return ge::GRAPH_FAILED); + + maxBlockNumPerBatch_ = opParamInfo_.blockTable.tensor->GetStorageShape().GetDim(1); + s2Size_ = maxBlockNumPerBatch_ * blockSize_; + OP_LOGI(context_->GetNodeName(), "maxBlockNumPerBatch_ is %d, blockSize_ is %d, s2Size_ is %d", + maxBlockNumPerBatch_, blockSize_, s2Size_); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QLIV2InfoParser::GetS2SizeForBatchContinuous() +{ + std::string layout_key(opParamInfo_.layOutKey); + if (kLayout_ == DataLayout::BSND) { + s2Size_ = opParamInfo_.key.shape->GetStorageShape().GetDim(DIM_IDX_ONE); + } else if (kLayout_ == DataLayout::TND) { + s2Size_ = opParamInfo_.key.shape->GetStorageShape().GetDim(DIM_IDX_ZERO); + } + OP_CHECK_IF((kLayout_ != DataLayout::BSND) && (kLayout_ != DataLayout::TND), + OP_LOGE_FOR_INVALID_VALUE(opName_, "layout_k", layout_key.c_str(), "BSND or TND"), + return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QLIV2InfoParser::GetS2Size() +{ + // 获取S2基准值 + // 1、BATCH_CONTINUOUS时, 从key的S轴获取 + // 3、PAGE_ATTENTION时, S2 = block_table.dim1 * block_size + if (kLayout_ == DataLayout::PA_BBND) { + return GetS2SizeForPageAttention(); + } + return GetS2SizeForBatchContinuous(); +} + +ge::graphStatus QLIV2InfoParser::ValidateInputShapesMatch() +{ + /* + TND: + query [T,N1,D], + key [BlockNum,BlockSize,N2,D], + weight [T,N1], + block_table [BatchSize, BatchMaxBlockNum], + act_seq_k [BatchSize] + act_seq_q [BatchSize], + out [T,N2,topk] + ---------------------- + BSND: + query [BatchSize,S1,N1,D], + key [BlockNum,BlockSize,N2,D], + weight [BatchSize,S1,N1], + block_table [BatchSize, BatchMaxBlockNum], + act_seq_k [BatchSize] + act_seq_q [BatchSize] 可选 + out [BatchSize,S1,N2,topk] + */ + uint32_t queryWeightsN1Dim = 1; + uint32_t outN2Dim = 1; + + if (qLayout_ == DataLayout::TND) { + // -----------------------check BatchSize------------------- + // bSize_ 来源于cu_seqlens_q (shape=[B+1], bSize_=B) + OP_CHECK_IF((kLayout_ == DataLayout::PA_BBND) && + ((opParamInfo_.sequsedK.tensor->GetShapeSize() != bSize_) || + (opParamInfo_.blockTable.tensor != nullptr && + opParamInfo_.blockTable.tensor->GetStorageShape().GetDim(0) != bSize_)), + OP_LOGE_FOR_INVALID_SHAPES_WITH_REASON( + opName_, "cu_seqlens_q, seqused_k and block_table", + Ops::Base::ToString(opParamInfo_.cuSeqLensQ.tensor->GetStorageShape()) + ", " + + Ops::Base::ToString(opParamInfo_.sequsedK.tensor->GetStorageShape()) + " and " + + Ops::Base::ToString(opParamInfo_.blockTable.tensor->GetStorageShape()), + "TND case, the dim 0 of cu_seqlens_q, seqused_k and block_table must be same"), + return ge::GRAPH_FAILED); + OP_CHECK_IF((kLayout_ == DataLayout::TND) && (opParamInfo_.cuSeqLensK.tensor->GetShapeSize() != bSize_ + 1), + OP_LOGE_FOR_INVALID_SHAPESIZES_WITH_REASON( + opName_, "cu_seqlens_q and cu_seqlens_k", + std::to_string(opParamInfo_.cuSeqLensQ.tensor->GetStorageShape().GetShapeSize()) + " and " + + std::to_string(opParamInfo_.cuSeqLensK.tensor->GetStorageShape().GetShapeSize()), + "TND case, the shape size of cu_seqlens_q and cu_seqlens_k must be same"), + return ge::GRAPH_FAILED); + // -----------------------check T------------------- + uint32_t qTsize = opParamInfo_.query.shape->GetStorageShape().GetDim(0); + OP_CHECK_IF((opParamInfo_.weights.shape->GetStorageShape().GetDim(0) != qTsize) || + (opParamInfo_.attenOut.shape->GetStorageShape().GetDim(0) != qTsize), + OP_LOGE_FOR_INVALID_SHAPES_WITH_REASON( + opName_, "q, w and sparse_indices", + Ops::Base::ToString(opParamInfo_.query.shape->GetStorageShape()) + ", " + + Ops::Base::ToString(opParamInfo_.weights.shape->GetStorageShape()) + " and " + + Ops::Base::ToString(opParamInfo_.attenOut.shape->GetStorageShape()), + "TND case q, w, sparse_values dim 0 are " + std::to_string(qTsize) + ", " + + std::to_string(opParamInfo_.weights.shape->GetStorageShape().GetDim(0)) + ", " + + std::to_string(opParamInfo_.attenOut.shape->GetStorageShape().GetDim(0)) + + " respectively, they must be same"), + return ge::GRAPH_FAILED); + if (npuArch_ == NpuArch::DAV_3510) { + if (*opParamInfo_.returnValue == 1 && opParamInfo_.sparseValues.shape != nullptr) { + OP_CHECK_IF((opParamInfo_.sparseValues.shape->GetStorageShape().GetDim(0) != qTsize), + OP_LOGE_FOR_INVALID_SHAPES_WITH_REASON( + opName_, "q and sparse_values", + Ops::Base::ToString(opParamInfo_.query.shape->GetStorageShape()) + " and " + + Ops::Base::ToString(opParamInfo_.sparseValues.shape->GetStorageShape()), + "TND case q and sparse_values dim 0 are " + std::to_string(qTsize) + ", " + + std::to_string(opParamInfo_.sparseValues.shape->GetStorageShape().GetDim(0)) + + " respectively, they must be same"), + return ge::GRAPH_FAILED); + } + if (opParamInfo_.outputIdxOffset.tensor != nullptr) { + OP_CHECK_IF((opParamInfo_.outputIdxOffset.tensor->GetStorageShape().GetDim(0) != qTsize), + OP_LOGE_FOR_INVALID_SHAPES_WITH_REASON( + opName_, "q and output_idx_offset", + Ops::Base::ToString(opParamInfo_.query.shape->GetStorageShape()) + " and " + + Ops::Base::ToString(opParamInfo_.outputIdxOffset.tensor->GetStorageShape()), + "TND case q and output_idx_offset dim 0 are " + std::to_string(qTsize) + " and " + + std::to_string(opParamInfo_.outputIdxOffset.tensor->GetStorageShape().GetDim(0)) + + " respectively, they must be same"), + return ge::GRAPH_FAILED); + } + } + } else { + // -----------------------check BatchSize------------------- + // bSize_ 来源于query + OP_CHECK_IF((kLayout_ == DataLayout::PA_BBND) && + ((opParamInfo_.weights.shape->GetStorageShape().GetDim(0) != bSize_) || + (opParamInfo_.blockTable.tensor != nullptr && + opParamInfo_.blockTable.tensor->GetStorageShape().GetDim(0) != bSize_) || + (opParamInfo_.sequsedK.tensor->GetShapeSize() != bSize_) || + (opParamInfo_.attenOut.shape->GetStorageShape().GetDim(0) != bSize_)), + OP_LOGE_FOR_INVALID_SHAPES_WITH_REASON( + opName_, "q, w, seqused_k, block_table and sparse_indices", + Ops::Base::ToString(opParamInfo_.query.shape->GetStorageShape()) + ", " + + Ops::Base::ToString(opParamInfo_.weights.shape->GetStorageShape()) + ", " + + Ops::Base::ToString(opParamInfo_.sequsedK.tensor->GetStorageShape()) + ", " + + Ops::Base::ToString(opParamInfo_.blockTable.tensor->GetStorageShape()) + " and " + + Ops::Base::ToString(opParamInfo_.attenOut.shape->GetStorageShape()), + "BSND case q, w, seqused_k, block_table, sparse_indices dim 0 are " + std::to_string(bSize_) + + ", " + std::to_string(opParamInfo_.weights.shape->GetStorageShape().GetDim(0)) + ", " + + std::to_string(opParamInfo_.sequsedK.tensor->GetStorageShape().GetDim(0)) + ", " + + std::to_string(opParamInfo_.blockTable.tensor->GetStorageShape().GetDim(0)) + ", " + + std::to_string(opParamInfo_.attenOut.shape->GetStorageShape().GetDim(0)) + + " respectively, they must be same"), + return ge::GRAPH_FAILED); + OP_CHECK_IF( + (kLayout_ != DataLayout::PA_BBND) && + ((opParamInfo_.weights.shape->GetStorageShape().GetDim(0) != bSize_) || + (opParamInfo_.sequsedK.tensor != nullptr && opParamInfo_.sequsedK.tensor->GetShapeSize() != bSize_) || + (opParamInfo_.attenOut.shape->GetStorageShape().GetDim(0) != bSize_)), + OP_LOGE_FOR_INVALID_SHAPES_WITH_REASON( + opName_, "q, w, seqused_k and sparse_indices", + Ops::Base::ToString(opParamInfo_.query.shape->GetStorageShape()) + ", " + + Ops::Base::ToString(opParamInfo_.weights.shape->GetStorageShape()) + ", " + + Ops::Base::ToString(opParamInfo_.sequsedK.tensor->GetStorageShape()) + " and " + + Ops::Base::ToString(opParamInfo_.attenOut.shape->GetStorageShape()), + "BSND case q, w, seqused_k, sparse_indices dim 0 are " + std::to_string(bSize_) + ", " + + std::to_string(opParamInfo_.weights.shape->GetStorageShape().GetDim(0)) + ", " + + std::to_string(opParamInfo_.sequsedK.tensor->GetStorageShape().GetDim(0)) + ", " + + std::to_string(opParamInfo_.attenOut.shape->GetStorageShape().GetDim(0)) + + " respectively, they must be same"), + return ge::GRAPH_FAILED); + OP_CHECK_IF( + (opParamInfo_.sequsedQ.tensor != nullptr) && (opParamInfo_.sequsedQ.tensor->GetShapeSize() != bSize_), + OP_LOGE_FOR_INVALID_SHAPES_WITH_REASON( + opName_, "q and seqused_q", + Ops::Base::ToString(opParamInfo_.query.shape->GetStorageShape()) + " and " + + Ops::Base::ToString(opParamInfo_.sequsedQ.tensor->GetStorageShape()), + "BSND case q, seqused_q dim 0 are " + std::to_string(bSize_) + ", " + + std::to_string(opParamInfo_.sequsedQ.tensor->GetStorageShape().GetDim(0)) + + " respectively, they must be same"), + return ge::GRAPH_FAILED); + // -----------------------check S1------------------- + OP_CHECK_IF((opParamInfo_.weights.shape->GetStorageShape().GetDim(1) != s1Size_) || + (opParamInfo_.attenOut.shape->GetStorageShape().GetDim(1) != s1Size_), + OP_LOGE_FOR_INVALID_SHAPES_WITH_REASON( + opName_, "q, w and sparse_indices", + Ops::Base::ToString(opParamInfo_.query.shape->GetStorageShape()) + ", " + + Ops::Base::ToString(opParamInfo_.weights.shape->GetStorageShape()) + " and " + + Ops::Base::ToString(opParamInfo_.attenOut.shape->GetStorageShape()), + "BSND case q, w and sparse_indices dim 1 are " + std::to_string(s1Size_) + ", " + + std::to_string(opParamInfo_.weights.shape->GetStorageShape().GetDim(1)) + ", " + + std::to_string(opParamInfo_.attenOut.shape->GetStorageShape().GetDim(1)) + + " respectively, they must be same"), + return ge::GRAPH_FAILED); + queryWeightsN1Dim = DIM_IDX_TWO; + outN2Dim = DIM_IDX_TWO; + } + // -----------------------check N1------------------- + OP_CHECK_IF((opParamInfo_.weights.shape->GetStorageShape().GetDim(queryWeightsN1Dim) != n1Size_), + OP_LOGE_FOR_INVALID_SHAPES_WITH_REASON( + opName_, "q and w", + Ops::Base::ToString(opParamInfo_.query.shape->GetStorageShape()) + " and " + + Ops::Base::ToString(opParamInfo_.weights.shape->GetStorageShape()), + "BSND case the head num of q, w are " + std::to_string(n1Size_) + ", " + + std::to_string(opParamInfo_.weights.shape->GetStorageShape().GetDim(queryWeightsN1Dim)) + + " respectively, they must be same"), + return ge::GRAPH_FAILED); + // -----------------------check D------------------- + OP_CHECK_IF( + ((kLayout_ != DataLayout::TND && opParamInfo_.key.shape->GetStorageShape().GetDim(DIM_IDX_THREE) != headDim_) || + (kLayout_ == DataLayout::TND && opParamInfo_.key.shape->GetStorageShape().GetDim(DIM_IDX_TWO) != headDim_)), + OP_LOGE_FOR_INVALID_SHAPES_WITH_REASON(opName_, "q and k", + Ops::Base::ToString(opParamInfo_.query.shape->GetStorageShape()) + + " and " + + Ops::Base::ToString(opParamInfo_.key.shape->GetStorageShape()), + "BSND case q, k last dim are " + std::to_string(headDim_) + ", " + + std::to_string(opParamInfo_.key.shape->GetStorageShape().GetDim( + (kLayout_ == DataLayout::TND) ? DIM_IDX_TWO : DIM_IDX_THREE)) + + " respectively, they must be same"), + return ge::GRAPH_FAILED); + // -----------------------check N2------------------- + OP_CHECK_IF((opParamInfo_.attenOut.shape->GetStorageShape().GetDim(outN2Dim) != n2Size_), + OP_LOGE_FOR_INVALID_SHAPES_WITH_REASON( + opName_, "k and sparse_indices", + Ops::Base::ToString(opParamInfo_.key.shape->GetStorageShape()) + " and " + + Ops::Base::ToString(opParamInfo_.attenOut.shape->GetStorageShape()), + "BSND case the head num of k, sparse_indices are " + std::to_string(n2Size_) + ", " + + std::to_string(opParamInfo_.attenOut.shape->GetStorageShape().GetDim(outN2Dim)) + + " respectively, they must be same"), + return ge::GRAPH_FAILED); + // -----------------------check sparse_count------------------- + OP_CHECK_IF((opParamInfo_.attenOut.shape->GetStorageShape().GetDim(outN2Dim + 1) != *opParamInfo_.sparseCount), + OP_LOGE_FOR_INVALID_SHAPES_WITH_REASON( + opName_, "sparse_count and sparse_indices", + Ops::Base::ToString(opParamInfo_.key.shape->GetStorageShape()) + " and " + + Ops::Base::ToString(opParamInfo_.attenOut.shape->GetStorageShape()), + "BSND case sparse_count, sparse_indices last dim are " + std::to_string(*opParamInfo_.sparseCount) + + ", " + std::to_string(opParamInfo_.attenOut.shape->GetStorageShape().GetDim(outN2Dim + 1)) + + " respectively, they must be same"), + return ge::GRAPH_FAILED); + // -----------------------check cmp_residual_k------------------- + if (opParamInfo_.cmpResidualK.tensor != nullptr) { + OP_CHECK_IF( + (opParamInfo_.cmpResidualK.tensor->GetShapeSize() != bSize_), + OP_LOGE_FOR_INVALID_SHAPESIZE_WITH_REASON( + opName_, "cmp_residual_k", std::to_string(opParamInfo_.cmpResidualK.tensor->GetShapeSize()), + "The shape size of cmp_residual_k must be equal to batch_size (" + std::to_string(bSize_) + ")"), + return ge::GRAPH_FAILED); + } + // -----------------------check sparse_values------------------ + if (npuArch_ == NpuArch::DAV_3510) { + if (*opParamInfo_.returnValue == 1) { + OP_CHECK_IF((opParamInfo_.sparseValues.shape->GetStorageShape().GetDim(outN2Dim) != n2Size_), + OP_LOGE_FOR_INVALID_SHAPES_WITH_REASON( + opName_, "k and sparse_values", + Ops::Base::ToString(opParamInfo_.key.shape->GetStorageShape()) + " and " + + Ops::Base::ToString(opParamInfo_.sparseValues.shape->GetStorageShape()), + "The head num of k and sparse_values must be same"), + return ge::GRAPH_FAILED); + OP_CHECK_IF( + (opParamInfo_.sparseValues.shape->GetStorageShape().GetDim(outN2Dim + 1) != *opParamInfo_.sparseCount), + OP_LOGE_FOR_INVALID_SHAPES_WITH_REASON( + opName_, "topk and sparse_values", + std::to_string(*opParamInfo_.sparseCount) + " and " + + Ops::Base::ToString(opParamInfo_.sparseValues.shape->GetStorageShape()), + "The last dim of sparse_values must be same as topk"), + return ge::GRAPH_FAILED); + } + } + // -----------------------check metadata------------------- + OP_CHECK_IF((opParamInfo_.metadata.tensor->GetShapeSize() != METADATA_LIMIT), + OP_LOGE_FOR_INVALID_SHAPESIZE_WITH_REASON( + opName_, "metadata", std::to_string(opParamInfo_.metadata.tensor->GetShapeSize()).c_str(), + "The shape size of metadata must be " + std::to_string(METADATA_LIMIT)), + return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QLIV2InfoParser::CheckScaleShape() +{ + uint32_t qShapeDim = opParamInfo_.query.shape->GetStorageShape().GetDimNum(); + uint32_t kShapeDim = opParamInfo_.key.shape->GetStorageShape().GetDimNum(); + uint32_t qDequantScaleShapeDim = opParamInfo_.query_dequant_scale.shape->GetStorageShape().GetDimNum(); + uint32_t kDequantScaleShapeDim = opParamInfo_.key_dequant_scale.shape->GetStorageShape().GetDimNum(); + bool isMxQuantMode = (*opParamInfo_.quantMode == QUANT_MODE_MXFP8) || (*opParamInfo_.quantMode == QUANT_MODE_MXFP4); + + if (*opParamInfo_.quantMode == QUANT_MODE_HIFLOAT8) { + OP_CHECK_IF(qDequantScaleShapeDim != 1, + OP_LOGE_FOR_INVALID_SHAPEDIM_WITH_REASON( + opName_, "q_descale", std::to_string(qDequantScaleShapeDim).c_str(), + "When quant_mode is 4, the dim num of q_descale should be 1"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(kDequantScaleShapeDim != 1, + OP_LOGE_FOR_INVALID_SHAPEDIM_WITH_REASON( + opName_, "k_descale", std::to_string(kDequantScaleShapeDim).c_str(), + "When quant_mode is 4, the dim num of k_descale should be 1"), + return ge::GRAPH_FAILED); + OP_CHECK_IF( + opParamInfo_.query_dequant_scale.shape->GetStorageShape().GetDim(0) != 1, + OP_LOGE_FOR_INVALID_SHAPE_WITH_REASON( + opName_, "q_descale", ToStringRaw(opParamInfo_.query_dequant_scale.shape->GetStorageShape()).c_str(), + "When quant_mode is 4, q_descale's shape[0] should be 1, but now is " + + std::to_string(opParamInfo_.query_dequant_scale.shape->GetStorageShape().GetDim(0))), + return ge::GRAPH_FAILED); + OP_CHECK_IF( + opParamInfo_.key_dequant_scale.shape->GetStorageShape().GetDim(0) != 1, + OP_LOGE_FOR_INVALID_SHAPE_WITH_REASON( + opName_, "k_descale", ToStringRaw(opParamInfo_.key_dequant_scale.shape->GetStorageShape()).c_str(), + "When quant_mode is 4, k_descale's shape[0] should be 1, but now is " + + std::to_string(opParamInfo_.key_dequant_scale.shape->GetStorageShape().GetDim(0))), + return ge::GRAPH_FAILED); + } else if (isMxQuantMode) { + OP_CHECK_IF(headDim_ % MX_SCALE_SHAPE_ALIGN != 0, + OP_LOGE_FOR_INVALID_SHAPE_WITH_REASON( + opName_, "q", ToStringRaw(opParamInfo_.query.shape->GetStorageShape()).c_str(), + "When quant_mode is " + std::to_string(*opParamInfo_.quantMode) + + ", head_dim should be a multiple of " + std::to_string(MX_SCALE_SHAPE_ALIGN) + + ", but now is " + std::to_string(headDim_)), + return ge::GRAPH_FAILED); + OP_CHECK_IF(qDequantScaleShapeDim != (qShapeDim + 1), + OP_LOGE_FOR_INVALID_SHAPEDIM_WITH_REASON( + opName_, "q_descale", std::to_string(qDequantScaleShapeDim).c_str(), + "When quant_mode is " + std::to_string(*opParamInfo_.quantMode) + + ", the dim num of q_descale should be " + std::to_string(qShapeDim + 1)), + return ge::GRAPH_FAILED); + OP_CHECK_IF(kDequantScaleShapeDim != (kShapeDim + 1), + OP_LOGE_FOR_INVALID_SHAPEDIM_WITH_REASON( + opName_, "k_descale", std::to_string(qDequantScaleShapeDim).c_str(), + "When quant_mode is " + std::to_string(*opParamInfo_.quantMode) + + ", the dim num of k_descale should be " + std::to_string(kShapeDim + 1)), + return ge::GRAPH_FAILED); + for (uint32_t i = 0; i < (qShapeDim - 1); i++) { + uint32_t dimValueQueryScale = opParamInfo_.query_dequant_scale.shape->GetStorageShape().GetDim(i); + uint32_t dimValueQuery = opParamInfo_.query.shape->GetStorageShape().GetDim(i); + OP_CHECK_IF( + dimValueQueryScale != dimValueQuery, + OP_LOGE_FOR_INVALID_SHAPES_WITH_REASON( + opName_, "q and q_descale", + Ops::Base::ToString(opParamInfo_.query.shape->GetStorageShape()) + " and " + + Ops::Base::ToString(opParamInfo_.query_dequant_scale.shape->GetStorageShape()), + "Q_descale's shape[" + std::to_string(i) + "] " + std::to_string(dimValueQueryScale) + + " and q's shape[" + std::to_string(i) + "] " + std::to_string(dimValueQuery) + " are not same"), + return ge::GRAPH_FAILED); + } + for (uint32_t i = 0; i < (kShapeDim - 1); i++) { + uint32_t dimValueKeyScale = opParamInfo_.key_dequant_scale.shape->GetStorageShape().GetDim(i); + uint32_t dimValueKey = opParamInfo_.key.shape->GetStorageShape().GetDim(i); + OP_CHECK_IF(dimValueKeyScale != dimValueKey, + OP_LOGE_FOR_INVALID_SHAPES_WITH_REASON( + opName_, "k and k_descale", + Ops::Base::ToString(opParamInfo_.key.shape->GetStorageShape()) + " and " + + Ops::Base::ToString(opParamInfo_.key_dequant_scale.shape->GetStorageShape()), + "K_descale's shape[" + std::to_string(i) + "] " + std::to_string(dimValueKeyScale) + + " and k's shape[" + std::to_string(i) + "] " + std::to_string(dimValueKey) + + " are not the same"), + return ge::GRAPH_FAILED); + } + uint32_t expectScaleD = headDim_ / MX_SCALE_SHAPE_ALIGN; + OP_CHECK_IF( + (opParamInfo_.query_dequant_scale.shape->GetStorageShape().GetDim(qShapeDim - 1) != expectScaleD) || + (opParamInfo_.query_dequant_scale.shape->GetStorageShape().GetDim(qShapeDim) != MX_E8M0_SCALE_PACK_NUM), + OP_LOGE_FOR_INVALID_SHAPES_WITH_REASON( + opName_, "q_descale", Ops::Base::ToString(opParamInfo_.query_dequant_scale.shape->GetStorageShape()), + "When quant_mode is " + std::to_string(*opParamInfo_.quantMode) + + ", q_descale's last dims should be [" + std::to_string(expectScaleD) + ", " + + std::to_string(MX_E8M0_SCALE_PACK_NUM) + "]"), + return ge::GRAPH_FAILED); + OP_CHECK_IF( + (opParamInfo_.key_dequant_scale.shape->GetStorageShape().GetDim(kShapeDim - 1) != expectScaleD) || + (opParamInfo_.key_dequant_scale.shape->GetStorageShape().GetDim(kShapeDim) != MX_E8M0_SCALE_PACK_NUM), + OP_LOGE_FOR_INVALID_SHAPES_WITH_REASON( + opName_, "k_descale", Ops::Base::ToString(opParamInfo_.key_dequant_scale.shape->GetStorageShape()), + "When quant_mode is " + std::to_string(*opParamInfo_.quantMode) + + ", k_descale's last dims should be [" + std::to_string(expectScaleD) + ", " + + std::to_string(MX_E8M0_SCALE_PACK_NUM) + "]"), + return ge::GRAPH_FAILED); + } else { + OP_CHECK_IF(qDequantScaleShapeDim != (qShapeDim - 1), + OP_LOGE_FOR_INVALID_SHAPEDIM(opName_, "q_descale", std::to_string(qDequantScaleShapeDim).c_str(), + std::to_string(qShapeDim - 1).c_str()), + return ge::GRAPH_FAILED); + OP_CHECK_IF(kDequantScaleShapeDim != (kShapeDim - 1), + OP_LOGE_FOR_INVALID_SHAPEDIM(opName_, "k_descale", std::to_string(kDequantScaleShapeDim).c_str(), + std::to_string(kShapeDim - 1).c_str()), + return ge::GRAPH_FAILED); + // check q scale + for (uint32_t i = 0; i < (qShapeDim - 1); i++) { + uint32_t dimValueQueryScale = opParamInfo_.query_dequant_scale.shape->GetStorageShape().GetDim(i); + uint32_t dimValueQuery = opParamInfo_.query.shape->GetStorageShape().GetDim(i); + OP_CHECK_IF(dimValueQueryScale != dimValueQuery, + OP_LOGE_FOR_INVALID_SHAPES_WITH_REASON( + opName_, "q_descale and q", + Ops::Base::ToString(opParamInfo_.query_dequant_scale.shape->GetStorageShape()) + " and " + + Ops::Base::ToString(opParamInfo_.query.shape->GetStorageShape()), + "q_descale's shape[" + std::to_string(i) + "] and q's shape[" + std::to_string(i) + + "] are not the same"), + return ge::GRAPH_FAILED); + } + // check k scale + for (uint32_t i = 0; i < (kShapeDim - 1); i++) { + uint32_t dimValueKeyScale = opParamInfo_.key_dequant_scale.shape->GetStorageShape().GetDim(i); + uint32_t dimValueKey = opParamInfo_.key.shape->GetStorageShape().GetDim(i); + OP_CHECK_IF(dimValueKeyScale != dimValueKey, + OP_LOGE_FOR_INVALID_SHAPES_WITH_REASON( + opName_, "k_descale and k", + Ops::Base::ToString(opParamInfo_.key_dequant_scale.shape->GetStorageShape()) + " and " + + Ops::Base::ToString(opParamInfo_.key.shape->GetStorageShape()), + "k_descale's shape[" + std::to_string(i) + "] and k's shape[" + std::to_string(i) + + "] are not the same"), + return ge::GRAPH_FAILED); + } + } + + return ge::GRAPH_SUCCESS; +} + +// key非连续校验：通过shape计算expected stride进行校验 +// PA_BBND时，只允许0轴非连续，其余轴必须连续 +// 非PA_BBND时，所有轴都必须连续 +ge::graphStatus QLIV2InfoParser::CheckKeyContiguous() const +{ + bool keyNonContiguous = false; + bool scaleNonContiguous = false; + // A5 PA_BBND: 0轴允许非连续，从1轴开始检查；非PA_BBND或A3: 从0轴开始检查 + // PA_BBND: axis 0 allows non-contiguous, check starts from axis 1 + // Non-PA_BBND: check starts from axis 0 + size_t checkStartIdx = (npuArch_ == NpuArch::DAV_3510 && kLayout_ == DataLayout::PA_BBND) ? 1 : 0; + if (!keyStridesVec_.empty() && opParamInfo_.key.shape != nullptr) { + auto &shape = opParamInfo_.key.shape->GetStorageShape(); + std::vector expectedStrides; + if (kLayout_ == DataLayout::BSND || kLayout_ == DataLayout::PA_BBND) { + expectedStrides = {shape.GetDim(1) * shape.GetDim(2) * shape.GetDim(3), shape.GetDim(2) * shape.GetDim(3), + shape.GetDim(3), 1}; + } else if (kLayout_ == DataLayout::TND) { + expectedStrides = {shape.GetDim(1) * shape.GetDim(2), shape.GetDim(2), 1}; + } + for (size_t i = checkStartIdx; i < expectedStrides.size(); ++i) { + if (i < keyStridesVec_.size() && keyStridesVec_[i] != expectedStrides[i]) { + keyNonContiguous = true; + break; + } + } + } + bool isMxQuantMode = (*opParamInfo_.quantMode == QUANT_MODE_MXFP8) || (*opParamInfo_.quantMode == QUANT_MODE_MXFP4); + if ((*opParamInfo_.quantMode != QUANT_MODE_HIFLOAT8) && !keyDequantScaleStridesVec_.empty() && + opParamInfo_.key_dequant_scale.shape != nullptr) { + auto &shape = opParamInfo_.key_dequant_scale.shape->GetStorageShape(); + std::vector expectedStrides; + if (isMxQuantMode) { + if (kLayout_ == DataLayout::BSND || kLayout_ == DataLayout::PA_BBND) { + expectedStrides = {shape.GetDim(1) * shape.GetDim(2) * shape.GetDim(3) * shape.GetDim(4), + shape.GetDim(2) * shape.GetDim(3) * shape.GetDim(4), + shape.GetDim(3) * shape.GetDim(4), shape.GetDim(4), 1}; + } else if (kLayout_ == DataLayout::TND) { + expectedStrides = {shape.GetDim(1) * shape.GetDim(2) * shape.GetDim(3), + shape.GetDim(2) * shape.GetDim(3), shape.GetDim(3), 1}; + } + } else { + if (kLayout_ == DataLayout::BSND || kLayout_ == DataLayout::PA_BBND) { + expectedStrides = {shape.GetDim(1) * shape.GetDim(2), shape.GetDim(2), 1}; + } else if (kLayout_ == DataLayout::TND) { + expectedStrides = {shape.GetDim(1), 1}; + } + } + for (size_t i = checkStartIdx; i < expectedStrides.size(); ++i) { + if (i < keyDequantScaleStridesVec_.size() && keyDequantScaleStridesVec_[i] != expectedStrides[i]) { + scaleNonContiguous = true; + break; + } + } + } + if (kLayout_ == DataLayout::PA_BBND) { + if (!keyStridesVec_.empty()) { + if (*opParamInfo_.quantMode == QUANT_MODE_MXFP4) { + OP_CHECK_IF( + keyStridesVec_[0] <= 0 || keyStridesVec_[0] % MXFP4_PACK_NUM != 0, + OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON( + opName_, "k", + "When quant_mode is 5 and layout_k is PA_BBND, key stride0 must be positive and satisfy " + "2-element FP4 packing alignment, but got " + + std::to_string(keyStridesVec_[0])), + return ge::GRAPH_FAILED); + } else { + OP_CHECK_IF(keyStridesVec_[0] <= 0, + OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON( + opName_, "k", + "When layout_k is PA_BBND, key stride0 must be positive, but got " + + std::to_string(keyStridesVec_[0])), + return ge::GRAPH_FAILED); + } + } + if (isMxQuantMode && !keyDequantScaleStridesVec_.empty()) { + OP_CHECK_IF( + keyDequantScaleStridesVec_[0] <= 0 || keyDequantScaleStridesVec_[0] % MX_E8M0_SCALE_PACK_NUM != 0, + OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON( + opName_, "k_descale", + "When quant_mode is 3 or 5 and layout_k is PA_BBND, key_dequant_scale stride0 must be positive " + "and satisfy 2-element E8M0 packing alignment, but got " + + std::to_string(keyDequantScaleStridesVec_[0])), + return ge::GRAPH_FAILED); + } + } + OP_CHECK_IF(keyNonContiguous || scaleNonContiguous, + OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON( + opName_, "k and k_descale", "k and k_descale only support non-continuous keying on the 0-axis"), + return ge::GRAPH_FAILED); + + return ge::GRAPH_SUCCESS; +} + +void QLIV2InfoParser::GenerateInfo(QLIV2TilingInfo &QLIV2Info) +{ + QLIV2Info.opName = opName_; + QLIV2Info.platformInfo = platformInfo_; + QLIV2Info.opParamInfo = opParamInfo_; + QLIV2Info.socVersion = socVersion_; + QLIV2Info.npuArch = npuArch_; + + QLIV2Info.bSize = bSize_; + QLIV2Info.n1Size = n1Size_; + QLIV2Info.n2Size = n2Size_; + QLIV2Info.s1Size = s1Size_; + QLIV2Info.s2Size = s2Size_; + QLIV2Info.gSize = gSize_; + + QLIV2Info.inputQType = inputQType_; + QLIV2Info.inputKType = inputKType_; + QLIV2Info.outputType = outputType_; + + QLIV2Info.blockSize = blockSize_; + QLIV2Info.maxBlockNumPerBatch = maxBlockNumPerBatch_; + + QLIV2Info.pageAttentionFlag = (kLayout_ == DataLayout::PA_BBND); + QLIV2Info.sparseMode = *opParamInfo_.sparseMode; + QLIV2Info.sparseCount = *opParamInfo_.sparseCount; + QLIV2Info.cmpRatio = *opParamInfo_.cmpRatio; + QLIV2Info.returnValue = *opParamInfo_.returnValue; + QLIV2Info.maxSeqlenQ = (opParamInfo_.maxSeqlenQ != nullptr) ? *opParamInfo_.maxSeqlenQ : -1; + + QLIV2Info.keyStridesVec = keyStridesVec_; + QLIV2Info.keyDequantScaleStridesVec = keyDequantScaleStridesVec_; + if (!keyStridesVec_.empty()) { + uint32_t keyStride0 = static_cast(keyStridesVec_[0]); + if (*opParamInfo_.quantMode == QUANT_MODE_MXFP4) { + // FP4 shape stride以逻辑元素计数，kernel侧以打包后的uint8为寻址单位。 + keyStride0 /= MXFP4_PACK_NUM; + } + QLIV2Info.keyStride0 = keyStride0; + } else { + QLIV2Info.keyStride0 = 0; // 非PA无需使用stride + } + if (!keyDequantScaleStridesVec_.empty()) { + QLIV2Info.keyDequantScaleStride0 = static_cast(keyDequantScaleStridesVec_[0]); + } else if ((*opParamInfo_.quantMode == QUANT_MODE_MXFP8) || (*opParamInfo_.quantMode == QUANT_MODE_MXFP4)) { + QLIV2Info.keyDequantScaleStride0 = static_cast(blockSize_) * (headDim_ / MX_SCALE_GROUP_SIZE); + } else { + QLIV2Info.keyDequantScaleStride0 = 0; + } + + QLIV2Info.inputQLayout = qLayout_; + QLIV2Info.inputKLayout = kLayout_; +} + +ge::graphStatus QLIV2InfoParser::ParseAndCheck(QLIV2TilingInfo &QLIV2Info) +{ + if (ge::GRAPH_SUCCESS != GetOpName() || ge::GRAPH_SUCCESS != GetNpuInfo() || ge::GRAPH_SUCCESS != GetOpParaInfo() || + ge::GRAPH_SUCCESS != CheckRequiredParaExistence()) { + return ge::GRAPH_FAILED; + } + + if (ge::GRAPH_SUCCESS != GetAndCheckInOutDataType() || ge::GRAPH_SUCCESS != GetQueryKeyAndOutLayout() || + ge::GRAPH_SUCCESS != GetAndCheckOptionalInput()) { + return ge::GRAPH_FAILED; + } + + if (ge::GRAPH_SUCCESS != CheckShapeDim() || ge::GRAPH_SUCCESS != GetN1Size() || + ge::GRAPH_SUCCESS != GetAndCheckN2Size() || ge::GRAPH_SUCCESS != GetGSize()) { + return ge::GRAPH_FAILED; + } + + if (ge::GRAPH_SUCCESS != GetBatchSize() || ge::GRAPH_SUCCESS != GetS1Size() || ge::GRAPH_SUCCESS != GetHeadDim() || + ge::GRAPH_SUCCESS != GetS2Size()) { + return ge::GRAPH_FAILED; + } + if (ge::GRAPH_SUCCESS != ValidateInputShapesMatch() || ge::GRAPH_SUCCESS != CheckScaleShape() || + ge::GRAPH_SUCCESS != CheckKeyContiguous()) { + return ge::GRAPH_FAILED; + } + + GenerateInfo(QLIV2Info); + + return ge::GRAPH_SUCCESS; +} + +// --------------------------TilingPrepare函数定义------------------------------------- +static ge::graphStatus TilingPrepareForQuantLightningIndexerV2(gert::TilingParseContext * /* context */) +{ + return ge::GRAPH_SUCCESS; +} + +// --------------------------QuantLightningIndexerV2Tiling类成员函数定义----------------------- +ge::graphStatus QuantLightningIndexerV2Tiling::DoTiling(QLIV2TilingInfo *tilingInfo) +{ + // -------------set blockdim----------------- + auto ascendcPlatform = platform_ascendc::PlatformAscendC(tilingInfo->platformInfo); + uint32_t aivNum = ascendcPlatform.GetCoreNumAiv(); + uint32_t aicNum = ascendcPlatform.GetCoreNumAic(); + uint32_t blockDim = ascendcPlatform.CalcTschBlockDim(aivNum, aicNum, aivNum); + context_->SetBlockDim(blockDim); + + // -------------set workspacesize----------------- + constexpr uint32_t MM1_RES_ELEM_SIZE = 4; // 4: fp32 + constexpr uint32_t DOUBLE_BUFFER = 2; // 双Buffer + constexpr uint32_t M_BASE_SIZE = 256; // m轴基本块大小 + constexpr uint32_t S2_BASE_SIZE = 2048; // S2轴基本块大小 + constexpr uint32_t V1_RES_ELEM_SIZE = 4; // 4: int32 + constexpr uint32_t V1_RES_ELEM_TYPE = 2; // 保留Index和Value 2种数据 + constexpr uint32_t V1_DECODE_PARAM_ELEM_SIZE = 8; // 8: int64 + constexpr uint32_t V1_DECODE_PARAM_NUM = 16; // Decode参数个数 + constexpr uint32_t V1_DECODE_DATA_NUM = 2; // Decode每个核需要存储头和尾部两块数据 + constexpr uint32_t S1_BASE_SIZE = 4; // S1轴基本块的大小 + constexpr uint32_t TOPK_MAX_SIZE = 2048; // TopK选取个数 + constexpr uint32_t TOPK_MAX_SIZE_950 = 8192; // A5 TopK最大选取个数 + uint64_t workspaceSize = ascendcPlatform.GetLibApiWorkSpaceSize(); + // 主流程需Workspace大小 + if (ascendcPlatform.GetCurNpuArch() == NpuArch::DAV_3510) { + constexpr uint32_t S1_BASE_SIZE_950 = 4; + constexpr uint32_t S2_BASE_SIZE_950 = 128; + workspaceSize += S1_BASE_SIZE_950 * ((tilingInfo->s2Size + S2_BASE_SIZE_950 - 1) / S2_BASE_SIZE_950) * + S2_BASE_SIZE_950 * sizeof(uint16_t) * aicNum; + // 临时存储Decode中间结果大小: 2(头/尾)*8(s1Base)*2(idx/value)*2048(K)*sizeof(int32)*24=6M + workspaceSize += + V1_DECODE_DATA_NUM * S1_BASE_SIZE * V1_RES_ELEM_TYPE * TOPK_MAX_SIZE_950 * V1_RES_ELEM_SIZE * aicNum; + // 临时存储Decode中间参数信息大小: 2(头/尾)*8(s1Base)*16(paramNum)*sizeof(int64_t)*24=48k + workspaceSize += V1_DECODE_DATA_NUM * S1_BASE_SIZE * V1_DECODE_PARAM_NUM * V1_DECODE_PARAM_ELEM_SIZE * aicNum; + } else { + uint32_t mm1ResSize = M_BASE_SIZE * S2_BASE_SIZE; + workspaceSize += mm1ResSize * MM1_RES_ELEM_SIZE * DOUBLE_BUFFER * aicNum; + // Decode流程(LD)需要Workspace大小 + // 临时存储Decode中间结果大小: 2(头/尾)*8(s1Base)*2(idx/value)*2048(K)*sizeof(int32)*24=6M + workspaceSize += + V1_DECODE_DATA_NUM * S1_BASE_SIZE * V1_RES_ELEM_TYPE * TOPK_MAX_SIZE * V1_RES_ELEM_SIZE * aicNum; + // 临时存储Decode中间参数信息大小: 2(头/尾)*8(s1Base)*16(paramNum)*sizeof(int64_t)*24=48k + workspaceSize += V1_DECODE_DATA_NUM * S1_BASE_SIZE * V1_DECODE_PARAM_NUM * V1_DECODE_PARAM_ELEM_SIZE * aicNum; + } + size_t *workSpaces = context_->GetWorkspaceSizes(1); + workSpaces[0] = workspaceSize; + + // -------------set tilingdata----------------- + tilingData_.set_bSize(tilingInfo->bSize); + tilingData_.set_s2Size(tilingInfo->s2Size); + tilingData_.set_s1Size(tilingInfo->s1Size); + tilingData_.set_sparseCount(tilingInfo->sparseCount); + tilingData_.set_gSize(tilingInfo->gSize); + tilingData_.set_blockSize(tilingInfo->blockSize); + tilingData_.set_maxBlockNumPerBatch(tilingInfo->maxBlockNumPerBatch); + tilingData_.set_sparseMode(tilingInfo->sparseMode); + tilingData_.set_cmpRatio(tilingInfo->cmpRatio); + tilingData_.set_returnValue(tilingInfo->returnValue); + tilingData_.set_maxSeqlenQ(tilingInfo->maxSeqlenQ); + tilingData_.set_keyStride0(tilingInfo->keyStride0); + tilingData_.set_keyDequantScaleStride0(tilingInfo->keyDequantScaleStride0); + tilingData_.set_quantMode(*tilingInfo->opParamInfo.quantMode); + tilingData_.set_usedCoreNum(blockDim); + tilingData_.SaveToBuffer(context_->GetRawTilingData()->GetData(), context_->GetRawTilingData()->GetCapacity()); + context_->GetRawTilingData()->SetDataSize(tilingData_.GetDataSize()); + + // -------------set tilingkey----------------- + // DT_Q, DT_KV, DT_OUT, PAGE_ATTENTION, FLASH_DECODE, LAYOUT_T, KV_LAYOUT_T + uint32_t inputQType = static_cast(tilingInfo->inputQType); + uint32_t inputKType = static_cast(tilingInfo->inputKType); + uint32_t outputType = static_cast(tilingInfo->outputType); + uint32_t pageAttentionFlag = static_cast(tilingInfo->pageAttentionFlag); + uint32_t inputQLayout = static_cast(tilingInfo->inputQLayout); + uint32_t inputKLayout = static_cast(tilingInfo->inputKLayout); + uint64_t tilingKey = + GET_TPL_TILING_KEY(inputQType, inputKType, outputType, pageAttentionFlag, inputQLayout, inputKLayout); + context_->SetTilingKey(tilingKey); + context_->SetScheduleMode(1); + + return ge::GRAPH_SUCCESS; +} + +// --------------------------Tiling函数定义--------------------------- +ge::graphStatus TilingForQuantLightningIndexerV2(gert::TilingContext *context) +{ + OP_CHECK_IF(context == nullptr, OP_LOGE("QuantLightningIndexerV2", "Tiling context is null."), + return ge::GRAPH_FAILED); + QLIV2TilingInfo QLIV2Info; + QLIV2InfoParser QLIV2InfoParser(context); + if (QLIV2InfoParser.ParseAndCheck(QLIV2Info) != ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } + QuantLightningIndexerV2Tiling QLIV2Tiling(context); + return QLIV2Tiling.DoTiling(&QLIV2Info); +} + +// --------------------------Tiling及函数TilingPrepare函数注册-------- +IMPL_OP_OPTILING(QuantLightningIndexerV2) + .Tiling(TilingForQuantLightningIndexerV2) + .TilingParse(TilingPrepareForQuantLightningIndexerV2); + +} // namespace optiling diff --git a/xllm_ops/attention/quant_lightning_indexer_v2/op_host/quant_lightning_indexer_v2_tiling.h b/xllm_ops/attention/quant_lightning_indexer_v2/op_host/quant_lightning_indexer_v2_tiling.h new file mode 100644 index 0000000..2052afa --- /dev/null +++ b/xllm_ops/attention/quant_lightning_indexer_v2/op_host/quant_lightning_indexer_v2_tiling.h @@ -0,0 +1,279 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file quant_lightning_indexer_v2_tiling.h + * \brief + */ + +#ifndef QUANT_LIGHTNING_INDEXER_V2_TILING_H +#define QUANT_LIGHTNING_INDEXER_V2_TILING_H + +#include "err/ops_err.h" +#include "exe_graph/runtime/tiling_context.h" +#include "platform/platform_info.h" +#include "register/op_def_registry.h" +#include "register/tilingdata_base.h" +#include "tiling/platform/platform_ascendc.h" +#include "tiling/tiling_api.h" +#include "op_host/tiling_util.h" + +namespace optiling { +// ------------------公共定义-------------------------- +struct TilingRequiredParaInfo { + const gert::CompileTimeTensorDesc *desc; + const gert::StorageShape *shape; +}; + +struct TilingOptionalParaInfo { + const gert::CompileTimeTensorDesc *desc; + const gert::Tensor *tensor; +}; + +enum class DataLayout : uint32_t { BSND = 0, TND = 1, PA_BBND = 2 }; + +// ------------------算子原型索引常量定义---------------- +// Inputs Index +constexpr uint32_t QUERY_INDEX = 0; +constexpr uint32_t KEY_INDEX = 1; +constexpr uint32_t WEIGTHS_INDEX = 2; +constexpr uint32_t QUERY_DEQUANT_SCALE_INDEX = 3; +constexpr uint32_t KEY_DEQUANT_SCALE_INDEX = 4; +constexpr uint32_t CU_SEQLENS_Q_INDEX = 5; +constexpr uint32_t CU_SEQLENS_K_INDEX = 6; +constexpr uint32_t SEQUSED_Q_INDEX = 7; +constexpr uint32_t SEQUSED_K_INDEX = 8; +constexpr uint32_t CMP_RESIDUAL_K_INDEX = 9; +constexpr uint32_t BLOCK_TABLE_INDEX = 10; +constexpr uint32_t OUTPUT_IDX_OFFSET_INDEX = 11; +constexpr uint32_t METADATA_INDEX = 12; +constexpr uint32_t SPARSE_INDICES_INDEX = 0; +constexpr uint32_t SPARSE_VALUES_INDEX = 1; +// Attributes Index +constexpr uint32_t ATTR_TOPK_INDEX = 0; +constexpr uint32_t ATTR_QUANT_MODE_INDEX = 1; +constexpr uint32_t ATTR_MAX_SEQLEN_Q_INDEX = 2; +constexpr uint32_t ATTR_QUERY_LAYOUT_INDEX = 3; +constexpr uint32_t ATTR_KEY_LAYOUT_INDEX = 4; +constexpr uint32_t ATTR_MASK_MODE_INDEX = 5; +constexpr uint32_t ATTR_CMP_RATIO_INDEX = 6; +constexpr uint32_t ATTR_RETURN_VALUE_INDEX = 7; + +// Dim Index +constexpr uint32_t DIM_IDX_ZERO = 0; +constexpr uint32_t DIM_IDX_ONE = 1; +constexpr uint32_t DIM_IDX_TWO = 2; +constexpr uint32_t DIM_IDX_THREE = 3; +// Dim Num +constexpr uint32_t DIM_NUM_TWO = 2; +constexpr uint32_t DIM_NUM_THREE = 3; +constexpr uint32_t DIM_NUM_FOUR = 4; +// 入参限制常量 +constexpr uint32_t HEAD_DIM_LIMIT = 128; +constexpr uint32_t SPARSE_LIMIT = 2048; +constexpr uint32_t SPARSE_LIMIT_8K = 8192; +constexpr uint32_t G_SIZE_LIMIT = 64; +constexpr uint32_t G_SIZE_LIMIT_32_950 = 32; +constexpr uint32_t BLOCK_SIZE_LIMIT = 1024; +constexpr uint32_t BLOCK_SIZE_FACTOR = 16; +constexpr uint32_t SPARSE_MODE_LOWER = 3; +constexpr uint32_t METADATA_LIMIT = 1024; +constexpr int32_t QUANT_MODE_FP8 = 1; +constexpr int32_t QUANT_MODE_INT8 = 2; +constexpr int32_t QUANT_MODE_MXFP8 = 3; +constexpr int32_t QUANT_MODE_HIFLOAT8 = 4; +constexpr int32_t QUANT_MODE_MXFP4 = 5; +constexpr uint32_t MX_SCALE_SHAPE_ALIGN = 64; +constexpr uint32_t MX_E8M0_SCALE_PACK_NUM = 2; // MX的E8M0 scale形状最后一维打包数为2 +constexpr uint32_t MXFP4_PACK_NUM = 2; // 每个uint8承载2个FP4 E2M1逻辑元素 +constexpr uint32_t MX_SCALE_GROUP_SIZE = 32; // MX量化每32个D维元素对应1个E8M0 scale + +// -----------算子TilingData定义--------------- +BEGIN_TILING_DATA_DEF(QLIV2TilingData) +TILING_DATA_FIELD_DEF(uint32_t, bSize) +TILING_DATA_FIELD_DEF(uint32_t, n2Size) +TILING_DATA_FIELD_DEF(uint32_t, gSize) +TILING_DATA_FIELD_DEF(uint32_t, s1Size) +TILING_DATA_FIELD_DEF(uint32_t, s2Size) +TILING_DATA_FIELD_DEF(uint32_t, sparseCount) +TILING_DATA_FIELD_DEF(uint32_t, usedCoreNum) +TILING_DATA_FIELD_DEF(uint32_t, blockSize) +TILING_DATA_FIELD_DEF(uint32_t, maxBlockNumPerBatch) +TILING_DATA_FIELD_DEF(uint32_t, sparseMode) +TILING_DATA_FIELD_DEF(uint32_t, cmpRatio) +TILING_DATA_FIELD_DEF(uint32_t, returnValue) +TILING_DATA_FIELD_DEF(int32_t, maxSeqlenQ) +TILING_DATA_FIELD_DEF(uint32_t, keyStride0) +TILING_DATA_FIELD_DEF(uint32_t, keyDequantScaleStride0) +TILING_DATA_FIELD_DEF(uint32_t, quantMode) +END_TILING_DATA_DEF +REGISTER_TILING_DATA_CLASS(QuantLightningIndexerV2, QLIV2TilingData) + +// -----------算子CompileInfo定义------------------- +struct QLIV2CompileInfo {}; + +// -----------算子Tiling入参结构体定义--------------- +struct QLIV2ParaInfo { + TilingRequiredParaInfo query = {nullptr, nullptr}; + TilingRequiredParaInfo key = {nullptr, nullptr}; + TilingRequiredParaInfo weights = {nullptr, nullptr}; + TilingRequiredParaInfo query_dequant_scale = {nullptr, nullptr}; + TilingRequiredParaInfo key_dequant_scale = {nullptr, nullptr}; + TilingOptionalParaInfo cuSeqLensQ = {nullptr, nullptr}; + TilingOptionalParaInfo cuSeqLensK = {nullptr, nullptr}; + TilingOptionalParaInfo sequsedQ = {nullptr, nullptr}; + TilingOptionalParaInfo sequsedK = {nullptr, nullptr}; + TilingOptionalParaInfo cmpResidualK = {nullptr, nullptr}; + TilingOptionalParaInfo blockTable = {nullptr, nullptr}; + TilingOptionalParaInfo outputIdxOffset = {nullptr, nullptr}; + TilingOptionalParaInfo metadata = {nullptr, nullptr}; + TilingRequiredParaInfo attenOut = {nullptr, nullptr}; + TilingRequiredParaInfo sparseValues = {nullptr, nullptr}; + + const int32_t *quantMode = nullptr; + const int32_t *maxSeqlenQ = nullptr; + const char *layOutQuery = nullptr; + const char *layOutKey = nullptr; + const int32_t *blockSize = nullptr; + const int32_t *sparseMode = nullptr; + const int32_t *sparseCount = nullptr; + const int32_t *cmpRatio = nullptr; + const int32_t *returnValue = nullptr; +}; + +// -----------算子Tiling入参信息类--------------- +class QLIV2TilingInfo { +public: + const char *opName = nullptr; + fe::PlatFormInfos *platformInfo = nullptr; + QLIV2ParaInfo opParamInfo; + // Base Param + platform_ascendc::SocVersion socVersion = platform_ascendc::SocVersion::ASCEND910B; + NpuArch npuArch = NpuArch::DAV_2201; + uint32_t bSize = 0; + uint32_t n1Size = 0; + uint32_t n2Size = 0; + uint32_t s1Size = 0; + int64_t s2Size = 0; + uint32_t qkHeadDim = 0; + uint32_t gSize = 0; + // PageAttention + bool pageAttentionFlag = false; + int32_t blockSize = 0; + uint32_t maxBlockNumPerBatch = 0; + // Mask + int32_t sparseMode = 0; + // Others Flag + uint32_t sparseCount = 0; + uint32_t cmpRatio = 1; + bool returnValue = false; + uint32_t keyStride0 = 0; + uint32_t keyDequantScaleStride0 = 0; + std::vector keyStridesVec; + std::vector keyDequantScaleStridesVec; + int32_t maxSeqlenQ = -1; + // DType + ge::DataType inputQType = ge::DT_FLOAT16; + ge::DataType inputKType = ge::DT_FLOAT16; + ge::DataType outputType = ge::DT_INT32; + // Layout + DataLayout inputQLayout = DataLayout::BSND; + DataLayout inputKLayout = DataLayout::PA_BBND; +}; + +// -----------算子Tiling入参信息解析及Check类--------------- +class QLIV2InfoParser { +public: + explicit QLIV2InfoParser(gert::TilingContext *context) : context_(context) {} + ~QLIV2InfoParser() = default; + + ge::graphStatus CheckRequiredInOutExistence() const; + ge::graphStatus CheckRequiredAttrExistence() const; + ge::graphStatus CheckRequiredParaExistence() const; + ge::graphStatus GetActualSeqLenSize(uint32_t &size, const gert::Tensor *tensor, + const std::string &actualSeqLenName) const; + ge::graphStatus GetOpName(); + ge::graphStatus GetNpuInfo(); + void GetOptionalInputParaInfo(); + void GetInputParaInfo(); + void GetOutputParaInfo(); + ge::graphStatus GetAttrParaInfo(); + ge::graphStatus CheckAttrParaInfo(); + ge::graphStatus GetOpParaInfo(); + ge::graphStatus ValidateInputShapesMatch(); + ge::graphStatus CheckScaleShape(); + ge::graphStatus GetAndCheckInOutDataType(); + ge::graphStatus GetBatchSize(); + ge::graphStatus GetHeadDim(); + ge::graphStatus GetS1Size(); + ge::graphStatus GetAndCheckOptionalInput(); + ge::graphStatus CheckShapeDim(); + ge::graphStatus GetAndCheckBlockSize(); + ge::graphStatus GetS2SizeForPageAttention(); + ge::graphStatus GetS2SizeForBatchContinuous(); + ge::graphStatus GetS2Size(); + ge::graphStatus GetQueryKeyAndOutLayout(); + ge::graphStatus GetN1Size(); + ge::graphStatus GetAndCheckN2Size(); + ge::graphStatus GetGSize(); + ge::graphStatus GetAttenMaskInfo(); + ge::graphStatus GetActualSeqInfo(); + ge::graphStatus CheckKeyContiguous() const; + void GenerateInfo(QLIV2TilingInfo &QLIV2Info); + ge::graphStatus ParseAndCheck(QLIV2TilingInfo &QLIV2Info); + +public: + gert::TilingContext *context_ = nullptr; + const char *opName_; + fe::PlatFormInfos *platformInfo_; + QLIV2ParaInfo opParamInfo_; + + // BaseParams + uint32_t bSize_ = 0; + uint32_t n1Size_ = 0; + uint32_t n2Size_ = 0; + uint32_t gSize_ = 0; + uint32_t s1Size_ = 0; + int64_t s2Size_ = 0; + uint32_t headDim_ = 0; + // Layout + DataLayout qLayout_ = DataLayout::BSND; + DataLayout kLayout_ = DataLayout::PA_BBND; + // PageAttention + uint32_t maxBlockNumPerBatch_ = 0; + int32_t blockSize_ = 0; + platform_ascendc::SocVersion socVersion_ = platform_ascendc::SocVersion::ASCEND910B; + NpuArch npuArch_ = NpuArch::DAV_2201; + ge::DataType inputQType_ = ge::DT_FLOAT16; + ge::DataType inputKType_ = ge::DT_FLOAT16; + ge::DataType weightsType_ = ge::DT_FLOAT16; + ge::DataType inputQueryScaleType_ = ge::DT_FLOAT16; + ge::DataType inputKeyScaleType_ = ge::DT_FLOAT16; + ge::DataType blockTableType_ = ge::DT_FLOAT16; + ge::DataType inputKRopeType_ = ge::DT_FLOAT16; + ge::DataType outputType_ = ge::DT_FLOAT16; + ge::DataType valuesOutType_ = ge::DT_BF16; + std::vector keyStridesVec_; + std::vector keyDequantScaleStridesVec_; +}; + +// ---------------算子Tiling类--------------- +class QuantLightningIndexerV2Tiling { +public: + explicit QuantLightningIndexerV2Tiling(gert::TilingContext *context) : context_(context) {}; + ge::graphStatus DoTiling(QLIV2TilingInfo *tilingInfo); + +private: + gert::TilingContext *context_ = nullptr; + QLIV2TilingData tilingData_; +}; + +} // namespace optiling +#endif // QUANT_LIGHTNING_INDEXER_V2_TILING_H diff --git a/xllm_ops/attention/quant_lightning_indexer_v2/op_kernel/arch22/quant_lightning_indexer_v2_common_arch22.h b/xllm_ops/attention/quant_lightning_indexer_v2/op_kernel/arch22/quant_lightning_indexer_v2_common_arch22.h new file mode 100644 index 0000000..24211d0 --- /dev/null +++ b/xllm_ops/attention/quant_lightning_indexer_v2/op_kernel/arch22/quant_lightning_indexer_v2_common_arch22.h @@ -0,0 +1,163 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file quant_lightning_indexer_v2_common_arch22.h + * \brief + */ +#ifndef QUANT_LIGHTNING_INDEXER_V2_COMMON_H +#define QUANT_LIGHTNING_INDEXER_V2_COMMON_H + +namespace QLIV2Common { + +// 与tiling的layout保持一致 +enum class LI_LAYOUT : uint32_t { + BSND = 0, + TND = 1, + PA_BBND = 2 +}; + +template +struct QLIV2Type { + using queryType = Q_T; + using keyType = K_T; + using outputType = OUT_T; + static constexpr bool pageAttention = PAGE_ATTENTION; + static constexpr LI_LAYOUT layout = Q_LAYOUT_T; + static constexpr LI_LAYOUT keyLayout = K_LAYOUT_T; +}; + +struct RunInfo { + uint32_t loop; + uint32_t bN2Idx; + uint32_t bIdx; + uint32_t n2Idx = 0; + uint32_t gS1Idx; + uint32_t s2Idx; + + uint32_t actS1Size = 1; + uint32_t actS2Size = 1; + uint32_t actS2SizeOrig = 1; + uint32_t actMBaseSize; + uint32_t actualSingleProcessSInnerSize; + uint32_t actualSingleProcessSInnerSizeAlign; + + uint64_t tensorQueryOffset; + uint64_t tensorKeyOffset; + uint64_t tensorKeyScaleOffset; + uint64_t tensorWeightsOffset; + uint64_t indiceOutOffset; + + bool isFirstS2InnerLoop; + bool isLastS2InnerLoop; + bool isValid = false; + bool isNeedLD = false; + uint32_t saveWorkSpaceIdx = 0; +}; + +struct ConstInfo { + // CUBE与VEC核间同步的模式 + static constexpr uint32_t FIA_SYNC_MODE2 = 2; + // BUFFER的字节数 + static constexpr uint32_t BUFFER_SIZE_BYTE_32B = 32; + static constexpr uint32_t BUFFER_SIZE_BYTE_64B = 64; + static constexpr uint32_t BUFFER_SIZE_BYTE_256B = 256; + static constexpr uint32_t BUFFER_SIZE_BYTE_512B = 512; + static constexpr uint32_t BUFFER_SIZE_BYTE_1K = 1024; + static constexpr uint32_t BUFFER_SIZE_BYTE_2K = 2048; + static constexpr uint32_t BUFFER_SIZE_BYTE_4K = 4096; + static constexpr uint32_t BUFFER_SIZE_BYTE_8K = 8192; + static constexpr uint32_t BUFFER_SIZE_BYTE_16K = 16384; + static constexpr uint32_t BUFFER_SIZE_BYTE_32K = 32768; + // 无效索引 + static constexpr int INVALID_IDX = -1; + + // CUBE和VEC的核间同步EventID + uint32_t syncC1V1 = 0U; + uint32_t syncC1V0 = 2U; + uint32_t syncV1C1 = 0U; + uint32_t syncV0C1 = 1U; + + // 基本块大小 + uint32_t mBaseSize = 1ULL; + uint32_t s1BaseSize = 1ULL; + uint32_t s2BaseSize = 1ULL; + + uint64_t batchSize = 0ULL; + uint64_t gSize = 0ULL; + uint64_t qHeadNum = 0ULL; + uint64_t kHeadNum; + uint64_t headDim; + uint64_t sparseCount; // topK选取大小 + uint64_t kSeqSize = 0ULL; // kv最大S长度 + uint64_t qSeqSize = 1ULL; // q最大S长度 + uint32_t kCacheBlockSize = 0; // PA场景的block size + uint32_t maxBlockNumPerBatch = 0; // PA场景的最大单batch block number + LI_LAYOUT outputLayout; // 输出的格式 + bool attenMaskFlag = false; + uint32_t cmpRatio = 1; // 压缩率 + + uint32_t actualLenQDims = 0U; // query的actualSeqLength 的维度 + uint32_t actualLenDims = 0U; // KV 的actualSeqLength 的维度 + uint32_t cmpResiduaKLenDims = 0U; // cmpResidualK的维度 + bool isAccumSeqS1 = false; // 是否累加模式 + bool isAccumSeqS2 = false; // 是否累加模式 + + uint32_t s2Start = 0U; + uint32_t s2End = 0U; + uint32_t bN2Start = 0U; + uint32_t bN2End = 0U; + uint32_t gS1Start = 0U; + uint32_t gS1End = 0U; + uint32_t coreEnable = 0U; +}; + +struct LdSplitCoreInfo { + bool isLdCoreEnable = false; // 当前核是否参与规约任务 + uint32_t saveWorkSpaceIdx = 0U; // 存放LD参数的地址 + uint32_t bn2Idx = 0U; // 归约任务 + uint32_t bIdx = 0U; + uint32_t n2Idx = 0U; + uint32_t mIdx = 0U; + uint32_t workspaceIdx = 0U; // 当前AIV核上规约任务的索引 + uint32_t workspaceNum = 0U; // 当前AIV核上规约任务的S2切分数量 + uint32_t mStart = 0U; + uint32_t mNum = 0U; + uint64_t indiceOutCoreOffset = 0U; // 最终输出索引搬出Topk的初始偏移地址 + }; + +template +__aicore__ inline T1 Align(T1 num, T2 rnd) +{ + return (((rnd) == 0) ? 0 : (((num) + (rnd) - 1) / (rnd) * (rnd))); +} + +template +__aicore__ inline T1 Min(T1 a, T2 b) +{ + return (a > b) ? (b) : (a); +} + +template +__aicore__ inline T1 Max(T1 a, T2 b) +{ + return (a > b) ? (a) : (b); +} + +template +__aicore__ inline T CeilDiv(T num, T rnd) +{ + return (((rnd) == 0) ? 0 : (((num) + (rnd)-1) / (rnd))); +} +} // namespace QLIV2Common + +#endif // QUANT_LIGHTNING_INDEXER_V2_COMMON_H \ No newline at end of file diff --git a/xllm_ops/attention/quant_lightning_indexer_v2/op_kernel/arch22/quant_lightning_indexer_v2_kernel_arch22.h b/xllm_ops/attention/quant_lightning_indexer_v2/op_kernel/arch22/quant_lightning_indexer_v2_kernel_arch22.h new file mode 100644 index 0000000..121229d --- /dev/null +++ b/xllm_ops/attention/quant_lightning_indexer_v2/op_kernel/arch22/quant_lightning_indexer_v2_kernel_arch22.h @@ -0,0 +1,814 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file quant_lightning_indexer_v2_kernel_arch22.h + * \brief + */ + +#ifndef QUANT_LIGHTNING_INDEXER_V2_KERNEL_H +#define QUANT_LIGHTNING_INDEXER_V2_KERNEL_H + +#include "kernel_operator.h" +#include "kernel_operator_list_tensor_intf.h" +#include "kernel_tiling/kernel_tiling.h" +#include "lib/matmul_intf.h" +#include "lib/matrix/matmul/tiling.h" +#include "quant_lightning_indexer_v2_common_arch22.h" +#include "quant_lightning_indexer_v2_service_vector_arch22.h" +#include "quant_lightning_indexer_v2_service_cube_arch22.h" +#include "../quant_lightning_indexer_v2_metadata.h" + +namespace QLIV2Kernel { +using namespace QLIV2Common; +using namespace QLIV2ServiceVec; +using namespace matmul; +using namespace optiling::detail; +using namespace optiling; +using AscendC::CacheMode; +using AscendC::CrossCoreSetFlag; +using AscendC::CrossCoreWaitFlag; + +// 由于S2循环前，RunInfo还没有赋值，使用TempLoopInfo临时存放B、N、S1轴相关的信息 +// 同时减少重复计算 +struct TempLoopInfo { + uint32_t bN2Idx = 0; + uint32_t bIdx = 0U; + uint32_t n2Idx = 0U; + uint32_t gS1Idx = 0U; + uint32_t gS1LoopEnd = 0U; // gS1方向循环的结束Idx + uint32_t s2LoopEnd = 0U; // S2方向循环的结束Idx + uint32_t actS1Size = 1ULL; // 当前Batch循环处理的S1轴的实际大小 + uint32_t actS2Size = 0ULL; + uint32_t actS2SizeOrig = 0ULL; + bool curActSeqLenIsZero = false; + bool needDealActS1LessThanS1 = false; // S1的实际长度小于shape的S1长度时，是否需要清理输出 + bool isNeedLD = false; // 该基本块是否需要LD + uint32_t actMBaseSize = 0U; // m轴(gS1)方向实际大小 + uint32_t mBasicSizeTail = 0U; // gS1方向循环的尾基本块大小 + uint32_t s2BasicSizeTail = 0U; // S2方向循环的尾基本块大小 + uint32_t validS2Len = 0U; +}; + +template +class QLIV2Preload { +public: + __aicore__ inline QLIV2Preload() {}; + __aicore__ inline void Init(__gm__ uint8_t *query, __gm__ uint8_t *key, __gm__ uint8_t *weights, + __gm__ uint8_t *queryScale, __gm__ uint8_t *keyScale, + __gm__ uint8_t *cuSeqlensQ, __gm__ uint8_t *cuSeqlensK, + __gm__ uint8_t *sequsedQ, __gm__ uint8_t *sequsedK, + __gm__ uint8_t *cmpResidualK, __gm__ uint8_t *blockTable, + __gm__ uint8_t *outputIdxOffset, __gm__ uint8_t *metadata, + __gm__ uint8_t *sparseIndices, __gm__ uint8_t *sparseValues, + __gm__ uint8_t *workspace, const QLIV2TilingData *__restrict tiling, TPipe *tPipe); + __aicore__ inline void Process(); + + // =================================类型定义区================================= + using Q_T = typename QLIV2T::queryType; + using K_T = typename QLIV2T::keyType; + using OUT_T = typename QLIV2T::outputType; + static constexpr bool PAGE_ATTENTION = QLIV2T::pageAttention; + static constexpr LI_LAYOUT Q_LAYOUT_T = QLIV2T::layout; + static constexpr LI_LAYOUT K_LAYOUT_T = QLIV2T::keyLayout; + + using MM1_OUT_T = float; + + QLIV2Matmul matmulService; + QLIV2Vector vectorService; + + // =================================常量区================================= + static constexpr uint32_t SYNC_C1_V1_FLAG = 4; + static constexpr uint32_t SYNC_V1_C1_FLAG = 5; + + static constexpr uint32_t M_BASE_SIZE = 256; + static constexpr uint32_t S2_BASE_SIZE = 2048; + static constexpr uint32_t HEAD_DIM = 128; + static constexpr uint32_t K_HEAD_NUM = 1; + static constexpr uint32_t GM_ALIGN_BYTES = 512; + static constexpr uint32_t LI_QUANT_PRELOAD_TASK_CACHE_SIZE = 2; + + // for workspace double + static constexpr uint32_t WS_DOUBLE = 2; + static constexpr uint32_t ELE_NUM_PER_BLOCK = 16; + + static constexpr int64_t LD_PREFETCH_LEN = 2; + +protected: + TPipe *pipe = nullptr; + + // offset + uint64_t queryCoreOffset = 0ULL; + uint64_t keyCoreOffset = 0ULL; + uint64_t keyScaleCoreOffset = 0ULL; + uint64_t weightsCoreOffset = 0ULL; + uint64_t indiceOutCoreOffset = 0ULL; + uint32_t coreZeroEnable = 1U; + + // ================================Global Buffer区================================= + GlobalTensor queryGm; + GlobalTensor keyGm; + GlobalTensor weightsGm; + GlobalTensor metadataGm; + GlobalTensor indiceOutGm; + GlobalTensor blockTableGm; + + GlobalTensor actualSeqLengthsGmQ; + GlobalTensor actualSeqLengthsGm; + GlobalTensor cmpResidualKGm; + + // ================================类成员变量==================================== + // aic、aiv核信息 + uint32_t tmpBlockIdx = 0U; + uint32_t aiCoreIdx = 0U; + + QLIV2Common::ConstInfo constInfo{}; + QLIV2Common::LdSplitCoreInfo ldInfo{}; + TempLoopInfo tempLoopInfo{}; + + // ================================Init functions================================== + __aicore__ inline void InitTilingData(const QLIV2TilingData *__restrict tilingData); + __aicore__ inline void InitBuffers(); + __aicore__ inline void InitActualSeqLen(__gm__ uint8_t *cuSeqlensQ, __gm__ uint8_t *cuSeqlensK, + __gm__ uint8_t *sequsedQ, __gm__ uint8_t *sequsedK, + __gm__ uint8_t *cmpResidualK); + // ================================Split Core================================ + __aicore__ inline void SplitCore(); + __aicore__ inline uint32_t GetS2BaseBlockNumOnMask(uint32_t s1gIdx, uint32_t actS1Size, uint32_t actS2SizeOrig, + uint32_t &validS2Len); + __aicore__ inline uint32_t GetTotalBaseBlockNum(); + // ================================Process functions================================ + __aicore__ inline void ProcessMain(); + __aicore__ inline void ProcessBaseBlock(uint32_t loop, uint64_t s2LoopIdx, + QLIV2Common::RunInfo + runInfo[LI_QUANT_PRELOAD_TASK_CACHE_SIZE]); + __aicore__ inline void ProcessInvalid(); + __aicore__ inline void ProcessDecode(); + // ================================Params Calc===================================== + __aicore__ inline void CalcGS1LoopParams(uint32_t bN2Idx); + __aicore__ inline void GetBN2Idx(uint32_t bN2Idx); + __aicore__ inline uint32_t GetActualSeqLen(uint32_t bIdx, uint32_t actualLenDims, bool isAccumSeq, + GlobalTensor &actualSeqLengthsGm, uint32_t defaultSeqLen); + __aicore__ inline uint32_t GetActualSeqLenKey(uint32_t bIdx, uint32_t actualLenDims, + uint32_t cmpResiduaKLenDims, bool isAccumSeq, + GlobalTensor &actualSeqLengthsGm, + GlobalTensor &cmpResidualKGm, + uint32_t defaultSeqLen, uint32_t cmpRatio); + __aicore__ inline void GetS1S2ActualSeqLen(uint32_t bIdx, uint32_t &actS1Size, uint32_t &actS2Size, + uint32_t &actS2SizeOrig); + __aicore__ inline void CalcS2LoopParams(uint32_t bN2LoopIdx, uint32_t gS1LoopIdx); + __aicore__ inline void CalcRunInfo(uint32_t loop, uint32_t s2LoopIdx, QLIV2Common::RunInfo &runInfo); + __aicore__ inline void DealActSeqLenIsZero(uint32_t bIdx, uint32_t n2Idx, uint32_t s1Start); +}; + +template +__aicore__ inline void QLIV2Preload::InitTilingData(const QLIV2TilingData *__restrict tilingData) +{ + constInfo.batchSize = tilingData->bSize; + constInfo.qHeadNum = constInfo.gSize = tilingData->gSize; + constInfo.kSeqSize = tilingData->s2Size; + constInfo.qSeqSize = tilingData->s1Size; + constInfo.attenMaskFlag = (tilingData->sparseMode == 3); + constInfo.kCacheBlockSize = tilingData->blockSize; + constInfo.maxBlockNumPerBatch = tilingData->maxBlockNumPerBatch; + constInfo.sparseCount = tilingData->sparseCount; + constInfo.cmpRatio = tilingData->cmpRatio; + constInfo.outputLayout = Q_LAYOUT_T; // 输出和输入形状一致 + if constexpr (Q_LAYOUT_T == LI_LAYOUT::TND) { + constInfo.isAccumSeqS1 = true; + } + if constexpr (K_LAYOUT_T == LI_LAYOUT::TND) { + constInfo.isAccumSeqS2 = true; + } + + constInfo.kHeadNum = K_HEAD_NUM; + constInfo.headDim = HEAD_DIM; + + constInfo.mBaseSize = M_BASE_SIZE; + constInfo.s2BaseSize = S2_BASE_SIZE; + constInfo.s1BaseSize = (constInfo.mBaseSize + constInfo.gSize - 1) / constInfo.gSize; +} + +template +__aicore__ inline void QLIV2Preload::InitBuffers() +{ + if ASCEND_IS_AIV { + vectorService.InitBuffers(pipe); + } else { + matmulService.InitBuffers(pipe); + } +} + +template +__aicore__ inline void QLIV2Preload::InitActualSeqLen(__gm__ uint8_t *cuSeqlensQ, __gm__ uint8_t *cuSeqlensK, + __gm__ uint8_t *sequsedQ, __gm__ uint8_t *sequsedK, + __gm__ uint8_t *cmpResidualK) +{ + // Q side: cu_seqlens (TND) or seqused (non-TND) + if constexpr (Q_LAYOUT_T == LI_LAYOUT::TND) { + if (cuSeqlensQ != nullptr) { + constInfo.actualLenQDims = constInfo.batchSize + 1; + actualSeqLengthsGmQ.SetGlobalBuffer((__gm__ uint32_t *)cuSeqlensQ, constInfo.actualLenQDims); + } else { + constInfo.actualLenQDims = 0; + } + } else { + if (sequsedQ != nullptr) { + constInfo.actualLenQDims = constInfo.batchSize; + actualSeqLengthsGmQ.SetGlobalBuffer((__gm__ uint32_t *)sequsedQ, constInfo.actualLenQDims); + } else { + constInfo.actualLenQDims = 0; + } + } + // K side: cu_seqlens (TND) or seqused (non-TND) + if constexpr (K_LAYOUT_T == LI_LAYOUT::TND) { + if (cuSeqlensK != nullptr) { + constInfo.actualLenDims = constInfo.batchSize + 1; + actualSeqLengthsGm.SetGlobalBuffer((__gm__ uint32_t *)cuSeqlensK, constInfo.actualLenDims); + } else { + constInfo.actualLenDims = 0; + } + } else { + if (sequsedK != nullptr) { + constInfo.actualLenDims = constInfo.batchSize; + actualSeqLengthsGm.SetGlobalBuffer((__gm__ uint32_t *)sequsedK, constInfo.actualLenDims); + } else { + constInfo.actualLenDims = 0; + } + } + // cmpResidualK获取 + if (cmpResidualK != nullptr) { + constInfo.cmpResiduaKLenDims = constInfo.batchSize; + cmpResidualKGm.SetGlobalBuffer((__gm__ uint32_t *)cmpResidualK, constInfo.batchSize); + } else { + constInfo.cmpResiduaKLenDims = 0; + } +} + +template +__aicore__ inline uint32_t QLIV2Preload::GetActualSeqLen(uint32_t bIdx, uint32_t actualLenDims, bool isAccumSeq, + GlobalTensor &actualSeqLengthsGm, + uint32_t defaultSeqLen) +{ + if (actualLenDims == 0) { + return defaultSeqLen; + } else if (isAccumSeq) { + // TND with cu_seqlens: length[i] = cu_seqlens[i+1] - cu_seqlens[i] + return actualSeqLengthsGm.GetValue(bIdx + 1) - actualSeqLengthsGm.GetValue(bIdx); + } else { + // non-TND with seqused: length[i] = seqused[i] + return actualSeqLengthsGm.GetValue(bIdx); + } +} + +template +__aicore__ inline uint32_t QLIV2Preload::GetActualSeqLenKey(uint32_t bIdx, uint32_t actualLenDims, + uint32_t cmpResiduaKLenDims, bool isAccumSeq, + GlobalTensor &actualSeqLengthsGm, + GlobalTensor &cmpResidualKGm, + uint32_t defaultSeqLen, uint32_t cmpRatio) +{ + uint32_t cmpResidualK; // 当前bidx对应的cmpResidualK + if (cmpResiduaKLenDims == 0) { + cmpResidualK = 0; + } else { + cmpResidualK = cmpResidualKGm.GetValue(bIdx); + } + if (actualLenDims == 0) { + return defaultSeqLen * cmpRatio + cmpResidualK; + } else if (isAccumSeq) { + // TND with cu_seqlens: length[i] = cu_seqlens[i+1] - cu_seqlens[i] + return (actualSeqLengthsGm.GetValue(bIdx + 1) - actualSeqLengthsGm.GetValue(bIdx)) * cmpRatio + cmpResidualK; + } else { + return (actualSeqLengthsGm.GetValue(bIdx)) * cmpRatio + cmpResidualK; + } +} + +template +__aicore__ inline void QLIV2Preload::GetS1S2ActualSeqLen(uint32_t bIdx, uint32_t &actS1Size, + uint32_t &actS2Size, uint32_t &actS2SizeOrig) +{ + actS1Size = GetActualSeqLen(bIdx, constInfo.actualLenQDims, constInfo.isAccumSeqS1, actualSeqLengthsGmQ, + constInfo.qSeqSize); + actS2SizeOrig = + GetActualSeqLenKey(bIdx, constInfo.actualLenDims, constInfo.cmpResiduaKLenDims, constInfo.isAccumSeqS2, + actualSeqLengthsGm, cmpResidualKGm, constInfo.kSeqSize, constInfo.cmpRatio); // 压缩前的actS2Size + actS2Size = actS2SizeOrig / constInfo.cmpRatio; // 真实使用的压缩后S2长度 +} + +template +__aicore__ inline uint32_t QLIV2Preload::GetS2BaseBlockNumOnMask(uint32_t s1gIdx, uint32_t actS1Size, + uint32_t actS2SizeOrig, uint32_t &validS2Len) +{ + if (actS2SizeOrig / constInfo.cmpRatio == 0) { + validS2Len = 0; + return 0; + } + uint32_t s1Offset = constInfo.s1BaseSize * s1gIdx; + int32_t validS2LenBase = static_cast(actS2SizeOrig) + - static_cast(actS1Size); // 压缩前的validS2LenBase + validS2Len = + (static_cast(s1Offset) + validS2LenBase + + static_cast(constInfo.s1BaseSize)) / + static_cast(constInfo.cmpRatio); + validS2Len = Min(validS2Len, static_cast(actS2SizeOrig) / constInfo.cmpRatio); + validS2Len = Max(validS2Len, 1); + return (validS2Len + constInfo.s2BaseSize - 1) / constInfo.s2BaseSize; +} + +template +__aicore__ inline uint32_t QLIV2Preload::GetTotalBaseBlockNum() +{ + uint32_t totalBlockNum = 0; + uint32_t actS1Size, actS2Size, actS2SizeOrig; + uint32_t s1GBaseNum, s2BaseNum; + uint32_t validS2Len = 0; + for (uint32_t bIdx = 0; bIdx < constInfo.batchSize; bIdx++) { + GetS1S2ActualSeqLen(bIdx, actS1Size, actS2Size, actS2SizeOrig); + s1GBaseNum = CeilDiv(actS1Size, constInfo.s1BaseSize); + if (!constInfo.attenMaskFlag) { + s2BaseNum = CeilDiv(actS2Size, constInfo.s2BaseSize); + totalBlockNum += s1GBaseNum * s2BaseNum * constInfo.kHeadNum; + continue; + } + for (uint32_t s1gIdx = 0; s1gIdx < s1GBaseNum; s1gIdx++) { + s2BaseNum = GetS2BaseBlockNumOnMask(s1gIdx, actS1Size, actS2SizeOrig, validS2Len); + totalBlockNum += s2BaseNum * constInfo.kHeadNum; + } + } + return totalBlockNum; +} + +// 多核版本，双闭区间。基本原则：计算每个核最少处理的块数, +// 剩余的部分前面的核每个核多处理一块 +template +__aicore__ void inline QLIV2Preload::SplitCore() +{ + constInfo.coreEnable = metadataGm.GetValue(GetAttrAbsIndex(aiCoreIdx, QLI_V2_CORE_ENABLE_INDEX, false)); + if (aiCoreIdx != 0) { + constInfo.bN2Start = metadataGm.GetValue(GetAttrAbsIndex(aiCoreIdx, QLI_V2_BN2_START_INDEX, false)); + constInfo.gS1Start = metadataGm.GetValue(GetAttrAbsIndex(aiCoreIdx, QLI_V2_M_START_INDEX, false)); + constInfo.s2Start = metadataGm.GetValue(GetAttrAbsIndex(aiCoreIdx, QLI_V2_S2_START_INDEX, false)); + } + constInfo.bN2End = metadataGm.GetValue(GetAttrAbsIndex(aiCoreIdx, QLI_V2_BN2_END_INDEX, false)); + constInfo.gS1End = metadataGm.GetValue(GetAttrAbsIndex(aiCoreIdx, QLI_V2_M_END_INDEX, false)); + constInfo.s2End = metadataGm.GetValue(GetAttrAbsIndex(aiCoreIdx, QLI_V2_S2_END_INDEX, false)); + + // 如果0核都没有启动，说明所有核都没启动 + coreZeroEnable = metadataGm.GetValue(GetAttrAbsIndex(0, QLI_V2_CORE_ENABLE_INDEX, false)); + + // LD 第一个workspace的索引 + uint32_t ldFirstWorkSpaceIndex = GetAttrAbsIndex(aiCoreIdx, QLI_V2_FIRST_QLD_V2_DATA_WORKSPACE_IDX_INDEX, false); + ldInfo.saveWorkSpaceIdx = metadataGm.GetValue(ldFirstWorkSpaceIndex); + if ASCEND_IS_AIV { + uint32_t vecCoreIdx = tmpBlockIdx; // 此时代表vectoe核数 + uint32_t ldCoreEnableIndex = GetAttrAbsIndex(vecCoreIdx, QLD_V2_CORE_ENABLE_INDEX, true); + ldInfo.isLdCoreEnable = metadataGm.GetValue(ldCoreEnableIndex); + + if (!ldInfo.isLdCoreEnable) { + return; + } + + // LD 参数信息 + uint32_t ldBn2IdxIndex = GetAttrAbsIndex(vecCoreIdx, QLD_V2_BN2_IDX_INDEX, true); + uint32_t ldMIdxIndex = GetAttrAbsIndex(vecCoreIdx, QLD_V2_M_IDX_INDEX, true); + uint32_t ldWorkspaceIdxIndex = GetAttrAbsIndex(vecCoreIdx, QLD_V2_WORKSPACE_IDX_INDEX, true); + uint32_t ldWorkspaceNumINDEX = GetAttrAbsIndex(vecCoreIdx, QLD_V2_WORKSPACE_NUM_INDEX, true); + uint32_t ldMstartIndex = GetAttrAbsIndex(vecCoreIdx, QLD_V2_M_START_INDEX, true); + uint32_t ldMNumIndex = GetAttrAbsIndex(vecCoreIdx, QLD_V2_M_NUM_INDEX, true); + + ldInfo.bn2Idx = metadataGm.GetValue(ldBn2IdxIndex); + ldInfo.bIdx = ldInfo.bn2Idx / constInfo.kHeadNum; + ldInfo.n2Idx = ldInfo.bn2Idx % constInfo.kHeadNum; + ldInfo.mIdx = metadataGm.GetValue(ldMIdxIndex); + ldInfo.workspaceIdx = metadataGm.GetValue(ldWorkspaceIdxIndex); + ldInfo.workspaceNum = metadataGm.GetValue(ldWorkspaceNumINDEX); + ldInfo.mStart = metadataGm.GetValue(ldMstartIndex); + ldInfo.mNum = metadataGm.GetValue(ldMNumIndex); + uint64_t actualSeqQPrefixSum = 0; + if constexpr (Q_LAYOUT_T == LI_LAYOUT::TND) { + actualSeqQPrefixSum = (ldInfo.bIdx <= 0) ? 0 : actualSeqLengthsGmQ.GetValue(ldInfo.bIdx); + } else { // BSND + actualSeqQPrefixSum = (ldInfo.bIdx <= 0) ? 0 : ldInfo.bIdx * constInfo.qSeqSize; + } + // 搬出Topk的初始偏移地址 + ldInfo.indiceOutCoreOffset = actualSeqQPrefixSum * constInfo.kHeadNum * constInfo.sparseCount + + static_cast(ldInfo.n2Idx) * constInfo.sparseCount + + static_cast(ldInfo.mIdx) * constInfo.s1BaseSize * + constInfo.kHeadNum * constInfo.sparseCount; + } +} + +template +__aicore__ inline void QLIV2Preload::DealActSeqLenIsZero(uint32_t bIdx, uint32_t n2Idx, uint32_t s1Start) +{ + if ASCEND_IS_AIV { + if (constInfo.outputLayout == LI_LAYOUT::TND) { + uint32_t tSize = actualSeqLengthsGmQ.GetValue(constInfo.batchSize); + uint32_t tBase = bIdx == 0 ? 0 : actualSeqLengthsGmQ.GetValue(bIdx); + uint32_t s1Count = tempLoopInfo.actS1Size; + + for (uint32_t s1Idx = s1Start; s1Idx < s1Count; s1Idx++) { + uint64_t indiceOutOffset = + (tBase + s1Idx) * constInfo.kHeadNum * constInfo.sparseCount + // T轴、s1轴偏移 + n2Idx * constInfo.sparseCount; // N2轴偏移 + vectorService.CleanInvalidOutput(indiceOutOffset); + } + } else if (constInfo.outputLayout == LI_LAYOUT::BSND) { + for (uint32_t s1Idx = s1Start; s1Idx < constInfo.qSeqSize; s1Idx++) { + // B,S1,N2,K + uint64_t indiceOutOffset = static_cast(bIdx) * constInfo.qSeqSize * + constInfo.kHeadNum * constInfo.sparseCount + + static_cast(s1Idx) * constInfo.kHeadNum * + constInfo.sparseCount + // B轴、S1轴偏移 + static_cast(n2Idx) * constInfo.sparseCount; // N2轴偏移 + vectorService.CleanInvalidOutput(indiceOutOffset); + } + } + } +} + +template +__aicore__ inline void QLIV2Preload::Init(__gm__ uint8_t *query, __gm__ uint8_t *key, __gm__ uint8_t *weights, + __gm__ uint8_t *queryScale, __gm__ uint8_t *keyScale, + __gm__ uint8_t *cuSeqlensQ, __gm__ uint8_t *cuSeqlensK, + __gm__ uint8_t *sequsedQ, __gm__ uint8_t *sequsedK, + __gm__ uint8_t *cmpResidualK, __gm__ uint8_t *blockTable, + __gm__ uint8_t *outputIdxOffset, __gm__ uint8_t *metadata, + __gm__ uint8_t *sparseIndices, __gm__ uint8_t *sparseValues, + __gm__ uint8_t *workspace, const QLIV2TilingData *__restrict tiling, + TPipe *tPipe) +{ + if ASCEND_IS_AIV { + tmpBlockIdx = GetBlockIdx(); // vec:0-47 + aiCoreIdx = tmpBlockIdx / 2; + } else { + tmpBlockIdx = GetBlockIdx(); // cube:0-23 + aiCoreIdx = tmpBlockIdx; + } + + InitTilingData(tiling); + InitActualSeqLen(cuSeqlensQ, cuSeqlensK, sequsedQ, sequsedK, cmpResidualK); + + if (metadata != nullptr) { + metadataGm.SetGlobalBuffer((__gm__ uint32_t *)metadata); + // 计算分核 + SplitCore(); + } + + pipe = tPipe; + // workspace 内存排布 + // |mm1ResGm(存S) + uint64_t offset = 0; + + // mm1开DoubleBuffer + GlobalTensor mm1ResGm; // 存放S + uint64_t singleCoreMm1ResSize = WS_DOUBLE * constInfo.s1BaseSize * constInfo.s2BaseSize * sizeof(MM1_OUT_T); + mm1ResGm.SetGlobalBuffer((__gm__ MM1_OUT_T *)(workspace + aiCoreIdx * singleCoreMm1ResSize)); + offset += GetBlockNum() * singleCoreMm1ResSize; + + GlobalTensor weightWorkspaceGm; // v1阶段处理w*scale后的结果 + uint64_t weightMemSize = BLOCK_CUBE * constInfo.mBaseSize * WS_DOUBLE * sizeof(half); + weightWorkspaceGm.SetGlobalBuffer((__gm__ half *)(workspace + offset + aiCoreIdx * weightMemSize)); + offset += GetBlockNum() * weightMemSize; + + // ld流程需要ws大小: [aicnum, 2, s1BaseSize, topkOut_*2] + // (aic, 8, 2, 2, 2048) + // (aic, s1_cube, 头尾, idx/value, K) + GlobalTensor vec1ResGm; // 存放TopK计算中间结果 + vec1ResGm.SetGlobalBuffer((__gm__ float *)(workspace + offset)); + offset += GetBlockNum() * constInfo.s1BaseSize * WS_DOUBLE * WS_DOUBLE * BASE_TOPK * sizeof(float); + + GlobalTensor qScaleGm; + GlobalTensor kScaleGm; + if ASCEND_IS_AIV { + vectorService.InitParams(constInfo, ldInfo, tiling); + indiceOutGm.SetGlobalBuffer((__gm__ int32_t *)sparseIndices); + weightsGm.SetGlobalBuffer((__gm__ half *)weights); + qScaleGm.SetGlobalBuffer((__gm__ half *)queryScale); + kScaleGm.SetGlobalBuffer((__gm__ half *)keyScale); + blockTableGm.SetGlobalBuffer((__gm__ int32_t *)blockTable); + vectorService.InitVecInputTensor(weightsGm, qScaleGm, kScaleGm, indiceOutGm, blockTableGm); + vectorService.InitVecWorkspaceTensor(weightWorkspaceGm, mm1ResGm, vec1ResGm); + } else { + matmulService.InitParams(constInfo); + queryGm.SetGlobalBuffer((__gm__ Q_T *)query); + if constexpr (PAGE_ATTENTION) { + blockTableGm.SetGlobalBuffer((__gm__ int32_t *)blockTable); + } + keyGm.SetGlobalBuffer((__gm__ K_T *)key); + matmulService.InitMm1GlobalTensor(blockTableGm, keyGm, queryGm, mm1ResGm, weightWorkspaceGm); + } + InitBuffers(); +} + +template +__aicore__ inline void QLIV2Preload::GetBN2Idx(uint32_t bN2Idx) +{ + tempLoopInfo.bN2Idx = bN2Idx; + tempLoopInfo.bIdx = bN2Idx / constInfo.kHeadNum; + tempLoopInfo.n2Idx = bN2Idx % constInfo.kHeadNum; +} + +template +__aicore__ inline void QLIV2Preload::CalcS2LoopParams(uint32_t bN2LoopIdx, uint32_t gS1LoopIdx) +{ + tempLoopInfo.gS1Idx = gS1LoopIdx; + tempLoopInfo.actMBaseSize = constInfo.mBaseSize; + uint32_t remainedGS1Size = tempLoopInfo.actS1Size * constInfo.gSize - tempLoopInfo.gS1Idx * constInfo.mBaseSize; + if (remainedGS1Size <= constInfo.mBaseSize && remainedGS1Size > 0) { + tempLoopInfo.actMBaseSize = tempLoopInfo.mBasicSizeTail; + } + + bool isEnd = (bN2LoopIdx + 1 == constInfo.bN2End) && (gS1LoopIdx + 1 == tempLoopInfo.gS1LoopEnd); + uint32_t s2BlockNum; + uint32_t validS2Len = 0; + if (constInfo.attenMaskFlag) { + s2BlockNum = GetS2BaseBlockNumOnMask(gS1LoopIdx, tempLoopInfo.actS1Size, tempLoopInfo.actS2SizeOrig, + tempLoopInfo.validS2Len); + } else { + s2BlockNum = (tempLoopInfo.actS2Size + constInfo.s2BaseSize - 1) / constInfo.s2BaseSize; + tempLoopInfo.validS2Len = tempLoopInfo.actS2Size; + } + tempLoopInfo.s2LoopEnd = (isEnd && constInfo.s2End != 0) ? constInfo.s2End : s2BlockNum; + if (constInfo.s2Start > 0 || tempLoopInfo.s2LoopEnd < s2BlockNum) { + tempLoopInfo.isNeedLD = true; + } else { + tempLoopInfo.isNeedLD = false; + } + tempLoopInfo.s2BasicSizeTail = tempLoopInfo.validS2Len % constInfo.s2BaseSize; + tempLoopInfo.s2BasicSizeTail = (tempLoopInfo.s2BasicSizeTail == 0) ? + constInfo.s2BaseSize : tempLoopInfo.s2BasicSizeTail; +} + +template +__aicore__ inline void QLIV2Preload::CalcGS1LoopParams(uint32_t bN2LoopIdx) +{ + GetBN2Idx(bN2LoopIdx); + GetS1S2ActualSeqLen(tempLoopInfo.bIdx, tempLoopInfo.actS1Size, tempLoopInfo.actS2Size, tempLoopInfo.actS2SizeOrig); + if ((tempLoopInfo.actS2Size == 0) || (tempLoopInfo.actS1Size == 0)) { + tempLoopInfo.curActSeqLenIsZero = true; + return; + } + tempLoopInfo.curActSeqLenIsZero = false; + tempLoopInfo.mBasicSizeTail = (tempLoopInfo.actS1Size * constInfo.gSize) % constInfo.mBaseSize; + tempLoopInfo.mBasicSizeTail = + (tempLoopInfo.mBasicSizeTail == 0) ? constInfo.mBaseSize : tempLoopInfo.mBasicSizeTail; + + uint32_t gS1SplitNum = (tempLoopInfo.actS1Size * constInfo.gSize + constInfo.mBaseSize - 1) / constInfo.mBaseSize; + tempLoopInfo.gS1LoopEnd = (bN2LoopIdx + 1 == constInfo.bN2End && constInfo.gS1End != 0) + ? constInfo.gS1End : gS1SplitNum; + if constexpr (Q_LAYOUT_T == LI_LAYOUT::BSND) { + if (tempLoopInfo.gS1LoopEnd == gS1SplitNum && constInfo.qSeqSize > tempLoopInfo.actS1Size) { + tempLoopInfo.needDealActS1LessThanS1 = true; + } + } +} + +template +__aicore__ inline void QLIV2Preload::CalcRunInfo(uint32_t loop, uint32_t s2LoopIdx, + QLIV2Common::RunInfo &runInfo) +{ + runInfo.loop = loop; + runInfo.bIdx = tempLoopInfo.bIdx; + runInfo.gS1Idx = tempLoopInfo.gS1Idx; + runInfo.s2Idx = s2LoopIdx; + runInfo.bN2Idx = tempLoopInfo.bN2Idx; + runInfo.isValid = s2LoopIdx < tempLoopInfo.s2LoopEnd; + runInfo.isNeedLD = tempLoopInfo.isNeedLD; + if (runInfo.isNeedLD && s2LoopIdx + 1 == tempLoopInfo.s2LoopEnd) { + runInfo.saveWorkSpaceIdx = ldInfo.saveWorkSpaceIdx; + ldInfo.saveWorkSpaceIdx++; + } + + if (!runInfo.isValid) { + return; + } + + runInfo.actS1Size = tempLoopInfo.actS1Size; + runInfo.actS2Size = tempLoopInfo.actS2Size; + runInfo.actS2SizeOrig = tempLoopInfo.actS2SizeOrig; + // 计算实际基本块size + runInfo.actMBaseSize = tempLoopInfo.actMBaseSize; + runInfo.actualSingleProcessSInnerSize = constInfo.s2BaseSize; + uint32_t s2SplitNum = (tempLoopInfo.validS2Len + constInfo.s2BaseSize - 1) / constInfo.s2BaseSize; + if (runInfo.s2Idx == s2SplitNum - 1) { + runInfo.actualSingleProcessSInnerSize = tempLoopInfo.s2BasicSizeTail; + } + runInfo.actualSingleProcessSInnerSizeAlign = + QLIV2Common::Align((uint32_t)runInfo.actualSingleProcessSInnerSize, + QLIV2Common::ConstInfo::BUFFER_SIZE_BYTE_32B); + + runInfo.isFirstS2InnerLoop = s2LoopIdx == constInfo.s2Start; + runInfo.isLastS2InnerLoop = (s2LoopIdx + 1 == tempLoopInfo.s2LoopEnd); + + if (runInfo.isFirstS2InnerLoop) { + uint64_t actualSeqQPrefixSum; + if constexpr (Q_LAYOUT_T == LI_LAYOUT::TND) { + actualSeqQPrefixSum = (runInfo.bIdx <= 0) ? 0 : actualSeqLengthsGmQ.GetValue(runInfo.bIdx); + } else { // BSND + actualSeqQPrefixSum = (runInfo.bIdx <= 0) ? 0 : runInfo.bIdx * constInfo.qSeqSize; + } + uint64_t tndBIdxOffset = actualSeqQPrefixSum * constInfo.qHeadNum * constInfo.headDim; + // B,S1,N1(N2,G),D + queryCoreOffset = tndBIdxOffset + runInfo.gS1Idx * constInfo.mBaseSize * constInfo.headDim; + // B,S1,N1(N2,G)/T,N1(N2,G) + weightsCoreOffset = actualSeqQPrefixSum * constInfo.qHeadNum + runInfo.n2Idx * constInfo.gSize; + // B,S1,N2,k/T,N2,k + indiceOutCoreOffset = + actualSeqQPrefixSum * constInfo.kHeadNum * constInfo.sparseCount + runInfo.n2Idx * constInfo.sparseCount; + } + uint64_t actualSeqKPrefixSum; + if constexpr (K_LAYOUT_T == LI_LAYOUT::TND) { // T N2 D, cu_seqlens_k + actualSeqKPrefixSum = (runInfo.bIdx <= 0) ? 0 : actualSeqLengthsGm.GetValue(runInfo.bIdx); + } else { + actualSeqKPrefixSum = (runInfo.bIdx <= 0) ? 0 : runInfo.bIdx * constInfo.kSeqSize; + } + uint64_t tndBIdxOffsetForK = actualSeqKPrefixSum * constInfo.kHeadNum * constInfo.headDim; + keyCoreOffset = tndBIdxOffsetForK + runInfo.s2Idx * constInfo.s2BaseSize * constInfo.kHeadNum * constInfo.headDim; + keyScaleCoreOffset = (actualSeqKPrefixSum + runInfo.s2Idx * constInfo.s2BaseSize) * constInfo.kHeadNum; + runInfo.tensorQueryOffset = queryCoreOffset; + runInfo.tensorKeyOffset = keyCoreOffset; + runInfo.tensorKeyScaleOffset = keyScaleCoreOffset; + runInfo.tensorWeightsOffset = weightsCoreOffset; + runInfo.indiceOutOffset = indiceOutCoreOffset; +} + +template +__aicore__ inline void QLIV2Preload::Process() +{ + // 没有计算任务，直接清理输出 + if (coreZeroEnable == 0) { + ProcessInvalid(); + return; + } + ProcessMain(); + + ProcessDecode(); +} + +template +__aicore__ inline void QLIV2Preload::ProcessInvalid() +{ + if ASCEND_IS_AIV { + uint32_t aivCoreNum = GetBlockNum() * 2; // 2 means c:v = 1:2 + uint64_t totalOutputSize = + constInfo.batchSize * constInfo.qSeqSize * constInfo.kHeadNum * constInfo.sparseCount; + uint64_t singleCoreSize = + QLIV2Common::Align((totalOutputSize + aivCoreNum - 1) / aivCoreNum, GM_ALIGN_BYTES / sizeof(OUT_T)); + uint64_t baseSize = tmpBlockIdx * singleCoreSize; + if (baseSize < totalOutputSize) { + uint64_t dealSize = + (baseSize + singleCoreSize <= totalOutputSize) ? singleCoreSize : totalOutputSize - baseSize; + GlobalTensor output = indiceOutGm[baseSize]; + AscendC::InitGlobalMemory(output, dealSize, constInfo.INVALID_IDX); + } + } +} + +template +__aicore__ inline void QLIV2Preload::ProcessMain() +{ + // 无任务核直接返回 + if (constInfo.coreEnable == 0) { + return; + } + + if ASCEND_IS_AIV { + vectorService.AllocEventID(); + CrossCoreSetFlag(constInfo.syncV1C1); + CrossCoreSetFlag(constInfo.syncV1C1); + } else { + matmulService.AllocEventID(); + CrossCoreSetFlag(constInfo.syncC1V0); + CrossCoreSetFlag(constInfo.syncC1V0); + } + + QLIV2Common::RunInfo runInfo[LI_QUANT_PRELOAD_TASK_CACHE_SIZE]; + + // 适配左闭右开 + if (constInfo.bN2Start == constInfo.bN2End) { + if (constInfo.gS1Start != constInfo.gS1End || constInfo.s2Start != constInfo.s2End) { + constInfo.bN2End += 1; + if (constInfo.s2End != 0) { + constInfo.gS1End += 1; + } + } + } else if ((constInfo.gS1End != 0) || (constInfo.s2End != 0)) { + constInfo.bN2End += 1; + if (constInfo.s2End != 0) { + constInfo.gS1End += 1; + } + } + + uint32_t gloop = 0; + for (uint32_t bN2LoopIdx = constInfo.bN2Start; bN2LoopIdx < constInfo.bN2End; bN2LoopIdx++) { + CalcGS1LoopParams(bN2LoopIdx); + if (tempLoopInfo.curActSeqLenIsZero) { + DealActSeqLenIsZero(tempLoopInfo.bIdx, tempLoopInfo.n2Idx, 0U); + + if ASCEND_IS_AIV { + if (bN2LoopIdx + 1 == constInfo.bN2End && gloop > 0) { + CrossCoreWaitFlag(constInfo.syncC1V1); + vectorService.ProcessVec1(runInfo[1 - gloop % LI_QUANT_PRELOAD_TASK_CACHE_SIZE]); + CrossCoreSetFlag( + constInfo.syncV1C1); // 反向同步 1 + } + } + continue; + } + for (uint32_t gS1LoopIdx = constInfo.gS1Start; gS1LoopIdx < tempLoopInfo.gS1LoopEnd; gS1LoopIdx++) { + CalcS2LoopParams(bN2LoopIdx, gS1LoopIdx); + bool isEnd = (bN2LoopIdx + 1 == constInfo.bN2End) && (gS1LoopIdx + 1 == tempLoopInfo.gS1LoopEnd); + uint32_t extraLoop = isEnd ? LI_QUANT_PRELOAD_TASK_CACHE_SIZE - 1 : 0; // 只preload一轮 + + for (uint32_t s2LoopIdx = constInfo.s2Start; + s2LoopIdx < (tempLoopInfo.s2LoopEnd + extraLoop); + s2LoopIdx++) { + ProcessBaseBlock(gloop, s2LoopIdx, runInfo); + ++gloop; + } + constInfo.s2Start = 0; + } + if (tempLoopInfo.needDealActS1LessThanS1) { + DealActSeqLenIsZero(tempLoopInfo.bIdx, tempLoopInfo.n2Idx, tempLoopInfo.actS1Size); + } + constInfo.gS1Start = 0; + } + + if ASCEND_IS_AIV { + vectorService.FreeEventID(); + CrossCoreWaitFlag(constInfo.syncC1V0); + CrossCoreWaitFlag(constInfo.syncC1V0); + } else { + matmulService.FreeEventID(); + CrossCoreWaitFlag(constInfo.syncV1C1); + CrossCoreWaitFlag(constInfo.syncV1C1); + } +} + +template +__aicore__ inline void QLIV2Preload::ProcessBaseBlock( + uint32_t loop, uint64_t s2LoopIdx, + QLIV2Common::RunInfo runInfo[LI_QUANT_PRELOAD_TASK_CACHE_SIZE]) +{ + int32_t curTaskId = loop % LI_QUANT_PRELOAD_TASK_CACHE_SIZE; + QLIV2Common::RunInfo &curRunInfo = runInfo[curTaskId]; + QLIV2Common::RunInfo &lastRunInfo = runInfo[1 - curTaskId]; + + CalcRunInfo(loop, s2LoopIdx, curRunInfo); + + if (curRunInfo.isValid) { + if ASCEND_IS_AIC { + if (curRunInfo.isFirstS2InnerLoop) { + CrossCoreWaitFlag(constInfo.syncV0C1); + } + CrossCoreWaitFlag(constInfo.syncV1C1); // 反向同步 1 + matmulService.ComputeMm1(curRunInfo); + CrossCoreSetFlag(constInfo.syncC1V1); + if (curRunInfo.isLastS2InnerLoop) { + // 反向同步 0 + CrossCoreSetFlag(constInfo.syncC1V0); + } + } else { + if (curRunInfo.isFirstS2InnerLoop) { + CrossCoreWaitFlag(constInfo.syncC1V0); // 反向同步 0 + vectorService.ProcessVec0(curRunInfo); + CrossCoreSetFlag(constInfo.syncV0C1); + } + } + } + + if (lastRunInfo.isValid) { + if ASCEND_IS_AIV { + CrossCoreWaitFlag(constInfo.syncC1V1); + vectorService.ProcessVec1(lastRunInfo); + CrossCoreSetFlag(constInfo.syncV1C1); // 反向同步 1 + } + lastRunInfo.isValid = false; + } +} + +template +__aicore__ inline void QLIV2Preload::ProcessDecode() +{ + if ASCEND_IS_AIV { + vectorService.InitLDBuffers(pipe); + ICachePreLoad(LD_PREFETCH_LEN); + SyncAll(); + if (ldInfo.isLdCoreEnable) { + vectorService.ProcessLD(); + } + } +} + +} // namespace QLIV2Kernel +#endif // QUANT_LIGHTNING_INDEXER_V2_KERNEL_H \ No newline at end of file diff --git a/xllm_ops/attention/quant_lightning_indexer_v2/op_kernel/arch22/quant_lightning_indexer_v2_service_cube_arch22.h b/xllm_ops/attention/quant_lightning_indexer_v2/op_kernel/arch22/quant_lightning_indexer_v2_service_cube_arch22.h new file mode 100644 index 0000000..b8670a2 --- /dev/null +++ b/xllm_ops/attention/quant_lightning_indexer_v2/op_kernel/arch22/quant_lightning_indexer_v2_service_cube_arch22.h @@ -0,0 +1,617 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file quant_lightning_indexer_v2_service_cube_arch22.h + * \brief use 5 buffer for matmul l1, better pipeline + */ +#ifndef QUANT_LIGHTNING_INDEXER_V2_SERVICE_CUBE_H +#define QUANT_LIGHTNING_INDEXER_V2_SERVICE_CUBE_H + +#include "kernel_operator.h" +#include "kernel_operator_list_tensor_intf.h" +#include "kernel_tiling/kernel_tiling.h" +#include "lib/matmul_intf.h" +#include "lib/matrix/matmul/tiling.h" +#include "quant_lightning_indexer_v2_common_arch22.h" + +namespace QLIV2Kernel { +using namespace QLIV2Common; +struct MmInfo { + int64_t s2L0LoopId; + int64_t s1gL0LoopId; + int64_t s2L0RealSize; + int64_t s2GmOffset; +}; + +template +class QLIV2Matmul { +public: + using Q_T = typename QLIV2T::queryType; + using K_T = typename QLIV2T::keyType; + + __aicore__ inline QLIV2Matmul() {}; + __aicore__ inline void InitBuffers(TPipe *pipe); + __aicore__ inline void InitMm1GlobalTensor(const GlobalTensor &blkTableGm, const GlobalTensor &keyGm, + const GlobalTensor &queryGm, const GlobalTensor &mm1ResGm, + const GlobalTensor &weightWorkspaceGm); + __aicore__ inline void InitParams(const ConstInfo &constInfo); + __aicore__ inline void AllocEventID(); + __aicore__ inline void FreeEventID(); + __aicore__ inline void ComputeMm1(const QLIV2Common::RunInfo &runInfo); + + static constexpr IsResetLoad3dConfig LOAD3DV2_CONFIG = {true, true}; // isSetFMatrix isSetPadding; + static constexpr uint64_t DOUBLE_BUF_NUM = 2; + static constexpr uint64_t L0AB_BUF_NUM = 4; + + static constexpr uint32_t KEY_MTE1_MTE2_EVENT = EVENT_ID2; + static constexpr uint32_t QW_MTE1_MTE2_EVENT = EVENT_ID5; // KEY_MTE1_MTE2_EVENT + DOUBLE_BUF_NUM; + static constexpr uint32_t M_MTE1_EVENT = EVENT_ID3; + static constexpr uint32_t M_FIX_EVENT = EVENT_ID0; + static constexpr uint32_t FIX_M_EVENT = EVENT_ID2; + static constexpr uint32_t FIX_MTE1_EVENT = EVENT_ID4; + + static constexpr uint64_t S8_BLOCK_CUBE = 32; + + static constexpr uint32_t MTE2_MTE1_EVENT = EVENT_ID2; + static constexpr uint32_t MTE1_M_EVENT = EVENT_ID2; + + static constexpr uint64_t D_BASIC_BLOCK = 128; + static constexpr uint64_t S1G_BASIC_BLOCK_L1 = 256; + + static constexpr uint64_t S1G_BASIC_BLOCK_L0 = 128; + static constexpr uint64_t S2_BASIC_BLOCK_L0 = 128; + + static constexpr uint64_t QUERY_BUFFER_OFFSET = S1G_BASIC_BLOCK_L1 * D_BASIC_BLOCK; + static constexpr uint64_t SL1_BUFFER_OFFSET = S1G_BASIC_BLOCK_L0 * S2_BASIC_BLOCK_L0; + static constexpr uint64_t KEY_BUFFER_OFFSET = S2_BASIC_BLOCK_L0 * D_BASIC_BLOCK; + static constexpr uint64_t WEIGHT_BUFFER_OFFSET = S1G_BASIC_BLOCK_L1 * BLOCK_CUBE; + static constexpr uint64_t L0AB_BUFFER_OFFSET_S8_16K = 16 * 1024; + static constexpr uint64_t L0AB_BUFFER_OFFSET_FP16_16K = 16 * 512; + static constexpr uint64_t L0C_BUFFER_OFFSET = 64 * 256; + +private: + __aicore__ inline void WeightDmaCopy(uint64_t s1gL1RealSize, const QLIV2Common::RunInfo &runInfo); + __aicore__ inline void LoadKeyToL0b(uint64_t s2L0RealSize); + __aicore__ inline void LoadQueryToL0a(uint64_t s1gL1Offset, uint64_t s1gL1RealSize, uint64_t s1gL0RealSize); + __aicore__ inline void QueryNd2Nz(uint64_t s1gL1RealSize, const QLIV2Common::RunInfo &runInfo); + __aicore__ inline void KeyNd2NzForPA(uint64_t s2L1RealSize, uint64_t s2GmOffset, + const QLIV2Common::RunInfo &runInfo); + __aicore__ inline void KeyNd2Nz(uint64_t s2L1RealSize, const MmInfo &mmInfo, const QLIV2Common::RunInfo &runInfo); + __aicore__ inline void FixpSToL1(uint64_t s1gL0RealSize, uint64_t s2L0RealSize); + __aicore__ inline void LoadSToL0b(uint64_t s1gL1RealSize, uint64_t s2L0RealSize, uint64_t sL1BufIdx, + int64_t mStartPt); + __aicore__ inline void LoadWeightToL0a(uint64_t s1gL1Offset); + __aicore__ inline void ComputeWs(uint64_t s1gL0RealSize, uint64_t s2L0RealSize, int64_t s1gOffset); + __aicore__ inline void FixpResToGm(uint64_t s1L0RealCount, uint64_t s2L0RealSize, uint64_t s1GmOffset, + uint64_t s2GmOffset, const QLIV2Common::RunInfo &runInfo); + __aicore__ inline void ComputeQk(uint64_t s1gL0RealSize, uint64_t s2L0RealSize); + __aicore__ inline void ProcessWs(uint64_t s1gL0RealSize, uint64_t s1gL1Offset, uint64_t sL1BufIdx, + const MmInfo &mmInfo, const QLIV2Common::RunInfo &runInfo); + __aicore__ inline void ProcessQk(uint64_t s1gL0RealSize, uint64_t s1gL1Offset, uint64_t s1L0LoopCnt, + const MmInfo &mmInfo, const QLIV2Common::RunInfo &runInfo); + __aicore__ inline void CalcMmInfo(MmInfo &mmInfo, uint64_t loopIdx, uint64_t s1L0LoopCnt, const MmInfo &lastMmInfo, + const QLIV2Common::RunInfo &runInfo); + static constexpr LI_LAYOUT Q_LAYOUT_T = QLIV2T::layout; + static constexpr LI_LAYOUT K_LAYOUT_T = QLIV2T::keyLayout; + GlobalTensor blkTableGm_; + GlobalTensor keyGm_; + GlobalTensor queryGm_; + GlobalTensor weightGm_; + GlobalTensor mm1ResGm_; + + TBuf bufQL1_; + LocalTensor queryL1_; + TBuf bufKeyL1_; + LocalTensor keyL1_; + TBuf bufWeightL1_; + LocalTensor weightL1_; + TBuf bufSL1_; + LocalTensor sL1_; + + TBuf bufL0A_; + LocalTensor l0a_; + TBuf bufL0B_; + LocalTensor l0b_; + + TBuf bufL0C_; + LocalTensor cL0_; + + uint64_t keyL1BufIdx_ = 0; + uint64_t qwL1Mte2BufIdx_ = 0; + uint64_t sL1BufIdx_ = 0; + uint64_t l0BufIdx_ = 0; + uint64_t l0cBufIdx_ = 0; + + ConstInfo constInfo_; +}; + +template +__aicore__ inline void QLIV2Matmul::InitParams(const ConstInfo &constInfo) +{ + constInfo_ = constInfo; +} + +template +__aicore__ inline void QLIV2Matmul::InitBuffers(TPipe *pipe) +{ + pipe->InitBuffer(bufQL1_, DOUBLE_BUF_NUM * S1G_BASIC_BLOCK_L1 * D_BASIC_BLOCK * sizeof(Q_T)); + queryL1_ = bufQL1_.Get(); + pipe->InitBuffer(bufKeyL1_, DOUBLE_BUF_NUM * S2_BASIC_BLOCK_L0 * D_BASIC_BLOCK * sizeof(K_T)); + keyL1_ = bufKeyL1_.Get(); + + pipe->InitBuffer(bufWeightL1_, DOUBLE_BUF_NUM * S1G_BASIC_BLOCK_L1 * BLOCK_CUBE * sizeof(half)); + weightL1_ = bufWeightL1_.Get(); + pipe->InitBuffer(bufSL1_, DOUBLE_BUF_NUM * S2_BASIC_BLOCK_L0 * S1G_BASIC_BLOCK_L0 * sizeof(half)); + sL1_ = bufSL1_.Get(); + + pipe->InitBuffer(bufL0A_, 64 * 1024); + l0a_ = bufL0A_.Get(); + pipe->InitBuffer(bufL0B_, 64 * 1024); + l0b_ = bufL0B_.Get(); + + pipe->InitBuffer(bufL0C_, 128 * 1024); + cL0_ = bufL0C_.Get(); +} + +template +__aicore__ inline void QLIV2Matmul::InitMm1GlobalTensor(const GlobalTensor &blkTableGm, + const GlobalTensor &keyGm, + const GlobalTensor &queryGm, + const GlobalTensor &mm1ResGm, + const GlobalTensor &weightWorkspaceGm) +{ + blkTableGm_ = blkTableGm; + keyGm_ = keyGm; + queryGm_ = queryGm; + mm1ResGm_ = mm1ResGm; + weightGm_ = weightWorkspaceGm; +} + +template +__aicore__ inline void QLIV2Matmul::ProcessWs(uint64_t s1gL0RealSize, uint64_t s1gL1Offset, uint64_t sL1BufIdx, + const MmInfo &mmInfo, const QLIV2Common::RunInfo &runInfo) +{ + WaitFlag(FIX_M_EVENT + l0cBufIdx_ % DOUBLE_BUF_NUM); + for (int64_t s1gOffset = 0; s1gOffset < s1gL0RealSize; s1gOffset += constInfo_.gSize) { + WaitFlag(M_MTE1_EVENT + l0BufIdx_ % L0AB_BUF_NUM); + LoadSToL0b(s1gL0RealSize, mmInfo.s2L0RealSize, sL1BufIdx, s1gOffset); + LoadWeightToL0a(s1gOffset + s1gL1Offset); + + ComputeWs(s1gL0RealSize, mmInfo.s2L0RealSize, s1gOffset); + + SetFlag(M_MTE1_EVENT + l0BufIdx_ % L0AB_BUF_NUM); + l0BufIdx_++; + } + + FixpResToGm(s1gL0RealSize / constInfo_.gSize, mmInfo.s2L0RealSize, s1gL1Offset / constInfo_.gSize, + mmInfo.s2L0LoopId * S2_BASIC_BLOCK_L0, runInfo); + SetFlag(FIX_M_EVENT + l0cBufIdx_ % DOUBLE_BUF_NUM); + l0cBufIdx_++; +} + +template +__aicore__ inline void QLIV2Matmul::ProcessQk(uint64_t s1gL0RealSize, uint64_t s1gL1Offset, + uint64_t s1L0LoopCnt, + const MmInfo &mmInfo, const QLIV2Common::RunInfo &runInfo) +{ + if (mmInfo.s1gL0LoopId == 0) { + WaitFlag(KEY_MTE1_MTE2_EVENT + keyL1BufIdx_ % DOUBLE_BUF_NUM); + if constexpr (K_LAYOUT_T == LI_LAYOUT::PA_BBND) { + KeyNd2NzForPA(mmInfo.s2L0RealSize, runInfo.s2Idx * constInfo_.s2BaseSize + mmInfo.s2GmOffset, runInfo); + } else { + KeyNd2Nz(mmInfo.s2L0RealSize, mmInfo, runInfo); + } + + SetFlag(MTE2_MTE1_EVENT); + WaitFlag(MTE2_MTE1_EVENT); + } + + WaitFlag(M_MTE1_EVENT + l0BufIdx_ % L0AB_BUF_NUM); + LoadQueryToL0a(s1gL1Offset, runInfo.actMBaseSize, s1gL0RealSize); + LoadKeyToL0b(mmInfo.s2L0RealSize); + + if (mmInfo.s1gL0LoopId + 1 >= s1L0LoopCnt) { + SetFlag(KEY_MTE1_MTE2_EVENT + keyL1BufIdx_ % DOUBLE_BUF_NUM); + keyL1BufIdx_++; + } + + WaitFlag(FIX_M_EVENT + l0cBufIdx_ % DOUBLE_BUF_NUM); + ComputeQk(s1gL0RealSize, mmInfo.s2L0RealSize); + SetFlag(M_MTE1_EVENT + l0BufIdx_ % L0AB_BUF_NUM); + + FixpSToL1(s1gL0RealSize, mmInfo.s2L0RealSize); + SetFlag(FIX_M_EVENT + l0cBufIdx_ % DOUBLE_BUF_NUM); + l0BufIdx_++; + l0cBufIdx_++; +} + +template +__aicore__ inline void QLIV2Matmul::CalcMmInfo(MmInfo &mmInfo, uint64_t loopIdx, uint64_t s1L0LoopCnt, + const MmInfo &lastMmInfo, const QLIV2Common::RunInfo &runInfo) +{ + mmInfo.s2L0LoopId = loopIdx / s1L0LoopCnt; + mmInfo.s1gL0LoopId = loopIdx % s1L0LoopCnt; + + if (mmInfo.s1gL0LoopId == 0) { + mmInfo.s2GmOffset = mmInfo.s2L0LoopId * S2_BASIC_BLOCK_L0; + mmInfo.s2L0RealSize = mmInfo.s2GmOffset + S2_BASIC_BLOCK_L0 > runInfo.actualSingleProcessSInnerSize + ? runInfo.actualSingleProcessSInnerSize - mmInfo.s2GmOffset + : S2_BASIC_BLOCK_L0; + } else { + mmInfo.s2L0RealSize = lastMmInfo.s2L0RealSize; + } +} + +template +__aicore__ inline void QLIV2Matmul::ComputeMm1(const QLIV2Common::RunInfo &runInfo) +{ + if (runInfo.isFirstS2InnerLoop) { + WaitFlag(QW_MTE1_MTE2_EVENT + qwL1Mte2BufIdx_ % DOUBLE_BUF_NUM); + QueryNd2Nz(runInfo.actMBaseSize, runInfo); // 256 * 128 // L1BasicBlock + WeightDmaCopy(runInfo.actMBaseSize, runInfo); + } + int64_t loopIdx = 0; + int64_t s2L0LoopCnt = CeilDiv(runInfo.actualSingleProcessSInnerSize, S2_BASIC_BLOCK_L0); // 2048取128 + int64_t s1L0LoopCnt = CeilDiv(runInfo.actMBaseSize, S1G_BASIC_BLOCK_L0); // 256取128 + int64_t s1gL1Offset[2] = {0, static_cast(S1G_BASIC_BLOCK_L0)}; + int64_t s1gL0RealSize[2] = {s1L0LoopCnt > 1 ? static_cast(S1G_BASIC_BLOCK_L0) : runInfo.actMBaseSize, + runInfo.actMBaseSize - s1gL1Offset[1]}; + MmInfo mmInfo[2]; + CalcMmInfo(mmInfo[loopIdx & 1], loopIdx, s1L0LoopCnt, mmInfo[(loopIdx + 1) & 1], runInfo); + + ProcessQk(s1gL0RealSize[mmInfo[loopIdx & 1].s1gL0LoopId % s1L0LoopCnt], + s1gL1Offset[mmInfo[loopIdx & 1].s1gL0LoopId % s1L0LoopCnt], s1L0LoopCnt, mmInfo[loopIdx & 1], + runInfo); + + SetFlag(FIX_MTE1_EVENT + sL1BufIdx_ % DOUBLE_BUF_NUM); + sL1BufIdx_++; + loopIdx++; + + while (loopIdx < s2L0LoopCnt * s1L0LoopCnt) { + CalcMmInfo(mmInfo[loopIdx & 1], loopIdx, s1L0LoopCnt, mmInfo[(loopIdx + 1) & 1], runInfo); + + ProcessQk(s1gL0RealSize[mmInfo[loopIdx & 1].s1gL0LoopId % s1L0LoopCnt], + s1gL1Offset[mmInfo[loopIdx & 1].s1gL0LoopId % s1L0LoopCnt], s1L0LoopCnt, mmInfo[loopIdx & 1], + runInfo); + + SetFlag(FIX_MTE1_EVENT + sL1BufIdx_ % DOUBLE_BUF_NUM); + sL1BufIdx_++; + + WaitFlag(FIX_MTE1_EVENT + sL1BufIdx_ % DOUBLE_BUF_NUM); + + ProcessWs(s1gL0RealSize[mmInfo[(loopIdx + 1) & 1].s1gL0LoopId % s1L0LoopCnt], + s1gL1Offset[mmInfo[(loopIdx + 1) & 1].s1gL0LoopId % s1L0LoopCnt], sL1BufIdx_, + mmInfo[(loopIdx + 1) & 1], runInfo); + loopIdx++; + } + + WaitFlag(FIX_MTE1_EVENT + (sL1BufIdx_ + 1) % DOUBLE_BUF_NUM); + + ProcessWs(s1gL0RealSize[mmInfo[(loopIdx + 1) & 1].s1gL0LoopId % s1L0LoopCnt], + s1gL1Offset[mmInfo[(loopIdx + 1) & 1].s1gL0LoopId % s1L0LoopCnt], sL1BufIdx_ - 1, + mmInfo[(loopIdx + 1) & 1], runInfo); + + if (runInfo.isLastS2InnerLoop) { + SetFlag(QW_MTE1_MTE2_EVENT + qwL1Mte2BufIdx_ % DOUBLE_BUF_NUM); + qwL1Mte2BufIdx_++; + } +} + +// blkNum, blkSize, N2, D +template +__aicore__ inline void QLIV2Matmul::KeyNd2NzForPA(uint64_t s2L1RealSize, uint64_t s2GmOffset, + const QLIV2Common::RunInfo &runInfo) +{ + uint64_t s2L1Offset = 0; + while (s2L1Offset < s2L1RealSize) { + uint64_t s2BlkId = (s2L1Offset + s2GmOffset) / constInfo_.kCacheBlockSize; + uint64_t s2BlkOffset = (s2L1Offset + s2GmOffset) % constInfo_.kCacheBlockSize; + uint64_t keyGmOffset = blkTableGm_.GetValue(runInfo.bIdx * constInfo_.maxBlockNumPerBatch + s2BlkId) * + constInfo_.kCacheBlockSize * constInfo_.kHeadNum * constInfo_.headDim + + s2BlkOffset * constInfo_.headDim; + uint64_t s2Mte2Size = s2L1RealSize - s2L1Offset; + s2Mte2Size = s2BlkOffset + s2Mte2Size >= constInfo_.kCacheBlockSize ? constInfo_.kCacheBlockSize - s2BlkOffset + : s2Mte2Size; + Nd2NzParams nd2nzPara; + nd2nzPara.ndNum = 1; + nd2nzPara.nValue = s2Mte2Size; // 行数 + nd2nzPara.dValue = constInfo_.headDim; + nd2nzPara.srcDValue = constInfo_.headDim; + nd2nzPara.dstNzC0Stride = CeilAlign(s2L1RealSize, (uint64_t)BLOCK_CUBE); // 对齐到16 单位block + nd2nzPara.dstNzNStride = 1; + nd2nzPara.srcNdMatrixStride = 0; + nd2nzPara.dstNzMatrixStride = 0; + DataCopy(keyL1_[(keyL1BufIdx_ % DOUBLE_BUF_NUM) * KEY_BUFFER_OFFSET + s2L1Offset * S8_BLOCK_CUBE], + keyGm_[keyGmOffset], nd2nzPara); + + s2L1Offset += s2Mte2Size; + } +} + +template +__aicore__ inline void QLIV2Matmul::KeyNd2Nz(uint64_t s2L1RealSize, const MmInfo &mmInfo, + const QLIV2Common::RunInfo &runInfo) +{ + uint64_t dStride = constInfo_.headDim; + if constexpr (K_LAYOUT_T == LI_LAYOUT::BSND || K_LAYOUT_T == LI_LAYOUT::TND) { + dStride = constInfo_.headDim * constInfo_.kHeadNum; // constInfo_.kHeadNum + } + Nd2NzParams nd2nzPara; + nd2nzPara.ndNum = 1; + nd2nzPara.nValue = s2L1RealSize; // 行数 + nd2nzPara.dValue = constInfo_.headDim; + nd2nzPara.srcDValue = dStride; + nd2nzPara.dstNzC0Stride = CeilAlign(s2L1RealSize, (uint64_t)BLOCK_CUBE); // 对齐到16 单位block + nd2nzPara.dstNzNStride = 1; + nd2nzPara.srcNdMatrixStride = 0; + nd2nzPara.dstNzMatrixStride = 0; + // 默认一块buf最多放两份 + DataCopy(keyL1_[(keyL1BufIdx_ % DOUBLE_BUF_NUM) * KEY_BUFFER_OFFSET], + keyGm_[runInfo.tensorKeyOffset + mmInfo.s2GmOffset * constInfo_.headDim], nd2nzPara); +} + +// batch, s1, g, 1 +template +__aicore__ inline void QLIV2Matmul::WeightDmaCopy(uint64_t s1gL1RealSize, const QLIV2Common::RunInfo &runInfo) +{ + DataCopyParams copyInParams; + copyInParams.blockCount = 1; + copyInParams.blockLen = s1gL1RealSize; + copyInParams.srcStride = 0; + copyInParams.dstStride = 0; + DataCopy(weightL1_[(qwL1Mte2BufIdx_ % DOUBLE_BUF_NUM) * WEIGHT_BUFFER_OFFSET], + weightGm_[runInfo.loop % DOUBLE_BUF_NUM * BLOCK_CUBE * constInfo_.mBaseSize], copyInParams); +} + +// batch, s1, n2, g, d +template +__aicore__ inline void QLIV2Matmul::QueryNd2Nz(uint64_t s1gL1RealSize, const QLIV2Common::RunInfo &runInfo) +{ + Nd2NzParams nd2nzPara; + nd2nzPara.ndNum = 1; + nd2nzPara.nValue = s1gL1RealSize; // 行数 + nd2nzPara.dValue = constInfo_.headDim; + nd2nzPara.srcDValue = constInfo_.headDim; + nd2nzPara.dstNzC0Stride = CeilAlign(s1gL1RealSize, (uint64_t)BLOCK_CUBE); // 对齐到16 单位block + nd2nzPara.dstNzNStride = 1; + nd2nzPara.srcNdMatrixStride = 0; + nd2nzPara.dstNzMatrixStride = 0; + // 默认一块buf最多放两份 + DataCopy(queryL1_[(qwL1Mte2BufIdx_ % DOUBLE_BUF_NUM) * QUERY_BUFFER_OFFSET], queryGm_[runInfo.tensorQueryOffset], + nd2nzPara); +} + +// s1g, d +template +__aicore__ inline void QLIV2Matmul::LoadQueryToL0a(uint64_t s1gL1Offset, uint64_t s1gL1RealSize, + uint64_t s1gL0RealSize) +{ + LoadData3DParamsV2 loadData3DParams; + // SetFmatrixParams + loadData3DParams.l1H = CeilDiv(s1gL1RealSize, BLOCK_CUBE); // Hin=M1=8 + loadData3DParams.l1W = BLOCK_CUBE; // Win=M0 + loadData3DParams.channelSize = constInfo_.headDim; // Cin=K + + loadData3DParams.padList[0] = 0; + loadData3DParams.padList[1] = 0; + loadData3DParams.padList[2] = 0; + loadData3DParams.padList[3] = 255; // 尾部数据不影响滑窗的结果 + + // SetLoadToA0Params + loadData3DParams.mExtension = CeilAlign(s1gL0RealSize, BLOCK_CUBE); // M height维度目的 + loadData3DParams.kExtension = constInfo_.headDim; // K width维度目的 + loadData3DParams.mStartPt = s1gL1Offset; + loadData3DParams.kStartPt = 0; + loadData3DParams.strideW = 1; + loadData3DParams.strideH = 1; + loadData3DParams.filterW = 1; + loadData3DParams.filterSizeW = (1 >> 8) & 255; + loadData3DParams.filterH = 1; + loadData3DParams.filterSizeH = (1 >> 8) & 255; + loadData3DParams.dilationFilterW = 1; + loadData3DParams.dilationFilterH = 1; + loadData3DParams.enTranspose = 0; + loadData3DParams.fMatrixCtrl = 0; + + LoadData(l0a_[(l0BufIdx_ % L0AB_BUF_NUM) * L0AB_BUFFER_OFFSET_S8_16K], + queryL1_[(qwL1Mte2BufIdx_ % DOUBLE_BUF_NUM) * QUERY_BUFFER_OFFSET], + loadData3DParams); +} + +// s1, g, s2 --> 2 * 64* 128 +template +__aicore__ inline void QLIV2Matmul::LoadSToL0b(uint64_t s1gL1RealSize, uint64_t s2L0RealSize, + uint64_t sL1BufIdx, + int64_t mStartPt) +{ + LoadData3DParamsV2 loadData3DParams; + // SetFmatrixParams + loadData3DParams.l1H = S1G_BASIC_BLOCK_L0 / BLOCK_CUBE; // Hin=M1=8 + loadData3DParams.l1W = BLOCK_CUBE; // Win=M0 + loadData3DParams.channelSize = CeilAlign(s2L0RealSize, BLOCK_CUBE); // Cin=K + + loadData3DParams.padList[0] = 0; + loadData3DParams.padList[1] = 0; + loadData3DParams.padList[2] = 0; + loadData3DParams.padList[3] = 255; // 尾部数据不影响滑窗的结果 + + // SetLoadToA0Params + loadData3DParams.mExtension = constInfo_.gSize; // M height维度目的 + loadData3DParams.kExtension = CeilAlign(s2L0RealSize, BLOCK_CUBE); // K width维度目的 + loadData3DParams.kStartPt = 0; + loadData3DParams.strideW = 1; + loadData3DParams.strideH = 1; + loadData3DParams.filterW = 1; + loadData3DParams.filterSizeW = (1 >> 8) & 255; + loadData3DParams.filterH = 1; + loadData3DParams.filterSizeH = (1 >> 8) & 255; + loadData3DParams.dilationFilterW = 1; + loadData3DParams.dilationFilterH = 1; + loadData3DParams.enTranspose = 1; + loadData3DParams.fMatrixCtrl = 0; + + loadData3DParams.mStartPt = mStartPt; + LoadData( + l0b_.template ReinterpretCast()[(l0BufIdx_ % L0AB_BUF_NUM) * L0AB_BUFFER_OFFSET_FP16_16K], + sL1_[(sL1BufIdx % DOUBLE_BUF_NUM) * SL1_BUFFER_OFFSET], loadData3DParams); +} + +// s1,g,1(16), 2,64,16 +template +__aicore__ inline void QLIV2Matmul::LoadWeightToL0a(uint64_t s1gL1Offset) +{ + LoadData2DParams loadData2DParams; + loadData2DParams.startIndex = 0; + loadData2DParams.repeatTimes = CeilDiv(constInfo_.gSize, BLOCK_CUBE); + loadData2DParams.srcStride = 1; + loadData2DParams.dstGap = 0; + loadData2DParams.ifTranspose = true; + LoadData(l0a_.template ReinterpretCast()[(l0BufIdx_ % L0AB_BUF_NUM) * L0AB_BUFFER_OFFSET_FP16_16K], + weightL1_[(qwL1Mte2BufIdx_ % DOUBLE_BUF_NUM) * WEIGHT_BUFFER_OFFSET + s1gL1Offset* BLOCK_CUBE], + loadData2DParams); +} + +// s2, d -> 128,128 +template +__aicore__ inline void QLIV2Matmul::LoadKeyToL0b(uint64_t s2L0RealSize) +{ + LoadData2DParams loadData2DParams; + loadData2DParams.startIndex = 0; + loadData2DParams.repeatTimes = CeilDiv(s2L0RealSize, BLOCK_CUBE) * CeilDiv(constInfo_.headDim, S8_BLOCK_CUBE); + loadData2DParams.srcStride = 1; + loadData2DParams.dstGap = 0; + loadData2DParams.ifTranspose = false; + LoadData(l0b_[(l0BufIdx_ % L0AB_BUF_NUM) * L0AB_BUFFER_OFFSET_S8_16K], + keyL1_[(keyL1BufIdx_ % DOUBLE_BUF_NUM) * KEY_BUFFER_OFFSET], loadData2DParams); +} + +// A: s1,g,1(16) B: s1,g,s2 C: s1, 1(16), s2 +template +__aicore__ inline void QLIV2Matmul::ComputeWs(uint64_t s1gL0RealSize, uint64_t s2L0RealSize, int64_t s1gOffset) +{ + SetFlag(MTE1_M_EVENT); + WaitFlag(MTE1_M_EVENT); + MmadParams mmadParams; + mmadParams.m = BLOCK_CUBE; + mmadParams.n = s2L0RealSize; + mmadParams.k = constInfo_.gSize; + mmadParams.cmatrixInitVal = true; + mmadParams.cmatrixSource = false; + Mmad(cL0_.template ReinterpretCast()[(l0cBufIdx_ % DOUBLE_BUF_NUM) * L0C_BUFFER_OFFSET + + s1gOffset * S2_BASIC_BLOCK_L0], + l0a_.template ReinterpretCast()[(l0BufIdx_ % L0AB_BUF_NUM) * L0AB_BUFFER_OFFSET_FP16_16K], + l0b_.template ReinterpretCast()[(l0BufIdx_ % L0AB_BUF_NUM) * L0AB_BUFFER_OFFSET_FP16_16K], + mmadParams); +} + +template +__aicore__ inline void QLIV2Matmul::ComputeQk(uint64_t s1gL0RealSize, uint64_t s2L0RealSize) +{ + SetFlag(MTE1_M_EVENT); + WaitFlag(MTE1_M_EVENT); + + MmadParams mmadParams; + mmadParams.m = CeilAlign(s1gL0RealSize, BLOCK_CUBE); + mmadParams.n = s2L0RealSize; + mmadParams.k = constInfo_.headDim; + mmadParams.cmatrixInitVal = true; + mmadParams.cmatrixSource = false; + Mmad(cL0_[(l0cBufIdx_ % DOUBLE_BUF_NUM) * L0C_BUFFER_OFFSET], + l0a_[(l0BufIdx_ % L0AB_BUF_NUM) * L0AB_BUFFER_OFFSET_S8_16K], + l0b_[(l0BufIdx_ % L0AB_BUF_NUM) * L0AB_BUFFER_OFFSET_S8_16K], mmadParams); + if ((mmadParams.m / 16) * (mmadParams.n / 16) < 10) { + PipeBarrier(); + } +} + +template +__aicore__ inline void QLIV2Matmul::FixpSToL1(uint64_t s1gL0RealSize, uint64_t s2L0RealSize) +{ + SetFlag(M_FIX_EVENT); + WaitFlag(M_FIX_EVENT); + DataCopyCO12DstParams params; + params.mSize = CeilAlign(s1gL0RealSize, BLOCK_CUBE); + params.nSize = CeilAlign(s2L0RealSize, BLOCK_CUBE); + params.dstStride = S1G_BASIC_BLOCK_L0; + params.srcStride = params.mSize; + params.quantPre = QuantMode_t::DEQF16; + params.reluPre = 1; + params.channelSplit = 0; + params.nz2ndEn = 0; + SetFixpipePreQuantFlag(0x3a800000); + DataCopy(sL1_[(sL1BufIdx_ % DOUBLE_BUF_NUM) * SL1_BUFFER_OFFSET], + cL0_[(l0cBufIdx_ % DOUBLE_BUF_NUM) * L0C_BUFFER_OFFSET], params); +} + +template +__aicore__ inline void QLIV2Matmul::FixpResToGm(uint64_t s1L0RealCount, uint64_t s2L0RealSize, + uint64_t s1GmOffset, + uint64_t s2GmOffset, const QLIV2Common::RunInfo &runInfo) +{ + SetFlag(M_FIX_EVENT); + WaitFlag(M_FIX_EVENT); + + AscendC::DataCopyCO12DstParams intriParams; + intriParams.mSize = 1; + intriParams.nSize = s2L0RealSize; + intriParams.dstStride = constInfo_.s2BaseSize; + intriParams.srcStride = 16; + // set mode according to dtype + intriParams.quantPre = QuantMode_t::NoQuant; + intriParams.nz2ndEn = true; + intriParams.reluPre = 0; + AscendC::SetFixpipeNz2ndFlag(s1L0RealCount, CeilDiv(constInfo_.gSize, BLOCK_CUBE) * S2_BASIC_BLOCK_L0 / BLOCK_CUBE, + 2048); + AscendC::DataCopy(mm1ResGm_[(runInfo.loop % 2) * constInfo_.mBaseSize / constInfo_.gSize * constInfo_.s2BaseSize + + s1GmOffset * intriParams.dstStride + s2GmOffset], + cL0_.template ReinterpretCast()[(l0cBufIdx_ % DOUBLE_BUF_NUM) * L0C_BUFFER_OFFSET], + intriParams); +} + +template +__aicore__ inline void QLIV2Matmul::AllocEventID() +{ + SetFlag(KEY_MTE1_MTE2_EVENT + 0); + SetFlag(KEY_MTE1_MTE2_EVENT + 1); + SetFlag(KEY_MTE1_MTE2_EVENT + 2); + + SetFlag(QW_MTE1_MTE2_EVENT + 0); + SetFlag(QW_MTE1_MTE2_EVENT + 1); + + SetFlag(M_MTE1_EVENT + 0); + SetFlag(M_MTE1_EVENT + 1); + SetFlag(M_MTE1_EVENT + 2); + SetFlag(M_MTE1_EVENT + 3); + + SetFlag(FIX_M_EVENT + 0); + SetFlag(FIX_M_EVENT + 1); +} + +template +__aicore__ inline void QLIV2Matmul::FreeEventID() +{ + WaitFlag(KEY_MTE1_MTE2_EVENT + 0); + WaitFlag(KEY_MTE1_MTE2_EVENT + 1); + WaitFlag(KEY_MTE1_MTE2_EVENT + 2); + + WaitFlag(QW_MTE1_MTE2_EVENT + 0); + WaitFlag(QW_MTE1_MTE2_EVENT + 1); + + WaitFlag(M_MTE1_EVENT + 0); + WaitFlag(M_MTE1_EVENT + 1); + WaitFlag(M_MTE1_EVENT + 2); + WaitFlag(M_MTE1_EVENT + 3); + + WaitFlag(FIX_M_EVENT + 0); + WaitFlag(FIX_M_EVENT + 1); +} +} // namespace QLIV2Kernel +#endif // QUANT_LIGHTNING_INDEXER_V2_SERVICE_CUBE_H \ No newline at end of file diff --git a/xllm_ops/attention/quant_lightning_indexer_v2/op_kernel/arch22/quant_lightning_indexer_v2_service_vector_arch22.h b/xllm_ops/attention/quant_lightning_indexer_v2/op_kernel/arch22/quant_lightning_indexer_v2_service_vector_arch22.h new file mode 100644 index 0000000..82d0713 --- /dev/null +++ b/xllm_ops/attention/quant_lightning_indexer_v2/op_kernel/arch22/quant_lightning_indexer_v2_service_vector_arch22.h @@ -0,0 +1,627 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file quant_lightning_indexer_v2_service_vector_arch22.h + * \brief + */ +#ifndef QUANT_LIGHTNING_INDEXER_V2_SERVICE_VECTOR_H +#define QUANT_LIGHTNING_INDEXER_V2_SERVICE_VECTOR_H + +#include "kernel_operator.h" +#include "kernel_operator_list_tensor_intf.h" +#include "kernel_tiling/kernel_tiling.h" +#include "lib/matmul_intf.h" +#include "lib/matrix/matmul/tiling.h" +#include "quant_lightning_indexer_v2_common_arch22.h" +#include "quant_lightning_indexer_v2_vector.h" + +namespace QLIV2Kernel { +using namespace QLIV2Common; +using namespace QLIV2ServiceVec; +constexpr uint32_t BASE_TOPK = 2048; +constexpr uint32_t BASE_TOPK_VALUE_IDX_SIZE = 4096; +constexpr uint32_t ELE_NUM_32 = 32; +constexpr uint32_t ELE_NUM_128 = 128; +constexpr uint32_t ELE_NUM_512 = 512; + +template +class QLIV2Vector { +public: + // =================================类型定义区================================= + static constexpr LI_LAYOUT Q_LAYOUT_T = QLIV2T::layout; + static constexpr LI_LAYOUT K_LAYOUT_T = QLIV2T::keyLayout; + static constexpr bool PAGE_ATTENTION = QLIV2T::pageAttention; + // MM输出数据类型, 当前只支持float + using MM1_OUT_T = float; + + __aicore__ inline QLIV2Vector() {}; + __aicore__ inline void ProcessVec0(const QLIV2Common::RunInfo &info); + __aicore__ inline void ProcessVec1(const QLIV2Common::RunInfo &info); + __aicore__ inline void InitBuffers(TPipe *pipe); + __aicore__ inline void InitParams(const struct QLIV2Common::ConstInfo &constInfo, + const struct QLIV2Common::LdSplitCoreInfo &ldInfo, + const QLIV2TilingData *__restrict tilingData); + __aicore__ inline void ProcessLD(); + __aicore__ inline void InitVecWorkspaceTensor(GlobalTensor vec0OutGm, GlobalTensor mm1ResGm, + GlobalTensor vec1ResGm); + __aicore__ inline void InitVecInputTensor(GlobalTensor weightsGm, GlobalTensor qScaleGm, + GlobalTensor kScaleGm, GlobalTensor indiceOutGm, + GlobalTensor blockTableGm); + __aicore__ inline void CleanInvalidOutput(int64_t invalidS1offset); + __aicore__ inline int32_t AlignS2(int32_t cuS2Len); + __aicore__ inline void AllocEventID(); + __aicore__ inline void FreeEventID(); + __aicore__ inline void InitLDBuffers(TPipe *pipe); + +protected: + GlobalTensor mm1ResGm; + GlobalTensor vec1ResGm; + GlobalTensor weightsGm; + GlobalTensor qScaleGm; + GlobalTensor kScaleGm; + GlobalTensor vec0OutGm; + GlobalTensor indiceOutGm; + GlobalTensor blockTableGm; + // =================================常量区================================= + +private: + __aicore__ inline void GetKeyScale(const QLIV2Common::RunInfo &runInfo, const LocalTensor &resUb, + int64_t batchId, int64_t startS2, int64_t getLen); + // ================================Local Buffer区==================================== + // queue + TQue inQueue_; + TQue outQueue_; + + // tmp buff for vector + TBuf sortOutBuf_; + TBuf indexBuf_; + TBuf tmpBuf_; + + // tmp buff for LD + TBuf<> ldToBeMrgBuf_; + TBuf<> ldTmpBuf_; + TBuf<> ldOutValueBuf_; + TBuf<> ldOutIdxBuf_; + + LocalTensor globalTopkIndice_; + LocalTensor globalTopkUb_; + + int32_t blockId_ = -1; + // para for vector + int32_t groupInner_ = 0; + int32_t globalTopkNum_ = 0; + int64_t blockS2StartIdx_ = 0; + int32_t gSize_ = 0; + int32_t kSeqSize_ = 0; + int32_t kHeadNum_ = 0; + int32_t qHeadNum_ = 0; + int32_t s1BaseSize_ = 0; + int32_t s2BaseSize_ = 0; + int32_t kCacheBlockSize_ = 0; + int32_t maxBlockNumPerBatch_ = 0; + + // para for LD + uint32_t mrgListNum_ = 4; + + struct QLIV2Common::ConstInfo constInfo_; + struct QLIV2Common::LdSplitCoreInfo ldInfo_; +}; + +template +__aicore__ inline void QLIV2Vector::GetKeyScale(const QLIV2Common::RunInfo &runInfo, + const LocalTensor &resUb, + int64_t batchId, int64_t startS2, int64_t getLen) +{ + // startS2一定能整除kCacheBlockSize_ + AscendC::DataCopyPadExtParams padParams{false, 0, 0, 0}; + AscendC::DataCopyExtParams copyInParams; + if constexpr (PAGE_ATTENTION) { + int32_t startBlockTableIdx = startS2 / kCacheBlockSize_; + int32_t startBlockTableOffset = startS2 % kCacheBlockSize_; + int32_t blockTableBatchOffset = batchId * maxBlockNumPerBatch_; + copyInParams.blockCount = 1; + copyInParams.srcStride = 0; + copyInParams.dstStride = 0; + copyInParams.rsv = 0; + int32_t resUbBaseOffset = 0; + if (startBlockTableOffset > 0) { + int32_t firstPartLen = + kCacheBlockSize_ - startBlockTableOffset > getLen ? getLen : kCacheBlockSize_ - startBlockTableOffset; + copyInParams.blockLen = firstPartLen * sizeof(half); + int32_t blockId = blockTableGm.GetValue(blockTableBatchOffset + startBlockTableIdx); + SetWaitFlag(HardEvent::S_MTE2); + AscendC::DataCopyPad(resUb, kScaleGm[blockId * kCacheBlockSize_ + startBlockTableOffset], + copyInParams, padParams); + startBlockTableIdx++; + getLen = getLen - firstPartLen; + resUbBaseOffset = firstPartLen; + } + int32_t getLoopNum = CeilDiv(getLen, kCacheBlockSize_); + copyInParams.blockLen = kCacheBlockSize_ * sizeof(half); + for (int32_t i = 0; i < getLoopNum; i++) { + if (i == getLoopNum - 1) { + copyInParams.blockLen = (getLen - i * kCacheBlockSize_) * sizeof(half); + } + int32_t blockId = blockTableGm.GetValue(blockTableBatchOffset + startBlockTableIdx + i); + SetWaitFlag(HardEvent::S_MTE2); + AscendC::DataCopyPad(resUb[resUbBaseOffset + i * kCacheBlockSize_], kScaleGm[blockId * kCacheBlockSize_], + copyInParams, padParams); + } + } else { + copyInParams.blockCount = 1; + copyInParams.blockLen = getLen * sizeof(half); + copyInParams.srcStride = 0; + copyInParams.dstStride = 0; + copyInParams.rsv = 0; + AscendC::DataCopyPad(resUb, kScaleGm[runInfo.tensorKeyScaleOffset], copyInParams, padParams); + } +} + +template +__aicore__ inline void QLIV2Vector::InitBuffers(TPipe *pipe) +{ + pipe->InitBuffer(inQueue_, 2, s2BaseSize_ * sizeof(float) * 2); // 32KB + pipe->InitBuffer(outQueue_, 1, BASE_TOPK * sizeof(float)); // 8 KB + pipe->InitBuffer(indexBuf_, s2BaseSize_ * sizeof(int32_t)); // 8 KB + pipe->InitBuffer(tmpBuf_, 64 * 1024); // 64KB + pipe->InitBuffer(sortOutBuf_, CeilDiv(s1BaseSize_, 2) * BASE_TOPK_VALUE_IDX_SIZE * sizeof(float)); // 32KB + + globalTopkIndice_ = indexBuf_.Get(); + globalTopkUb_ = sortOutBuf_.Get(); + globalTopkNum_ = 0; + + // 基本块执行前初始化UB和GM + // step1. 初始化一个有序索引 0 - s2BaseSize_ + ArithProgression(globalTopkIndice_, 0, 1, s2BaseSize_); + // step2. globalTopkUb_ [CeilDiv(s1BaseSize_, 2), BASE_TOPK, 2] -inf,-1 + InitSortOutBuf(globalTopkUb_, CeilDiv(s1BaseSize_, 2) * BASE_TOPK_VALUE_IDX_SIZE); +} + +template +__aicore__ inline void QLIV2Vector::InitLDBuffers(TPipe *pipe) +{ + pipe->Reset(); + pipe->InitBuffer(ldToBeMrgBuf_, BASE_TOPK_VALUE_IDX_SIZE * mrgListNum_ * sizeof(float)); + pipe->InitBuffer(ldTmpBuf_, BASE_TOPK_VALUE_IDX_SIZE * mrgListNum_ * sizeof(float)); + pipe->InitBuffer(ldOutValueBuf_, BASE_TOPK * sizeof(float)); + pipe->InitBuffer(ldOutIdxBuf_, BASE_TOPK * sizeof(int32_t)); +} + +template +__aicore__ inline void QLIV2Vector::InitParams(const struct QLIV2Common::ConstInfo &constInfo, + const struct QLIV2Common::LdSplitCoreInfo &ldInfo, + const QLIV2TilingData *__restrict tilingData) +{ + this->constInfo_ = constInfo; + this->ldInfo_ = ldInfo; + blockS2StartIdx_ = 0; + gSize_ = constInfo.gSize; + kSeqSize_ = constInfo.kSeqSize; + // define N2 para + kHeadNum_ = constInfo.kHeadNum; + qHeadNum_ = constInfo.qHeadNum; + // define MMBase para + s1BaseSize_ = constInfo.s1BaseSize; // 4 + s2BaseSize_ = constInfo.s2BaseSize; // 2048 + kCacheBlockSize_ = constInfo.kCacheBlockSize; + maxBlockNumPerBatch_ = constInfo.maxBlockNumPerBatch; + blockId_ = GetBlockIdx(); +} + +template +__aicore__ inline void QLIV2Vector::InitVecInputTensor(GlobalTensor weightsGm, + GlobalTensor qScaleGm, + GlobalTensor kScaleGm, + GlobalTensor indiceOutGm, + GlobalTensor blockTableGm) +{ + this->weightsGm = weightsGm; + this->qScaleGm = qScaleGm; + this->kScaleGm = kScaleGm; + this->indiceOutGm = indiceOutGm; + this->blockTableGm = blockTableGm; +} + +template +__aicore__ inline void QLIV2Vector::InitVecWorkspaceTensor(GlobalTensor vec0OutGm, + GlobalTensor mm1ResGm, + GlobalTensor vec1ResGm) +{ + this->mm1ResGm = mm1ResGm; + this->vec1ResGm = vec1ResGm; + this->vec0OutGm = vec0OutGm; +} + +template +__aicore__ inline void QLIV2Vector::AllocEventID() +{ +} + +template +__aicore__ inline void QLIV2Vector::FreeEventID() +{ +} + +template +__aicore__ inline void QLIV2Vector::CleanInvalidOutput(int64_t invalidS1offset) +{ + // init -1 and copy to output + LocalTensor valueULocal = outQueue_.AllocTensor(); + LocalTensor idxULocal1 = valueULocal.template ReinterpretCast(); + Duplicate(idxULocal1, constInfo_.INVALID_IDX, constInfo_.sparseCount); + outQueue_.EnQue(valueULocal); + valueULocal = outQueue_.DeQue(); + QLIV2ServiceVec::CopyOut(indiceOutGm[invalidS1offset], idxULocal1, constInfo_.sparseCount); + outQueue_.FreeTensor(valueULocal); +} + +template +__aicore__ inline void QLIV2Vector::ProcessVec0(const QLIV2Common::RunInfo &info) +{ + // 只需要一个v核做 + if (blockId_ % 2 != 0) { + return; + } + int32_t cuBaseS1Idx = info.gS1Idx * s1BaseSize_; + // 计算输出w基地址偏移 偶数循环 -> 0 + aic_offset 奇数循环 -> 4*64 + aic_offset + int64_t vec0OutGmOffset = (info.loop % 2) * ((s1BaseSize_ * gSize_ * BLOCK_CUBE)); + // 计算输入weight的地址偏移，qScale的地址偏移与weight相同 + int64_t weightGmOffset = info.tensorWeightsOffset + cuBaseS1Idx * qHeadNum_; + // 当前需要计算的S1行数，处理尾块场景 + int32_t cuS1ProcNum = cuBaseS1Idx + s1BaseSize_ > info.actS1Size ? info.actS1Size % s1BaseSize_ : s1BaseSize_; + int32_t cuProcEleNum = cuS1ProcNum * gSize_; + + LocalTensor inWeightsUb = inQueue_.AllocTensor(); + LocalTensor inQScaleUb = inWeightsUb[cuProcEleNum]; + AscendC::DataCopyPadExtParams padParams{false, 0, 0, 0}; + AscendC::DataCopyExtParams copyInParams; + copyInParams.blockCount = 1; + copyInParams.blockLen = cuProcEleNum * sizeof(half); + copyInParams.srcStride = 0; + copyInParams.dstStride = 0; + copyInParams.rsv = 0; + AscendC::DataCopyPad(inWeightsUb, weightsGm[weightGmOffset], copyInParams, padParams); + AscendC::DataCopyPad(inQScaleUb, qScaleGm[weightGmOffset], copyInParams, padParams); + + inQueue_.EnQue(inWeightsUb); + inWeightsUb = inQueue_.DeQue(); + AscendC::Mul(inWeightsUb, inWeightsUb, inQScaleUb, cuProcEleNum); + PipeBarrier(); + LocalTensor resUb = outQueue_.AllocTensor(); + AscendC::Brcb(resUb, inWeightsUb, static_cast(cuProcEleNum / 8), {1, 8}); + inQueue_.FreeTensor(inWeightsUb); + + outQueue_.EnQue(resUb); + resUb = outQueue_.DeQue(); + AscendC::DataCopyParams copyOutParams; + copyOutParams.blockCount = 1; + copyOutParams.blockLen = cuProcEleNum * BLOCK_CUBE * sizeof(half); + copyOutParams.srcStride = 0; + copyOutParams.dstStride = 0; + AscendC::DataCopyPad(vec0OutGm[vec0OutGmOffset], resUb, copyOutParams); + outQueue_.FreeTensor(resUb); +} + +template +__aicore__ inline int32_t QLIV2Vector::AlignS2(int32_t cuS2Len) +{ + // 限制：当前cuS2Len最大为2048，暂不考虑更长 + // 该函数目的是将cuS2Len对齐到形如 32*(4^n)*m 的形式 (m ∈ [1, 3])，方便后续sort/merge + if (cuS2Len <= ELE_NUM_128) { + return Align(cuS2Len, ELE_NUM_32); + } else if (cuS2Len <= ELE_NUM_512) { + return Align(cuS2Len, ELE_NUM_128); + } else { + return Align(cuS2Len, ELE_NUM_512); + } +} + +template +__aicore__ inline void QLIV2Vector::ProcessVec1(const QLIV2Common::RunInfo &info) +{ + int32_t cuBaseS1Idx = info.gS1Idx * s1BaseSize_; + int32_t cuBaseS2Idx = info.s2Idx * s2BaseSize_; + + // 计算基本块基地址偏移 偶数循环 -> 0 + aic_offset 奇数循环 -> 4*2048 + aic_offset + int64_t mmGmOffset = (info.loop % 2) * (s1BaseSize_ * s2BaseSize_); + + // cuS1BeginIdxPerAiv: 每个AIV的S1起始偏移 + int32_t cuS1BeginIdxPerAiv = cuBaseS1Idx; + int32_t cuS1ProcNum = + cuS1BeginIdxPerAiv + s1BaseSize_ > info.actS1Size ? info.actS1Size % s1BaseSize_ : s1BaseSize_; + // cuS1ProcNumPerAiv: 每个AIv的S1计算量 + int32_t cuS1ProcNumPerAiv = blockId_ % 2 == 0 ? CeilDiv(cuS1ProcNum, 2) : (cuS1ProcNum / 2); + cuS1BeginIdxPerAiv += (blockId_ % 2) * CeilDiv(cuS1ProcNum, 2); + // 基本块基地址偏移奇数核加一个S1地址偏移 + mmGmOffset += (blockId_ % 2) * CeilDiv(cuS1ProcNum, 2) * s2BaseSize_; + // 非首个基本块, M(S1)轴发生切换需要初始化 + if (info.loop != 0 && info.s2Idx == 0) { + // globalTopkUb_ value,index=-inf,-1 + InitSortOutBuf(globalTopkUb_, CeilDiv(s1BaseSize_, 2) * BASE_TOPK_VALUE_IDX_SIZE); + blockS2StartIdx_ = 0; + } else if (info.loop == 0) { + blockS2StartIdx_ = info.s2Idx; + } + // cuRealAcSeq: 当前基本块S1对应的AcSeq + int32_t cuRealAcSeq = info.actS2Size; + int32_t cuRealAcSeqCount = 0; + if (constInfo_.attenMaskFlag) { + // attenMask true场景 + cuRealAcSeq = info.actS2SizeOrig - info.actS1Size + cuS1BeginIdxPerAiv; + } + int32_t cuRealAcSeqIni = cuRealAcSeq; + + // LD输出S1方向偏移，保证2个Vector输出的内容连续 + uint32_t ldS1Offset = (blockId_ % 2 == 0) ? s1BaseSize_ / 2 - cuS1ProcNumPerAiv : 0; + for (int innerS1Idx = 0; innerS1Idx < cuS1ProcNumPerAiv; innerS1Idx++) { + if (constInfo_.attenMaskFlag) { + cuRealAcSeqCount += 1; + cuRealAcSeq = (cuRealAcSeqCount + cuRealAcSeqIni) / static_cast(constInfo_.cmpRatio); + } + int32_t cuS2Len = cuBaseS2Idx + s2BaseSize_ >= cuRealAcSeq ? cuRealAcSeq - cuBaseS2Idx : s2BaseSize_; + int32_t cuS1Idx = cuS1BeginIdxPerAiv + innerS1Idx; + // 当前vec1ResGm对应S1的位置 + uint64_t wsOffset = static_cast(info.saveWorkSpaceIdx) * s1BaseSize_ * BASE_TOPK_VALUE_IDX_SIZE + + static_cast(cuS1Idx - cuBaseS1Idx) * BASE_TOPK_VALUE_IDX_SIZE; + if (cuRealAcSeq > 0 && cuS2Len > 0) { + int32_t cuS2LenVecAlign = AlignS2(cuS2Len); + LocalTensor mmInUb = inQueue_.AllocTensor(); + LocalTensor kScaleUb = mmInUb[cuS2LenVecAlign]; + LocalTensor kScaleTUb = kScaleUb.template ReinterpretCast()[cuS2LenVecAlign]; + AscendC::DataCopyPadExtParams padParams{false, 0, 0, 0}; + AscendC::DataCopyPadExtParams padTParams{false, 0, 0, 0}; + AscendC::DataCopyExtParams copyInParams; + copyInParams.blockCount = 1; + copyInParams.blockLen = cuS2Len * sizeof(float); + copyInParams.srcStride = 0; + copyInParams.dstStride = 0; + copyInParams.rsv = 0; + AscendC::DataCopyPad(mmInUb, mm1ResGm[mmGmOffset + innerS1Idx * s2BaseSize_], copyInParams, padParams); + GetKeyScale(info, kScaleTUb, info.bIdx, cuBaseS2Idx, cuS2Len); + inQueue_.EnQue(mmInUb); + mmInUb = inQueue_.DeQue(); + AscendC::Cast(kScaleUb, kScaleTUb, RoundMode::CAST_NONE, cuS2Len); + PipeBarrier(); + AscendC::Mul(mmInUb, mmInUb, kScaleUb, cuS2Len); + PipeBarrier(); + LocalTensor sortBuff = tmpBuf_.Get(); + LocalTensor sortScoreUb = sortBuff; + LocalTensor sortIndiceUb = sortBuff[cuS2LenVecAlign]; + PipeBarrier(); + Duplicate(sortScoreUb.template ReinterpretCast(), QLIV2ServiceVec::NEG_INF, cuS2LenVecAlign); + PipeBarrier(); + Adds(sortScoreUb, mmInUb, 0.0f, cuS2Len); + PipeBarrier(); + inQueue_.FreeTensor(mmInUb); + LocalTensor sortIndiceUbInt = sortIndiceUb.template ReinterpretCast(); + // 无效数据索引填充为-1 + if (cuS2LenVecAlign != cuS2Len) { + Duplicate(sortIndiceUbInt, -1, cuS2LenVecAlign); + PipeBarrier(); + } + Adds(sortIndiceUbInt, globalTopkIndice_, static_cast(cuBaseS2Idx), cuS2Len); + PipeBarrier(); + LocalTensor tmpSortBuf = sortBuff[2 * cuS2LenVecAlign]; + QLIV2ServiceVec::SortAll(sortBuff, tmpSortBuf, cuS2LenVecAlign); + PipeBarrier(); + QLIV2ServiceVec::MergeSort(globalTopkUb_[innerS1Idx * BASE_TOPK_VALUE_IDX_SIZE], BASE_TOPK, sortBuff, + cuS2LenVecAlign, tmpSortBuf); + PipeBarrier(); + bool isS2End = cuBaseS2Idx + s2BaseSize_ >= cuRealAcSeq; + bool needCopyOutGm = blockS2StartIdx_ == 0 && isS2End; + // 中间结果保存 + if (needCopyOutGm && !info.isNeedLD) { + LocalTensor idxULocal = outQueue_.AllocTensor(); + ExtractIndex(idxULocal, + globalTopkUb_[innerS1Idx * BASE_TOPK_VALUE_IDX_SIZE].template ReinterpretCast(), + BASE_TOPK); + PipeBarrier(); + InitSortOutBuf(globalTopkUb_[innerS1Idx * BASE_TOPK_VALUE_IDX_SIZE], BASE_TOPK_VALUE_IDX_SIZE); + outQueue_.EnQue(idxULocal); + idxULocal = outQueue_.DeQue(); + QLIV2ServiceVec::CopyOut(indiceOutGm[info.indiceOutOffset + cuS1Idx * constInfo_.sparseCount], + idxULocal.template ReinterpretCast(), constInfo_.sparseCount); + outQueue_.FreeTensor(idxULocal); + } + // LD拷贝到当前vector对应S1的位置 + if (info.isNeedLD && info.isLastS2InnerLoop) { // 当前核存在归约任务 且是最后处理的一段 + AscendC::DataCopyExtParams copyWsParams; + copyWsParams.blockLen = BASE_TOPK_VALUE_IDX_SIZE * sizeof(float); + copyWsParams.srcStride = 0; + copyWsParams.dstStride = 0; + copyWsParams.blockCount = 1; + SetWaitFlag(HardEvent::V_MTE3); + AscendC::DataCopyPad(vec1ResGm[wsOffset], + globalTopkUb_[innerS1Idx * BASE_TOPK_VALUE_IDX_SIZE], copyWsParams); + SetWaitFlag(HardEvent::MTE3_V); + InitSortOutBuf(globalTopkUb_[innerS1Idx * BASE_TOPK_VALUE_IDX_SIZE], BASE_TOPK_VALUE_IDX_SIZE); + PipeBarrier(); + } + } else if (cuRealAcSeq <= 0) { + // 无效长度处理 + if (info.isNeedLD && info.isLastS2InnerLoop) { + PipeBarrier(); + InitSortOutBuf(globalTopkUb_[innerS1Idx * BASE_TOPK_VALUE_IDX_SIZE], BASE_TOPK_VALUE_IDX_SIZE); + SetWaitFlag(HardEvent::V_MTE3); + QLIV2ServiceVec::CopyOut(vec1ResGm[wsOffset], globalTopkUb_[innerS1Idx * BASE_TOPK_VALUE_IDX_SIZE], + BASE_TOPK_VALUE_IDX_SIZE); + SetWaitFlag(HardEvent::MTE3_V); + } else { + CleanInvalidOutput(info.indiceOutOffset + cuS1Idx * constInfo_.sparseCount); + } + } else if (cuS2Len <= 0) { + // LD拷贝到当前vector对应S1的位置 + if (info.isNeedLD && info.isLastS2InnerLoop) { // 当前核存在归约任务 且是最后处理的一段 + SetWaitFlag(HardEvent::V_MTE3); + QLIV2ServiceVec::CopyOut(vec1ResGm[wsOffset], globalTopkUb_[innerS1Idx * BASE_TOPK_VALUE_IDX_SIZE], + BASE_TOPK_VALUE_IDX_SIZE); + SetWaitFlag(HardEvent::MTE3_V); + InitSortOutBuf(globalTopkUb_[innerS1Idx * BASE_TOPK_VALUE_IDX_SIZE], BASE_TOPK_VALUE_IDX_SIZE); + PipeBarrier(); + } + } + } + + // BNSD场景无效S1 输出-1 + if (Q_LAYOUT_T == LI_LAYOUT::BSND) { + // 最后一个S1的基本块, 需要 >= info.actS1Size + bool isS1LoopEnd = (cuBaseS1Idx + s1BaseSize_) >= info.actS1Size; + int32_t invalidS1Num = constInfo_.qSeqSize - info.actS1Size; + // blockS2StartIdx_ == 0 控制S2从开始的核去做冗余清理 + if (invalidS1Num > 0 && isS1LoopEnd && blockS2StartIdx_ == 0) { + int32_t s1NumPerAiv = blockId_ % 2 == 0 ? CeilDiv(invalidS1Num, 2) : (invalidS1Num / 2); + int32_t s1OffsetPerAiv = info.actS1Size + (blockId_ % 2) * CeilDiv(invalidS1Num, 2); + for (int innerS1Idx = 0; innerS1Idx < s1NumPerAiv; innerS1Idx++) { + CleanInvalidOutput(info.indiceOutOffset + (s1OffsetPerAiv + innerS1Idx) * constInfo_.sparseCount); + } + } + + int32_t invalidS1Num2 = info.actS1Size - info.actS2SizeOrig; + if (invalidS1Num2 > 0 && isS1LoopEnd && blockS2StartIdx_ == 0 && constInfo_.attenMaskFlag) { + int32_t s1NumPerAiv = blockId_ % 2 == 0 ? CeilDiv(invalidS1Num2, 2) : (invalidS1Num2 / 2); + int32_t s1OffsetPerAiv = (blockId_ % 2) * CeilDiv(invalidS1Num2, 2); + for (int innerS1Idx = 0; innerS1Idx < s1NumPerAiv; innerS1Idx++) { + CleanInvalidOutput((info.bN2Idx * constInfo_.qSeqSize + s1OffsetPerAiv + innerS1Idx) * + constInfo_.sparseCount); + } + } + } + + if (info.isLastS2InnerLoop) { + // S2最后一个Loop后, 下一个基本块初始从0开始 + blockS2StartIdx_ = 0; + } +} + +template +__aicore__ inline void QLIV2Vector::ProcessLD() +{ + LocalTensor curValueIdxUb = ldToBeMrgBuf_.Get(); + LocalTensor tmpUb = ldTmpBuf_.Get(); + + LocalTensor outValueUb = ldOutValueBuf_.Get(); + LocalTensor outIdxUb = ldOutIdxBuf_.Get(); + + AscendC::DataCopyParams copyOutParams; + copyOutParams.blockCount = 1; + copyOutParams.blockLen = constInfo_.sparseCount * sizeof(int32_t); // bytes + copyOutParams.srcStride = 0; + copyOutParams.dstStride = 0; + + AscendC::DataCopyPadExtParams indexValuePadParams{true, 0, 0, 0}; + AscendC::DataCopyExtParams indexValueParams; + indexValueParams.blockLen = BASE_TOPK_VALUE_IDX_SIZE * sizeof(int32_t); // bytes + indexValueParams.srcStride = 3 * BASE_TOPK_VALUE_IDX_SIZE * sizeof(int32_t); + indexValueParams.dstStride = 0; + + uint32_t ldProWorkspaceNum = ldInfo_.workspaceNum; + uint32_t ldProcessLen = 4; // 4: 4块归约任务做一次Merge + uint32_t ldProcessNum = (ldProWorkspaceNum - 1) / (ldProcessLen - 1); + uint32_t ldTailLen = ldProWorkspaceNum - (ldProcessNum * (ldProcessLen - 1) + 1); // 尾块长度 + + for (uint32_t j = 0; j < ldInfo_.mNum; j++) { + SetWaitFlag(HardEvent::V_MTE2); + // 拷贝第一块归约任务 + indexValueParams.blockCount = 1; + uint64_t wsOffsetIni = static_cast(ldInfo_.workspaceIdx) * s1BaseSize_ * BASE_TOPK_VALUE_IDX_SIZE + + static_cast(ldInfo_.mStart + j) * BASE_TOPK_VALUE_IDX_SIZE; + AscendC::DataCopyPad(curValueIdxUb, vec1ResGm[wsOffsetIni], indexValueParams, indexValuePadParams); + uint64_t valueOffset = BASE_TOPK_VALUE_IDX_SIZE; + SetWaitFlag(HardEvent::MTE2_V); + + // 处理等于4的部分 + for (uint32_t i = 0; i < ldProcessNum; i++) { + // LD处理偏移 + uint64_t wsOffset = static_cast(ldInfo_.workspaceIdx) * s1BaseSize_ * BASE_TOPK_VALUE_IDX_SIZE + + static_cast(ldInfo_.mStart + j) * BASE_TOPK_VALUE_IDX_SIZE+ + static_cast(i * (ldProcessLen - 1) + 1) * + s1BaseSize_ * BASE_TOPK_VALUE_IDX_SIZE; + indexValueParams.blockCount = ldProcessLen - 1; // 拷贝4块进行merge + + SetWaitFlag(HardEvent::V_MTE2); + AscendC::DataCopyPad(curValueIdxUb[valueOffset], vec1ResGm[wsOffset], + indexValueParams, indexValuePadParams); + // merge参数 + AscendC::MrgSort4Info params; + params.elementLengths[0] = BASE_TOPK; + params.elementLengths[1] = BASE_TOPK; + params.elementLengths[2] = BASE_TOPK; + params.elementLengths[3] = BASE_TOPK; + params.ifExhaustedSuspension = true; + params.validBit = 0b1111; + params.repeatTimes = 1; + + SetWaitFlag(HardEvent::MTE2_V); + AscendC::MrgSortSrcList srcList; + srcList.src1 = curValueIdxUb[0]; + srcList.src2 = curValueIdxUb[BASE_TOPK_VALUE_IDX_SIZE]; + srcList.src3 = curValueIdxUb[2 * BASE_TOPK_VALUE_IDX_SIZE]; + srcList.src4 = curValueIdxUb[3 * BASE_TOPK_VALUE_IDX_SIZE]; + MrgSort(tmpUb, srcList, params); + PipeBarrier(); + DataCopy(curValueIdxUb, tmpUb, BASE_TOPK_VALUE_IDX_SIZE); + PipeBarrier(); + } + + // 处理不等于4的部分 + if (ldTailLen != 0) { + // 搬运尾块 + uint64_t wsOffsetTail = static_cast(ldInfo_.workspaceIdx) * s1BaseSize_ + * BASE_TOPK_VALUE_IDX_SIZE + + static_cast(ldInfo_.mStart + j) * BASE_TOPK_VALUE_IDX_SIZE + + static_cast(ldProcessNum * (ldProcessLen - 1) + 1) * s1BaseSize_ + * BASE_TOPK_VALUE_IDX_SIZE; + indexValueParams.blockCount = ldTailLen; + SetWaitFlag(HardEvent::V_MTE2); + AscendC::DataCopyPad(curValueIdxUb[valueOffset], vec1ResGm[wsOffsetTail], + indexValueParams, indexValuePadParams); + SetWaitFlag(HardEvent::MTE2_V); + AscendC::MrgSort4Info params; + params.elementLengths[0] = BASE_TOPK; + params.elementLengths[1] = BASE_TOPK; + params.elementLengths[2] = BASE_TOPK; + params.elementLengths[3] = BASE_TOPK; + params.ifExhaustedSuspension = true; + if (ldTailLen == 1) { + params.validBit = 0b0011; + } else if (ldTailLen == 2) { + params.validBit = 0b0111; + } + params.repeatTimes = 1; + + AscendC::MrgSortSrcList srcList; + srcList.src1 = curValueIdxUb[0]; + srcList.src2 = curValueIdxUb[BASE_TOPK_VALUE_IDX_SIZE]; + srcList.src3 = curValueIdxUb[2 * BASE_TOPK_VALUE_IDX_SIZE]; + srcList.src4 = curValueIdxUb[3 * BASE_TOPK_VALUE_IDX_SIZE]; + PipeBarrier(); + MrgSort(tmpUb, srcList, params); + PipeBarrier(); + DataCopy(curValueIdxUb, tmpUb, BASE_TOPK_VALUE_IDX_SIZE); + PipeBarrier(); + } + + // 搬出 + Extract(outValueUb, outIdxUb, curValueIdxUb, (BASE_TOPK / 32)); + PipeBarrier(); + InitSortOutBuf(curValueIdxUb, BASE_TOPK_VALUE_IDX_SIZE); + LocalTensor idxULocal1 = outIdxUb.template ReinterpretCast(); + SetWaitFlag(HardEvent::V_MTE3); + uint64_t outOffset = ldInfo_.indiceOutCoreOffset + + (ldInfo_.mStart + j) * constInfo_.kHeadNum * constInfo_.sparseCount; + AscendC::DataCopyPad(indiceOutGm[outOffset], idxULocal1, copyOutParams); + SetWaitFlag(HardEvent::MTE3_V); + } + SetWaitFlag(HardEvent::MTE3_V); +} + +} // namespace QLIV2Kernel +#endif // QUANT_LIGHTNING_INDEXER_V2_SERVICE_VECTOR_H \ No newline at end of file diff --git a/xllm_ops/attention/quant_lightning_indexer_v2/op_kernel/arch22/quant_lightning_indexer_v2_vector.h b/xllm_ops/attention/quant_lightning_indexer_v2/op_kernel/arch22/quant_lightning_indexer_v2_vector.h new file mode 100644 index 0000000..39a36af --- /dev/null +++ b/xllm_ops/attention/quant_lightning_indexer_v2/op_kernel/arch22/quant_lightning_indexer_v2_vector.h @@ -0,0 +1,193 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file quant_lightning_indexer_v2_vector.h + * \brief + */ +#ifndef QUANT_LIGHTNING_INDEXER_V2_VECTOR_H +#define QUANT_LIGHTNING_INDEXER_V2_VECTOR_H + +#include "kernel_operator.h" +#include "quant_lightning_indexer_v2_vector.h" + +namespace QLIV2ServiceVec { +using namespace AscendC; + +constexpr int32_t NEG_INF = 0xFF800000; +constexpr int32_t INVALID_INDEX = -1; +constexpr uint8_t VEC_REPEAT_MAX = 255; +constexpr uint8_t B32_VEC_ELM_NUM = 64; +constexpr uint8_t B32_BLOCK_ALIGN_NUM = 8; +constexpr uint8_t B32_VEC_REPEAT_STRIDE = 8; +constexpr uint64_t VEC_REPEAT_BYTES = 256; +constexpr int32_t CONST_TWO = 2; +constexpr int64_t VALUE_AND_INDEX_NUM = 2; +constexpr int64_t BLOCK_BYTES = 32; +constexpr int64_t MRG_QUE_0 = 0; +constexpr int64_t MRG_QUE_1 = 1; +constexpr int64_t MRG_QUE_2 = 2; +constexpr int64_t MRG_QUE_3 = 3; +constexpr int64_t MRG_BLOCK_2 = 2; +constexpr int64_t MRG_BLOCK_3 = 3; +constexpr int64_t MRG_BLOCK_4 = 4; + +template +__aicore__ inline void CopyOut(const GlobalTensor &dstGm, const LocalTensor &srcUb, int64_t copyCount) +{ + AscendC::DataCopyParams dataCopyOutyParams; + dataCopyOutyParams.blockCount = 1; + dataCopyOutyParams.blockLen = copyCount * sizeof(T); + dataCopyOutyParams.srcStride = 0; + dataCopyOutyParams.dstStride = 0; + AscendC::DataCopyPad(dstGm, srcUb, dataCopyOutyParams); +} + +/** + src: 传入的初始化空间 + eleNum: 需要初始化的元素个数需为64整数倍，元素将被初始化为交错排布的-inf，-1 + */ +__aicore__ inline void InitSortOutBuf(const LocalTensor &src, int64_t eleNum) +{ + uint64_t mask1[2] = {0x5555555555555555, 0}; + uint64_t mask0[2] = {0xaaaaaaaaaaaaaaaa, 0}; + int64_t repeatNum = eleNum / B32_VEC_ELM_NUM; + int64_t forLoop = repeatNum / VEC_REPEAT_MAX; + int64_t forRemain = repeatNum % VEC_REPEAT_MAX; + for (int i = 0; i < forLoop; i++) { + AscendC::Duplicate(src.template ReinterpretCast(), NEG_INF, mask1, VEC_REPEAT_MAX, 1, + B32_VEC_REPEAT_STRIDE); + AscendC::Duplicate(src.template ReinterpretCast(), INVALID_INDEX, mask0, VEC_REPEAT_MAX, 1, + B32_VEC_REPEAT_STRIDE); + } + if (forRemain > 0) { + AscendC::Duplicate(src.template ReinterpretCast()[forLoop * VEC_REPEAT_MAX * B32_VEC_ELM_NUM], NEG_INF, + mask1, forRemain, 1, B32_VEC_REPEAT_STRIDE); + AscendC::Duplicate(src.template ReinterpretCast()[forLoop * VEC_REPEAT_MAX * B32_VEC_ELM_NUM], + INVALID_INDEX, mask0, forRemain, 1, B32_VEC_REPEAT_STRIDE); + } + AscendC::PipeBarrier(); +} + +/** + src: logits和索引，前logitsNum为logits，后logitsNum为索引 + tmp: 计算使用到的临时空间，大小与src一致 + logitsNum: 排序的元素个数, 暂只支持[128,256,384,512,1024,2048] + */ +__aicore__ inline void SortAll(LocalTensor &src, LocalTensor &tmp, int64_t logitsNum) +{ + int64_t sort32Repeats = logitsNum / BLOCK_BYTES; + AscendC::Sort32(tmp, src, src[logitsNum].ReinterpretCast(), sort32Repeats); + AscendC::PipeBarrier(); + + int64_t mrgGroups = sort32Repeats; + int64_t mrgElements = BLOCK_BYTES; + int64_t i = 0; + AscendC::LocalTensor srcTensor; + AscendC::LocalTensor dstTensor; + while (true) { + if (i % CONST_TWO == 0) { + srcTensor = tmp; + dstTensor = src; + } else { + srcTensor = src; + dstTensor = tmp; + } + AscendC::MrgSort4Info params; + params.elementLengths[0] = mrgElements; + params.elementLengths[MRG_QUE_1] = mrgElements; + params.elementLengths[MRG_QUE_2] = mrgElements; + params.elementLengths[MRG_QUE_3] = mrgElements; + params.ifExhaustedSuspension = false; + params.validBit = 0b1111; + + AscendC::MrgSortSrcList srcList; + srcList.src1 = srcTensor[0]; + srcList.src2 = srcTensor[MRG_QUE_1 * VALUE_AND_INDEX_NUM * mrgElements]; + srcList.src3 = srcTensor[MRG_QUE_2 * VALUE_AND_INDEX_NUM * mrgElements]; + srcList.src4 = srcTensor[MRG_QUE_3 * VALUE_AND_INDEX_NUM * mrgElements]; + if (mrgGroups <= MRG_BLOCK_4) { + params.repeatTimes = 1; + if (mrgGroups == 1) { + break; + } else if (mrgGroups == MRG_BLOCK_2) { + params.validBit = 0b0011; + } else if (mrgGroups == MRG_BLOCK_3) { + params.validBit = 0b0111; + } else if (mrgGroups == MRG_BLOCK_4) { + params.validBit = 0b1111; + } + AscendC::MrgSort(dstTensor, srcList, params); + i += 1; + break; + } else { + params.repeatTimes = mrgGroups / MRG_BLOCK_4; + AscendC::MrgSort(dstTensor, srcList, params); + i += 1; + mrgElements = mrgElements * MRG_BLOCK_4; + mrgGroups = mrgGroups / MRG_BLOCK_4; + } + AscendC::PipeBarrier(); + } + if (i % CONST_TWO == 0) { + AscendC::DataCopy(src, tmp, logitsNum * VALUE_AND_INDEX_NUM); + AscendC::PipeBarrier(); + } +} + +/** + mrgDst: 合并进的Tensor + mrgSrc: 待合并的Tensor + tmpTensor：空间为mrgDst+mrgSrc + */ +__aicore__ inline void MergeSort(const LocalTensor &mrgDst, int32_t mrgDstNum, LocalTensor &mrgSrc, + int32_t mrgSrcNum, LocalTensor &tmpTensor) +{ + AscendC::MrgSort4Info params; + params.elementLengths[0] = mrgSrcNum; + params.elementLengths[1] = mrgDstNum; + params.ifExhaustedSuspension = false; + params.validBit = 0b0011; + params.repeatTimes = 1; + + AscendC::MrgSortSrcList srcList; + srcList.src1 = mrgSrc; + srcList.src2 = mrgDst; + + AscendC::MrgSort(tmpTensor, srcList, params); + AscendC::PipeBarrier(); + AscendC::DataCopy(mrgDst, tmpTensor, mrgDstNum * VALUE_AND_INDEX_NUM); + AscendC::PipeBarrier(); +} + +__aicore__ inline void ExtractIndex(const LocalTensor &idxULocal, const LocalTensor &sortLocal, + int64_t extractNum) +{ + AscendC::GatherMaskParams gatherMaskParams; + gatherMaskParams.repeatTimes = Ceil(extractNum * sizeof(float) * VALUE_AND_INDEX_NUM, VEC_REPEAT_BYTES); + gatherMaskParams.src0BlockStride = 1; + gatherMaskParams.src0RepeatStride = B32_VEC_REPEAT_STRIDE; + gatherMaskParams.src1RepeatStride = 0; + uint64_t rsvdCnt = 0; // 用于保存筛选后保留下来的元素个数 + uint8_t src1Pattern = 2; // 固定模式2,表示筛选出奇数索引的数 + AscendC::GatherMask(idxULocal, sortLocal, src1Pattern, false, static_cast(0), gatherMaskParams, rsvdCnt); + AscendC::PipeBarrier(); +} + +template +__aicore__ inline void SetWaitFlag(HardEvent evt) +{ + event_t eventId = static_cast(GetTPipePtr()->FetchEventID(evt)); + AscendC::SetFlag(eventId); + AscendC::WaitFlag(eventId); +} + +} // namespace QLIV2ServiceVec +#endif // QUANT_LIGHTNING_INDEXER_V2_VECTOR_H diff --git a/xllm_ops/attention/quant_lightning_indexer_v2/op_kernel/arch35/quant_lightning_indexer_v2_common_arch35.h b/xllm_ops/attention/quant_lightning_indexer_v2/op_kernel/arch35/quant_lightning_indexer_v2_common_arch35.h new file mode 100644 index 0000000..f91c1c1 --- /dev/null +++ b/xllm_ops/attention/quant_lightning_indexer_v2/op_kernel/arch35/quant_lightning_indexer_v2_common_arch35.h @@ -0,0 +1,232 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file quant_lightning_indexer_v2_common_arch35.h + * \brief + */ +#ifndef QUANT_LIGHTNING_INDEXER_V2_COMMON_H +#define QUANT_LIGHTNING_INDEXER_V2_COMMON_H +using namespace AscendC; +namespace QLIV2Common { +using FP8E4M3 = fp8_e4m3fn_t; +using FP4E2M1 = fp4x2_e2m1_t; +using FP8E8M0 = fp8_e8m0_t; + +constexpr uint32_t MX_SCALE_GROUP_SIZE = 32; // 每32个D维度元素一个scale +constexpr uint32_t FP8_TWO = 2; // 2个FP8E8M0打包成1个BF16 +constexpr uint32_t FP4_PACK_NUM = 2; // 2个FP4E2M1打包成1个字节存储 + +// 与tiling的layout保持一致 +enum class LI_LAYOUT : uint32_t { + BSND = 0, + TND = 1, + PA_BBND = 2 +}; + +template +struct QLIV2Type { + static_assert((std::is_same_v && + (std::is_same_v || std::is_same_v)) || + (std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v), + "Invalid combination of QK_T and SCORE_T"); + using rawQueryType = Q_T; + using rawKeyType = K_T; + + static constexpr bool isMxFp8 = std::is_same_v && std::is_same_v && + std::is_same_v; + static constexpr bool isMxFp4 = std::is_same_v && std::is_same_v && + std::is_same_v; + static constexpr bool isMx = isMxFp8 || isMxFp4; + + using queryType = std::conditional_t; + using keyType = std::conditional_t; + using queryKeyType = QK_T; + using scoreType = SCORE_T; + using outputType = OUT_T; + using scaleType = SCALE_T; + using weightType = WEIGHT_T; + + static constexpr bool pageAttention = PAGE_ATTENTION; + static constexpr LI_LAYOUT layout = Q_LAYOUT_T; + static constexpr LI_LAYOUT keyLayout = K_LAYOUT_T; + static constexpr bool isWeightFP16 = std::is_same_v; +}; + +struct RunInfo { + uint32_t loop; + uint32_t bN2Idx; + uint32_t bIdx; + uint32_t n2Idx = 0; + uint32_t gS1Idx; + uint32_t s2Idx; + uint32_t s2Start; + uint32_t s2LoopEnd; + uint32_t validS2Len; + uint32_t qScaleLoop; + uint32_t kScaleLoop; + + uint32_t actS1Size = 1; + uint32_t actS2Size = 1; + uint32_t actS2SizeOrig = 1; + uint32_t actMBaseSize; + uint32_t actualSingleProcessSInnerSize; + uint32_t actualSingleProcessSInnerSizeAlign; + uint32_t curCuSeqlensQ; + uint32_t curCuSeqlensK; + uint32_t curSequsedQ; + uint32_t curSequsedK; + uint64_t tensorQueryOffset; + uint64_t tensorKeyOffset; + uint64_t tensorQScaleOffset; // MX场景qScale在Cube核使用的偏移 + uint64_t tensorKeyScaleOffset; + uint64_t tensorWeightsOffset; + uint64_t indiceOutOffset; + uint64_t valueOutOffset; + uint64_t outputIdxCoreOffset; + + bool isFirstS2InnerLoop; + bool isLastS2InnerLoop; + bool isAllLoopEnd = false; + bool isValid = false; + bool isNeedLD = false; + bool isOutputIdxOffsetValid = false; + bool needTndPadding = false; + uint32_t saveWorkSpaceIdx = 0; +}; + +struct ConstInfo { + // CUBE与VEC核间同步的模式 + static constexpr uint32_t QLIV2_SYNC_MODE4 = 4; + static constexpr uint32_t AIV0_AIV1_OFFSET = 16; + static constexpr uint32_t CROSS_VC_EVENT = 0; + static constexpr uint32_t CROSS_CV_EVENT = 2; + // BUFFER的字节数 + static constexpr uint32_t BUFFER_SIZE_BYTE_32B = 32; + static constexpr uint32_t BUFFER_SIZE_BYTE_64B = 64; + static constexpr uint32_t BUFFER_SIZE_BYTE_256B = 256; + static constexpr uint32_t BUFFER_SIZE_BYTE_512B = 512; + static constexpr uint32_t BUFFER_SIZE_BYTE_1K = 1024; + static constexpr uint32_t BUFFER_SIZE_BYTE_2K = 2048; + static constexpr uint32_t BUFFER_SIZE_BYTE_4K = 4096; + static constexpr uint32_t BUFFER_SIZE_BYTE_8K = 8192; + static constexpr uint32_t BUFFER_SIZE_BYTE_16K = 16384; + static constexpr uint32_t BUFFER_SIZE_BYTE_32K = 32768; + // 无效索引 + static constexpr int INVALID_IDX = -1; + static constexpr uint16_t NEG_INF_BFLOAT = 0xFF80; + + // CUBE和VEC的核间同步EventID + uint32_t syncC1V1 = 0U; + uint32_t syncC1V0 = 2U; + uint32_t syncV1C1 = 0U; + uint32_t syncV0C1 = 1U; + + // 基本块大小 + uint32_t mBaseSize = 1ULL; + uint32_t mBaseSizeMax = 1ULL; + uint32_t s1BaseSize = 1ULL; + uint32_t s2BaseSize = 1ULL; + + uint64_t batchSize = 0ULL; + uint64_t gSize = 0ULL; + uint64_t qHeadNum = 0ULL; + uint64_t kHeadNum; + uint64_t headDim; + uint64_t sparseCount; // topK选取大小 + uint64_t kSeqSize = 0ULL; // kv最大S长度 + uint64_t qSeqSize = 1ULL; // q最大S长度 + uint32_t kCacheBlockSize = 0; // PA场景的block size + uint32_t maxBlockNumPerBatch = 0; // PA场景的最大单batch block number + LI_LAYOUT outputLayout; // 输出的格式 + bool attenMaskFlag = false; + uint32_t cmpRatio = 1; + uint32_t keyStride0 = 0; + uint32_t keyDequantScaleStride0 = 0; + int32_t maxSeqlenQ = -1; + uint32_t quantMode = 1; // quant模式，默认为1 + + uint32_t actualLenQDims = 0U; // query的actualSeqLength 的维度 + uint32_t actualLenDims = 0U; // KV 的actualSeqLength 的维度 + uint32_t cmpResiduaKLenDims = 0U; // cmpResidualK的维度 + bool isAccumSeqS1 = false; // 是否累加模式 + bool isAccumSeqS2 = false; // 是否累加模式 + bool isLDOpen = false; + bool returnValue = false; +}; + +struct LdSplitCoreInfo { + bool isLdCoreEnable = false; // 当前核是否参与规约任务 + uint32_t saveWorkSpaceIdx = 0U; // 存放LD参数的地址 + uint32_t bn2Idx = 0U; // 归约任务 + uint32_t bIdx = 0U; + uint32_t n2Idx = 0U; + uint32_t mIdx = 0U; + uint32_t workspaceIdx = 0U; // 当前AIV核上规约任务的索引 + uint32_t workspaceNum = 0U; // 当前AIV核上规约任务的S2切分数量 + uint32_t mStart = 0U; + uint32_t mNum = 0U; + uint64_t indiceOutCoreOffset = 0U; // 最终输出索引搬出Topk的初始偏移地址 +}; + +struct SplitCoreInfo { + uint32_t s2Start = 0U; // S2的起始位置 + uint32_t s2End = 0U; // S2循环index上限 + uint32_t bN2Start = 0U; + uint32_t bN2End = 0U; + uint32_t gS1Start = 0U; + uint32_t gS1End = 0U; + bool isLD = false; // 当前核是否需要进行Decode归约任务 + bool isCoreEnable = false; +}; + +template +__aicore__ inline T Align(T num, T rnd) +{ + return (((rnd) == 0) ? 0 : (((num) + (rnd)-1) / (rnd) * (rnd))); +} + +template +__aicore__ inline T1 Min(T1 a, T2 b) +{ + return (a > b) ? (b) : (a); +} + +template +__aicore__ inline T1 Max(T1 a, T2 b) +{ + return (a > b) ? (a) : (b); +} + +template +__aicore__ inline T CeilDiv(T num, T rnd) +{ + return (((rnd) == 0) ? 0 : (((num) + (rnd)-1) / (rnd))); +} +} // namespace QLIV2Common + +// bank冲突优化 +// david 256KB bank layout +// shape ( bank_depth ( banks bank_groups block)) (512 ( 2 8 32)) +// stride (banks*bank_groups*block (bank_groups*block block 1)) (512 (256 32 1)) +#define UB_BLOCK 32 // 32B +#define UB_BANK_GROUPS 8 +#define UB_BANKS 2 +#define UB_BANK_DEPTH 512 + +#define UB_BANK_GROUP_STRIDE UB_BLOCK // 32B +#define UB_BANK_STRIDE (UB_BANK_GROUPS * UB_BLOCK) // 256B +#define UB_BANK_DEPTH_STRIDE (UB_BANKS * UB_BANK_GROUPS * UB_BLOCK) // 512B + +#endif // QUANT_LIGHTNING_INDEXER_V2_COMMON_H diff --git a/xllm_ops/attention/quant_lightning_indexer_v2/op_kernel/arch35/quant_lightning_indexer_v2_kernel_arch35.h b/xllm_ops/attention/quant_lightning_indexer_v2/op_kernel/arch35/quant_lightning_indexer_v2_kernel_arch35.h new file mode 100644 index 0000000..ae4fdeb --- /dev/null +++ b/xllm_ops/attention/quant_lightning_indexer_v2/op_kernel/arch35/quant_lightning_indexer_v2_kernel_arch35.h @@ -0,0 +1,855 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file quant_lightning_indexer_v2_kernel_arch35.h + * \brief + */ + +#ifndef QUANT_LIGHTNING_INDEXER_V2_KERNEL_H +#define QUANT_LIGHTNING_INDEXER_V2_KERNEL_H + +#include "kernel_operator.h" +#include "kernel_operator_list_tensor_intf.h" +#include "kernel_tiling/kernel_tiling.h" +#include "lib/matmul_intf.h" +#include "lib/matrix/matmul/tiling.h" +#include "quant_lightning_indexer_v2_common_arch35.h" +#include "quant_lightning_indexer_v2_service_vector_arch35.h" +#include "quant_lightning_indexer_v2_service_cube_arch35.h" +#include "../quant_lightning_indexer_v2_metadata.h" + +namespace QLIV2Kernel { +using namespace QLIV2Common; +using namespace matmul; +using namespace optiling; +using namespace optiling::detail; +using AscendC::CacheMode; +using AscendC::CrossCoreSetFlag; +using AscendC::CrossCoreWaitFlag; + +// 由于S2循环前，RunInfo还没有赋值，使用TempLoopInfo临时存放B、N、S1轴相关的信息 +// 同时减少重复计算 +struct TempLoopInfo { + uint32_t bN2Idx = 0; + uint32_t bIdx = 0U; + uint32_t n2Idx = 0U; + uint32_t gS1Idx = 0U; + uint32_t gS1LoopEnd = 0U; // gS1方向循环的结束Idx + uint32_t s2LoopEnd = 0U; // S2方向循环的结束Idx + uint32_t actS1Size = 1ULL; // 当前Batch循环处理的S1轴的实际大小 + uint32_t actS2Size = 0ULL; + uint32_t actS2SizeOrig = 0ULL; // 压缩前s2 + bool curActSeqLenIsZero = false; + bool needDealActS1LessThanS1 = false; // S1的实际长度小于shape的S1长度时，是否需要清理输出 + uint32_t actMBaseSize = 0U; // m轴(gS1)方向实际大小 + uint32_t mBasicSizeTail = 0U; // gS1方向循环的尾基本块大小 + uint32_t s2BasicSizeTail = 0U; // S2方向循环的尾基本块大小 + uint32_t validS2Len = 0U; + bool isNeedLD = false; // 该基本块是否需要LD +}; + +template +class QLIV2Preload { +public: + __aicore__ inline QLIV2Preload(){}; + __aicore__ inline void Init(__gm__ uint8_t *query, __gm__ uint8_t *key, __gm__ uint8_t *weights, + __gm__ uint8_t *queryScale, __gm__ uint8_t *keyScale, __gm__ uint8_t *cuSeqlensQ, + __gm__ uint8_t *cuSeqlensK, __gm__ uint8_t *sequsedQ, __gm__ uint8_t *sequsedK, + __gm__ uint8_t *cmpResidualK, __gm__ uint8_t *blockTable, + __gm__ uint8_t *outputIdxOffset, __gm__ uint8_t *metadata, + __gm__ uint8_t *sparseIndices, __gm__ uint8_t *sparseValues, __gm__ uint8_t *workspace, + const QLIV2TilingData *__restrict tiling, TPipe *tPipe); + __aicore__ inline void Process(); + + // =================================类型定义区================================= + using Q_T = typename QLIV2T::queryType; + using K_T = typename QLIV2T::keyType; + using OUT_T = typename QLIV2T::outputType; + static constexpr bool PAGE_ATTENTION = QLIV2T::pageAttention; + static constexpr LI_LAYOUT Q_LAYOUT_T = QLIV2T::layout; + static constexpr LI_LAYOUT K_LAYOUT_T = QLIV2T::keyLayout; + using W_T = typename QLIV2T::weightType; + using SCORE_T = typename QLIV2T::scoreType; + using SCALE_T = typename QLIV2T::scaleType; + using WEIGHT_T = typename QLIV2T::weightType; + static constexpr bool IS_MX = QLIV2T::isMx; + + QLIV2Matmul matmulService; + QLIV2Vector vectorService; + + // =================================常量区================================= + static constexpr uint32_t SYNC_C1_V1_FLAG = 4; + static constexpr uint32_t SYNC_V1_C1_FLAG = 5; + + static constexpr uint32_t M_BASE_SIZE = 256; + static constexpr uint32_t M_BASE_SIZE_SMALL = 128; + static constexpr uint32_t S1_BASE_SIZE = 4; + static constexpr uint32_t S1_BASE_SIZE_SMALL = 2; + static constexpr uint32_t S2_BASE_SIZE = 128; + static constexpr uint32_t HEAD_DIM = 128; + static constexpr uint32_t K_HEAD_NUM = 1; + static constexpr uint32_t GM_ALIGN_BYTES = 512; + static constexpr uint32_t TOPK_6K = 6144; + static constexpr int64_t LD_PREFETCH_LEN = 2; + // for workspace double + static constexpr uint32_t WS_DOUBLE = 2; + +protected: + TPipe *pipe = nullptr; + + // offset + uint64_t queryCoreOffset = 0ULL; + uint64_t keyCoreOffset = 0ULL; + uint64_t qScaleCoreOffset = 0ULL; // MX场景qScale偏移 + uint64_t keyScaleCoreOffset = 0ULL; + uint64_t weightsCoreOffset = 0ULL; + uint64_t indiceOutCoreOffset = 0ULL; + uint64_t valueOutCoreOffset = 0ULL; + uint64_t outputIdxCoreOffset = 0ULL; + bool isUsedCoreEqZero = false; + bool isOutputIdxOffsetValid = false; + bool hasCuSeqlensQ = false; + bool hasCuSeqlensK = false; + bool hasSequsedQ = false; + bool hasSequsedK = false; + bool hasCmpResidualK = false; + // ================================Global Buffer区================================= + GlobalTensor queryGm; + GlobalTensor keyGm; + GlobalTensor weightsGm; + GlobalTensor mxQueryScaleGmBf16; + GlobalTensor mxKeyScaleGmBf16; + GlobalTensor qScaleGm; + GlobalTensor kScaleGm; + GlobalTensor metadataGm; + + GlobalTensor indiceOutGm; + GlobalTensor valueOutGm; + GlobalTensor blockTableGm; + GlobalTensor outputIdxOffsetGm; + GlobalTensor cuSeqlensQGm; + GlobalTensor cuSeqlensKGm; + GlobalTensor sequsedQGm; + GlobalTensor sequsedKGm; + GlobalTensor cmpResidualKGm; + + // ================================类成员变量==================================== + // aic、aiv核信息 + uint32_t tmpBlockIdx = 0U; + uint32_t aiCoreIdx = 0U; + uint32_t usedCoreNum = 0U; + + QLIV2Common::ConstInfo constInfo{}; + TempLoopInfo tempLoopInfo{}; + QLIV2Common::SplitCoreInfo splitCoreInfo{}; + QLIV2Common::LdSplitCoreInfo ldInfo{}; + + // ================================Init functions================================== + __aicore__ inline void InitTilingData(const QLIV2TilingData *__restrict tilingData); + __aicore__ inline void InitBuffers(); + __aicore__ inline void InitActualSeqLen(__gm__ uint8_t *cuSeqlensQ, __gm__ uint8_t *cuSeqlensK, + __gm__ uint8_t *sequsedQ, __gm__ uint8_t *sequsedK, + __gm__ uint8_t *cmpResidualK); + // ================================Split Core================================ + __aicore__ inline void SplitCoreByAICPU(uint32_t cubeCoreIdx, uint32_t vecCoreIdx, + GlobalTensor &metadataGm); + __aicore__ inline uint32_t GetS2BaseBlockNumOnMask(uint32_t s1gIdx, uint32_t actS1Size, uint32_t actS2SizeOrig); + // ================================Process functions================================ + __aicore__ inline void ProcessMain(); + __aicore__ inline void ProcessBaseBlock(uint32_t loop, uint64_t s2LoopIdx, QLIV2Common::RunInfo runInfo, + uint32_t qScaleLoop, uint32_t kScaleLoop); + __aicore__ inline void ProcessDecode(); + __aicore__ inline void ProcessInvalid(); + // ================================Params Calc===================================== + __aicore__ inline void CalcGS1LoopParams(uint32_t bN2Idx); + __aicore__ inline void GetBN2Idx(uint32_t bN2Idx); + __aicore__ inline uint32_t GetActualSeqLen(uint32_t bIdx, uint32_t actualLenDims, bool isAccumSeq, + GlobalTensor &cuSeqlensQGm, GlobalTensor &sequsedQGm, + uint32_t defaultSeqLen); + __aicore__ inline uint32_t GetActualSeqLenKey(uint32_t bIdx, uint32_t actualLenDims, uint32_t cmpResiduaKLenDims, + bool isAccumSeq, GlobalTensor &cuSeqlensKGm, + GlobalTensor &sequsedKGm, + GlobalTensor &cmpResidualKGm, uint32_t defaultSeqLen, + uint32_t cmpRatio); + __aicore__ inline void GetS1S2ActualSeqLen(uint32_t bIdx, uint32_t &actS1Size, uint32_t &actS2Size, + uint32_t &actS2SizeOrig); + __aicore__ inline void CalcS2LoopParams(uint32_t bN2LoopIdx, uint32_t gS1LoopIdx); + __aicore__ inline void CalcRunInfo(uint32_t loop, uint32_t s2LoopIdx, QLIV2Common::RunInfo &runInfo, + uint32_t qScaleLoop, uint32_t kScaleLoop); + __aicore__ inline void DealActSeqLenIsZero(uint32_t bIdx, uint32_t n2Idx, uint32_t s1Start); +}; + +template +__aicore__ inline void QLIV2Preload::InitTilingData(const QLIV2TilingData *__restrict tilingData) +{ + usedCoreNum = tilingData->usedCoreNum; + constInfo.batchSize = tilingData->bSize; + constInfo.qHeadNum = constInfo.gSize = tilingData->gSize; + constInfo.kSeqSize = tilingData->s2Size; + constInfo.qSeqSize = tilingData->s1Size; + constInfo.attenMaskFlag = (tilingData->sparseMode == 3); + constInfo.kCacheBlockSize = tilingData->blockSize; + constInfo.maxBlockNumPerBatch = tilingData->maxBlockNumPerBatch; + constInfo.sparseCount = tilingData->sparseCount; + constInfo.cmpRatio = tilingData->cmpRatio; + constInfo.keyStride0 = tilingData->keyStride0; + constInfo.keyDequantScaleStride0 = tilingData->keyDequantScaleStride0; + constInfo.maxSeqlenQ = tilingData->maxSeqlenQ; + constInfo.quantMode = tilingData->quantMode; + constInfo.outputLayout = Q_LAYOUT_T; // 输出和输入形状一致 + if (Q_LAYOUT_T == LI_LAYOUT::TND) { + constInfo.isAccumSeqS1 = true; + } + if (K_LAYOUT_T == LI_LAYOUT::TND) { + constInfo.isAccumSeqS2 = true; + } + + constInfo.kHeadNum = K_HEAD_NUM; + constInfo.headDim = HEAD_DIM; + if (constInfo.sparseCount > TOPK_6K) { + constInfo.s1BaseSize = S1_BASE_SIZE_SMALL; + constInfo.mBaseSizeMax = M_BASE_SIZE_SMALL; + } else { + constInfo.s1BaseSize = S1_BASE_SIZE; + constInfo.mBaseSizeMax = M_BASE_SIZE; + } + constInfo.mBaseSize = constInfo.s1BaseSize * constInfo.gSize; + constInfo.s2BaseSize = S2_BASE_SIZE; + constInfo.returnValue = tilingData->returnValue; +} + +template +__aicore__ inline void QLIV2Preload::InitBuffers() +{ + if ASCEND_IS_AIV { + vectorService.InitBuffers(pipe); + } else { + matmulService.InitBuffers(pipe); + } +} + +template +__aicore__ inline void QLIV2Preload::InitActualSeqLen(__gm__ uint8_t *cuSeqlensQ, __gm__ uint8_t *cuSeqlensK, + __gm__ uint8_t *sequsedQ, __gm__ uint8_t *sequsedK, + __gm__ uint8_t *cmpResidualK) +{ + if (cuSeqlensQ != nullptr) { + cuSeqlensQGm.SetGlobalBuffer((__gm__ uint32_t *)cuSeqlensQ); + hasCuSeqlensQ = true; + } + if (cuSeqlensK != nullptr) { + cuSeqlensKGm.SetGlobalBuffer((__gm__ uint32_t *)cuSeqlensK); + hasCuSeqlensK = true; + } + if (sequsedQ != nullptr) { + sequsedQGm.SetGlobalBuffer((__gm__ uint32_t *)sequsedQ); + hasSequsedQ = true; + } + if (sequsedK != nullptr) { + sequsedKGm.SetGlobalBuffer((__gm__ uint32_t *)sequsedK); + hasSequsedK = true; + } + if (cmpResidualK != nullptr) { + cmpResidualKGm.SetGlobalBuffer((__gm__ uint32_t *)cmpResidualK); + hasCmpResidualK = true; + } +} + +template +__aicore__ inline uint32_t QLIV2Preload::GetActualSeqLen(uint32_t bIdx, uint32_t actualLenDims, bool isAccumSeq, + GlobalTensor &cuSeqlensQGm, + GlobalTensor &sequsedQGm, + uint32_t defaultSeqLen) +{ + if (hasSequsedQ) { + return sequsedQGm.GetValue(bIdx); + } else if (hasCuSeqlensQ) { + return cuSeqlensQGm.GetValue(bIdx + 1) - cuSeqlensQGm.GetValue(bIdx); + } else { + return defaultSeqLen; + } +} + +template +__aicore__ inline uint32_t QLIV2Preload::GetActualSeqLenKey(uint32_t bIdx, uint32_t actualLenDims, + uint32_t cmpResiduaKLenDims, bool isAccumSeq, + GlobalTensor &cuSeqlensKGm, + GlobalTensor &sequsedKGm, + GlobalTensor &cmpResidualKGm, + uint32_t defaultSeqLen, uint32_t cmpRatio) +{ + uint32_t residual = hasCmpResidualK ? cmpResidualKGm.GetValue(bIdx) : 0; + if (hasSequsedK) { + return sequsedKGm.GetValue(bIdx) * cmpRatio + residual; + } else if (hasCuSeqlensK) { + return (cuSeqlensKGm.GetValue(bIdx + 1) - cuSeqlensKGm.GetValue(bIdx)) * cmpRatio + residual; + } else { + return defaultSeqLen * cmpRatio + residual; + } +} + +template +__aicore__ inline void QLIV2Preload::GetS1S2ActualSeqLen(uint32_t bIdx, uint32_t &actS1Size, + uint32_t &actS2Size, uint32_t &actS2SizeOrig) +{ + actS1Size = GetActualSeqLen(bIdx, constInfo.actualLenQDims, constInfo.isAccumSeqS1, cuSeqlensQGm, sequsedQGm, + constInfo.qSeqSize); + actS2SizeOrig = GetActualSeqLenKey(bIdx, constInfo.actualLenDims, constInfo.cmpResiduaKLenDims, + constInfo.isAccumSeqS2, cuSeqlensKGm, sequsedKGm, cmpResidualKGm, + constInfo.kSeqSize, constInfo.cmpRatio); // 压缩前的actS2Size + actS2Size = actS2SizeOrig / constInfo.cmpRatio; // 真实使用的压缩后S2长度 +} + +template +__aicore__ inline uint32_t QLIV2Preload::GetS2BaseBlockNumOnMask(uint32_t s1gIdx, uint32_t actS1Size, + uint32_t actS2SizeOrig) +{ + if (actS2SizeOrig / constInfo.cmpRatio == 0) { + return 0; + } + uint32_t s1Offset = constInfo.s1BaseSize * s1gIdx; + int32_t validS2LenBase = + static_cast(actS2SizeOrig) - static_cast(actS1Size); // 压缩前的validS2LenBase + int32_t validS2Len = + (static_cast(s1Offset) + validS2LenBase + static_cast(constInfo.s1BaseSize)) / + static_cast(constInfo.cmpRatio); + validS2Len = Min(validS2Len, static_cast(actS2SizeOrig) / constInfo.cmpRatio); + validS2Len = Max(validS2Len, 1); + tempLoopInfo.validS2Len = validS2Len; + return (validS2Len + constInfo.s2BaseSize - 1) / constInfo.s2BaseSize; +} + +template +__aicore__ inline void QLIV2Preload::SplitCoreByAICPU(uint32_t cubeCoreIdx, uint32_t vecCoreIdx, + GlobalTensor &metadataGm) +{ + uint32_t liCoreEnableIndex = GetAttrAbsIndex(cubeCoreIdx, QLI_V2_CORE_ENABLE_INDEX); + uint32_t bN2StartIndex = GetAttrAbsIndex(cubeCoreIdx, QLI_V2_BN2_START_INDEX); + uint32_t mStartIndex = GetAttrAbsIndex(cubeCoreIdx, QLI_V2_M_START_INDEX); + uint32_t s2StartIndex = GetAttrAbsIndex(cubeCoreIdx, QLI_V2_S2_START_INDEX); + uint32_t bN2EndIndex = GetAttrAbsIndex(cubeCoreIdx, QLI_V2_BN2_END_INDEX); + uint32_t mEndIndex = GetAttrAbsIndex(cubeCoreIdx, QLI_V2_M_END_INDEX); + uint32_t s2EndIndex = GetAttrAbsIndex(cubeCoreIdx, QLI_V2_S2_END_INDEX); + + uint32_t liZeroCoreEnableIndex = GetAttrAbsIndex(0, QLI_V2_CORE_ENABLE_INDEX); + if (metadataGm.GetValue(liZeroCoreEnableIndex) == 0) { + isUsedCoreEqZero = true; + } + if (metadataGm.GetValue(liCoreEnableIndex) == 0) { + splitCoreInfo.isCoreEnable = false; + return; + } else { + splitCoreInfo.isCoreEnable = true; + } + + splitCoreInfo.bN2Start = metadataGm.GetValue(bN2StartIndex); + splitCoreInfo.gS1Start = metadataGm.GetValue(mStartIndex); + splitCoreInfo.s2Start = metadataGm.GetValue(s2StartIndex); + splitCoreInfo.bN2End = metadataGm.GetValue(bN2EndIndex); + splitCoreInfo.gS1End = metadataGm.GetValue(mEndIndex); + splitCoreInfo.s2End = metadataGm.GetValue(s2EndIndex); + + if (splitCoreInfo.s2End != 0) { + // 此时只需要s2End往前退一格，bN2End和gS1End都不变 + splitCoreInfo.s2End = splitCoreInfo.s2End - 1; + } else { + if (splitCoreInfo.gS1End != 0) { + // splitCoreInfo.gS1End != 0 splitCoreInfo.s2End == 0 时，gS1End需要往前退一格, bN2End不变 + // 此时需要使用bIdx获取实际Actal S2来计算出 s2End + splitCoreInfo.gS1End = splitCoreInfo.gS1End - 1; + // 需要获取当前的Actaul S2 + uint32_t bIdx = splitCoreInfo.bN2End / constInfo.kHeadNum; + uint32_t actS1Size, actS2Size, actS2SizeOrig; + GetS1S2ActualSeqLen(bIdx, actS1Size, actS2Size, actS2SizeOrig); + // s2的切块数量 + uint32_t s2BaseNum; + if (constInfo.attenMaskFlag) { + s2BaseNum = GetS2BaseBlockNumOnMask(splitCoreInfo.gS1End, actS1Size, actS2SizeOrig); + } else { + s2BaseNum = CeilDiv(actS2Size, constInfo.s2BaseSize); + } + splitCoreInfo.s2End = s2BaseNum - 1; + } else { + // splitCoreInfo.gS1End == 0 splitCoreInfo.s2End == 0 时，bN2End需要往前退一格 + // 此时需要使用bIdx获取实际Actal S1和S2来计算出 gS1End 和 s2End + splitCoreInfo.bN2End = splitCoreInfo.bN2End - 1; + + // 需要获取当前的Actaul S1 S2 + uint32_t bIdx = splitCoreInfo.bN2End / constInfo.kHeadNum; + uint32_t actS1Size, actS2Size, actS2SizeOrig; + GetS1S2ActualSeqLen(bIdx, actS1Size, actS2Size, actS2SizeOrig); + + // s1的切块数量 + uint32_t s1GBaseNum = CeilDiv(actS1Size, constInfo.s1BaseSize); + splitCoreInfo.gS1End = s1GBaseNum - 1; + + // s2的切块数量 + uint32_t s2BaseNum; + if (constInfo.attenMaskFlag) { + s2BaseNum = GetS2BaseBlockNumOnMask(splitCoreInfo.gS1End, actS1Size, actS2SizeOrig); + } else { + s2BaseNum = CeilDiv(actS2Size, constInfo.s2BaseSize); + } + splitCoreInfo.s2End = s2BaseNum - 1; + } + } + uint32_t ldFirstWorkSpaceIndex = + GetAttrAbsIndex(cubeCoreIdx, QLI_V2_FIRST_QLD_V2_DATA_WORKSPACE_IDX_INDEX, false); // LD 第一个workspace的索引 + ldInfo.saveWorkSpaceIdx = metadataGm.GetValue(ldFirstWorkSpaceIndex); + if ASCEND_IS_AIV { + uint32_t ldCoreEnableIndex = GetAttrAbsIndex(vecCoreIdx, QLD_V2_CORE_ENABLE_INDEX, true); + ldInfo.isLdCoreEnable = metadataGm.GetValue(ldCoreEnableIndex); + + if (!ldInfo.isLdCoreEnable) { + return; + } + + // LD 参数信息 + uint32_t ldBn2IdxIndex = GetAttrAbsIndex(vecCoreIdx, QLD_V2_BN2_IDX_INDEX, true); + uint32_t ldMIdxIndex = GetAttrAbsIndex(vecCoreIdx, QLD_V2_M_IDX_INDEX, true); + uint32_t ldWorkspaceIdxIndex = GetAttrAbsIndex(vecCoreIdx, QLD_V2_WORKSPACE_IDX_INDEX, true); + uint32_t ldWorkspaceNumINDEX = GetAttrAbsIndex(vecCoreIdx, QLD_V2_WORKSPACE_NUM_INDEX, true); + uint32_t ldMstartIndex = GetAttrAbsIndex(vecCoreIdx, QLD_V2_M_START_INDEX, true); + uint32_t ldMNumIndex = GetAttrAbsIndex(vecCoreIdx, QLD_V2_M_NUM_INDEX, true); + + ldInfo.bn2Idx = metadataGm.GetValue(ldBn2IdxIndex); + ldInfo.bIdx = ldInfo.bn2Idx / constInfo.kHeadNum; + ldInfo.n2Idx = ldInfo.bn2Idx % constInfo.kHeadNum; + ldInfo.mIdx = metadataGm.GetValue(ldMIdxIndex); + ldInfo.workspaceIdx = metadataGm.GetValue(ldWorkspaceIdxIndex); + ldInfo.workspaceNum = metadataGm.GetValue(ldWorkspaceNumINDEX); + ldInfo.mStart = metadataGm.GetValue(ldMstartIndex); + ldInfo.mNum = metadataGm.GetValue(ldMNumIndex); + uint64_t actualSeqQPrefixSum = 0; + if constexpr (Q_LAYOUT_T == LI_LAYOUT::TND) { + actualSeqQPrefixSum = cuSeqlensQGm.GetValue(ldInfo.bIdx); + } else { // BSND + actualSeqQPrefixSum = (ldInfo.bIdx <= 0) ? 0 : static_cast(ldInfo.bIdx) * constInfo.qSeqSize; + } + ldInfo.indiceOutCoreOffset = actualSeqQPrefixSum * constInfo.kHeadNum * constInfo.sparseCount + + static_cast(ldInfo.n2Idx) * constInfo.sparseCount + + static_cast(ldInfo.mIdx) * constInfo.s1BaseSize * constInfo.kHeadNum * + constInfo.sparseCount; // 搬出Topk的初始偏移地址 + } +} + +template +__aicore__ inline void QLIV2Preload::DealActSeqLenIsZero(uint32_t bIdx, uint32_t n2Idx, uint32_t s1Start) +{ + if ASCEND_IS_AIV { + if (constInfo.outputLayout == LI_LAYOUT::TND) { + uint32_t tBase = cuSeqlensQGm.GetValue(bIdx); + uint32_t s1Count = cuSeqlensQGm.GetValue(bIdx + 1) - tBase; + + for (uint32_t s1Idx = s1Start; s1Idx < s1Count; s1Idx++) { + uint64_t indiceOutOffset = + (static_cast(tBase) + s1Idx) * constInfo.kHeadNum * constInfo.sparseCount + + static_cast(n2Idx) * constInfo.sparseCount; // N2轴偏移 + vectorService.CleanInvalidOutput(indiceOutOffset); + } + } else if (constInfo.outputLayout == LI_LAYOUT::BSND) { + for (uint32_t s1Idx = s1Start; s1Idx < constInfo.qSeqSize; s1Idx++) { + // B,S1,N2,K + uint64_t indiceOutOffset = + static_cast(bIdx) * constInfo.qSeqSize * constInfo.kHeadNum * constInfo.sparseCount + + static_cast(s1Idx) * constInfo.kHeadNum * constInfo.sparseCount + + static_cast(n2Idx) * constInfo.sparseCount; // N2轴偏移 + vectorService.CleanInvalidOutput(indiceOutOffset); + } + } + } +} + +template +__aicore__ inline void QLIV2Preload::Init( + __gm__ uint8_t *query, __gm__ uint8_t *key, __gm__ uint8_t *weights, __gm__ uint8_t *queryScale, + __gm__ uint8_t *keyScale, __gm__ uint8_t *cuSeqlensQ, __gm__ uint8_t *cuSeqlensK, __gm__ uint8_t *sequsedQ, + __gm__ uint8_t *sequsedK, __gm__ uint8_t *cmpResidualK, __gm__ uint8_t *blockTable, __gm__ uint8_t *outputIdxOffset, + __gm__ uint8_t *metadata, __gm__ uint8_t *sparseIndices, __gm__ uint8_t *sparseValues, __gm__ uint8_t *workspace, + const QLIV2TilingData *__restrict tiling, TPipe *tPipe) +{ + if ASCEND_IS_AIV { + tmpBlockIdx = GetBlockIdx(); // vec:0-47 + aiCoreIdx = tmpBlockIdx / 2; + } else { + tmpBlockIdx = GetBlockIdx(); // cube:0-23 + aiCoreIdx = tmpBlockIdx; + } + + InitTilingData(tiling); + InitActualSeqLen(cuSeqlensQ, cuSeqlensK, sequsedQ, sequsedK, cmpResidualK); + + // 获取分核信息 + metadataGm.SetGlobalBuffer((__gm__ uint32_t *)metadata); + SplitCoreByAICPU(aiCoreIdx, tmpBlockIdx, metadataGm); + + pipe = tPipe; + + uint64_t offset = 0; + uint32_t topkCountAlign16_ = QLIV2Common::Align(constInfo.sparseCount, (uint64_t)16); // topkCount对齐到16 + // vec 把整个s2的score存储在GM，大小为s1BaseSize * 16K * 4 + GlobalTensor scoreGm; // 存放vec核写出的score + uint64_t singleCoreScoreSize = constInfo.s1BaseSize * + QLIV2Common::Align((uint64_t)constInfo.kSeqSize, (uint64_t)constInfo.s2BaseSize) * + sizeof(SCORE_T); + scoreGm.SetGlobalBuffer((__gm__ SCORE_T *)(workspace + aiCoreIdx * singleCoreScoreSize)); + offset += GetBlockNum() * singleCoreScoreSize; + // vec 存储需要LD的s1对应的s2的score与index， + // 大小为s1BaseSize * sparseCount * 2，一个核内最多有两个s1BaseSize需要LD + GlobalTensor ldScoreGm; // 存放进行LD的s2 score + ldScoreGm.SetGlobalBuffer((__gm__ SCORE_T *)(workspace + offset)); + offset += static_cast(GetBlockNum()) * constInfo.s1BaseSize * topkCountAlign16_ * 2 * sizeof(SCORE_T); + GlobalTensor ldIndexGm; // 存放进行LD的s2 Index + ldIndexGm.SetGlobalBuffer((__gm__ int32_t *)(workspace + offset)); + offset += static_cast(GetBlockNum()) * constInfo.s1BaseSize * topkCountAlign16_ * 2 * sizeof(int32_t); + + if ASCEND_IS_AIV { + indiceOutGm.SetGlobalBuffer((__gm__ int32_t *)sparseIndices); + weightsGm.SetGlobalBuffer((__gm__ WEIGHT_T *)weights); + blockTableGm.SetGlobalBuffer((__gm__ int32_t *)blockTable); + valueOutGm.SetGlobalBuffer((__gm__ bfloat16_t *)sparseValues); + if (outputIdxOffset != nullptr) { + isOutputIdxOffsetValid = true; + outputIdxOffsetGm.SetGlobalBuffer((__gm__ int32_t *)outputIdxOffset); + } + if constexpr (IS_MX) { + vectorService.InitVecInputTensor(weightsGm, indiceOutGm, blockTableGm, valueOutGm, outputIdxOffsetGm); + } else { + GlobalTensor qScaleGmForVec; + GlobalTensor kScaleGmForVec; + qScaleGmForVec.SetGlobalBuffer((__gm__ SCALE_T *)queryScale); + kScaleGmForVec.SetGlobalBuffer((__gm__ SCALE_T *)keyScale); + vectorService.InitVecInputTensor(weightsGm, indiceOutGm, blockTableGm, valueOutGm, outputIdxOffsetGm, + qScaleGmForVec, kScaleGmForVec); + } + vectorService.InitVecWorkspaceTensor(scoreGm, ldScoreGm, ldIndexGm); + vectorService.InitParams(constInfo, ldInfo, tiling); + } else { + matmulService.InitParams(constInfo); + queryGm.SetGlobalBuffer((__gm__ Q_T *)query); + if constexpr (PAGE_ATTENTION) { + blockTableGm.SetGlobalBuffer((__gm__ int32_t *)blockTable); + } + keyGm.SetGlobalBuffer((__gm__ K_T *)key); + if constexpr (IS_MX) { + mxQueryScaleGmBf16.SetGlobalBuffer((__gm__ bfloat16_t *)queryScale); + mxKeyScaleGmBf16.SetGlobalBuffer((__gm__ bfloat16_t *)keyScale); + matmulService.InitMm1GlobalTensor(blockTableGm, keyGm, queryGm, mxKeyScaleGmBf16, mxQueryScaleGmBf16); + } else { + matmulService.InitMm1GlobalTensor(blockTableGm, keyGm, queryGm); + } + } + InitBuffers(); +} + +template +__aicore__ inline void QLIV2Preload::GetBN2Idx(uint32_t bN2Idx) +{ + tempLoopInfo.bN2Idx = bN2Idx; + tempLoopInfo.bIdx = bN2Idx / constInfo.kHeadNum; + tempLoopInfo.n2Idx = bN2Idx % constInfo.kHeadNum; +} + +template +__aicore__ inline void QLIV2Preload::CalcS2LoopParams(uint32_t bN2LoopIdx, uint32_t gS1LoopIdx) +{ + tempLoopInfo.gS1Idx = gS1LoopIdx; + tempLoopInfo.actMBaseSize = constInfo.mBaseSize; + uint32_t remainedGS1Size = tempLoopInfo.actS1Size * constInfo.gSize - tempLoopInfo.gS1Idx * constInfo.mBaseSize; + if (remainedGS1Size <= constInfo.mBaseSize && remainedGS1Size > 0) { + tempLoopInfo.actMBaseSize = tempLoopInfo.mBasicSizeTail; + } + + bool isEnd = (bN2LoopIdx == splitCoreInfo.bN2End) && (gS1LoopIdx == splitCoreInfo.gS1End); + uint32_t s2BlockNum; + if (constInfo.attenMaskFlag) { + s2BlockNum = GetS2BaseBlockNumOnMask(gS1LoopIdx, tempLoopInfo.actS1Size, tempLoopInfo.actS2SizeOrig); + } else { + tempLoopInfo.validS2Len = tempLoopInfo.actS2Size; + s2BlockNum = (tempLoopInfo.actS2Size + constInfo.s2BaseSize - 1) / constInfo.s2BaseSize; + } + tempLoopInfo.s2LoopEnd = isEnd ? splitCoreInfo.s2End : s2BlockNum - 1; + if (splitCoreInfo.s2Start > 0 || tempLoopInfo.s2LoopEnd < s2BlockNum - 1) { + tempLoopInfo.isNeedLD = true; + } else { + tempLoopInfo.isNeedLD = false; + } +} + +template +__aicore__ inline void QLIV2Preload::CalcGS1LoopParams(uint32_t bN2LoopIdx) +{ + GetBN2Idx(bN2LoopIdx); + GetS1S2ActualSeqLen(tempLoopInfo.bIdx, tempLoopInfo.actS1Size, tempLoopInfo.actS2Size, tempLoopInfo.actS2SizeOrig); + if ((tempLoopInfo.actS2Size == 0) || (tempLoopInfo.actS1Size == 0)) { + tempLoopInfo.curActSeqLenIsZero = true; + return; + } + tempLoopInfo.curActSeqLenIsZero = false; + tempLoopInfo.s2BasicSizeTail = tempLoopInfo.actS2Size % constInfo.s2BaseSize; + tempLoopInfo.s2BasicSizeTail = + (tempLoopInfo.s2BasicSizeTail == 0) ? constInfo.s2BaseSize : tempLoopInfo.s2BasicSizeTail; + tempLoopInfo.mBasicSizeTail = (tempLoopInfo.actS1Size * constInfo.gSize) % constInfo.mBaseSize; + tempLoopInfo.mBasicSizeTail = + (tempLoopInfo.mBasicSizeTail == 0) ? constInfo.mBaseSize : tempLoopInfo.mBasicSizeTail; + + uint32_t gS1SplitNum = (tempLoopInfo.actS1Size * constInfo.gSize + constInfo.mBaseSize - 1) / constInfo.mBaseSize; + tempLoopInfo.gS1LoopEnd = (bN2LoopIdx == splitCoreInfo.bN2End) ? splitCoreInfo.gS1End : gS1SplitNum - 1; + if constexpr (Q_LAYOUT_T == LI_LAYOUT::BSND) { + if (tempLoopInfo.gS1LoopEnd == gS1SplitNum - 1 && constInfo.qSeqSize > tempLoopInfo.actS1Size) { + tempLoopInfo.needDealActS1LessThanS1 = true; + } + } +} + +template +__aicore__ inline void QLIV2Preload::CalcRunInfo(uint32_t loop, uint32_t s2LoopIdx, + QLIV2Common::RunInfo &runInfo, uint32_t qScaleLoop, + uint32_t kScaleLoop) +{ + runInfo.loop = loop; + runInfo.bIdx = tempLoopInfo.bIdx; + runInfo.gS1Idx = tempLoopInfo.gS1Idx; + runInfo.s2Idx = s2LoopIdx; + runInfo.bN2Idx = tempLoopInfo.bN2Idx; + runInfo.isValid = s2LoopIdx <= tempLoopInfo.s2LoopEnd; + runInfo.validS2Len = tempLoopInfo.validS2Len; + runInfo.qScaleLoop = qScaleLoop; + runInfo.kScaleLoop = kScaleLoop; + runInfo.isNeedLD = tempLoopInfo.isNeedLD; + if (runInfo.isNeedLD && s2LoopIdx == tempLoopInfo.s2LoopEnd) { + runInfo.saveWorkSpaceIdx = ldInfo.saveWorkSpaceIdx; + ldInfo.saveWorkSpaceIdx++; + } + + if (!runInfo.isValid) { + return; // 需要验证， v1 时候需要runInfo + } + + runInfo.actS1Size = tempLoopInfo.actS1Size; + runInfo.actS2Size = tempLoopInfo.actS2Size; + runInfo.actS2SizeOrig = tempLoopInfo.actS2SizeOrig; + // 计算实际基本块size + runInfo.actMBaseSize = tempLoopInfo.actMBaseSize; + runInfo.actualSingleProcessSInnerSize = constInfo.s2BaseSize; + uint32_t s2SplitNum = (tempLoopInfo.actS2Size + constInfo.s2BaseSize - 1) / constInfo.s2BaseSize; + if (runInfo.s2Idx == s2SplitNum - 1) { + runInfo.actualSingleProcessSInnerSize = tempLoopInfo.s2BasicSizeTail; + } + runInfo.actualSingleProcessSInnerSizeAlign = QLIV2Common::Align((uint32_t)runInfo.actualSingleProcessSInnerSize, + QLIV2Common::ConstInfo::BUFFER_SIZE_BYTE_32B); + + runInfo.isFirstS2InnerLoop = s2LoopIdx == splitCoreInfo.s2Start; + runInfo.isLastS2InnerLoop = s2LoopIdx == tempLoopInfo.s2LoopEnd; + runInfo.isAllLoopEnd = (runInfo.bN2Idx == splitCoreInfo.bN2End) && (runInfo.gS1Idx == splitCoreInfo.gS1End) && + (runInfo.s2Idx == splitCoreInfo.s2End); + runInfo.isOutputIdxOffsetValid = isOutputIdxOffsetValid; + uint64_t qkHeadDim = constInfo.headDim; + if constexpr (QLIV2T::isMxFp4) { + qkHeadDim = constInfo.headDim / FP4_PACK_NUM; + } + if (runInfo.isFirstS2InnerLoop) { + uint64_t actualSeqQPrefixSum; + if constexpr (Q_LAYOUT_T == LI_LAYOUT::TND) { + actualSeqQPrefixSum = cuSeqlensQGm.GetValue(runInfo.bIdx); + if (hasSequsedQ) { + uint32_t curSequsedQ = sequsedQGm.GetValue(runInfo.bIdx); + uint32_t nextPrefixSum = cuSeqlensQGm.GetValue(runInfo.bIdx + 1); + uint32_t curCuLensQ = nextPrefixSum - actualSeqQPrefixSum; + if (curSequsedQ < curCuLensQ) { + runInfo.needTndPadding = true; + runInfo.curCuSeqlensQ = curCuLensQ; + runInfo.curSequsedQ = curSequsedQ; + } + } + } else { // BSND + actualSeqQPrefixSum = (runInfo.bIdx <= 0) ? 0 : static_cast(runInfo.bIdx) * constInfo.qSeqSize; + } + uint64_t tndBIdxOffset = actualSeqQPrefixSum * constInfo.qHeadNum * qkHeadDim; + // B,S1,N1(N2,G),D + queryCoreOffset = tndBIdxOffset + runInfo.gS1Idx * constInfo.mBaseSize * qkHeadDim; + // B,S1,N1(N2,G)/T,N1(N2,G) + weightsCoreOffset = actualSeqQPrefixSum * constInfo.qHeadNum + runInfo.n2Idx * constInfo.gSize; + // B,S1,N2,k/T,N2,k + indiceOutCoreOffset = + actualSeqQPrefixSum * constInfo.kHeadNum * constInfo.sparseCount + runInfo.n2Idx * constInfo.sparseCount; + valueOutCoreOffset = + actualSeqQPrefixSum * constInfo.kHeadNum * constInfo.sparseCount + runInfo.n2Idx * constInfo.sparseCount; + outputIdxCoreOffset = + (actualSeqQPrefixSum + runInfo.gS1Idx * constInfo.s1BaseSize) * constInfo.kHeadNum + runInfo.n2Idx; + if constexpr (IS_MX) { + // MX: qScale offset, shape [B, S1, N1, D/64, 2] + uint64_t qScalePrefixSum = + actualSeqQPrefixSum * constInfo.qHeadNum * (constInfo.headDim / MX_SCALE_GROUP_SIZE); + uint64_t qScaleS1Offset = + static_cast(runInfo.gS1Idx) * constInfo.mBaseSize * (constInfo.headDim / MX_SCALE_GROUP_SIZE); + qScaleCoreOffset = qScalePrefixSum + qScaleS1Offset; + } + } + uint64_t actualSeqKPrefixSum; + if constexpr (K_LAYOUT_T == LI_LAYOUT::TND) { // T N2 D, cu_seqlens_k + actualSeqKPrefixSum = cuSeqlensKGm.GetValue(runInfo.bIdx); + } else { + actualSeqKPrefixSum = (runInfo.bIdx <= 0) ? 0 : runInfo.bIdx * constInfo.kSeqSize; + } + uint64_t tndBIdxOffsetForK = actualSeqKPrefixSum * constInfo.kHeadNum * qkHeadDim; + keyCoreOffset = tndBIdxOffsetForK + runInfo.s2Idx * constInfo.s2BaseSize * constInfo.kHeadNum * qkHeadDim; + uint64_t keyScaleS2Offset = actualSeqKPrefixSum + static_cast(runInfo.s2Idx) * constInfo.s2BaseSize; + if constexpr (IS_MX) { + // MX: kScale offset, shape [B, S2, N2, D/64, 2] + keyScaleCoreOffset = keyScaleS2Offset * constInfo.kHeadNum * (constInfo.headDim / MX_SCALE_GROUP_SIZE); + } else { + keyScaleCoreOffset = keyScaleS2Offset * constInfo.kHeadNum; + } + runInfo.tensorQueryOffset = queryCoreOffset; + runInfo.tensorKeyOffset = keyCoreOffset; + runInfo.tensorQScaleOffset = qScaleCoreOffset; + runInfo.tensorKeyScaleOffset = keyScaleCoreOffset; + runInfo.tensorWeightsOffset = weightsCoreOffset; + runInfo.indiceOutOffset = indiceOutCoreOffset; + runInfo.valueOutOffset = valueOutCoreOffset; + runInfo.outputIdxCoreOffset = outputIdxCoreOffset; +} + +template +__aicore__ inline void QLIV2Preload::Process() +{ + if (isUsedCoreEqZero) { + // 没有计算任务，直接清理输出 + ProcessInvalid(); + return; + } + + ProcessMain(); + + ProcessDecode(); +} + +template +__aicore__ inline void QLIV2Preload::ProcessInvalid() +{ + if ASCEND_IS_AIV { + uint32_t aivCoreNum = GetBlockNum() * 2; // 2 means c:v = 1:2 + uint64_t totalOutputSize = static_cast(constInfo.batchSize) * constInfo.qSeqSize * + constInfo.kHeadNum * constInfo.sparseCount; + uint64_t singleCoreSize = + QLIV2Common::Align((totalOutputSize + aivCoreNum - 1) / aivCoreNum, GM_ALIGN_BYTES / sizeof(OUT_T)); + uint64_t baseSize = tmpBlockIdx * singleCoreSize; + if (baseSize < totalOutputSize) { + uint64_t dealSize = + (baseSize + singleCoreSize <= totalOutputSize) ? singleCoreSize : totalOutputSize - baseSize; + GlobalTensor output = indiceOutGm[baseSize]; + AscendC::InitGlobalMemory(output, dealSize, constInfo.INVALID_IDX); + if (constInfo.returnValue) { + event_t eventIDMTE3ToV = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE3_V)); + SetFlag(eventIDMTE3ToV); + WaitFlag(eventIDMTE3ToV); + + GlobalTensor valueOutGmTmp; + valueOutGmTmp.SetGlobalBuffer((__gm__ uint16_t *)valueOutGm.GetPhyAddr()); + GlobalTensor valueOut = valueOutGmTmp[baseSize]; + + AscendC::InitGlobalMemory(valueOut, dealSize, constInfo.NEG_INF_BFLOAT); + } + } + } +} + +template +__aicore__ inline void QLIV2Preload::ProcessMain() +{ + if (!splitCoreInfo.isCoreEnable) { + return; + } + + if ASCEND_IS_AIV { + vectorService.AllocEventID(); + CrossCoreSetFlag(QLIV2Common::ConstInfo::CROSS_VC_EVENT + 0); + CrossCoreSetFlag(QLIV2Common::ConstInfo::CROSS_VC_EVENT + 1); + } else { + matmulService.AllocEventID(); + } + + QLIV2Common::RunInfo runInfo; + uint32_t gloop = 0; + uint32_t qScaleLoop = 0; + for (uint32_t bN2LoopIdx = splitCoreInfo.bN2Start; bN2LoopIdx <= splitCoreInfo.bN2End; bN2LoopIdx++) { + CalcGS1LoopParams(bN2LoopIdx); + if (tempLoopInfo.curActSeqLenIsZero) { + DealActSeqLenIsZero(tempLoopInfo.bIdx, tempLoopInfo.n2Idx, 0U); + continue; + } + for (uint32_t gS1LoopIdx = splitCoreInfo.gS1Start; gS1LoopIdx <= tempLoopInfo.gS1LoopEnd; gS1LoopIdx++) { + CalcS2LoopParams(bN2LoopIdx, gS1LoopIdx); + uint32_t kScaleLoop = 0; + runInfo.s2Start = splitCoreInfo.s2Start; + runInfo.s2LoopEnd = tempLoopInfo.s2LoopEnd; + for (int s2LoopIdx = splitCoreInfo.s2Start; s2LoopIdx <= tempLoopInfo.s2LoopEnd; s2LoopIdx++) { + if ((s2LoopIdx - splitCoreInfo.s2Start) % 16 == 0) { + ++kScaleLoop; + } + ProcessBaseBlock(gloop, s2LoopIdx, runInfo, qScaleLoop, kScaleLoop); + ++gloop; + } + ++qScaleLoop; + splitCoreInfo.s2Start = 0; + } + if (tempLoopInfo.needDealActS1LessThanS1) { + DealActSeqLenIsZero(tempLoopInfo.bIdx, tempLoopInfo.n2Idx, tempLoopInfo.actS1Size); + } + splitCoreInfo.gS1Start = 0; + } + + if ASCEND_IS_AIV { + vectorService.FreeEventID(); + } else { + matmulService.FreeEventID(); + CrossCoreWaitFlag(QLIV2Common::ConstInfo::CROSS_VC_EVENT + + 0); + CrossCoreWaitFlag(QLIV2Common::ConstInfo::CROSS_VC_EVENT + + 1); + } +} + +template +__aicore__ inline void QLIV2Preload::ProcessBaseBlock(uint32_t loop, uint64_t s2LoopIdx, + QLIV2Common::RunInfo runInfo, uint32_t qScaleLoop, + uint32_t kScaleLoop) +{ + CalcRunInfo(loop, s2LoopIdx, runInfo, qScaleLoop, kScaleLoop); + if ASCEND_IS_AIC { + matmulService.ComputeMm1(runInfo); + } else { + if (runInfo.needTndPadding) { + vectorService.DoTndPadding(runInfo); + } + vectorService.ProcessVec1(runInfo); + if (runInfo.isLastS2InnerLoop) { // 本核s2last + vectorService.ProcessTopK(runInfo); + } + } +} + +template +__aicore__ inline void QLIV2Preload::ProcessDecode() +{ + if ASCEND_IS_AIV { + vectorService.InitLDBuffers(pipe, ldInfo); + ICachePreLoad(LD_PREFETCH_LEN); + SyncAll(); + if (ldInfo.isLdCoreEnable) { + vectorService.ProcessLD(); + } + } +} + +} // namespace QLIV2Kernel +#endif // QUANT_LIGHTNING_INDEXER_V2_KERNEL_H diff --git a/xllm_ops/attention/quant_lightning_indexer_v2/op_kernel/arch35/quant_lightning_indexer_v2_service_cube_arch35.h b/xllm_ops/attention/quant_lightning_indexer_v2/op_kernel/arch35/quant_lightning_indexer_v2_service_cube_arch35.h new file mode 100644 index 0000000..3478151 --- /dev/null +++ b/xllm_ops/attention/quant_lightning_indexer_v2/op_kernel/arch35/quant_lightning_indexer_v2_service_cube_arch35.h @@ -0,0 +1,757 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file quant_lightning_indexer_v2_service_cube_arch35.h + * \brief use 5 buffer for matmul l1, better pipeline + */ +#ifndef QUANT_LIGHTNING_INDEXER_V2_SERVICE_CUBE_H +#define QUANT_LIGHTNING_INDEXER_V2_SERVICE_CUBE_H + +#include "kernel_operator.h" +#include "kernel_operator_list_tensor_intf.h" +#include "kernel_tiling/kernel_tiling.h" +#include "lib/matmul_intf.h" +#include "lib/matrix/matmul/tiling.h" +#include "quant_lightning_indexer_v2_common_arch35.h" + +namespace QLIV2Kernel { +using namespace QLIV2Common; +template +class QLIV2Matmul { +public: + using Q_T = typename QLIV2T::queryType; + using K_T = typename QLIV2T::keyType; + using QK_T = typename QLIV2T::queryKeyType; + using SCALE_T = typename QLIV2T::scaleType; + // MXFP4的GM/L1按uint8_t存packed字节，Load2DMX源操作数语义使用原始E2M1类型。 + using MX_DATA_SRC_T = std::conditional_t; + // Load2DMX要求：fp8_e4m3fn_t源操作数的dst使用mx_fp8_e4m3_t；其他类型dst与src一致。 + using MX_DATA_DST_T = std::conditional_t; + using L0_Q_T = std::conditional_t; + using L0_K_T = std::conditional_t; + using CL0_T = std::conditional_t, int32_t, float>; + + __aicore__ inline QLIV2Matmul(){}; + __aicore__ inline void InitBuffers(TPipe *pipe); + __aicore__ inline void InitMm1GlobalTensor(const GlobalTensor &blkTableGm, const GlobalTensor &keyGm, + const GlobalTensor &queryGm, + const GlobalTensor &keyScaleGmBf16 = {}, + const GlobalTensor &queryScaleGmBf16 = {}); + __aicore__ inline void InitParams(const ConstInfo &constInfo); + __aicore__ inline void AllocEventID(); + __aicore__ inline void FreeEventID(); + __aicore__ inline void ComputeMm1(const QLIV2Common::RunInfo &runInfo); + + static constexpr IsResetLoad3dConfig LOAD3DV2_CONFIG = {true, true}; // isSetFMatrix isSetPadding; + static constexpr uint64_t KEY_BUF_NUM = 3; + static constexpr uint64_t QUERY_BUF_NUM = 2; + static constexpr uint64_t L0_BUF_NUM = 2; + + static constexpr uint32_t KEY_MTE1_MTE2_EVENT = EVENT_ID2; + static constexpr uint32_t QUERY_MTE1_MTE2_EVENT = EVENT_ID5; // KEY_MTE1_MTE2_EVENT + KEY_BUF_NUM; + static constexpr uint32_t M_MTE1_EVENT = EVENT_ID3; + + static constexpr uint32_t MTE2_MTE1_EVENT = EVENT_ID2; + static constexpr uint32_t MTE1_M_EVENT = EVENT_ID2; + static constexpr uint32_t FIX_M_EVENT = EVENT_ID2; + static constexpr uint32_t M_FIX_EVENT = EVENT_ID3; + + static constexpr uint64_t M_BASIC_BLOCK = 256; + static constexpr uint64_t M_BASIC_BLOCK_SMALL = 128; + static constexpr uint64_t D_BASIC_BLOCK = 128; + + static constexpr uint64_t M_BASIC_BLOCK_L0 = 256; + static constexpr uint64_t D_BASIC_BLOCK_L0 = 128; + static constexpr uint64_t S2_BASIC_BLOCK_L0 = 128; + + static constexpr uint64_t FP8_BLOCK_CUBE = 32; + static constexpr uint64_t BLOCK_CUBE = 16; + static constexpr FixpipeConfig QLIV2_CFG_ROW_MAJOR_UB = { + CO2Layout::ROW_MAJOR, + // ROW_MAJOR: 使能NZ2ND，输出ND格式; true: 用户指定目的地址是否是UB + true}; + + static constexpr uint64_t QUERY_BUFFER_OFFSET = M_BASIC_BLOCK * D_BASIC_BLOCK; + static constexpr uint64_t L0AB_BUFFER_OFFSET = M_BASIC_BLOCK_L0 * D_BASIC_BLOCK_L0; + static constexpr uint64_t L0C_BUFFER_OFFSET = M_BASIC_BLOCK_L0 * S2_BASIC_BLOCK_L0; + // Key L0 ping-pong is sized by S2, not by the (potentially larger) M block. + static constexpr uint64_t KEY_L0_BUFFER_OFFSET = S2_BASIC_BLOCK_L0 * D_BASIC_BLOCK_L0; + // MX qScale L1乒乓缓冲区步长：M基本块 * 固定D基本块 / 32 + static constexpr uint64_t QUERY_SCALE_BUFFER_OFFSET = M_BASIC_BLOCK * D_BASIC_BLOCK / MX_SCALE_GROUP_SIZE; + // MXFP4打包路径每次32字节搬入/scale对应64个逻辑D元素 + static constexpr uint64_t MX_LOAD_SCALE_ALIGN = MX_SCALE_GROUP_SIZE * FP4_PACK_NUM; + +protected: + __aicore__ inline void Fixp(uint64_t s1gGmOffset, uint64_t s2GmOffset, uint64_t s1gL0RealSize, + uint64_t s2L0RealSize, uint64_t s1gL1SizeAlign2G, const QLIV2Common::RunInfo &runInfo); + __aicore__ inline void ComputeL0c(uint64_t s1gL0RealSize, uint64_t s2L0RealSize, + const QLIV2Common::RunInfo &runInfo); + __aicore__ inline void LoadKeyToL0b(uint64_t s2L0Offset, uint64_t s2L1RealSize, uint64_t s2L0RealSize, + const QLIV2Common::RunInfo &runInfo); + __aicore__ inline void LoadQueryToL0a(uint64_t s1gL1Offset, uint64_t s1gL1RealSize, uint64_t s1gL0RealSize, + const QLIV2Common::RunInfo &runInfo); + __aicore__ inline void QueryNd2Nz(uint64_t s1gL1RealSize, uint64_t s1gL1Offset, + const QLIV2Common::RunInfo &runInfo); + __aicore__ inline void KeyNd2Nz(uint64_t s2L1RealSize, uint64_t s2GmOffset, const QLIV2Common::RunInfo &runInfo); + __aicore__ inline void KeyNd2NzForPA(uint64_t s2L1RealSize, uint64_t s2GmOffset, + const QLIV2Common::RunInfo &runInfo); + __aicore__ inline void LoadQScaleToL1(uint64_t s1gL1RealSize, uint64_t s1gGmOffset, + const QLIV2Common::RunInfo &runInfo); + __aicore__ inline void LoadKScaleToL1(uint64_t s2L1RealSize, uint64_t s2GmOffset, + const QLIV2Common::RunInfo &runInfo); + GlobalTensor blkTableGm_; + GlobalTensor keyGm_; + GlobalTensor queryGm_; + GlobalTensor mxKeyScaleGmBf16_; + GlobalTensor mxQueryScaleGmBf16_; + + TBuf bufQL1_; + LocalTensor queryL1_; + TBuf bufKeyL1_; + LocalTensor keyL1_; + TBuf bufQScaleL1_; + LocalTensor queryScaleL1_; + TBuf bufKeyScaleL1_; + LocalTensor keyScaleL1_; + + TBuf bufQL0_; + LocalTensor queryL0_; + TBuf bufKeyL0_; + LocalTensor keyL0_; + + TBuf bufL0C_; + LocalTensor cL0_; + + TBuf bufUB_; + LocalTensor mm1ResUB_; + + uint64_t keyL1BufIdx_ = 0; + uint64_t queryL1Mte2BufIdx_ = 0; + uint64_t queryL1Mte1BufIdx_ = 0; + uint64_t l0BufIdx_ = 0; + + bool isKeyCacheValid_ = false; // L1中是否有可复用的数据 + uint64_t keyGmStart_ = 0; // L1数据对应的GM S2偏移 + uint64_t keyLoadedSize_ = 0; // L1中实际加载的S2元素数量 + uint64_t s2BasicBlock_ = 128; + uint64_t qkHeadDim_ = 128; // Q/K单行GM搬入宽度；MXFP4为打包后的headDim/2，其他场景为headDim + uint64_t scaleHeadDim_ = 4; // MX scale单行元素数，即headDim/32，MXFP8/MXFP4共用 + uint64_t keyBufferOffset_ = 16384; // Key L1乒乓缓冲区步长，s2BasicBlock_ * D_BASIC_BLOCK + uint64_t keyScaleBufferOffset_ = 512; // Key scale L1乒乓缓冲区步长，s2BasicBlock_ * D_BASIC_BLOCK / 32 + + ConstInfo constInfo_; + uint64_t queryBufferOffset_ = 0; + uint64_t l0abBufferOffset_ = 0; + uint64_t l0cBufferOffset_ = 0; + uint64_t queryScaleBufferOffset = 0; + +private: + static constexpr bool PAGE_ATTENTION = QLIV2T::pageAttention; + static constexpr bool IS_MX = QLIV2T::isMx; + static constexpr bool IS_MXFP4 = QLIV2T::isMxFp4; +}; + +template +__aicore__ inline void QLIV2Matmul::InitParams(const ConstInfo &constInfo) +{ + constInfo_ = constInfo; + s2BasicBlock_ = (constInfo_.maxSeqlenQ <= 4 && constInfo_.maxSeqlenQ >= 0) ? 256 : 128; + qkHeadDim_ = constInfo_.headDim; + if constexpr (IS_MXFP4) { + qkHeadDim_ = constInfo_.headDim / FP4_PACK_NUM; + } + scaleHeadDim_ = constInfo_.headDim / MX_SCALE_GROUP_SIZE; + keyBufferOffset_ = s2BasicBlock_ * D_BASIC_BLOCK; + keyScaleBufferOffset_ = s2BasicBlock_ * D_BASIC_BLOCK / MX_SCALE_GROUP_SIZE; + queryBufferOffset_ = constInfo_.mBaseSizeMax * D_BASIC_BLOCK; + l0abBufferOffset_ = constInfo_.mBaseSizeMax * D_BASIC_BLOCK_L0; + l0cBufferOffset_ = constInfo_.mBaseSizeMax * S2_BASIC_BLOCK_L0; + queryScaleBufferOffset = constInfo_.mBaseSizeMax * D_BASIC_BLOCK / MX_SCALE_GROUP_SIZE; +} + +template +__aicore__ inline void QLIV2Matmul::InitBuffers(TPipe *pipe) +{ + pipe->InitBuffer(bufUB_, 2 * CeilDiv(constInfo_.mBaseSizeMax, 2) * constInfo_.s2BaseSize * sizeof(QK_T)); + // 大小：2(开dB) * 2 * 64 * 128 * 4 = 128KB + mm1ResUB_ = bufUB_.Get(); + pipe->InitBuffer(bufQL1_, QUERY_BUF_NUM * constInfo_.mBaseSizeMax * D_BASIC_BLOCK * sizeof(Q_T)); + queryL1_ = bufQL1_.Get(); + pipe->InitBuffer(bufKeyL1_, KEY_BUF_NUM * keyBufferOffset_ * sizeof(K_T)); + keyL1_ = bufKeyL1_.Get(); + if constexpr (IS_MX) { + pipe->InitBuffer(bufQScaleL1_, QUERY_BUF_NUM * queryScaleBufferOffset * sizeof(SCALE_T)); + queryScaleL1_ = bufQScaleL1_.Get(); + pipe->InitBuffer(bufKeyScaleL1_, KEY_BUF_NUM * keyScaleBufferOffset_ * sizeof(SCALE_T)); + keyScaleL1_ = bufKeyScaleL1_.Get(); + } + + pipe->InitBuffer(bufQL0_, L0_BUF_NUM * constInfo_.mBaseSizeMax * D_BASIC_BLOCK_L0 * sizeof(L0_Q_T)); + queryL0_ = bufQL0_.Get(); + pipe->InitBuffer(bufKeyL0_, L0_BUF_NUM * D_BASIC_BLOCK_L0 * S2_BASIC_BLOCK_L0 * sizeof(L0_K_T)); + keyL0_ = bufKeyL0_.Get(); + + pipe->InitBuffer(bufL0C_, L0_BUF_NUM * constInfo_.mBaseSizeMax * S2_BASIC_BLOCK_L0 * sizeof(float)); + cL0_ = bufL0C_.Get(); +} + +template +__aicore__ inline void QLIV2Matmul::InitMm1GlobalTensor(const GlobalTensor &blkTableGm, + const GlobalTensor &keyGm, + const GlobalTensor &queryGm, + const GlobalTensor &keyScaleGmBf16, + const GlobalTensor &queryScaleGmBf16) +{ + blkTableGm_ = blkTableGm; + keyGm_ = keyGm; + queryGm_ = queryGm; + if constexpr (IS_MX) { + mxKeyScaleGmBf16_ = keyScaleGmBf16; + mxQueryScaleGmBf16_ = queryScaleGmBf16; + } +} + +template +__aicore__ inline void QLIV2Matmul::ComputeMm1(const QLIV2Common::RunInfo &runInfo) +{ + CrossCoreWaitFlag(QLIV2Common::ConstInfo::CROSS_VC_EVENT + + runInfo.loop % 2); + CrossCoreWaitFlag( + QLIV2Common::ConstInfo::CROSS_VC_EVENT + runInfo.loop % 2 + QLIV2Common::ConstInfo::AIV0_AIV1_OFFSET); + uint64_t s2GmBaseOffset = runInfo.s2Idx * constInfo_.s2BaseSize; + uint64_t s1gProcessSize = runInfo.actMBaseSize; + uint64_t s2ProcessSize = runInfo.actualSingleProcessSInnerSize; + if (s2BasicBlock_ == 128) { + for (uint64_t s2GmOffset = 0; s2GmOffset < s2ProcessSize; s2GmOffset += s2BasicBlock_) { + WaitFlag(KEY_MTE1_MTE2_EVENT + keyL1BufIdx_ % KEY_BUF_NUM); + uint64_t s2L1RealSize = + s2GmOffset + s2BasicBlock_ > s2ProcessSize ? s2ProcessSize - s2GmOffset : s2BasicBlock_; + if (PAGE_ATTENTION) { + KeyNd2NzForPA(s2L1RealSize, s2GmBaseOffset + s2GmOffset, runInfo); + if constexpr (IS_MX) { + // MX: PA路径需要绝对偏移 + LoadKScaleToL1(s2L1RealSize, s2GmBaseOffset + s2GmOffset, runInfo); + } + } else { + KeyNd2Nz(s2L1RealSize, s2GmOffset, runInfo); + if constexpr (IS_MX) { + // MX: 非PA路径只传循环内相对偏移（tensorKeyScaleOffset 已含 s2GmBaseOffset） + LoadKScaleToL1(s2L1RealSize, s2GmOffset, runInfo); + } + } + + SetFlag(MTE2_MTE1_EVENT); + WaitFlag(MTE2_MTE1_EVENT); + // s1gProcessSize当前必定不会超过2倍的s1g basic block + for (uint64_t s1gGmOffset = 0; s1gGmOffset < s1gProcessSize; s1gGmOffset += constInfo_.mBaseSizeMax) { + uint64_t s1gL1RealSize = s1gGmOffset + constInfo_.mBaseSizeMax > s1gProcessSize ? + s1gProcessSize - s1gGmOffset : + constInfo_.mBaseSizeMax; + uint64_t s1gL1SizeAlign2G = CeilAlign(s1gL1RealSize, 2 * constInfo_.gSize); + uint64_t s1gL1SizeAlign = CeilAlign(s1gL1SizeAlign2G, BLOCK_CUBE); + if (runInfo.isFirstS2InnerLoop && s2GmOffset == 0) { + queryL1Mte2BufIdx_++; + queryL1Mte1BufIdx_ = queryL1Mte2BufIdx_; + WaitFlag(QUERY_MTE1_MTE2_EVENT + queryL1Mte2BufIdx_ % QUERY_BUF_NUM); + QueryNd2Nz(s1gL1RealSize, s1gGmOffset, runInfo); + if constexpr (IS_MX) { + LoadQScaleToL1(s1gL1RealSize, s1gGmOffset, runInfo); + } + SetFlag(MTE2_MTE1_EVENT); + WaitFlag(MTE2_MTE1_EVENT); + } else { + queryL1Mte1BufIdx_ = + queryL1Mte2BufIdx_ - (CeilDiv(s1gProcessSize, constInfo_.mBaseSizeMax) - 1 - (s1gGmOffset > 0)); + } + for (uint64_t s2L1Offset = 0; s2L1Offset < s2L1RealSize; s2L1Offset += S2_BASIC_BLOCK_L0) { + uint64_t s2L0RealSize = + s2L1Offset + S2_BASIC_BLOCK_L0 > s2L1RealSize ? s2L1RealSize - s2L1Offset : S2_BASIC_BLOCK_L0; + for (uint64_t s1gL1Offset = 0; s1gL1Offset < s1gL1SizeAlign; + s1gL1Offset += constInfo_.mBaseSizeMax) { + WaitFlag(M_MTE1_EVENT + l0BufIdx_ % L0_BUF_NUM); + uint64_t s1gL0RealSize = s1gL1Offset + constInfo_.mBaseSizeMax > s1gL1SizeAlign ? + s1gL1SizeAlign - s1gL1Offset : + constInfo_.mBaseSizeMax; + LoadQueryToL0a(s1gL1Offset, s1gL1SizeAlign, s1gL0RealSize, runInfo); + LoadKeyToL0b(s2L1Offset, s2L1RealSize, s2L0RealSize, runInfo); + + SetFlag(MTE1_M_EVENT); + WaitFlag(MTE1_M_EVENT); + + WaitFlag(FIX_M_EVENT + l0BufIdx_ % L0_BUF_NUM); + ComputeL0c(s1gL0RealSize, s2L0RealSize, runInfo); + + SetFlag(M_MTE1_EVENT + l0BufIdx_ % L0_BUF_NUM); + + Fixp(s1gGmOffset + s1gL1Offset, s2GmOffset + s2L1Offset, s1gL0RealSize, s2L0RealSize, + s1gL1SizeAlign2G, runInfo); + SetFlag(FIX_M_EVENT + l0BufIdx_ % L0_BUF_NUM); + l0BufIdx_++; + } + } + if (s2GmOffset + s2BasicBlock_ >= s2ProcessSize && runInfo.isLastS2InnerLoop) { + SetFlag(QUERY_MTE1_MTE2_EVENT + queryL1Mte1BufIdx_ % QUERY_BUF_NUM); + } + } + SetFlag(KEY_MTE1_MTE2_EVENT + keyL1BufIdx_ % KEY_BUF_NUM); + keyL1BufIdx_++; + } + } else if (s2BasicBlock_ == 256) { + // 第一个s2循环 keycache置为false + if (runInfo.isFirstS2InnerLoop) { + isKeyCacheValid_ = false; + } + for (uint64_t s2GmOffset = 0; s2GmOffset < s2ProcessSize; s2GmOffset += s2BasicBlock_) { + // 缓存命中不需要进行key的搬运 + bool keyCacheHit = isKeyCacheValid_ && (s2GmBaseOffset >= keyGmStart_) && + (s2GmBaseOffset + s2ProcessSize <= keyGmStart_ + keyLoadedSize_); + if (!keyCacheHit) { + WaitFlag(KEY_MTE1_MTE2_EVENT + keyL1BufIdx_ % KEY_BUF_NUM); + // 缓存未命中，需要从GM搬到L1 min（256, 剩余s2） + uint64_t s2TotalRemainNum = runInfo.actS2Size - s2GmBaseOffset; + uint64_t s2L1LoadSize = (s2TotalRemainNum < s2BasicBlock_) ? s2TotalRemainNum : s2BasicBlock_; + if (PAGE_ATTENTION) { + KeyNd2NzForPA(s2L1LoadSize, s2GmBaseOffset + s2GmOffset, runInfo); + if constexpr (IS_MX) { + LoadKScaleToL1(s2L1LoadSize, s2GmBaseOffset + s2GmOffset, runInfo); + } + } else { + KeyNd2Nz(s2L1LoadSize, s2GmOffset, runInfo); + if constexpr (IS_MX) { + LoadKScaleToL1(s2L1LoadSize, s2GmOffset, runInfo); + } + } + + SetFlag(MTE2_MTE1_EVENT); + WaitFlag(MTE2_MTE1_EVENT); + + isKeyCacheValid_ = true; + keyGmStart_ = s2GmBaseOffset; + keyLoadedSize_ = s2L1LoadSize; + } + uint64_t l1S2Offset = s2GmBaseOffset - keyGmStart_; + uint64_t l1TotalSize = keyLoadedSize_; + // s1gProcessSize当前必定不会超过2倍的s1g basic block + for (uint64_t s1gGmOffset = 0; s1gGmOffset < s1gProcessSize; s1gGmOffset += constInfo_.mBaseSizeMax) { + uint64_t s1gL1RealSize = s1gGmOffset + constInfo_.mBaseSizeMax > s1gProcessSize ? + s1gProcessSize - s1gGmOffset : + constInfo_.mBaseSizeMax; + uint64_t s1gL1SizeAlign2G = CeilAlign(s1gL1RealSize, 2 * constInfo_.gSize); + uint64_t s1gL1SizeAlign = CeilAlign(s1gL1SizeAlign2G, BLOCK_CUBE); + if (runInfo.isFirstS2InnerLoop && s2GmOffset == 0) { + queryL1Mte2BufIdx_++; + queryL1Mte1BufIdx_ = queryL1Mte2BufIdx_; + WaitFlag(QUERY_MTE1_MTE2_EVENT + queryL1Mte2BufIdx_ % QUERY_BUF_NUM); + QueryNd2Nz(s1gL1RealSize, s1gGmOffset, runInfo); + if constexpr (IS_MX) { + LoadQScaleToL1(s1gL1RealSize, s1gGmOffset, runInfo); + } + SetFlag(MTE2_MTE1_EVENT); + WaitFlag(MTE2_MTE1_EVENT); + } else { + queryL1Mte1BufIdx_ = + queryL1Mte2BufIdx_ - (CeilDiv(s1gProcessSize, constInfo_.mBaseSizeMax) - 1 - (s1gGmOffset > 0)); + } + uint64_t s2Boundry = l1S2Offset + s2ProcessSize; + for (uint64_t s2L1Offset = l1S2Offset; s2L1Offset < s2Boundry; s2L1Offset += S2_BASIC_BLOCK_L0) { + uint64_t s2L0RealSize = s2L1Offset + S2_BASIC_BLOCK_L0 > l1S2Offset + s2ProcessSize ? + l1S2Offset + s2ProcessSize - s2L1Offset : + S2_BASIC_BLOCK_L0; + for (uint64_t s1gOffset = 0; s1gOffset < s1gL1SizeAlign; s1gOffset += constInfo_.mBaseSizeMax) { + WaitFlag(M_MTE1_EVENT + l0BufIdx_ % L0_BUF_NUM); + uint64_t s1gL0RealSize = s1gOffset + constInfo_.mBaseSizeMax > s1gL1SizeAlign ? + s1gL1SizeAlign - s1gOffset : + constInfo_.mBaseSizeMax; + LoadQueryToL0a(s1gOffset, s1gL1SizeAlign, s1gL0RealSize, runInfo); + LoadKeyToL0b(s2L1Offset, l1TotalSize, s2L0RealSize, runInfo); + + SetFlag(MTE1_M_EVENT); + WaitFlag(MTE1_M_EVENT); + + WaitFlag(FIX_M_EVENT + l0BufIdx_ % L0_BUF_NUM); + ComputeL0c(s1gL0RealSize, s2L0RealSize, runInfo); + + SetFlag(M_MTE1_EVENT + l0BufIdx_ % L0_BUF_NUM); + + Fixp(s1gGmOffset + s1gOffset, (s2L1Offset - l1S2Offset), s1gL0RealSize, s2L0RealSize, + s1gL1SizeAlign2G, runInfo); + SetFlag(FIX_M_EVENT + l0BufIdx_ % L0_BUF_NUM); + l0BufIdx_++; + } + } + if (s2GmOffset + s2BasicBlock_ >= s2ProcessSize && runInfo.isLastS2InnerLoop) { + SetFlag(QUERY_MTE1_MTE2_EVENT + queryL1Mte1BufIdx_ % QUERY_BUF_NUM); + } + } + bool l1FullyUsed = (s2GmBaseOffset + s2ProcessSize >= keyGmStart_ + keyLoadedSize_); + if (l1FullyUsed || runInfo.isLastS2InnerLoop) { + SetFlag(KEY_MTE1_MTE2_EVENT + keyL1BufIdx_ % KEY_BUF_NUM); + keyL1BufIdx_++; + isKeyCacheValid_ = false; + } + } + } + + CrossCoreSetFlag(QLIV2Common::ConstInfo::CROSS_CV_EVENT + + runInfo.loop % 2); + CrossCoreSetFlag( + QLIV2Common::ConstInfo::CROSS_CV_EVENT + runInfo.loop % 2 + QLIV2Common::ConstInfo::AIV0_AIV1_OFFSET); +} + +template +__aicore__ inline void QLIV2Matmul::KeyNd2Nz(uint64_t s2L1RealSize, uint64_t s2GmOffset, + const QLIV2Common::RunInfo &runInfo) +{ + Nd2NzParams nd2nzPara; + nd2nzPara.ndNum = 1; + nd2nzPara.nValue = s2L1RealSize; // 行数 + nd2nzPara.dValue = qkHeadDim_; + nd2nzPara.srcDValue = qkHeadDim_; + nd2nzPara.dstNzC0Stride = CeilAlign(s2L1RealSize, (uint64_t)BLOCK_CUBE); // 对齐到16 单位block + nd2nzPara.dstNzNStride = 1; + nd2nzPara.srcNdMatrixStride = 0; + nd2nzPara.dstNzMatrixStride = 0; + // 默认一块buf最多放两份 + DataCopy(keyL1_[(keyL1BufIdx_ % KEY_BUF_NUM) * keyBufferOffset_], + keyGm_[runInfo.tensorKeyOffset + s2GmOffset * qkHeadDim_], nd2nzPara); +} + +// blkNum, blkSize, N2, D +template +__aicore__ inline void QLIV2Matmul::KeyNd2NzForPA(uint64_t s2L1RealSize, uint64_t s2GmOffset, + const QLIV2Common::RunInfo &runInfo) +{ + uint64_t s2L1Offset = 0; + while (s2L1Offset < s2L1RealSize) { + uint64_t s2BlkId = (s2L1Offset + s2GmOffset) / constInfo_.kCacheBlockSize; + uint64_t s2BlkOffset = (s2L1Offset + s2GmOffset) % constInfo_.kCacheBlockSize; + uint64_t keyGmOffset = + blkTableGm_.GetValue(runInfo.bIdx * constInfo_.maxBlockNumPerBatch + s2BlkId) * constInfo_.keyStride0 + + s2BlkOffset * qkHeadDim_; + + uint64_t s2Mte2Size = s2L1RealSize - s2L1Offset; + s2Mte2Size = s2BlkOffset + s2Mte2Size >= constInfo_.kCacheBlockSize ? constInfo_.kCacheBlockSize - s2BlkOffset : + s2Mte2Size; + Nd2NzParams nd2nzPara; + nd2nzPara.ndNum = 1; + nd2nzPara.nValue = s2Mte2Size; // 行数 + nd2nzPara.dValue = qkHeadDim_; + nd2nzPara.srcDValue = qkHeadDim_; + nd2nzPara.dstNzC0Stride = CeilAlign(s2L1RealSize, (uint64_t)BLOCK_CUBE); // 对齐到16 单位block + nd2nzPara.dstNzNStride = 1; + nd2nzPara.srcNdMatrixStride = 0; + nd2nzPara.dstNzMatrixStride = 0; + DataCopy(keyL1_[(keyL1BufIdx_ % KEY_BUF_NUM) * keyBufferOffset_ + s2L1Offset * FP8_BLOCK_CUBE], + keyGm_[keyGmOffset], nd2nzPara); + + s2L1Offset += s2Mte2Size; + } +} + +// batch, s1, n2, g, d +template +__aicore__ inline void QLIV2Matmul::QueryNd2Nz(uint64_t s1gL1RealSize, uint64_t s1gGmOffset, + const QLIV2Common::RunInfo &runInfo) +{ + uint64_t dstNzC0Stride = CeilAlign(s1gL1RealSize, constInfo_.gSize * 2); + Nd2NzParams nd2nzPara; + nd2nzPara.ndNum = 1; + nd2nzPara.nValue = s1gL1RealSize; // 行数 + nd2nzPara.dValue = qkHeadDim_; + nd2nzPara.srcDValue = qkHeadDim_; + nd2nzPara.dstNzC0Stride = CeilAlign(dstNzC0Stride, (uint64_t)BLOCK_CUBE); + nd2nzPara.dstNzNStride = 1; + nd2nzPara.srcNdMatrixStride = 0; + nd2nzPara.dstNzMatrixStride = 0; + // 默认一块buf最多放两份 + DataCopy(queryL1_[(queryL1Mte2BufIdx_ % QUERY_BUF_NUM) * queryBufferOffset_], + queryGm_[runInfo.tensorQueryOffset + s1gGmOffset * qkHeadDim_], nd2nzPara); +} + +template +__aicore__ inline void QLIV2Matmul::LoadQScaleToL1(uint64_t s1gL1RealSize, uint64_t s1gGmOffset, + const QLIV2Common::RunInfo &runInfo) +{ + uint64_t scaleOffsetInBuf = (queryL1Mte2BufIdx_ % QUERY_BUF_NUM) * queryScaleBufferOffset; + LocalTensor scaleL1 = queryScaleL1_[scaleOffsetInBuf].template ReinterpretCast(); + uint32_t scalePerToken = scaleHeadDim_; + Dn2NzParams dn2Nzparam; + dn2Nzparam.dnNum = 1; + dn2Nzparam.nValue = scalePerToken / FP8_TWO; + dn2Nzparam.dValue = s1gL1RealSize; + dn2Nzparam.srcDnMatrixStride = 0; + dn2Nzparam.srcDValue = scalePerToken / FP8_TWO; + dn2Nzparam.dstNzC0Stride = scalePerToken / FP8_TWO; + dn2Nzparam.dstNzNStride = 1; + dn2Nzparam.dstNzMatrixStride = 0; + uint64_t gmOffset = runInfo.tensorQScaleOffset + s1gGmOffset * scalePerToken; + DataCopy(scaleL1, mxQueryScaleGmBf16_[gmOffset / FP8_TWO], dn2Nzparam); +} + +template +__aicore__ inline void QLIV2Matmul::LoadKScaleToL1(uint64_t s2L1RealSize, uint64_t s2GmOffset, + const QLIV2Common::RunInfo &runInfo) +{ + uint64_t scaleOffsetInBuf = (keyL1BufIdx_ % KEY_BUF_NUM) * keyScaleBufferOffset_; + LocalTensor scaleL1 = keyScaleL1_[scaleOffsetInBuf].template ReinterpretCast(); + uint32_t scalePerToken = scaleHeadDim_; + + if constexpr (PAGE_ATTENTION) { + uint64_t s2L1Offset = 0; + while (s2L1Offset < s2L1RealSize) { + uint64_t s2BlkId = (s2L1Offset + s2GmOffset) / constInfo_.kCacheBlockSize; + uint64_t s2BlkOffset = (s2L1Offset + s2GmOffset) % constInfo_.kCacheBlockSize; + uint64_t physicalBlkId = blkTableGm_.GetValue(runInfo.bIdx * constInfo_.maxBlockNumPerBatch + s2BlkId); + // kScale GM偏移：physicalBlkId * stride0 + s2BlkOffset * scalePerToken + uint64_t gmOffset = physicalBlkId * constInfo_.keyDequantScaleStride0 + s2BlkOffset * scalePerToken; + + uint64_t s2Mte2Size = s2L1RealSize - s2L1Offset; + s2Mte2Size = s2BlkOffset + s2Mte2Size >= constInfo_.kCacheBlockSize ? + constInfo_.kCacheBlockSize - s2BlkOffset : + s2Mte2Size; + + Dn2NzParams dn2Nzparam; + dn2Nzparam.dnNum = 1; + dn2Nzparam.nValue = scalePerToken / FP8_TWO; + dn2Nzparam.dValue = s2Mte2Size; + dn2Nzparam.srcDnMatrixStride = 0; + dn2Nzparam.srcDValue = scalePerToken / FP8_TWO; + dn2Nzparam.dstNzC0Stride = scalePerToken / FP8_TWO; + dn2Nzparam.dstNzNStride = 1; + dn2Nzparam.dstNzMatrixStride = 0; + DataCopy(scaleL1[s2L1Offset * scalePerToken / FP8_TWO], mxKeyScaleGmBf16_[gmOffset / FP8_TWO], dn2Nzparam); + + s2L1Offset += s2Mte2Size; + } + } else { + Dn2NzParams dn2Nzparam; + dn2Nzparam.dnNum = 1; + dn2Nzparam.nValue = scalePerToken / FP8_TWO; + dn2Nzparam.dValue = s2L1RealSize; + dn2Nzparam.srcDnMatrixStride = 0; + dn2Nzparam.srcDValue = scalePerToken / FP8_TWO; + dn2Nzparam.dstNzC0Stride = scalePerToken / FP8_TWO; + dn2Nzparam.dstNzNStride = 1; + dn2Nzparam.dstNzMatrixStride = 0; + uint64_t gmOffset = runInfo.tensorKeyScaleOffset + s2GmOffset * scalePerToken; + DataCopy(scaleL1, mxKeyScaleGmBf16_[gmOffset / FP8_TWO], dn2Nzparam); + } +} + +template +__aicore__ inline void QLIV2Matmul::LoadQueryToL0a(uint64_t s1gL1Offset, uint64_t s1gL1RealSize, + uint64_t s1gL0RealSize, const QLIV2Common::RunInfo &runInfo) +{ + LoadData2DParamsV2 loadData2DParamsV2; + loadData2DParamsV2.mStartPosition = CeilDiv(s1gL1Offset, BLOCK_CUBE); + loadData2DParamsV2.kStartPosition = 0; + loadData2DParamsV2.mStep = CeilDiv(s1gL0RealSize, BLOCK_CUBE); + loadData2DParamsV2.kStep = + IS_MXFP4 ? constInfo_.headDim / MX_LOAD_SCALE_ALIGN : CeilDiv(constInfo_.headDim, FP8_BLOCK_CUBE); + loadData2DParamsV2.srcStride = CeilDiv(s1gL1RealSize, BLOCK_CUBE); + loadData2DParamsV2.dstStride = CeilDiv(s1gL0RealSize, BLOCK_CUBE); + loadData2DParamsV2.ifTranspose = false; + + if constexpr (IS_MX) { + // MX: 使用LoadData Mx变体，同时加载数据和scale到L0 + LoadData2DMxParams loadDataMxParams; + loadDataMxParams.xStartPosition = CeilDiv(s1gL1Offset, BLOCK_CUBE); + loadDataMxParams.yStartPosition = 0; + loadDataMxParams.xStep = CeilDiv(s1gL0RealSize, BLOCK_CUBE); + loadDataMxParams.yStep = scaleHeadDim_ / FP8_TWO; + loadDataMxParams.srcStride = loadDataMxParams.yStep; + loadDataMxParams.dstStride = loadDataMxParams.yStep; + + uint64_t queryDataOffsetInBuf = (queryL1Mte1BufIdx_ % QUERY_BUF_NUM) * queryBufferOffset_; + uint64_t queryScaleOffsetInBuf = (queryL1Mte1BufIdx_ % QUERY_BUF_NUM) * queryScaleBufferOffset; + LocalTensor queryL0Tensor = queryL0_[(l0BufIdx_ % L0_BUF_NUM) * l0abBufferOffset_]; + LocalTensor queryScaleL1Tensor = + queryScaleL1_[queryScaleOffsetInBuf].template ReinterpretCast(); + if constexpr (IS_MXFP4) { + // MXFP4: Q data本身按E2M1参与计算，而queryL1_是uint8_t类型 + LocalTensor queryL1MxTensor = + queryL1_[queryDataOffsetInBuf].template ReinterpretCast(); + LoadData(queryL0Tensor, queryL1MxTensor, queryScaleL1Tensor, loadData2DParamsV2, loadDataMxParams); + } else { + LocalTensor queryL1MxTensor = queryL1_[queryDataOffsetInBuf]; + LoadData(queryL0Tensor, queryL1MxTensor, queryScaleL1Tensor, loadData2DParamsV2, loadDataMxParams); + } + } else { + LoadData(queryL0_[(l0BufIdx_ % L0_BUF_NUM) * l0abBufferOffset_], + queryL1_[(queryL1Mte1BufIdx_ % QUERY_BUF_NUM) * queryBufferOffset_], loadData2DParamsV2); + } +} + +template +__aicore__ inline void QLIV2Matmul::LoadKeyToL0b(uint64_t s2L1Offset, uint64_t s2L1RealSize, + uint64_t s2L0RealSize, const QLIV2Common::RunInfo &runInfo) +{ + LoadData2DParamsV2 loadData2DParamsV2; + loadData2DParamsV2.mStartPosition = CeilDiv(s2L1Offset, BLOCK_CUBE); + loadData2DParamsV2.kStartPosition = 0; + loadData2DParamsV2.mStep = CeilDiv(s2L0RealSize, BLOCK_CUBE); + loadData2DParamsV2.kStep = + IS_MXFP4 ? constInfo_.headDim / MX_LOAD_SCALE_ALIGN : CeilDiv(constInfo_.headDim, FP8_BLOCK_CUBE); + loadData2DParamsV2.srcStride = CeilDiv(s2L1RealSize, BLOCK_CUBE); + loadData2DParamsV2.dstStride = CeilDiv(s2L0RealSize, BLOCK_CUBE); + loadData2DParamsV2.ifTranspose = false; + + if constexpr (IS_MX) { + // MX: 使用LoadData Mx变体，同时加载数据和scale到L0 + LoadData2DMxParams loadDataMxParams; + loadDataMxParams.xStartPosition = CeilDiv(s2L1Offset, BLOCK_CUBE); + loadDataMxParams.yStartPosition = 0; + loadDataMxParams.xStep = CeilDiv(s2L0RealSize, BLOCK_CUBE); + loadDataMxParams.yStep = scaleHeadDim_ / FP8_TWO; + loadDataMxParams.srcStride = loadDataMxParams.yStep; + loadDataMxParams.dstStride = loadDataMxParams.yStep; + + uint64_t keyDataOffsetInBuf = (keyL1BufIdx_ % KEY_BUF_NUM) * keyBufferOffset_; + uint64_t keyScaleOffsetInBuf = (keyL1BufIdx_ % KEY_BUF_NUM) * keyScaleBufferOffset_; + LocalTensor keyL0Tensor = keyL0_[(l0BufIdx_ % L0_BUF_NUM) * KEY_L0_BUFFER_OFFSET]; + LocalTensor keyScaleL1Tensor = + keyScaleL1_[keyScaleOffsetInBuf].template ReinterpretCast(); + if constexpr (IS_MXFP4) { + LocalTensor keyL1MxTensor = + keyL1_[keyDataOffsetInBuf].template ReinterpretCast(); + LoadData(keyL0Tensor, keyL1MxTensor, keyScaleL1Tensor, loadData2DParamsV2, loadDataMxParams); + } else { + LocalTensor keyL1MxTensor = keyL1_[keyDataOffsetInBuf]; + LoadData(keyL0Tensor, keyL1MxTensor, keyScaleL1Tensor, loadData2DParamsV2, loadDataMxParams); + } + } else { + LoadData(keyL0_[(l0BufIdx_ % L0_BUF_NUM) * KEY_L0_BUFFER_OFFSET], + keyL1_[(keyL1BufIdx_ % KEY_BUF_NUM) * keyBufferOffset_], loadData2DParamsV2); + } +} + +template +__aicore__ inline void QLIV2Matmul::ComputeL0c(uint64_t s1gL0RealSize, uint64_t s2L0RealSize, + const QLIV2Common::RunInfo &runInfo) +{ + MmadParams mmadParams; + mmadParams.m = CeilAlign(s1gL0RealSize, BLOCK_CUBE); + mmadParams.n = s2L0RealSize; + mmadParams.k = constInfo_.headDim; + mmadParams.cmatrixInitVal = true; + mmadParams.cmatrixSource = false; + LocalTensor queryL0Tensor = queryL0_[(l0BufIdx_ % L0_BUF_NUM) * l0abBufferOffset_]; + LocalTensor keyL0Tensor = keyL0_[(l0BufIdx_ % L0_BUF_NUM) * KEY_L0_BUFFER_OFFSET]; + Mmad(cL0_[(l0BufIdx_ % L0_BUF_NUM) * l0cBufferOffset_], queryL0Tensor, keyL0Tensor, mmadParams); + if ((mmadParams.m / 16) * (mmadParams.n / 16) < 10) { + PipeBarrier(); + } +} + +template +__aicore__ inline void QLIV2Matmul::Fixp(uint64_t s1gGmOffset, uint64_t s2GmOffset, uint64_t s1gL0RealSize, + uint64_t s2L0RealSize, uint64_t s1gSizeAlign2G, + const QLIV2Common::RunInfo &runInfo) +{ + SetFlag(M_FIX_EVENT + l0BufIdx_ % L0_BUF_NUM); + WaitFlag(M_FIX_EVENT + l0BufIdx_ % L0_BUF_NUM); + + if constexpr (std::is_same_v || std::is_same_v) { + // s1gL0RealSize：2*gSize(128)对齐, 最大256 + // s2L0RealSize <= S2_BASIC_BLOCK_L0, 未约束 + uint32_t nSize = (s2L0RealSize + 7) >> 3 << 3; // 32B对齐 + FixpipeParamsC310 fixpipeParams; + // 固定参数 + fixpipeParams.mSize = s1gSizeAlign2G; + fixpipeParams.srcStride = (s1gL0RealSize + 1) >> 1 << 1; // 已16对齐 + fixpipeParams.dstStride = UB_BANK_DEPTH_STRIDE / sizeof(QK_T); // 落到同一个bank + // 双目标模式，按M维度拆分，M/2*N写入每个UB，M必须为2的倍数 + fixpipeParams.dualDstCtl = 1; + + // nSize已保证N方向32B对齐 + if (nSize <= (256 / sizeof(float))) { + // N方向小于一个bank(256B), 只需搬一个ND块, 且不用补齐 + fixpipeParams.nSize = nSize; + fixpipeParams.params.ndNum = 1; + fixpipeParams.params.srcNdStride = 0; + fixpipeParams.params.dstNdStride = 0; + } else { + // N方向在(256B, 512B]范围， 直接按512B搬, 注意此时不能开unitflag + fixpipeParams.nSize = S2_BASIC_BLOCK_L0 / 2; // 分2个ND搬, S2_BASIC_BLOCK_L0不为128会有问题 + fixpipeParams.params.ndNum = 2; + fixpipeParams.params.srcNdStride = ((fixpipeParams.mSize + 15) / 16) * fixpipeParams.nSize; + fixpipeParams.params.dstNdStride = + constInfo_.s2BaseSize * constInfo_.mBaseSizeMax / 2; // s2BasicBlock_ * M_BASE_SIZE / 2 + } + Fixpipe(mm1ResUB_[(runInfo.loop % 2) * constInfo_.s2BaseSize / 2], + // 未考虑s1gGmOffset和s2GmOffset，将matmul结果从L0C搬运到UB + cL0_[(l0BufIdx_ % L0_BUF_NUM) * l0cBufferOffset_], fixpipeParams); + } else { + uint32_t nSize = CeilAlign(s2L0RealSize, static_cast(UB_BLOCK / sizeof(QK_T))); + // 有效数据不足16行，只需输出部分行即可; L0C上bmm1结果矩阵M方向size必须是偶数 + uint32_t mSize = s1gSizeAlign2G; + // L0C上matmul结果相邻连续数据片断间隔, 单位为16 * sizeof(T) + uint32_t srcStride = ((mSize + 15) / 16) * 16; + FixpipeParamsC310 fixpipeParams; // L0C->UB + fixpipeParams.nSize = nSize; + fixpipeParams.mSize = mSize / 2; // M方向每个AIV一半 + fixpipeParams.srcStride = srcStride; + fixpipeParams.dstStride = UB_BANK_DEPTH_STRIDE / sizeof(QK_T); // 落到同一个bank + fixpipeParams.params.ndNum = 1; + fixpipeParams.params.srcNdStride = 0; + fixpipeParams.params.dstNdStride = 0; + // F322BF16和ReLU属于随路功能，不能与dualDstCtl同时使用，分别写入两个SUB BLOCK。 + fixpipeParams.dualDstCtl = 0; + fixpipeParams.quantPre = F322BF16; + fixpipeParams.reluEn = true; + fixpipeParams.subBlockId = 0; + Fixpipe(mm1ResUB_[(runInfo.loop % 2) * (UB_BANK_STRIDE / sizeof(QK_T))], + cL0_[(l0BufIdx_ % L0_BUF_NUM) * l0cBufferOffset_], fixpipeParams); + + fixpipeParams.subBlockId = 1; + Fixpipe( + mm1ResUB_[(runInfo.loop % 2) * (UB_BANK_STRIDE / sizeof(QK_T))], + cL0_[(l0BufIdx_ % L0_BUF_NUM) * l0cBufferOffset_ + mSize / 2 * 16], fixpipeParams); + } +} + +template +__aicore__ inline void QLIV2Matmul::AllocEventID() +{ + SetMMLayoutTransform(true); + SetFlag(KEY_MTE1_MTE2_EVENT + 0); + SetFlag(KEY_MTE1_MTE2_EVENT + 1); + SetFlag(KEY_MTE1_MTE2_EVENT + 2); + + SetFlag(QUERY_MTE1_MTE2_EVENT + 0); + SetFlag(QUERY_MTE1_MTE2_EVENT + 1); + + SetFlag(M_MTE1_EVENT + 0); + SetFlag(M_MTE1_EVENT + 1); + + SetFlag(FIX_M_EVENT + 0); + SetFlag(FIX_M_EVENT + 1); +} + +template +__aicore__ inline void QLIV2Matmul::FreeEventID() +{ + SetMMLayoutTransform(false); + WaitFlag(KEY_MTE1_MTE2_EVENT + 0); + WaitFlag(KEY_MTE1_MTE2_EVENT + 1); + WaitFlag(KEY_MTE1_MTE2_EVENT + 2); + + WaitFlag(QUERY_MTE1_MTE2_EVENT + 0); + WaitFlag(QUERY_MTE1_MTE2_EVENT + 1); + + WaitFlag(M_MTE1_EVENT + 0); + WaitFlag(M_MTE1_EVENT + 1); + + WaitFlag(FIX_M_EVENT + 0); + WaitFlag(FIX_M_EVENT + 1); +} +} // namespace QLIV2Kernel +#endif // QUANT_LIGHTNING_INDEXER_V2_SERVICE_CUBE_H diff --git a/xllm_ops/attention/quant_lightning_indexer_v2/op_kernel/arch35/quant_lightning_indexer_v2_service_vector_arch35.h b/xllm_ops/attention/quant_lightning_indexer_v2/op_kernel/arch35/quant_lightning_indexer_v2_service_vector_arch35.h new file mode 100644 index 0000000..33700cf --- /dev/null +++ b/xllm_ops/attention/quant_lightning_indexer_v2/op_kernel/arch35/quant_lightning_indexer_v2_service_vector_arch35.h @@ -0,0 +1,1121 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file quant_lightning_indexer_v2_service_vector_arch35.h + * \brief + */ +#ifndef QUANT_LIGHTNING_INDEXER_V2_SERVICE_VECTOR_H +#define QUANT_LIGHTNING_INDEXER_V2_SERVICE_VECTOR_H + +#include "kernel_operator.h" +#include "kernel_operator_list_tensor_intf.h" +#include "kernel_tiling/kernel_tiling.h" +#include "lib/matmul_intf.h" +#include "lib/matrix/matmul/tiling.h" +#include "quant_lightning_indexer_v2_common_arch35.h" +#include "../arch35/vf/quant_lightning_indexer_v2_vector1.h" +#include "../arch35/vf/quant_lightning_indexer_v2_topk.h" + +namespace QLIV2Kernel { +using namespace QLIV2Common; +constexpr uint32_t TRUNK_LEN_16K = 16384; +// LD阶段UB目标占用（248KB = 256KB UB - 8KB 系统保留） +constexpr uint32_t LD_UB_TARGET_BYTES = 248 * 1024; +constexpr uint32_t TRUNK_LEN_12K = 12288; +constexpr uint32_t TRUNK_LEN_8K = 8192; +constexpr uint32_t TRUNK_LEN_4K = 4096; +constexpr uint32_t TRUNK_LEN_2K = 2048; +constexpr uint32_t TOPK_2K = 2048; +constexpr uint32_t TOPK_3K = 3072; +constexpr uint32_t TOPK_4K = 4096; +constexpr uint32_t TOPK_5K = 5120; +constexpr uint32_t TOPK_6K = 6144; +constexpr uint32_t DUPSIZE = 256; // duplicate所需额外地址空间 + +template +class QLIV2Vector { +public: + // =================================类型定义区================================= + static constexpr LI_LAYOUT Q_LAYOUT_T = QLIV2T::layout; + static constexpr LI_LAYOUT K_LAYOUT_T = QLIV2T::keyLayout; + static constexpr bool PAGE_ATTENTION = QLIV2T::pageAttention; + + using QK_T = typename QLIV2T::queryKeyType; + using SCORE_T = typename QLIV2T::scoreType; + using WEIGHT_T = typename QLIV2T::weightType; + static constexpr bool IS_MX = QLIV2T::isMx; + static constexpr bool IS_MXFP4 = QLIV2T::isMxFp4; + using SCALE_T = typename QLIV2T::scaleType; + static constexpr bool IS_WEIGHT_FP16 = QLIV2T::isWeightFP16; + __aicore__ inline QLIV2Vector(){}; + __aicore__ inline void ProcessVec1(const QLIV2Common::RunInfo &info); + __aicore__ inline void ProcessTopK(const QLIV2Common::RunInfo &info); + __aicore__ inline void ProcessLD(); + __aicore__ inline void InitBuffers(TPipe *pipe); + __aicore__ inline void InitParams(const struct QLIV2Common::ConstInfo &constInfo, + const struct QLIV2Common::LdSplitCoreInfo &ldInfo, + const QLIV2TilingData *__restrict tilingData); + __aicore__ inline void InitVecWorkspaceTensor(GlobalTensor scoreGm, GlobalTensor ldScoreGm, + GlobalTensor ldIndexGm); + __aicore__ inline void InitVecInputTensor(GlobalTensor weightsGm, GlobalTensor indiceOutGm, + GlobalTensor blockTableGm, GlobalTensor valueOutGm, + GlobalTensor outputIdxOffsetGm, + GlobalTensor qScaleGm = {}, GlobalTensor kScaleGm = {}); + __aicore__ inline void CleanInvalidOutput(int64_t invalidS1offset); + __aicore__ inline void AllocEventID(); + __aicore__ inline void FreeEventID(); + __aicore__ inline void InitLDBuffers(TPipe *pipe, const struct QLIV2Common::LdSplitCoreInfo &ldInfo); + __aicore__ inline void DoTndPadding(const QLIV2Common::RunInfo &runInfo); + +protected: + GlobalTensor scoreGm; + GlobalTensor weightsGm; + GlobalTensor ldScoreGm; + GlobalTensor ldIndexGm; + GlobalTensor qScaleGm; + GlobalTensor kScaleGm; + GlobalTensor indiceOutGm; + GlobalTensor valueOutGm; + GlobalTensor blockTableGm; + GlobalTensor outputIdxOffsetGm; + // =================================常量区================================= + static constexpr uint32_t VEC1_V_MTE2_EVENT_KSCALE = EVENT_ID0; + static constexpr uint32_t VEC1_MTE2_V_EVENT_KSCALE = EVENT_ID1; + static constexpr uint32_t VEC1_V_MTE3_EVENT = EVENT_ID2; + static constexpr uint32_t VEC1_MTE3_V_EVENT = EVENT_ID3; + static constexpr uint32_t VEC1_V_MTE2_EVENT_QSCALE = EVENT_ID6; + static constexpr uint32_t VEC1_MTE2_V_EVENT_QSCALE = EVENT_ID3; + static constexpr uint32_t TOPK_V_MTE2_EVENT = EVENT_ID4; + static constexpr uint32_t TOPK_MTE2_V_EVENT = EVENT_ID5; + static constexpr uint32_t TOPK_V_MTE3_EVENT = EVENT_ID6; + static constexpr uint32_t TOPK_MTE3_V_EVENT = EVENT_ID7; + + static constexpr uint32_t KSCALE_S_MTE2_EVENT = EVENT_ID7; + static constexpr uint32_t MTE3_MTE2_EVENT = EVENT_ID0; + static constexpr uint32_t V_MTE2_EVENT1 = EVENT_ID2; + static constexpr uint32_t V_MTE2_EVENT2 = EVENT_ID3; + static constexpr uint32_t V_MTE2_EVENT3 = EVENT_ID5; + static constexpr uint32_t MTE3_V_EVENT = EVENT_ID6; + + static constexpr uint64_t BLOCK_CUBE = 16; + +private: + __aicore__ inline void GetKeyScale(const QLIV2Common::RunInfo &runInfo, LocalTensor &kScaleUB, + int64_t batchId, int64_t startS2, int64_t getLen); + // ================================Local Buffer区==================================== + + // tmp buff for vector + TBuf resMm1Buf_; + LocalTensor resMm1UB_; + // tmp buff for weight + TBuf weightBuf_; + LocalTensor weightUB_; + // tmp buff for weight cast float + TBuf weightTempBuf_; + LocalTensor weightTempUB_; + // tmp buff for kScale + TBuf kScaleBuf_; + LocalTensor kScaleUB_; + // tmp buff for qScale + TBuf qScaleBuf_; + LocalTensor qScaleUB_; + // tmp buff for out + TBuf outBuf_; + LocalTensor vec1OutUB_; + TBuf valueOutBuf_; + LocalTensor valueOutLocal_; + // tmp buff for LD + TBuf ldValueBuf_; + LocalTensor ldValueLocal_; // SCORE_T + TBuf ldIndexBuf_; + LocalTensor ldIndexLocal_; + + TBuf topkIndexBuf_; + LocalTensor topkIndexLocal_; + + TBuf topkValueBuf_; + LocalTensor topkValueLocal_; + + // tmp buff for topk + TBuf mrgValueBuf_; + LocalTensor mrgValueLocal_; + + TBuf indicesOutBuf_; + LocalTensor indicesOutLocal_; + + TBuf scoreOutBuf_; + LocalTensor scoreOutLocal_; + + TBuf topkSharedTmpBuf_; + LocalTensor topkSharedTmpLocal_; + + int32_t blockId_ = -1; + // para for vector + int32_t groupInner_ = 0; + int32_t globalTopkNum_ = 0; + int64_t blockS2StartIdx_ = 0; + int32_t gSize_ = 0; + int32_t kSeqSize_ = 0; + int32_t kHeadNum_ = 0; + int32_t qHeadNum_ = 0; + int32_t s1BaseSize_ = 0; + int32_t s2BaseSize_ = 0; + int32_t kCacheBlockSize_ = 0; + int32_t maxBlockNumPerBatch_ = 0; + uint32_t topkCount_ = 0; + uint32_t topkCountAlign256_ = 0; // topkCount对齐到256(直方图需要)，支持topk泛化 + uint32_t topkCountAlign16_ = 0; // LD读取到UB，需要满足32B对齐 + float globalQScale_ = 1.0f; // quantMode=4时全局query scale + float globalKScale_ = 1.0f; // quantMode=4时全局key scale + uint32_t trunkLen_ = 0; // ProcessTopK（非LD路径）每次处理的s2长度 + uint32_t trunkLenLd_ = 0; // ProcessLD路径每次处理的s2长度，根据sparseCount动态计算以填满UB + bool returnValueFlag = false; + + struct QLIV2Common::ConstInfo constInfo_; + struct QLIV2Common::LdSplitCoreInfo ldInfo_; + topk::LITopk topkOp_; +}; + +template +__aicore__ inline void QLIV2Vector::InitBuffers(TPipe *pipe) +{ + // 大小：2(开dB) 2 * 64 * 128 * 4 = 128KB + pipe->InitBuffer(resMm1Buf_, 2 * CeilDiv(constInfo_.mBaseSizeMax, 2) * s2BaseSize_ * sizeof(QK_T)); + resMm1UB_ = resMm1Buf_.Get(); // qk + // weight buffer按WEIGHT_T访问；UB_BANK_DEPTH_STRIDE的单位已经是字节，无需再乘数据类型大小。 + pipe->InitBuffer(weightBuf_, 2 * CeilDiv(s1BaseSize_, 2) * UB_BANK_DEPTH_STRIDE); + weightUB_ = weightBuf_.Get(); // weight + pipe->InitBuffer(weightTempBuf_, 2 * CeilDiv(s1BaseSize_, 2) * UB_BANK_DEPTH_STRIDE); + weightTempUB_ = weightTempBuf_.Get(); + // 大小：2(开dB) * 128 * 4 = 1KB + pipe->InitBuffer(kScaleBuf_, 2 * s2BaseSize_ * 16 * sizeof(SCALE_T)); + kScaleUB_ = kScaleBuf_.Get(); // kScale + // 大小：2(开dB) * 2 * 64 * 4 = 1KB + pipe->InitBuffer(qScaleBuf_, + 2 * CeilDiv(s1BaseSize_, 2) * UB_BANK_DEPTH_STRIDE); + qScaleUB_ = qScaleBuf_.Get(); // qScale + // 大小：2(开dB) * 2 * 128 * 4 = 2KB + pipe->InitBuffer(outBuf_, 2 * CeilDiv(s1BaseSize_, 2) * s2BaseSize_ * sizeof(SCORE_T)); + vec1OutUB_ = outBuf_.Get(); // out + // Topk + pipe->InitBuffer(mrgValueBuf_, (topkCountAlign256_ + trunkLen_) * sizeof(SCORE_T)); + // 大小：(topkCountAlign256_ + 每次排序长度) * sizeof(SCORE_T) + mrgValueLocal_ = mrgValueBuf_.Get(); + valueOutLocal_ = mrgValueBuf_.Get(); + + pipe->InitBuffer(indicesOutBuf_, topkCountAlign256_ * sizeof(uint32_t) + DUPSIZE); + // 大小：(topkCountAlign256_ + 64) * 4 64:duplicate刷-1需要额外空间 + indicesOutLocal_ = indicesOutBuf_.Get(); + + pipe->InitBuffer(scoreOutBuf_, topkCountAlign256_ * sizeof(SCORE_T) + DUPSIZE); + // (topkCountAlign256_ + 64) * sizeof(SCORE_T) 64:duplicate刷-1额外空间 + scoreOutLocal_ = scoreOutBuf_.Get(); + + uint64_t topkSharedTmpSize = topkOp_.GetSharedTmpBufferSize(); + pipe->InitBuffer(topkSharedTmpBuf_, topkSharedTmpSize); + topkSharedTmpLocal_ = topkSharedTmpBuf_.Get(); + topkOp_.InitBuffers(topkSharedTmpLocal_, indicesOutLocal_); + + // MX场景不走vector侧kScale路径 + if constexpr (!IS_MX) { + Duplicate(kScaleUB_, (SCALE_T)(0), 2 * s2BaseSize_ * 16); + } +} + +template +__aicore__ inline void QLIV2Vector::InitLDBuffers(TPipe *pipe, + const struct QLIV2Common::LdSplitCoreInfo &ldInfo) +{ + pipe->Reset(); + + // 根据sparseCount动态计算LD的trunk长度，使LD阶段UB占用尽量接近248KB且不超限 + // + // UB总占用 = Σ各buffer，按 A=topkCountAlign256_、L=trunkLenLd_ 归并后： + // Total = 20·A + 8·L + fixedBytes + // A系数20 = ldIndex(4)+indicesOut(4)+mrgValue(2)+scoreOut(2)+valueOut(2)+topkTmp(6) + // L系数8 = ldIndex(4)+mrgValue(2)+topkTmp(2) + // fixedBytes = topk内部固定区 + 三输出buffer的+64余量(见下方命名常量) + + // topk内部固定区: histograms(256)+idxHigh(256)+idxLow(256)+nkValue(64)，均为uint32 + constexpr uint32_t TOPK_FIXED_BYTES = (3 * 256 + 64) * sizeof(uint32_t); + // 三输出buffer的+64余量(用于duplicate刷-1/刷0): indicesOut(uint32)+scoreOut(SCORE_T)+valueOut(bf16) + constexpr uint32_t OUTPUT_PAD_BYTES = 64 * (sizeof(uint32_t) + sizeof(SCORE_T) + sizeof(bfloat16_t)); + const uint32_t fixedBytes = TOPK_FIXED_BYTES + OUTPUT_PAD_BYTES; + + uint32_t availForTrunk = (LD_UB_TARGET_BYTES > 20 * topkCountAlign256_ + fixedBytes) ? + (LD_UB_TARGET_BYTES - 20 * topkCountAlign256_ - fixedBytes) : + 0; + trunkLenLd_ = availForTrunk / 8; + // trunkLenLd_向下对齐到256，确保LD阶段UB不超限且满足topk直方图对齐要求 + trunkLenLd_ = (trunkLenLd_ / 256) * 256; + + // 更新topkOp_的trunkLen，使GetSharedTmpBufferSize返回正确尺寸 + topkOp_.Init(topkCount_, trunkLenLd_); + + pipe->InitBuffer(ldIndexBuf_, (topkCountAlign256_ + trunkLenLd_) * sizeof(uint32_t)); + pipe->InitBuffer(indicesOutBuf_, (topkCountAlign256_ + 64) * sizeof(uint32_t)); + pipe->InitBuffer(mrgValueBuf_, (topkCountAlign256_ + trunkLenLd_) * sizeof(SCORE_T)); + pipe->InitBuffer(scoreOutBuf_, (topkCountAlign256_ + 64) * sizeof(SCORE_T)); + pipe->InitBuffer(valueOutBuf_, (topkCountAlign256_ + 64) * sizeof(bfloat16_t)); + // 补齐topk共享tmp buffer分配：pipe->Reset()后原topkSharedTmpBuf_已释放， + // 需重新分配并重新绑定topkOp_内部指针，避免悬垂指针 + // Re-allocate topk shared tmp buffer: pipe->Reset() freed the original + // topkSharedTmpBuf_; re-allocate and rebind topkOp_ internal pointers + pipe->InitBuffer(topkSharedTmpBuf_, topkOp_.GetSharedTmpBufferSize()); + + ldIndexLocal_ = ldIndexBuf_.Get(); + indicesOutLocal_ = indicesOutBuf_.Get(); + mrgValueLocal_ = mrgValueBuf_.Get(); + scoreOutLocal_ = scoreOutBuf_.Get(); + valueOutLocal_ = valueOutBuf_.Get(); + topkSharedTmpLocal_ = topkSharedTmpBuf_.Get(); + topkOp_.InitBuffers(topkSharedTmpLocal_, indicesOutLocal_); +} + +template +__aicore__ inline void QLIV2Vector::InitParams(const struct QLIV2Common::ConstInfo &constInfo, + const struct QLIV2Common::LdSplitCoreInfo &ldInfo, + const QLIV2TilingData *__restrict tilingData) +{ + this->constInfo_ = constInfo; + this->ldInfo_ = ldInfo; + blockS2StartIdx_ = 0; + gSize_ = constInfo.gSize; + kSeqSize_ = constInfo.kSeqSize; + // define N2 para + kHeadNum_ = constInfo.kHeadNum; + qHeadNum_ = constInfo.qHeadNum; + // define MMBase para + s1BaseSize_ = constInfo.s1BaseSize; // 4 + s2BaseSize_ = constInfo.s2BaseSize; // 128 + kCacheBlockSize_ = constInfo.kCacheBlockSize; + maxBlockNumPerBatch_ = constInfo.maxBlockNumPerBatch; + returnValueFlag = constInfo.returnValue; + blockId_ = GetBlockIdx(); + trunkLen_ = constInfo.sparseCount <= TOPK_2K ? TRUNK_LEN_16K : + constInfo.sparseCount <= TOPK_3K ? TRUNK_LEN_12K : + constInfo.sparseCount <= TOPK_4K ? TRUNK_LEN_8K : + constInfo.sparseCount <= TOPK_5K ? TRUNK_LEN_4K : + constInfo.sparseCount <= TOPK_6K ? TRUNK_LEN_2K : + TRUNK_LEN_12K; + topkCount_ = constInfo.sparseCount; + topkCountAlign256_ = QLIV2Common::Align(constInfo.sparseCount, (uint64_t)256); // topkCount对齐到256 + topkCountAlign16_ = QLIV2Common::Align(constInfo.sparseCount, (uint64_t)16); // topkCount对齐到16 + topkOp_.Init(topkCount_, trunkLen_); + + if constexpr (!IS_MX) { + if (constInfo_.quantMode == 4) { + globalQScale_ = qScaleGm.GetValue(0); + globalKScale_ = kScaleGm.GetValue(0); + } + } +} + +template +__aicore__ inline void QLIV2Vector::InitVecInputTensor( + GlobalTensor weightsGm, GlobalTensor indiceOutGm, GlobalTensor blockTableGm, + GlobalTensor valueOutGm, GlobalTensor outputIdxOffsetGm, GlobalTensor qScaleGm, + GlobalTensor kScaleGm) +{ + this->weightsGm = weightsGm; + this->indiceOutGm = indiceOutGm; + this->blockTableGm = blockTableGm; + this->valueOutGm = valueOutGm; + this->outputIdxOffsetGm = outputIdxOffsetGm; + if constexpr (!IS_MX) { + this->qScaleGm = qScaleGm; + this->kScaleGm = kScaleGm; + } +} + +template +__aicore__ inline void QLIV2Vector::InitVecWorkspaceTensor(GlobalTensor scoreGm, + GlobalTensor ldScoreGm, + GlobalTensor ldIndexGm) +{ + this->scoreGm = scoreGm; // resucesum*k + this->ldScoreGm = ldScoreGm; + this->ldIndexGm = ldIndexGm; +} + +template +__aicore__ inline void QLIV2Vector::AllocEventID() +{ + if constexpr (!IS_MX) { + SetFlag(VEC1_V_MTE2_EVENT_KSCALE + 0); + SetFlag(VEC1_V_MTE2_EVENT_KSCALE + 1); + } + SetFlag(VEC1_MTE3_V_EVENT + 0); + SetFlag(VEC1_MTE3_V_EVENT + 1); + SetFlag(VEC1_V_MTE2_EVENT_QSCALE + 0); + SetFlag(VEC1_V_MTE2_EVENT_QSCALE + 1); + + SetFlag(TOPK_V_MTE2_EVENT); + SetFlag(TOPK_MTE3_V_EVENT); + SetFlag(V_MTE2_EVENT1); +} + +template +__aicore__ inline void QLIV2Vector::FreeEventID() +{ + if constexpr (!IS_MX) { + WaitFlag(VEC1_V_MTE2_EVENT_KSCALE + 0); + WaitFlag(VEC1_V_MTE2_EVENT_KSCALE + 1); + } + WaitFlag(VEC1_MTE3_V_EVENT + 0); + WaitFlag(VEC1_MTE3_V_EVENT + 1); + WaitFlag(VEC1_V_MTE2_EVENT_QSCALE + 0); + WaitFlag(VEC1_V_MTE2_EVENT_QSCALE + 1); + + WaitFlag(TOPK_V_MTE2_EVENT); + WaitFlag(TOPK_MTE3_V_EVENT); + WaitFlag(V_MTE2_EVENT1); +} + +template +__aicore__ inline void QLIV2Vector::CleanInvalidOutput(int64_t invalidS1Offset) +{ + // init -1 and copy to output + uint64_t dealSize = constInfo_.sparseCount; + GlobalTensor indexOutput = indiceOutGm[invalidS1Offset]; + AscendC::InitGlobalMemory(indexOutput, dealSize, constInfo_.INVALID_IDX); + + if (returnValueFlag) { + WaitFlag(TOPK_MTE3_V_EVENT); + Duplicate(valueOutLocal_.template ReinterpretCast(), constInfo_.NEG_INF_BFLOAT, + constInfo_.sparseCount); + + SetFlag(TOPK_V_MTE3_EVENT); + WaitFlag(TOPK_V_MTE3_EVENT); + + AscendC::DataCopyParams copyOutValueParams; + copyOutValueParams.blockCount = 1; + copyOutValueParams.blockLen = constInfo_.sparseCount * sizeof(bfloat16_t); + copyOutValueParams.srcStride = 0; + copyOutValueParams.dstStride = 0; + AscendC::DataCopyPad(valueOutGm[invalidS1Offset], valueOutLocal_, copyOutValueParams); + SetFlag(TOPK_MTE3_V_EVENT); + } +} + +template +__aicore__ inline void QLIV2Vector::DoTndPadding(const QLIV2Common::RunInfo &runInfo) +{ + uint32_t paddingLen = runInfo.curCuSeqlensQ - runInfo.curSequsedQ; + uint64_t paddingOffset = + runInfo.indiceOutOffset + runInfo.curSequsedQ * constInfo_.kHeadNum * constInfo_.sparseCount; + uint64_t dealSize = paddingLen * constInfo_.kHeadNum * constInfo_.sparseCount; + GlobalTensor indiceOutPaddingStart = indiceOutGm[paddingOffset]; + AscendC::InitGlobalMemory(indiceOutPaddingStart, dealSize, constInfo_.INVALID_IDX); + + if (constInfo_.returnValue) { + SetFlag(MTE3_V_EVENT); + WaitFlag(MTE3_V_EVENT); + + GlobalTensor valueOutGmTmp; + valueOutGmTmp.SetGlobalBuffer((__gm__ uint16_t *)valueOutGm.GetPhyAddr()); + GlobalTensor valueOut = valueOutGmTmp[paddingOffset]; + + AscendC::InitGlobalMemory(valueOut, dealSize, constInfo_.NEG_INF_BFLOAT); + SetFlag(MTE3_V_EVENT); + WaitFlag(MTE3_V_EVENT); + } +} +template +__aicore__ inline void QLIV2Vector::GetKeyScale(const QLIV2Common::RunInfo &runInfo, + LocalTensor &kScaleUB, + int64_t batchId, int64_t startS2, + int64_t getLen) +{ + // startS2一定能整除kCacheBlockSize_ + AscendC::DataCopyPadExtParams padParams{false, 0, 0, 0}; + AscendC::DataCopyExtParams copyInParams; + if constexpr (PAGE_ATTENTION) { + int32_t startBlockTableIdx = startS2 / kCacheBlockSize_; + int32_t startBlockTableOffset = startS2 % kCacheBlockSize_; + int32_t blockTableBatchOffset = batchId * maxBlockNumPerBatch_; + copyInParams.blockCount = 1; + copyInParams.srcStride = 0; + copyInParams.dstStride = 0; + copyInParams.rsv = 0; + int32_t resUbBaseOffset = 0; + if (startBlockTableOffset > 0) { + int32_t firstPartLen = + kCacheBlockSize_ - startBlockTableOffset > getLen ? getLen : kCacheBlockSize_ - startBlockTableOffset; + copyInParams.blockLen = firstPartLen * sizeof(SCALE_T); + int32_t blockId = blockTableGm.GetValue(blockTableBatchOffset + startBlockTableIdx); + SetFlag(KSCALE_S_MTE2_EVENT); + WaitFlag(KSCALE_S_MTE2_EVENT); + AscendC::DataCopyPad(kScaleUB[16 * (runInfo.kScaleLoop % 2) * s2BaseSize_], + kScaleGm[blockId * constInfo_.keyDequantScaleStride0 + startBlockTableOffset], + copyInParams, padParams); + startBlockTableIdx++; + getLen = getLen - firstPartLen; + resUbBaseOffset = firstPartLen; + } + int32_t getLoopNum = CeilDiv(getLen, kCacheBlockSize_); + copyInParams.blockLen = kCacheBlockSize_ * sizeof(SCALE_T); + for (int32_t i = 0; i < getLoopNum; i++) { + if (i == getLoopNum - 1) { + copyInParams.blockLen = (getLen - i * kCacheBlockSize_) * sizeof(SCALE_T); + } + int32_t blockId = blockTableGm.GetValue(blockTableBatchOffset + startBlockTableIdx + i); + SetFlag(KSCALE_S_MTE2_EVENT); + WaitFlag(KSCALE_S_MTE2_EVENT); + AscendC::DataCopyPad( + kScaleUB[16 * (runInfo.kScaleLoop % 2) * s2BaseSize_ + resUbBaseOffset + i * kCacheBlockSize_], + kScaleGm[blockId * constInfo_.keyDequantScaleStride0], copyInParams, padParams); + } + } else { + copyInParams.blockCount = 1; + copyInParams.blockLen = getLen * sizeof(SCALE_T); + copyInParams.srcStride = 0; + copyInParams.dstStride = 0; + copyInParams.rsv = 0; + AscendC::DataCopyPad(kScaleUB[16 * (runInfo.kScaleLoop % 2) * s2BaseSize_], + kScaleGm[runInfo.tensorKeyScaleOffset], copyInParams, padParams); + } +} + +template +__aicore__ inline void QLIV2Vector::ProcessVec1(const QLIV2Common::RunInfo &info) +{ + auto pingpong = (info.loop % 2); + auto qScalepingpong = (info.qScaleLoop % 2); + auto kScalepingpong = (info.kScaleLoop % 2); + auto s1BaseSizePerAIV = CeilDiv(s1BaseSize_, 2); + int64_t curS1Idx = info.gS1Idx * s1BaseSize_; + int64_t curS2Idx = info.s2Idx * s2BaseSize_; + int64_t curS1ProcNum = curS1Idx + s1BaseSize_ > info.actS1Size ? info.actS1Size % s1BaseSize_ : s1BaseSize_; + int64_t curAivS1Idx = curS1Idx + (blockId_ % 2) * CeilDiv(curS1ProcNum, 2); + int64_t curAivS1ProcNum = (blockId_ % 2 == 0) ? CeilDiv(curS1ProcNum, 2) : curS1ProcNum / 2; + + if (curAivS1ProcNum == 0) { + CrossCoreWaitFlag( + QLIV2Common::ConstInfo::CROSS_CV_EVENT + pingpong); // V核等C核计算完mm1，mm1Res已搬运到UB + CrossCoreSetFlag( + QLIV2Common::ConstInfo::CROSS_VC_EVENT + pingpong); // V核处理完，通知C核可以把mm1Res搬运到UB + return; + } + + if (info.isFirstS2InnerLoop) { + WaitFlag(VEC1_V_MTE2_EVENT_QSCALE + qScalepingpong); + // weightsGm --> weightUB_ + int64_t weightGmOffset = info.tensorWeightsOffset + curAivS1Idx * kHeadNum_ * gSize_; + DataCopyPadExtParams padWeightsParams{false, 0, 0, 0}; + DataCopyExtParams qwDataCopyExtParams; + qwDataCopyExtParams.blockCount = curAivS1ProcNum; + qwDataCopyExtParams.blockLen = gSize_ * sizeof(WEIGHT_T); + qwDataCopyExtParams.srcStride = 0; + qwDataCopyExtParams.dstStride = (UB_BANK_DEPTH_STRIDE - qwDataCopyExtParams.blockLen) / 32; + DataCopyPad(weightUB_[qScalepingpong * (UB_BANK_STRIDE / sizeof(WEIGHT_T))], weightsGm[weightGmOffset], + qwDataCopyExtParams, padWeightsParams); + + if constexpr (!IS_MX) { + if (constInfo_.quantMode != 4) { + // qScaleGm --> qScaleUB_ + DataCopyPadExtParams padQScaleParams{false, 0, 0, 0}; + DataCopyExtParams qScaleDataCopyExtParams; + qScaleDataCopyExtParams.blockCount = curAivS1ProcNum; + qScaleDataCopyExtParams.blockLen = gSize_ * sizeof(SCALE_T); + qScaleDataCopyExtParams.srcStride = 0; + qScaleDataCopyExtParams.dstStride = (UB_BANK_DEPTH_STRIDE - qScaleDataCopyExtParams.blockLen) / 32; + DataCopyPad(qScaleUB_[qScalepingpong * (UB_BANK_STRIDE / sizeof(SCALE_T))], qScaleGm[weightGmOffset], + qScaleDataCopyExtParams, padQScaleParams); + } + } + + SetFlag(VEC1_MTE2_V_EVENT_QSCALE + qScalepingpong); + WaitFlag(VEC1_MTE2_V_EVENT_QSCALE + qScalepingpong); + } + + if constexpr (!IS_MX) { + if (((info.s2Idx - info.s2Start) % 16 == 0) && (constInfo_.quantMode != 4)) { + WaitFlag(VEC1_V_MTE2_EVENT_KSCALE + kScalepingpong); + uint32_t getLen = 16 * s2BaseSize_ > (info.validS2Len - info.s2Idx * s2BaseSize_) ? + info.validS2Len - info.s2Idx * s2BaseSize_ : + 16 * s2BaseSize_; + // kScaleGm --> kScaleUB_ + GetKeyScale(info, kScaleUB_, info.bIdx, curS2Idx, getLen); + SetFlag(VEC1_MTE2_V_EVENT_KSCALE + kScalepingpong); + WaitFlag(VEC1_MTE2_V_EVENT_KSCALE + kScalepingpong); + } + } + + WaitFlag(VEC1_MTE3_V_EVENT + pingpong); + + // CV同步 + CrossCoreWaitFlag( + QLIV2Common::ConstInfo::CROSS_CV_EVENT + info.loop % 2); // V核等C核计算完mm1，mm1Res已搬运到UB + + static_assert(std::is_same_v); + auto outBase = vec1OutUB_[pingpong * (UB_BANK_STRIDE / sizeof(SCORE_T))]; + auto weightBase = weightUB_[qScalepingpong * (UB_BANK_STRIDE / sizeof(WEIGHT_T))]; + auto weightTempBase = weightTempUB_[qScalepingpong * (UB_BANK_STRIDE / sizeof(float))]; + + auto qkBase = resMm1UB_[pingpong * (UB_BANK_STRIDE / sizeof(QK_T))]; + auto qkVLstride = (UB_BANK_DEPTH_STRIDE / sizeof(QK_T)) / 2 * constInfo_.mBaseSizeMax; + if constexpr (IS_MXFP4) { + vector1::BatchMulWeightAndReduceSumMXFP4(outBase, UB_BANK_DEPTH_STRIDE / sizeof(SCORE_T), qkBase, qkVLstride, + (uint32_t)(gSize_ * UB_BANK_DEPTH_STRIDE / sizeof(QK_T)), weightBase, + UB_BANK_DEPTH_STRIDE / sizeof(WEIGHT_T), gSize_, curAivS1ProcNum); + } else if constexpr (IS_MX) { + vector1::BatchMulWeightAndReduceSumMX(outBase, UB_BANK_DEPTH_STRIDE / sizeof(SCORE_T), qkBase, qkVLstride, + (uint32_t)(gSize_ * UB_BANK_DEPTH_STRIDE / sizeof(QK_T)), weightBase, + UB_BANK_DEPTH_STRIDE / sizeof(WEIGHT_T), weightTempBase, gSize_, + curAivS1ProcNum); + } else if constexpr (IS_WEIGHT_FP16) { + auto qScaleBase = qScaleUB_[qScalepingpong * (UB_BANK_STRIDE / sizeof(WEIGHT_T))]; + auto kScaleBase = kScaleUB_[kScalepingpong * 16 * s2BaseSize_ + + ((info.s2Idx - info.s2Start) % 16) * s2BaseSize_]; + vector1::BatchMulWeightAndReduceSum(outBase, UB_BANK_DEPTH_STRIDE / sizeof(SCORE_T), + qkBase, qkVLstride, + (uint32_t)(gSize_ * UB_BANK_DEPTH_STRIDE / sizeof(QK_T)), + weightBase, UB_BANK_DEPTH_STRIDE / sizeof(WEIGHT_T), weightTempBase, + kScaleBase, (uint32_t)0, + qScaleBase, UB_BANK_DEPTH_STRIDE / sizeof(SCALE_T), + gSize_, curAivS1ProcNum); + } else if (constInfo_.quantMode == 4) { // 4: per_tensor量化 + // quantMode为4时不适用sacle的UB + float kScaleValue = globalKScale_; + float qScaleValue = globalQScale_; + vector1::BatchMulWeightAndReduceSumPerTensor( + outBase, UB_BANK_DEPTH_STRIDE / sizeof(SCORE_T), qkBase, qkVLstride, + (uint32_t)(gSize_ * UB_BANK_DEPTH_STRIDE / sizeof(QK_T)), weightBase, + UB_BANK_DEPTH_STRIDE / sizeof(WEIGHT_T), weightTempBase, kScaleValue, qScaleValue, gSize_, curAivS1ProcNum); + } else if (constInfo_.quantMode == 1) { + auto qScaleBase = qScaleUB_[qScalepingpong * (UB_BANK_STRIDE / sizeof(float))]; + auto kScaleBase = + kScaleUB_[kScalepingpong * 16 * s2BaseSize_ + ((info.s2Idx - info.s2Start) % 16) * s2BaseSize_]; + vector1::BatchMulWeightAndReduceSum(outBase, UB_BANK_DEPTH_STRIDE / sizeof(SCORE_T), qkBase, qkVLstride, + (uint32_t)(gSize_ * UB_BANK_DEPTH_STRIDE / sizeof(QK_T)), weightBase, + UB_BANK_DEPTH_STRIDE / sizeof(WEIGHT_T), weightTempBase, kScaleBase, + (uint32_t)0, qScaleBase, UB_BANK_DEPTH_STRIDE / sizeof(float), gSize_, + curAivS1ProcNum); + } + if (info.isFirstS2InnerLoop) { + SetFlag(VEC1_V_MTE2_EVENT_QSCALE + qScalepingpong); + } + if constexpr (!IS_MX) { + if (((info.s2Idx - info.s2Start) % 16 == 0) && (constInfo_.quantMode != 4)) { + SetFlag(VEC1_V_MTE2_EVENT_KSCALE + kScalepingpong); + } + } + SetFlag(VEC1_V_MTE3_EVENT + pingpong); + WaitFlag(VEC1_V_MTE3_EVENT + pingpong); + // outUB_ ---> scoreGm + int64_t vec1OutGmOffset = + blockId_ % 2 == 0 ? + curS2Idx : + s1BaseSizePerAIV * QLIV2Common::Align((uint64_t)constInfo_.kSeqSize, (uint64_t)s2BaseSize_) + curS2Idx; + DataCopyExtParams copyOutParams; + copyOutParams.blockCount = curAivS1ProcNum; + copyOutParams.blockLen = s2BaseSize_ * sizeof(SCORE_T); + copyOutParams.srcStride = (UB_BANK_DEPTH_STRIDE - copyOutParams.blockLen) / 32; + copyOutParams.dstStride = + (QLIV2Common::Align((uint64_t)constInfo_.kSeqSize, (uint64_t)s2BaseSize_) - s2BaseSize_) * sizeof(SCORE_T); + DataCopyPad(scoreGm[vec1OutGmOffset], outBase, copyOutParams); + SetFlag(VEC1_MTE3_V_EVENT + pingpong); + CrossCoreSetFlag( + QLIV2Common::ConstInfo::CROSS_VC_EVENT + pingpong); // V核处理完，通知C核可以把mm1Res搬运到UB +} + +template +__aicore__ inline void QLIV2Vector::ProcessTopK(const QLIV2Common::RunInfo &info) +{ + SetFlag(MTE3_MTE2_EVENT); + WaitFlag(MTE3_MTE2_EVENT); + + int64_t curS1Idx = info.gS1Idx * s1BaseSize_; + int64_t curS2StartIdx = info.s2Start * s2BaseSize_; + int64_t curS2EndIdx = (info.s2LoopEnd + 1) * s2BaseSize_; + int64_t curS1ProcNum = curS1Idx + s1BaseSize_ > info.actS1Size ? info.actS1Size % s1BaseSize_ : s1BaseSize_; + int64_t curAivS1Idx = curS1Idx + (blockId_ % 2) * CeilDiv(curS1ProcNum, 2); + int64_t curAivS1ProcNum = (blockId_ % 2 == 0) ? CeilDiv(curS1ProcNum, 2) : curS1ProcNum / 2; + + // LD 需要搬运到的起始位置 32字节 + int64_t ldDstOffset = (info.isNeedLD) ? curS2StartIdx : 0; + AscendC::DataCopyExtParams copyInParams; + copyInParams.blockCount = 1; + copyInParams.srcStride = 0; + copyInParams.dstStride = 0; + copyInParams.rsv = 0; + + AscendC::DataCopyParams copyOutParams; + copyOutParams.blockCount = 1; + copyOutParams.blockLen = topkCount_ * sizeof(uint32_t); // bytes + copyOutParams.srcStride = 0; + copyOutParams.dstStride = 0; + + AscendC::DataCopyParams ldCopyOutParams; // ldIndicesCopyOutParams + ldCopyOutParams.blockCount = 1; + ldCopyOutParams.blockLen = topkCountAlign16_ * sizeof(uint32_t); // bytes + ldCopyOutParams.srcStride = 0; + ldCopyOutParams.dstStride = 0; + + AscendC::DataCopyParams ldCopyScoreOutParams; + ldCopyScoreOutParams.blockCount = 1; + ldCopyScoreOutParams.blockLen = topkCountAlign16_ * sizeof(SCORE_T); // bytes + ldCopyScoreOutParams.srcStride = 0; + ldCopyScoreOutParams.dstStride = 0; + + int32_t cuRealAcSeq = info.actS2Size; + if (constInfo_.attenMaskFlag) { + cuRealAcSeq = info.actS2SizeOrig - info.actS1Size + curAivS1Idx + 1; + } + + int32_t validAllS2Len = cuRealAcSeq; + for (uint32_t i = 0; i < curAivS1ProcNum; i++) { + if (i > 0) { + SetFlag(MTE3_MTE2_EVENT); + WaitFlag(MTE3_MTE2_EVENT); + } + uint32_t rowIdx = blockId_ % 2 * CeilDiv(curS1ProcNum, 2) + i; + uint32_t vecOffset = blockId_ % 2 * CeilDiv(s1BaseSize_, 2) + i; + int64_t outputIdxOffset = 0; + if (info.isOutputIdxOffsetValid) { + outputIdxOffset = outputIdxOffsetGm.GetValue(info.outputIdxCoreOffset + rowIdx * kHeadNum_); + } + + SCORE_T zero = 0; + int32_t neg = -1; + uint32_t scoreAlign = sizeof(SCORE_T) == 4 ? 8 : 16; // int32对应UB对齐数为8，int16需要16 + if (constInfo_.attenMaskFlag) { + validAllS2Len = ((int32_t)i + cuRealAcSeq) / static_cast(constInfo_.cmpRatio); + } + int32_t validS2Len = validAllS2Len; + if (info.isNeedLD) { + // 当前核处理的s2长度validS2Len + validS2Len = Min((info.s2LoopEnd + 1) * s2BaseSize_, validAllS2Len) - curS2StartIdx; + } + + uint64_t offset = info.saveWorkSpaceIdx * s1BaseSize_ * topkCountAlign16_ + rowIdx * topkCountAlign16_; + if (validS2Len <= 0 && !info.isNeedLD) { + WaitFlag(TOPK_MTE3_V_EVENT); + Duplicate(indicesOutLocal_.ReinterpretCast(), neg, topkCount_); + SetFlag(TOPK_V_MTE3_EVENT); + WaitFlag(TOPK_V_MTE3_EVENT); + AscendC::DataCopyPad(indiceOutGm[info.indiceOutOffset + (curS1Idx + rowIdx) * topkCount_], + indicesOutLocal_.ReinterpretCast(), copyOutParams); + SetFlag(TOPK_MTE3_V_EVENT); + if (returnValueFlag) { + WaitFlag(TOPK_MTE3_V_EVENT); + + Duplicate(valueOutLocal_.template ReinterpretCast(), constInfo_.NEG_INF_BFLOAT, topkCount_); + + SetFlag(TOPK_V_MTE3_EVENT); + WaitFlag(TOPK_V_MTE3_EVENT); + + AscendC::DataCopyParams copyOutValueParams; + copyOutValueParams.blockCount = 1; + copyOutValueParams.blockLen = topkCount_ * sizeof(bfloat16_t); + copyOutValueParams.srcStride = 0; + copyOutValueParams.dstStride = 0; + AscendC::DataCopyPad(valueOutGm[info.valueOutOffset + (curS1Idx + rowIdx) * topkCount_], valueOutLocal_, + copyOutValueParams); + SetFlag(TOPK_MTE3_V_EVENT); + } + continue; + } else if (validS2Len <= 0 && info.isNeedLD) { + WaitFlag(TOPK_MTE3_V_EVENT); + Duplicate(indicesOutLocal_.ReinterpretCast(), neg, topkCountAlign16_); + Duplicate(scoreOutLocal_, zero, topkCountAlign16_); + SetFlag(TOPK_V_MTE3_EVENT); + WaitFlag(TOPK_V_MTE3_EVENT); + // 将每一行S1的Topk结果存放在ldScoreGm和ldIndexGm中 + AscendC::DataCopyPad(ldScoreGm[offset], scoreOutLocal_, ldCopyScoreOutParams); + AscendC::DataCopyPad(ldIndexGm[offset], indicesOutLocal_.ReinterpretCast(), ldCopyOutParams); + SetFlag(TOPK_MTE3_V_EVENT); + continue; + } + + WaitFlag(TOPK_V_MTE2_EVENT); + WaitFlag(TOPK_MTE3_V_EVENT); + + AscendC::DataCopyPadExtParams padParams{true, 0, 0, 0}; + if (validS2Len >= topkCount_) { + uint32_t s2LoopNum = (validS2Len + trunkLen_ - 1) / trunkLen_; + bool useSingleLoop = + (s2LoopNum == 1) || ((topkCount_ > trunkLen_) && (validS2Len <= (uint32_t)topkCountAlign256_)); + if (useSingleLoop) { + uint32_t validS2LenAlign = QLIV2Common::Align(validS2Len, (int32_t)256); + Duplicate(mrgValueLocal_[validS2Len / 256 * 256], zero, validS2LenAlign - validS2Len / 256 * 256); + SetFlag(V_MTE2_EVENT3); + WaitFlag(V_MTE2_EVENT3); + copyInParams.blockLen = validS2Len * sizeof(SCORE_T); // byte + AscendC::DataCopyPadExtParams padParams{true, 0, 0, 0}; + AscendC::DataCopyPad( + mrgValueLocal_, + scoreGm[vecOffset * QLIV2Common::Align((uint64_t)constInfo_.kSeqSize, (uint64_t)s2BaseSize_) + + ldDstOffset], + copyInParams, padParams); + SetFlag(TOPK_MTE2_V_EVENT); + WaitFlag(TOPK_MTE2_V_EVENT); + topkOp_.TopK(mrgValueLocal_, indicesOutLocal_, scoreOutLocal_, validS2LenAlign, 0, 1, info.isNeedLD, + returnValueFlag, outputIdxOffset); + } else { + uint32_t outputIdxOffsetTmp = 0; + uint32_t actS2LoopNum = 0; + if (topkCount_ > trunkLen_) { + actS2LoopNum = 1 + (validS2Len - topkCountAlign256_ + trunkLen_ - 1) / trunkLen_; + } else { + actS2LoopNum = (validS2Len + trunkLen_ - 1) / trunkLen_; + } + for (uint32_t loopIdx = 0; loopIdx < actS2LoopNum; loopIdx++) { + if (loopIdx == actS2LoopNum - 1) { + outputIdxOffsetTmp = outputIdxOffset; + } + if (loopIdx == 0) { + if (topkCount_ > trunkLen_) { + copyInParams.blockLen = topkCountAlign256_ * sizeof(SCORE_T); // byte + AscendC::DataCopyPad(scoreOutLocal_, + scoreGm[vecOffset * QLIV2Common::Align((uint64_t)constInfo_.kSeqSize, + (uint64_t)s2BaseSize_) + + ldDstOffset], + copyInParams, padParams); + SetFlag(TOPK_MTE2_V_EVENT); + WaitFlag(TOPK_MTE2_V_EVENT); + AscendC::CreateVecIndex(indicesOutLocal_.ReinterpretCast(), (int32_t)zero, + topkCountAlign256_); + AscendC::CreateVecIndex(topkSharedTmpLocal_.ReinterpretCast(), (int32_t)zero, + topkCountAlign256_); + } else { + copyInParams.blockLen = trunkLen_ * sizeof(SCORE_T); // byte + AscendC::DataCopyPad(mrgValueLocal_, + scoreGm[vecOffset * QLIV2Common::Align((uint64_t)constInfo_.kSeqSize, + (uint64_t)s2BaseSize_) + + ldDstOffset], + copyInParams, padParams); + SetFlag(TOPK_MTE2_V_EVENT); + WaitFlag(TOPK_MTE2_V_EVENT); + topkOp_.TopK(mrgValueLocal_, indicesOutLocal_, scoreOutLocal_, trunkLen_, loopIdx, + actS2LoopNum, info.isNeedLD, returnValueFlag, outputIdxOffsetTmp); + } + + continue; + } + SetFlag(V_MTE2_EVENT2); + WaitFlag(V_MTE2_EVENT2); + uint32_t validTrunkLen = 0; + uint32_t offset = 0; + if (topkCount_ > trunkLen_) { + validTrunkLen = (topkCountAlign256_ + (loopIdx - 1) * trunkLen_ + trunkLen_) > validS2Len ? + (validS2Len - topkCountAlign256_) % trunkLen_ : + trunkLen_; + offset = vecOffset * QLIV2Common::Align((uint64_t)constInfo_.kSeqSize, (uint64_t)s2BaseSize_) + + topkCountAlign256_ + (loopIdx - 1) * trunkLen_ + ldDstOffset; + } else { + validTrunkLen = + (loopIdx * trunkLen_ + trunkLen_) > validS2Len ? validS2Len % trunkLen_ : trunkLen_; + offset = vecOffset * QLIV2Common::Align((uint64_t)constInfo_.kSeqSize, (uint64_t)s2BaseSize_) + + loopIdx * trunkLen_ + ldDstOffset; + } + + AscendC::DataCopy(mrgValueLocal_, scoreOutLocal_, topkCountAlign256_); + // topk如果没有对齐到256，则把topkCountAlign256_ - topkCount_部分刷0 + // 如果topk > trunkLen，第一轮调用topk是直接拷贝的，不需要刷0 + bool isZeroPadding = (topkCount_ > trunkLen_) ? (loopIdx > 1) : true; + if (topkCountAlign256_ != topkCount_ && isZeroPadding) { + uint64_t mask[2] = {0, 0}; // 0: mask初始化，代表对应位置不需要刷0 + mask[0] = ~0; + mask[0] = mask[0] << (topkCount_ % 64); + PipeBarrier(); + // 把topkCount_对齐到64刷0，此处由于duplicate的限制mask[0]刷64个数 + Duplicate(mrgValueLocal_[topkCount_ / 64 * 64], zero, mask, 1, 1, 0); + PipeBarrier(); + // 把topk剩余对齐到256的部分刷0 + Duplicate(mrgValueLocal_[topkCount_ / 64 * 64 + 64], zero, + topkCountAlign256_ - (topkCount_ / 64 * 64 + 64)); + SetFlag(V_MTE2_EVENT3); + WaitFlag(V_MTE2_EVENT3); + } + copyInParams.blockLen = validTrunkLen * sizeof(SCORE_T); // byte + // TOPK 直方图一次必须计算256，输入处理数据需要和256对齐 + if ((topkCountAlign256_ + validTrunkLen) % 256 != 0) { + Duplicate(mrgValueLocal_[topkCountAlign256_ + validTrunkLen / 256 * 256], zero, + QLIV2Common::Align(validTrunkLen, (uint32_t)256) - validTrunkLen / 256 * 256); + SetFlag(V_MTE2_EVENT3); + WaitFlag(V_MTE2_EVENT3); + } + WaitFlag(V_MTE2_EVENT1); + AscendC::DataCopyPad(mrgValueLocal_[topkCountAlign256_], scoreGm[offset], copyInParams, padParams); + SetFlag(TOPK_MTE2_V_EVENT); + WaitFlag(TOPK_MTE2_V_EVENT); + topkOp_.TopK(mrgValueLocal_, indicesOutLocal_, scoreOutLocal_, + QLIV2Common::Align(topkCountAlign256_ + validTrunkLen, (uint32_t)256), loopIdx, + actS2LoopNum, info.isNeedLD, returnValueFlag, outputIdxOffsetTmp); + SetFlag(V_MTE2_EVENT1); + } + } + } else { + AscendC::CreateVecIndex(indicesOutLocal_.ReinterpretCast(), (int32_t)zero, validS2Len); + if (outputIdxOffset != 0) { + topkb16gather::IndicesAddOffset(indicesOutLocal_, outputIdxOffset, constInfo_.sparseCount); + } + // LD或returnValue场景需要将score搬入scoreOutLocal_ + if (info.isNeedLD || returnValueFlag) { + copyInParams.blockLen = QLIV2Common::Align(validS2Len, (int32_t)32) * sizeof(SCORE_T); + AscendC::DataCopyPad( + scoreOutLocal_, + scoreGm[vecOffset * QLIV2Common::Align((uint64_t)constInfo_.kSeqSize, (uint64_t)s2BaseSize_) + + ldDstOffset], + copyInParams, padParams); + SetFlag(TOPK_MTE2_V_EVENT); + WaitFlag(TOPK_MTE2_V_EVENT); + } + } + + if (!info.isNeedLD) { + if (validS2Len < topkCount_) { + uint64_t mask[2] = {0, 0}; + mask[0] = ~0; + mask[0] = mask[0] << (validS2Len % 8); // 将`mask[0]`左移`validS2Len % 8`位 对齐 + PipeBarrier(); + // repeatTime 每次读取连续的8个datablock（每个block32Bytes，共256Bytes） + // dstBlockStride 单次迭代内，矢量目的操作数不同datablock间地址步长。 + Duplicate(indicesOutLocal_.ReinterpretCast()[validS2Len / 8 * 8], neg, mask, 1, 1, 0); + } + + if (validS2Len / 8 * 8 + 64 < topkCount_) { + PipeBarrier(); + Duplicate(indicesOutLocal_.ReinterpretCast()[validS2Len / 8 * 8 + 64], neg, + topkCount_ - (validS2Len / 8 * 8 + 64)); + } + SetFlag(TOPK_V_MTE2_EVENT); + if (returnValueFlag) { + WaitFlag(TOPK_V_MTE2_EVENT); + vector1::UIntToFloatReturnValue(valueOutLocal_, scoreOutLocal_, topkCountAlign256_); + // 无效值刷-inf + if (validS2Len < topkCount_) { + // mask[1]=0：bit模式高64位不参与，第一段只刷低64个bf16元素，剩余由下方count版Duplicate补刷 + uint64_t mask[2] = {0, 0}; + mask[0] = ~0; + mask[0] = mask[0] << (validS2Len % 16); + PipeBarrier(); + Duplicate(valueOutLocal_.template ReinterpretCast()[validS2Len / 16 * 16], + constInfo_.NEG_INF_BFLOAT, mask, 1, 1, 0); + } + + if (validS2Len / 16 * 16 + 64 < topkCount_) { + PipeBarrier(); + Duplicate(valueOutLocal_.template ReinterpretCast()[validS2Len / 16 * 16 + 64], + constInfo_.NEG_INF_BFLOAT, topkCount_ - (validS2Len / 16 * 16 + 64)); + } + SetFlag(TOPK_V_MTE2_EVENT); + } + + SetFlag(TOPK_V_MTE3_EVENT); + WaitFlag(TOPK_V_MTE3_EVENT); + AscendC::DataCopyPad(indiceOutGm[info.indiceOutOffset + (curS1Idx + rowIdx) * topkCount_], + indicesOutLocal_.ReinterpretCast(), copyOutParams); + if (returnValueFlag) { + AscendC::DataCopyParams copyOutValueParams; + copyOutValueParams.blockCount = 1; + copyOutValueParams.blockLen = topkCount_ * sizeof(bfloat16_t); // bytes + copyOutValueParams.srcStride = 0; + copyOutValueParams.dstStride = 0; + AscendC::DataCopyPad(valueOutGm[info.valueOutOffset + (curS1Idx + rowIdx) * topkCount_], valueOutLocal_, + copyOutValueParams); + } + SetFlag(TOPK_MTE3_V_EVENT); + } else { + PipeBarrier(); + AscendC::Adds(indicesOutLocal_.ReinterpretCast(), indicesOutLocal_.ReinterpretCast(), + static_cast(curS2StartIdx), topkCountAlign16_); + if (validS2Len < topkCount_) { + uint64_t mask[2] = {0, 0}; + mask[0] = ~0; + mask[0] = mask[0] << (validS2Len % scoreAlign); + PipeBarrier(); + Duplicate(scoreOutLocal_[validS2Len / scoreAlign * scoreAlign], zero, mask, 1, 1, 0); + uint64_t maskI[2] = {0, 0}; + maskI[0] = ~0; + maskI[0] = maskI[0] << (validS2Len % 8); + PipeBarrier(); + Duplicate(indicesOutLocal_.ReinterpretCast()[validS2Len / 8 * 8], neg, maskI, 1, 1, 0); + } + if (validS2Len / scoreAlign * scoreAlign + 64 < topkCount_) { + PipeBarrier(); + Duplicate(scoreOutLocal_[validS2Len / scoreAlign * scoreAlign + 64], zero, + topkCount_ - (validS2Len / scoreAlign * scoreAlign + 64)); + } + if (validS2Len / 8 * 8 + 64 < topkCount_) { + PipeBarrier(); + Duplicate(indicesOutLocal_.ReinterpretCast()[validS2Len / 8 * 8 + 64], neg, + topkCount_ - (validS2Len / 8 * 8 + 64)); + } + if (topkCountAlign16_ != topkCount_) { + uint64_t mask[2] = {0, 0}; + mask[0] = ~0; + mask[0] = mask[0] << (topkCount_ % scoreAlign); + PipeBarrier(); + Duplicate(scoreOutLocal_[topkCount_ / scoreAlign * scoreAlign], zero, mask, 1, 1, 0); + uint64_t maskIndices[2] = {0, 0}; + maskIndices[0] = ~0; + maskIndices[0] = maskIndices[0] << (topkCount_ % 8); + PipeBarrier(); + Duplicate(indicesOutLocal_.ReinterpretCast()[topkCount_ / 8 * 8], neg, maskIndices, 1, 1, 0); + } + + SetFlag(TOPK_V_MTE2_EVENT); + SetFlag(TOPK_V_MTE3_EVENT); + WaitFlag(TOPK_V_MTE3_EVENT); + // 将每一行S1的Topk结果存放在ldScoreGm和ldIndexGm中 + AscendC::DataCopyPad(ldScoreGm[offset], scoreOutLocal_, ldCopyScoreOutParams); + AscendC::DataCopyPad(ldIndexGm[offset], indicesOutLocal_.ReinterpretCast(), ldCopyOutParams); + SetFlag(TOPK_MTE3_V_EVENT); + } + } +} + +template +__aicore__ inline void QLIV2Vector::ProcessLD() +{ + AscendC::DataCopyParams copyOutParams; + copyOutParams.blockCount = 1; + copyOutParams.blockLen = topkCount_ * sizeof(uint32_t); // bytes + copyOutParams.srcStride = 0; + copyOutParams.dstStride = 0; + + AscendC::DataCopyParams copyOutValueParams; + copyOutValueParams.blockCount = 1; + copyOutValueParams.blockLen = topkCount_ * sizeof(bfloat16_t); + copyOutValueParams.srcStride = 0; + copyOutValueParams.dstStride = 0; + + uint32_t copyBytes = topkCountAlign256_; // 32B对齐 topkCountAlign256_ + + AscendC::DataCopyPadExtParams scorePadParams{true, 0, 0, 0}; + AscendC::DataCopyPadExtParams indexPadParams{true, 0, 0, 0}; + AscendC::DataCopyExtParams ldScoreParams; + ldScoreParams.blockLen = topkCountAlign16_ * sizeof(SCORE_T); // bytes + ldScoreParams.srcStride = (s1BaseSize_ - 1) * topkCountAlign16_ * sizeof(SCORE_T); // 两个基本块之间的距离 + ldScoreParams.dstStride = 0; + AscendC::DataCopyExtParams ldIndexParams; + ldIndexParams.blockLen = topkCountAlign16_ * sizeof(uint32_t); // bytes + ldIndexParams.srcStride = (s1BaseSize_ - 1) * topkCountAlign16_ * sizeof(uint32_t); + ldIndexParams.dstStride = 0; + + uint64_t ldProcessLen = ldInfo_.workspaceNum * topkCountAlign16_; + if (ldProcessLen <= 0) { + return; + } + uint64_t mrgValueLen = trunkLenLd_ + topkCountAlign256_; + uint32_t ldWorkspaceNum = (ldProcessLen > mrgValueLen) ? (trunkLenLd_ / topkCountAlign16_) : ldInfo_.workspaceNum; + // 搬运次数，一次搬运ldworkspaceNum块 + uint32_t ldProcessNum = CeilDiv(ldInfo_.workspaceNum, ldWorkspaceNum); + uint32_t ldProcessOffset = 0; + + SCORE_T zero = 0; + int32_t neg = -1; + uint32_t zero32 = 0; + uint32_t scoreAlign = sizeof(SCORE_T) == 4 ? 8 : 16; // int32对应UB对齐数为8，int16需要16 + uint32_t ldProWorkspaceNum = ldInfo_.workspaceNum; + SetFlag(TOPK_MTE3_V_EVENT); + for (uint32_t j = 0; j < ldInfo_.mNum; j++) { + WaitFlag(TOPK_MTE3_V_EVENT); + for (uint32_t i = 0; i < ldProcessNum; i++) { + // 读取数据段过长 + if (ldProcessNum > 1) { + ldProWorkspaceNum = (i == ldProcessNum - 1) ? ldInfo_.workspaceNum - i * ldWorkspaceNum : + ldWorkspaceNum; // 当前搬运块数 + } + ldProcessOffset = (i != 0) ? topkCountAlign256_ : 0; + PipeBarrier(); + // 索引全部刷-1 value全部刷0 + Duplicate(mrgValueLocal_[ldProcessOffset], zero, topkCountAlign256_ + trunkLenLd_ - ldProcessOffset); + Duplicate(ldIndexLocal_.ReinterpretCast()[ldProcessOffset], neg, + topkCountAlign256_ + trunkLenLd_ - ldProcessOffset); + + ldScoreParams.blockCount = ldProWorkspaceNum; + ldIndexParams.blockCount = ldProWorkspaceNum; + + int32_t s2Len = topkCountAlign16_ * ldProWorkspaceNum + ldProcessOffset; + uint32_t s2LenAlign = QLIV2Common::Align(s2Len, (int32_t)256); // 寄存器需要256对齐 + uint64_t LDGmOffset = ldInfo_.workspaceIdx * s1BaseSize_ * topkCountAlign16_ + + topkCountAlign16_ * (ldInfo_.mStart + j) + + i * ldWorkspaceNum * s1BaseSize_ * topkCountAlign16_; // 加入LD处理偏移 + SetFlag(TOPK_V_MTE2_EVENT); + WaitFlag(TOPK_V_MTE2_EVENT); + AscendC::DataCopyPad(mrgValueLocal_[ldProcessOffset], ldScoreGm[LDGmOffset], ldScoreParams, scorePadParams); + AscendC::DataCopyPad(ldIndexLocal_[ldProcessOffset], ldIndexGm[LDGmOffset].ReinterpretCast(), + ldIndexParams, indexPadParams); + + SetFlag(TOPK_MTE2_V_EVENT); + WaitFlag(TOPK_MTE2_V_EVENT); + + // 对非对齐索引和值刷0 -1 + if (s2LenAlign != s2Len) { + uint64_t mask[2] = {0, 0}; + mask[0] = ~0; + mask[0] = mask[0] << (s2Len % 64); + Duplicate(mrgValueLocal_[s2Len / 64 * 64], zero, mask, 1, 1, 0); + // 把s2Len对齐到64刷0，此处由于duplicate的限制mask[0]刷64个数 + Duplicate(ldIndexLocal_.ReinterpretCast()[s2Len / 64 * 64], neg, mask, 1, 1, 0); + } + if (s2Len / 64 * 64 + 64 < s2LenAlign) { + PipeBarrier(); + Duplicate(mrgValueLocal_[s2Len / 64 * 64 + 64], zero, s2LenAlign - (s2Len / 64 * 64 + 64)); + PipeBarrier(); + Duplicate(ldIndexLocal_.ReinterpretCast()[s2Len / 64 * 64 + 64], neg, + s2LenAlign - (s2Len / 64 * 64 + 64)); + } + + PipeBarrier(); + topkOp_.LdTopK(mrgValueLocal_, ldIndexLocal_, indicesOutLocal_, scoreOutLocal_, s2LenAlign, j, + ldProcessNum); + if (topkCountAlign256_ != topkCount_) { + uint64_t mask[2] = {0, 0}; + mask[0] = ~0; + mask[0] = mask[0] << (topkCount_ % 64); + PipeBarrier(); + Duplicate(scoreOutLocal_[topkCount_ / 64 * 64], zero, mask, 1, 1, 0); + uint64_t maskIndices[2] = {0, 0}; + maskIndices[0] = ~0; + maskIndices[0] = maskIndices[0] << (topkCount_ % 64); + PipeBarrier(); + Duplicate(indicesOutLocal_.ReinterpretCast()[topkCount_ / 64 * 64], neg, maskIndices, 1, 1, 0); + } + if (topkCount_ / 64 * 64 + 64 < topkCountAlign256_) { + PipeBarrier(); + Duplicate(scoreOutLocal_[topkCount_ / 64 * 64 + 64], zero, + topkCountAlign256_ - (topkCount_ / 64 * 64 + 64)); + PipeBarrier(); + Duplicate(indicesOutLocal_.ReinterpretCast()[topkCount_ / 64 * 64 + 64], neg, + topkCountAlign256_ - (topkCount_ / 64 * 64 + 64)); + } + + PipeBarrier(); + AscendC::DataCopy(ldIndexLocal_, indicesOutLocal_, copyBytes); + AscendC::DataCopy(mrgValueLocal_, scoreOutLocal_, copyBytes); + } + if (returnValueFlag) { + PipeBarrier(); + // 可排序键还原为 bf16，并将无效位（score==0，对应 index==-1）刷为 -inf + // Convert sortable key to bf16 and flush invalid (score==0, index==-1) positions to -inf + vector1::UIntToFloatReturnValueWithInfMask(valueOutLocal_, scoreOutLocal_, topkCountAlign256_, + constInfo_.NEG_INF_BFLOAT); + } + SetFlag(TOPK_V_MTE3_EVENT); + WaitFlag(TOPK_V_MTE3_EVENT); + uint64_t indiceOutGmOffset = + ldInfo_.indiceOutCoreOffset + (ldInfo_.mStart + j) * constInfo_.kHeadNum * topkCount_; + AscendC::DataCopyPad(indiceOutGm[indiceOutGmOffset], indicesOutLocal_.ReinterpretCast(), + copyOutParams); + if (returnValueFlag) { + AscendC::DataCopyPad(valueOutGm[indiceOutGmOffset], valueOutLocal_, copyOutValueParams); + } + SetFlag(TOPK_MTE3_V_EVENT); + } + WaitFlag(TOPK_MTE3_V_EVENT); +} +} // namespace QLIV2Kernel +#endif // QUANT_LIGHTNING_INDEXER_V2_SERVICE_VECTOR_H diff --git a/xllm_ops/attention/quant_lightning_indexer_v2/op_kernel/arch35/vf/quant_lightning_indexer_v2_topk.h b/xllm_ops/attention/quant_lightning_indexer_v2/op_kernel/arch35/vf/quant_lightning_indexer_v2_topk.h new file mode 100644 index 0000000..c715bd3 --- /dev/null +++ b/xllm_ops/attention/quant_lightning_indexer_v2/op_kernel/arch35/vf/quant_lightning_indexer_v2_topk.h @@ -0,0 +1,221 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file quant_lightning_indexer_v2_topk.h + * \brief + */ +#ifndef QUANT_LIGHTNING_INDEXER_V2_TOPK_H +#define QUANT_LIGHTNING_INDEXER_V2_TOPK_H + +#include "kernel_operator.h" +#include "vf_topk.h" +#include "vf_topk_16_gather_quant_v2.h" + +namespace topk { +template +class LITopk { +public: + __aicore__ inline void operator()(LocalTensor& outputIdxLocal, + LocalTensor& inputLocal, + uint32_t s2SeqLen) + { + } +}; + +template<> +class LITopk { +public: + static __aicore__ inline uint32_t GetSharedTmpBufferSize(uint32_t topK) + { + return 2 * topK * sizeof(uint32_t) + 5 * 256 * sizeof(uint32_t) + 64 * sizeof(uint32_t) + + (topK + 64) * sizeof(uint32_t); // for output value tensor + } + + static __aicore__ inline uint32_t GetIndexBufferSize(uint32_t topK) + { + return (topK + 64) * sizeof(uint32_t); + } + + __aicore__ inline void Init(uint32_t topK) + { + this->topK = topK; + } + + __aicore__ inline void InitBuffers(LocalTensor& sharedTmpBuffer) + { + tmpIdxLocal = sharedTmpBuffer[0]; + tmpValueLocal = tmpIdxLocal[topK]; + histogramsLocal = tmpValueLocal[topK]; + idx0Local = histogramsLocal[256]; + idx1Local = idx0Local[256]; + idx2Local = idx1Local[256]; + idx3Local = idx2Local[256]; + nkValueLocal = idx3Local[256]; + outputValueLocal = nkValueLocal[64]; + } + + __aicore__ inline void operator()(LocalTensor& outputIdxLocal, + LocalTensor& inputLocal, + uint32_t s2SeqLen) + { + topkb32::LiTopKVF(outputIdxLocal, // filter阶段使用输出value Buf topK * 4B + outputValueLocal, // filter阶段使用输出 Idx Buf topK * 4B + inputLocal, // 输入 s2SeqLen * 4B + tmpIdxLocal, // filter阶段使用暂存index Buf topK * 4B + tmpValueLocal, // filter阶段使用暂存value Buf topK * 4B + histogramsLocal, // 直方图的临时Buf 256 * 4B + idx0Local, // 输入数据第1个8位Buf 256 * 4B + idx1Local, // 输入数据第2个8位Buf 256 * 4B + idx2Local, // 输入数据第3个8位Buf 256 * 4B + idx3Local, // 输入数据第4个8位Buf 256 * 4B + nkValueLocal, // next_k 暂存Buf 64 * 4B + topK, // topk数量 + s2SeqLen); // 输入元素总数 + } +private: + LocalTensor tmpIdxLocal; // filter阶段使用暂存index Buf topK * 4B + LocalTensor tmpValueLocal; // filter阶段使用暂存value Buf topK * 4B + LocalTensor histogramsLocal; // 直方图的临时Buf 256 * 4B + LocalTensor idx0Local; // 输入数据第1个8位Buf 256 * 4B + LocalTensor idx1Local; // 输入数据第2个8位Buf 256 * 4B + LocalTensor idx2Local; // 输入数据第3个8位Buf 256 * 4B + LocalTensor idx3Local; // 输入数据第4个8位Buf 256 * 4B + LocalTensor nkValueLocal; // next_k 暂存Buf 64 * 4B + LocalTensor outputValueLocal; // 输出value tensor + uint32_t topK; +}; + +template<> +class LITopk { +public: + __aicore__ inline uint32_t GetSharedTmpBufferSize() + { + // 2 * QLIV2Common::Align(topK, (uint32_t)256): 两块hisIndexLocal; + // 3 * 256: histogramsLocal idxHighLocal idxLowLocal; 64: nkValueLocal + uint64_t bufferSize1 = (2 * QLIV2Common::Align(topK, (uint32_t)256) + 3 * 256 + 64) * sizeof(uint32_t); + // QLIV2Common::Align(topK, (uint32_t)256) + trunkLen：tmpIndexLocal + uint64_t bufferSize2 = (QLIV2Common::Align(topK, (uint32_t)256) + trunkLen) * sizeof(uint16_t); + uint64_t reuseBufferSize = QLIV2Common::Align(topK, (uint32_t)256) * sizeof(uint32_t); + return bufferSize1 + bufferSize2 - reuseBufferSize; + } + + __aicore__ inline void Init(uint32_t topK, uint32_t trunkLen) + { + this->topK = topK; + this->trunkLen = trunkLen; + } + + __aicore__ inline void InitBuffers(LocalTensor& sharedTmpBuffer, LocalTensor& indicesOutLocal) + { + LocalTensor hisIndexLocal1 = indicesOutLocal; + LocalTensor hisIndexLocal2 = sharedTmpBuffer[0]; + hisIndexLocal[0] = hisIndexLocal1; + hisIndexLocal[1] = hisIndexLocal2; + histogramsLocal = hisIndexLocal2[QLIV2Common::Align(topK, (uint32_t)256)]; + idxHighLocal = histogramsLocal[256]; + idxLowLocal = idxHighLocal[256]; + nkValueLocal = idxLowLocal[256]; + LocalTensor tmpIndexLocalTmp = nkValueLocal[64]; + tmpIndexLocal = tmpIndexLocalTmp.template ReinterpretCast(); + } + + __aicore__ inline void TopK(LocalTensor& mrgValueLocal, LocalTensor& indicesOutLocal, + LocalTensor& hisValueLocal, uint32_t s2SeqLen, uint32_t loopIdx, + uint32_t s2LoopNum, bool isNeedLD, bool returnValueFlag, + uint32_t outputIdxOffset) + { + // true: 开启返回hisValueLocal + if (s2LoopNum == 1) { + if (isNeedLD || returnValueFlag) { + topkb16gather::LiTopKVF(tmpIndexLocal, hisValueLocal, mrgValueLocal, histogramsLocal, + idxHighLocal, idxLowLocal, nkValueLocal, topK, s2SeqLen); + } else { + topkb16gather::LiTopKVF(tmpIndexLocal, hisValueLocal, mrgValueLocal, histogramsLocal, + idxHighLocal, idxLowLocal, nkValueLocal, topK, s2SeqLen); + } + PipeBarrier(); + Cast(indicesOutLocal, tmpIndexLocal, RoundMode::CAST_NONE, topK); + if (outputIdxOffset != 0) { + topkb16gather::IndicesAddOffset(indicesOutLocal, outputIdxOffset, topK); + } + return; + } + if (loopIdx == 0 && !isNeedLD) { + topkb16gather::LiTopKVF(tmpIndexLocal, hisValueLocal, mrgValueLocal, histogramsLocal, idxHighLocal, + idxLowLocal, nkValueLocal, topK, s2SeqLen); + PipeBarrier(); + Cast(hisIndexLocal[(loopIdx + 1) % 2], tmpIndexLocal, RoundMode::CAST_NONE, topK); + } else if (loopIdx != 0 && !isNeedLD) { + topkb16gather::LiTopKVF(tmpIndexLocal, hisValueLocal, mrgValueLocal, histogramsLocal, idxHighLocal, + idxLowLocal, nkValueLocal, topK, s2SeqLen); + PipeBarrier(); + uint32_t curProcess = topK < trunkLen ? loopIdx * trunkLen - QLIV2Common::Align(topK, (uint32_t)256) : + (loopIdx - 1) * trunkLen; + topkb16gather::LiTopKGatherVF(hisIndexLocal[(loopIdx + 1) % 2], hisValueLocal, mrgValueLocal, + tmpIndexLocal, hisIndexLocal[loopIdx % 2], topK, + curProcess, s2SeqLen); + if (loopIdx == s2LoopNum - 1) { + PipeBarrier(); + if ((loopIdx + 1) % 2 == 1) { // 2:pingpong + AscendC::DataCopy(indicesOutLocal, hisIndexLocal[(loopIdx + 1) % 2], + QLIV2Common::Align(topK, (uint32_t)256)); + } + } + } + + if (loopIdx == 0 && isNeedLD) { + topkb16gather::LiTopKVF(tmpIndexLocal, hisValueLocal, mrgValueLocal, histogramsLocal, idxHighLocal, + idxLowLocal, nkValueLocal, topK, s2SeqLen); + PipeBarrier(); + Cast(hisIndexLocal[(loopIdx + 1) % 2], tmpIndexLocal, RoundMode::CAST_NONE, topK); + PipeBarrier(); + AscendC::DataCopy(indicesOutLocal, hisIndexLocal[(loopIdx + 1) % 2], + QLIV2Common::Align(topK, (uint32_t)256)); + } else if (loopIdx != 0 && isNeedLD) { + topkb16gather::LiTopKVF(tmpIndexLocal, hisValueLocal, mrgValueLocal, histogramsLocal, idxHighLocal, + idxLowLocal, nkValueLocal, topK, s2SeqLen); + PipeBarrier(); + uint32_t curProcess = topK < trunkLen ? loopIdx * trunkLen - QLIV2Common::Align(topK, (uint32_t)256) : + (loopIdx - 1) * trunkLen; + topkb16gather::LiTopKGatherVF(hisIndexLocal[(loopIdx + 1) % 2], hisValueLocal, mrgValueLocal, + tmpIndexLocal, hisIndexLocal[loopIdx % 2], + topK, curProcess, s2SeqLen); + PipeBarrier(); + AscendC::DataCopy(indicesOutLocal, hisIndexLocal[(loopIdx + 1) % 2], + QLIV2Common::Align(topK, (uint32_t)256)); + } + if (outputIdxOffset != 0) { + topkb16gather::IndicesAddOffset(indicesOutLocal, outputIdxOffset, topK); + } + } + __aicore__ inline void LdTopK(LocalTensor& mrgValueLocal, LocalTensor indexLocal, + LocalTensor& indicesOutLocal, + LocalTensor& hisValueLocal, uint32_t s2SeqLen, uint32_t loopIdx, + uint32_t s2LoopNum) + { + topkb16gather::LiTopKVF(tmpIndexLocal, hisValueLocal, mrgValueLocal, histogramsLocal, idxHighLocal, + idxLowLocal, nkValueLocal, topK, s2SeqLen); + PipeBarrier(); + topkb16gather::LiTopKLDGatherVF(indicesOutLocal, tmpIndexLocal, indexLocal, topK); + } + +private: + LocalTensor hisIndexLocal[2]; // 每trunkLen长度的s2选出的topK个索引 + LocalTensor histogramsLocal; // 直方图的临时Buf 256 * 4B + LocalTensor idxHighLocal; // 输入数据高8位Buf 256 * 4B + LocalTensor idxLowLocal; // 输入数据低8位Buf 256 * 4B + LocalTensor nkValueLocal; // next_k 暂存Buf 64 * 4B + LocalTensor tmpIndexLocal; // 每trunkLen + topK的临时index + uint32_t topK = 512; + uint32_t trunkLen = 16384; +}; +} +#endif // QUANT_LIGHTNING_INDEXER_V2_TOPK_H diff --git a/xllm_ops/attention/quant_lightning_indexer_v2/op_kernel/arch35/vf/quant_lightning_indexer_v2_vector1.h b/xllm_ops/attention/quant_lightning_indexer_v2/op_kernel/arch35/vf/quant_lightning_indexer_v2_vector1.h new file mode 100644 index 0000000..dbf662e --- /dev/null +++ b/xllm_ops/attention/quant_lightning_indexer_v2/op_kernel/arch35/vf/quant_lightning_indexer_v2_vector1.h @@ -0,0 +1,1502 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file quant_lightning_indexer_v2_vector1.h + * \brief + */ +#ifndef QUANT_LIGHTNING_INDEXER_V2_VECTOR1_H +#define QUANT_LIGHTNING_INDEXER_V2_VECTOR1_H + +#include "kernel_operator.h" +#if __has_include("../../../lightning_indexer_v2/arch35/vf/common/lightning_indexer_v2_vector1_base.h") +#include "../../../lightning_indexer_v2/arch35/vf/common/lightning_indexer_v2_vector1_base.h" +#else +#include "../../../../lightning_indexer_v2/op_kernel/arch35/vf/common/lightning_indexer_v2_vector1_base.h" +#endif + +namespace vector1 { +__simd_vf__ void UIntToFloatReturnValueVF(__ubuf__ bfloat16_t *outBuf, __ubuf__ uint16_t *inBuf, uint16_t vfLoop) +{ + MicroAPI::RegTensor regIn; + MicroAPI::RegTensor regOut; + MicroAPI::MaskReg maskAllB16 = MicroAPI::CreateMask(); + + for (uint16_t i = 0; i < vfLoop; ++i) { + MicroAPI::LoadAlign(regIn, inBuf + i * 128); + + liV2Vector1::UIntSortConstCtx uint16Ctx; + liV2Vector1::InitUIntSortConstCtx(uint16Ctx, maskAllB16); + + liV2Vector1::UIntToSortableKey(regOut, regIn, uint16Ctx, maskAllB16); + + MicroAPI::StoreAlign(outBuf + i * 128, regOut, maskAllB16); + } +} + +__aicore__ inline void UIntToFloatReturnValue(const LocalTensor &out_, const LocalTensor &in, + const uint32_t topK) +{ + __ubuf__ bfloat16_t *outBuf = (__ubuf__ bfloat16_t *)out_.GetPhyAddr(); + __ubuf__ uint16_t *inBuf = (__ubuf__ uint16_t *)in.GetPhyAddr(); + const uint16_t repeatSize16 = 128; + uint16_t topkLoopNum = (topK + repeatSize16 - 1) / repeatSize16; + UIntToFloatReturnValueVF(outBuf, inBuf, topkLoopNum); +} + +// 可排序键还原为 bf16 返回值，并将无效位（score==0）刷为 -inf +__simd_vf__ void UIntToFloatReturnValueWithInfMaskVF(__ubuf__ bfloat16_t *valueOutBuf, __ubuf__ uint16_t *scoreOutBuf, + uint16_t vfLoop, uint16_t negInfBits) +{ + MicroAPI::RegTensor regIn; + MicroAPI::RegTensor regOut; + MicroAPI::RegTensor regNegInf; + MicroAPI::RegTensor regZero; + MicroAPI::MaskReg maskAllB16 = MicroAPI::CreateMask(); + MicroAPI::MaskReg maskInvalid; + + // 常量寄存器初始化：-inf(0xFF80) 与 0（用于识别无效位） + MicroAPI::Duplicate(regNegInf, negInfBits, maskAllB16); + MicroAPI::Duplicate(regZero, (uint16_t)0, maskAllB16); + + liV2Vector1::UIntSortConstCtx uint16Ctx; + liV2Vector1::InitUIntSortConstCtx(uint16Ctx, maskAllB16); + + for (uint16_t i = 0; i < vfLoop; ++i) { + MicroAPI::LoadAlign(regIn, scoreOutBuf + i * 128); + // 比较得无效位掩码：score==0 即无效位（可排序键性质，真实值永不为0） + MicroAPI::Compare(maskInvalid, regIn, regZero, maskAllB16); + // 逆变换：可排序键 → bf16 值 + liV2Vector1::UIntToSortableKey(regOut, regIn, uint16Ctx, maskAllB16); + // 无效位覆盖为 -inf，有效位保留还原值 + MicroAPI::Select((MicroAPI::RegTensor &)regOut, regNegInf, (MicroAPI::RegTensor &)regOut, + maskInvalid); + MicroAPI::StoreAlign(valueOutBuf + i * 128, regOut, maskAllB16); + } +} + +__aicore__ inline void UIntToFloatReturnValueWithInfMask(const LocalTensor &valueOutLocal, + const LocalTensor &scoreOutLocal, + const uint32_t topK, const uint16_t negInfBits) +{ + __ubuf__ bfloat16_t *valueOutBuf = (__ubuf__ bfloat16_t *)valueOutLocal.GetPhyAddr(); + __ubuf__ uint16_t *scoreOutBuf = (__ubuf__ uint16_t *)scoreOutLocal.GetPhyAddr(); + const uint16_t repeatSize16 = 128; + uint16_t topkLoopNum = (topK + repeatSize16 - 1) / repeatSize16; + UIntToFloatReturnValueWithInfMaskVF(valueOutBuf, scoreOutBuf, topkLoopNum, negInfBits); +} + +__simd_callee__ inline void LoadKScaleFP16(AscendC::MicroAPI::RegTensor (®KScaleFP16)[2], + AscendC::MicroAPI::RegTensor (®KScale)[2], + AscendC::MicroAPI::MaskReg &maskAllB16, + __ubuf__ half *kScale_) +{ + constexpr static MicroAPI::CastTrait castTraitFP16ToFP32 = {MicroAPI::RegLayout::ZERO, MicroAPI::SatMode::UNKNOWN, + MicroAPI::MaskMergeMode::ZEROING, RoundMode::UNKNOWN}; + AscendC::MicroAPI::LoadAlign(regKScaleFP16[0], kScale_); + AscendC::MicroAPI::LoadAlign(regKScaleFP16[1], kScale_ + 64); + AscendC::MicroAPI::Cast(regKScale[0], regKScaleFP16[0], maskAllB16); + AscendC::MicroAPI::Cast(regKScale[1], regKScaleFP16[1], maskAllB16); +} + +__simd_callee__ inline void CastFP32ToFP16ToFP32(AscendC::MicroAPI::RegTensor (®QK0)[2], + AscendC::MicroAPI::RegTensor (®QK0Half)[2], + AscendC::MicroAPI::MaskReg &maskAllB32) +{ + AscendC::MicroAPI::MaskReg maskAllB16 = AscendC::MicroAPI::CreateMask(); + constexpr static MicroAPI::CastTrait castTraitFP32ToFP16 = {MicroAPI::RegLayout::ZERO, MicroAPI::SatMode::NO_SAT, + MicroAPI::MaskMergeMode::ZEROING, RoundMode::CAST_RINT}; + constexpr static MicroAPI::CastTrait castTraitFP16ToFP32 = {MicroAPI::RegLayout::ZERO, MicroAPI::SatMode::UNKNOWN, + MicroAPI::MaskMergeMode::ZEROING, RoundMode::UNKNOWN}; + float mulsScalar = 1.0 / 1024; + + MicroAPI::Muls(regQK0[0], regQK0[0], mulsScalar, maskAllB32); + MicroAPI::Muls(regQK0[1], regQK0[1], mulsScalar, maskAllB32); + + MicroAPI::Cast(regQK0Half[0], regQK0[0], maskAllB32); + MicroAPI::Cast(regQK0Half[1], regQK0[1], maskAllB32); + + MicroAPI::Cast(regQK0[0], regQK0Half[0], maskAllB16); + MicroAPI::Cast(regQK0[1], regQK0Half[1], maskAllB16); +} + +// int32 in uint16 out +__simd_vf__ void MulWeightAndReduceSumInt32GSizeOddVF(__ubuf__ uint16_t *out, __ubuf__ int32_t *qk, + uint32_t qkVLStride, __ubuf__ half *weight, + __ubuf__ half *kScale, __ubuf__ half *qScale, + uint16_t gSize) +{ + MicroAPI::RegTensor regwBrc; + MicroAPI::RegTensor regQK[2]; + MicroAPI::RegTensor regQKHalf[2]; + MicroAPI::RegTensor regQKInt32[2]; + MicroAPI::RegTensor regW; + MicroAPI::RegTensor regWFP16; + MicroAPI::RegTensor regWFP16Temp; + MicroAPI::RegTensor regQScale; + MicroAPI::RegTensor regQScaleFP16; + MicroAPI::RegTensor regKScale[2]; + MicroAPI::RegTensor regKScaleFP16[2]; + MicroAPI::RegTensor regSum0[2]; + MicroAPI::RegTensor regSum1[2]; + MicroAPI::MaskReg maskAllB32 = MicroAPI::CreateMask(); + MicroAPI::MaskReg maskAllB16 = MicroAPI::CreateMask(); + + liV2Vector1::FloatSortConstCtx bf16Ctx; + liV2Vector1::InitFloatSortConstCtx(bf16Ctx, maskAllB16); + + constexpr static MicroAPI::CastTrait castTraitF32ToF16_EVEN = {MicroAPI::RegLayout::ZERO, + MicroAPI::SatMode::NO_SAT, + MicroAPI::MaskMergeMode::MERGING, + RoundMode::CAST_ROUND}; + constexpr static MicroAPI::CastTrait castTraitF32ToF16_ODD = {MicroAPI::RegLayout::ONE, + MicroAPI::SatMode::NO_SAT, + MicroAPI::MaskMergeMode::ZEROING, + RoundMode::CAST_ROUND}; + constexpr static MicroAPI::CastTrait castTraitF16ToF32 = {MicroAPI::RegLayout::ZERO, + MicroAPI::SatMode::UNKNOWN, + MicroAPI::MaskMergeMode::ZEROING, + RoundMode::UNKNOWN}; + constexpr static MicroAPI::CastTrait castTraitInt32ToFP32 = {MicroAPI::RegLayout::UNKNOWN, + MicroAPI::SatMode::NO_SAT, + MicroAPI::MaskMergeMode::ZEROING, + RoundMode::CAST_ROUND}; + constexpr static MicroAPI::CastTrait castTraitF32ToF16 = {MicroAPI::RegLayout::ZERO, + MicroAPI::SatMode::NO_SAT, + MicroAPI::MaskMergeMode::ZEROING, + RoundMode::CAST_RINT}; + MicroAPI::LoadAlign(regWFP16, weight); + MicroAPI::LoadAlign(regQScaleFP16, qScale); + MicroAPI::Cast(regW, regWFP16, maskAllB16); + MicroAPI::Cast(regQScale, regQScaleFP16, maskAllB16); + MicroAPI::Mul(regW, regW, regQScale, maskAllB32); + MicroAPI::Cast(regWFP16Temp, regW, maskAllB32); + MicroAPI::Cast(regW, regWFP16Temp, maskAllB16); + liV2Vector1::DuplicateZero(regSum0, maskAllB32); + liV2Vector1::DuplicateZero(regSum1, maskAllB32); + + LoadKScaleFP16(regKScaleFP16, regKScale, maskAllB16, kScale); + // float mulsScalar = 1.0f / 1024; + // unroll2 + for (uint16_t i = (uint16_t)(0); (uint16_t)(i + 1) < gSize; i += 2) { + MicroAPI::LoadAlign(regQKInt32[0], qk + 128 * i); + MicroAPI::LoadAlign(regQKInt32[1], qk + 128 * i + qkVLStride); + MicroAPI::Cast(regQK[0], regQKInt32[0], maskAllB32); + MicroAPI::Cast(regQK[1], regQKInt32[1], maskAllB32); + + CastFP32ToFP16ToFP32(regQK, regQKHalf, maskAllB32); + + liV2Vector1::BroadcastLane(regwBrc, regW, i); + liV2Vector1::WeightedAccum(regSum0, regQK, regwBrc, maskAllB32); + + MicroAPI::LoadAlign(regQKInt32[0], qk + 128 * i + 128); + MicroAPI::LoadAlign(regQKInt32[1], qk + 128 * i + 128 + qkVLStride); + MicroAPI::Cast(regQK[0], regQKInt32[0], maskAllB32); + MicroAPI::Cast(regQK[1], regQKInt32[1], maskAllB32); + + CastFP32ToFP16ToFP32(regQK, regQKHalf, maskAllB32); + + liV2Vector1::BroadcastLane(regwBrc, regW, i + 1); + liV2Vector1::WeightedAccum(regSum1, regQK, regwBrc, maskAllB32); + } + + MicroAPI::LoadAlign(regQKInt32[0], qk + 128 * (gSize - 1)); + MicroAPI::LoadAlign(regQKInt32[1], qk + 128 * (gSize - 1) + qkVLStride); + MicroAPI::Cast(regQK[0], regQKInt32[0], maskAllB32); + MicroAPI::Cast(regQK[1], regQKInt32[1], maskAllB32); + + CastFP32ToFP16ToFP32(regQK, regQKHalf, maskAllB32); + + liV2Vector1::BroadcastLane(regwBrc, regW, gSize - 1); + liV2Vector1::WeightedAccum(regSum0, regQK, regwBrc, maskAllB32); + + MicroAPI::Add(regSum0[0], regSum0[0], regSum1[0], maskAllB32); + MicroAPI::Add(regSum0[1], regSum0[1], regSum1[1], maskAllB32); + + MicroAPI::Mul(regSum0[0], regSum0[0], regKScale[0], maskAllB32); + MicroAPI::Mul(regSum0[1], regSum0[1], regKScale[1], maskAllB32); + + MicroAPI::RegTensor regSumBF16; + // interleave cast ==> regSum[1] high regSum[0] low + MicroAPI::DeInterleave(regSum0[0], regSum0[1], regSum0[0], regSum0[1]); + MicroAPI::Cast(regSumBF16, regSum0[1], maskAllB32); + MicroAPI::Cast(regSumBF16, regSum0[0], maskAllB32); + + MicroAPI::RegTensor regOut; + liV2Vector1::FloatToSortableKey(regOut, regSumBF16, bf16Ctx, maskAllB16); + // normal store + MicroAPI::StoreAlign(out, regOut, maskAllB16); +} + +// float in uint16 out +__simd_vf__ void MulWeightAndReduceSumF32GSizeOddVF(__ubuf__ uint16_t *out, __ubuf__ float *qk, uint32_t qkVLStride, + __ubuf__ float *weight, __ubuf__ float *kScale, + __ubuf__ float *qScale, uint16_t gSize) +{ + MicroAPI::RegTensor regwBrc; + MicroAPI::RegTensor regQK[2]; + MicroAPI::RegTensor regW; + + MicroAPI::RegTensor regQScale; + MicroAPI::RegTensor regKScale[2]; + MicroAPI::RegTensor regSum0[2]; + MicroAPI::RegTensor regSum1[2]; + MicroAPI::MaskReg maskAllB32 = MicroAPI::CreateMask(); + MicroAPI::MaskReg maskAllB16 = MicroAPI::CreateMask(); + + liV2Vector1::FloatSortConstCtx bf16Ctx; + liV2Vector1::InitFloatSortConstCtx(bf16Ctx, maskAllB16); + + constexpr static MicroAPI::CastTrait castTraitF32ToF16_EVEN = { + MicroAPI::RegLayout::ZERO, MicroAPI::SatMode::NO_SAT, MicroAPI::MaskMergeMode::MERGING, RoundMode::CAST_ROUND}; + constexpr static MicroAPI::CastTrait castTraitF32ToF16_ODD = { + MicroAPI::RegLayout::ONE, MicroAPI::SatMode::NO_SAT, MicroAPI::MaskMergeMode::ZEROING, RoundMode::CAST_ROUND}; + + MicroAPI::LoadAlign(regW, weight); + MicroAPI::LoadAlign(regQScale, qScale); + MicroAPI::Mul(regW, regW, regQScale, maskAllB32); + + liV2Vector1::DuplicateZero(regSum0, maskAllB32); + liV2Vector1::DuplicateZero(regSum1, maskAllB32); + + MicroAPI::LoadAlign(regKScale[0], kScale); + MicroAPI::LoadAlign(regKScale[1], kScale + 64); + + // unroll2 + for (uint16_t i = (uint16_t)(0); (uint16_t)(i + 1) < gSize; i += 2) { + MicroAPI::LoadAlign(regQK[0], qk + 128 * i); // RowStride是128, 行都落在一个bank上 + MicroAPI::LoadAlign(regQK[1], qk + 128 * i + qkVLStride); + liV2Vector1::BroadcastLane(regwBrc, regW, i); + liV2Vector1::WeightedAccum(regSum0, regQK, regwBrc, maskAllB32); + + MicroAPI::LoadAlign(regQK[0], qk + 128 * i + 128); + MicroAPI::LoadAlign(regQK[1], qk + 128 * i + 128 + qkVLStride); + liV2Vector1::BroadcastLane(regwBrc, regW, i + 1); + liV2Vector1::WeightedAccum(regSum1, regQK, regwBrc, maskAllB32); + } + + MicroAPI::LoadAlign(regQK[0], qk + 128 * (gSize - 1)); // RowStride是128, 行都落在一个bank上 + MicroAPI::LoadAlign(regQK[1], qk + 128 * (gSize - 1) + qkVLStride); + liV2Vector1::BroadcastLane(regwBrc, regW, (gSize - 1)); + liV2Vector1::WeightedAccum(regSum0, regQK, regwBrc, maskAllB32); + + MicroAPI::Add(regSum0[0], regSum0[0], regSum1[0], maskAllB32); + MicroAPI::Add(regSum0[1], regSum0[1], regSum1[1], maskAllB32); + + MicroAPI::Mul(regSum0[0], regSum0[0], regKScale[0], maskAllB32); + MicroAPI::Mul(regSum0[1], regSum0[1], regKScale[1], maskAllB32); + + MicroAPI::RegTensor regSumBF16; + // interleave cast ==> regSum[1] high regSum[0] low + MicroAPI::DeInterleave(regSum0[0], regSum0[1], regSum0[0], regSum0[1]); + MicroAPI::Cast(regSumBF16, regSum0[1], maskAllB32); + MicroAPI::Cast(regSumBF16, regSum0[0], maskAllB32); + + MicroAPI::RegTensor regOut; + liV2Vector1::FloatToSortableKey(regOut, regSumBF16, bf16Ctx, maskAllB16); + // normal store + MicroAPI::StoreAlign(out, regOut, maskAllB16); +} + +// int32 in uint16 out +__simd_vf__ void MulWeightAndReduceSumInt32GSizeEvenVF(__ubuf__ uint16_t *out, __ubuf__ int32_t *qk, + uint32_t qkVLStride, __ubuf__ half *weight, + __ubuf__ half *kScale, __ubuf__ half *qScale, + uint16_t gSize) +{ + MicroAPI::RegTensor regwBrc; + MicroAPI::RegTensor regQK[2]; + MicroAPI::RegTensor regQKHalf[2]; + MicroAPI::RegTensor regQKInt32[2]; + MicroAPI::RegTensor regW; + MicroAPI::RegTensor regWFP16; + MicroAPI::RegTensor regWFP16Temp; + MicroAPI::RegTensor regQScale; + MicroAPI::RegTensor regQScaleFP16; + MicroAPI::RegTensor regKScale[2]; + MicroAPI::RegTensor regKScaleFP16[2]; + MicroAPI::RegTensor regSum0[2]; + MicroAPI::RegTensor regSum1[2]; + MicroAPI::MaskReg maskAllB32 = MicroAPI::CreateMask(); + MicroAPI::MaskReg maskAllB16 = MicroAPI::CreateMask(); + + liV2Vector1::FloatSortConstCtx bf16Ctx; + liV2Vector1::InitFloatSortConstCtx(bf16Ctx, maskAllB16); + + constexpr static MicroAPI::CastTrait castTraitF32ToF16_EVEN = {MicroAPI::RegLayout::ZERO, + MicroAPI::SatMode::NO_SAT, + MicroAPI::MaskMergeMode::MERGING, + RoundMode::CAST_ROUND}; + constexpr static MicroAPI::CastTrait castTraitF32ToF16_ODD = {MicroAPI::RegLayout::ONE, + MicroAPI::SatMode::NO_SAT, + MicroAPI::MaskMergeMode::ZEROING, + RoundMode::CAST_ROUND}; + constexpr static MicroAPI::CastTrait castTraitF16ToF32 = {MicroAPI::RegLayout::ZERO, + MicroAPI::SatMode::UNKNOWN, + MicroAPI::MaskMergeMode::ZEROING, + RoundMode::UNKNOWN}; + constexpr static MicroAPI::CastTrait castTraitInt32ToFP32 = {MicroAPI::RegLayout::UNKNOWN, + MicroAPI::SatMode::NO_SAT, + MicroAPI::MaskMergeMode::ZEROING, + RoundMode::CAST_ROUND}; + constexpr static MicroAPI::CastTrait castTraitF32ToF16 = {MicroAPI::RegLayout::ZERO, + MicroAPI::SatMode::NO_SAT, + MicroAPI::MaskMergeMode::ZEROING, + RoundMode::CAST_RINT}; + MicroAPI::LoadAlign(regWFP16, weight); + MicroAPI::LoadAlign(regQScaleFP16, qScale); + MicroAPI::Cast(regW, regWFP16, maskAllB16); + MicroAPI::Cast(regQScale, regQScaleFP16, maskAllB16); + MicroAPI::Mul(regW, regW, regQScale, maskAllB32); + MicroAPI::Cast(regWFP16Temp, regW, maskAllB32); + MicroAPI::Cast(regW, regWFP16Temp, maskAllB16); + liV2Vector1::DuplicateZero(regSum0, maskAllB32); + liV2Vector1::DuplicateZero(regSum1, maskAllB32); + + LoadKScaleFP16(regKScaleFP16, regKScale, maskAllB16, kScale); + // float mulsScalar = 1.0f / 1024; + // unroll2 + for (uint16_t i = (uint16_t)(0); i < gSize; i += 2) { + MicroAPI::LoadAlign(regQKInt32[0], qk + 128 * i); + MicroAPI::LoadAlign(regQKInt32[1], qk + 128 * i + qkVLStride); + MicroAPI::Cast(regQK[0], regQKInt32[0], maskAllB32); + MicroAPI::Cast(regQK[1], regQKInt32[1], maskAllB32); + + CastFP32ToFP16ToFP32(regQK, regQKHalf, maskAllB32); + + liV2Vector1::BroadcastLane(regwBrc, regW, i); + liV2Vector1::WeightedAccum(regSum0, regQK, regwBrc, maskAllB32); + + MicroAPI::LoadAlign(regQKInt32[0], qk + 128 * i + 128); + MicroAPI::LoadAlign(regQKInt32[1], qk + 128 * i + 128 + qkVLStride); + MicroAPI::Cast(regQK[0], regQKInt32[0], maskAllB32); + MicroAPI::Cast(regQK[1], regQKInt32[1], maskAllB32); + + CastFP32ToFP16ToFP32(regQK, regQKHalf, maskAllB32); + + liV2Vector1::BroadcastLane(regwBrc, regW, i + 1); + liV2Vector1::WeightedAccum(regSum1, regQK, regwBrc, maskAllB32); + } + + MicroAPI::Add(regSum0[0], regSum0[0], regSum1[0], maskAllB32); + MicroAPI::Add(regSum0[1], regSum0[1], regSum1[1], maskAllB32); + + MicroAPI::Mul(regSum0[0], regSum0[0], regKScale[0], maskAllB32); + MicroAPI::Mul(regSum0[1], regSum0[1], regKScale[1], maskAllB32); + + MicroAPI::RegTensor regSumBF16; + // interleave cast ==> regSum[1] high regSum[0] low + MicroAPI::DeInterleave(regSum0[0], regSum0[1], regSum0[0], regSum0[1]); + MicroAPI::Cast(regSumBF16, regSum0[1], maskAllB32); + MicroAPI::Cast(regSumBF16, regSum0[0], maskAllB32); + + MicroAPI::RegTensor regOut; + liV2Vector1::FloatToSortableKey(regOut, regSumBF16, bf16Ctx, maskAllB16); + // normal store + MicroAPI::StoreAlign(out, regOut, maskAllB16); +} + +__simd_vf__ void MulWeightAndReduceSumF32GSizeEvenVF(__ubuf__ uint16_t *out, __ubuf__ float *qk, uint32_t qkVLStride, + __ubuf__ float *weight, __ubuf__ float *kScale, + __ubuf__ float *qScale, uint16_t gSize) +{ + MicroAPI::RegTensor regwBrc; + MicroAPI::RegTensor regQK[2]; + MicroAPI::RegTensor regW; + + MicroAPI::RegTensor regQScale; + MicroAPI::RegTensor regKScale[2]; + MicroAPI::RegTensor regSum0[2]; + MicroAPI::RegTensor regSum1[2]; + MicroAPI::MaskReg maskAllB32 = MicroAPI::CreateMask(); + MicroAPI::MaskReg maskAllB16 = MicroAPI::CreateMask(); + + liV2Vector1::FloatSortConstCtx bf16Ctx; + liV2Vector1::InitFloatSortConstCtx(bf16Ctx, maskAllB16); + + constexpr static MicroAPI::CastTrait castTraitF32ToF16_EVEN = { + MicroAPI::RegLayout::ZERO, MicroAPI::SatMode::NO_SAT, MicroAPI::MaskMergeMode::MERGING, RoundMode::CAST_ROUND}; + constexpr static MicroAPI::CastTrait castTraitF32ToF16_ODD = { + MicroAPI::RegLayout::ONE, MicroAPI::SatMode::NO_SAT, MicroAPI::MaskMergeMode::ZEROING, RoundMode::CAST_ROUND}; + + MicroAPI::LoadAlign(regW, weight); + MicroAPI::LoadAlign(regQScale, qScale); + MicroAPI::Mul(regW, regW, regQScale, maskAllB32); + + liV2Vector1::DuplicateZero(regSum0, maskAllB32); + liV2Vector1::DuplicateZero(regSum1, maskAllB32); + + MicroAPI::LoadAlign(regKScale[0], kScale); + MicroAPI::LoadAlign(regKScale[1], kScale + 64); + + // unroll2 + for (uint16_t i = (uint16_t)(0); i < gSize; i += 2) { + MicroAPI::LoadAlign(regQK[0], qk + 128 * i); // RowStride是128, 行都落在一个bank上 + MicroAPI::LoadAlign(regQK[1], qk + 128 * i + qkVLStride); + liV2Vector1::BroadcastLane(regwBrc, regW, i); + liV2Vector1::WeightedAccum(regSum0, regQK, regwBrc, maskAllB32); + + MicroAPI::LoadAlign(regQK[0], qk + 128 * i + 128); + MicroAPI::LoadAlign(regQK[1], qk + 128 * i + 128 + qkVLStride); + liV2Vector1::BroadcastLane(regwBrc, regW, i + 1); + liV2Vector1::WeightedAccum(regSum1, regQK, regwBrc, maskAllB32); + } + + MicroAPI::Add(regSum0[0], regSum0[0], regSum1[0], maskAllB32); + MicroAPI::Add(regSum0[1], regSum0[1], regSum1[1], maskAllB32); + + MicroAPI::Mul(regSum0[0], regSum0[0], regKScale[0], maskAllB32); + MicroAPI::Mul(regSum0[1], regSum0[1], regKScale[1], maskAllB32); + + MicroAPI::RegTensor regSumBF16; + // interleave cast ==> regSum[1] high regSum[0] low + MicroAPI::DeInterleave(regSum0[0], regSum0[1], regSum0[0], regSum0[1]); + MicroAPI::Cast(regSumBF16, regSum0[1], maskAllB32); + MicroAPI::Cast(regSumBF16, regSum0[0], maskAllB32); + + MicroAPI::RegTensor regOut; + liV2Vector1::FloatToSortableKey(regOut, regSumBF16, bf16Ctx, maskAllB16); + // normal store + MicroAPI::StoreAlign(out, regOut, maskAllB16); +} + +__aicore__ inline void MulWeightAndReduceSum(const LocalTensor &out_, // out [S2Base] [128 ] + const LocalTensor &qk_, // q*k^t [G, S2Base] [64 128] + const uint32_t qkVLStride, + const LocalTensor &weight_, // w [G] [64 ] + const LocalTensor &kScale_, // kScale [S2Base] [128 ] + const LocalTensor &qScale_, // qScale [G] [64 ] + const int gSize) // G 64 +{ + __ubuf__ uint16_t *out = (__ubuf__ uint16_t *)out_.GetPhyAddr(); + __ubuf__ float *weight = (__ubuf__ float *)weight_.GetPhyAddr(); + __ubuf__ float *qScale = (__ubuf__ float *)qScale_.GetPhyAddr(); + __ubuf__ float *kScale = (__ubuf__ float *)kScale_.GetPhyAddr(); + __ubuf__ float *qk = (__ubuf__ float *)qk_.GetPhyAddr(); + if (gSize % 2 == 0) { + MulWeightAndReduceSumF32GSizeEvenVF(out, qk, qkVLStride, weight, kScale, qScale, (uint16_t)gSize); + } else { + MulWeightAndReduceSumF32GSizeOddVF(out, qk, qkVLStride, weight, kScale, qScale, (uint16_t)gSize); + } +} + +__aicore__ inline void MulWeightAndReduceSum(const LocalTensor &out_, // out [S2Base] [128 ] + const LocalTensor &qk_, // q*k^t [G, S2Base] [64 128] + const uint32_t qkVLStride, + const LocalTensor &weight_, // w [G] [64 ] + const LocalTensor &kScale_, // kScale [S2Base] [128 ] + const LocalTensor &qScale_, // qScale [G] [64 ] + const int gSize) // G 64 +{ + __ubuf__ uint16_t *out = (__ubuf__ uint16_t *)out_.GetPhyAddr(); + __ubuf__ half *weight = (__ubuf__ half *)weight_.GetPhyAddr(); + __ubuf__ half *qScale = (__ubuf__ half *)qScale_.GetPhyAddr(); + __ubuf__ half *kScale = (__ubuf__ half *)kScale_.GetPhyAddr(); + __ubuf__ int32_t *qk = (__ubuf__ int32_t *)qk_.GetPhyAddr(); + if (gSize % 2 != 0) { + MulWeightAndReduceSumInt32GSizeOddVF(out, qk, qkVLStride, weight, kScale, qScale, (uint16_t)gSize); + } else { + MulWeightAndReduceSumInt32GSizeEvenVF(out, qk, qkVLStride, weight, kScale, qScale, (uint16_t)gSize); + } +} + +// bfloat16_t in uint16 out +__simd_vf__ void MulWeightAndReduceSumB16VF(__ubuf__ uint16_t *out, __ubuf__ bfloat16_t *qk, __ubuf__ float *weight, + __ubuf__ float *kScale, __ubuf__ float *qScale, uint16_t gSize) +{ + MicroAPI::RegTensor regQK[4]; + MicroAPI::RegTensor regQKB16[2]; + MicroAPI::RegTensor regW; + MicroAPI::RegTensor regwBrc[2]; + MicroAPI::RegTensor regQScale; + MicroAPI::RegTensor regKScale[2]; + MicroAPI::RegTensor regSum[2]; + + MicroAPI::MaskReg maskAllB32 = MicroAPI::CreateMask(); + MicroAPI::MaskReg maskAllB16 = MicroAPI::CreateMask(); + + MicroAPI::RegTensor regSumBF16; + + liV2Vector1::FloatSortConstCtx bf16Ctx; + liV2Vector1::InitFloatSortConstCtx(bf16Ctx, maskAllB16); + + using CastTrait = MicroAPI::CastTrait; + static constexpr CastTrait castTraitB162B32_EVEN = {MicroAPI::RegLayout::ZERO, MicroAPI::SatMode::UNKNOWN, + MicroAPI::MaskMergeMode::ZEROING, RoundMode::UNKNOWN}; + static constexpr CastTrait castTraitB162B32_ODD = {MicroAPI::RegLayout::ONE, MicroAPI::SatMode::UNKNOWN, + MicroAPI::MaskMergeMode::ZEROING, RoundMode::UNKNOWN}; + + constexpr static CastTrait castTraitF32ToF16_EVEN = {MicroAPI::RegLayout::ZERO, MicroAPI::SatMode::NO_SAT, + MicroAPI::MaskMergeMode::MERGING, RoundMode::CAST_ROUND}; + constexpr static CastTrait castTraitF32ToF16_ODD = {MicroAPI::RegLayout::ONE, MicroAPI::SatMode::NO_SAT, + MicroAPI::MaskMergeMode::ZEROING, RoundMode::CAST_ROUND}; + + MicroAPI::LoadAlign(regW, weight); + MicroAPI::LoadAlign(regQScale, qScale); + MicroAPI::Mul(regW, regW, regQScale, maskAllB32); + MicroAPI::StoreAlign(weight, regW, maskAllB32); + MicroAPI::LocalMemBar(); + + liV2Vector1::DuplicateZero(regSum, maskAllB32); + + // interleave load + MicroAPI::LoadAlign(regKScale[0], regKScale[1], kScale); + + // Duplicate + Gather方法劣化 + // Relu在cube随路做 + for (uint16_t i = (uint16_t)(0); i < gSize; i++) { + // RowStride是256, 行都落在一个bank上 + MicroAPI::LoadAlign(regQKB16[0], qk + 256 * i); + MicroAPI::LoadAlign(regwBrc[0], weight + i); + // interleave cast + MicroAPI::Cast(regQK[0], regQKB16[0], maskAllB16); + MicroAPI::Cast(regQK[1], regQKB16[0], maskAllB16); + MicroAPI::MulAddDst(regSum[0], regQK[0], regwBrc[0], maskAllB32); + MicroAPI::MulAddDst(regSum[1], regQK[1], regwBrc[0], maskAllB32); + } + + MicroAPI::Mul(regSum[0], regSum[0], regKScale[0], maskAllB32); + MicroAPI::Mul(regSum[1], regSum[1], regKScale[1], maskAllB32); + // interleave cast back + MicroAPI::Cast(regSumBF16, regSum[1], maskAllB32); + MicroAPI::Cast(regSumBF16, regSum[0], maskAllB32); + + MicroAPI::RegTensor regOut; + liV2Vector1::FloatToSortableKey(regOut, regSumBF16, bf16Ctx, maskAllB16); + // norm load + MicroAPI::StoreAlign(out, regOut, maskAllB16); +} + +__aicore__ inline void MulWeightAndReduceSum(const LocalTensor &out_, // out [S2Base] [128 ] + const LocalTensor &qk_, // q*k^t [G, S2Base] [64 128] + const uint32_t qkVLStride, // unused for bfloat16 + const LocalTensor &weight_, // w [G] [64 ] + const LocalTensor &kScale_, // kScale [S2Base] [128 ] + const LocalTensor &qScale_, // qScale [G] [64 ] + const int gSize) // G 64 +{ + __ubuf__ uint16_t *out = (__ubuf__ uint16_t *)out_.GetPhyAddr(); + __ubuf__ float *weight = (__ubuf__ float *)weight_.GetPhyAddr(); + __ubuf__ float *qScale = (__ubuf__ float *)qScale_.GetPhyAddr(); + __ubuf__ bfloat16_t *qk = (__ubuf__ bfloat16_t *)qk_.GetPhyAddr(); + __ubuf__ float *kScale = (__ubuf__ float *)kScale_.GetPhyAddr(); + MulWeightAndReduceSumB16VF(out, qk, weight, kScale, qScale, (uint16_t)gSize); +} + +// 计算S1=2 +// float in uint16 out +__simd_vf__ void MulWeightAndReduceSum2F32VF(__ubuf__ uint16_t *out0, __ubuf__ uint16_t *out1, __ubuf__ float *qk0, + __ubuf__ float *qk1, uint32_t qkVLStride, __ubuf__ float *weight0, + __ubuf__ float *weight1, __ubuf__ float *weightTemp, + __ubuf__ float *qScale0, __ubuf__ float *qScale1, __ubuf__ float *kScale0, + uint16_t gSize) +{ + MicroAPI::RegTensor regwBrc[2]; + MicroAPI::RegTensor regQK0[2]; + MicroAPI::RegTensor regQK1[2]; + MicroAPI::RegTensor regW[2]; + + MicroAPI::RegTensor regQScale[2]; + MicroAPI::RegTensor regKScale[2]; + MicroAPI::RegTensor regSum0[2]; + MicroAPI::RegTensor regSum1[2]; + MicroAPI::MaskReg maskAllB32 = MicroAPI::CreateMask(); + MicroAPI::MaskReg maskAllB16 = MicroAPI::CreateMask(); + + liV2Vector1::FloatSortConstCtx bf16Ctx; + liV2Vector1::InitFloatSortConstCtx(bf16Ctx, maskAllB16); + + constexpr static MicroAPI::CastTrait castTraitF32ToF16_EVEN = { + MicroAPI::RegLayout::ZERO, MicroAPI::SatMode::NO_SAT, MicroAPI::MaskMergeMode::MERGING, RoundMode::CAST_ROUND}; + constexpr static MicroAPI::CastTrait castTraitF32ToF16_ODD = { + MicroAPI::RegLayout::ONE, MicroAPI::SatMode::NO_SAT, MicroAPI::MaskMergeMode::ZEROING, RoundMode::CAST_ROUND}; + + MicroAPI::LoadAlign(regW[0], weight0); + MicroAPI::LoadAlign(regW[1], weight1); + MicroAPI::LoadAlign(regQScale[0], qScale0); + MicroAPI::LoadAlign(regQScale[1], qScale1); + MicroAPI::Mul(regW[0], regW[0], regQScale[0], maskAllB32); + MicroAPI::Mul(regW[1], regW[1], regQScale[1], maskAllB32); + // regW[0]与weight1混合使用 + MicroAPI::StoreAlign(weightTemp, regW[1], maskAllB32); + MicroAPI::LocalMemBar(); + liV2Vector1::DuplicateZero(regSum0, maskAllB32); + liV2Vector1::DuplicateZero(regSum1, maskAllB32); + + MicroAPI::LoadAlign(regKScale[0], kScale0); + MicroAPI::LoadAlign(regKScale[1], kScale0 + 64); + + for (uint16_t i = (uint16_t)(0); i < gSize; i++) { + MicroAPI::LoadAlign(regQK0[0], qk0 + 128 * i); + MicroAPI::LoadAlign(regQK0[1], qk0 + 128 * i + qkVLStride); + MicroAPI::LoadAlign(regQK1[0], qk1 + 128 * i); + MicroAPI::LoadAlign(regQK1[1], qk1 + 128 * i + qkVLStride); + // 混合使用对整体性能更好 + liV2Vector1::BroadcastLane(regwBrc[0], regW[0], i); + // Weight无bank冲突，用LoadAlign来提取weight标量 + // 地址空间处理：原 BroadcastLane(ptr) 内联为 LoadAlign BRC，避免 __ubuf__ 传给 __local_mem__ 参数 + MicroAPI::LoadAlign(regwBrc[1], weightTemp + i); + MicroAPI::Relu(regQK0[0], regQK0[0], maskAllB32); + MicroAPI::Relu(regQK0[1], regQK0[1], maskAllB32); + MicroAPI::Relu(regQK1[0], regQK1[0], maskAllB32); + MicroAPI::Relu(regQK1[1], regQK1[1], maskAllB32); + MicroAPI::MulAddDst(regSum0[0], regQK0[0], regwBrc[0], maskAllB32); + MicroAPI::MulAddDst(regSum0[1], regQK0[1], regwBrc[0], maskAllB32); + MicroAPI::MulAddDst(regSum1[0], regQK1[0], regwBrc[1], maskAllB32); + MicroAPI::MulAddDst(regSum1[1], regQK1[1], regwBrc[1], maskAllB32); + } + + // Apply kScale scaling + MicroAPI::Mul(regSum0[0], regSum0[0], regKScale[0], maskAllB32); + MicroAPI::Mul(regSum0[1], regSum0[1], regKScale[1], maskAllB32); + MicroAPI::Mul(regSum1[0], regSum1[0], regKScale[0], maskAllB32); + MicroAPI::Mul(regSum1[1], regSum1[1], regKScale[1], maskAllB32); + + // Convert to bfloat16 and store output channel + MicroAPI::RegTensor regSumBF16[2]; + MicroAPI::RegTensor regOut[2]; + MicroAPI::DeInterleave(regSum0[0], regSum0[1], regSum0[0], regSum0[1]); + MicroAPI::DeInterleave(regSum1[0], regSum1[1], regSum1[0], regSum1[1]); + MicroAPI::Cast(regSumBF16[0], regSum0[1], maskAllB32); + MicroAPI::Cast(regSumBF16[1], regSum1[1], maskAllB32); + MicroAPI::Cast(regSumBF16[0], regSum0[0], maskAllB32); + MicroAPI::Cast(regSumBF16[1], regSum1[0], maskAllB32); + + liV2Vector1::FloatX2ToSortableKey(regOut[0], regOut[1], regSumBF16[0], regSumBF16[1], bf16Ctx, + maskAllB16); + MicroAPI::StoreAlign(out0, regOut[0], maskAllB16); + MicroAPI::StoreAlign(out1, regOut[1], maskAllB16); +} + +// 计算S1=2 +// int32 in uint16 out +__simd_vf__ void MulWeightAndReduceSum2Int32VF(__ubuf__ uint16_t *out0, __ubuf__ uint16_t *out1, + __ubuf__ int32_t *qk0, __ubuf__ int32_t *qk1, + uint32_t qkVLStride, __ubuf__ half *weight0, + __ubuf__ half *weight1, __ubuf__ float *weightTemp, + __ubuf__ half *qScale0, __ubuf__ half *qScale1, + __ubuf__ half *kScale0, uint16_t gSize) +{ + MicroAPI::RegTensor regwBrc[2]; + MicroAPI::RegTensor regQK0[2]; + MicroAPI::RegTensor regQK1[2]; + MicroAPI::RegTensor regQK0Half[2]; + MicroAPI::RegTensor regQK1Half[2]; + MicroAPI::RegTensor regQK0Int32[2]; + MicroAPI::RegTensor regQK1Int32[2]; + MicroAPI::RegTensor regW[2]; + MicroAPI::RegTensor regWFP16[2]; + MicroAPI::RegTensor regWFP16Temp[2]; + MicroAPI::RegTensor regQScale[2]; + MicroAPI::RegTensor regQScaleFP16[2]; + MicroAPI::RegTensor regKScale[2]; + MicroAPI::RegTensor regKScaleFP16[2]; + MicroAPI::RegTensor regSum0[2]; + MicroAPI::RegTensor regSum1[2]; + MicroAPI::MaskReg maskAllB32 = MicroAPI::CreateMask(); + MicroAPI::MaskReg maskAllB16 = MicroAPI::CreateMask(); + + liV2Vector1::FloatSortConstCtx bf16Ctx; + liV2Vector1::InitFloatSortConstCtx(bf16Ctx, maskAllB16); + + constexpr static MicroAPI::CastTrait castTraitF32ToF16_EVEN = {MicroAPI::RegLayout::ZERO, + MicroAPI::SatMode::NO_SAT, + MicroAPI::MaskMergeMode::MERGING, + RoundMode::CAST_ROUND}; + constexpr static MicroAPI::CastTrait castTraitF32ToF16_ODD = {MicroAPI::RegLayout::ONE, + MicroAPI::SatMode::NO_SAT, + MicroAPI::MaskMergeMode::ZEROING, + RoundMode::CAST_ROUND}; + constexpr static MicroAPI::CastTrait castTraitF16ToF32 = {MicroAPI::RegLayout::ZERO, + MicroAPI::SatMode::UNKNOWN, + MicroAPI::MaskMergeMode::ZEROING, + RoundMode::UNKNOWN}; + constexpr static MicroAPI::CastTrait castTraitF32ToF16 = {MicroAPI::RegLayout::ZERO, + MicroAPI::SatMode::NO_SAT, + MicroAPI::MaskMergeMode::ZEROING, + RoundMode::CAST_RINT}; + constexpr static MicroAPI::CastTrait castTraitInt32ToFP32 = {MicroAPI::RegLayout::UNKNOWN, + MicroAPI::SatMode::NO_SAT, + MicroAPI::MaskMergeMode::ZEROING, + RoundMode::CAST_ROUND}; + MicroAPI::LoadAlign(regWFP16[0], weight0); + MicroAPI::LoadAlign(regWFP16[1], weight1); + MicroAPI::LoadAlign(regQScaleFP16[0], qScale0); + MicroAPI::LoadAlign(regQScaleFP16[1], qScale1); + MicroAPI::Cast(regW[0], regWFP16[0], maskAllB16); + MicroAPI::Cast(regW[1], regWFP16[1], maskAllB16); + MicroAPI::Cast(regQScale[0], regQScaleFP16[0], maskAllB16); + MicroAPI::Cast(regQScale[1], regQScaleFP16[1], maskAllB16); + + MicroAPI::Mul(regW[0], regW[0], regQScale[0], maskAllB32); + MicroAPI::Mul(regW[1], regW[1], regQScale[1], maskAllB32); + + MicroAPI::Cast(regWFP16Temp[0], regW[0], maskAllB32); + MicroAPI::Cast(regW[0], regWFP16Temp[0], maskAllB16); + MicroAPI::Cast(regWFP16Temp[1], regW[1], maskAllB32); + MicroAPI::Cast(regW[1], regWFP16Temp[1], maskAllB16); + // regW[0]与weight1混合使用 + MicroAPI::StoreAlign(weightTemp, regW[1], maskAllB32); + MicroAPI::LocalMemBar(); + liV2Vector1::DuplicateZero(regSum0, maskAllB32); + liV2Vector1::DuplicateZero(regSum1, maskAllB32); + + LoadKScaleFP16(regKScaleFP16, regKScale, maskAllB16, kScale0); + // float mulsScalar = 1.0 / 1024; + + for (uint16_t i = (uint16_t)(0); i < gSize; i++) { + MicroAPI::LoadAlign(regQK0Int32[0], qk0 + 128 * i); + MicroAPI::Cast(regQK0[0], regQK0Int32[0], maskAllB32); + MicroAPI::LoadAlign(regQK0Int32[1], qk0 + 128 * i + qkVLStride); + MicroAPI::Cast(regQK0[1], regQK0Int32[1], maskAllB32); + MicroAPI::LoadAlign(regQK1Int32[0], qk1 + 128 * i); + MicroAPI::Cast(regQK1[0], regQK1Int32[0], maskAllB32); + MicroAPI::LoadAlign(regQK1Int32[1], qk1 + 128 * i + qkVLStride); + MicroAPI::Cast(regQK1[1], regQK1Int32[1], maskAllB32); + // 混合使用对整体性能更好 + liV2Vector1::BroadcastLane(regwBrc[0], regW[0], i); + // Weight无bank冲突，用LoadAlign来提取weight标量 + // 地址空间处理：原 BroadcastLane(ptr) 内联为 LoadAlign BRC，避免 __ubuf__ 传给 __local_mem__ 参数 + MicroAPI::LoadAlign(regwBrc[1], weightTemp + i); + MicroAPI::Relu(regQK0[0], regQK0[0], maskAllB32); + MicroAPI::Relu(regQK0[1], regQK0[1], maskAllB32); + MicroAPI::Relu(regQK1[0], regQK1[0], maskAllB32); + MicroAPI::Relu(regQK1[1], regQK1[1], maskAllB32); + + CastFP32ToFP16ToFP32(regQK0, regQK0Half, maskAllB32); + CastFP32ToFP16ToFP32(regQK1, regQK1Half, maskAllB32); + + MicroAPI::MulAddDst(regSum0[0], regQK0[0], regwBrc[0], maskAllB32); + MicroAPI::MulAddDst(regSum0[1], regQK0[1], regwBrc[0], maskAllB32); + MicroAPI::MulAddDst(regSum1[0], regQK1[0], regwBrc[1], maskAllB32); + MicroAPI::MulAddDst(regSum1[1], regQK1[1], regwBrc[1], maskAllB32); + } + + // Apply kScale scaling + MicroAPI::Mul(regSum0[0], regSum0[0], regKScale[0], maskAllB32); + MicroAPI::Mul(regSum0[1], regSum0[1], regKScale[1], maskAllB32); + MicroAPI::Mul(regSum1[0], regSum1[0], regKScale[0], maskAllB32); + MicroAPI::Mul(regSum1[1], regSum1[1], regKScale[1], maskAllB32); + + // Convert to bfloat16 and store output channel + MicroAPI::RegTensor regSumBF16[2]; + MicroAPI::RegTensor regOut[2]; + MicroAPI::DeInterleave(regSum0[0], regSum0[1], regSum0[0], regSum0[1]); + MicroAPI::DeInterleave(regSum1[0], regSum1[1], regSum1[0], regSum1[1]); + MicroAPI::Cast(regSumBF16[0], regSum0[1], maskAllB32); + MicroAPI::Cast(regSumBF16[1], regSum1[1], maskAllB32); + MicroAPI::Cast(regSumBF16[0], regSum0[0], maskAllB32); + MicroAPI::Cast(regSumBF16[1], regSum1[0], maskAllB32); + + liV2Vector1::FloatX2ToSortableKey(regOut[0], regOut[1], + regSumBF16[0], regSumBF16[1], bf16Ctx, maskAllB16); + MicroAPI::StoreAlign(out0, regOut[0], maskAllB16); + MicroAPI::StoreAlign(out1, regOut[1], maskAllB16); +} + +__aicore__ inline void MulWeightAndReduceSum2(const LocalTensor &out_, // out [2, S2Base] [128 ] + uint32_t outStride, + const LocalTensor &qk_, // q*k^t [2, G, S2Base] [64 128] + uint32_t qkVLStride, uint32_t qkStride, + const LocalTensor &weight_, // w [2, G] [64 ] + uint32_t weightStride, const LocalTensor &weightTemp_, + const LocalTensor &kScale_, // kScale [S2Base] [128 ] + uint32_t kScaleStride, + const LocalTensor &qScale_, // qScale [2, G] [64 ] + uint32_t qScaleStride, + const int gSize) // G 64 +{ + __ubuf__ float *weight0 = (__ubuf__ float *)weight_.GetPhyAddr(); + __ubuf__ float *weightTemp = (__ubuf__ float *)weightTemp_.GetPhyAddr(); + __ubuf__ float *qScale0 = (__ubuf__ float *)qScale_.GetPhyAddr(); + __ubuf__ float *kScale0 = (__ubuf__ float *)kScale_.GetPhyAddr(); + __ubuf__ float *qk0 = (__ubuf__ float *)qk_.GetPhyAddr(); + __ubuf__ uint16_t *out0 = (__ubuf__ uint16_t *)out_.GetPhyAddr(); + + __ubuf__ float *weight1 = weight0 + weightStride; + __ubuf__ float *qScale1 = qScale0 + qScaleStride; + __ubuf__ float *qk1 = qk0 + qkStride; + // kScaleStride is zero + __ubuf__ uint16_t *out1 = out0 + outStride; + + MulWeightAndReduceSum2F32VF(out0, out1, qk0, qk1, qkVLStride, weight0, weight1, weightTemp, qScale0, qScale1, + kScale0, (uint16_t)gSize); +} + +__aicore__ inline void MulWeightAndReduceSum2(const LocalTensor &out_, // out [2, S2Base] [128 ] + uint32_t outStride, + const LocalTensor &qk_, // q*k^t [2, G, S2Base] [64 128] + uint32_t qkVLStride, + uint32_t qkStride, + const LocalTensor &weight_, // w [2, G] [64 ] + uint32_t weightStride, + const LocalTensor &weightTemp_, + const LocalTensor &kScale_, // kScale [S2Base] [128 ] + uint32_t kScaleStride, + const LocalTensor &qScale_, // qScale [2, G] [64 ] + uint32_t qScaleStride, + const int gSize) // G 64 +{ + __ubuf__ half *weight0 = (__ubuf__ half *)weight_.GetPhyAddr(); + __ubuf__ float *weightTemp = (__ubuf__ float *)weightTemp_.GetPhyAddr(); + __ubuf__ half *qScale0 = (__ubuf__ half *)qScale_.GetPhyAddr(); + __ubuf__ half *kScale0 = (__ubuf__ half *)kScale_.GetPhyAddr(); + __ubuf__ int32_t *qk0 = (__ubuf__ int32_t *)qk_.GetPhyAddr(); + __ubuf__ uint16_t *out0 = (__ubuf__ uint16_t *)out_.GetPhyAddr(); + + __ubuf__ half *weight1 = weight0 + weightStride; + __ubuf__ half *qScale1 = qScale0 + qScaleStride; + __ubuf__ int32_t *qk1 = qk0 + qkStride; + // kScaleStride is zero + __ubuf__ uint16_t *out1 = out0 + outStride; + + MulWeightAndReduceSum2Int32VF(out0, out1, qk0, qk1, qkVLStride, weight0, weight1, weightTemp, + qScale0, qScale1, kScale0, (uint16_t)gSize); +} + +// 计算S1=2 +// bfloat16 in uint16 out +__simd_vf__ void MulWeightAndReduceSum2B16VF(__ubuf__ uint16_t *out0, __ubuf__ uint16_t *out1, + __ubuf__ bfloat16_t *qk0, + __ubuf__ bfloat16_t *qk1, __ubuf__ float *weight0, + __ubuf__ float *weight1, + __ubuf__ float *weightTemp0, __ubuf__ float *weightTemp1, + __ubuf__ float *qScale0, __ubuf__ float *qScale1, + __ubuf__ float *kScale0, uint16_t gSize) +{ + MicroAPI::RegTensor regwBrc[2]; + MicroAPI::RegTensor regQK0[2]; + MicroAPI::RegTensor regQK1[2]; + MicroAPI::RegTensor regW[2]; + MicroAPI::RegTensor regQKB16[2]; + + MicroAPI::RegTensor regQScale[2]; + MicroAPI::RegTensor regKScale[2]; + MicroAPI::RegTensor regSum0[2]; + MicroAPI::RegTensor regSum1[2]; + MicroAPI::MaskReg maskAllB32 = MicroAPI::CreateMask(); + MicroAPI::MaskReg maskAllB16 = MicroAPI::CreateMask(); + + liV2Vector1::FloatSortConstCtx bf16Ctx; + liV2Vector1::InitFloatSortConstCtx(bf16Ctx, maskAllB16); + + using CastTrait = MicroAPI::CastTrait; + static constexpr CastTrait castTraitB162B32_EVEN = {MicroAPI::RegLayout::ZERO, MicroAPI::SatMode::UNKNOWN, + MicroAPI::MaskMergeMode::ZEROING, RoundMode::UNKNOWN}; + static constexpr CastTrait castTraitB162B32_ODD = {MicroAPI::RegLayout::ONE, MicroAPI::SatMode::UNKNOWN, + MicroAPI::MaskMergeMode::ZEROING, RoundMode::UNKNOWN}; + + constexpr static MicroAPI::CastTrait castTraitF32ToF16_EVEN = { + MicroAPI::RegLayout::ZERO, MicroAPI::SatMode::NO_SAT, MicroAPI::MaskMergeMode::MERGING, RoundMode::CAST_ROUND}; + constexpr static MicroAPI::CastTrait castTraitF32ToF16_ODD = { + MicroAPI::RegLayout::ONE, MicroAPI::SatMode::NO_SAT, MicroAPI::MaskMergeMode::ZEROING, RoundMode::CAST_ROUND}; + + MicroAPI::LoadAlign(regW[0], weight0); + MicroAPI::LoadAlign(regW[1], weight1); + MicroAPI::LoadAlign(regQScale[0], qScale0); + MicroAPI::LoadAlign(regQScale[1], qScale1); + MicroAPI::Mul(regW[0], regW[0], regQScale[0], maskAllB32); + MicroAPI::Mul(regW[1], regW[1], regQScale[1], maskAllB32); + // 读写依赖，寄存器可以保序 + MicroAPI::StoreAlign(weightTemp0, regW[0], maskAllB32); + MicroAPI::StoreAlign(weightTemp1, regW[1], maskAllB32); + liV2Vector1::DuplicateZero(regSum0, maskAllB32); + liV2Vector1::DuplicateZero(regSum1, maskAllB32); + + // interleave load + MicroAPI::LoadAlign(regKScale[0], regKScale[1], kScale0); + + for (uint16_t i = (uint16_t)(0); i < gSize; i++) { + // RowStride是256, 行都落在一个bank上 + MicroAPI::LoadAlign(regQKB16[0], qk0 + 256 * i); + // RowStride是256, 行都落在一个bank上 + MicroAPI::LoadAlign(regQKB16[1], qk1 + 256 * i); + MicroAPI::LoadAlign(regwBrc[0], weightTemp0 + i); + MicroAPI::LoadAlign(regwBrc[1], weightTemp1 + i); + // interleave cast + MicroAPI::Cast(regQK0[0], regQKB16[0], maskAllB32); + MicroAPI::Cast(regQK0[1], regQKB16[0], maskAllB32); + MicroAPI::Cast(regQK1[0], regQKB16[1], maskAllB32); + MicroAPI::Cast(regQK1[1], regQKB16[1], maskAllB32); + MicroAPI::MulAddDst(regSum0[0], regQK0[0], regwBrc[0], maskAllB32); + MicroAPI::MulAddDst(regSum0[1], regQK0[1], regwBrc[0], maskAllB32); + MicroAPI::MulAddDst(regSum1[0], regQK1[0], regwBrc[1], maskAllB32); + MicroAPI::MulAddDst(regSum1[1], regQK1[1], regwBrc[1], maskAllB32); + } + + // Apply kScale scaling + MicroAPI::Mul(regSum0[0], regSum0[0], regKScale[0], maskAllB32); + MicroAPI::Mul(regSum0[1], regSum0[1], regKScale[1], maskAllB32); + MicroAPI::Mul(regSum1[0], regSum1[0], regKScale[0], maskAllB32); + MicroAPI::Mul(regSum1[1], regSum1[1], regKScale[1], maskAllB32); + + // Convert to bfloat16 and store output channel + MicroAPI::RegTensor regSumBF16[2]; + MicroAPI::RegTensor regOut[2]; + MicroAPI::Cast(regSumBF16[0], regSum0[1], maskAllB32); + MicroAPI::Cast(regSumBF16[1], regSum1[1], maskAllB32); + MicroAPI::Cast(regSumBF16[0], regSum0[0], maskAllB32); + MicroAPI::Cast(regSumBF16[1], regSum1[0], maskAllB32); + + liV2Vector1::FloatX2ToSortableKey(regOut[0], regOut[1], regSumBF16[0], regSumBF16[1], bf16Ctx, + maskAllB16); + MicroAPI::StoreAlign(out0, regOut[0], maskAllB16); + MicroAPI::StoreAlign(out1, regOut[1], maskAllB16); +} + +__aicore__ inline void MulWeightAndReduceSum2(const LocalTensor &out_, // out [2, S2Base] [128 ] + uint32_t outStride, + const LocalTensor &qk_, // q*k^t [2, G, S2Base] [64 128] + uint32_t qkVLStride, + uint32_t qkStride, // gSize * 256 + const LocalTensor &weight_, // w [2, G] [64 ] + uint32_t weightStride, const LocalTensor &weightTemp_, + const LocalTensor &kScale_, // kScale [S2Base] [128 ] + uint32_t kScaleStride, + const LocalTensor &qScale_, // qScale [2, G] [64 ] + uint32_t qScaleStride, + const int gSize) // G 64 +{ + __ubuf__ float *weight0 = (__ubuf__ float *)weight_.GetPhyAddr(); + __ubuf__ float *weightTemp0 = (__ubuf__ float *)weightTemp_.GetPhyAddr(); + __ubuf__ float *qScale0 = (__ubuf__ float *)qScale_.GetPhyAddr(); + __ubuf__ float *kScale0 = (__ubuf__ float *)kScale_.GetPhyAddr(); + __ubuf__ bfloat16_t *qk0 = (__ubuf__ bfloat16_t *)qk_.GetPhyAddr(); + __ubuf__ uint16_t *out0 = (__ubuf__ uint16_t *)out_.GetPhyAddr(); + + __ubuf__ float *weightTemp1 = weightTemp0 + weightStride; + __ubuf__ float *weight1 = weight0 + weightStride; + __ubuf__ float *qScale1 = qScale0 + qScaleStride; + __ubuf__ bfloat16_t *qk1 = qk0 + qkStride; + // kScaleStride is zero + __ubuf__ uint16_t *out1 = out0 + outStride; + + MulWeightAndReduceSum2B16VF(out0, out1, qk0, qk1, weight0, weight1, weightTemp0, weightTemp1, qScale0, qScale1, + kScale0, (uint16_t)gSize); +} + +template +__aicore__ inline void BatchMulWeightAndReduceSum(const LocalTensor &out_, // out [S2Base] [128 ] + uint32_t outStride, + const LocalTensor &qk_, // q*k^t [G, S2Base] [64 128] + uint32_t qkVLStride, uint32_t qkStride, + const LocalTensor &weight_, // w [G] [64 ] + uint32_t weightStride, const LocalTensor &weightTemp_, + const LocalTensor &kScale_, // kScale [S2Base] [128 ] + uint32_t kScaleStride, + const LocalTensor &qScale_, // qScale [G] [64 ] + uint32_t qScaleStride, + const int gSize, // G 64 + const int batch) +{ + // 暂只支持这两种情况, 后续改成循环 + if (batch != 2 && batch != 1) { + return; + } + if (batch == 2) { + MulWeightAndReduceSum2(out_, outStride, qk_, qkVLStride, qkStride, weight_, weightStride, weightTemp_, kScale_, + kScaleStride, qScale_, qScaleStride, gSize); + } else { + MulWeightAndReduceSum(out_, qk_, qkVLStride, weight_, kScale_, qScale_, gSize); + } +} + +// per_tensor与MX共用的weight加权归约实现，WITH_SCALE控制是否额外应用scalar scale +// float in uint16 out +template +__simd_vf__ void MulWeightAndReduceSumOptionalScaleGSizeEvenVF(__ubuf__ uint16_t *out, __ubuf__ float *qk, + uint32_t qkVLStride, __ubuf__ float *weight, + float kScaleValue, float qScaleValue, uint16_t gSize) +{ + if constexpr (!WITH_SCALE) { + (void)kScaleValue; + (void)qScaleValue; + } + + MicroAPI::RegTensor regwBrc; + MicroAPI::RegTensor regQK[2]; + MicroAPI::RegTensor regW; + MicroAPI::RegTensor regSum0[2]; + MicroAPI::RegTensor regSum1[2]; + MicroAPI::MaskReg maskAllB32 = MicroAPI::CreateMask(); + MicroAPI::MaskReg maskAllB16 = MicroAPI::CreateMask(); + + liV2Vector1::FloatSortConstCtx bf16Ctx; + liV2Vector1::InitFloatSortConstCtx(bf16Ctx, maskAllB16); + + constexpr static MicroAPI::CastTrait castTraitF32ToF16_EVEN = { + MicroAPI::RegLayout::ZERO, MicroAPI::SatMode::NO_SAT, MicroAPI::MaskMergeMode::MERGING, RoundMode::CAST_ROUND}; + constexpr static MicroAPI::CastTrait castTraitF32ToF16_ODD = { + MicroAPI::RegLayout::ONE, MicroAPI::SatMode::NO_SAT, MicroAPI::MaskMergeMode::ZEROING, RoundMode::CAST_ROUND}; + + MicroAPI::LoadAlign(regW, weight); + if constexpr (WITH_SCALE) { + MicroAPI::Muls(regW, regW, qScaleValue, maskAllB32); + } + + liV2Vector1::DuplicateZero(regSum0, maskAllB32); + liV2Vector1::DuplicateZero(regSum1, maskAllB32); + + // unroll2 + for (uint16_t i = (uint16_t)(0); i < gSize; i += 2) { + MicroAPI::LoadAlign(regQK[0], qk + 128 * i); // RowStride是128, 行都落在一个bank上 + MicroAPI::LoadAlign(regQK[1], qk + 128 * i + qkVLStride); + liV2Vector1::BroadcastLane(regwBrc, regW, i); + liV2Vector1::WeightedAccum(regSum0, regQK, regwBrc, maskAllB32); + + MicroAPI::LoadAlign(regQK[0], qk + 128 * i + 128); + MicroAPI::LoadAlign(regQK[1], qk + 128 * i + 128 + qkVLStride); + liV2Vector1::BroadcastLane(regwBrc, regW, i + 1); + liV2Vector1::WeightedAccum(regSum1, regQK, regwBrc, maskAllB32); + } + + MicroAPI::Add(regSum0[0], regSum0[0], regSum1[0], maskAllB32); + MicroAPI::Add(regSum0[1], regSum0[1], regSum1[1], maskAllB32); + + if constexpr (WITH_SCALE) { + MicroAPI::Muls(regSum0[0], regSum0[0], kScaleValue, maskAllB32); + MicroAPI::Muls(regSum0[1], regSum0[1], kScaleValue, maskAllB32); + } + + MicroAPI::RegTensor regSumBF16; + // interleave cast ==> regSum[1] high regSum[0] low + MicroAPI::DeInterleave(regSum0[0], regSum0[1], regSum0[0], regSum0[1]); + MicroAPI::Cast(regSumBF16, regSum0[1], maskAllB32); + MicroAPI::Cast(regSumBF16, regSum0[0], maskAllB32); + + MicroAPI::RegTensor regOut; + liV2Vector1::FloatToSortableKey(regOut, regSumBF16, bf16Ctx, maskAllB16); + // normal store + MicroAPI::StoreAlign(out, regOut, maskAllB16); +} + +template +__simd_vf__ void MulWeightAndReduceSumOptionalScaleGSizeOddVF(__ubuf__ uint16_t *out, __ubuf__ float *qk, + uint32_t qkVLStride, __ubuf__ float *weight, + float kScaleValue, float qScaleValue, uint16_t gSize) +{ + if constexpr (!WITH_SCALE) { + (void)kScaleValue; + (void)qScaleValue; + } + + MicroAPI::RegTensor regwBrc; + MicroAPI::RegTensor regQK[2]; + MicroAPI::RegTensor regW; + MicroAPI::RegTensor regSum0[2]; + MicroAPI::RegTensor regSum1[2]; + MicroAPI::MaskReg maskAllB32 = MicroAPI::CreateMask(); + MicroAPI::MaskReg maskAllB16 = MicroAPI::CreateMask(); + + liV2Vector1::FloatSortConstCtx bf16Ctx; + liV2Vector1::InitFloatSortConstCtx(bf16Ctx, maskAllB16); + + constexpr static MicroAPI::CastTrait castTraitF32ToF16_EVEN = { + MicroAPI::RegLayout::ZERO, MicroAPI::SatMode::NO_SAT, MicroAPI::MaskMergeMode::MERGING, RoundMode::CAST_ROUND}; + constexpr static MicroAPI::CastTrait castTraitF32ToF16_ODD = { + MicroAPI::RegLayout::ONE, MicroAPI::SatMode::NO_SAT, MicroAPI::MaskMergeMode::ZEROING, RoundMode::CAST_ROUND}; + + MicroAPI::LoadAlign(regW, weight); + if constexpr (WITH_SCALE) { + MicroAPI::Muls(regW, regW, qScaleValue, maskAllB32); + } + + liV2Vector1::DuplicateZero(regSum0, maskAllB32); + liV2Vector1::DuplicateZero(regSum1, maskAllB32); + + // unroll2 + for (uint16_t i = (uint16_t)(0); (uint16_t)(i + 1) < gSize; i += 2) { + MicroAPI::LoadAlign(regQK[0], qk + 128 * i); // RowStride是128, 行都落在一个bank上 + MicroAPI::LoadAlign(regQK[1], qk + 128 * i + qkVLStride); + liV2Vector1::BroadcastLane(regwBrc, regW, i); + liV2Vector1::WeightedAccum(regSum0, regQK, regwBrc, maskAllB32); + + MicroAPI::LoadAlign(regQK[0], qk + 128 * i + 128); + MicroAPI::LoadAlign(regQK[1], qk + 128 * i + 128 + qkVLStride); + liV2Vector1::BroadcastLane(regwBrc, regW, i + 1); + liV2Vector1::WeightedAccum(regSum1, regQK, regwBrc, maskAllB32); + } + + MicroAPI::LoadAlign(regQK[0], qk + 128 * (gSize - 1)); // RowStride是128, 行都落在一个bank上 + MicroAPI::LoadAlign(regQK[1], qk + 128 * (gSize - 1) + qkVLStride); + liV2Vector1::BroadcastLane(regwBrc, regW, gSize - 1); + liV2Vector1::WeightedAccum(regSum0, regQK, regwBrc, maskAllB32); + + MicroAPI::Add(regSum0[0], regSum0[0], regSum1[0], maskAllB32); + MicroAPI::Add(regSum0[1], regSum0[1], regSum1[1], maskAllB32); + + if constexpr (WITH_SCALE) { + MicroAPI::Muls(regSum0[0], regSum0[0], kScaleValue, maskAllB32); + MicroAPI::Muls(regSum0[1], regSum0[1], kScaleValue, maskAllB32); + } + + MicroAPI::RegTensor regSumBF16; + // interleave cast ==> regSum[1] high regSum[0] low + MicroAPI::DeInterleave(regSum0[0], regSum0[1], regSum0[0], regSum0[1]); + MicroAPI::Cast(regSumBF16, regSum0[1], maskAllB32); + MicroAPI::Cast(regSumBF16, regSum0[0], maskAllB32); + + MicroAPI::RegTensor regOut; + liV2Vector1::FloatToSortableKey(regOut, regSumBF16, bf16Ctx, maskAllB16); + // normal store + MicroAPI::StoreAlign(out, regOut, maskAllB16); +} + +template +__aicore__ inline void MulWeightAndReduceSumOptionalScaleImpl( + const LocalTensor &out_, // out [S2Base] [128 ] + const LocalTensor &qk_, // q*k^t [G, S2Base] [64 128] + const uint32_t qkVLStride, + const LocalTensor &weight_, // w [G] [64 ] + const float kScaleValue, // kScale scalar + const float qScaleValue, // qScale scalar + const int gSize) // G 64 +{ + __ubuf__ uint16_t *out = (__ubuf__ uint16_t *)out_.GetPhyAddr(); + __ubuf__ float *weight = (__ubuf__ float *)weight_.GetPhyAddr(); + __ubuf__ float *qk = (__ubuf__ float *)qk_.GetPhyAddr(); + if (gSize % 2 == 0) { + MulWeightAndReduceSumOptionalScaleGSizeEvenVF(out, qk, qkVLStride, weight, kScaleValue, qScaleValue, + (uint16_t)gSize); + } else { + MulWeightAndReduceSumOptionalScaleGSizeOddVF(out, qk, qkVLStride, weight, kScaleValue, qScaleValue, + (uint16_t)gSize); + } +} + +__aicore__ inline void MulWeightAndReduceSumPerTensor(const LocalTensor &out_, // out [S2Base] [128 ] + const LocalTensor &qk_, // q*k^t [G, S2Base] [64 128] + const uint32_t qkVLStride, + const LocalTensor &weight_, // w [G] [64 ] + const float kScaleValue, // kScale scalar + const float qScaleValue, // qScale scalar + const int gSize) // G 64 +{ + MulWeightAndReduceSumOptionalScaleImpl(out_, qk_, qkVLStride, weight_, kScaleValue, qScaleValue, gSize); +} + +__aicore__ inline void MulWeightAndReduceSumMX(const LocalTensor &out_, // out [S2Base] [128 ] + const LocalTensor &qk_, // q*k^t [G, S2Base] [64 128] + const uint32_t qkVLStride, + const LocalTensor &weight_, // w [G] [64 ] + const int gSize) // G 64 +{ + MulWeightAndReduceSumOptionalScaleImpl(out_, qk_, qkVLStride, weight_, 1.0f, 1.0f, gSize); +} + +// 计算S1=2 +// float in uint16 out +template +__simd_vf__ void MulWeightAndReduceSumOptionalScale2VF(__ubuf__ uint16_t *out0, __ubuf__ uint16_t *out1, + __ubuf__ float *qk0, __ubuf__ float *qk1, uint32_t qkVLStride, + __ubuf__ float *weight0, __ubuf__ float *weight1, + __ubuf__ float *weightTemp, float kScaleValue, float qScaleValue, + uint16_t gSize) +{ + if constexpr (!WITH_SCALE) { + (void)kScaleValue; + (void)qScaleValue; + } + + MicroAPI::RegTensor regwBrc[2]; + MicroAPI::RegTensor regQK0[2]; + MicroAPI::RegTensor regQK1[2]; + MicroAPI::RegTensor regW[2]; + + MicroAPI::RegTensor regSum0[2]; + MicroAPI::RegTensor regSum1[2]; + MicroAPI::MaskReg maskAllB32 = MicroAPI::CreateMask(); + MicroAPI::MaskReg maskAllB16 = MicroAPI::CreateMask(); + + liV2Vector1::FloatSortConstCtx bf16Ctx; + liV2Vector1::InitFloatSortConstCtx(bf16Ctx, maskAllB16); + + constexpr static MicroAPI::CastTrait castTraitF32ToF16_EVEN = { + MicroAPI::RegLayout::ZERO, MicroAPI::SatMode::NO_SAT, MicroAPI::MaskMergeMode::MERGING, RoundMode::CAST_ROUND}; + constexpr static MicroAPI::CastTrait castTraitF32ToF16_ODD = { + MicroAPI::RegLayout::ONE, MicroAPI::SatMode::NO_SAT, MicroAPI::MaskMergeMode::ZEROING, RoundMode::CAST_ROUND}; + + MicroAPI::LoadAlign(regW[0], weight0); + MicroAPI::LoadAlign(regW[1], weight1); + if constexpr (WITH_SCALE) { + MicroAPI::Muls(regW[0], regW[0], qScaleValue, maskAllB32); + MicroAPI::Muls(regW[1], regW[1], qScaleValue, maskAllB32); + } + // regW[0]与weight1混合使用 + MicroAPI::StoreAlign(weightTemp, regW[1], maskAllB32); + MicroAPI::LocalMemBar(); + liV2Vector1::DuplicateZero(regSum0, maskAllB32); + liV2Vector1::DuplicateZero(regSum1, maskAllB32); + + for (uint16_t i = (uint16_t)(0); i < gSize; i++) { + MicroAPI::LoadAlign(regQK0[0], qk0 + 128 * i); + MicroAPI::LoadAlign(regQK0[1], qk0 + 128 * i + qkVLStride); + MicroAPI::LoadAlign(regQK1[0], qk1 + 128 * i); + MicroAPI::LoadAlign(regQK1[1], qk1 + 128 * i + qkVLStride); + // 混合使用对整体性能更好 + liV2Vector1::BroadcastLane(regwBrc[0], regW[0], i); + // Weight无bank冲突，用LoadAlign来提取weight标量 + // 地址空间处理：原 BroadcastLane(ptr) 内联为 LoadAlign BRC，避免 __ubuf__ 传给 __local_mem__ 参数 + MicroAPI::LoadAlign(regwBrc[1], weightTemp + i); + MicroAPI::Relu(regQK0[0], regQK0[0], maskAllB32); + MicroAPI::Relu(regQK0[1], regQK0[1], maskAllB32); + MicroAPI::Relu(regQK1[0], regQK1[0], maskAllB32); + MicroAPI::Relu(regQK1[1], regQK1[1], maskAllB32); + MicroAPI::MulAddDst(regSum0[0], regQK0[0], regwBrc[0], maskAllB32); + MicroAPI::MulAddDst(regSum0[1], regQK0[1], regwBrc[0], maskAllB32); + MicroAPI::MulAddDst(regSum1[0], regQK1[0], regwBrc[1], maskAllB32); + MicroAPI::MulAddDst(regSum1[1], regQK1[1], regwBrc[1], maskAllB32); + } + + if constexpr (WITH_SCALE) { + MicroAPI::Muls(regSum0[0], regSum0[0], kScaleValue, maskAllB32); + MicroAPI::Muls(regSum0[1], regSum0[1], kScaleValue, maskAllB32); + MicroAPI::Muls(regSum1[0], regSum1[0], kScaleValue, maskAllB32); + MicroAPI::Muls(regSum1[1], regSum1[1], kScaleValue, maskAllB32); + } + + // Convert to bfloat16 and store output channel + MicroAPI::RegTensor regSumBF16[2]; + MicroAPI::RegTensor regOut[2]; + MicroAPI::DeInterleave(regSum0[0], regSum0[1], regSum0[0], regSum0[1]); + MicroAPI::DeInterleave(regSum1[0], regSum1[1], regSum1[0], regSum1[1]); + MicroAPI::Cast(regSumBF16[0], regSum0[1], maskAllB32); + MicroAPI::Cast(regSumBF16[1], regSum1[1], maskAllB32); + MicroAPI::Cast(regSumBF16[0], regSum0[0], maskAllB32); + MicroAPI::Cast(regSumBF16[1], regSum1[0], maskAllB32); + + liV2Vector1::FloatX2ToSortableKey(regOut[0], regOut[1], regSumBF16[0], regSumBF16[1], bf16Ctx, + maskAllB16); + MicroAPI::StoreAlign(out0, regOut[0], maskAllB16); + MicroAPI::StoreAlign(out1, regOut[1], maskAllB16); +} + +template +__aicore__ inline void MulWeightAndReduceSumOptionalScale2Impl( + const LocalTensor &out_, // out [2, S2Base] [128 ] + uint32_t outStride, + const LocalTensor &qk_, // q*k^t [2, G, S2Base] [64 128] + uint32_t qkVLStride, uint32_t qkStride, + const LocalTensor &weight_, // w [2, G] [64 ] + uint32_t weightStride, const LocalTensor &weightTemp_, + const float kScaleValue, // kScale scalar + const float qScaleValue, // qScale scalar for batch 0和1 + const int gSize) // G 64 +{ + __ubuf__ float *weight0 = (__ubuf__ float *)weight_.GetPhyAddr(); + __ubuf__ float *weightTemp = (__ubuf__ float *)weightTemp_.GetPhyAddr(); + __ubuf__ float *qk0 = (__ubuf__ float *)qk_.GetPhyAddr(); + __ubuf__ uint16_t *out0 = (__ubuf__ uint16_t *)out_.GetPhyAddr(); + + __ubuf__ float *weight1 = weight0 + weightStride; + __ubuf__ float *qk1 = qk0 + qkStride; + __ubuf__ uint16_t *out1 = out0 + outStride; + + MulWeightAndReduceSumOptionalScale2VF(out0, out1, qk0, qk1, qkVLStride, weight0, weight1, weightTemp, + kScaleValue, qScaleValue, (uint16_t)gSize); +} + +__aicore__ inline void MulWeightAndReduceSumPerTensor2( + const LocalTensor &out_, // out [2, S2Base] [128 ] + uint32_t outStride, + const LocalTensor &qk_, // q*k^t [2, G, S2Base] [64 128] + uint32_t qkVLStride, uint32_t qkStride, + const LocalTensor &weight_, // w [2, G] [64 ] + uint32_t weightStride, const LocalTensor &weightTemp_, + const float kScaleValue, // kScale scalar + const float qScaleValue, // qScale scalar for batch 0和1 + const int gSize) // G 64 +{ + MulWeightAndReduceSumOptionalScale2Impl(out_, outStride, qk_, qkVLStride, qkStride, weight_, weightStride, + weightTemp_, kScaleValue, qScaleValue, gSize); +} + +__aicore__ inline void MulWeightAndReduceSumMX2(const LocalTensor &out_, // out [2, S2Base] [128 ] + uint32_t outStride, + const LocalTensor &qk_, // q*k^t [2, G, S2Base] [64 128] + uint32_t qkVLStride, uint32_t qkStride, + const LocalTensor &weight_, // w [2, G] [64 ] + uint32_t weightStride, const LocalTensor &weightTemp_, + const int gSize) // G 64 +{ + MulWeightAndReduceSumOptionalScale2Impl(out_, outStride, qk_, qkVLStride, qkStride, weight_, weightStride, + weightTemp_, 1.0f, 1.0f, gSize); +} + +__simd_callee__ inline void CastWeightToBf16(AscendC::MicroAPI::RegTensor &dst, __ubuf__ float *src, + AscendC::MicroAPI::MaskReg &maskAllB32) +{ + using CastTrait = AscendC::MicroAPI::CastTrait; + static constexpr CastTrait castTraitF32ToBf16 = {AscendC::MicroAPI::RegLayout::ZERO, + AscendC::MicroAPI::SatMode::NO_SAT, + AscendC::MicroAPI::MaskMergeMode::ZEROING, RoundMode::CAST_RINT}; + AscendC::MicroAPI::RegTensor regWeightF32; + AscendC::MicroAPI::LoadAlign(regWeightF32, src); + AscendC::MicroAPI::Cast(dst, regWeightF32, maskAllB32); + // Cast结果按B32 lane落位，在目标寄存器内压紧为连续BF16，供BroadcastLane按元素索引。 + AscendC::MicroAPI::Pack( + (AscendC::MicroAPI::RegTensor &)dst, (AscendC::MicroAPI::RegTensor &)dst); +} + +__simd_vf__ void MulWeightAndReduceSumMXFP4VF(__ubuf__ uint16_t *out, __ubuf__ bfloat16_t *qk, __ubuf__ float *weight, + uint16_t gSize) +{ + constexpr uint32_t BF16_QK_ROW_STRIDE = UB_BANK_DEPTH_STRIDE / sizeof(bfloat16_t); + AscendC::MicroAPI::RegTensor regQK; + AscendC::MicroAPI::RegTensor regWeight; + AscendC::MicroAPI::RegTensor regWeightBrc; + AscendC::MicroAPI::RegTensor regSum; + AscendC::MicroAPI::MaskReg maskAllB16 = + AscendC::MicroAPI::CreateMask(); + AscendC::MicroAPI::MaskReg maskAllB32 = AscendC::MicroAPI::CreateMask(); + + liV2Vector1::FloatSortConstCtx bf16Ctx; + liV2Vector1::InitFloatSortConstCtx(bf16Ctx, maskAllB16); + + CastWeightToBf16(regWeight, weight, maskAllB32); + AscendC::MicroAPI::Duplicate(regSum, bfloat16_t(0.0f), maskAllB16); + + for (uint16_t i = 0; i < gSize; i++) { + AscendC::MicroAPI::LoadAlign(regQK, qk + BF16_QK_ROW_STRIDE * i); + liV2Vector1::BroadcastLane(regWeightBrc, regWeight, i); + AscendC::MicroAPI::MulAddDst(regSum, regQK, regWeightBrc, maskAllB16); + } + + AscendC::MicroAPI::RegTensor regOut; + liV2Vector1::FloatToSortableKey(regOut, regSum, bf16Ctx, maskAllB16); + AscendC::MicroAPI::StoreAlign(out, regOut, maskAllB16); +} + +__aicore__ inline void MulWeightAndReduceSumMXFP4(const LocalTensor &out_, const LocalTensor &qk_, + const uint32_t qkVLStride, const LocalTensor &weight_, + const int gSize) +{ + (void)qkVLStride; + __ubuf__ uint16_t *out = (__ubuf__ uint16_t *)out_.GetPhyAddr(); + __ubuf__ bfloat16_t *qk = (__ubuf__ bfloat16_t *)qk_.GetPhyAddr(); + __ubuf__ float *weight = (__ubuf__ float *)weight_.GetPhyAddr(); + MulWeightAndReduceSumMXFP4VF(out, qk, weight, static_cast(gSize)); +} + +__simd_vf__ void MulWeightAndReduceSumMXFP4TwoRowsVF(__ubuf__ uint16_t *out0, __ubuf__ uint16_t *out1, + __ubuf__ bfloat16_t *qk0, __ubuf__ bfloat16_t *qk1, + __ubuf__ float *weight0, __ubuf__ float *weight1, uint16_t gSize) +{ + constexpr uint32_t BF16_QK_ROW_STRIDE = UB_BANK_DEPTH_STRIDE / sizeof(bfloat16_t); + AscendC::MicroAPI::RegTensor regQK0; + AscendC::MicroAPI::RegTensor regQK1; + AscendC::MicroAPI::RegTensor regWeight[2]; + AscendC::MicroAPI::RegTensor regWeightBrc[2]; + AscendC::MicroAPI::RegTensor regSum[2]; + AscendC::MicroAPI::MaskReg maskAllB16 = + AscendC::MicroAPI::CreateMask(); + AscendC::MicroAPI::MaskReg maskAllB32 = + AscendC::MicroAPI::CreateMask(); + + liV2Vector1::FloatSortConstCtx bf16Ctx; + liV2Vector1::InitFloatSortConstCtx(bf16Ctx, maskAllB16); + + CastWeightToBf16(regWeight[0], weight0, maskAllB32); + CastWeightToBf16(regWeight[1], weight1, maskAllB32); + AscendC::MicroAPI::Duplicate(regSum[0], bfloat16_t(0.0f), maskAllB16); + AscendC::MicroAPI::Duplicate(regSum[1], bfloat16_t(0.0f), maskAllB16); + + for (uint16_t i = 0; i < gSize; i++) { + AscendC::MicroAPI::LoadAlign(regQK0, qk0 + BF16_QK_ROW_STRIDE * i); + AscendC::MicroAPI::LoadAlign(regQK1, qk1 + BF16_QK_ROW_STRIDE * i); + liV2Vector1::BroadcastLane(regWeightBrc[0], regWeight[0], i); + liV2Vector1::BroadcastLane(regWeightBrc[1], regWeight[1], i); + AscendC::MicroAPI::MulAddDst(regSum[0], regQK0, regWeightBrc[0], maskAllB16); + AscendC::MicroAPI::MulAddDst(regSum[1], regQK1, regWeightBrc[1], maskAllB16); + } + + AscendC::MicroAPI::RegTensor regOut[2]; + liV2Vector1::FloatX2ToSortableKey(regOut[0], regOut[1], regSum[0], regSum[1], bf16Ctx, maskAllB16); + AscendC::MicroAPI::StoreAlign(out0, regOut[0], maskAllB16); + AscendC::MicroAPI::StoreAlign(out1, regOut[1], maskAllB16); +} + +__aicore__ inline void MulWeightAndReduceSumMXFP4TwoRows(const LocalTensor &out_, uint32_t outStride, + const LocalTensor &qk_, uint32_t qkVLStride, + uint32_t qkStride, const LocalTensor &weight_, + uint32_t weightStride, const int gSize) +{ + (void)qkVLStride; + __ubuf__ uint16_t *out0 = (__ubuf__ uint16_t *)out_.GetPhyAddr(); + __ubuf__ uint16_t *out1 = out0 + outStride; + __ubuf__ bfloat16_t *qk0 = (__ubuf__ bfloat16_t *)qk_.GetPhyAddr(); + __ubuf__ bfloat16_t *qk1 = qk0 + qkStride; + __ubuf__ float *weight0 = (__ubuf__ float *)weight_.GetPhyAddr(); + __ubuf__ float *weight1 = weight0 + weightStride; + MulWeightAndReduceSumMXFP4TwoRowsVF(out0, out1, qk0, qk1, weight0, weight1, static_cast(gSize)); +} + +template +__aicore__ inline void BatchMulWeightAndReduceSumMXFP4(const LocalTensor &out_, uint32_t outStride, + const LocalTensor &qk_, uint32_t qkVLStride, + uint32_t qkStride, const LocalTensor &weight_, + uint32_t weightStride, const int gSize, const int batch) +{ + static_assert(std::is_same_v); + static_assert(std::is_same_v); + static_assert(std::is_same_v); + if (batch == 2) { + MulWeightAndReduceSumMXFP4TwoRows(out_, outStride, qk_, qkVLStride, qkStride, weight_, weightStride, gSize); + } else if (batch == 1) { + MulWeightAndReduceSumMXFP4(out_, qk_, qkVLStride, weight_, gSize); + } +} +template +__aicore__ inline void BatchMulWeightAndReduceSumMX(const LocalTensor &out_, + uint32_t outStride, // out [S2Base] [128 ] + const LocalTensor &qk_, // q*k^t [G, S2Base] [64 128] + uint32_t qkVLStride, uint32_t qkStride, + const LocalTensor &weight_, // w [G] [64 ] + uint32_t weightStride, const LocalTensor &weightTemp_, + const int gSize, const int batch) +{ + // 暂只支持这两种情况, 后续改成循环 + if (batch != 2 && batch != 1) { + return; + } + if (batch == 2) { + MulWeightAndReduceSumMX2(out_, outStride, qk_, qkVLStride, qkStride, + weight_, weightStride, weightTemp_, gSize); + } else { + MulWeightAndReduceSumMX(out_, qk_, qkVLStride, weight_, gSize); + } +} + +template +__aicore__ inline void BatchMulWeightAndReduceSumPerTensor( + const LocalTensor &out_, + uint32_t outStride, // out [S2Base] [128 ] + const LocalTensor &qk_, // q*k^t [G, S2Base] [64 128] + uint32_t qkVLStride, uint32_t qkStride, + const LocalTensor &weight_, // w [G] [64 ] + uint32_t weightStride, const LocalTensor &weightTemp_, const float kScaleValue, const float qScaleValue, + const int gSize, const int batch) +{ + // 暂只支持这两种情况, 后续改成循环 + if (batch != 2 && batch != 1) { + return; + } + if (batch == 2) { + MulWeightAndReduceSumPerTensor2(out_, outStride, qk_, qkVLStride, qkStride, weight_, weightStride, weightTemp_, + kScaleValue, qScaleValue, gSize); + } else { + MulWeightAndReduceSumPerTensor(out_, qk_, qkVLStride, weight_, kScaleValue, qScaleValue, gSize); + } +} + +} // namespace vector1 + +#endif // QUANT_LIGHTNING_INDEXER_V2_VECTOR1_H diff --git a/xllm_ops/attention/quant_lightning_indexer_v2/op_kernel/arch35/vf/vf_topk.h b/xllm_ops/attention/quant_lightning_indexer_v2/op_kernel/arch35/vf/vf_topk.h new file mode 100644 index 0000000..aab1e2b --- /dev/null +++ b/xllm_ops/attention/quant_lightning_indexer_v2/op_kernel/arch35/vf/vf_topk.h @@ -0,0 +1,609 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. +  */ + +/*! +* \file vf_topk.h +* \brief +*/ + +#ifndef VF_TOP_K_H +#define VF_TOP_K_H + +namespace topkb32 { +__simd_callee__ inline void StoreHistogramResult(__ubuf__ uint32_t* histogramsBuf, + MicroAPI::RegTensor& cout0, + MicroAPI::RegTensor& cout1, + MicroAPI::MaskReg& pregB16, + MicroAPI::MaskReg& pregB32) +{ + MicroAPI::RegTensor cout0U32Even; + MicroAPI::RegTensor cout0U32Odd; + MicroAPI::RegTensor cout1U32Even; + MicroAPI::RegTensor cout1U32Odd; + + static constexpr MicroAPI::CastTrait CAST_TRAIT_UINT16_TOUINT32_EVEN = {MicroAPI::RegLayout::ZERO, + MicroAPI::SatMode::UNKNOWN, MicroAPI::MaskMergeMode::ZEROING, RoundMode::UNKNOWN}; + + static constexpr MicroAPI::CastTrait CAST_TRAIT_UINT16_TOUINT32_ODD = {MicroAPI::RegLayout::ONE, + MicroAPI::SatMode::UNKNOWN, MicroAPI::MaskMergeMode::ZEROING, RoundMode::UNKNOWN}; + + MicroAPI::Cast(cout0U32Even, cout0, pregB16); + MicroAPI::Cast(cout0U32Odd, cout0, pregB16); + MicroAPI::Cast(cout1U32Even, cout1, pregB16); + MicroAPI::Cast(cout1U32Odd, cout1, pregB16); + + MicroAPI::StoreAlign(histogramsBuf, + cout0U32Even, cout0U32Odd, pregB32); + MicroAPI::StoreAlign(histogramsBuf + 128, + cout1U32Even, cout1U32Odd, pregB32); +} + +__simd_callee__ inline void FindTargetBinAndUpdateNextK(__ubuf__ uint32_t* idxBuf, + __ubuf__ uint32_t* nkValueBuf, + __ubuf__ uint32_t* histogramsBuf, + MicroAPI::RegTensor& btmK, + MicroAPI::MaskReg& pregB32) +{ + MicroAPI::ClearSpr(); + + MicroAPI::UnalignRegForStore alignIdx; + + for (uint16_t i = 0; i < (uint16_t)(4); ++i) { + MicroAPI::RegTensor idxC; + MicroAPI::RegTensor cout; + MicroAPI::RegTensor sqzIdx; + + MicroAPI::MaskReg pregGE = MicroAPI::CreateMask(); + + MicroAPI::Arange(idxC, i * 64); + MicroAPI::LoadAlign(cout, histogramsBuf + i * 64); + MicroAPI::Compare(pregGE, cout, btmK, pregB32); + MicroAPI::Squeeze( + sqzIdx, (MicroAPI::RegTensor&)idxC, pregGE); + MicroAPI::StoreUnAlign(idxBuf, sqzIdx, alignIdx); + } + MicroAPI::StoreUnAlignPost(idxBuf, alignIdx); + + MicroAPI::LocalMemBar(); + + MicroAPI::RegTensor idx; + MicroAPI::LoadAlign(idx, idxBuf); + + MicroAPI::RegTensor idxAll1; + MicroAPI::RegTensor idxPrev; + MicroAPI::RegTensor prevBinValue; + MicroAPI::Duplicate(idxAll1, 1); + + MicroAPI::RegTensor zeroAll; + MicroAPI::Duplicate(zeroAll, 0); + + MicroAPI::MaskReg pregZero = MicroAPI::CreateMask(); + MicroAPI::Compare(pregZero, idx, zeroAll, pregB32); + MicroAPI::Sub(idxPrev, idx, (MicroAPI::RegTensor&)idxAll1, pregB32); + MicroAPI::ShiftRights(idxPrev, idxPrev, (int16_t)24, pregB32); + + MicroAPI::Gather(prevBinValue, histogramsBuf, idxPrev, pregB32); + MicroAPI::Select(prevBinValue, zeroAll, prevBinValue, pregZero); + + MicroAPI::RegTensor nextK; + MicroAPI::Sub(nextK, btmK, prevBinValue, pregB32); + MicroAPI::StoreAlign(nkValueBuf, nextK, pregB32); +} + +template +__simd_vf__ void HistogramsFirstVFImpl(__ubuf__ uint32_t* histogramsBuf, + __ubuf__ uint32_t* inputBuf, + uint16_t vfLoop, bool init) +{ + MicroAPI::MaskReg pregB32 = MicroAPI::CreateMask(); + MicroAPI::MaskReg pregB16 = MicroAPI::CreateMask(); + MicroAPI::MaskReg pregB8 = MicroAPI::CreateMask(); + + // 计算直方图cout0 0-127 cout1 128-255 + MicroAPI::RegTensor cout0; + MicroAPI::RegTensor cout1; + MicroAPI::Duplicate(cout0, 0); + MicroAPI::Duplicate(cout1, 0); + + MicroAPI::RegTensor vreg0; + MicroAPI::RegTensor vreg1; + MicroAPI::RegTensor vreg2; + MicroAPI::RegTensor vreg3; + + // 32bit 高16bit + MicroAPI::RegTensor vreg0U16; + // 32bit 低16bit + MicroAPI::RegTensor vreg1U16; + MicroAPI::RegTensor vreg2U16; + MicroAPI::RegTensor vreg3U16; + + for (uint16_t i = 0; i < vfLoop; ++i) { + MicroAPI::LoadAlign(vreg1U16, vreg0U16, inputBuf + i * 256); + MicroAPI::LoadAlign( + vreg3U16, vreg2U16, inputBuf + (i * 256) + 128); + + MicroAPI::DeInterleave(vreg1, vreg0, + (MicroAPI::RegTensor&)vreg0U16, + (MicroAPI::RegTensor&)vreg2U16); + + MicroAPI::Histograms(cout0, vreg0, pregB8); + MicroAPI::Histograms(cout1, vreg0, pregB8); + } + + StoreHistogramResult(histogramsBuf, cout0, cout1, pregB16, pregB32); +} + +__simd_vf__ void FindFirstTargetBinVFImpl(__ubuf__ uint32_t* idx0Buf, + __ubuf__ uint32_t* nkValueBuf, __ubuf__ uint32_t* + histogramsBuf, uint32_t bottomK) +{ + MicroAPI::MaskReg pregB32 = MicroAPI::CreateMask(); + + MicroAPI::RegTensor btmK; + MicroAPI::Duplicate(btmK, bottomK); + + FindTargetBinAndUpdateNextK(idx0Buf, nkValueBuf, histogramsBuf, btmK, pregB32); +} + +template +__simd_vf__ void HistogramsSecondVFImpl(__ubuf__ uint32_t* histogramsBuf, + __ubuf__ uint32_t* inputBuf, __ubuf__ uint32_t* idx0Buf, + uint16_t vfLoop, bool init) +{ + MicroAPI::MaskReg pregB32 = MicroAPI::CreateMask(); + MicroAPI::MaskReg pregB16 = MicroAPI::CreateMask(); + MicroAPI::MaskReg pregB8 = MicroAPI::CreateMask(); + + // 计算直方图0-127 128-255 + MicroAPI::RegTensor cout0; + MicroAPI::RegTensor cout1; + MicroAPI::Duplicate(cout0, 0); + MicroAPI::Duplicate(cout1, 0); + + MicroAPI::RegTensor idx0; + // 0x000000fc -> 0xfcfcfcfc + MicroAPI::LoadAlign(idx0, idx0Buf); + + MicroAPI::RegTensor vreg0U16; + MicroAPI::RegTensor vreg1U16; + MicroAPI::RegTensor vreg2U16; + MicroAPI::RegTensor vreg3U16; + + MicroAPI::RegTensor vreg0; + MicroAPI::RegTensor vreg1; + MicroAPI::RegTensor vreg2; + MicroAPI::RegTensor vreg3; + + for (uint16_t i = 0; i < vfLoop; ++i) { + MicroAPI::LoadAlign(vreg1U16, + vreg0U16, inputBuf + i * 256); + MicroAPI::LoadAlign(vreg3U16, + vreg2U16, inputBuf + (i * 256) + 128); + + MicroAPI::DeInterleave(vreg1, vreg0, + (MicroAPI::RegTensor&)vreg0U16, + (MicroAPI::RegTensor&)vreg2U16); + + MicroAPI::MaskReg pregEQ = MicroAPI::CreateMask(); + MicroAPI::Compare(pregEQ, vreg0, (MicroAPI::RegTensor&)idx0, pregB8); + + MicroAPI::Histograms(cout0, vreg1, pregEQ); + MicroAPI::Histograms(cout1, vreg1, pregEQ); + } + + StoreHistogramResult(histogramsBuf, cout0, cout1, pregB16, pregB32); +} + +// kValue新的bottomK +__simd_vf__ void FindSecondTargetBinVFImpl(__ubuf__ uint32_t* idx1Buf, + __ubuf__ uint32_t* nkValueBuf, __ubuf__ uint32_t* kValue, + __ubuf__ uint32_t* histogramsBuf) +{ + MicroAPI::MaskReg pregB32 = MicroAPI::CreateMask(); + + MicroAPI::RegTensor btmK1; + MicroAPI::LoadAlign(btmK1, kValue); + + FindTargetBinAndUpdateNextK(idx1Buf, nkValueBuf, histogramsBuf, btmK1, pregB32); +} + +template +__simd_vf__ void HistogramsThirdVFImpl(__ubuf__ uint32_t* histogramsBuf, + __ubuf__ uint32_t* inputBuf, __ubuf__ uint32_t* idx0Buf, + __ubuf__ uint32_t* idx1Buf, uint16_t vfLoop, bool init) +{ + MicroAPI::MaskReg pregB32 = MicroAPI::CreateMask(); + MicroAPI::MaskReg pregB16 = MicroAPI::CreateMask(); + MicroAPI::MaskReg pregB8 = MicroAPI::CreateMask(); + + // 计算直方图0-127 128-255 + MicroAPI::RegTensor cout0; + MicroAPI::RegTensor cout1; + MicroAPI::Duplicate(cout0, 0); + MicroAPI::Duplicate(cout1, 0); + + MicroAPI::RegTensor idx0; + MicroAPI::RegTensor idx1; + // 0x000000fc -> 0xfcfcfcfc + MicroAPI::LoadAlign(idx0, idx0Buf); + MicroAPI::LoadAlign(idx1, idx1Buf); + + MicroAPI::RegTensor vreg0; + MicroAPI::RegTensor vreg1; + MicroAPI::RegTensor vreg2; + MicroAPI::RegTensor vreg3; + + MicroAPI::RegTensor vreg0U16; + MicroAPI::RegTensor vreg1U16; + MicroAPI::RegTensor vreg2U16; + MicroAPI::RegTensor vreg3U16; + + for (uint16_t i = 0; i < vfLoop; ++i) { + MicroAPI::LoadAlign(vreg1U16, + vreg0U16, inputBuf + i * 256); + MicroAPI::LoadAlign(vreg3U16, + vreg2U16, inputBuf + (i * 256) + 128); + + MicroAPI::DeInterleave(vreg1, vreg0, (MicroAPI::RegTensor&)vreg0U16, + (MicroAPI::RegTensor&)vreg2U16); + MicroAPI::DeInterleave(vreg3, vreg2, (MicroAPI::RegTensor&)vreg1U16, + (MicroAPI::RegTensor&)vreg3U16); + + MicroAPI::MaskReg pregEQ0 = MicroAPI::CreateMask(); + MicroAPI::MaskReg pregEQ1 = MicroAPI::CreateMask(); + MicroAPI::Compare(pregEQ0, vreg0, (MicroAPI::RegTensor&)idx0, pregB8); + MicroAPI::Compare(pregEQ1, vreg1, (MicroAPI::RegTensor&)idx1, pregB8); + + MicroAPI::MaskReg pregEQ = MicroAPI::CreateMask(); + MicroAPI::And(pregEQ, pregEQ0, pregEQ1, pregB8); + + MicroAPI::Histograms(cout0, vreg2, pregEQ); + MicroAPI::Histograms(cout1, vreg2, pregEQ); + } + + StoreHistogramResult(histogramsBuf, cout0, cout1, pregB16, pregB32); +} + +__simd_vf__ void FindThirdTargetBinVFImpl(__ubuf__ uint32_t* idx2Buf, + __ubuf__ uint32_t* nkValueBuf, __ubuf__ uint32_t* kValue, + __ubuf__ uint32_t* histogramsBuf) +{ + MicroAPI::MaskReg pregB32 = MicroAPI::CreateMask(); + + MicroAPI::RegTensor btmK2; + MicroAPI::LoadAlign(btmK2, kValue); + + FindTargetBinAndUpdateNextK(idx2Buf, nkValueBuf, histogramsBuf, btmK2, pregB32); +} + +template +__simd_vf__ void HistogramsLastVFImpl(__ubuf__ uint32_t* histogramsBuf, + __ubuf__ uint32_t* inputBuf, __ubuf__ uint32_t* idx0Buf, + __ubuf__ uint32_t* idx1Buf, __ubuf__ uint32_t* idx2Buf, + uint16_t vfLoop, bool init) +{ + MicroAPI::MaskReg pregB32 = MicroAPI::CreateMask(); + MicroAPI::MaskReg pregB16 = MicroAPI::CreateMask(); + MicroAPI::MaskReg pregB8 = MicroAPI::CreateMask(); + + MicroAPI::RegTensor idx0; + MicroAPI::RegTensor idx1; + MicroAPI::RegTensor idx2; + // 0x000000fc -> 0xfcfcfcfc + MicroAPI::LoadAlign(idx0, idx0Buf); + MicroAPI::LoadAlign(idx1, idx1Buf); + MicroAPI::LoadAlign(idx2, idx2Buf); + + // 计算直方图0-127 128-255 + MicroAPI::RegTensor cout0; + MicroAPI::RegTensor cout1; + MicroAPI::Duplicate(cout0, 0); + MicroAPI::Duplicate(cout1, 0); + + MicroAPI::RegTensor vreg0U16; + MicroAPI::RegTensor vreg1U16; + MicroAPI::RegTensor vreg2U16; + MicroAPI::RegTensor vreg3U16; + + MicroAPI::RegTensor vreg0; + MicroAPI::RegTensor vreg1; + MicroAPI::RegTensor vreg2; + MicroAPI::RegTensor vreg3; + + for (uint16_t i = 0; i < vfLoop; ++i) { + MicroAPI::LoadAlign(vreg1U16, vreg0U16, inputBuf + i * 256); + MicroAPI::LoadAlign(vreg3U16, + vreg2U16, inputBuf + (i * 256) + 128); + + MicroAPI::DeInterleave(vreg1, vreg0, + (MicroAPI::RegTensor&)vreg0U16, + (MicroAPI::RegTensor&)vreg2U16); + MicroAPI::DeInterleave(vreg3, vreg2, + (MicroAPI::RegTensor&)vreg1U16, + (MicroAPI::RegTensor&)vreg3U16); + + MicroAPI::MaskReg pregEQ0 = MicroAPI::CreateMask(); + MicroAPI::MaskReg pregEQ1 = MicroAPI::CreateMask(); + MicroAPI::MaskReg pregEQ2 = MicroAPI::CreateMask(); + MicroAPI::Compare(pregEQ0, vreg0, (MicroAPI::RegTensor&)idx0, pregB8); + MicroAPI::Compare(pregEQ1, vreg1, (MicroAPI::RegTensor&)idx1, pregB8); + MicroAPI::Compare(pregEQ2, vreg2, (MicroAPI::RegTensor&)idx2, pregB8); + + MicroAPI::MaskReg pregEQ0And1 = MicroAPI::CreateMask(); + MicroAPI::MaskReg pregEQAll = MicroAPI::CreateMask(); + MicroAPI::And(pregEQ0And1, pregEQ0, pregEQ1, pregB8); + MicroAPI::And(pregEQAll, pregEQ0And1, pregEQ2, pregB8); + + MicroAPI::Histograms(cout0, vreg3, pregEQAll); + MicroAPI::Histograms(cout1, vreg3, pregEQAll); + } + + StoreHistogramResult(histogramsBuf, cout0, cout1, pregB16, pregB32); +} + +__simd_vf__ void FindKthVFImpl(__ubuf__ uint32_t* kValue, + __ubuf__ uint32_t* histogramsBuf, __ubuf__ uint32_t* idx0Buf, + __ubuf__ uint32_t* idx1Buf, __ubuf__ uint32_t* idx2Buf, + __ubuf__ uint32_t* idx3Buf) +{ + MicroAPI::MaskReg pregB32 = MicroAPI::CreateMask(); + + MicroAPI::ClearSpr(); + + MicroAPI::UnalignRegForStore alignIdx3; + + MicroAPI::RegTensor btmK3; + MicroAPI::LoadAlign(btmK3, kValue); + + for (uint16_t i = 0; i < (uint16_t)(4); ++i) { + MicroAPI::RegTensor idxC; + MicroAPI::RegTensor cout; + MicroAPI::RegTensor sqzIdx3; + + MicroAPI::MaskReg pregGE = MicroAPI::CreateMask(); + + MicroAPI::Arange(idxC, i * 64); + MicroAPI::LoadAlign(cout, histogramsBuf + i * 64); + MicroAPI::Compare(pregGE, cout, btmK3, pregB32); + MicroAPI::Squeeze(sqzIdx3, + (MicroAPI::RegTensor&)idxC, pregGE); + MicroAPI::StoreUnAlign(idx3Buf, sqzIdx3, alignIdx3); + } + MicroAPI::StoreUnAlignPost(idx3Buf, alignIdx3); + + MicroAPI::LocalMemBar(); + + MicroAPI::RegTensor idx0; + MicroAPI::RegTensor idx1; + MicroAPI::RegTensor idx2; + MicroAPI::RegTensor idx3; + MicroAPI::LoadAlign(idx0, idx0Buf); + MicroAPI::LoadAlign(idx1, idx1Buf); + MicroAPI::LoadAlign(idx2, idx2Buf); + MicroAPI::LoadAlign(idx3, idx3Buf); + + MicroAPI::ShiftLefts(idx0, idx0, (int16_t)24, pregB32); + MicroAPI::ShiftLefts(idx1, idx1, (int16_t)16, pregB32); + MicroAPI::ShiftLefts(idx2, idx2, (int16_t)8, pregB32); + + // ADD + MicroAPI::Add(idx0, idx0, idx1, pregB32); + MicroAPI::Add(idx0, idx0, idx2, pregB32); + MicroAPI::Add(idx0, idx0, idx3, pregB32); + + MicroAPI::StoreAlign(kValue, idx0, pregB32); +} + +__simd_vf__ void FindIdxGTOutputVFImpl(__ubuf__ uint32_t* outputIdxBuf, + __ubuf__ uint32_t* inputBuf, uint32_t beginIdx, + __ubuf__ uint32_t* kValue, uint16_t vfLoop) +{ + MicroAPI::MaskReg pregB32 = MicroAPI::CreateMask(); + + MicroAPI::ClearSpr(); + + MicroAPI::UnalignRegForStore alignIdx; + + MicroAPI::RegTensor kthValue; + MicroAPI::LoadAlign(kthValue, kValue); + + MicroAPI::RegTensor vregInput; + + for (uint16_t i = 0; i < (uint16_t)(vfLoop); ++i) { + MicroAPI::RegTensor idxC; + MicroAPI::Arange(idxC, beginIdx + i * 64); + + MicroAPI::LoadAlign(vregInput, inputBuf + i * 64); + + MicroAPI::MaskReg poutGT = MicroAPI::CreateMask(); + + MicroAPI::RegTensor sqzIdxOut; + MicroAPI::Compare(poutGT, vregInput, kthValue, pregB32); + + MicroAPI::Squeeze(sqzIdxOut, + (MicroAPI::RegTensor&)idxC, poutGT); + MicroAPI::StoreUnAlign(outputIdxBuf, sqzIdxOut, alignIdx); + } + MicroAPI::StoreUnAlignPost(outputIdxBuf, alignIdx); +} + +__simd_vf__ void FindIdxEQOutputVFImpl(__ubuf__ uint32_t* outputIdxBuf, + __ubuf__ uint32_t* inputBuf, uint32_t beginIdx, + __ubuf__ uint32_t* kValue) +{ + MicroAPI::MaskReg pregB32 = MicroAPI::CreateMask(); + + MicroAPI::UnalignRegForStore alignIdx; + + MicroAPI::RegTensor kthValue; + MicroAPI::LoadAlign(kthValue, kValue); + + MicroAPI::RegTensor vregInput; + + MicroAPI::RegTensor idxC; + MicroAPI::Arange(idxC, beginIdx); + + MicroAPI::LoadAlign(vregInput, inputBuf); + + MicroAPI::MaskReg poutEQ = MicroAPI::CreateMask(); + + MicroAPI::RegTensor sqzIdxOut; + MicroAPI::Compare(poutEQ, vregInput, kthValue, pregB32); + + MicroAPI::Squeeze(sqzIdxOut, + (MicroAPI::RegTensor&)idxC, poutEQ); + MicroAPI::StoreUnAlign(outputIdxBuf, sqzIdxOut, alignIdx); + MicroAPI::StoreUnAlignPost(outputIdxBuf, alignIdx); +} + +__simd_vf__ void FindValueGTOutputVFImpl(__ubuf__ uint32_t* outputValueBuf, + __ubuf__ uint32_t* inputBuf, __ubuf__ uint32_t* kValue, + uint16_t vfLoop) +{ + MicroAPI::MaskReg pregB32 = MicroAPI::CreateMask(); + + MicroAPI::ClearSpr(); + + MicroAPI::UnalignRegForStore alignValue; + + MicroAPI::RegTensor kthValue; + MicroAPI::LoadAlign(kthValue, kValue); + + MicroAPI::RegTensor vregInput; + + for (uint16_t i = 0; i < (uint16_t)(vfLoop); ++i) { + MicroAPI::LoadAlign(vregInput, inputBuf + i * 64); + + MicroAPI::MaskReg poutGT = MicroAPI::CreateMask(); + + MicroAPI::RegTensor sqzValueOut; + MicroAPI::Compare(poutGT, vregInput, kthValue, pregB32); + + MicroAPI::Squeeze(sqzValueOut, vregInput, poutGT); + MicroAPI::StoreUnAlign(outputValueBuf, + sqzValueOut, alignValue); + } + MicroAPI::StoreUnAlignPost(outputValueBuf, alignValue); +} + +__simd_vf__ void FindValueEQOutputVFImpl(__ubuf__ uint32_t* outputValueBuf, + __ubuf__ uint32_t* inputBuf, __ubuf__ uint32_t* kValue) +{ + MicroAPI::MaskReg pregB32 = MicroAPI::CreateMask(); + + MicroAPI::UnalignRegForStore alignValue; + + MicroAPI::RegTensor kthValue; + MicroAPI::LoadAlign(kthValue, kValue); + + MicroAPI::RegTensor vregInput; + + MicroAPI::LoadAlign(vregInput, inputBuf); + + MicroAPI::MaskReg poutEQ = MicroAPI::CreateMask(); + + MicroAPI::RegTensor sqzValueOut; + MicroAPI::Compare(poutEQ, vregInput, kthValue, pregB32); + + MicroAPI::Squeeze(sqzValueOut, vregInput, poutEQ); + MicroAPI::StoreUnAlign(outputValueBuf, sqzValueOut, alignValue); + MicroAPI::StoreUnAlignPost(outputValueBuf, alignValue); +} + +__aicore__ inline void LiTopKVF(const LocalTensor& outputIdxLocal, + const LocalTensor& outputValueLocal, + const LocalTensor& inputLocal, + const LocalTensor& tmpIdxLocal, + const LocalTensor& tmpValueLocal, + const LocalTensor& histogramsLocal, + const LocalTensor& idx0Local, + const LocalTensor& idx1Local, + const LocalTensor& idx2Local, + const LocalTensor& idx3Local, + const LocalTensor& nkValueLocal, + uint32_t topK, + uint32_t s2SeqLen) +{ + __ubuf__ uint32_t* outputIdxBuf = (__ubuf__ uint32_t*)outputIdxLocal.GetPhyAddr(); + __ubuf__ uint32_t* outputValueBuf = (__ubuf__ uint32_t*)outputValueLocal.GetPhyAddr(); + __ubuf__ uint32_t* inputBuf = (__ubuf__ uint32_t*)inputLocal.GetPhyAddr(); + __ubuf__ uint32_t* tmpIdxBuf = (__ubuf__ uint32_t*)tmpIdxLocal.GetPhyAddr(); + __ubuf__ uint32_t* tmpValueBuf = (__ubuf__ uint32_t*)tmpValueLocal.GetPhyAddr(); + __ubuf__ uint32_t* histogramsBuf = (__ubuf__ uint32_t*)histogramsLocal.GetPhyAddr(); + __ubuf__ uint32_t* idx0Buf = (__ubuf__ uint32_t*)idx0Local.GetPhyAddr(); + __ubuf__ uint32_t* idx1Buf = (__ubuf__ uint32_t*)idx1Local.GetPhyAddr(); + __ubuf__ uint32_t* idx2Buf = (__ubuf__ uint32_t*)idx2Local.GetPhyAddr(); + __ubuf__ uint32_t* idx3Buf = (__ubuf__ uint32_t*)idx3Local.GetPhyAddr(); + __ubuf__ uint32_t* nkValueBuf = (__ubuf__ uint32_t*)nkValueLocal.GetPhyAddr(); + + uint32_t bottomK = s2SeqLen - topK + 1; + uint32_t beginIdx = 0; + bool flag = true; + + const uint16_t repeatSize8 = 256; + const uint16_t repeatSize32 = 64; + + uint16_t histogramsLoopNum = (s2SeqLen + repeatSize8 - 1) / repeatSize8; + uint16_t inputLoopNum = (s2SeqLen + repeatSize32 - 1) / repeatSize32; + uint16_t topkLoopNum = (topK + 64 - 1) / 64; + + // find kth-value + HistogramsFirstVFImpl(histogramsBuf, inputBuf, histogramsLoopNum, flag); + FindFirstTargetBinVFImpl(idx0Buf, nkValueBuf, histogramsBuf, bottomK); + HistogramsSecondVFImpl(histogramsBuf, inputBuf, idx0Buf, histogramsLoopNum, flag); + FindSecondTargetBinVFImpl(idx1Buf, nkValueBuf, nkValueBuf, histogramsBuf); + HistogramsThirdVFImpl(histogramsBuf, inputBuf, idx0Buf, idx1Buf, histogramsLoopNum, flag); + FindThirdTargetBinVFImpl(idx2Buf, nkValueBuf, nkValueBuf, histogramsBuf); + HistogramsLastVFImpl(histogramsBuf, inputBuf, idx0Buf, idx1Buf, idx2Buf, histogramsLoopNum, flag); + FindKthVFImpl(nkValueBuf, histogramsBuf, idx0Buf, idx1Buf, idx2Buf, idx3Buf); + + // filter + // 输出大于k-value的值value + FindValueGTOutputVFImpl(outputValueBuf, inputBuf, nkValueBuf, inputLoopNum); + // value-当前偏移大于k-value的值在AR特殊寄存器中的有效字节数 + int64_t arValueNum = AscendC::GetSpr(); + // value-剩余需要输出等于k-value的数量 + int64_t remainValueNum = topK - (arValueNum / sizeof(uint32_t)); + for (uint16_t i = 0; i < inputLoopNum; ++i) { + int64_t arValueNumPerLoop = AscendC::GetSpr(); + if (((arValueNumPerLoop - arValueNum) / sizeof(uint32_t)) < remainValueNum) { + // 调用一次查找等于k-value情况的过程；64: 单次循环处理的元素块大小 + FindValueEQOutputVFImpl(outputValueBuf, inputBuf + i * 64, nkValueBuf); + } else { + break; + } + } + + // 输出大于k-value的值idx + FindIdxGTOutputVFImpl(outputIdxBuf, inputBuf, (uint32_t)(0), nkValueBuf, inputLoopNum); + // idx-当前偏移大于k-value的值在AR特殊寄存器中的有效字节数 + int64_t arIdxNum = AscendC::GetSpr(); + int64_t remainIdxNum = topK - (arIdxNum / sizeof(uint32_t)); + for (uint16_t i = 0; i < inputLoopNum; ++i) { + int64_t arIdxNumPerLoop = AscendC::GetSpr(); + if (((arIdxNumPerLoop - arIdxNum) / sizeof(uint32_t)) < remainIdxNum) { + // 调用一次查找等于k-value情况的过程 + beginIdx = i * 64; // 64: 块起始偏移量 + FindIdxEQOutputVFImpl(outputIdxBuf, inputBuf + i * 64, beginIdx, nkValueBuf); // 64: 块起始偏移量 + } else { + break; + } + } +} +} +#endif diff --git a/xllm_ops/attention/quant_lightning_indexer_v2/op_kernel/arch35/vf/vf_topk_16_gather_quant_v2.h b/xllm_ops/attention/quant_lightning_indexer_v2/op_kernel/arch35/vf/vf_topk_16_gather_quant_v2.h new file mode 100644 index 0000000..34b7753 --- /dev/null +++ b/xllm_ops/attention/quant_lightning_indexer_v2/op_kernel/arch35/vf/vf_topk_16_gather_quant_v2.h @@ -0,0 +1,514 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! +* \file vf_topk_16_gather_quant_v2.h +* \brief +*/ + +#ifndef VF_TOPK_16_GATHER_QUANT_V2_H +#define VF_TOPK_16_GATHER_QUANT_V2_H + +namespace topkb16gather { + +template +__simd_vf__ void HistogramsHighVFImpl(__ubuf__ uint32_t* histogramsBuf, __ubuf__ uint16_t* inputBuf, uint16_t vfLoop, + bool init) +{ + MicroAPI::MaskReg pregB32 = MicroAPI::CreateMask(); + MicroAPI::MaskReg pregB16 = MicroAPI::CreateMask(); + MicroAPI::MaskReg pregB8 = MicroAPI::CreateMask(); + + // 计算直方图cout0 0-127 cout1 128-255 + MicroAPI::RegTensor cout0; + MicroAPI::RegTensor cout1; + MicroAPI::Duplicate(cout0, 0); + MicroAPI::Duplicate(cout1, 0); + + MicroAPI::RegTensor cout0U32Even; + MicroAPI::RegTensor cout0U32Odd; + MicroAPI::RegTensor cout1U32Even; + MicroAPI::RegTensor cout1U32Odd; + + MicroAPI::RegTensor vregHigh; + MicroAPI::RegTensor vregLow; + + static constexpr MicroAPI::CastTrait CAST_TRAIT_UINT16_TOUINT32_EVEN = {MicroAPI::RegLayout::ZERO, + MicroAPI::SatMode::UNKNOWN, MicroAPI::MaskMergeMode::ZEROING, RoundMode::UNKNOWN}; + + static constexpr MicroAPI::CastTrait CAST_TRAIT_UINT16_TOUINT32_ODD = {MicroAPI::RegLayout::ONE, + MicroAPI::SatMode::UNKNOWN, MicroAPI::MaskMergeMode::ZEROING, RoundMode::UNKNOWN}; + + for (uint16_t i = 0; i < vfLoop; ++i) { + MicroAPI::LoadAlign(vregLow, vregHigh, inputBuf + i * 256); + + MicroAPI::Histograms(cout0, (MicroAPI::RegTensor&)vregHigh, + pregB8); + MicroAPI::Histograms(cout1, (MicroAPI::RegTensor&)vregHigh, + pregB8); + } + + MicroAPI::Cast(cout0U32Even, cout0, pregB16); + MicroAPI::Cast(cout0U32Odd, cout0, pregB16); + MicroAPI::Cast(cout1U32Even, cout1, pregB16); + MicroAPI::Cast(cout1U32Odd, cout1, pregB16); + + MicroAPI::StoreAlign(histogramsBuf, cout0U32Even, cout0U32Odd, + pregB32); + MicroAPI::StoreAlign(histogramsBuf + 128, cout1U32Even, cout1U32Odd, + pregB32); +} + +__simd_vf__ void FindHighTargetBinVFImpl(__ubuf__ uint32_t* idxHighBuf, __ubuf__ uint32_t* nkValueBuf, + __ubuf__ uint32_t* histogramsBuf, uint32_t bottomK) +{ + MicroAPI::MaskReg pregB32 = MicroAPI::CreateMask(); + + MicroAPI::MaskReg pregGE; + + MicroAPI::ClearSpr(); + + MicroAPI::UnalignRegForStore alignIdxHigh; + + MicroAPI::RegTensor btmK; + MicroAPI::Duplicate(btmK, bottomK); + + MicroAPI::RegTensor idxC; + MicroAPI::RegTensor cout; + MicroAPI::RegTensor sqzIdxHigh; + + for (uint16_t i = 0; i < (uint16_t)(4); ++i) { + MicroAPI::Arange(idxC, i * 64); + + MicroAPI::LoadAlign(cout, histogramsBuf + i * 64); + + MicroAPI::Compare(pregGE, cout, btmK, pregB32); + + MicroAPI::Squeeze(sqzIdxHigh, + (MicroAPI::RegTensor&)idxC, pregGE); + MicroAPI::StoreUnAlign(idxHighBuf, sqzIdxHigh, alignIdxHigh); + } + MicroAPI::StoreUnAlignPost(idxHighBuf, alignIdxHigh); + + MicroAPI::LocalMemBar(); + + MicroAPI::RegTensor idxHigh; + MicroAPI::LoadAlign(idxHigh, idxHighBuf); + + MicroAPI::RegTensor idxAll1; + MicroAPI::RegTensor idxPrev0; + MicroAPI::RegTensor prevBinValue; + MicroAPI::Duplicate(idxAll1, 1); + + MicroAPI::RegTensor zeroAll; + MicroAPI::Duplicate(zeroAll, 0); + + MicroAPI::MaskReg preg0 = MicroAPI::CreateMask(); + MicroAPI::Compare(preg0, idxHigh, zeroAll, pregB32); + MicroAPI::Sub(idxPrev0, idxHigh, (MicroAPI::RegTensor&)idxAll1, pregB32); + MicroAPI::ShiftRights(idxPrev0, idxPrev0, (int16_t)24, pregB32); + + MicroAPI::Gather(prevBinValue, histogramsBuf, idxPrev0, pregB32); + MicroAPI::Select(prevBinValue, zeroAll, prevBinValue, preg0); + + MicroAPI::RegTensor nextK; + MicroAPI::Sub(nextK, btmK, prevBinValue, pregB32); + MicroAPI::StoreAlign(nkValueBuf, nextK, pregB32); +} + +template +__simd_vf__ void HistogramsLowVFImpl(__ubuf__ uint32_t* histogramsBuf, __ubuf__ uint16_t* inputBuf, + __ubuf__ uint32_t* idxHighBuf, uint16_t vfLoop, bool init) +{ + MicroAPI::MaskReg pregB32 = MicroAPI::CreateMask(); + MicroAPI::MaskReg pregB16 = MicroAPI::CreateMask(); + MicroAPI::MaskReg pregB8 = MicroAPI::CreateMask(); + + MicroAPI::MaskReg pregEQ; + + // 计算直方图0-127 128-255 + MicroAPI::RegTensor cout0; + MicroAPI::RegTensor cout1; + MicroAPI::Duplicate(cout0, 0); + MicroAPI::Duplicate(cout1, 0); + + MicroAPI::RegTensor cout0U32Even; + MicroAPI::RegTensor cout0U32Odd; + MicroAPI::RegTensor cout1U32Even; + MicroAPI::RegTensor cout1U32Odd; + + MicroAPI::RegTensor idxHigh; + MicroAPI::LoadAlign(idxHigh, idxHighBuf); + + MicroAPI::RegTensor vregHigh; + MicroAPI::RegTensor vregLow; + + static constexpr MicroAPI::CastTrait CAST_TRAIT_UINT16_TOUINT32_EVEN = {MicroAPI::RegLayout::ZERO, + MicroAPI::SatMode::UNKNOWN, MicroAPI::MaskMergeMode::ZEROING, RoundMode::UNKNOWN}; + + static constexpr MicroAPI::CastTrait CAST_TRAIT_UINT16_TOUINT32_ODD = {MicroAPI::RegLayout::ONE, + MicroAPI::SatMode::UNKNOWN, MicroAPI::MaskMergeMode::ZEROING, RoundMode::UNKNOWN}; + + for (uint16_t i = 0; i < vfLoop; ++i) { + MicroAPI::LoadAlign(vregLow, vregHigh, inputBuf + i * 256); + + MicroAPI::Compare(pregEQ, (MicroAPI::RegTensor&)vregHigh, + (MicroAPI::RegTensor&)idxHigh, pregB8); + + MicroAPI::Histograms(cout0, (MicroAPI::RegTensor&)vregLow, + pregEQ); + MicroAPI::Histograms(cout1, (MicroAPI::RegTensor&)vregLow, + pregEQ); + } + + MicroAPI::Cast(cout0U32Even, cout0, pregB16); + MicroAPI::Cast(cout0U32Odd, cout0, pregB16); + MicroAPI::Cast(cout1U32Even, cout1, pregB16); + MicroAPI::Cast(cout1U32Odd, cout1, pregB16); + + MicroAPI::StoreAlign(histogramsBuf, cout0U32Even, cout0U32Odd, + pregB32); + MicroAPI::StoreAlign(histogramsBuf + 128, cout1U32Even, cout1U32Odd, + pregB32); +} + +__simd_vf__ void FindKthVFImpl(__ubuf__ uint32_t* kValue, __ubuf__ uint32_t* histogramsBuf, + __ubuf__ uint32_t* idxHighBuf, __ubuf__ uint32_t* idxLowBuf) +{ + MicroAPI::MaskReg pregB32 = MicroAPI::CreateMask(); + MicroAPI::MaskReg pregB16 = MicroAPI::CreateMask(); + + MicroAPI::MaskReg pregGE; + + MicroAPI::ClearSpr(); + + MicroAPI::UnalignRegForStore alignIdxLow; + + MicroAPI::RegTensor btmK; + MicroAPI::LoadAlign(btmK, kValue); + + MicroAPI::RegTensor idxC; + MicroAPI::RegTensor cout; + MicroAPI::RegTensor sqzIdxLow; + + for (uint16_t i = 0; i < (uint16_t)(4); ++i) { + MicroAPI::Arange(idxC, i * 64); + + MicroAPI::LoadAlign(cout, histogramsBuf + i * 64); + + MicroAPI::Compare(pregGE, cout, btmK, pregB32); + + MicroAPI::Squeeze(sqzIdxLow, + (MicroAPI::RegTensor&)idxC, pregGE); + MicroAPI::StoreUnAlign(idxLowBuf, sqzIdxLow, alignIdxLow); + } + MicroAPI::StoreUnAlignPost(idxLowBuf, alignIdxLow); + + MicroAPI::LocalMemBar(); + + MicroAPI::RegTensor idxHigh; + MicroAPI::RegTensor idxLow; + MicroAPI::LoadAlign(idxHigh, idxHighBuf); + MicroAPI::LoadAlign(idxLow, idxLowBuf); + + MicroAPI::RegTensor idxTmp; + MicroAPI::Duplicate(idxTmp, 0xff00); + + MicroAPI::And(idxHigh, idxHigh, (MicroAPI::RegTensor&)idxTmp, pregB32); + + MicroAPI::RegTensor idxK; + MicroAPI::Add(idxK, idxHigh, idxLow, pregB16); + + MicroAPI::StoreAlign(kValue, idxK, pregB32); +} + +/** + 输出所有大于的kth-value的Index + */ +__simd_vf__ void FindIdxGTOutputVFImpl(__ubuf__ uint16_t* outputIdxBuf, __ubuf__ uint16_t* inputValueBuf, + uint16_t beginIdx, __ubuf__ uint32_t* kValue, uint16_t vfLoop) +{ + MicroAPI::MaskReg pregB16 = MicroAPI::CreateMask(); + + MicroAPI::MaskReg poutGT; + + MicroAPI::ClearSpr(); + + MicroAPI::UnalignRegForStore alignIdx; + + MicroAPI::RegTensor kthValue; + MicroAPI::LoadAlign(kthValue, kValue); + + MicroAPI::RegTensor vregInput; + MicroAPI::RegTensor idxC; + MicroAPI::RegTensor sqzIdxOut; + + for (uint16_t i = 0; i < (uint16_t)(vfLoop); ++i) { + MicroAPI::Arange(idxC, beginIdx + i * 128); + + MicroAPI::LoadAlign(vregInput, inputValueBuf + i * 128); + + MicroAPI::Compare(poutGT, vregInput, (MicroAPI::RegTensor&)kthValue, pregB16); + + MicroAPI::Squeeze(sqzIdxOut, + (MicroAPI::RegTensor&)idxC, poutGT); + MicroAPI::StoreUnAlign(outputIdxBuf, sqzIdxOut, alignIdx); + } + MicroAPI::StoreUnAlignPost(outputIdxBuf, alignIdx); +} + +/** + 输出所有等于的kth-value的Index + */ +__simd_vf__ void FindIdxEQOutputVFImpl(__ubuf__ uint16_t* outputIdxBuf, __ubuf__ uint16_t* inputValueBuf, + uint16_t beginIdx, __ubuf__ uint32_t* kValue, uint16_t vfLoop) +{ + MicroAPI::MaskReg pregB16 = MicroAPI::CreateMask(); + + MicroAPI::MaskReg poutEQ; + + MicroAPI::UnalignRegForStore alignIdx; + + MicroAPI::RegTensor kthValue; + MicroAPI::LoadAlign(kthValue, kValue); + + MicroAPI::RegTensor vregInput; + MicroAPI::RegTensor idxC; + MicroAPI::RegTensor sqzIdxOut; + + for (uint16_t i = 0; i < (uint16_t)(vfLoop); ++i) { + MicroAPI::Arange(idxC, beginIdx + i * 128); + + MicroAPI::LoadAlign(vregInput, inputValueBuf + i * 128); + + MicroAPI::Compare(poutEQ, vregInput, (MicroAPI::RegTensor&)kthValue, pregB16); + + MicroAPI::Squeeze(sqzIdxOut, + (MicroAPI::RegTensor&)idxC, poutEQ); + MicroAPI::StoreUnAlign(outputIdxBuf, sqzIdxOut, alignIdx); + } + MicroAPI::StoreUnAlignPost(outputIdxBuf, alignIdx); +} + +/** + 输出最终的Value + */ +__simd_vf__ void FindValueOutputVFImpl(__ubuf__ uint16_t* outputValueBuf, __ubuf__ uint16_t* inputValueBuf, + __ubuf__ uint16_t* tmpIdxBuf, uint16_t vfLoop) +{ + MicroAPI::MaskReg pregB16 = MicroAPI::CreateMask(); + + MicroAPI::RegTensor tmpIdx; + MicroAPI::RegTensor outputValue; + + for (uint16_t i = 0; i < (uint16_t)(vfLoop); ++i) { + MicroAPI::LoadAlign(tmpIdx, tmpIdxBuf + i * 128); + + MicroAPI::Gather(outputValue, inputValueBuf, tmpIdx, pregB16); + + MicroAPI::StoreAlign(outputValueBuf + i * 128, outputValue, pregB16); + } +} + +/** + 输出最终的Idx + */ +__simd_vf__ void FindRealIndexVFImpl(__ubuf__ uint32_t* outputIdxBuf, __ubuf__ uint16_t* tmpIdxBuf, + __ubuf__ uint32_t* hisIdxBuf, uint32_t topK, uint32_t loopIndex, uint16_t vfLoop) +{ + MicroAPI::MaskReg pregB32 = MicroAPI::CreateMask(); + + MicroAPI::MaskReg pregNow; + MicroAPI::MaskReg pregHis; + + MicroAPI::RegTensor tmpIdx; + MicroAPI::RegTensor outputGatherIdx; + MicroAPI::RegTensor outputAddsIdx; + + for (uint16_t i = 0; i < (uint16_t)(vfLoop); ++i) { + MicroAPI::LoadAlign(tmpIdx, tmpIdxBuf + i * 64); + + MicroAPI::Compares(pregNow, (MicroAPI::RegTensor&)tmpIdx, topK - 1, pregB32); + MicroAPI::Xor(pregHis, pregNow, pregB32, pregB32); + + MicroAPI::Gather(outputGatherIdx, hisIdxBuf, (MicroAPI::RegTensor&)tmpIdx, pregHis); + MicroAPI::Adds(outputAddsIdx, (MicroAPI::RegTensor&)tmpIdx, loopIndex, pregNow); + + MicroAPI::Add(outputGatherIdx, outputGatherIdx, outputAddsIdx, pregB32); + + MicroAPI::StoreAlign(outputIdxBuf + i * 64, outputGatherIdx, pregB32); + } +} + +__simd_vf__ void IndicesAddOffsetVF(__ubuf__ uint32_t* indicesOutBuf, uint32_t outputIdxOffset, uint32_t vfLoop) +{ + MicroAPI::MaskReg pregB32 = MicroAPI::CreateMask(); + + MicroAPI::RegTensor outIndices; + + for (uint16_t i = 0; i < (uint16_t)(vfLoop); ++i) { + MicroAPI::LoadAlign(outIndices, indicesOutBuf + i * 64); + MicroAPI::Adds(outIndices, outIndices, outputIdxOffset, pregB32); + MicroAPI::StoreAlign(indicesOutBuf + i * 64, outIndices, pregB32); + } +} + +__aicore__ inline void IndicesAddOffset(const LocalTensor& indicesOutLocal, + uint32_t outputIdxOffset, uint32_t topK) +{ + __ubuf__ uint32_t* indicesOutBuf = (__ubuf__ uint32_t*)indicesOutLocal.GetPhyAddr(); + const uint16_t repeatSize32 = 64; + uint16_t topkLoopNum32 = (topK + repeatSize32 - 1) / repeatSize32; + IndicesAddOffsetVF(indicesOutBuf, outputIdxOffset, topkLoopNum32); +} + +/** + * @brief LiTopKVF 对一个validLen的输入进行topk算法，输出idx_tmp + * @param tmpIdxLocal Temp阶段输出的TopKIndex;如果s2SeqLen < 16K作为最终输出 validLen * 2B + * @param outputValueLocal 如果s2SeqLen > 16K并且是首轮输出Value topK * 2B + * @param inputValueLocal 输入Value validLen * 2B + * @param histogramsLocal 直方图 256 * 4B + * @param idxHighLocal 目标桶高八位 256 * 4B + * @param idxLowLocal 目标桶低八位 256 * 4B + * @param nkValueLocal 存储next_k的值 64 * 4B + * @param topK topK元素 + * @param validLen 有效元素个数:QLIV2Common::Align(topkCountAlign256_ + validTrunkLen, (uint32_t)256) + */ +template // 是否输出VALUE +__aicore__ inline void LiTopKVF(const LocalTensor& tmpIdxLocal, + const LocalTensor& outputValueLocal, + const LocalTensor& inputValueLocal, + const LocalTensor& histogramsLocal, + const LocalTensor& idxHighLocal, + const LocalTensor& idxLowLocal, + const LocalTensor& nkValueLocal, + uint32_t topK, + uint32_t validLen) +{ + __ubuf__ uint16_t* tmpIdxBuf = (__ubuf__ uint16_t*)tmpIdxLocal.GetPhyAddr(); + __ubuf__ uint16_t* outputValueBuf = (__ubuf__ uint16_t*)outputValueLocal.GetPhyAddr(); + __ubuf__ uint16_t* inputValueBuf = (__ubuf__ uint16_t*)inputValueLocal.GetPhyAddr(); + __ubuf__ uint32_t* histogramsBuf = (__ubuf__ uint32_t*)histogramsLocal.GetPhyAddr(); + __ubuf__ uint32_t* idxHighBuf = (__ubuf__ uint32_t*)idxHighLocal.GetPhyAddr(); + __ubuf__ uint32_t* idxLowBuf = (__ubuf__ uint32_t*)idxLowLocal.GetPhyAddr(); + __ubuf__ uint32_t* nkValueBuf = (__ubuf__ uint32_t*)nkValueLocal.GetPhyAddr(); + + uint32_t bottomK = validLen - topK + 1; + uint32_t beginIdx = 0; + bool flag = true; + + const uint16_t repeatSize8 = 256; + const uint16_t repeatSize16 = 128; + const uint16_t repeatSize32 = 64; + + uint16_t histogramsLoopNum = (validLen + repeatSize8 - 1) / repeatSize8; + uint16_t inputLoopNum = (validLen + repeatSize16 - 1) / repeatSize16; + uint16_t topkLoopNum = (topK + repeatSize32 - 1) / repeatSize32; + uint16_t topkLoopNum16 = (topK + repeatSize16 - 1) / repeatSize16; + + // find kth-value + HistogramsHighVFImpl(histogramsBuf, inputValueBuf, histogramsLoopNum, flag); + FindHighTargetBinVFImpl(idxHighBuf, nkValueBuf, histogramsBuf, bottomK); + + HistogramsLowVFImpl(histogramsBuf, inputValueBuf, idxHighBuf, histogramsLoopNum, flag); + FindKthVFImpl(nkValueBuf, histogramsBuf, idxHighBuf, idxLowBuf); + + // filter + AscendC::Duplicate(tmpIdxLocal, (uint16_t)(0), QLIV2Common::Align(topK, (uint32_t)128)); + // 输出大于k-value的值idx + FindIdxGTOutputVFImpl(tmpIdxBuf, inputValueBuf, (uint32_t)(0), nkValueBuf, inputLoopNum); + // 输出等于k-value的值idx + FindIdxEQOutputVFImpl(tmpIdxBuf, inputValueBuf, (uint32_t)(0), nkValueBuf, inputLoopNum); + + // 是否输出Value + if constexpr (ISOUTVALUE) { + FindValueOutputVFImpl(outputValueBuf, inputValueBuf, tmpIdxBuf, topkLoopNum16); + } +} + +/** + LD:输出最终的Idx +*/ +__simd_vf__ void FindLDRealIndexVFImpl(__ubuf__ uint32_t* outputIdxBuf, __ubuf__ uint16_t* tmpIdxBuf, + __ubuf__ uint32_t* hisIdxBuf, uint16_t vfLoop) +{ + MicroAPI::MaskReg pregB32 = MicroAPI::CreateMask(); + + MicroAPI::RegTensor tmpIdx; + MicroAPI::RegTensor outputIdx; + + for (uint16_t i = 0; i < (uint16_t)(vfLoop); ++i) { + MicroAPI::LoadAlign(tmpIdx, tmpIdxBuf + i * 64); + + MicroAPI::Gather(outputIdx, hisIdxBuf, (MicroAPI::RegTensor&)tmpIdx, pregB32); + + MicroAPI::StoreAlign(outputIdxBuf + i * 64, outputIdx, pregB32); + } +} + +/** + * @brief 通过idx_tmp gather出实际的TopKIndex，s2SeqLen > 16K才会执行 + * @param outputIdxLocal 输出Idx 有效:topK * 2B + * @param outputValueLocal 输出Value topK * 2B(以后需要输出实际value使用) + * @param inputValueLocal 输入Value validLen * 2B + * @param tmpIdxLocal 本轮tmpIdx输入 validLen * 2B (0 ~ validLen - 1) + * @param hisIdxLocal 上一轮实际Idx输入 有效:topK * 4B + * @param topK topK元素个数 + * @param loopBasicIdx 当前循环需要加上得基准Index + * @param validLen 有效元素个数 + */ +__aicore__ inline void LiTopKGatherVF(const LocalTensor& outputIdxLocal, + const LocalTensor& outputValueLocal, + const LocalTensor& inputValueLocal, + const LocalTensor& tmpIdxLocal, + const LocalTensor& hisIdxLocal, + uint32_t topK, + uint32_t loopBasicIdx, + uint32_t validLen) +{ + __ubuf__ uint32_t* outputIdxBuf = (__ubuf__ uint32_t*)outputIdxLocal.GetPhyAddr(); + __ubuf__ uint16_t* outputValueBuf = (__ubuf__ uint16_t*)outputValueLocal.GetPhyAddr(); + __ubuf__ uint16_t* inputValueBuf = (__ubuf__ uint16_t*)inputValueLocal.GetPhyAddr(); + __ubuf__ uint16_t* tmpIdxBuf = (__ubuf__ uint16_t*)tmpIdxLocal.GetPhyAddr(); + __ubuf__ uint32_t* hisIdxBuf = (__ubuf__ uint32_t*)hisIdxLocal.GetPhyAddr(); + + const uint16_t repeatSize32 = 64; + const uint16_t repeatSize16 = 128; + uint16_t topkLoopNum16 = (topK + repeatSize16 - 1) / repeatSize16; + uint16_t topkLoopNum32 = (topK + repeatSize32 - 1) / repeatSize32; + + FindRealIndexVFImpl(outputIdxBuf, tmpIdxBuf, hisIdxBuf, topK, loopBasicIdx, topkLoopNum32); +} + +/** + LD:gather最终的Idx +*/ +__aicore__ inline void LiTopKLDGatherVF(const LocalTensor& outputIdxLocal, // 输出Idx topK * 2B + const LocalTensor& tmpIdxLocal, // 本轮tmpIdx输入 validLen * 2B + const LocalTensor& hisIdxLocal, // 上一轮Idx输入 topK * 4B + uint32_t topK) // topK元素个数 +{ + __ubuf__ uint32_t* outputIdxBuf = (__ubuf__ uint32_t*)outputIdxLocal.GetPhyAddr(); + __ubuf__ uint16_t* tmpIdxBuf = (__ubuf__ uint16_t*)tmpIdxLocal.GetPhyAddr(); + __ubuf__ uint32_t* hisIdxBuf = (__ubuf__ uint32_t*)hisIdxLocal.GetPhyAddr(); + + const uint16_t repeatSize32 = 64; + const uint16_t repeatSize16 = 128; + uint16_t topkLoopNum16 = (topK + repeatSize16 - 1) / repeatSize16; + uint16_t topkLoopNum32 = (topK + repeatSize32 - 1) / repeatSize32; + + FindLDRealIndexVFImpl(outputIdxBuf, tmpIdxBuf, hisIdxBuf, topkLoopNum32); +} +} +#endif // VF_TOPK_16_GATHER_QUANT_V2_H \ No newline at end of file diff --git a/xllm_ops/attention/quant_lightning_indexer_v2/op_kernel/quant_lightning_indexer_v2.cpp b/xllm_ops/attention/quant_lightning_indexer_v2/op_kernel/quant_lightning_indexer_v2.cpp new file mode 100644 index 0000000..8c9e7fc --- /dev/null +++ b/xllm_ops/attention/quant_lightning_indexer_v2/op_kernel/quant_lightning_indexer_v2.cpp @@ -0,0 +1,78 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file quant_lightning_indexer_v2.cpp + * \brief + */ + +#include "kernel_operator.h" +#include "lib/matmul_intf.h" +#if (__CCE_AICORE__ == 310) +#include "arch35/quant_lightning_indexer_v2_kernel_arch35.h" +#else +#include "arch22/quant_lightning_indexer_v2_kernel_arch22.h" +#endif +#include "quant_lightning_indexer_v2_template_tiling_key.h" +using namespace QLIV2Kernel; +using namespace optiling::detail; + +#define INVOKE_LI_NO_KFC_OP_IMPL(templateClass, ...) \ + do { \ + templateClass> op; \ + op.Init(query, key, weights, queryScale, keyScale, cuSeqlensQ, cuSeqlensK, sequsedQ, sequsedK, cmpResidualK, \ + blockTable, outputIdxOffset, metadata, sparseIndices, sparseValues, user, tiling_data, &tPipe); \ + op.Process(); \ + } while (0) + +template +__global__ __aicore__ void quant_lightning_indexer_v2( + __gm__ uint8_t *query, __gm__ uint8_t *key, __gm__ uint8_t *weights, __gm__ uint8_t *queryScale, + __gm__ uint8_t *keyScale, __gm__ uint8_t *cuSeqlensQ, __gm__ uint8_t *cuSeqlensK, __gm__ uint8_t *sequsedQ, + __gm__ uint8_t *sequsedK, __gm__ uint8_t *cmpResidualK, __gm__ uint8_t *blockTable, __gm__ uint8_t *outputIdxOffset, + __gm__ uint8_t *metadata, __gm__ uint8_t *sparseIndices, __gm__ uint8_t *sparseValues, __gm__ uint8_t *workspace, + __gm__ uint8_t *tiling) +{ + TPipe tPipe; + __gm__ uint8_t *user = GetUserWorkspace(workspace); + KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_MIX_AIC_1_2); + GET_TILING_DATA_WITH_STRUCT(QLIV2TilingData, tiling_data_in, tiling); + const QLIV2TilingData *__restrict tiling_data = &tiling_data_in; + +#if (__CCE_AICORE__ == 310) + constexpr uint32_t QUANT_MODE_FP8 = 1; + constexpr uint32_t QUANT_MODE_INT8 = 2; + constexpr uint32_t QUANT_MODE_MXFP8 = 3; + constexpr uint32_t QUANT_MODE_HIFLOAT8 = 4; + constexpr uint32_t QUANT_MODE_MXFP4 = 5; + if (tiling_data->quantMode == QUANT_MODE_FP8) { + INVOKE_LI_NO_KFC_OP_IMPL(QLIV2Preload, fp8_e4m3fn_t, fp8_e4m3fn_t, float, uint16_t, int32_t, PAGE_ATTENTION, + LI_LAYOUT(Q_LAYOUT_T), LI_LAYOUT(K_LAYOUT_T), float, float, float); + } else if (tiling_data->quantMode == QUANT_MODE_MXFP8) { + INVOKE_LI_NO_KFC_OP_IMPL(QLIV2Preload, fp8_e4m3fn_t, fp8_e4m3fn_t, float, uint16_t, int32_t, PAGE_ATTENTION, + LI_LAYOUT(Q_LAYOUT_T), LI_LAYOUT(K_LAYOUT_T), fp8_e8m0_t, float, float); + } else if (tiling_data->quantMode == QUANT_MODE_HIFLOAT8) { + INVOKE_LI_NO_KFC_OP_IMPL(QLIV2Preload, hifloat8, hifloat8, float, uint16_t, int32_t, PAGE_ATTENTION, + LI_LAYOUT(Q_LAYOUT_T), LI_LAYOUT(K_LAYOUT_T), float, float, float); + } else if (tiling_data->quantMode == QUANT_MODE_MXFP4) { + INVOKE_LI_NO_KFC_OP_IMPL(QLIV2Preload, fp4x2_e2m1_t, fp4x2_e2m1_t, bfloat16_t, uint16_t, int32_t, + PAGE_ATTENTION, LI_LAYOUT(Q_LAYOUT_T), LI_LAYOUT(K_LAYOUT_T), fp8_e8m0_t, float, + float); + } else if (tiling_data->quantMode == QUANT_MODE_INT8) { + INVOKE_LI_NO_KFC_OP_IMPL(QLIV2Preload, int8_t, int8_t, int32_t, uint16_t, int32_t, + PAGE_ATTENTION, LI_LAYOUT(Q_LAYOUT_T), LI_LAYOUT(K_LAYOUT_T), + half, half, int32_t); + } + +#else + INVOKE_LI_NO_KFC_OP_IMPL(QLIV2Preload, int8_t, int8_t, float, uint16_t, int32_t, PAGE_ATTENTION, + LI_LAYOUT(Q_LAYOUT_T), LI_LAYOUT(K_LAYOUT_T)); +#endif +} diff --git a/xllm_ops/attention/quant_lightning_indexer_v2/op_kernel/quant_lightning_indexer_v2_metadata.h b/xllm_ops/attention/quant_lightning_indexer_v2/op_kernel/quant_lightning_indexer_v2_metadata.h new file mode 100644 index 0000000..17a97a1 --- /dev/null +++ b/xllm_ops/attention/quant_lightning_indexer_v2/op_kernel/quant_lightning_indexer_v2_metadata.h @@ -0,0 +1,79 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! +* \file quant_lightning_indexer_v2_metadata.h +* \brief +*/ + +#ifndef QUANT_LIGHTNING_INDEXER_V2_METADATA_H +#define QUANT_LIGHTNING_INDEXER_V2_METADATA_H + +#include + +namespace optiling { + +// Constants +inline constexpr uint32_t AIC_CORE_MAX_NUM = 36; +inline constexpr uint32_t AIV_CORE_MAX_NUM = 72; +constexpr uint32_t QLI_V2_METADATA_TOTAL_SIZE = 1024; +using QLI_V2_METADATA_T = int32_t; + +inline constexpr uint32_t QLI_V2_METADATA_SIZE = 8; +inline constexpr uint32_t QLD_V2_METADATA_SIZE = 8; + +// LI Metadata Index Definitions +inline constexpr uint32_t QLI_V2_CORE_ENABLE_INDEX = 0; +inline constexpr uint32_t QLI_V2_BN2_START_INDEX = 1; +inline constexpr uint32_t QLI_V2_M_START_INDEX = 2; +inline constexpr uint32_t QLI_V2_S2_START_INDEX = 3; +inline constexpr uint32_t QLI_V2_BN2_END_INDEX = 4; +inline constexpr uint32_t QLI_V2_M_END_INDEX = 5; +inline constexpr uint32_t QLI_V2_S2_END_INDEX = 6; +inline constexpr uint32_t QLI_V2_FIRST_QLD_V2_DATA_WORKSPACE_IDX_INDEX = 7; + +// LD Metadata Index Definitions +inline constexpr uint32_t QLD_V2_CORE_ENABLE_INDEX = 0; +inline constexpr uint32_t QLD_V2_BN2_IDX_INDEX = 1; +inline constexpr uint32_t QLD_V2_M_IDX_INDEX = 2; +inline constexpr uint32_t QLD_V2_WORKSPACE_IDX_INDEX = 3; +inline constexpr uint32_t QLD_V2_WORKSPACE_NUM_INDEX = 4; +inline constexpr uint32_t QLD_V2_M_START_INDEX = 5; +inline constexpr uint32_t QLD_V2_M_NUM_INDEX = 6; + + /** + * @brief 获取属性的绝对索引 + * @param coreIdx 核索引 + * @param metaIdx 元数据索引 + * @param isAIV 是否为AIV数据，默认为false + * @return 返回属性的绝对索引 + */ +#ifdef __CCE_AICORE__ +__aicore__ inline uint32_t GetAttrAbsIndex(uint32_t coreIdx, uint32_t metaIdx, bool isAIV = false) +{ + if (isAIV) { + return QLI_V2_METADATA_SIZE * AIC_CORE_MAX_NUM + QLD_V2_METADATA_SIZE * coreIdx + metaIdx; + } else { + return QLI_V2_METADATA_SIZE * coreIdx + metaIdx; + } +} +#endif + +namespace detail { + struct QliV2Metadata { + uint32_t qliV2Metadata[AIC_CORE_MAX_NUM][QLI_V2_METADATA_SIZE]; + uint32_t qldV2Metadata[AIV_CORE_MAX_NUM][QLD_V2_METADATA_SIZE]; + }; +}; + +static_assert(QLI_V2_METADATA_TOTAL_SIZE * sizeof(QLI_V2_METADATA_T) >= sizeof(detail::QliV2Metadata)); +}; + +#endif // QUANT_LIGHTNING_INDEXER_V2_METADATA_H diff --git a/xllm_ops/attention/quant_lightning_indexer_v2/op_kernel/quant_lightning_indexer_v2_template_tiling_key.h b/xllm_ops/attention/quant_lightning_indexer_v2/op_kernel/quant_lightning_indexer_v2_template_tiling_key.h new file mode 100644 index 0000000..03c8dd2 --- /dev/null +++ b/xllm_ops/attention/quant_lightning_indexer_v2/op_kernel/quant_lightning_indexer_v2_template_tiling_key.h @@ -0,0 +1,134 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file quant_lightning_indexer_v2_template_tiling_key.h + * \brief + */ + +#ifndef QUANT_LIGHTNING_INDEXER_V2_TEMPLATE_TILING_KEY_H +#define QUANT_LIGHTNING_INDEXER_V2_TEMPLATE_TILING_KEY_H + +#include "ascendc/host_api/tiling/template_argument.h" + +#define QLIV2_TPL_INT8 2 +#define QLIV2_TPL_INT32 3 +#define QLIV2_TPL_FLOAT8_E4M3FN 36 +#define QLIV2_TPL_HIFLOAT8 34 +#define QLIV2_TPL_FLOAT4_E2M1 40 +#define QLIV2_LAYOUT_BSND 0 +#define QLIV2_LAYOUT_TND 1 +#define QLIV2_LAYOUT_PA_BBND 2 + +#define ASCENDC_TPL_4_BW 4 + +// 模板参数支持的范围定义 +#if (__CCE_AICORE__ == 310) +ASCENDC_TPL_ARGS_DECL(QuantLightningIndexerV2, // 算子OpType + ASCENDC_TPL_DTYPE_DECL(DT_Q, QLIV2_TPL_FLOAT8_E4M3FN, QLIV2_TPL_HIFLOAT8, + QLIV2_TPL_FLOAT4_E2M1, QLIV2_TPL_INT8), + ASCENDC_TPL_DTYPE_DECL(DT_K, QLIV2_TPL_FLOAT8_E4M3FN, QLIV2_TPL_HIFLOAT8, + QLIV2_TPL_FLOAT4_E2M1, QLIV2_TPL_INT8), + ASCENDC_TPL_DTYPE_DECL(DT_OUT, QLIV2_TPL_INT32), ASCENDC_TPL_BOOL_DECL(PAGE_ATTENTION, 1, 0), + ASCENDC_TPL_UINT_DECL(Q_LAYOUT_T, ASCENDC_TPL_4_BW, ASCENDC_TPL_UI_LIST, QLIV2_LAYOUT_BSND, + QLIV2_LAYOUT_TND), + ASCENDC_TPL_UINT_DECL(K_LAYOUT_T, ASCENDC_TPL_4_BW, ASCENDC_TPL_UI_LIST, QLIV2_LAYOUT_BSND, + QLIV2_LAYOUT_TND, QLIV2_LAYOUT_PA_BBND), ); +// 支持的模板参数组合 +// 用于调用GET_TPL_TILING_KEY获取TilingKey时，接口内部校验TilingKey是否合法 +ASCENDC_TPL_SEL( + ASCENDC_TPL_ARGS_SEL(ASCENDC_TPL_DTYPE_SEL(DT_Q, QLIV2_TPL_FLOAT8_E4M3FN), + ASCENDC_TPL_DTYPE_SEL(DT_K, QLIV2_TPL_FLOAT8_E4M3FN), + ASCENDC_TPL_DTYPE_SEL(DT_OUT, QLIV2_TPL_INT32), ASCENDC_TPL_BOOL_SEL(PAGE_ATTENTION, 1), + ASCENDC_TPL_UINT_SEL(Q_LAYOUT_T, ASCENDC_TPL_UI_LIST, QLIV2_LAYOUT_BSND, QLIV2_LAYOUT_TND), + ASCENDC_TPL_UINT_SEL(K_LAYOUT_T, ASCENDC_TPL_UI_LIST, QLIV2_LAYOUT_PA_BBND), ), + ASCENDC_TPL_ARGS_SEL(ASCENDC_TPL_DTYPE_SEL(DT_Q, QLIV2_TPL_FLOAT8_E4M3FN), + ASCENDC_TPL_DTYPE_SEL(DT_K, QLIV2_TPL_FLOAT8_E4M3FN), + ASCENDC_TPL_DTYPE_SEL(DT_OUT, QLIV2_TPL_INT32), ASCENDC_TPL_BOOL_SEL(PAGE_ATTENTION, 0), + ASCENDC_TPL_UINT_SEL(Q_LAYOUT_T, ASCENDC_TPL_UI_LIST, QLIV2_LAYOUT_BSND), + ASCENDC_TPL_UINT_SEL(K_LAYOUT_T, ASCENDC_TPL_UI_LIST, QLIV2_LAYOUT_BSND), ), + ASCENDC_TPL_ARGS_SEL(ASCENDC_TPL_DTYPE_SEL(DT_Q, QLIV2_TPL_FLOAT8_E4M3FN), + ASCENDC_TPL_DTYPE_SEL(DT_K, QLIV2_TPL_FLOAT8_E4M3FN), + ASCENDC_TPL_DTYPE_SEL(DT_OUT, QLIV2_TPL_INT32), ASCENDC_TPL_BOOL_SEL(PAGE_ATTENTION, 0), + ASCENDC_TPL_UINT_SEL(Q_LAYOUT_T, ASCENDC_TPL_UI_LIST, QLIV2_LAYOUT_TND), + ASCENDC_TPL_UINT_SEL(K_LAYOUT_T, ASCENDC_TPL_UI_LIST, QLIV2_LAYOUT_TND), ), + + ASCENDC_TPL_ARGS_SEL(ASCENDC_TPL_DTYPE_SEL(DT_Q, QLIV2_TPL_HIFLOAT8), + ASCENDC_TPL_DTYPE_SEL(DT_K, QLIV2_TPL_HIFLOAT8), + ASCENDC_TPL_DTYPE_SEL(DT_OUT, QLIV2_TPL_INT32), ASCENDC_TPL_BOOL_SEL(PAGE_ATTENTION, 1), + ASCENDC_TPL_UINT_SEL(Q_LAYOUT_T, ASCENDC_TPL_UI_LIST, QLIV2_LAYOUT_BSND, QLIV2_LAYOUT_TND), + ASCENDC_TPL_UINT_SEL(K_LAYOUT_T, ASCENDC_TPL_UI_LIST, QLIV2_LAYOUT_PA_BBND), ), + ASCENDC_TPL_ARGS_SEL(ASCENDC_TPL_DTYPE_SEL(DT_Q, QLIV2_TPL_HIFLOAT8), + ASCENDC_TPL_DTYPE_SEL(DT_K, QLIV2_TPL_HIFLOAT8), + ASCENDC_TPL_DTYPE_SEL(DT_OUT, QLIV2_TPL_INT32), ASCENDC_TPL_BOOL_SEL(PAGE_ATTENTION, 0), + ASCENDC_TPL_UINT_SEL(Q_LAYOUT_T, ASCENDC_TPL_UI_LIST, QLIV2_LAYOUT_BSND), + ASCENDC_TPL_UINT_SEL(K_LAYOUT_T, ASCENDC_TPL_UI_LIST, QLIV2_LAYOUT_BSND), ), + ASCENDC_TPL_ARGS_SEL(ASCENDC_TPL_DTYPE_SEL(DT_Q, QLIV2_TPL_HIFLOAT8), + ASCENDC_TPL_DTYPE_SEL(DT_K, QLIV2_TPL_HIFLOAT8), + ASCENDC_TPL_DTYPE_SEL(DT_OUT, QLIV2_TPL_INT32), ASCENDC_TPL_BOOL_SEL(PAGE_ATTENTION, 0), + ASCENDC_TPL_UINT_SEL(Q_LAYOUT_T, ASCENDC_TPL_UI_LIST, QLIV2_LAYOUT_TND), + ASCENDC_TPL_UINT_SEL(K_LAYOUT_T, ASCENDC_TPL_UI_LIST, QLIV2_LAYOUT_TND), ), + + ASCENDC_TPL_ARGS_SEL(ASCENDC_TPL_DTYPE_SEL(DT_Q, QLIV2_TPL_FLOAT4_E2M1), + ASCENDC_TPL_DTYPE_SEL(DT_K, QLIV2_TPL_FLOAT4_E2M1), + ASCENDC_TPL_DTYPE_SEL(DT_OUT, QLIV2_TPL_INT32), ASCENDC_TPL_BOOL_SEL(PAGE_ATTENTION, 1), + ASCENDC_TPL_UINT_SEL(Q_LAYOUT_T, ASCENDC_TPL_UI_LIST, QLIV2_LAYOUT_BSND, QLIV2_LAYOUT_TND), + ASCENDC_TPL_UINT_SEL(K_LAYOUT_T, ASCENDC_TPL_UI_LIST, QLIV2_LAYOUT_PA_BBND), ), + ASCENDC_TPL_ARGS_SEL(ASCENDC_TPL_DTYPE_SEL(DT_Q, QLIV2_TPL_FLOAT4_E2M1), + ASCENDC_TPL_DTYPE_SEL(DT_K, QLIV2_TPL_FLOAT4_E2M1), + ASCENDC_TPL_DTYPE_SEL(DT_OUT, QLIV2_TPL_INT32), ASCENDC_TPL_BOOL_SEL(PAGE_ATTENTION, 0), + ASCENDC_TPL_UINT_SEL(Q_LAYOUT_T, ASCENDC_TPL_UI_LIST, QLIV2_LAYOUT_BSND), + ASCENDC_TPL_UINT_SEL(K_LAYOUT_T, ASCENDC_TPL_UI_LIST, QLIV2_LAYOUT_BSND), ), + ASCENDC_TPL_ARGS_SEL(ASCENDC_TPL_DTYPE_SEL(DT_Q, QLIV2_TPL_FLOAT4_E2M1), + ASCENDC_TPL_DTYPE_SEL(DT_K, QLIV2_TPL_FLOAT4_E2M1), + ASCENDC_TPL_DTYPE_SEL(DT_OUT, QLIV2_TPL_INT32), ASCENDC_TPL_BOOL_SEL(PAGE_ATTENTION, 0), + ASCENDC_TPL_UINT_SEL(Q_LAYOUT_T, ASCENDC_TPL_UI_LIST, QLIV2_LAYOUT_TND), + ASCENDC_TPL_UINT_SEL(K_LAYOUT_T, ASCENDC_TPL_UI_LIST, QLIV2_LAYOUT_TND), ), + ASCENDC_TPL_ARGS_SEL(ASCENDC_TPL_DTYPE_SEL(DT_Q, QLIV2_TPL_INT8), + ASCENDC_TPL_DTYPE_SEL(DT_K, QLIV2_TPL_INT8), + ASCENDC_TPL_DTYPE_SEL(DT_OUT, QLIV2_TPL_INT32), ASCENDC_TPL_BOOL_SEL(PAGE_ATTENTION, 1), + ASCENDC_TPL_UINT_SEL(Q_LAYOUT_T, ASCENDC_TPL_UI_LIST, QLIV2_LAYOUT_BSND, QLIV2_LAYOUT_TND), + ASCENDC_TPL_UINT_SEL(K_LAYOUT_T, ASCENDC_TPL_UI_LIST, QLIV2_LAYOUT_PA_BBND), ), + ASCENDC_TPL_ARGS_SEL(ASCENDC_TPL_DTYPE_SEL(DT_Q, QLIV2_TPL_INT8), + ASCENDC_TPL_DTYPE_SEL(DT_K, QLIV2_TPL_INT8), + ASCENDC_TPL_DTYPE_SEL(DT_OUT, QLIV2_TPL_INT32), ASCENDC_TPL_BOOL_SEL(PAGE_ATTENTION, 0), + ASCENDC_TPL_UINT_SEL(Q_LAYOUT_T, ASCENDC_TPL_UI_LIST, QLIV2_LAYOUT_BSND), + ASCENDC_TPL_UINT_SEL(K_LAYOUT_T, ASCENDC_TPL_UI_LIST, QLIV2_LAYOUT_BSND), ), + ASCENDC_TPL_ARGS_SEL(ASCENDC_TPL_DTYPE_SEL(DT_Q, QLIV2_TPL_INT8), + ASCENDC_TPL_DTYPE_SEL(DT_K, QLIV2_TPL_INT8), + ASCENDC_TPL_DTYPE_SEL(DT_OUT, QLIV2_TPL_INT32), ASCENDC_TPL_BOOL_SEL(PAGE_ATTENTION, 0), + ASCENDC_TPL_UINT_SEL(Q_LAYOUT_T, ASCENDC_TPL_UI_LIST, QLIV2_LAYOUT_TND), + ASCENDC_TPL_UINT_SEL(K_LAYOUT_T, ASCENDC_TPL_UI_LIST, QLIV2_LAYOUT_TND),)); +#else +ASCENDC_TPL_ARGS_DECL(QuantLightningIndexerV2, // 算子OpType + ASCENDC_TPL_DTYPE_DECL(DT_Q, QLIV2_TPL_INT8), ASCENDC_TPL_DTYPE_DECL(DT_K, QLIV2_TPL_INT8), + ASCENDC_TPL_DTYPE_DECL(DT_OUT, QLIV2_TPL_INT32), ASCENDC_TPL_BOOL_DECL(PAGE_ATTENTION, 1, 0), + ASCENDC_TPL_UINT_DECL(Q_LAYOUT_T, ASCENDC_TPL_4_BW, ASCENDC_TPL_UI_LIST, QLIV2_LAYOUT_BSND, + QLIV2_LAYOUT_TND), + ASCENDC_TPL_UINT_DECL(K_LAYOUT_T, ASCENDC_TPL_4_BW, ASCENDC_TPL_UI_LIST, QLIV2_LAYOUT_BSND, + QLIV2_LAYOUT_TND, QLIV2_LAYOUT_PA_BBND), ); +// 支持的模板参数组合 +// 用于调用GET_TPL_TILING_KEY获取TilingKey时，接口内部校验TilingKey是否合法 +ASCENDC_TPL_SEL( + ASCENDC_TPL_ARGS_SEL(ASCENDC_TPL_DTYPE_SEL(DT_Q, QLIV2_TPL_INT8), ASCENDC_TPL_DTYPE_SEL(DT_K, QLIV2_TPL_INT8), + ASCENDC_TPL_DTYPE_SEL(DT_OUT, QLIV2_TPL_INT32), ASCENDC_TPL_BOOL_SEL(PAGE_ATTENTION, 1), + ASCENDC_TPL_UINT_SEL(Q_LAYOUT_T, ASCENDC_TPL_UI_LIST, QLIV2_LAYOUT_BSND, QLIV2_LAYOUT_TND), + ASCENDC_TPL_UINT_SEL(K_LAYOUT_T, ASCENDC_TPL_UI_LIST, QLIV2_LAYOUT_PA_BBND), ), + ASCENDC_TPL_ARGS_SEL(ASCENDC_TPL_DTYPE_SEL(DT_Q, QLIV2_TPL_INT8), ASCENDC_TPL_DTYPE_SEL(DT_K, QLIV2_TPL_INT8), + ASCENDC_TPL_DTYPE_SEL(DT_OUT, QLIV2_TPL_INT32), ASCENDC_TPL_BOOL_SEL(PAGE_ATTENTION, 0), + ASCENDC_TPL_UINT_SEL(Q_LAYOUT_T, ASCENDC_TPL_UI_LIST, QLIV2_LAYOUT_BSND), + ASCENDC_TPL_UINT_SEL(K_LAYOUT_T, ASCENDC_TPL_UI_LIST, QLIV2_LAYOUT_BSND), ), + ASCENDC_TPL_ARGS_SEL(ASCENDC_TPL_DTYPE_SEL(DT_Q, QLIV2_TPL_INT8), ASCENDC_TPL_DTYPE_SEL(DT_K, QLIV2_TPL_INT8), + ASCENDC_TPL_DTYPE_SEL(DT_OUT, QLIV2_TPL_INT32), ASCENDC_TPL_BOOL_SEL(PAGE_ATTENTION, 0), + ASCENDC_TPL_UINT_SEL(Q_LAYOUT_T, ASCENDC_TPL_UI_LIST, QLIV2_LAYOUT_TND), + ASCENDC_TPL_UINT_SEL(K_LAYOUT_T, ASCENDC_TPL_UI_LIST, QLIV2_LAYOUT_TND), ), ); +#endif + +#endif // QUANT_LIGHTNING_INDEXER_V2_TEMPLATE_TILING_KEY_H diff --git a/xllm_ops/attention/quant_lightning_indexer_v2/tests/assets/impl/compare.py b/xllm_ops/attention/quant_lightning_indexer_v2/tests/assets/impl/compare.py new file mode 100644 index 0000000..abafc94 --- /dev/null +++ b/xllm_ops/attention/quant_lightning_indexer_v2/tests/assets/impl/compare.py @@ -0,0 +1,194 @@ +#!/usr/bin/python +# -*- coding: utf-8 -*- +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +"""TTK result adapter for the QuantLightningIndexer V2 pytest TopK comparison.""" + +import importlib.util +import logging +import sys +import threading +from pathlib import Path + +import numpy as np +import torch + + +class PytestV2TopKComparator: + """Run the pytest V2 TopK compare with replay-safe data from the TestSpec.""" + + def __init__(self): + self.module = None + self.lock = threading.Lock() + + def load_module(self): + if self.module is not None: + return self.module + with self.lock: + if self.module is not None: + return self.module + pytest_dir = Path(__file__).resolve().parents[2] / "pytest" + module_path = pytest_dir / "result_compare_method.py" + module_name = "qli_v2_ttk_pytest_compare" + inserted = str(pytest_dir) not in sys.path + original_basic_config = logging.basicConfig + if inserted: + sys.path.insert(0, str(pytest_dir)) + try: + logging.basicConfig = lambda *args, **kwargs: None + spec = importlib.util.spec_from_file_location(module_name, module_path) + if spec is None or spec.loader is None: + raise ImportError(f"cannot create import spec for {module_path}") + module = importlib.util.module_from_spec(spec) + sys.modules[module_name] = module + spec.loader.exec_module(module) + self.module = module + except Exception as exc: + sys.modules.pop(module_name, None) + raise RuntimeError( + "Failed to load QuantLightningIndexerV2 pytest compare; " + f"module={module_path.resolve()}; " + f"original error: {type(exc).__name__}: {exc}" + ) from exc + finally: + logging.basicConfig = original_basic_config + if inserted and str(pytest_dir) in sys.path: + sys.path.remove(str(pytest_dir)) + return self.module + + @staticmethod + def to_torch(value): + if value is None: + return None + if torch.is_tensor(value): + return value.detach().cpu().clone() + array = np.array(value, copy=True, order="C") + dtype_name = str(array.dtype) + custom_dtypes = { + "bfloat16": (np.uint16, torch.bfloat16), + "float8_e4m3fn": (np.uint8, torch.float8_e4m3fn), + "float8_e5m2": (np.uint8, torch.float8_e5m2), + } + if dtype_name in custom_dtypes: + storage_dtype, torch_dtype = custom_dtypes[dtype_name] + storage = np.ascontiguousarray(array).view(storage_dtype) + return torch.from_numpy(storage).view(torch_dtype).reshape(array.shape) + return torch.from_numpy(array) + + @staticmethod + def result_dict(result, stage): + if not isinstance(result, (list, tuple)) or len(result) < 2: + raise ValueError(f"pytest {stage} returned invalid result: {result!r}") + status, precision = result[:2] + passed = str(status).strip().lower() == "pass" + return { + "pass": passed, + "precision": float(precision), + "error_info": None + if passed + else (f"pytest QuantLightningIndexerV2 {stage} returned {status!r}"), + } + + def compare(self, *outputs, compare_data=None): + if compare_data is None: + raise ValueError( + "QuantLightningIndexerV2 pytest compare data is unavailable" + ) + if len(outputs) < 2 or len(outputs) % 2 != 0: + return { + "pass": False, + "precision": "invalid", + "error_info": "compare expects NPU outputs followed by golden outputs", + } + params = compare_data.get("params") + topk_value = compare_data.get("topk_value") + if params is None or topk_value is None: + raise ValueError( + "QuantLightningIndexerV2 pytest compare data lacks params or topk_value" + ) + half = len(outputs) // 2 + npu_outputs = outputs[:half] + golden_outputs = outputs[half:] + if tuple(getattr(npu_outputs[0], "shape", ())) != tuple( + getattr(golden_outputs[0], "shape", ()) + ): + return { + "pass": False, + "precision": "shape_mismatch", + "error_info": ( + "index output shape mismatch: " + f"npu={getattr(npu_outputs[0], 'shape', None)}, " + f"golden={getattr(golden_outputs[0], 'shape', None)}" + ), + } + return_value = bool(params[-2]) + if return_value and half < 2: + return { + "pass": False, + "precision": "missing_output", + "error_info": "return_value is enabled but the NPU sparse-value output is missing", + } + npu_values = npu_outputs[1] if half > 1 else torch.empty(0) + golden_values = golden_outputs[1] if half > 1 else torch.empty(0) + if return_value and tuple(getattr(npu_values, "shape", ())) != tuple( + getattr(golden_values, "shape", ()) + ): + return { + "pass": False, + "precision": "shape_mismatch", + "error_info": ( + "sparse-value output shape mismatch: " + f"npu={getattr(npu_values, 'shape', None)}, " + f"golden={getattr(golden_values, 'shape', None)}" + ), + } + npu_indices = self.to_torch(npu_outputs[0]) + npu_values = self.to_torch(npu_values) + golden_values = self.to_torch(golden_values) + if return_value: + npu_values, sort_order = npu_values.sort(dim=-1, descending=True) + npu_indices = torch.gather(npu_indices, dim=-1, index=sort_order) + golden_indices = self.to_torch(golden_outputs[0]) + topk_value = self.to_torch(topk_value) + output_idx_offset = self.to_torch(compare_data.get("output_idx_offset")) + golden_values_for_index = golden_values.detach().cpu().float().numpy() + npu_values_for_index = npu_values.detach().cpu().float().numpy() + module = self.load_module() + index_result = module.check_result( + golden_indices, + npu_indices, + topk_value, + output_idx_offset, + params, + golden_values_for_index, + npu_values_for_index, + ) + results = [self.result_dict(index_result, "index compare")] + if return_value: + value_result = module.check_result_return_value( + golden_values, + npu_values, + params, + golden_indices, + npu_indices, + topk_value, + output_idx_offset, + ) + results.append(self.result_dict(value_result, "sparse-value compare")) + return results + + +COMPARATOR = PytestV2TopKComparator() + + +def compare(*outputs, compare_data=None): + """Compare V2 TopK outputs with the canonical pytest policy.""" + return COMPARATOR.compare(*outputs, compare_data=compare_data) diff --git a/xllm_ops/attention/quant_lightning_indexer_v2/tests/assets/impl/golden.py b/xllm_ops/attention/quant_lightning_indexer_v2/tests/assets/impl/golden.py new file mode 100644 index 0000000..1972db5 --- /dev/null +++ b/xllm_ops/attention/quant_lightning_indexer_v2/tests/assets/impl/golden.py @@ -0,0 +1,212 @@ +#!/usr/bin/python +# -*- coding: utf-8 -*- +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +"""CPU Golden adapter for QuantLightningIndexer V2 TTK cases.""" + +import importlib.util +import sys +from pathlib import Path + +import numpy as np +import torch + +PYTEST_MODULE_NAME = "qli_v2_pytest_golden" +PYTEST_MODULE_FILE = "quant_lightning_indexer_v2_golden.py" + + +class CaseDataStore: + """Share pytest data in-process and return compact metadata API inputs.""" + + def __init__(self): + self.case_data = {} + self.active_testcase_name = None + + def clear(self): + self.case_data.clear() + self.active_testcase_name = None + + def put(self, testcase_name, data): + if testcase_name is not None: + self.case_data[str(testcase_name)] = data + + def get(self, testcase_name): + if testcase_name is None: + return None + return self.case_data.get(str(testcase_name)) + + def discard(self, data): + for testcase_name, stored in tuple(self.case_data.items()): + if stored is data: + self.case_data.pop(testcase_name, None) + + +CASE_DATA = CaseDataStore() + + +def load_pytest_golden(): + """Load the pytest CPU reference only when the Golden stage needs it.""" + if PYTEST_MODULE_NAME in sys.modules: + return sys.modules[PYTEST_MODULE_NAME] + pytest_dir = Path(__file__).resolve().parents[2] / "pytest" + path = pytest_dir / PYTEST_MODULE_FILE + inserted = str(pytest_dir) not in sys.path + if inserted: + sys.path.insert(0, str(pytest_dir)) + try: + spec = importlib.util.spec_from_file_location(PYTEST_MODULE_NAME, path) + if spec is None or spec.loader is None: + raise ImportError(f"cannot create import spec for {path}") + module = importlib.util.module_from_spec(spec) + sys.modules[PYTEST_MODULE_NAME] = module + spec.loader.exec_module(module) + return module + except Exception as exc: + sys.modules.pop(PYTEST_MODULE_NAME, None) + raise RuntimeError( + "Failed to load QuantLightningIndexer V2 pytest Golden module; " + f"module={path.resolve()}; original error: {type(exc).__name__}: {exc}" + ) from exc + finally: + if inserted: + sys.path.remove(str(pytest_dir)) + + +def get_case_data(testcase_name): + return CASE_DATA.get(testcase_name) + + +def materialize_golden(data): + if data.get("cpu_result") is None: + load_pytest_golden().generate_cpu_golden(data) + return data + + +def activate_case_data(testcase_name): + data = CASE_DATA.get(testcase_name) + if data is None: + raise RuntimeError( + "QuantLightningIndexer V2 Golden requires pytest data from the input stage" + ) + CASE_DATA.active_testcase_name = str(testcase_name) + return materialize_golden(data) + + +def get_compare_data(testcase_name): + if testcase_name is None: + testcase_name = CASE_DATA.active_testcase_name + if testcase_name is None: + return None + data = CASE_DATA.get(testcase_name) + return None if data is None else materialize_golden(data) + + +def set_compare_data(testcase_name, data): + name = str(testcase_name) + CASE_DATA.active_testcase_name = name + CASE_DATA.case_data[name] = data + + +def discard_compare_data(data): + CASE_DATA.discard(data) + CASE_DATA.active_testcase_name = None + + +def cpu_quant_lightning_indexer_v2( + query, + key, + weights, + query_dequant_scale, + key_dequant_scale, + topk, + quant_mode, + *, + return_value=0, + testcase_name=None, + **kwargs, +): + """Materialize Golden from the exact pytest data produced by input.""" + del query, key, weights, query_dequant_scale, key_dequant_scale + del topk, quant_mode, kwargs + data = activate_case_data(testcase_name) + if int(return_value): + sparse_value = data["cpu_topk_value"] + else: + sparse_value = torch.zeros(0, dtype=data["topk_value"].dtype) + return data["cpu_result"], sparse_value + + +def cpu_aclnn_qli_v2( + query, + key, + weights, + query_dequant_scale, + key_dequant_scale, + cu_seqlens_q, + cu_seqlens_k, + seqused_q, + seqused_k, + cmp_residual_k, + block_table, + output_idx_offset, + metadata, + topk, + quant_mode, + max_seqlen_q, + layout_q, + layout_k, + mask_mode, + cmp_ratio, + return_value, + sparse_indices_out, + sparse_values_out, + testcase_name=None, + **kwargs, +): + """Return the pytest Golden for the ACLNN C API parameter order.""" + del ( + cu_seqlens_q, + cu_seqlens_k, + seqused_q, + seqused_k, + cmp_residual_k, + block_table, + output_idx_offset, + metadata, + max_seqlen_q, + layout_q, + layout_k, + mask_mode, + cmp_ratio, + sparse_indices_out, + ) + sparse_indices, sparse_values = cpu_quant_lightning_indexer_v2( + query, + key, + weights, + query_dequant_scale, + key_dequant_scale, + topk, + quant_mode, + return_value=return_value, + testcase_name=testcase_name, + **kwargs, + ) + if not int(return_value): + if sparse_values_out is None: + raise ValueError("ACLNN QLI_V2 requires the sparseValuesOut tensor slot") + if torch.is_tensor(sparse_values_out): + sparse_values = torch.zeros( + tuple(sparse_values_out.shape), dtype=sparse_values_out.dtype + ) + else: + sparse_values = np.zeros_like(np.asarray(sparse_values_out)) + return sparse_indices, sparse_values diff --git a/xllm_ops/attention/quant_lightning_indexer_v2/tests/assets/impl/inputs.py b/xllm_ops/attention/quant_lightning_indexer_v2/tests/assets/impl/inputs.py new file mode 100644 index 0000000..6519b95 --- /dev/null +++ b/xllm_ops/attention/quant_lightning_indexer_v2/tests/assets/impl/inputs.py @@ -0,0 +1,933 @@ +#!/usr/bin/python +# -*- coding: utf-8 -*- +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +"""Input customization for QuantLightningIndexer V2 TTK cases.""" + +import importlib.util +import inspect +import sys +from pathlib import Path + +import numpy as np +import torch + +QUANT_MODE_MXFP4 = 5 +QUANT_MODE_MXFP8 = 3 + + +def restore_mx_input_dtypes(query, key, query_scale, key_scale, quant_mode): + """Restore packed MX dtypes needed by the existing replay compare path.""" + quant_mode = int(quant_mode) + if quant_mode == QUANT_MODE_MXFP8: + qk_dtype = getattr(torch, "float8_e4m3fn", None) + elif quant_mode == QUANT_MODE_MXFP4: + qk_dtype = getattr(torch, "float4_e2m1fn_x2", None) + else: + return query, key, query_scale, key_scale + scale_dtype = getattr(torch, "float8_e8m0fnu", None) + if qk_dtype is None or scale_dtype is None: + raise RuntimeError("current PyTorch does not provide the requested MX dtype") + if query.dtype == torch.uint8: + query = query.view(qk_dtype) + if key.dtype == torch.uint8: + key = key.view(qk_dtype) + if query_scale.dtype == torch.uint8: + query_scale = query_scale.view(scale_dtype) + if key_scale.dtype == torch.uint8: + key_scale = key_scale.view(scale_dtype) + return query, key, query_scale, key_scale + + +class QuantLightningIndexerV2InputAdapter: + """Translate a TTK case and reuse the pytest input/golden generator.""" + + @staticmethod + def module_load_error(stage, path, exc): + return RuntimeError( + "Failed to load QuantLightningIndexerV2 module; " + f"stage={stage}; module={path.resolve()}; " + f"original error: {type(exc).__name__}: {exc}" + ) + + def __init__(self): + self.pytest_golden = None + self.pytest_normalizer = None + + @staticmethod + def load_golden_store(): + name = "qli_v2_ttk_golden" + if name in sys.modules: + return sys.modules[name] + path = Path(__file__).with_name("golden.py") + try: + spec = importlib.util.spec_from_file_location(name, path) + if spec is None or spec.loader is None: + raise ImportError(f"cannot create import spec for {path}") + module = importlib.util.module_from_spec(spec) + sys.modules[name] = module + spec.loader.exec_module(module) + except Exception as exc: + sys.modules.pop(name, None) + raise QuantLightningIndexerV2InputAdapter.module_load_error( + "assets Golden store", path, exc + ) from exc + return module + + @staticmethod + def load_metadata_protocol(): + name = "qli_v2_ttk_metadata_protocol" + if name in sys.modules: + return sys.modules[name] + path = Path(__file__).with_name("metadata_protocol.py") + try: + spec = importlib.util.spec_from_file_location(name, path) + if spec is None or spec.loader is None: + raise ImportError(f"cannot create import spec for {path}") + module = importlib.util.module_from_spec(spec) + sys.modules[name] = module + spec.loader.exec_module(module) + except Exception as exc: + sys.modules.pop(name, None) + raise QuantLightningIndexerV2InputAdapter.module_load_error( + "assets metadata protocol", path, exc + ) from exc + return module + + def load_pytest_golden(self): + if self.pytest_golden is not None: + return self.pytest_golden + pytest_dir = Path(__file__).resolve().parents[2] / "pytest" + path = pytest_dir / "quant_lightning_indexer_v2_golden.py" + name = "qli_v2_pytest_golden" + inserted = str(pytest_dir) not in sys.path + if inserted: + sys.path.insert(0, str(pytest_dir)) + try: + if name in sys.modules: + module = sys.modules[name] + else: + spec = importlib.util.spec_from_file_location(name, path) + if spec is None or spec.loader is None: + raise ImportError(f"cannot create import spec for {path}") + module = importlib.util.module_from_spec(spec) + sys.modules[name] = module + spec.loader.exec_module(module) + self.pytest_golden = module + return module + except Exception as exc: + sys.modules.pop(name, None) + raise self.module_load_error("pytest Golden", path, exc) from exc + finally: + if inserted: + sys.path.remove(str(pytest_dir)) + + def load_pytest_normalizer(self): + if self.pytest_normalizer is not None: + return self.pytest_normalizer + pytest_dir = Path(__file__).resolve().parents[2] / "pytest" + path = pytest_dir / "qliv2_parameter_normalization.py" + name = "qli_v2_pytest_normalizer" + try: + spec = importlib.util.spec_from_file_location(name, path) + if spec is None or spec.loader is None: + raise ImportError(f"cannot create import spec for {path}") + module = importlib.util.module_from_spec(spec) + sys.modules[name] = module + spec.loader.exec_module(module) + self.pytest_normalizer = module + return module + except Exception as exc: + sys.modules.pop(name, None) + raise self.module_load_error( + "pytest parameter normalizer", path, exc + ) from exc + + @staticmethod + def list_value(kwargs, name): + value = kwargs.get(f"{name}_values") + if value is None: + return None + if torch.is_tensor(value): + value = value.detach().cpu().reshape(-1).tolist() + elif isinstance(value, np.ndarray): + value = value.reshape(-1).tolist() + return [int(item) for item in value] + + @staticmethod + def tensor_dtype(tensor): + if torch.is_tensor(tensor): + return tensor.dtype + if "hifloat8" in str(tensor.dtype): + return torch.uint8 + return torch.from_numpy(np.asarray(tensor)).dtype + + @staticmethod + def to_cpu_tensor(tensor): + if tensor is None: + return None + if torch.is_tensor(tensor): + return tensor.detach().cpu() + array = np.asarray(tensor) + if "hifloat8" in str(array.dtype): + array = array.view(np.uint8) + return torch.from_numpy(np.array(array, copy=True)) + + @staticmethod + def prefix_lengths(value): + if not value: + return [] + return [ + int(value[index + 1]) - int(value[index]) for index in range(len(value) - 1) + ] + + @staticmethod + def data_range(input_ranges, index): + if ( + input_ranges + and index < len(input_ranges) + and input_ranges[index] is not None + ): + return repr(list(input_ranges[index])) + return None + + @staticmethod + def qk_dtype_name(tensor, quant_mode): + if quant_mode == QUANT_MODE_MXFP8: + return "FLOAT8_E4M3FN" + if quant_mode == QUANT_MODE_MXFP4: + return "FLOAT4_E2M1FN_X2" + dtype = QuantLightningIndexerV2InputAdapter.tensor_dtype(tensor) + dtype_name = str(tensor.dtype) + if dtype == torch.int8: + return "INT8" + if "float8_e4m3fn" in dtype_name: + return "FLOAT8_E4M3FN" + if dtype == torch.uint8: + return "HIFLOAT8" + raise ValueError(f"unsupported QLI_V2 q/k dtype: {tensor.dtype}") + + @staticmethod + def dequant_dtype_name(tensor, quant_mode): + if quant_mode in (QUANT_MODE_MXFP8, QUANT_MODE_MXFP4): + return "FLOAT8_E8M0FNU" + dtype = QuantLightningIndexerV2InputAdapter.tensor_dtype(tensor) + if dtype == torch.float16: + return "FP16" + if dtype == torch.float32: + return "FP32" + raise ValueError(f"unsupported QLI_V2 dequant dtype: {tensor.dtype}") + + @staticmethod + def weight_dtype_name(tensor): + mapping = { + torch.int8: "INT8", + torch.uint8: "UINT8", + torch.float16: "FP16", + torch.float32: "FP32", + torch.bfloat16: "BF16", + } + dtype = QuantLightningIndexerV2InputAdapter.tensor_dtype(tensor) + if dtype in mapping: + return mapping[dtype] + if "float8_e4m3fn" in str(tensor.dtype): + return "FLOAT8_E4M3FN" + raise ValueError(f"unsupported QLI_V2 weight dtype: {tensor.dtype}") + + @staticmethod + def pytest_uses_weight_dtype(pytest_golden): + parameters = inspect.signature( + pytest_golden.GeneralizedQLIV2.__init__ + ).parameters + return "weight_dtype" in parameters + + def geometry( + self, query, key, layout_query, layout_key, cu_q, cu_k, seq_q, seq_k, quant_mode + ): + q_lengths = self.prefix_lengths(cu_q) or (seq_q or []) + k_lengths = self.prefix_lengths(cu_k) or (seq_k or []) + if layout_query == "BSND": + batch_size, q_seq, q_head_num, head_dim = [ + int(item) for item in query.shape + ] + q_t_size = 0 + elif layout_query == "TND": + q_t_size, q_head_num, head_dim = [int(item) for item in query.shape] + batch_size = len(q_lengths) + q_seq = max(q_lengths, default=q_t_size) + else: + raise ValueError(f"unsupported QLI_V2 query layout: {layout_query}") + + if quant_mode == QUANT_MODE_MXFP4: + head_dim *= 2 + + if layout_key == "BSND": + _, k_seq, k_head_num, _ = [int(item) for item in key.shape] + k_t_size = 0 + block_size = 0 + block_num = 0 + elif layout_key == "TND": + k_t_size, k_head_num, _ = [int(item) for item in key.shape] + k_seq = max(k_lengths, default=k_t_size) + block_size = 0 + block_num = 0 + elif layout_key == "PA_BBND": + block_num, block_size, k_head_num, _ = [int(item) for item in key.shape] + capacity = block_num * block_size + per_batch_capacity = ( + capacity // batch_size if batch_size > 0 else block_size + ) + k_seq = max(k_lengths, default=per_batch_capacity) + k_t_size = 0 + else: + raise ValueError(f"unsupported QLI_V2 key layout: {layout_key}") + return ( + batch_size, + q_seq, + k_seq, + q_t_size, + k_t_size, + q_head_num, + k_head_num, + head_dim, + block_size, + block_num, + ) + + def build_case_params( + self, + query, + key, + weights, + query_dequant_scale, + layout_query, + layout_key, + kwargs, + include_weight_dtype=False, + ): + cu_q = self.list_value(kwargs, "cu_seqlens_q") + cu_k = self.list_value(kwargs, "cu_seqlens_k") + seq_q = self.list_value(kwargs, "seqused_q") + seq_k = self.list_value(kwargs, "seqused_k") + residual = self.list_value(kwargs, "cmp_residual_k") + quant_mode = kwargs.get("quant_mode") + quant_mode = None if quant_mode is None else int(quant_mode) + geometry = self.geometry( + query, + key, + layout_query, + layout_key, + cu_q, + cu_k, + seq_q, + seq_k, + quant_mode, + ) + input_ranges = ( + kwargs.get("qli_input_ranges") or kwargs.get("input_ranges") or () + ) + output_range = kwargs.get("output_idx_offset_range") + output_range = None if output_range is None else repr(list(output_range)) + max_seqlen_q = kwargs.get("max_seqlen_q") + max_seqlen_q = None if max_seqlen_q is None else int(max_seqlen_q) + dtype_params = (self.qk_dtype_name(query, quant_mode),) + if include_weight_dtype: + dtype_params += (self.weight_dtype_name(weights),) + dtype_params += ( + self.dequant_dtype_name(query_dequant_scale, quant_mode), + "INT32", + ) + params = ( + geometry + + dtype_params + + ( + cu_q, + cu_k, + seq_q, + seq_k, + residual, + max_seqlen_q, + quant_mode, + layout_query, + layout_key, + kwargs.get("sparse_count"), + kwargs.get("sparse_mode"), + self.data_range(input_ranges, 0), + self.data_range(input_ranges, 1), + self.data_range(input_ranges, 2), + self.data_range(input_ranges, 3), + self.data_range(input_ranges, 4), + kwargs.get("cmp_ratio"), + kwargs.get("return_value"), + output_range, + ) + ) + return self.load_pytest_normalizer().normalize_qliv2_params(params) + + @staticmethod + def copy_tensor(dst, src, name, packed_dtype=None): + if dst is None: + if src is not None: + raise ValueError( + f"{name} is absent from CSV but pytest generator produced a tensor" + ) + return + if src is None: + raise ValueError( + f"{name} is present in CSV but pytest generator returned None" + ) + src_cpu = QuantLightningIndexerV2InputAdapter.to_cpu_tensor(src) + src_cpu = src_cpu.contiguous() + if packed_dtype is not None: + src_cpu = src_cpu.to(packed_dtype) + if ( + torch.is_tensor(dst) + and dst.dtype == torch.uint8 + and src_cpu.element_size() == 1 + ): + src_cpu = src_cpu.view(torch.uint8) + if tuple(dst.shape) != tuple(src_cpu.shape): + raise ValueError( + f"{name} shape mismatch: TTK={tuple(dst.shape)} " + f"pytest={tuple(src_cpu.shape)}" + ) + if torch.is_tensor(dst): + src_tensor = torch.as_tensor(src_cpu) + dst.copy_(src_tensor.to(dtype=dst.dtype, device=dst.device)) + return + + dst_array = np.asarray(dst) + src_array = np.asarray(src_cpu) + if "hifloat8" in str(dst_array.dtype): + np.copyto(dst_array.view(np.uint8), src_array.view(np.uint8)) + else: + np.copyto(dst_array, src_array.astype(dst_array.dtype, copy=False)) + + @staticmethod + def tensor_values(tensor): + if tensor is None: + return None + tensor = QuantLightningIndexerV2InputAdapter.to_cpu_tensor(tensor) + return [int(value) for value in tensor.reshape(-1).tolist()] + + @staticmethod + def unpack_mxfp4(tensor, fp4_values): + """Unpack two FP4 E2M1 values stored in each uint8 byte.""" + packed = tensor.view(torch.uint8).contiguous() + codes = torch.stack( + (packed & 0x0F, packed >> 4), + dim=-1, + ).flatten(-2) + return fp4_values[codes.to(torch.long)] + + @staticmethod + def restore_paged_tensor(tensor, block_table, batch_size, sequence_length): + """Restore a paged key or scale tensor for the pytest compare model.""" + physical = QuantLightningIndexerV2InputAdapter.to_cpu_tensor(tensor) + table = QuantLightningIndexerV2InputAdapter.to_cpu_tensor(block_table).to( + torch.int64 + ) + if physical.ndim < 3: + raise ValueError( + "paged tensor must have at least 3 dimensions, " + f"got shape {tuple(physical.shape)}" + ) + block_size, head_num = int(physical.shape[1]), int(physical.shape[2]) + trailing = tuple(int(dim) for dim in physical.shape[3:]) + logical = torch.zeros( + (batch_size, head_num, sequence_length, *trailing), + dtype=physical.dtype, + ) + for batch_idx in range(batch_size): + for logical_block, block_id_value in enumerate(table[batch_idx].tolist()): + if block_id_value < 0: + continue + if block_id_value >= physical.shape[0]: + raise ValueError( + f"block id {block_id_value} exceeds paged block count" + ) + start = logical_block * block_size + if start >= sequence_length: + break + count = min(block_size, sequence_length - start) + block = physical[block_id_value, :count] + permutation = (1, 0, *range(2, block.ndim)) + logical[batch_idx, :, start : start + count] = block.permute( + *permutation + ) + return logical + + @staticmethod + def normalize_compare_attributes(compare_context): + attributes = dict(compare_context.attributes) + aliases = { + "topk": "sparse_count", + "mask_mode": "sparse_mode", + "layout_q": "layout_query", + "layout_k": "layout_key", + "quantMode": "quant_mode", + "maxSeqlenQ": "max_seqlen_q", + "layoutQOptional": "layout_query", + "layoutKOptional": "layout_key", + "maskMode": "sparse_mode", + "cmpRatio": "cmp_ratio", + "returnValue": "return_value", + } + for source, target in aliases.items(): + if target not in attributes and source in attributes: + attributes[target] = attributes[source] + return attributes + + def rebuild_compare_data(self, compare_context): + """Rebuild only the pytest TopK compare context from replayed inputs.""" + tensors = tuple(compare_context.input_tensors or ()) + if len(tensors) < 12: + raise ValueError( + "QuantLightningIndexerV2 compare context requires twelve input slots" + ) + ( + query, + key, + weights, + query_scale, + key_scale, + cu_q, + cu_k, + seq_q, + seq_k, + residual, + block_table, + offset, + ) = tensors[:12] + attributes = self.normalize_compare_attributes(compare_context) + for name, tensor in ( + ("cu_seqlens_q", cu_q), + ("cu_seqlens_k", cu_k), + ("seqused_q", seq_q), + ("seqused_k", seq_k), + ("cmp_residual_k", residual), + ): + values = self.tensor_values(tensor) + if values is not None: + attributes[f"{name}_values"] = values + + layout_q = attributes.get("layout_query") + layout_k = attributes.get("layout_key") + if layout_q is None or layout_k is None: + raise ValueError( + "QLI_V2 replay compare requires layout_query and layout_key " + "from attributes" + ) + pytest_golden = self.load_pytest_golden() + uses_weight_dtype = self.pytest_uses_weight_dtype(pytest_golden) + params = self.build_case_params( + query, + key, + weights, + query_scale, + layout_q, + layout_k, + attributes, + include_weight_dtype=uses_weight_dtype, + ) + if uses_weight_dtype: + ( + batch_size, + q_seq, + k_seq, + q_t_size, + k_t_size, + q_head_num, + k_head_num, + head_dim, + block_size, + block_num, + _, + _, + _, + _, + cu_q_values, + cu_k_values, + seq_q_values, + seq_k_values, + residual_values, + max_seqlen_q, + quant_mode, + _, + _, + sparse_count, + sparse_mode, + _, + _, + _, + _, + _, + cmp_ratio, + return_value, + _, + ) = params + else: + ( + batch_size, + q_seq, + k_seq, + q_t_size, + k_t_size, + q_head_num, + k_head_num, + head_dim, + block_size, + block_num, + _, + _, + _, + cu_q_values, + cu_k_values, + seq_q_values, + seq_k_values, + residual_values, + max_seqlen_q, + quant_mode, + _, + _, + sparse_count, + sparse_mode, + _, + _, + _, + _, + _, + cmp_ratio, + return_value, + _, + ) = params + + q_lengths = ( + self.prefix_lengths(cu_q_values) + if layout_q == "TND" + else (seq_q_values or [q_seq] * batch_size) + ) + k_lengths = ( + self.prefix_lengths(cu_k_values) + if layout_k == "TND" + else (seq_k_values or [k_seq] * batch_size) + ) + residual_for_cpu = ( + [0] * batch_size + if cmp_ratio == 1 or sparse_mode == 0 + else list(residual_values) + ) + + def as_int_tensor(value): + return None if value is None else torch.tensor(value, dtype=torch.int32) + + cu_q_cpu = as_int_tensor(cu_q_values) + cu_k_cpu = as_int_tensor(cu_k_values) + seq_q_cpu = as_int_tensor(seq_q_values) + seq_k_cpu = as_int_tensor(seq_k_values) + query, key, query_scale, key_scale = restore_mx_input_dtypes( + self.to_cpu_tensor(query), + self.to_cpu_tensor(key), + self.to_cpu_tensor(query_scale), + self.to_cpu_tensor(key_scale), + quant_mode, + ) + weights = self.to_cpu_tensor(weights) + block_table = self.to_cpu_tensor(block_table) + offset = self.to_cpu_tensor(offset) + qk_dtype = query.dtype + if quant_mode == QUANT_MODE_MXFP4: + query = self.unpack_mxfp4(query, pytest_golden.FP4_E2M1_VALUES) + key = self.unpack_mxfp4(key, pytest_golden.FP4_E2M1_VALUES) + model_args = [ + batch_size, + q_seq, + k_seq, + q_t_size, + k_t_size, + q_head_num, + k_head_num, + head_dim, + block_size, + block_num, + qk_dtype, + ] + if uses_weight_dtype: + model_args.append(weights.dtype) + model_args.extend( + [ + query_scale.dtype, + torch.int32, + cu_q_cpu, + cu_k_cpu, + q_lengths, + k_lengths, + residual_for_cpu, + max_seqlen_q, + quant_mode, + layout_q, + layout_k, + sparse_count, + sparse_mode, + None, + None, + None, + None, + None, + cmp_ratio, + return_value, + ] + ) + model = pytest_golden.GeneralizedQLIV2(*model_args) + + key_for_cpu = key + key_scale_for_cpu = key_scale + if layout_k == "PA_BBND": + if block_table is None or not k_lengths: + raise ValueError( + "PA_BBND compare context requires block_table and seqused_k" + ) + sequence_length = max(k_lengths) + key_for_cpu = self.restore_paged_tensor( + key, block_table, batch_size, sequence_length + ) + if quant_mode != 4: + key_scale_for_cpu = self.restore_paged_tensor( + key_scale, block_table, batch_size, sequence_length + ) + + query_scale_for_cpu = query_scale + if quant_mode == 4: + query_scale_for_cpu = torch.full( + tuple(query.shape[:-1]), + query_scale.reshape(-1)[0].item(), + dtype=query_scale.dtype, + ) + if layout_k == "PA_BBND": + key_scale_shape = (batch_size, k_head_num, max(k_lengths)) + else: + key_scale_shape = tuple(key.shape[:-1]) + key_scale_for_cpu = torch.full( + key_scale_shape, + key_scale.reshape(-1)[0].item(), + dtype=key_scale.dtype, + ) + + _, scores, _ = model.forward( + query, + key_for_cpu, + weights, + query_scale_for_cpu, + key_scale_for_cpu, + cu_q_cpu, + cu_k_cpu, + seq_q_cpu, + seq_k_cpu, + block_table, + offset, + ) + return { + "params": params, + "scores": scores, + "topk_value": scores, + "output_idx_offset": None if offset is None else offset.detach().cpu(), + "score_layout": layout_q, + "cu_seqlens_q": cu_q_cpu, + } + + def customize( + self, + query, + key, + weights, + query_dequant_scale, + key_dequant_scale, + cu_seqlens_q, + cu_seqlens_k, + seqused_q, + seqused_k, + cmp_residual_k, + block_table, + output_idx_offset, + layout_query, + layout_key, + kwargs, + ): + pytest_golden = self.load_pytest_golden() + quant_mode = int(kwargs.get("quant_mode") or 1) + params = self.build_case_params( + query, + key, + weights, + query_dequant_scale, + layout_query, + layout_key, + kwargs, + include_weight_dtype=self.pytest_uses_weight_dtype(pytest_golden), + ) + data = pytest_golden.generate_qliv2_test_data(params, generate_golden=False) + for name, dst, src_name in ( + ("query", query, "query"), + ("key", key, "key"), + ("weights", weights, "weights"), + ("query_dequant_scale", query_dequant_scale, "query_dequant_scale"), + ("key_dequant_scale", key_dequant_scale, "key_dequant_scale"), + ("cu_seqlens_q", cu_seqlens_q, "cu_seqlens_query"), + ("cu_seqlens_k", cu_seqlens_k, "cu_seqlens_key"), + ("seqused_q", seqused_q, "seqused_q"), + ("seqused_k", seqused_k, "seqused_k"), + ("cmp_residual_k", cmp_residual_k, "cmp_residual_k_for_npu"), + ("block_table", block_table, "block_table"), + ("output_idx_offset", output_idx_offset, "output_idx_offset"), + ): + packed_dtype = ( + torch.float8_e4m3fn + if quant_mode == QUANT_MODE_MXFP8 and name in ("query", "key") + else None + ) + self.copy_tensor(dst, data.get(src_name), name, packed_dtype) + return data + + +INPUT_ADAPTER = QuantLightningIndexerV2InputAdapter() + + +def rebuild_qli_v2_compare_data(compare_context): + return INPUT_ADAPTER.rebuild_compare_data(compare_context) + + +def zero_metadata(metadata): + if torch.is_tensor(metadata): + metadata.zero_() + else: + metadata[...] = 0 + + +def generate_qli_v2_inputs( + query, + key, + weights, + query_dequant_scale, + key_dequant_scale, + topk, + quant_mode, + *, + cu_seqlens_q=None, + cu_seqlens_k=None, + seqused_q=None, + seqused_k=None, + cmp_residual_k=None, + block_table=None, + output_idx_offset=None, + metadata=None, + max_seqlen_q=-1, + layout_q="BSND", + layout_k="BSND", + mask_mode=0, + cmp_ratio=1, + return_value=0, + **kwargs, +): + """Populate pytest-derived inputs; metadata is filled by npu_preprocess.""" + if metadata is None: + raise ValueError("QLI_V2 direct API CSV must reserve the metadata tensor slot") + params = dict(kwargs) + params.update( + { + "sparse_count": topk, + "quant_mode": quant_mode, + "max_seqlen_q": max_seqlen_q, + "layout_query": layout_q, + "layout_key": layout_k, + "sparse_mode": mask_mode, + "cmp_ratio": cmp_ratio, + "return_value": return_value, + } + ) + data = INPUT_ADAPTER.customize( + query, + key, + weights, + query_dequant_scale, + key_dequant_scale, + cu_seqlens_q, + cu_seqlens_k, + seqused_q, + seqused_k, + cmp_residual_k, + block_table, + output_idx_offset, + layout_q, + layout_k, + params, + ) + zero_metadata(metadata) + case_data = INPUT_ADAPTER.load_golden_store().CASE_DATA + testcase_name = params.get("testcase_name") + case_data.put(testcase_name, data) + INPUT_ADAPTER.load_metadata_protocol().save_metadata_inputs( + "quant_lightning_indexer_v2", testcase_name, data.get("metadata_input") + ) + return data + + +def generate_aclnn_qli_v2_inputs( + query, + key, + weights, + query_dequant_scale, + key_dequant_scale, + cu_seqlens_q, + cu_seqlens_k, + seqused_q, + seqused_k, + cmp_residual_k, + block_table, + output_idx_offset, + metadata, + topk, + quant_mode, + max_seqlen_q, + layout_q, + layout_k, + mask_mode, + cmp_ratio, + return_value, + sparse_indices_out, + sparse_values_out, + **kwargs, +): + """Map the ACLNN C signature to the canonical pytest input adapter.""" + del sparse_indices_out, sparse_values_out + return generate_qli_v2_inputs( + query, + key, + weights, + query_dequant_scale, + key_dequant_scale, + topk, + quant_mode, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + seqused_q=seqused_q, + seqused_k=seqused_k, + cmp_residual_k=cmp_residual_k, + block_table=block_table, + output_idx_offset=output_idx_offset, + metadata=metadata, + max_seqlen_q=max_seqlen_q, + layout_q=layout_q, + layout_k=layout_k, + mask_mode=mask_mode, + cmp_ratio=cmp_ratio, + return_value=return_value, + **kwargs, + ) diff --git a/xllm_ops/attention/quant_lightning_indexer_v2/tests/assets/impl/metadata_inputs.py b/xllm_ops/attention/quant_lightning_indexer_v2/tests/assets/impl/metadata_inputs.py new file mode 100644 index 0000000..ca4167f --- /dev/null +++ b/xllm_ops/attention/quant_lightning_indexer_v2/tests/assets/impl/metadata_inputs.py @@ -0,0 +1,100 @@ +#!/usr/bin/python +# -*- coding: utf-8 -*- +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- +"""Input customization for standalone QuantLightningIndexerMetadata cases.""" + +import importlib.util +import sys +from pathlib import Path + +import numpy as np + + +_VECTOR_NAMES = ( + "cu_seqlens_q", + "cu_seqlens_k", + "seqused_q", + "seqused_k", + "cmp_residual_k", +) + + +def load_metadata_protocol(): + name = "qli_v2_ttk_metadata_protocol" + if name in sys.modules: + return sys.modules[name] + path = Path(__file__).with_name("metadata_protocol.py") + spec = importlib.util.spec_from_file_location(name, path) + if spec is None or spec.loader is None: + raise ImportError(f"cannot create import spec for {path}") + module = importlib.util.module_from_spec(spec) + sys.modules[name] = module + try: + spec.loader.exec_module(module) + except Exception: + sys.modules.pop(name, None) + raise + return module + + +def load_sidecar_values(kwargs): + metadata_input = load_metadata_protocol().load_metadata_inputs( + "quant_lightning_indexer_v2", kwargs.get("testcase_name") + ) + return {} if metadata_input is None else metadata_input + + +def copy_values(target, values, name): + if target is None: + return + if values is None: + raise ValueError(f"QuantLightningIndexerMetadata requires {name}_values") + source = np.asarray(values, dtype=np.int32) + if tuple(source.shape) != tuple(target.shape): + raise ValueError( + f"QuantLightningIndexerMetadata {name} shape mismatch: " + f"CSV={tuple(target.shape)}, values={tuple(source.shape)}" + ) + if hasattr(target, "copy_"): + import torch + + target.copy_(torch.as_tensor(source, dtype=target.dtype, device=target.device)) + else: + np.copyto(target, source.astype(target.dtype, copy=False)) + + +def generate_quant_lightning_indexer_metadata_inputs( + num_heads_q, + num_heads_k, + head_dim, + topk, + quant_mode, + *, + cu_seqlens_q=None, + cu_seqlens_k=None, + seqused_q=None, + seqused_k=None, + cmp_residual_k=None, + **kwargs, +): + """Copy explicit descriptor vectors into metadata API input tensors.""" + del num_heads_q, num_heads_k, head_dim, topk, quant_mode + sidecar = load_sidecar_values(kwargs) + tensors = ( + cu_seqlens_q, + cu_seqlens_k, + seqused_q, + seqused_k, + cmp_residual_k, + ) + for name, tensor in zip(_VECTOR_NAMES, tensors): + values = sidecar[name] if name in sidecar else kwargs.get(f"{name}_values") + copy_values(tensor, values, name) diff --git a/xllm_ops/attention/quant_lightning_indexer_v2/tests/assets/impl/metadata_protocol.py b/xllm_ops/attention/quant_lightning_indexer_v2/tests/assets/impl/metadata_protocol.py new file mode 100644 index 0000000..59bef59 --- /dev/null +++ b/xllm_ops/attention/quant_lightning_indexer_v2/tests/assets/impl/metadata_protocol.py @@ -0,0 +1,225 @@ +#!/usr/bin/python +# -*- coding: utf-8 -*- +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- +"""Private five-operator bridge for TTK manual-data prepare and replay. + +The protocol is intentionally environment- and file-based. It does not import +TTK or depend on a runner/context object. TTK continues to own its case files; +the assets only keep a compact sidecar outside each case directory and replace +an existing metadata input atomically after a zero-placeholder replay. +""" + +import hashlib +import os +import re +import tempfile +from pathlib import Path + +import numpy as np +import torch + + +ENV_NAME = "MANUAL_DATA_DIRS" +PROTOCOL_VERSION = 1 +SIDECAR_DIRECTORY = ".ttk_asset_metadata" +_SAFE_CASE_NAME = re.compile(r"[^A-Za-z0-9_.-]+") +_SAFE_OPERATOR = re.compile(r"[A-Za-z0-9_.-]+") +_FORMATS = ("bin", "npy", "pt") + + +def manual_data_roots(): + """Resolve the optional private roots from one shell environment value.""" + value = os.getenv(ENV_NAME) + if not value: + return () + roots = [] + seen = set() + for item in value.split(os.pathsep): + if not item.strip(): + continue + root = Path(item).expanduser().resolve() + if root not in seen: + roots.append(root) + seen.add(root) + return tuple(roots) + + +def case_directory_name(testcase_name): + """Mirror TTK's stable case-directory rule without importing TTK.""" + name = str(testcase_name) + safe = _SAFE_CASE_NAME.sub("_", name).strip("._") or "case" + if safe == name and len(safe) <= 120: + return safe + digest = hashlib.sha256(name.encode("utf-8")).hexdigest()[:12] + return f"{safe[:96]}-{digest}" + + +def validate_operator(operator): + if _SAFE_OPERATOR.fullmatch(str(operator)) is None: + raise ValueError(f"invalid metadata protocol operator: {operator!r}") + + +def clone_to_cpu(value): + if torch.is_tensor(value): + return value.detach().cpu().clone() + if isinstance(value, np.ndarray): + return torch.from_numpy(np.array(value, copy=True)) + if isinstance(value, dict): + return {name: clone_to_cpu(item) for name, item in value.items()} + if isinstance(value, list): + return [clone_to_cpu(item) for item in value] + if isinstance(value, tuple): + return tuple(clone_to_cpu(item) for item in value) + if value is None or isinstance(value, (str, bool, int, float)): + return value + if hasattr(value, "item"): + return value.item() + raise TypeError(f"unsupported metadata sidecar value: {type(value).__name__}") + + +def build_sidecar_path(root, operator, testcase_name): + validate_operator(operator) + return ( + root + / SIDECAR_DIRECTORY + / str(operator) + / f"{case_directory_name(testcase_name)}.pt" + ) + + +def save_metadata_inputs(operator, testcase_name, metadata_input): + """Atomically save exact CPU metadata arguments during input preparation.""" + roots = manual_data_roots() + if not roots or testcase_name is None: + return None + if not isinstance(metadata_input, dict): + raise ValueError(f"{operator} pytest data lacks metadata_input") + + path = build_sidecar_path(roots[0], operator, testcase_name) + if path.parent.is_symlink(): + raise ValueError( + f"metadata sidecar directory must not be a symlink: {path.parent}" + ) + path.parent.mkdir(parents=True, exist_ok=True) + payload = { + "version": PROTOCOL_VERSION, + "operator": str(operator), + "testcase_name": str(testcase_name), + "metadata_input": clone_to_cpu(metadata_input), + } + temporary = path.with_name(f".{path.name}.{os.getpid()}.tmp") + try: + torch.save(payload, temporary) + os.replace(temporary, path) + finally: + temporary.unlink(missing_ok=True) + return path + + +def load_pt(path): + try: + return torch.load(path, map_location="cpu", weights_only=True) + except TypeError: + return torch.load(path, map_location="cpu") + + +def load_metadata_inputs(operator, testcase_name): + """Load the first matching sidecar, or return None for fallback derivation.""" + if testcase_name is None: + return None + for root in manual_data_roots(): + path = build_sidecar_path(root, operator, testcase_name) + if not path.exists(): + continue + if path.is_symlink() or not path.is_file(): + raise ValueError(f"metadata sidecar must be a regular file: {path}") + payload = load_pt(path) + if ( + not isinstance(payload, dict) + or payload.get("version") != PROTOCOL_VERSION + or payload.get("operator") != str(operator) + or payload.get("testcase_name") != str(testcase_name) + or not isinstance(payload.get("metadata_input"), dict) + ): + raise ValueError(f"incompatible metadata sidecar: {path}") + return payload["metadata_input"] + return None + + +def metadata_is_materialized(metadata): + """Treat a nonzero metadata slot as authoritative in every execution mode.""" + if metadata is None: + return False + if torch.is_tensor(metadata): + return bool(torch.count_nonzero(metadata).item()) + return bool(np.count_nonzero(np.asarray(metadata))) + + +def metadata_to_array(metadata): + if torch.is_tensor(metadata): + return metadata.detach().cpu().contiguous().numpy() + return np.ascontiguousarray(np.asarray(metadata)) + + +def find_metadata_file(root, testcase_name, metadata_index): + case_dir = root / case_directory_name(testcase_name) + if not case_dir.exists(): + return None + if case_dir.is_symlink() or not case_dir.is_dir(): + raise ValueError( + f"manual-data testcase path must be a regular directory: {case_dir}" + ) + for file_format in _FORMATS: + matches = tuple(case_dir.glob(f"input_{int(metadata_index)}_*.{file_format}")) + if len(matches) > 1: + raise RuntimeError( + f"expected one metadata input[{metadata_index}], found {len(matches)} in {case_dir}" + ) + if matches: + path = matches[0] + if path.is_symlink() or not path.is_file(): + raise ValueError(f"metadata input must be a regular file: {path}") + return path + return None + + +def write_array(path, array): + with tempfile.NamedTemporaryFile( + prefix=f".{path.name}.", suffix=".tmp", dir=path.parent, delete=False + ) as stream: + temporary = Path(stream.name) + try: + if path.suffix == ".bin": + array.tofile(temporary) + elif path.suffix == ".npy": + with temporary.open("wb") as stream: + np.save(stream, array, allow_pickle=False) + elif path.suffix == ".pt": + torch.save(torch.from_numpy(np.array(array, copy=True)), temporary) + else: + raise ValueError(f"unsupported metadata input format: {path.suffix}") + os.replace(temporary, path) + finally: + temporary.unlink(missing_ok=True) + + +def rewrite_metadata_input(operator, testcase_name, metadata_index, metadata): + """Replace the existing zero placeholder after successful metadata execution.""" + validate_operator(operator) + if testcase_name is None: + return None + for root in manual_data_roots(): + path = find_metadata_file(root, testcase_name, metadata_index) + if path is None: + continue + write_array(path, metadata_to_array(metadata)) + return path + return None diff --git a/xllm_ops/attention/quant_lightning_indexer_v2/tests/assets/impl/npu_preprocess.py b/xllm_ops/attention/quant_lightning_indexer_v2/tests/assets/impl/npu_preprocess.py new file mode 100644 index 0000000..508f532 --- /dev/null +++ b/xllm_ops/attention/quant_lightning_indexer_v2/tests/assets/impl/npu_preprocess.py @@ -0,0 +1,276 @@ +#!/usr/bin/python +# -*- coding: utf-8 -*- +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- +"""Populate QuantLightningIndexer V2 metadata without TTK state coupling.""" + +import importlib.util +import logging +import sys +from pathlib import Path + +import torch + + +OPERATOR = "quant_lightning_indexer_v2" +METADATA_INDEX = 12 +QUANT_MODE_MXFP8 = 3 +QUANT_MODE_MXFP4 = 5 + + +def load_metadata_protocol(): + name = "qli_v2_ttk_metadata_protocol" + if name in sys.modules: + return sys.modules[name] + path = Path(__file__).with_name("metadata_protocol.py") + spec = importlib.util.spec_from_file_location(name, path) + if spec is None or spec.loader is None: + raise ImportError(f"cannot create import spec for {path}") + module = importlib.util.module_from_spec(spec) + sys.modules[name] = module + try: + spec.loader.exec_module(module) + except Exception: + sys.modules.pop(name, None) + raise + return module + + +def get_attribute(kwargs, name, default=None, aliases=()): + for key in (name, f"pytest_{name}", *aliases): + value = kwargs.get(key) + if value is not None: + return value + return default + + +def get_values(kwargs, name, tensor): + value = kwargs.get(f"{name}_values") + if value is None: + value = tensor + if value is None: + return None + if torch.is_tensor(value): + return value.detach().cpu().reshape(-1).tolist() + return [int(item) for item in value] + + +def max_sequence(prefix, used, fallback): + if used: + return max(int(value) for value in used) + if prefix and len(prefix) > 1: + return max( + int(prefix[index + 1]) - int(prefix[index]) + for index in range(len(prefix) - 1) + ) + return int(fallback) + + +def restore_mx_dtypes(query, key, query_scale, key_scale, quant_mode): + """Support older uint8-storage CSVs; native TTK dtypes pass through.""" + quant_mode = int(quant_mode) + if quant_mode == QUANT_MODE_MXFP8: + qk_dtype = getattr(torch, "float8_e4m3fn", None) + elif quant_mode == QUANT_MODE_MXFP4: + qk_dtype = getattr(torch, "float4_e2m1fn_x2", None) + else: + return + scale_dtype = getattr(torch, "float8_e8m0fnu", None) + if qk_dtype is None or scale_dtype is None: + raise RuntimeError("current PyTorch does not provide the requested MX dtype") + for name, tensor, dtype in ( + ("query", query, qk_dtype), + ("key", key, qk_dtype), + ("query_dequant_scale", query_scale, scale_dtype), + ("key_dequant_scale", key_scale, scale_dtype), + ): + if tensor.dtype == dtype: + continue + if tensor.dtype != torch.uint8: + raise TypeError( + f"QLI_V2 {name} must use {dtype} or uint8 storage, got {tensor.dtype}" + ) + tensor.data = tensor.data.view(dtype) + + +def build_metadata_arguments( + query, key, topk, quant_mode, layout_q, layout_k, mask_mode, cmp_ratio, kwargs +): + q_shape = tuple(int(value) for value in query.shape) + k_shape = tuple(int(value) for value in key.shape) + num_heads_q = q_shape[2] if layout_q == "BSND" else q_shape[1] + num_heads_k = k_shape[1] if layout_k == "TND" else k_shape[2] + head_dim = q_shape[-1] * (2 if int(quant_mode) == QUANT_MODE_MXFP4 else 1) + cu_q = kwargs.get("cu_seqlens_q") + cu_k = kwargs.get("cu_seqlens_k") + seq_q = kwargs.get("seqused_q") + seq_k = kwargs.get("seqused_k") + cu_q_values = get_values(kwargs, "cu_seqlens_q", cu_q) + cu_k_values = get_values(kwargs, "cu_seqlens_k", cu_k) + seq_q_values = get_values(kwargs, "seqused_q", seq_q) + seq_k_values = get_values(kwargs, "seqused_k", seq_k) + + batch_size = get_attribute(kwargs, "batch_size") + if batch_size is None: + if seq_q_values is not None: + batch_size = len(seq_q_values) + elif cu_q_values is not None: + batch_size = len(cu_q_values) - 1 + elif layout_q == "BSND": + batch_size = q_shape[0] + else: + batch_size = 0 + + q_fallback = q_shape[1] if layout_q == "BSND" else q_shape[0] + if layout_k == "BSND": + k_fallback = k_shape[1] + elif layout_k == "TND": + k_fallback = k_shape[0] + else: + k_fallback = int(get_attribute(kwargs, "max_seqlen_k", k_shape[1])) + + return { + "num_heads_q": int( + get_attribute(kwargs, "num_heads_q", num_heads_q, ("pytest_q_head_num",)) + ), + "num_heads_k": int( + get_attribute(kwargs, "num_heads_k", num_heads_k, ("pytest_k_head_num",)) + ), + "head_dim": int(get_attribute(kwargs, "head_dim", head_dim)), + "topk": int(topk), + "quant_mode": int(quant_mode), + "cu_seqlens_q": cu_q, + "cu_seqlens_k": cu_k, + "seqused_q": seq_q, + "seqused_k": seq_k, + "cmp_residual_k": kwargs.get("cmp_residual_k"), + "batch_size": int(batch_size), + "max_seqlen_q": int( + get_attribute( + kwargs, + "metadata_max_seqlen_q", + max_sequence(cu_q_values, seq_q_values, q_fallback), + ) + ), + "max_seqlen_k": int( + get_attribute( + kwargs, + "metadata_max_seqlen_k", + max_sequence(cu_k_values, seq_k_values, k_fallback), + ) + ), + "layout_q": str(layout_q), + "layout_k": str(layout_k), + "mask_mode": int(mask_mode), + "cmp_ratio": int(cmp_ratio), + } + + +def move_to_device(value, target): + if value is None: + return None + if torch.is_tensor(value): + return value.to(device=target.device) + return torch.as_tensor(value, device=target.device) + + +def run_metadata(arguments, metadata): + return torch.ops.cann_ops_transformer.quant_lightning_indexer_metadata( + int(arguments["num_heads_q"]), + int(arguments["num_heads_k"]), + int(arguments["head_dim"]), + int(arguments["topk"]), + int(arguments["quant_mode"]), + cu_seqlens_q=move_to_device(arguments.get("cu_seqlens_q"), metadata), + cu_seqlens_k=move_to_device(arguments.get("cu_seqlens_k"), metadata), + seqused_q=move_to_device(arguments.get("seqused_q"), metadata), + seqused_k=move_to_device(arguments.get("seqused_k"), metadata), + cmp_residual_k=move_to_device(arguments.get("cmp_residual_k"), metadata), + batch_size=int(arguments["batch_size"]), + max_seqlen_q=int(arguments["max_seqlen_q"]), + max_seqlen_k=int(arguments["max_seqlen_k"]), + layout_q=str(arguments["layout_q"]), + layout_k=str(arguments["layout_k"]), + mask_mode=int(arguments["mask_mode"]), + cmp_ratio=int(arguments["cmp_ratio"]), + ) + + +def run( + query, + key, + weights, + query_dequant_scale, + key_dequant_scale, + topk, + quant_mode, + *, + cu_seqlens_q=None, + cu_seqlens_k=None, + seqused_q=None, + seqused_k=None, + cmp_residual_k=None, + metadata=None, + layout_q="BSND", + layout_k="BSND", + mask_mode=0, + cmp_ratio=1, + **kwargs, +): + """Generate metadata once, or reuse a nonzero manual-data input.""" + del weights + if metadata is None: + raise ValueError("QuantLightningIndexer V2 npu_preprocess requires metadata") + restore_mx_dtypes(query, key, query_dequant_scale, key_dequant_scale, quant_mode) + arguments_kwargs = dict(kwargs) + arguments_kwargs.update( + { + "cu_seqlens_q": cu_seqlens_q, + "cu_seqlens_k": cu_seqlens_k, + "seqused_q": seqused_q, + "seqused_k": seqused_k, + "cmp_residual_k": cmp_residual_k, + } + ) + protocol = load_metadata_protocol() + testcase_name = kwargs.get("testcase_name") + if protocol.metadata_is_materialized(metadata): + logging.info("[%s] reuse nonzero QLI_V2 metadata input", testcase_name) + return None + arguments = protocol.load_metadata_inputs(OPERATOR, testcase_name) + if arguments is not None: + source = "manual-data sidecar" + else: + arguments = build_metadata_arguments( + query, + key, + topk, + quant_mode, + layout_q, + layout_k, + mask_mode, + cmp_ratio, + arguments_kwargs, + ) + source = "main API fallback (sidecar unavailable)" + logging.info("[%s] build QLI_V2 metadata from %s", testcase_name, source) + generated = run_metadata(arguments, metadata) + if tuple(metadata.shape) != tuple(generated.shape): + raise ValueError( + "QLI_V2 metadata shape mismatch: " + f"placeholder={tuple(metadata.shape)}, generated={tuple(generated.shape)}" + ) + metadata.copy_(generated.to(dtype=metadata.dtype, device=metadata.device)) + rewritten = protocol.rewrite_metadata_input( + OPERATOR, testcase_name, METADATA_INDEX, metadata + ) + if rewritten is not None: + logging.info("[%s] rewrote QLI_V2 metadata input: %s", testcase_name, rewritten) + return None diff --git a/xllm_ops/attention/quant_lightning_indexer_v2/tests/assets/spec.py b/xllm_ops/attention/quant_lightning_indexer_v2/tests/assets/spec.py new file mode 100644 index 0000000..9892247 --- /dev/null +++ b/xllm_ops/attention/quant_lightning_indexer_v2/tests/assets/spec.py @@ -0,0 +1,98 @@ +#!/usr/bin/python +# -*- coding: utf-8 -*- +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +"""TestSpec adapter for QuantLightningIndexer V2 TTK assets.""" + +import importlib.util +import sys +from pathlib import Path + +ASSET_IMPL_DIR = Path(__file__).with_name("impl") + + +def load_impl_module(stem): + name = f"qli_v2_ttk_{stem}" + if name in sys.modules: + return sys.modules[name] + path = ASSET_IMPL_DIR / f"{stem}.py" + try: + spec = importlib.util.spec_from_file_location(name, path) + if spec is None or spec.loader is None: + raise ImportError(f"cannot create import spec for {path}") + module = importlib.util.module_from_spec(spec) + sys.modules[name] = module + spec.loader.exec_module(module) + except Exception as exc: + sys.modules.pop(name, None) + raise RuntimeError( + "Failed to load QuantLightningIndexerV2 assets module; " + f"stage=impl/{stem}; module={path.resolve()}; " + f"original error: {type(exc).__name__}: {exc}" + ) from exc + return module + + +npu_preprocess_module = load_impl_module("npu_preprocess") +golden_module = load_impl_module("golden") +inputs_module = load_impl_module("inputs") +metadata_inputs_module = load_impl_module("metadata_inputs") +compare_module = load_impl_module("compare") + + +class QuantLightningIndexerV2Spec: + golden = golden_module.cpu_quant_lightning_indexer_v2 + customize_inputs = inputs_module.generate_qli_v2_inputs + npu_preprocess = npu_preprocess_module.run + tolerance = { + "float16": {"standard": "stat_rel_err"}, + "bfloat16": {"standard": "stat_rel_err"}, + "float8_e4m3fn": {"standard": "stat_rel_err"}, + } + + def compare(*outputs, compare_context=None, **kwargs): + del kwargs + testcase_name = ( + None if compare_context is None else compare_context.testcase_name + ) + data = golden_module.get_compare_data(testcase_name) + if data is None: + if compare_context is None: + raise RuntimeError( + "QuantLightningIndexerV2 pytest compare requires compare_context" + ) + data = inputs_module.rebuild_qli_v2_compare_data(compare_context) + golden_module.set_compare_data(compare_context.testcase_name, data) + try: + return compare_module.compare(*outputs, compare_data=data) + finally: + golden_module.discard_compare_data(data) + + +class AclnnQuantLightningIndexerV2Spec(QuantLightningIndexerV2Spec): + golden = golden_module.cpu_aclnn_qli_v2 + customize_inputs = inputs_module.generate_aclnn_qli_v2_inputs + npu_preprocess = None + + +class QuantLightningIndexerMetadataSpec: + customize_inputs = ( + metadata_inputs_module.generate_quant_lightning_indexer_metadata_inputs + ) + + +__spec__ = { + "torch.ops.cann_ops_transformer.quant_lightning_indexer": "QuantLightningIndexerV2Spec", + "torch.ops.cann_ops_transformer.quant_lightning_indexer_metadata": ( + "QuantLightningIndexerMetadataSpec" + ), + "aclnnQuantLightningIndexerV2": "AclnnQuantLightningIndexerV2Spec", +} diff --git a/xllm_ops/attention/quant_lightning_indexer_v2/tests/pytest/README.md b/xllm_ops/attention/quant_lightning_indexer_v2/tests/pytest/README.md new file mode 100644 index 0000000..107a204 --- /dev/null +++ b/xllm_ops/attention/quant_lightning_indexer_v2/tests/pytest/README.md @@ -0,0 +1,259 @@ +# quant_lightning_indexer_v2算子测试框架 + +## 功能说明 + +基于pytest测试框架，实现quant_lightning_indexer_v2算子的功能验证： + +- **CPU侧**：复现算子功能用以生成golden数据 +- **NPU侧**：通过TorchNPU进行算子直调获取实际数据 +- **精度对比**：进行CPU与NPU结果的精度对比验证算子功能 +- **双模式执行隔离**：支持直接pytest多进程执行和shell层进程隔离两种批量模式 +- **PT 复跑模式**：`batch -P <目录>` 直接执行已有 PT；增加 `-E` 才会先生成 PT +- **可控批跑**：`-P` 统一指定 PT 生成和读取目录，支持结果路径、case 名和 1-based 序号选择 +- **性能采集**：支持挂载msprof采集算子性能数据并汇总输出 +- **运行模式切换**：支持eager直接调用和graph（torch.compile + torchair）两种算子调用模式 + +## 当前实现范围 + +### 参数限制 + +- **数据格式**: + - **query_layout**：BSND、TND + - **key_layout**: PA_BBND、BSND、TND + +- **数据类型**: + - **qk_dtype**: FLOAT8_E4M3FN、INT8、HIFLOAT8、FLOAT4_E2M1 + - **dequant_dtype**: FP32（Ascend950）、FP16（Ascend910_93）、FLOAT8_E8M0（MXFP8/MXFP4） + - **actual_seq_dtype**: INT32 + +### PyTorch MX类型约定 + +- `quant_mode`为3/5时，`query_dequant_scale`和`key_dequant_scale`必须为`torch.float8_e8m0fnu`。 +- `quant_mode`为5时，`query`和`key`的算子逻辑数据类型为`FLOAT4_E2M1`（文档简写为`float4_e2m1`）： + - PyTorch提供`torch.float4_e2m1fn_x2`时，优先使用该原生打包类型；名称中的`x2`表示每个物理字节打包两个E2M1逻辑元素。 + - PyTorch未提供该类型时，使用`torch.uint8`承载已打包数据，封装层会将其按`ACL_FLOAT4_E2M1`传入算子。 + +- **运行模式**: + - **eager**：直接调用 `torch.ops.cann_ops_transformer.quant_lightning_indexer` + - **graph**：通过 `torch.compile` + `torchair` 后端编译执行（需torchair支持） + +### 环境配置 + +#### 前置要求 + +1、 确认TorchNPU为最新版本 +2、 激活CANN包和自定义算子包 +3、 graph模式需要安装torchair编译器后端 + +#### custom包调用 + +支持custom包调用 + +## 文件结构 + +#### pytest文件结构说明 + +- test_run.sh # 执行脚本，支持single/batch/batch_exec三种命令 +- batch_isolated_run.sh # 批量隔离执行脚本（shell层进程隔离+msprof性能采集） +- quant_lightning_indexer_v2_golden.py # cpu侧算子golden实现 +- quant_lightning_indexer_v2_acl_graph.py # graph模式torchair后端实现 +- result_compare_method.py # cpu golden与npu输出精度对比 +- qliv2_test_utils.py # case选择、稳定命名和结果表公共逻辑 +- collect_perf_data.py # msprof性能数据收集与汇总 +- pytest.ini # 创建测试标记 + +单用例测试： + +- test_quant_lightning_indexer_v2_single.py # pytest测试单用例运行主程序 +- test_quant_lightning_indexer_v2_paramset.py # 单用例入参配置，按芯片型号自动选择用例 + +批量测试： + +- test_quant_lightning_indexer_v2_batch.py # 用例批量测试主程序并生成excel文件保存结果 +- ./batch/quant_lightning_indexer_v2_pt_loadprocess.py # 读取pt文件并调用算子获取npu输出 +- ./batch/quant_lightning_indexer_v2_pt_save.py # 读取excel表格批量生成用例pt文件 +- ./batch/list_pt_from_excel.py # 从Excel提取Testcase_Name并按名匹配pt文件（batch_exec模式用） + +## 架构说明 + +- **数据生成入口**：`generate_qliv2_test_data` 复用原有 batch 数据链生成输入和 CPU golden，不调用 metadata 或主算子；参数准备仍可查询设备信息 +- **single 模式**：直接执行配置用例；指定 `--save-pt` 时保存并执行同一份实际输入，避免二次随机生成 +- **batch 模式**：`-P` 是唯一 PT 目录；有 `-E` 时先生成再执行，无 `-E` 时直接执行已有 PT +- **batch_exec 模式**：按 Excel 的 `Testcase_Name` 筛选已有 PT，仅执行 NPU 和精度对比 +- 两路共用 `_qliv2_prepare_tensors_and_metadata` 和 `_qliv2_run_compiled_graph`，统一使用 `fullgraph=False` +- 结果表会先落盘；index 或 return value 精度结果为 `Failed` 时，pytest 随后以非零状态退出 + +## 使用方法 + +在pytest文件夹路径下执行： + +### 运行测试用例 + +#### 单用例调测 + +1、手动配置test_quant_lightning_indexer_v2_paramset.py的ENABLED_PARAMS参数 + +2、执行指令： + +``` bash +bash test_run.sh single +bash test_run.sh single --save-pt ./single_pt -O ./result/single.xlsx +bash test_run.sh single -M graph --save-pt ./single_pt +``` + +#### 用例的批量生成与测试 + +##### 方式A：test_run.sh 批量执行 + +1、excel路径下存放用例excel表格 + +`-P` 同时指定 PT 的保存目录和读取目录，默认是当前 pytest 目录下的 `pt_path`。 + +##### 直接执行已有 PT + +不传 `-E` 时不读取 Excel、不重新生成 PT，只执行 NPU 和 compare： + +3、执行指令： + +``` bash +bash test_run.sh batch -P ./pt_path +bash test_run.sh batch -P ./pt_path -O ./result/rerun.xlsx +bash test_run.sh batch -P ./pt_path -C case_b,case_a # 按名称和给定顺序 +bash test_run.sh batch -P ./pt_path -I 3,1,5-7 # 按自然排序后的序号 +bash test_run.sh batch -P ./pt_path -M graph +``` + +4、配置区默认值： + +| 变量 | 默认值 | 命令行参数 | 说明 | +|---|---|---|---| +| DEFAULT_EXCEL | `./excel/test_cases.xlsx` | `-E` | Excel 用例表格路径（**必须指定具体文件名**，不支持通配符如 `./excel/*`） | +| DEFAULT_PT_PATH | `./pt_path` | `-P` | pt 文件存放目录 | +| （无） | `Sheet1` | `-S` | Excel Sheet 页名 | +| （无） | `eager` | `-M` | 运行模式（eager/graph） | + +#### 根据 Excel 表格筛选已有 pt 批量执行（batch_exec 模式） +> +> 仅重新执行 NPU 测试和精度对比，不重新生成 pt 文件。适用于已有 pt 文件、只需更新精度结果的场景。 + +增加 `-E` 后，脚本先把 Excel 用例生成到 `-P`，再从同一个目录执行： + +2、执行指令： + +``` bash +bash test_run.sh batch -E ./excel/test_cases.xlsx -P ./pt_path +bash test_run.sh batch -E ./excel/test_cases.xlsx -S Sheet1 -P ./pt_path +bash test_run.sh batch -E ./excel/test_cases.xlsx -P ./pt_path -O ./result/batch.xlsx +``` + +3、执行流程： + +- 从 Excel 表格读取 `Testcase_Name` 列 +- 按 `.pt` 在 pt_path 下匹配对应的 .pt 文件 +- 仅对匹配到的 .pt 文件执行 NPU 测试和精度对比 +- 生成 `result.xlsx` 测试结果表格 +- 如果 Excel 中某条用例无对应的 .pt 文件，会输出警告并跳过该用例 + +4、与 `batch` 模式的区别： + +| | `batch` | `batch_exec` | +|---|---|---| +| pt 生成 | 每次重新生成 | 跳过 | +| 执行速度 | 较慢（含 pt 生成） | 较快 | +| 适用场景 | 首次运行 / 参数变更 | 精度复测 / 仅 NPU 结果更新 | + +##### 方式B：手工分步执行 + +1、生成pt文件： + +``` bash +python3 batch/quant_lightning_indexer_v2_pt_save.py excel/test_cases.xlsx pt_path +python3 batch/quant_lightning_indexer_v2_pt_save.py excel/test_cases.xlsx pt_path --sheet Sheet1 # 指定 Sheet 页 +``` + +2、替换测试脚本路径： + +``` bash +QLIV2_TESTCASE_DIR=pt_path QLIV2_RESULT_PATH=result.xlsx \ +python3 -m pytest -rA -s test_quant_lightning_indexer_v2_batch.py -v -m ci +``` + +3、执行测试： + +``` bash +python3 -m pytest -rA -s test_quant_lightning_indexer_v2_batch.py -v -m ci -W ignore::UserWarning -W ignore::DeprecationWarning +``` + +4、恢复测试脚本： + +``` bash +cp test_quant_lightning_indexer_v2_batch.py.bak test_quant_lightning_indexer_v2_batch.py +``` + +##### 方式C：批量隔离执行（推荐用于性能采集） + +对每条用例单独拉起一个pytest进程，实现进程间完全隔离，避免单条用例崩溃影响其他用例。 + +``` bash +bash batch_isolated_run.sh ./pt_path 0 # 不采集性能 +bash batch_isolated_run.sh ./pt_path 1 # 采集性能（挂载msprof） +bash batch_isolated_run.sh ./pt_path 0 graph # graph模式 + 不采集性能 +bash batch_isolated_run.sh ./pt_path 1 graph # graph模式 + 性能采集 +``` + +## Excel 用例表格式 + +`excel/test_cases.xlsx` 需包含以下列（Sheet1）： + +| 列名 | 类型 | 示例 | +|---|---|---| +| Testcase_Name | str | `test_case_01` | +| batch_size | int | `8` | +| q_seq | int | `15` | +| k_seq | int | `111` | +| q_t_size | int | `8` | +| k_t_size | int | `15` | +| q_head_num | int | `64` | +| k_head_num | int | `1` | +| head_dim | int | `128` | +| block_size | int | `512` | +| block_num | int | `8` | +| qk_dtype | str | `FLOAT8_E4M3FN` / `INT8` / `HIFLOAT8` / `FLOAT4_E2M1` | +| dequant_dtype | str | `FP32` / `FP16` / `FLOAT8_E8M0` | +| actual_seq_dtype | str | `INT32` | +| cu_seqlens_q | None/str | `None` 或 `"[0, 1]"` | +| cu_seqlens_k | None/str | `None` 或 `"[0, 1]"` | +| seqused_q | None/str | `None` 或 `"[3,3,3,3,3,3,3,3]"` | +| seqused_k | str | `"[28,24,80,96,47,76,0,111]"` | +| cmp_residual_k | None/str | `None` 或 `"[0,0,0,0,0,0,0,0]"`（cmp_ratio>1时必填）| +| max_seqlen_q | int | `-1` | +| quant_mode | int | `1` / `2` / `4` | +| layout_query | str | `BSND` / `TND` | +| layout_key | str | `PA_BBND` | +| sparse_count | int | `512` | +| sparse_mode | int | `0` / `3` | +| query_datarange | str | `"[-448,448]"` | +| key_datarange | str | `"[-20,20]"` | +| weights_datarange | str | `"[-123,123]"` | +| q_scale_datarange | str | `"[0,255]"` | +| k_scale_datarange | str | `"[0,65504]"` | +| cmp_ratio | int | `1` / `4` | +| return_value | int | `0` / `1` | +| output_idx_offset | None/str | `None` 或列表字符串 | + +**注意事项**： + +- `dequant_dtype`：Ascend950的`quant_mode=3/5`仅支持`FLOAT8_E8M0`，其他量化模式支持`FP32`；Ascend910_93支持`FP16` +- `cmp_ratio > 1`且`sparse_mode != 0`时，`cmp_residual_k`必填（长度=batch_size的列表） +- `return_value=1`时，`output_idx_offset`需提供有效值 +- Ascend910_93要求`quant_mode=2` + +## 输出文件 + +| 文件 | 说明 | +|---|---| +| `result.xlsx` | 测试结果（精度、参数等） | +| `result_perf.xlsx` | 测试结果 + 性能数据（仅msprof模式） | +| `batch_summary.log` | 批量执行详细日志 | +| `batch_fail_list.log` | 失败用例清单 | +| `PROF_*/` | msprof性能原始数据目录 | diff --git a/xllm_ops/attention/quant_lightning_indexer_v2/tests/pytest/batch/list_pt_from_excel.py b/xllm_ops/attention/quant_lightning_indexer_v2/tests/pytest/batch/list_pt_from_excel.py new file mode 100644 index 0000000..c4c6e3d --- /dev/null +++ b/xllm_ops/attention/quant_lightning_indexer_v2/tests/pytest/batch/list_pt_from_excel.py @@ -0,0 +1,65 @@ +#!/usr/bin/python +# -*- coding: utf-8 -*- +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +import os +import sys +import argparse +import pandas as pd + + +def list_pt_from_excel(excel_path, sheet_name, pt_dir): + if not os.path.exists(excel_path): + print(f"ERROR: Excel file not found: {excel_path}", file=sys.stderr) + sys.exit(1) + + if not os.path.isdir(pt_dir): + print(f"ERROR: pt directory not found: {pt_dir}", file=sys.stderr) + sys.exit(1) + + df = pd.read_excel(excel_path, sheet_name=sheet_name) + if "Testcase_Name" not in df.columns: + print(f"ERROR: Column 'Testcase_Name' not found in sheet '{sheet_name}'", file=sys.stderr) + sys.exit(1) + + pt_files = [] + missing = [] + for case_name in df["Testcase_Name"]: + pt_path = os.path.join(pt_dir, f"{case_name}.pt") + if os.path.isfile(pt_path): + pt_files.append(pt_path) + else: + missing.append(case_name) + + if missing: + print(f"WARNING: {len(missing)} cases have no matching .pt file: {missing}", file=sys.stderr) + + if not pt_files: + print(f"ERROR: No matching .pt files found for any case in Excel", file=sys.stderr) + sys.exit(1) + + print(",".join(pt_files)) + + +def main(): + parser = argparse.ArgumentParser( + description="Extract Testcase_Name from Excel and map to .pt files in pt_dir" + ) + parser.add_argument("excel_path", type=str, help="Path to Excel file") + parser.add_argument("pt_dir", type=str, help="Directory containing .pt files") + parser.add_argument("--sheet", "-S", type=str, default="Sheet1", help="Sheet name (default: Sheet1)") + args = parser.parse_args() + + list_pt_from_excel(args.excel_path, args.sheet, args.pt_dir) + + +if __name__ == "__main__": + main() diff --git a/xllm_ops/attention/quant_lightning_indexer_v2/tests/pytest/batch/quant_lightning_indexer_v2_pt_loadprocess.py b/xllm_ops/attention/quant_lightning_indexer_v2/tests/pytest/batch/quant_lightning_indexer_v2_pt_loadprocess.py new file mode 100644 index 0000000..ebffb1a --- /dev/null +++ b/xllm_ops/attention/quant_lightning_indexer_v2/tests/pytest/batch/quant_lightning_indexer_v2_pt_loadprocess.py @@ -0,0 +1,220 @@ +#!/usr/bin/python +# -*- coding: utf-8 -*- +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +import os +import pandas as pd +import numpy as np +import torch +import torch_npu +import pytest +import random +import math +import ast +import cann_ops_transformer +from qliv2_parameter_normalization import normalize_qliv2_params + +QUANT_MODE_MXFP4 = 5 + + +def test_qliv2_process(filepath, device_id=0): + # 加载测试数据 + test_data = torch.load(filepath, map_location="cpu") + + params = normalize_qliv2_params(test_data["params"]) + cpu_result = test_data["cpu_result"] + topk_value = test_data["topk_value"] + cpu_topk_value = test_data["cpu_topk_value"] + print("执行用例：", filepath) + torch_npu.npu.set_device(device_id) + + quant_mode = test_data["quant_mode"] + if quant_mode == QUANT_MODE_MXFP4: + query = test_data["query"].view(torch.uint8).npu() + key = test_data["key"].view(torch.uint8).npu() + if "blockFusion" in test_data and test_data["blockFusion"] is not None: + blockFusion = test_data["blockFusion"].view(torch.uint8).npu() + elif params[10] == "FLOAT8_E4M3FN" or params[10] == torch.float8_e4m3fn: + query = test_data["query"].to(dtype=torch.float8_e4m3fn).npu() + key = test_data["key"].to(dtype=torch.float8_e4m3fn).npu() + if "blockFusion" in test_data and test_data["blockFusion"] is not None: + blockFusion = test_data["blockFusion"] + if blockFusion.dtype == torch.uint8: + blockFusion = blockFusion.view(torch.float8_e4m3fn) + else: + blockFusion = blockFusion.to(dtype=torch.float8_e4m3fn) + blockFusion = blockFusion.npu() + else: + query = test_data["query"].npu() + key = test_data["key"].npu() + if "blockFusion" in test_data and test_data["blockFusion"] is not None: + blockFusion = test_data["blockFusion"].npu() + + max_seqlen_q = params[19] + return_value = params[31] + weights = test_data["weights"].npu() + query_dequant_scale = test_data["query_dequant_scale"].npu() + key_dequant_scale = test_data["key_dequant_scale"].npu() + if "blockFusion" in test_data and test_data["blockFusion"] is not None: + block_num = params[9] + block_size = params[8] + head_dim = params[7] + k_head_num = params[6] + k_head_num = int(k_head_num) + head_dim = int(head_dim) + block_size = int(block_size) + block_num = int(block_num) + dequant_dtype_str = params[12] + if dequant_dtype_str == "FP16" or dequant_dtype_str == torch.float16: + dequant_dtype = torch.float16 + elif dequant_dtype_str == "FP32" or dequant_dtype_str == torch.float32: + dequant_dtype = torch.float32 + else: + dequant_dtype = torch.float16 + key = blockFusion[:, : block_size * k_head_num * head_dim].view( + block_num, block_size, k_head_num, head_dim + ) + key_dequant_scale = ( + blockFusion[:, block_size * k_head_num * head_dim :] + .view(dequant_dtype) + .view(block_num, block_size, k_head_num) + ) + if test_data["seqused_q"] is not None: + seqused_q = test_data["seqused_q"].npu() + else: + seqused_q = None + if test_data["seqused_k"] is not None: + seqused_k = test_data["seqused_k"].npu() + else: + seqused_k = None + if test_data["output_idx_offset"] is not None: + output_idx_offset = test_data["output_idx_offset"].npu() + else: + output_idx_offset = None + if test_data["cu_seqlens_query"] is not None: + cu_seqlens_query = test_data["cu_seqlens_query"].npu() + else: + cu_seqlens_query = None + if test_data["cu_seqlens_key"] is not None: + cu_seqlens_key = test_data["cu_seqlens_key"].npu() + else: + cu_seqlens_key = None + if test_data["block_table"] is not None: + block_table = test_data["block_table"].npu() + else: + block_table = None + layout_query = test_data["layout_query"] + layout_key = test_data["layout_key"] + sparse_count = test_data["sparse_count"] + sparse_mode = test_data["sparse_mode"] + cmp_ratio = test_data["cmp_ratio"] + if test_data["cmp_residual_k_for_npu"] is not None: + cmp_residual_k_for_npu = test_data["cmp_residual_k_for_npu"].npu() + else: + cmp_residual_k_for_npu = None + + max_seqlen_q_meta = test_data["max_seqlen_q_meta"] + max_seqlen_k_meta = test_data["max_seqlen_k_meta"] + metadata = torch.ops.cann_ops_transformer.quant_lightning_indexer_metadata( + cu_seqlens_q=cu_seqlens_query, + cu_seqlens_k=cu_seqlens_key, + seqused_q=seqused_q, + seqused_k=seqused_k, + cmp_residual_k=cmp_residual_k_for_npu, + batch_size=params[0], + max_seqlen_q=max_seqlen_q_meta, + max_seqlen_k=max_seqlen_k_meta, + num_heads_q=params[5], + num_heads_k=params[6], + head_dim=params[7], + topk=sparse_count, + quant_mode=quant_mode, + mask_mode=sparse_mode, + layout_q=layout_query, + layout_k=layout_key, + cmp_ratio=cmp_ratio, + ) + metadata = metadata.npu() + + # 调用qli算子 + npu_result, npu_value = torch.ops.cann_ops_transformer.quant_lightning_indexer( + query, + key, + weights, + query_dequant_scale, + key_dequant_scale, + cu_seqlens_q=cu_seqlens_query, + cu_seqlens_k=cu_seqlens_key, + seqused_q=seqused_q, + seqused_k=seqused_k, + cmp_residual_k=cmp_residual_k_for_npu, + output_idx_offset=output_idx_offset, + max_seqlen_q=max_seqlen_q, + block_table=block_table, + metadata=metadata, + quant_mode=quant_mode, + layout_q=layout_query, + layout_k=layout_key, + topk=sparse_count, + mask_mode=sparse_mode, + cmp_ratio=cmp_ratio, + return_value=return_value, + ) + + torch.npu.synchronize() + npu_topk_value = npu_value + if return_value: + if npu_topk_value.shape != npu_result.shape: + raise RuntimeError( + "sparse_values and sparse_indices must have the same shape when return_value=1, " + f"but got {tuple(npu_topk_value.shape)} and {tuple(npu_result.shape)}" + ) + npu_topk_value, npu_sort_order = npu_topk_value.sort(dim=-1, descending=True) + npu_result = torch.gather(npu_result, dim=-1, index=npu_sort_order) + return ( + cpu_result, + npu_result, + topk_value, + cpu_topk_value, + npu_topk_value, + output_idx_offset, + params, + ) + + +def test_qliv2_process_graph(filepath, device_id=0): + """ + graph 模式：从 .pt 文件加载 pre-computed tensor，走 torch.compile + torchair 后端执行算子， + 跳过 generate_qliv2_test_data 的随机数据重新生成和 CPU golden 重算。 + 与 eager 模式共用相同的 .pt 数据，仅算子调用路径不同（compile vs eager）。 + """ + import quant_lightning_indexer_v2_acl_graph + + test_data = torch.load(filepath, map_location="cpu") + params = normalize_qliv2_params(test_data["params"]) + output_idx_offset = test_data.get("output_idx_offset", None) + + torch_npu.npu.set_device(device_id) + cpu_result, npu_result, topk_value, cpu_topk_value, npu_topk_value = ( + quant_lightning_indexer_v2_acl_graph.qliv2_output_acl_graph_from_pt( + params, test_data + ) + ) + + return ( + cpu_result, + npu_result, + topk_value, + cpu_topk_value, + npu_topk_value, + output_idx_offset, + params, + ) diff --git a/xllm_ops/attention/quant_lightning_indexer_v2/tests/pytest/batch/quant_lightning_indexer_v2_pt_save.py b/xllm_ops/attention/quant_lightning_indexer_v2/tests/pytest/batch/quant_lightning_indexer_v2_pt_save.py new file mode 100644 index 0000000..7a10ab0 --- /dev/null +++ b/xllm_ops/attention/quant_lightning_indexer_v2/tests/pytest/batch/quant_lightning_indexer_v2_pt_save.py @@ -0,0 +1,189 @@ +#!/usr/bin/python +# -*- coding: utf-8 -*- +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +import os +from functools import partial +from quant_lightning_indexer_v2_golden import generate_qliv2_test_data +import pandas as pd +import numpy as np +import torch + +try: + import torch_npu +except ImportError: + torch_npu = None +import pytest +import random +import math +import ast +import argparse + + +def load_excel_test_cases(excel_file_path: str, sheetname: str): + """ + 从 Excel 文件加载测试用例。 + + 参数: + excel_file_path (str): Excel 文件的路径。 + sheetname (str, optional): 工作表名称。若未提供，则默认 'Sheet1'。 + + 返回: + list[tuple]: 测试用例元组列表，每个元组包含 20+ 个字段。 + 若失败或跳过，则返回空列表。 + """ + # 优先使用传入的 sheetname，否则尝试从环境变量获取 + if sheetname is None: + sheetname = "Sheet1" + + # 检查文件是否存在 + if not os.path.exists(excel_file_path): + pytest.skip(f"Excel file not found: {excel_file_path}", allow_module_level=True) + + try: + # 读取 Excel 文件的指定 sheet + df = pd.read_excel(excel_file_path, sheet_name=sheetname) + df = df.replace({np.nan: None, pd.NA: None}) + + # 定义必需的列名 + required_columns = [ + "Testcase_Name", + "batch_size", + "q_seq", + "k_seq", + "q_t_size", + "k_t_size", + "q_head_num", + "k_head_num", + "head_dim", + "block_size", + "block_num", + "qk_dtype", + "dequant_dtype", + "actual_seq_dtype", + "cu_seqlens_q", + "cu_seqlens_k", + "seqused_q", + "seqused_k", + "cmp_residual_k", + "max_seqlen_q", + "quant_mode", + "layout_query", + "layout_key", + "sparse_count", + "sparse_mode", + "query_datarange", + "key_datarange", + "weights_datarange", + "q_scale_datarange", + "k_scale_datarange", + "cmp_ratio", + "return_value", + "output_idx_offset", + ] + + # 检查是否缺少必要列 + missing_cols = [col for col in required_columns if col not in df.columns] + if missing_cols: + pytest.skip( + f"Missing required columns in Excel: {missing_cols}", + allow_module_level=True, + ) + + # 构建测试用例列表 + test_cases = [] + for _, row in df.iterrows(): + test_cases.append( + ( + row["Testcase_Name"], + row["batch_size"], + row["q_seq"], + row["k_seq"], + row["q_t_size"], + row["k_t_size"], + row["q_head_num"], + row["k_head_num"], + row["head_dim"], + row["block_size"], + row["block_num"], + row["qk_dtype"], + row["weight_dtype"] + if "weight_dtype" in row and row["weight_dtype"] is not None + else row["dequant_dtype"], + row["dequant_dtype"], + row["actual_seq_dtype"], + row["cu_seqlens_q"], + row["cu_seqlens_k"], + row["seqused_q"], + row["seqused_k"], + row["cmp_residual_k"], + row["max_seqlen_q"], + row["quant_mode"], + row["layout_query"], + row["layout_key"], + row["sparse_count"], + row["sparse_mode"], + row["query_datarange"], + row["key_datarange"], + row["weights_datarange"], + row["q_scale_datarange"], + row["k_scale_datarange"], + row["cmp_ratio"], + row["return_value"], + row["output_idx_offset"], + ) + ) + + return test_cases + + except Exception as e: + pytest.skip(f"Failed to read Excel file: {e}", allow_module_level=True) + return None + + +def save_test_case(test_cases, file_path): + print("正在保存pt文件...") + # 创建输出目录 + os.makedirs(file_path, exist_ok=True) + + for idx, case in enumerate(test_cases): + try: + case_name = case[0] + output_tensors = generate_qliv2_test_data(case[1:]) + # 生成文件名 + input_filename = f"{case_name}.pt" + input_filepath = os.path.join(file_path, input_filename) + + # 保存数据 + torch.save(output_tensors, input_filepath) + print(f"测试用例已保存到: {input_filepath}") + + except Exception as e: + print(f"[失败] 生成 pt 文件失败: {case[0]} (索引: {idx})") + print(f"错误详情: {e}") + + +def main(): + parser = argparse.ArgumentParser(description="qliv2_pt_save.py 接收路径参数") + parser.add_argument("path1", type=str, help="第一个路径") + parser.add_argument("path2", type=str, help="第二个路径") + parser.add_argument( + "--sheet", "-S", type=str, default="Sheet1", help="Sheet 页名（默认: Sheet1）" + ) + args = parser.parse_args() + path1 = args.path1 + path2 = args.path2 + testcase = load_excel_test_cases(path1, args.sheet) + save_test_case(testcase, path2) + + +if __name__ == "__main__": + main() diff --git a/xllm_ops/attention/quant_lightning_indexer_v2/tests/pytest/batch/replace_path.py b/xllm_ops/attention/quant_lightning_indexer_v2/tests/pytest/batch/replace_path.py new file mode 100644 index 0000000..515d2d4 --- /dev/null +++ b/xllm_ops/attention/quant_lightning_indexer_v2/tests/pytest/batch/replace_path.py @@ -0,0 +1,38 @@ +#!/usr/bin/python +# -*- coding: utf-8 -*- +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +import fileinput +import sys + +def replace_paths_in_test_file(test_file_path, path): + """ + 替换test_quant_lightning_indexer_v2_batch.py中的占位符为实际路径 + :param test_file_path: test_quant_lightning_indexer_v2_batch.py的路径 + :param path1: 实际路径 + """ + try: + # 逐行替换占位符 + with fileinput.FileInput(test_file_path, inplace=True, backup='.bak') as f: + for line in f: + # 替换__PATH__为实际路径 + line = line.replace('__PATH__', path) + # 输出替换后的行（inplace=True会自动写回文件） + print(line, end='') + print(f" 已成功替换 {test_file_path} 中的路径") + except Exception as e: + print(f" 替换路径失败：{e}") + sys.exit(1) + +if __name__ == '__main__': + test_file = sys.argv[1] + path = sys.argv[2] + replace_paths_in_test_file(test_file, path) \ No newline at end of file diff --git a/xllm_ops/attention/quant_lightning_indexer_v2/tests/pytest/batch_isolated_run.sh b/xllm_ops/attention/quant_lightning_indexer_v2/tests/pytest/batch_isolated_run.sh new file mode 100644 index 0000000..689d5ed --- /dev/null +++ b/xllm_ops/attention/quant_lightning_indexer_v2/tests/pytest/batch_isolated_run.sh @@ -0,0 +1,142 @@ +#!/bin/bash +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- +# ----------------------------------------------------------------------------------------------------------- + +# 批量隔离执行脚本: +# 1. 获取指定路径下的所有 .pt 用例 +# 2. 对每条用例单独拉起一个 pytest 进程执行, 进程间完全隔离 +# - 某条用例 device 越界/崩溃不会影响后续用例 +# - msprof 挂载在单用例进程上, 性能数据采集互不干扰 +# 用法: +# bash batch_isolated_run.sh [用例目录] [是否msprof采集: 0|1] [运行模式: eager|graph] +# 示例: +# bash batch_isolated_run.sh # 默认 pt_path 目录, 不采集性能, eager模式 +# bash batch_isolated_run.sh pt_path 1 # 启用 msprof 采集性能, eager模式 +# bash batch_isolated_run.sh pt_path 0 graph # graph模式, 不采集性能 +# ----------------------------------------------------------------------------------------------------------- + +set -o pipefail + +# Ctrl+C / SIGTERM 中断处理: 递归杀所有子进程后退出 +_cleanup_on_interrupt() { + echo -e "\n\n[中断] 收到终止信号，正在清理子进程..." | tee -a "$SUMMARY_LOG" + pkill -TERM -P $$ 2>/dev/null + sleep 2 + pkill -KILL -P $$ 2>/dev/null + exit 130 +} +trap _cleanup_on_interrupt SIGINT SIGTERM + +TEST_SCRIPT="test_quant_lightning_indexer_v2_batch.py" +TESTCASE_DIR="${1:-./pt_path}" +USE_MSPROF="${2:-0}" +RUN_MODE="${3:-eager}" +RESULT_XLSX="result.xlsx" +SUMMARY_LOG="batch_summary.log" +FAIL_LOG="batch_fail_list.log" + +# 清理旧文件 +[ -f "$RESULT_XLSX" ] && rm -f "$RESULT_XLSX" +[ -f "${RESULT_XLSX%.xlsx}_perf.xlsx" ] && rm -f "${RESULT_XLSX%.xlsx}_perf.xlsx" +rm -f "${RESULT_XLSX%.xlsx}_perf.xlsx.tmp.xlsx" +: > "$SUMMARY_LOG" +: > "$FAIL_LOG" + +# 清理旧的 PROF 文件夹, 避免与本次运行的数据混淆 +# 匹配 PROF_* 和 PROF_*_test_case_* 两种命名格式 +for _prof_dir in PROF_*/; do + [ -d "$_prof_dir" ] && rm -rf "$_prof_dir" +done + +# 1. 获取指定路径下的所有用例路径 +if [ ! -d "$TESTCASE_DIR" ]; then + echo "错误: 用例目录不存在: $TESTCASE_DIR" + exit 1 +fi + +mapfile -t CASE_FILES < <(find "$TESTCASE_DIR" -maxdepth 1 -name "*.pt" | sort) +TOTAL=${#CASE_FILES[@]} +if [ "$TOTAL" -eq 0 ]; then + echo "错误: 目录 $TESTCASE_DIR 下未找到 .pt 用例" + exit 1 +fi + +echo "共发现 $TOTAL 条用例, 目录: $TESTCASE_DIR , msprof采集: $USE_MSPROF , 运行模式: $RUN_MODE" +echo "开始隔离批量执行..." | tee -a "$SUMMARY_LOG" + +PASS=0 +FAIL=0 +FAIL_LIST=() + +# 2. 对每条用例单独调用一次测试脚本, 独立进程 +i=0 +for case_file in "${CASE_FILES[@]}"; do + i=$((i+1)) + case_name=$(basename "$case_file") + echo -e "\n===== [$i/$TOTAL] 执行用例: $case_name =====" | tee -a "$SUMMARY_LOG" + + if [ "$USE_MSPROF" = "1" ]; then + RUN_CMD="QLIV2_TESTCASE_PATH=\"${case_file}\" QLIV2_RUN_MODE=\"${RUN_MODE}\" msprof python3 -m pytest -rA -s ${TEST_SCRIPT} -v -m ci -W ignore::UserWarning -W ignore::DeprecationWarning" + else + RUN_CMD="QLIV2_TESTCASE_PATH=\"${case_file}\" QLIV2_RUN_MODE=\"${RUN_MODE}\" python3 -m pytest -rA -s ${TEST_SCRIPT} -v -m ci -W ignore::UserWarning -W ignore::DeprecationWarning" + fi + + eval "$RUN_CMD" 2>&1 | grep -v "^ninja: no work to do\.$" | tee -a "$SUMMARY_LOG" + status=${PIPESTATUS[0]} + + if [ "$status" -eq 0 ]; then + PASS=$((PASS+1)) + echo "[PASS] $case_name" | tee -a "$SUMMARY_LOG" + # 增量收集性能数据（每条用例跑完立即写入 result_perf.xlsx） + if [ "$USE_MSPROF" = "1" ]; then + sync + python3 collect_perf_data.py --incremental --test_result_path "$RESULT_XLSX" 2>&1 | tee -a "$SUMMARY_LOG" + # 重命名 PROF 文件夹，防止下一条用例的 msprof 覆盖 + _latest_prof=$(ls -dt PROF_*/ 2>/dev/null | head -1) + if [ -n "$_latest_prof" ]; then + _new_name="${_latest_prof%/}_${case_name%.pt}" + mv "$_latest_prof" "$_new_name" 2>/dev/null + fi + fi + else + FAIL=$((FAIL+1)) + FAIL_LIST+=("$case_name") + echo "[FAIL] $case_name" | tee -a "$SUMMARY_LOG" + echo "$case_name" >> "$FAIL_LOG" + fi +done + +# 3. 最终汇总（批量模式兜底，确保所有用例的性能数据都已收集） +if [ "$USE_MSPROF" = "1" ]; then + echo -e "\n========== 性能数据汇总校验 ==========" | tee -a "$SUMMARY_LOG" + python3 collect_perf_data.py --test_result_path "$RESULT_XLSX" 2>&1 | tee -a "$SUMMARY_LOG" +fi + +# 汇总 +echo -e "\n========== 批量执行汇总 ==========" | tee -a "$SUMMARY_LOG" +echo "总计: $TOTAL 通过: $PASS 失败: $FAIL" | tee -a "$SUMMARY_LOG" +if [ "$FAIL" -gt 0 ]; then + echo "失败用例:" | tee -a "$SUMMARY_LOG" + for f in "${FAIL_LIST[@]}"; do + echo " - $f" | tee -a "$SUMMARY_LOG" + done +fi +echo "详细日志: $SUMMARY_LOG" +echo "失败清单: $FAIL_LOG" +echo "结果表格: $RESULT_XLSX" +if [ "$USE_MSPROF" = "1" ]; then + echo "性能表格: ${RESULT_XLSX%.xlsx}_perf.xlsx" +fi + +if [ "$FAIL" -gt 0 ]; then + exit 1 +fi +exit 0 diff --git a/xllm_ops/attention/quant_lightning_indexer_v2/tests/pytest/collect_perf_data.py b/xllm_ops/attention/quant_lightning_indexer_v2/tests/pytest/collect_perf_data.py new file mode 100644 index 0000000..fe5f015 --- /dev/null +++ b/xllm_ops/attention/quant_lightning_indexer_v2/tests/pytest/collect_perf_data.py @@ -0,0 +1,254 @@ +#!/usr/bin/python +# -*- coding: utf-8 -*- +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +import os +import glob +import pandas as pd +import argparse + + +def extract_qliv2_row(prof_folder): + prof_output = os.path.join(prof_folder, "mindstudio_profiler_output") + if not os.path.isdir(prof_output): + return None + + csv_files = glob.glob(os.path.join(prof_output, "op_summary*.csv")) + if not csv_files: + return None + + df = pd.read_csv(csv_files[0]) + target = df[df['Op Name'] == 'QuantLightningIndexerV2'] + if target.empty: + return None + + row = target.iloc[0].to_dict() + return row + + +def collect_and_save_incremental(prof_folder, test_result_path): + """增量模式：读取单条用例的性能数据，追加到 result_perf.xlsx""" + if not os.path.exists(test_result_path): + print(f"结果文件不存在: {test_result_path}") + return None + + df_result = pd.read_excel(test_result_path) + if df_result.empty: + return None + + row_data = extract_qliv2_row(prof_folder) + if row_data is None: + return None + + last_idx = df_result.shape[0] - 1 + case_name = df_result.iloc[last_idx]["case_name"] + + perf_dict = {last_idx: row_data} + perf_df = pd.DataFrame.from_dict(perf_dict, orient='index') + + overlap = set(df_result.columns) & set(perf_df.columns) + if overlap: + rename_map = {col: f"op_{col}" for col in overlap} + perf_df = perf_df.rename(columns=rename_map) + + perf_path = test_result_path.replace(".xlsx", "_perf.xlsx") + tmp_path = perf_path + ".tmp.xlsx" + + if os.path.exists(perf_path): + df_perf = pd.read_excel(perf_path) + # 以 df_result 为基准重建，保留已有 perf 列，追加新行 + perf_cols = [c for c in df_perf.columns if c not in df_result.columns] + df_merged = df_result.copy() + for col in perf_cols: + df_merged[col] = None + for idx in range(min(len(df_result), len(df_perf))): + if col in df_perf.columns and idx < len(df_perf): + val = df_perf.at[idx, col] + if not (isinstance(val, float) and pd.isna(val)): + df_merged.at[idx, col] = val + for col in perf_df.columns: + if col not in df_merged.columns: + df_merged[col] = None + df_merged.loc[last_idx, col] = perf_df.loc[last_idx, col] + df_merged.to_excel(tmp_path, index=False) + os.replace(tmp_path, perf_path) + else: + df_result_with_perf = pd.concat([df_result, perf_df], axis=1) + df_result_with_perf.to_excel(tmp_path, index=False) + os.replace(tmp_path, perf_path) + + print(f" [perf] {case_name} Task Duration: {row_data['Task Duration(us)']}us -> {perf_path}") + return row_data + + +def collect_all(test_result_path, is_compare=False, perf_golden_path="perf_golden.xlsx"): + """批量模式：收集所有 PROF 文件夹的性能数据""" + if not os.path.exists(test_result_path): + print(f"结果文件不存在: {test_result_path}") + return + + df_b = pd.read_excel(test_result_path) + valid_mask = df_b["result"] != "NPU ERROR" + valid_count = valid_mask.sum() + + if valid_count == 0: + print("没有有效用例,跳过性能数据收集") + return + + prof_folders = sorted( + [d for d in os.listdir('.') if os.path.isdir(d) and d.startswith('PROF')], + key=lambda x: os.path.getmtime(x) + ) + + print(f"============= 开始收集性能数据 =============") + print(f"有效用例数: {valid_count}, PROF文件夹数: {len(prof_folders)}") + + if len(prof_folders) == 0: + print("未找到PROF文件夹, 跳过性能数据收集") + return + + if len(prof_folders) != valid_count: + print(f"警告: PROF文件夹数量({len(prof_folders)})与有效用例数({valid_count})不一致") + + perf_rows = {} + prof_idx = 0 + + for i in range(df_b.shape[0]): + if not valid_mask.iloc[i]: + continue + + if prof_idx >= len(prof_folders): + print(f" [{i}] PROF文件夹不足, 跳过剩余用例") + break + + prof = prof_folders[prof_idx] + case_name = df_b.iloc[i]["case_name"] + row_data = extract_qliv2_row(prof) + + if row_data is not None: + perf_rows[i] = row_data + print(f" [{prof_idx+1}] {case_name} -> {prof} (Task Duration: {row_data['Task Duration(us)']}us)") + else: + print(f" [{prof_idx+1}] {case_name}: 未找到QuantLightningIndexer数据") + + prof_idx += 1 + + if not perf_rows: + print("未收集到任何性能数据") + return + + perf_df = pd.DataFrame.from_dict(perf_rows, orient='index') + overlap = set(df_b.columns) & set(perf_df.columns) + if overlap: + rename_map = {col: f"op_{col}" for col in overlap} + perf_df = perf_df.rename(columns=rename_map) + print(f"op_summary列名冲突已重命名: {list(rename_map.values())}") + + df_b = pd.concat([df_b, perf_df], axis=1) + + if is_compare: + try: + df_c = pd.read_excel(perf_golden_path) + except Exception as e: + print(f"读取基线数据失败: {e}") + is_compare = False + + if is_compare: + perf_threshold = 10 + perf_fail_list = [] + df_b["qliv2_perf_diff"] = "" + df_b["perf_result"] = "" + + for i in perf_rows.keys(): + try: + cur_sas = df_b.at[i, "Task Duration(us)"] + golden_sas = df_c.iloc[i]["Task Duration(us)"] + sas_diff = float(cur_sas) - float(golden_sas) + df_b.at[i, "qliv2_perf_diff"] = sas_diff + + if abs(sas_diff) > perf_threshold: + df_b.at[i, "perf_result"] = "Failed" + perf_fail_list.append(df_b.iloc[i]["case_name"]) + else: + df_b.at[i, "perf_result"] = "Pass" + except Exception as e: + print(f" 基线对比出错 (行{i}): {e}") + + if perf_fail_list: + print(f"性能不达标用例: {perf_fail_list}") + + new_path = test_result_path.replace(".xlsx", "_perf.xlsx") + tmp_path = new_path + ".tmp.xlsx" + if os.path.exists(new_path): + existing = pd.read_excel(new_path) + existing = existing.reindex(range(df_b.shape[0])) + for col in df_b.columns: + for idx in range(df_b.shape[0]): + if pd.isna(existing.at[idx, col]): + existing.at[idx, col] = df_b.at[idx, col] + op_cols = [c for c in existing.columns if c.startswith("op_") or c == "Task Duration(us)"] + filled_count = existing[op_cols].dropna(how='all').shape[0] if op_cols else 0 + if filled_count >= len(perf_rows): + print(f"增量模式已收集 {filled_count}/{len(perf_rows)} 条性能数据，跳过批量补采") + return + print(f"增量模式仅 {filled_count}/{len(perf_rows)} 条，批量补采覆盖") + for col in perf_df.columns: + if col in existing.columns: + for idx in perf_rows: + if pd.isna(existing.at[idx, col]): + existing.at[idx, col] = perf_df.at[idx, col] + else: + existing[col] = None + for idx in perf_rows: + existing.at[idx, col] = perf_df.at[idx, col] + existing.to_excel(tmp_path, index=False) + os.replace(tmp_path, new_path) + return + + df_b.to_excel(tmp_path, index=False) + os.replace(tmp_path, new_path) + + print(f"\n性能数据已保存: {new_path}") + print(f"共拼接 {len(perf_rows)} 条用例的 op_summary 全字段数据") + print(f"============= 性能数据收集完成 =============") + + +def main(): + parser = argparse.ArgumentParser(description='收集性能数据') + parser.add_argument('--test_result_path', type=str, default="result.xlsx") + parser.add_argument('--incremental', action='store_true', default=False, + help='增量模式：只收集最新一条用例的性能数据') + parser.add_argument('--prof_folder', type=str, default=None, + help='增量模式下指定 PROF 文件夹路径') + parser.add_argument('--is_compare', action='store_true', default=False) + parser.add_argument('--perf_golden_path', type=str, default="perf_golden.xlsx") + args = parser.parse_args() + + if args.incremental: + if args.prof_folder: + prof = args.prof_folder + else: + prof_folders = sorted( + [d for d in os.listdir('.') if os.path.isdir(d) and d.startswith('PROF')], + key=lambda x: os.path.getmtime(x) + ) + if not prof_folders: + print("未找到PROF文件夹") + return + prof = prof_folders[-1] + + collect_and_save_incremental(prof, args.test_result_path) + else: + collect_all(args.test_result_path, args.is_compare, args.perf_golden_path) + + +if __name__ == "__main__": + main() diff --git a/xllm_ops/attention/quant_lightning_indexer_v2/tests/pytest/pytest.ini b/xllm_ops/attention/quant_lightning_indexer_v2/tests/pytest/pytest.ini new file mode 100644 index 0000000..339c185 --- /dev/null +++ b/xllm_ops/attention/quant_lightning_indexer_v2/tests/pytest/pytest.ini @@ -0,0 +1,4 @@ +[pytest] +markers = + ci: mark a test as a CI test + graph: marks tests as graph mode compilation tests \ No newline at end of file diff --git a/xllm_ops/attention/quant_lightning_indexer_v2/tests/pytest/qliv2_parameter_normalization.py b/xllm_ops/attention/quant_lightning_indexer_v2/tests/pytest/qliv2_parameter_normalization.py new file mode 100644 index 0000000..ba60e19 --- /dev/null +++ b/xllm_ops/attention/quant_lightning_indexer_v2/tests/pytest/qliv2_parameter_normalization.py @@ -0,0 +1,67 @@ +#!/usr/bin/python +# -*- coding: utf-8 -*- +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +"""Pure QLI_V2 parameter normalization shared by pytest and TTK adapters.""" + + +def normalize_qliv2_params(params): + """Apply the batch pytest scalar conversions without generating test data.""" + values = tuple(params) + if len(values) not in (32, 33): + raise ValueError( + f"QLI_V2 parameter count mismatch: got {len(values)}, expected 32 or 33" + ) + + has_weight_dtype = len(values) == 33 + if has_weight_dtype: + ( + batch_size, q_seq, k_seq, q_t_size, k_t_size, q_head_num, + k_head_num, head_dim, block_size, block_num, qk_dtype, + weight_dtype, dequant_dtype, actual_seq_dtype, cu_seqlens_q, + cu_seqlens_k, seqused_q, seqused_k, cmp_residual_k, max_seqlen_q, + quant_mode, layout_query, layout_key, sparse_count, sparse_mode, + query_datarange, key_datarange, weights_datarange, + q_scale_datarange, k_scale_datarange, cmp_ratio, return_value, + output_idx_offset, + ) = values + else: + ( + batch_size, q_seq, k_seq, q_t_size, k_t_size, q_head_num, + k_head_num, head_dim, block_size, block_num, qk_dtype, + dequant_dtype, actual_seq_dtype, cu_seqlens_q, cu_seqlens_k, + seqused_q, seqused_k, cmp_residual_k, max_seqlen_q, quant_mode, + layout_query, layout_key, sparse_count, sparse_mode, + query_datarange, key_datarange, weights_datarange, + q_scale_datarange, k_scale_datarange, cmp_ratio, return_value, + output_idx_offset, + ) = values + weight_dtype = dequant_dtype + + q_t_size = 0 if q_t_size is None else int(q_t_size) + k_t_size = 0 if k_t_size is None else int(k_t_size) + block_size = 0 if block_size is None else int(block_size) + block_num = 0 if block_num is None else int(block_num) + max_seqlen_q = -1 if max_seqlen_q is None else int(max_seqlen_q) + quant_mode = 1 if quant_mode is None else int(quant_mode) + if has_weight_dtype and weight_dtype is None: + weight_dtype = dequant_dtype + + normalized = ( + int(batch_size), int(q_seq), int(k_seq), q_t_size, k_t_size, + int(q_head_num), int(k_head_num), int(head_dim), block_size, block_num, + qk_dtype, dequant_dtype, actual_seq_dtype, cu_seqlens_q, cu_seqlens_k, + seqused_q, seqused_k, cmp_residual_k, max_seqlen_q, quant_mode, + layout_query, layout_key, int(sparse_count), int(sparse_mode), + query_datarange, key_datarange, weights_datarange, q_scale_datarange, + k_scale_datarange, int(cmp_ratio), int(return_value), output_idx_offset, + ) + return normalized[:11] + (weight_dtype,) + normalized[11:] diff --git a/xllm_ops/attention/quant_lightning_indexer_v2/tests/pytest/qliv2_test_utils.py b/xllm_ops/attention/quant_lightning_indexer_v2/tests/pytest/qliv2_test_utils.py new file mode 100644 index 0000000..19159ef --- /dev/null +++ b/xllm_ops/attention/quant_lightning_indexer_v2/tests/pytest/qliv2_test_utils.py @@ -0,0 +1,223 @@ +#!/usr/bin/python +# -*- coding: utf-8 -*- +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +"""NPU-independent case selection, naming, and result helpers for QLI_V2 tests.""" + +import re +from pathlib import Path + +import pandas as pd + + +PARAM_NAMES = ( + "batch_size", + "q_seq", + "k_seq", + "q_t_size", + "k_t_size", + "q_head_num", + "k_head_num", + "head_dim", + "block_size", + "block_num", + "qk_dtype", + "weight_dtype", + "dequant_dtype", + "actual_seq_dtype", + "cu_seq_q", + "cu_seq_k", + "act_seq_q", + "act_seq_k", + "cmp_residual_k", + "max_seqlen_q", + "quant_mode", + "layout_query", + "layout_key", + "sparse_count", + "sparse_mode", + "query_datarange", + "key_datarange", + "weights_datarange", + "q_scale_datarange", + "k_scale_datarange", + "cmp_ratio", + "return_value", + "output_idx_offset", +) + + +def ensure_comparison_passed( + case_name, + result, + fulfill_percent, + result_return_value="N/A", + fulfill_percent_return_value=0, +): + """Raise a serializable error when an accuracy comparison fails.""" + failures = [] + if result != "Pass": + failures.append(f"index result={result}, fulfill_percent={fulfill_percent}") + if result_return_value not in ("N/A", "Pass"): + failures.append( + "value result=" + f"{result_return_value}, fulfill_percent={fulfill_percent_return_value}" + ) + if failures: + raise AssertionError( + f"accuracy comparison failed for {case_name}: " + "; ".join(failures) + ) + + +class QliV2CaseSelector: + """Resolve an ordered subset of PT cases by explicit name or one-based index.""" + + @staticmethod + def natural_key(path): + return [ + int(part) if part.isdigit() else part.lower() + for part in re.split(r"(\d+)", Path(path).name) + ] + + @staticmethod + def parse_indexes(expression, total): + if not expression: + return [] + indexes = [] + for token in str(expression).split(","): + token = token.strip() + if not token: + continue + if "-" in token: + start_text, end_text = token.split("-", 1) + start, end = int(start_text), int(end_text) + if end < start: + raise ValueError(f"invalid descending case index range: {token}") + indexes.extend(range(start, end + 1)) + else: + indexes.append(int(token)) + invalid = [index for index in indexes if index < 1 or index > total] + if invalid: + raise ValueError(f"case indexes out of range 1..{total}: {invalid}") + return indexes + + @classmethod + def resolve(cls, pt_dir, explicit_files="", case_names="", case_indexes=""): + if explicit_files: + candidates = [ + Path(item.strip()) for item in explicit_files.split(",") if item.strip() + ] + else: + directory = Path(pt_dir) + if not directory.is_dir(): + raise ValueError(f"PT directory does not exist: {directory}") + candidates = sorted(directory.glob("*.pt"), key=cls.natural_key) + + missing = [str(path) for path in candidates if not path.is_file()] + if missing: + raise ValueError(f"PT files do not exist: {missing}") + if not candidates: + raise ValueError(f"no PT cases found in: {pt_dir}") + if case_names and case_indexes: + raise ValueError("case names and case indexes cannot be specified together") + + if case_names: + by_name = {path.stem: path for path in candidates} + selected = [] + unknown = [] + for item in case_names.split(","): + name = Path(item.strip()).stem + if not name: + continue + if name not in by_name: + unknown.append(name) + else: + selected.append(by_name[name]) + if unknown: + raise ValueError(f"unknown case names: {unknown}") + candidates = selected + elif case_indexes: + indexes = cls.parse_indexes(case_indexes, len(candidates)) + candidates = [candidates[index - 1] for index in indexes] + + return [str(path) for path in candidates] + + +class QliV2ResultWriter: + """Build stable case names and append rows using the batch result schema.""" + + @staticmethod + def case_name(params, explicit_name=None): + if explicit_name: + normalized = re.sub(r"[^A-Za-z0-9_.-]+", "_", str(explicit_name)) + normalized = normalized.strip("._-") + if not normalized: + raise ValueError("case name has no usable filename characters") + return normalized + + values = list(params) + readable = ( + f"QLI_B{values[0]}_S1{values[1]}_S2{values[2]}_" + f"N1{values[5]}_N2{values[6]}_D{values[7]}_" + f"{values[21]}_{values[22]}_{values[10]}_" + f"QM{values[20]}_SM{values[24]}_CR{values[30]}_" + f"K{values[23]}_RV{values[31]}" + ) + return re.sub(r"[^A-Za-z0-9_.-]+", "_", readable) + + @staticmethod + def row( + case_name, + params, + result, + fulfill_percent, + result_return_value="N/A", + fulfill_percent_return_value=0, + ): + values = list(params) + if len(values) != len(PARAM_NAMES): + raise ValueError( + f"QLI_V2 parameter count mismatch: got {len(values)}, expected {len(PARAM_NAMES)}" + ) + row = {"case_name": case_name} + row.update(dict(zip(PARAM_NAMES, values))) + row.update( + { + "result": result, + "fulfill_percent": fulfill_percent, + "result_return_value": result_return_value, + "fulfill_percent_return_value": fulfill_percent_return_value, + } + ) + return row + + @staticmethod + def append(path, row): + output = Path(path) + output.parent.mkdir(parents=True, exist_ok=True) + if output.exists(): + frame = pd.read_excel(output) + expected_columns = list(row.keys()) + legacy_columns = [ + name for name in expected_columns if name != "return_value" + ] + if list(frame.columns) == legacy_columns: + frame["return_value"] = None + frame = frame[expected_columns] + elif list(frame.columns) != expected_columns: + raise ValueError( + "result columns do not match existing Excel: " + f"existing={list(frame.columns)}, current={list(row.keys())}" + ) + frame = pd.concat([frame, pd.DataFrame([row])], ignore_index=True) + else: + frame = pd.DataFrame([row]) + frame.to_excel(output, index=False) diff --git a/xllm_ops/attention/quant_lightning_indexer_v2/tests/pytest/quant_lightning_indexer_v2_acl_graph.py b/xllm_ops/attention/quant_lightning_indexer_v2/tests/pytest/quant_lightning_indexer_v2_acl_graph.py new file mode 100644 index 0000000..bda599c --- /dev/null +++ b/xllm_ops/attention/quant_lightning_indexer_v2/tests/pytest/quant_lightning_indexer_v2_acl_graph.py @@ -0,0 +1,279 @@ +#!/usr/bin/python +# -*- coding: utf-8 -*- +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +import test +import torch +import torch_npu +import pytest +import torchair +import torch.nn as nn +from torchair.configs.compiler_config import CompilerConfig +import cann_ops_transformer +import quant_lightning_indexer_v2_golden + +QUANT_MODE_MXFP4 = 5 + + +class QLIV2Network(nn.Module): + def __init__(self): + super(QLIV2Network, self).__init__() + + def forward( + self, + query, + key, + weights, + query_dequant_scale, + key_dequant_scale, + cu_seqlens_q, + cu_seqlens_k, + seqused_q, + seqused_k, + cmp_residual_k, + output_idx_offset, + max_seqlen_q, + block_table, + metadata, + quant_mode, + layout_query, + layout_key, + sparse_count, + sparse_mode, + cmp_ratio, + return_value, + ): + return torch.ops.cann_ops_transformer.quant_lightning_indexer( + query, + key, + weights, + query_dequant_scale, + key_dequant_scale, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + seqused_q=seqused_q, + seqused_k=seqused_k, + cmp_residual_k=cmp_residual_k, + output_idx_offset=output_idx_offset, + max_seqlen_q=max_seqlen_q, + block_table=block_table, + metadata=metadata, + quant_mode=quant_mode, + layout_q=layout_query, + layout_k=layout_key, + topk=sparse_count, + mask_mode=sparse_mode, + cmp_ratio=cmp_ratio, + return_value=return_value, + ) + + +def _qliv2_prepare_tensors_and_metadata(params, tensor_dict): + """ + 统一处理 tensor 准备和 metadata 构造（共用逻辑）。 + 兼容两个来源：generate_qliv2_test_data 返回值和 .pt 文件加载，二者都在 CPU。 + """ + qk_dtype = params[10] + quant_mode = tensor_dict["quant_mode"] + + if quant_mode == QUANT_MODE_MXFP4: + # TorchAir通过foreach_copy搬运图输入，当前不支持FP4 shell dtype；使用相同存储的 + # packed uint8视图，C++入口会根据quant_mode恢复ACL_FLOAT4_E2M1语义。 + query = tensor_dict["query"].view(torch.uint8).npu() + key = tensor_dict["key"].view(torch.uint8).npu() + if "blockFusion" in tensor_dict and tensor_dict["blockFusion"] is not None: + blockFusion = tensor_dict["blockFusion"].view(torch.uint8).npu() + elif qk_dtype == "FLOAT8_E4M3FN" or qk_dtype == torch.float8_e4m3fn: + query = tensor_dict["query"].to(dtype=torch.float8_e4m3fn).npu() + key = tensor_dict["key"].to(dtype=torch.float8_e4m3fn).npu() + if "blockFusion" in tensor_dict and tensor_dict["blockFusion"] is not None: + blockFusion = tensor_dict["blockFusion"] + if blockFusion.dtype == torch.uint8: + blockFusion = blockFusion.view(torch.float8_e4m3fn) + else: + blockFusion = blockFusion.to(dtype=torch.float8_e4m3fn) + blockFusion = blockFusion.npu() + else: + query = tensor_dict["query"].npu() + key = tensor_dict["key"].npu() + if "blockFusion" in tensor_dict and tensor_dict["blockFusion"] is not None: + blockFusion = tensor_dict["blockFusion"].npu() + + weights = tensor_dict["weights"].npu() + query_dequant_scale = tensor_dict["query_dequant_scale"].npu() + key_dequant_scale = tensor_dict["key_dequant_scale"].npu() + + if "blockFusion" in tensor_dict and tensor_dict["blockFusion"] is not None: + block_num = int(params[9]) + block_size = int(params[8]) + head_dim = int(params[7]) + k_head_num = int(params[6]) + dequant_dtype_str = params[12] + if dequant_dtype_str == "FP16" or dequant_dtype_str == torch.float16: + dequant_dtype = torch.float16 + elif dequant_dtype_str == "FP32" or dequant_dtype_str == torch.float32: + dequant_dtype = torch.float32 + else: + dequant_dtype = torch.float16 + key = blockFusion[:, : block_size * k_head_num * head_dim].view( + block_num, block_size, k_head_num, head_dim + ) + key_dequant_scale = ( + blockFusion[:, block_size * k_head_num * head_dim :] + .view(dequant_dtype) + .view(block_num, block_size, k_head_num) + ) + + cu_seqlens_query = ( + tensor_dict["cu_seqlens_query"].npu() + if tensor_dict["cu_seqlens_query"] is not None + else None + ) + cu_seqlens_key = ( + tensor_dict["cu_seqlens_key"].npu() + if tensor_dict["cu_seqlens_key"] is not None + else None + ) + seqused_q = ( + tensor_dict["seqused_q"].npu() if tensor_dict["seqused_q"] is not None else None + ) + seqused_k = ( + tensor_dict["seqused_k"].npu() if tensor_dict["seqused_k"] is not None else None + ) + output_idx_offset = ( + tensor_dict["output_idx_offset"].npu() + if tensor_dict["output_idx_offset"] is not None + else None + ) + block_table = ( + tensor_dict["block_table"].npu() + if tensor_dict["block_table"] is not None + else None + ) + cmp_residual_k_for_npu = ( + tensor_dict["cmp_residual_k_for_npu"].npu() + if tensor_dict.get("cmp_residual_k_for_npu") is not None + else None + ) + + layout_query = tensor_dict["layout_query"] + layout_key = tensor_dict["layout_key"] + sparse_count = tensor_dict["sparse_count"] + sparse_mode = tensor_dict["sparse_mode"] + cmp_ratio = tensor_dict["cmp_ratio"] + max_seqlen_q_meta = tensor_dict["max_seqlen_q_meta"] + max_seqlen_k_meta = tensor_dict["max_seqlen_k_meta"] + + q_head_num = int(params[5]) + k_head_num = int(params[6]) + head_dim = int(params[7]) + batch_size = int(params[0]) + + metadata = torch.ops.cann_ops_transformer.quant_lightning_indexer_metadata( + cu_seqlens_q=cu_seqlens_query, + cu_seqlens_k=cu_seqlens_key, + seqused_q=seqused_q, + seqused_k=seqused_k, + cmp_residual_k=cmp_residual_k_for_npu, + batch_size=batch_size, + max_seqlen_q=max_seqlen_q_meta, + max_seqlen_k=max_seqlen_k_meta, + num_heads_q=q_head_num, + num_heads_k=k_head_num, + head_dim=head_dim, + topk=sparse_count, + quant_mode=quant_mode, + mask_mode=sparse_mode, + layout_q=layout_query, + layout_k=layout_key, + cmp_ratio=cmp_ratio, + ) + metadata = metadata.npu() + + run_args = { + "query": query, + "key": key, + "weights": weights, + "query_dequant_scale": query_dequant_scale, + "key_dequant_scale": key_dequant_scale, + "cu_seqlens_q": cu_seqlens_query, + "cu_seqlens_k": cu_seqlens_key, + "seqused_q": seqused_q, + "seqused_k": seqused_k, + "cmp_residual_k": cmp_residual_k_for_npu, + "output_idx_offset": output_idx_offset, + "max_seqlen_q": params[19], + "block_table": block_table, + "metadata": metadata, + "quant_mode": quant_mode, + "layout_query": layout_query, + "layout_key": layout_key, + "sparse_count": sparse_count, + "sparse_mode": sparse_mode, + "cmp_ratio": cmp_ratio, + "return_value": params[31], + } + return run_args + + +def _qliv2_run_compiled_graph(run_args): + """ + 通过 torch.compile + torchair 后端执行算子（共用逻辑）。 + """ + config = CompilerConfig() + config.mode = "reduce-overhead" + npu_backend = torchair.get_npu_backend(compiler_config=config) + torch._dynamo.reset() + npu_mode = torch.compile( + QLIV2Network().npu(), fullgraph=False, backend=npu_backend, dynamic=False + ) + npu_result, npu_topk_value = npu_mode(**run_args) + torch.npu.synchronize() + if run_args["return_value"]: + if npu_topk_value.shape != npu_result.shape: + raise RuntimeError( + "sparse_values and sparse_indices must have the same shape when return_value=1, " + f"but got {tuple(npu_topk_value.shape)} and {tuple(npu_result.shape)}" + ) + npu_topk_value, npu_sort_order = npu_topk_value.sort(dim=-1, descending=True) + npu_result = torch.gather(npu_result, dim=-1, index=npu_sort_order) + return npu_result, npu_topk_value + + +def qliv2_output_acl_graph(params): + """ + graph 模式入口（single 用例使用）：即时生成随机 tensor + CPU golden，再走 torch.compile 执行。 + """ + print("acl_graph") + tensor_dict = quant_lightning_indexer_v2_golden.generate_qliv2_test_data(params) + cpu_result = tensor_dict["cpu_result"] + topk_value = tensor_dict["topk_value"] + cpu_topk_value = tensor_dict["cpu_topk_value"] + + run_args = _qliv2_prepare_tensors_and_metadata(params, tensor_dict) + npu_result, npu_topk_value = _qliv2_run_compiled_graph(run_args) + + return cpu_result, npu_result, topk_value, cpu_topk_value, npu_topk_value + + +def qliv2_output_acl_graph_from_pt(params, tensor_dict): + """ + graph 模式入口（batch 用例使用）：从 .pt 文件加载 pre-computed tensor，走 torch.compile 执行。 + 跳过 generate_qliv2_test_data 的随机数据重新生成和 CPU golden 重算。 + """ + cpu_result = tensor_dict["cpu_result"] + topk_value = tensor_dict["topk_value"] + cpu_topk_value = tensor_dict["cpu_topk_value"] + + run_args = _qliv2_prepare_tensors_and_metadata(params, tensor_dict) + npu_result, npu_topk_value = _qliv2_run_compiled_graph(run_args) + + return cpu_result, npu_result, topk_value, cpu_topk_value, npu_topk_value diff --git a/xllm_ops/attention/quant_lightning_indexer_v2/tests/pytest/quant_lightning_indexer_v2_golden.py b/xllm_ops/attention/quant_lightning_indexer_v2/tests/pytest/quant_lightning_indexer_v2_golden.py new file mode 100644 index 0000000..b313425 --- /dev/null +++ b/xllm_ops/attention/quant_lightning_indexer_v2/tests/pytest/quant_lightning_indexer_v2_golden.py @@ -0,0 +1,3116 @@ +#!/usr/bin/python +# -*- coding: utf-8 -*- +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +import test +import torch + +try: + import torch_npu +except ImportError: + torch_npu = None +import pytest +import random +import numpy as np +import math +import ctypes +import copy +import ast +from qliv2_parameter_normalization import normalize_qliv2_params + +try: + import cann_ops_transformer +except ImportError: + cann_ops_transformer = None + +DISCONTINUOUS_KEYS = True # key非连续 +DEFAULT_SPLIT_S1 = False # golden切分S1Flag +DEFAULT_S1SIZE = 2048 # s1切分基本块大小 + +FP32_FRACTION_BITS = 23 # fp32尾数位数 + +HIF8_EXP_ZERO_THRESHOLD = -23 # 边界值 +HIF8_EXP_DML_MIN = -22 # DML最小指数 +HIF8_EXP_DML_MAX = -15 # DML最大指数 +HIF8_EXP_D0 = 0 # D0指数值 +HIF8_EXP_D1_BOUNDARY = 1 # D1指数值 +HIF8_EXP_D2_MIN, HIF8_EXP_D2_MAX = 2, 3 # D2指数范围 +HIF8_EXP_D3_MIN, HIF8_EXP_D3_MAX = 4, 7 # D3指数范围 +HIF8_EXP_D4_MIN, HIF8_EXP_D4_MAX = 8, 15 # D4指数范围 + +HIF8_DOT_DML = 0 # DML: Denormal Low, 指数范围 -22 ~ -16, 0位尾数 +HIF8_DOT_D0 = 1 # D0: 指数为0，3位尾数（最高精度） +HIF8_DOT_D1 = 2 # D1: 指数为±1，3位尾数 +HIF8_DOT_D2 = 4 # D2: 指数为±2 ~ ±3，3位尾数 +HIF8_DOT_D3 = 8 # D3: 指数为±4 ~ ±7，2位尾数 +HIF8_DOT_D4 = 12 # D4: 指数为±8 ~ ±15，1位尾数（最低精度） +HIF8_DOT_INVALID = -1 # 无效状态 + +HIF8_FRAC_BITS_DML = 0 # DML档位尾数位数 +HIF8_FRAC_BITS_D0 = 3 # D0档位尾数位数 +HIF8_FRAC_BITS_D1 = 3 # D1档位尾数位数 +HIF8_FRAC_BITS_D2 = 3 # D2档位尾数位数 +HIF8_FRAC_BITS_D3 = 2 # D3档位尾数位数 +HIF8_FRAC_BITS_D4 = 1 # D4档位尾数位数 + +HIF8_EXP_BITS_DML = 3 # DML档位指数位数 +HIF8_EXP_BITS_D0 = 0 # D0档位指数位数 +HIF8_EXP_BITS_D1 = 1 # D1档位指数位数 +HIF8_EXP_BITS_D2 = 2 # D2档位指数位数 +HIF8_EXP_BITS_D3 = 3 # D3档位指数位数 +HIF8_EXP_BITS_D4 = 4 # D4档位指数位数 + +HIF8_ZERO = 0 +HIF8_NAN = 128 # 0b10000000, NaN +HIF8_NEG_INF = 239 # 0b11101111, -inf +HIF8_NEG_MAX = 238 # 0b11101110, 负极大值 +HIF8_POS_INF = 111 # 0b01101111, +inf +HIF8_POS_MAX = 110 # 0b01101110, 正极大值 + +HIF8_SIGN_MASK = 128 # 0b10000000, 符号位掩码 +HIF8_DOT_MASK = 120 # 0b01110000, dot值掩码 +HIF8_FRAC_MASK_3BIT = 7 # 0b00000111, 3位尾数掩码（D0/D1/D2） +HIF8_FRAC_MASK_2BIT = 3 # 0b00000011, 2位尾数掩码（D3） +HIF8_FRAC_MASK_1BIT = 1 # 0b00000001, 1位尾数掩码（D4） +HIF8_EXP_MASK_DML = 7 # 0b00000111, DML指数掩码（bit0-2） +HIF8_EXP_MASK_D4 = 30 # 0b00011110, D4指数掩码（bit1-4） +HIF8_EXP_MASK_D3 = 28 # 0b00011100, D3指数掩码（bit2-4） +HIF8_EXP_MASK_D2 = 24 # 0b00011000, D2指数掩码（bit3-4） +HIF8_EXP_SIGN_MASK_D1 = 8 # 0b00001000, D1指数掩码（bit3） + +HIF8_DOT_BIT_SHIFT = 3 # Dot值在HiF8中的起始位置(bit3) +HIF8_DML_EXP_OFFSET = 23 # DML指数偏移值 +HIF8_OVERFLOW_SCALE = 1.25 # 溢出阈值缩放因子 +HIF8_MAX_FINITE_VALUE = 32768 # 最大有限值（非饱和模式下的边界值, 2^15 + +SSR_T14_MASK = 16383 # 0b0011 1111 1111 1111, 14位低位掩码 +SSR_F14_OFFSET = 8192 # 0b0010 0000 0000 0000, F14偏移值 +SSR_DML_SHIFT = 10 # SSR舍入移位值 +SSR_RESERVED_BITS = 14 # SSR舍入保留位数 +HYBRID_ROUND_EXP_THRESHOLD = 4 # 混合舍入的指数分界点 + +QUANT_MODE_MXFP8 = 3 +QUANT_MODE_HIFLOAT8 = 4 +QUANT_MODE_MXFP4 = 5 +MX_SCALE_SHAPE_ALIGN = 64 +MX_SCALE_PACK_NUM = 2 +MX_SCALE_GROUP_SIZE = MX_SCALE_SHAPE_ALIGN // MX_SCALE_PACK_NUM +FP4_PACK_NUM = 2 +E8M0_ONE_VALUE = 127 +BF16_SIGNIFICAND_BITS = 8 +BF16_MIN_NORMAL = 2.0**-126 +BF16_MIN_SUBNORMAL = 2.0**-133 +FP4_E2M1_VALUES = torch.tensor( + [ + 0.0, + 0.5, + 1.0, + 1.5, + 2.0, + 3.0, + 4.0, + 6.0, + -0.0, + -0.5, + -1.0, + -1.5, + -2.0, + -3.0, + -4.0, + -6.0, + ], + dtype=torch.float32, +) +MXFP4_TORCH_DTYPE = torch.float4_e2m1fn_x2 + + +def is_mx_quant_mode(quant_mode): + return quant_mode in (QUANT_MODE_MXFP8, QUANT_MODE_MXFP4) + + +def is_mxfp4_quant_mode(quant_mode): + return quant_mode == QUANT_MODE_MXFP4 + + +def round_fp64_to_bf16_rne(value): + if value.dtype != torch.float64: + raise TypeError( + f"round_fp64_to_bf16_rne expects float64 input, got {value.dtype}" + ) + + mantissa, exponent = torch.frexp(value) + rounded_normal = torch.ldexp( + torch.round(mantissa * (1 << BF16_SIGNIFICAND_BITS)), + exponent - BF16_SIGNIFICAND_BITS, + ) + rounded_subnormal = torch.round(value / BF16_MIN_SUBNORMAL) * BF16_MIN_SUBNORMAL + rounded = torch.where( + value.abs() < BF16_MIN_NORMAL, rounded_subnormal, rounded_normal + ) + rounded = torch.where(torch.isfinite(value), rounded, value) + + # The value is already on the BF16 grid, so both casts below are exact. + return rounded.to(torch.float32).to(torch.bfloat16) + + +def reduce_mxfp4_weighted_qk(weight_matrix, qk_matrix): + output_shape = (weight_matrix.shape[0], weight_matrix.shape[1], qk_matrix.shape[2]) + acc_bf16 = torch.zeros( + output_shape, dtype=torch.bfloat16, device=weight_matrix.device + ) + + for g_idx in range(weight_matrix.shape[2]): + weight_fp64 = weight_matrix[:, :, g_idx : g_idx + 1].to(torch.float64) + qk_fp64 = qk_matrix[:, g_idx : g_idx + 1, :].to(torch.float64) + # Model a BF16-destination FMA with one BF16 rounding per G. + fma_fp64 = acc_bf16.to(torch.float64) + weight_fp64 * qk_fp64 + acc_bf16 = round_fp64_to_bf16_rne(fma_fp64) + return acc_bf16 + + +def reduce_mxfp8_weighted_qk(weight_matrix, qk_matrix): + output_shape = (weight_matrix.shape[0], weight_matrix.shape[1], qk_matrix.shape[2]) + acc_fp32 = torch.zeros( + output_shape, dtype=torch.float32, device=weight_matrix.device + ) + + for g_idx in range(weight_matrix.shape[2]): + weight_fp64 = weight_matrix[:, :, g_idx : g_idx + 1].to(torch.float64) + qk_fp64 = qk_matrix[:, g_idx : g_idx + 1, :].to(torch.float64) + # Model an FP32-destination MulAddDst: compute the fused multiply-add in + # FP64, then round once to the FP32 destination after every G. + fma_fp64 = acc_fp32.to(torch.float64) + weight_fp64 * qk_fp64 + acc_fp32 = fma_fp64.to(torch.float32) + return acc_fp32 + + +def get_qk_physical_head_dim(head_dim, quant_mode): + return head_dim // FP4_PACK_NUM if is_mxfp4_quant_mode(quant_mode) else head_dim + + +def e8m0_raw_to_float(raw): + raw_float = raw.to(torch.float32) + scale = torch.pow(2.0, raw_float - E8M0_ONE_VALUE) + scale = torch.where(raw == 255, torch.full_like(scale, float("nan")), scale) + return scale + + +def validate_mx_scale_dtype(scale_dtype): + if scale_dtype != torch.float8_e8m0fnu: + raise TypeError( + "MXFP8/MXFP4 require dequant_dtype=torch.float8_e8m0fnu, " + f"but got {scale_dtype}" + ) + + +def make_e8m0_tensor_from_raw(raw, scale_dtype, to_npu=True): + validate_mx_scale_dtype(scale_dtype) + scale = raw.contiguous().view(scale_dtype) + return scale.npu() if to_npu else scale + + +def make_mx_e8m0_scale_pair( + base_shape, tail_shape, scale_range, scale_dtype, to_npu=True +): + validate_mx_scale_dtype(scale_dtype) + range_min = float(scale_range[0]) + range_max = float(scale_range[1]) + if not math.isfinite(range_min) or not math.isfinite(range_max): + raise ValueError(f"E8M0 scale range must be finite, got {scale_range}") + if range_min <= 0 or range_min > range_max: + raise ValueError( + f"E8M0 scale range must satisfy 0 < min <= max, got {scale_range}" + ) + + # datarange表示真实scale值；在生成器内部转换为E8M0编码（编码e对应2^(e-127)）。 + range_min_code = max(0, math.ceil(math.log2(range_min)) + E8M0_ONE_VALUE) + range_max_code = min(254, math.floor(math.log2(range_max)) + E8M0_ONE_VALUE) + while ( + range_min_code <= 254 + and math.ldexp(1.0, range_min_code - E8M0_ONE_VALUE) < range_min + ): + range_min_code += 1 + while ( + range_max_code >= 0 + and math.ldexp(1.0, range_max_code - E8M0_ONE_VALUE) > range_max + ): + range_max_code -= 1 + if range_min_code > range_max_code: + raise ValueError( + f"E8M0 scale range contains no representable value: {scale_range}" + ) + + # 沿连续存储顺序循环取值，保证不同token和不同D group使用可区分的scale。 + scale_shape = tuple(base_shape) + tuple(tail_shape) + raw_count = range_max_code - range_min_code + 1 + scale_raw = torch.arange(math.prod(scale_shape), dtype=torch.int64) + scale_raw = ( + (scale_raw % raw_count + range_min_code).reshape(scale_shape).to(torch.uint8) + ) + + cpu_scale = e8m0_raw_to_float(scale_raw) + return make_e8m0_tensor_from_raw(scale_raw, scale_dtype, to_npu), cpu_scale + + +def validate_mxfp4_dtype(data_dtype): + if data_dtype != MXFP4_TORCH_DTYPE: + raise TypeError( + f"MXFP4 requires qk_dtype=torch.float4_e2m1fn_x2, but got {data_dtype}" + ) + + +def make_mxfp4_tensor_pair(logical_shape, data_range, data_dtype, to_npu=True): + validate_mxfp4_dtype(data_dtype) + range_min = float(data_range[0]) + range_max = float(data_range[1]) + if ( + not math.isfinite(range_min) + or not math.isfinite(range_max) + or range_min > range_max + ): + raise ValueError(f"invalid FP4 E2M1 numeric range: {data_range}") + if logical_shape[-1] % FP4_PACK_NUM != 0: + raise ValueError(f"MXFP4 head_dim must be even, got {logical_shape[-1]}") + + valid_codes = ( + torch.nonzero( + (FP4_E2M1_VALUES >= range_min) & (FP4_E2M1_VALUES <= range_max), + as_tuple=False, + ) + .flatten() + .to(torch.uint8) + ) + if valid_codes.numel() == 0: + raise ValueError( + f"FP4 E2M1 range contains no representable value: {data_range}" + ) + + raw_indices = torch.randint( + valid_codes.numel(), tuple(logical_shape), dtype=torch.long + ) + raw = valid_codes[raw_indices] + raw_rows = raw.reshape(-1, logical_shape[-1]) + coverage_count = min(valid_codes.numel(), logical_shape[-1]) + coverage = torch.arange(coverage_count).unsqueeze(0) + torch.arange( + raw_rows.shape[0] + ).unsqueeze(1) + raw_rows[:, :coverage_count] = valid_codes[coverage % valid_codes.numel()] + packed = raw[..., 0::2] | (raw[..., 1::2] << 4) + packed_mxfp4 = packed.contiguous().view(data_dtype) + cpu_ref = FP4_E2M1_VALUES[raw.to(torch.long)] + return (packed_mxfp4.npu() if to_npu else packed_mxfp4), cpu_ref + + +def make_e8m0_zero(shape, scale_dtype): + validate_mx_scale_dtype(scale_dtype) + return torch.zeros(shape, dtype=torch.uint8).view(scale_dtype) + + +class GeneralizedQLIV2: + def __init__( + self, + batch_size, + q_seq, + k_seq, + q_t_size, + k_t_size, + q_head_num, + k_head_num, + head_dim, + block_size, + block_num, + qk_dtype, + weight_dtype, + dequant_dtype, + actual_seq_dtype, + cu_seqlens_q, + cu_seqlens_k, + act_seq_q, + act_seq_k, + cmp_residual_k, + max_seqlen_q, + quant_mode, + layout_query, + layout_key, + sparse_count, + sparse_mode, + query_datarange, + key_datarange, + weights_datarange, + q_scale_datarange, + k_scale_datarange, + cmp_ratio, + return_value, + split_s1=DEFAULT_SPLIT_S1, + s1size=DEFAULT_S1SIZE, + ): + self.batch_size = batch_size + self.q_seq = q_seq + self.k_seq = k_seq + self.q_t_size = q_t_size + self.k_t_size = k_t_size + self.q_head_num = q_head_num + self.k_head_num = k_head_num + self.group_size = q_head_num // k_head_num + self.head_dim = head_dim + self.block_size = block_size + self.block_num = block_num + self.qk_dtype = qk_dtype + self.weight_dtype = weight_dtype + self.dequant_dtype = dequant_dtype + self.actual_seq_dtype = actual_seq_dtype + self.cu_seqlens_q = cu_seqlens_q + self.cu_seqlens_k = cu_seqlens_k + self.seqused_q = act_seq_q + self.seqused_k = act_seq_k + self.act_seq_q = act_seq_q + self.act_seq_k = act_seq_k + self.cmp_residual_k = cmp_residual_k + self.max_seqlen_q = max_seqlen_q + self.quant_mode = quant_mode + self.layout_query = layout_query + self.layout_key = layout_key + self.sparse_count = sparse_count + self.sparse_mode = sparse_mode + self.cmp_ratio = cmp_ratio + self.w_dtype = weight_dtype + self.return_value = return_value + self.split_s1 = split_s1 # 是否切分S1轴 / Whether to split the S1 axis + self.s1size = s1size # S1轴切分块大小 / S1 axis chunk size + + if layout_query == "BSND": + self.q_shape = [batch_size, q_seq, q_head_num, head_dim] + self.w_shape = [batch_size, q_seq, q_head_num] + self.q_tnd_flag = 0 + elif layout_query == "TND": + self.q_shape = [q_t_size, q_head_num, head_dim] + self.w_shape = [q_t_size, q_head_num] + self.q_tnd_flag = 1 + + if layout_key == "BSND": + self.k_shape = [batch_size, k_seq, k_head_num, head_dim] + elif layout_key == "TND": + self.k_shape = [k_t_size, k_head_num, head_dim] + + if layout_query == "BSND": + self.out_shape = [batch_size, q_seq, k_head_num, sparse_count] + self.output_idx_offset_shape = [batch_size, q_seq, k_head_num] + elif layout_query == "TND": + self.out_shape = [q_t_size, k_head_num, sparse_count] + self.output_idx_offset_shape = [q_t_size, k_head_num] + + def cal_atten_bnsd(self, output_idx_offset): + batch_size = self.batch_size + qs = self.q_seq + ks = self.k_seq + n1 = self.q_head_num + n2 = self.k_head_num + cu_seqlens_q = self.cu_seqlens_q + cu_seqlens_k = self.cu_seqlens_k + seqused_q = self.seqused_q + seqused_k = self.seqused_k + cmp_residual_k = self.cmp_residual_k + q_bnsd_tensor = self.q_bnsd_tensor + k_bnsd_tensor = self.k_bnsd_tensor + wt_bnsd_tensor = self.wt_bnsd_tensor + mask_tensor = self.m_tensor + q_scale_bnsd_tensor = self.q_scale_bnsd_tensor + k_scale_bnsd_tensor = self.k_scale_bnsd_tensor + cmp_ratio = self.cmp_ratio + + out_shape_bnsd = copy.deepcopy(self.q_bnsd_shape) + out_shape_bnsd[1] = n2 + out_shape_bnsd[-1] = self.sparse_count + + out_shape_bnss = copy.deepcopy(self.q_bnsd_shape) + out_shape_bnss[1] = n2 + # out_shape_bnss[-1] = math.floor(max(actualSeqLengths_k)) + out_shape_bnss[-1] = math.floor(max(seqused_k)) if seqused_k is not None else ks + + y = torch.full(out_shape_bnsd, -1, dtype=torch.int32) + y_value = torch.full(out_shape_bnss, -float("inf"), dtype=torch.float32) + y_value_np = np.full(out_shape_bnsd, -np.inf, dtype=np.float32) + + prefix = 0 + for b_idx in range(batch_size): + if self.layout_query == "TND": + if seqused_q is not None: + curr_actualSeq_q = seqused_q[b_idx] + else: + # 已被处理为shape为(B,)的tensor + curr_actualSeq_q = cu_seqlens_q[b_idx] + elif self.layout_query == "BSND": + if seqused_q is not None: + curr_actualSeq_q = seqused_q[b_idx] + else: + curr_actualSeq_q = qs + + if self.layout_key == "TND": + if seqused_k is not None: + curr_actualSeq_k = seqused_k[b_idx] + else: + curr_actualSeq_k = cu_seqlens_k[b_idx] + elif self.layout_key == "PA_BBND": + curr_actualSeq_k = seqused_k[b_idx] + elif self.layout_key == "BSND": + if seqused_k is not None: + curr_actualSeq_k = seqused_k[b_idx] + else: + curr_actualSeq_k = ks + self.cur_actseq_q = curr_actualSeq_q + self.cur_actseq_k = curr_actualSeq_k + + self.cur_b_idx = b_idx + + if self.split_s1: + # 切分S1轴以减小中间结果内存占用 + # Split S1 axis to reduce intermediate result memory usage + num_s1_chunks = ( + math.ceil(curr_actualSeq_q / self.s1size) + if curr_actualSeq_q > 0 + else 1 + ) + for s1_chunk_idx in range(num_s1_chunks): + s1_start = s1_chunk_idx * self.s1size + s1_end = min(s1_start + self.s1size, curr_actualSeq_q) + cur_chunk_s1 = s1_end - s1_start + + self.cur_q = q_bnsd_tensor[ + b_idx : (b_idx + 1), :, s1_start:s1_end, : + ] + self.cur_k = k_bnsd_tensor[ + b_idx : (b_idx + 1), :, :curr_actualSeq_k, : + ] + self.cur_wt = wt_bnsd_tensor[ + b_idx : (b_idx + 1), :, s1_start:s1_end, : + ] + self.cur_q_scale = q_scale_bnsd_tensor[ + b_idx : (b_idx + 1), :, s1_start:s1_end, : + ] + self.cur_k_scale = k_scale_bnsd_tensor[ + b_idx : (b_idx + 1), :, :curr_actualSeq_k + ] + if self.sparse_mode != 0: + self.cur_m = mask_tensor[ + b_idx : (b_idx + 1), s1_start:s1_end, :curr_actualSeq_k + ] + + if cur_chunk_s1 != 0: + actual_selected_count = min(curr_actualSeq_k, self.sparse_count) + if is_mx_quant_mode(self.quant_mode): + ( + y[ + b_idx : (b_idx + 1), + :, + s1_start:s1_end, + :actual_selected_count, + ], + y_value[ + b_idx : (b_idx + 1), + :, + s1_start:s1_end, + :curr_actualSeq_k, + ], + ) = self.cal_atten_per_batch_mx(b_idx) + elif self.qk_dtype == torch.int8: + ( + y[ + b_idx : (b_idx + 1), + :, + s1_start:s1_end, + :actual_selected_count, + ], + y_value[ + b_idx : (b_idx + 1), + :, + s1_start:s1_end, + :curr_actualSeq_k, + ], + ) = self.cal_atten_per_batch_int8(b_idx) + elif self.qk_dtype == torch.float8_e4m3fn: + ( + y[ + b_idx : (b_idx + 1), + :, + s1_start:s1_end, + :actual_selected_count, + ], + y_value[ + b_idx : (b_idx + 1), + :, + s1_start:s1_end, + :curr_actualSeq_k, + ], + ) = self.cal_atten_per_batch_fp8(b_idx) + elif self.qk_dtype == torch.uint8: + ( + y[ + b_idx : (b_idx + 1), + :, + s1_start:s1_end, + :actual_selected_count, + ], + y_value[ + b_idx : (b_idx + 1), + :, + s1_start:s1_end, + :curr_actualSeq_k, + ], + ) = self.cal_atten_per_batch_hifp8(b_idx) + if output_idx_offset is not None: + if self.layout_query == "TND": + offset = output_idx_offset.flatten()[ + prefix + s1_start : prefix + s1_end + ].reshape(1, -1, 1) + else: + offset = output_idx_offset.flatten()[ + b_idx * qs + s1_start : b_idx * qs + s1_end + ].reshape(1, -1, 1) + offset_mask = ( + y[ + b_idx : (b_idx + 1), + :, + s1_start:s1_end, + :actual_selected_count, + ] + != -1 + ) + y[ + b_idx : (b_idx + 1), + :, + s1_start:s1_end, + :actual_selected_count, + ] += offset * offset_mask + y_value_np[ + b_idx : (b_idx + 1), + :, + s1_start:s1_end, + :actual_selected_count, + ] = -np.sort(-y_value.numpy())[ + b_idx : (b_idx + 1), + :, + s1_start:s1_end, + :actual_selected_count, + ] + y[ + b_idx : (b_idx + 1), + :, + curr_actualSeq_q:, + : min(curr_actualSeq_k, self.sparse_count), + ] = -1 + else: + self.cur_q = q_bnsd_tensor[b_idx : (b_idx + 1), :, :curr_actualSeq_q, :] + self.cur_k = k_bnsd_tensor[b_idx : (b_idx + 1), :, :curr_actualSeq_k, :] + self.cur_wt = wt_bnsd_tensor[ + b_idx : (b_idx + 1), :, :curr_actualSeq_q, : + ] + self.cur_q_scale = q_scale_bnsd_tensor[ + b_idx : (b_idx + 1), :, :curr_actualSeq_q, : + ] + self.cur_k_scale = k_scale_bnsd_tensor[ + b_idx : (b_idx + 1), :, :curr_actualSeq_k + ] + if self.sparse_mode != 0: + self.cur_m = mask_tensor[ + b_idx : (b_idx + 1), :curr_actualSeq_q, :curr_actualSeq_k + ] + + if curr_actualSeq_q != 0: + actual_selected_count = min(curr_actualSeq_k, self.sparse_count) + if is_mx_quant_mode(self.quant_mode): + ( + y[ + b_idx : (b_idx + 1), + :, + :curr_actualSeq_q, + :actual_selected_count, + ], + y_value[ + b_idx : (b_idx + 1), + :, + :curr_actualSeq_q, + :curr_actualSeq_k, + ], + ) = self.cal_atten_per_batch_mx(b_idx) + elif self.qk_dtype == torch.int8: + ( + y[ + b_idx : (b_idx + 1), + :, + :curr_actualSeq_q, + :actual_selected_count, + ], + y_value[ + b_idx : (b_idx + 1), + :, + :curr_actualSeq_q, + :curr_actualSeq_k, + ], + ) = self.cal_atten_per_batch_int8(b_idx) + elif self.qk_dtype == torch.float8_e4m3fn: + ( + y[ + b_idx : (b_idx + 1), + :, + :curr_actualSeq_q, + :actual_selected_count, + ], + y_value[ + b_idx : (b_idx + 1), + :, + :curr_actualSeq_q, + :curr_actualSeq_k, + ], + ) = self.cal_atten_per_batch_fp8(b_idx) + elif self.qk_dtype == torch.uint8: + ( + y[ + b_idx : (b_idx + 1), + :, + :curr_actualSeq_q, + :actual_selected_count, + ], + y_value[ + b_idx : (b_idx + 1), + :, + :curr_actualSeq_q, + :curr_actualSeq_k, + ], + ) = self.cal_atten_per_batch_hifp8(b_idx) + y[ + b_idx : (b_idx + 1), + :, + curr_actualSeq_q:, + :actual_selected_count, + ] = -1 + if output_idx_offset is not None: + if self.layout_query == "TND": + offset = output_idx_offset.flatten()[ + prefix : prefix + curr_actualSeq_q + ].reshape(1, -1, 1) + else: + offset = output_idx_offset.flatten()[ + b_idx * qs : b_idx * qs + curr_actualSeq_q + ].reshape(1, -1, 1) + offset_mask = ( + y[ + b_idx : (b_idx + 1), + :, + :curr_actualSeq_q, + :actual_selected_count, + ] + != -1 + ) + y[ + b_idx : (b_idx + 1), + :, + :curr_actualSeq_q, + :actual_selected_count, + ] += offset * offset_mask + y_value_np[ + b_idx : (b_idx + 1), + :, + :curr_actualSeq_q, + :actual_selected_count, + ] = -np.sort(-y_value.numpy())[ + b_idx : (b_idx + 1), + :, + :curr_actualSeq_q, + :actual_selected_count, + ] + else: + pass + if self.layout_query == "TND": + prefix += cu_seqlens_q[b_idx] + return y, y_value, y_value_np + + def trans_shape_to_bnsd( + self, + tensor, + shape, + layout, + headnums=None, + act_seq=None, + is_weights=False, + tensor_name=None, + ): + if layout in ["BSND"]: + B = shape[0] + S = shape[1] + N = shape[2] + D = 1 + if is_weights: + tensor = torch.unsqueeze(tensor, dim=-1) + else: + D = shape[3] + tensor = tensor.reshape(B, S, N, D).permute(0, 2, 1, 3) + return tensor, [B, N, S, D] + elif layout == "BSN": + print("shape", shape) + B = shape[0] + S = shape[1] + N = shape[2] + if is_weights: + D = 1 + tensor = torch.unsqueeze(tensor, dim=-1) # 补D轴 + tensor = tensor.reshape(B, S, N, D).permute(0, 2, 1, 3) + return tensor, [B, N, S, D] + else: + tensor = tensor.reshape(B, S, N).permute(0, 2, 1) + return tensor, [B, N, S] + elif layout in ["TND"]: + T = shape[0] + N = shape[1] + D = 1 + if is_weights: + tensor = torch.unsqueeze(tensor, dim=-1) + else: + D = shape[2] + B = len(act_seq) + S = max(act_seq) + new_tensor = torch.zeros((B, N, S, D), dtype=tensor.dtype) + t_start = 0 + for b_index in range(B): + act_s = act_seq[b_index] + t_end = t_start + act_s + if act_s == 0: + continue + for n_index in range(N): + new_tensor[b_index, n_index, 0:act_s, :] = tensor[ + t_start:t_end, n_index, : + ] + t_start += act_s + return new_tensor, [B, N, S, D] + elif layout == "TN": + T = shape[0] + N = shape[1] + D = 1 + B = len(act_seq) + S = max(act_seq) + new_tensor = torch.zeros((B, N, S), dtype=tensor.dtype) + t_start = 0 + for b_index in range(B): + act_s = act_seq[b_index] + t_end = t_start + act_s + if act_s == 0: + continue + for n_index in range(N): + new_tensor[b_index, n_index, 0:act_s] = tensor[ + t_start:t_end, n_index + ] + t_start += act_s + return new_tensor, [B, N, S] + else: + return tensor, shape + + def trans_tnd_actseq(self, list): + list_len = len(list) + if list_len == 0: + raise ValueError("TND情况下 act_seq需要必传") + list_new = [] + list_new.append(list[0]) + for i in range(list_len - 1): + new_item = list[i + 1] - list[i] + if new_item >= 0: + list_new.append(new_item) + else: + raise ValueError( + f"TND情况下 act_seq_len 为非递减数列 act_seq_len={list}" + ) + return list_new + + def cal_atten_per_batch_mx(self, b_idx): + cur_q = self.cur_q.to(dtype=torch.float32) + cur_k = self.cur_k.to(dtype=torch.float32) + cur_wt = self.cur_wt.to(dtype=torch.float32) + cur_q_scale = self.cur_q_scale.to(dtype=torch.float32) + cur_k_scale = self.cur_k_scale.to(dtype=torch.float32) + sparse_count = self.sparse_count + sparse_mode = self.sparse_mode + if is_mxfp4_quant_mode(self.quant_mode): + head_dim = cur_q.shape[-1] + if head_dim % MX_SCALE_GROUP_SIZE != 0: + raise ValueError( + f"MXFP4 head_dim must be divisible by {MX_SCALE_GROUP_SIZE}, but got {head_dim}" + ) + + cur_q = cur_q * cur_q_scale + cur_k = cur_k * cur_k_scale + qk_bmm_res = torch.bmm( + cur_q.squeeze(0), cur_k.permute(0, 1, 3, 2).squeeze(0) + ).unsqueeze(0) + + qk_relu_out = qk_bmm_res.clamp_min(0.0).to(torch.bfloat16) + weight_matrix = cur_wt.to(torch.bfloat16).permute(0, 2, 3, 1).squeeze(0) + qk_matrix = qk_relu_out.permute(0, 2, 1, 3).squeeze(0) + brc_vmul_matrix = reduce_mxfp4_weighted_qk(weight_matrix, qk_matrix) + brc_vmul = brc_vmul_matrix.unsqueeze(0) + else: + cur_q = cur_q * cur_q_scale + cur_k = cur_k * cur_k_scale + qk_bmm_res = torch.bmm( + cur_q.squeeze(0), cur_k.permute(0, 1, 3, 2).squeeze(0) + ).unsqueeze(0) + qk_relu_out = qk_bmm_res.to(dtype=torch.float32).clamp_min(0.0) + weight_matrix = cur_wt.permute(0, 2, 3, 1).squeeze(0) + qk_matrix = qk_relu_out.permute(0, 2, 1, 3).squeeze(0) + brc_vmul = reduce_mxfp8_weighted_qk(weight_matrix, qk_matrix).unsqueeze(0) + temp_b, temp_s1, temp_n1, temp_s2 = brc_vmul.shape + temp_n2 = self.k_head_num + actual_selected_count = min(temp_s2, sparse_count) + reduce_sum = brc_vmul.reshape(temp_b, temp_n2, temp_s1, temp_s2) + + if sparse_mode == 3: + cur_m = self.cur_m + cur_m_broadcasted = cur_m.reshape(1, 1, temp_s1, temp_s2) + cur_m_broadcasted = torch.broadcast_to( + cur_m_broadcasted, (1, temp_n2, temp_s1, temp_s2) + ) + reduce_sum[cur_m_broadcasted.to(dtype=torch.bool)] = -torch.inf + to_be_sort_ele = reduce_sum.clone().to(torch.bfloat16) + b_sorted_indices = torch.full(to_be_sort_ele.shape, -1, dtype=torch.int32) + if sparse_mode == 3: + for i in range(temp_s1): + row_mask = cur_m_broadcasted[0, 0, i, :].to(dtype=torch.bool) + true_indices = torch.where(~row_mask)[0] + row_ele = to_be_sort_ele[0, 0, i, true_indices] + indices = torch.arange(len(row_ele), device=row_ele.device) + sorted_vals, sorted_idx = torch.sort( + torch.stack([-row_ele, indices], dim=1), dim=0, stable=True + ) + b_sorted_indices[0, 0, i, true_indices] = true_indices[ + sorted_idx[:, 0] + ].to(torch.int32) + else: + for i in range(temp_s1): + row_ele = to_be_sort_ele[0, 0, i, :] + indices = torch.arange(len(row_ele), device=row_ele.device) + sorted_vals, sorted_idx = torch.sort( + torch.stack([-row_ele, indices], dim=1), dim=0, stable=True + ) + b_sorted_indices[0, 0, i, :] = sorted_idx[:, 0] + topk_indices = b_sorted_indices[..., :actual_selected_count] + return topk_indices, to_be_sort_ele + + def cal_atten_per_batch_hifp8(self, b_idx): + cur_q = self.cur_q + cur_k = self.cur_k + cur_wt = self.cur_wt.to(dtype=torch.float32) + cur_q_scale = self.cur_q_scale.to(dtype=torch.float32) + cur_k_scale = self.cur_k_scale.to(dtype=torch.float32) + sparse_count = self.sparse_count + sparse_mode = self.sparse_mode + cmp_ratio = self.cmp_ratio + cur_q = trans_hifuint8_tensor_to_float(cur_q) + cur_k = trans_hifuint8_tensor_to_float(cur_k) + qk_bmm_res = torch.bmm( + cur_q.squeeze(0), cur_k.permute(0, 1, 3, 2).squeeze(0) + ).unsqueeze(0) + cur_w = cur_wt * cur_q_scale + qk_relu_out = (qk_bmm_res.to(dtype=torch.float32)).clamp_min(0.0) + brc_vmul = torch.bmm( + cur_w.permute(0, 2, 3, 1).to(dtype=torch.float32).squeeze(0), + qk_relu_out.permute(0, 2, 1, 3).to(dtype=torch.float32).squeeze(0), + ).unsqueeze(0) + temp_b, temp_s1, temp_n1, temp_s2 = brc_vmul.shape + temp_g = self.group_size + temp_n2 = self.k_head_num + temp_b_idx = self.cur_b_idx + actual_selected_count = min(temp_s2, sparse_count) + reduce_sum = brc_vmul.reshape(temp_b, temp_n2, temp_s1, temp_s2) + reduce_sum[0, :, :, :] = reduce_sum[0, :, :, :] * cur_k_scale + + if sparse_mode == 3: + cur_m = self.cur_m + cur_m_broadcasted = cur_m.reshape(1, 1, temp_s1, temp_s2) + cur_m_broadcasted = torch.broadcast_to( + cur_m_broadcasted, (1, temp_n2, temp_s1, temp_s2) + ) + # 根据布尔矩阵置-inf + reduce_sum[cur_m_broadcasted.to(dtype=torch.bool)] = -torch.inf + to_be_sort_ele = reduce_sum.clone() + to_be_sort_ele = to_be_sort_ele.to(torch.bfloat16) + # 稳定排序 + b_sorted_indices = torch.full(to_be_sort_ele.shape, -1, dtype=torch.int32) + if sparse_mode == 3: + for i in range(temp_s1): + row_mask = cur_m_broadcasted[0, 0, i, :].to(dtype=torch.bool) + true_indices = torch.where(~row_mask)[0] + row_ele = to_be_sort_ele[0, 0, i, true_indices] + indices = torch.arange(len(row_ele), device=row_ele.device) + + sorted_vals, sorted_idx = torch.sort( + torch.stack([-row_ele, indices], dim=1), dim=0, stable=True + ) + b_sorted_indices[0, 0, i, true_indices] = true_indices[ + sorted_idx[:, 0] + ].to(torch.int32) + else: + for i in range(temp_s1): + row_ele = to_be_sort_ele[0, 0, i, :] + indices = torch.arange(len(row_ele), device=row_ele.device) + sorted_vals, sorted_idx = torch.sort( + torch.stack([-row_ele, indices], dim=1), dim=0, stable=True + ) + b_sorted_indices[0, 0, i, :] = sorted_idx[:, 0] + topk_indices = b_sorted_indices[..., :actual_selected_count] + return topk_indices, to_be_sort_ele + + def cal_atten_per_batch_fp8(self, b_idx): + cur_q = self.cur_q + cur_k = self.cur_k + cur_wt = self.cur_wt.to(dtype=torch.float32) + cur_q_scale = self.cur_q_scale.to(dtype=torch.float32) + cur_k_scale = self.cur_k_scale.to(dtype=torch.float32) + sparse_count = self.sparse_count + sparse_mode = self.sparse_mode + cmp_ratio = self.cmp_ratio + qk_bmm_res = torch.bmm( + cur_q.to(dtype=torch.float32).squeeze(0), + cur_k.to(dtype=torch.float32).permute(0, 1, 3, 2).squeeze(0), + ).unsqueeze(0) + cur_w = cur_wt * cur_q_scale + qk_relu_out = (qk_bmm_res.to(dtype=torch.float32)).clamp_min(0.0) + brc_vmul = torch.bmm( + cur_w.permute(0, 2, 3, 1).to(dtype=torch.float32).squeeze(0), + qk_relu_out.permute(0, 2, 1, 3).to(dtype=torch.float32).squeeze(0), + ).unsqueeze(0) + temp_b, temp_s1, temp_n1, temp_s2 = brc_vmul.shape + temp_g = self.group_size + temp_n2 = self.k_head_num + temp_b_idx = self.cur_b_idx + actual_selected_count = min(temp_s2, sparse_count) + reduce_sum = brc_vmul.reshape(temp_b, temp_n2, temp_s1, temp_s2) + reduce_sum[0, :, :, :] = reduce_sum[0, :, :, :] * cur_k_scale + + if sparse_mode == 3: + cur_m = self.cur_m + cur_m_broadcasted = cur_m.reshape(1, 1, temp_s1, temp_s2) + cur_m_broadcasted = torch.broadcast_to( + cur_m_broadcasted, (1, temp_n2, temp_s1, temp_s2) + ) + # 根据布尔矩阵置-inf + reduce_sum[cur_m_broadcasted.to(dtype=torch.bool)] = -torch.inf + to_be_sort_ele = reduce_sum.clone() + to_be_sort_ele = to_be_sort_ele.to(torch.bfloat16) + # 稳定排序 + b_sorted_indices = torch.full(to_be_sort_ele.shape, -1, dtype=torch.int32) + if sparse_mode == 3: + for i in range(temp_s1): + row_mask = cur_m_broadcasted[0, 0, i, :].to(dtype=torch.bool) + true_indices = torch.where(~row_mask)[0] + row_ele = to_be_sort_ele[0, 0, i, true_indices] + indices = torch.arange(len(row_ele), device=row_ele.device) + + sorted_vals, sorted_idx = torch.sort( + torch.stack([-row_ele, indices], dim=1), dim=0, stable=True + ) + b_sorted_indices[0, 0, i, true_indices] = true_indices[ + sorted_idx[:, 0] + ].to(torch.int32) + else: + for i in range(temp_s1): + row_ele = to_be_sort_ele[0, 0, i, :] + indices = torch.arange(len(row_ele), device=row_ele.device) + sorted_vals, sorted_idx = torch.sort( + torch.stack([-row_ele, indices], dim=1), dim=0, stable=True + ) + b_sorted_indices[0, 0, i, :] = sorted_idx[:, 0] + topk_indices = b_sorted_indices[..., :actual_selected_count] + return topk_indices, to_be_sort_ele + + def cal_atten_per_batch_int8(self, b_idx): + cur_q = self.cur_q + cur_k = self.cur_k + cur_wt = self.cur_wt.to(dtype=torch.float16) + cur_q_scale = self.cur_q_scale.to(dtype=torch.float16) + cur_k_scale = self.cur_k_scale.to(dtype=torch.float16) + sparse_count = self.sparse_count + sparse_mode = self.sparse_mode + cmp_ratio = self.cmp_ratio + qk_bmm_res = torch.bmm( + cur_q.to(dtype=torch.int32).squeeze(0), + cur_k.to(dtype=torch.int32).permute(0, 1, 3, 2).squeeze(0), + ).unsqueeze(0) + cur_w = cur_wt * cur_q_scale + qk_relu_out = ( + (qk_bmm_res.to(dtype=torch.float32) / 1024.0) + .clamp_min(0.0) + .to(torch.float16) + ) + brc_vmul = torch.bmm( + cur_w.permute(0, 2, 3, 1).to(dtype=torch.float32).squeeze(0), + qk_relu_out.permute(0, 2, 1, 3).to(dtype=torch.float32).squeeze(0), + ).unsqueeze(0) + temp_b, temp_s1, temp_n1, temp_s2 = brc_vmul.shape + temp_g = self.group_size + temp_n2 = self.k_head_num + temp_b_idx = self.cur_b_idx + actual_selected_count = min(temp_s2, sparse_count) + reduce_sum = brc_vmul.reshape(temp_b, temp_n2, temp_s1, temp_s2) + reduce_sum[0, :, :, :] = reduce_sum[0, :, :, :] * cur_k_scale + + if sparse_mode == 3: + cur_m = self.cur_m + cur_m_broadcasted = cur_m.reshape(1, 1, temp_s1, temp_s2) + cur_m_broadcasted = torch.broadcast_to( + cur_m_broadcasted, (1, temp_n2, temp_s1, temp_s2) + ) + # 根据布尔矩阵置-inf + reduce_sum[cur_m_broadcasted.to(dtype=torch.bool)] = -torch.inf + + to_be_sort_ele = reduce_sum.clone() + # 稳定排序 + b_sorted_indices = torch.full(to_be_sort_ele.shape, -1, dtype=torch.int32) + if sparse_mode == 3: + for i in range(temp_s1): + row_mask = cur_m_broadcasted[0, 0, i, :].to(dtype=torch.bool) + true_indices = torch.where(~row_mask)[0] + row_ele = to_be_sort_ele[0, 0, i, true_indices] + indices = torch.arange(len(row_ele), device=row_ele.device) + + sorted_vals, sorted_idx = torch.sort( + torch.stack([-row_ele, indices], dim=1), dim=0, stable=True + ) + b_sorted_indices[0, 0, i, true_indices] = true_indices[ + sorted_idx[:, 0] + ].to(torch.int32) + else: + for i in range(temp_s1): + row_ele = to_be_sort_ele[0, 0, i, :] + indices = torch.arange(len(row_ele), device=row_ele.device) + sorted_vals, sorted_idx = torch.sort( + torch.stack([-row_ele, indices], dim=1), dim=0, stable=True + ) + b_sorted_indices[0, 0, i, :] = sorted_idx[:, 0] + topk_indices = b_sorted_indices[..., :actual_selected_count] + return topk_indices, to_be_sort_ele + + def trans_bnsd_to_layout(self, tensor, shape, layout, act_q=None): + # 此时的输出D轴是K轴 + if layout == "BSH": + output = tensor.permute(0, 2, 1, 3).contiguous().view(shape) + return output + elif layout == "BSND": + output = tensor.permute(0, 2, 1, 3).contiguous() + return output + elif layout in ["BSND_NBSD", "BNSD_NBSD", "BSH_NBSD"]: + output = tensor.permute(1, 0, 2, 3).contiguous() + return output + elif layout in ["TND", "TND_NTD"]: + T = sum(act_q) + B = tensor.shape[0] + N = tensor.shape[1] + D = tensor.shape[3] + output = torch.full(size=(T, N, D), fill_value=-1, dtype=tensor.dtype) + t_start = 0 + for b_index in range(B): + act_s = act_q[b_index] + t_end = t_start + act_s + if act_s == 0: + continue + for n_index in range(N): + output[t_start:t_end, n_index, :] = tensor[ + b_index, n_index, :act_s, : + ] + t_start += act_s + if layout == "TND_NTD": + output = output.permute(1, 0, 2).contiguous() + return output + else: + return tensor + + def broadcast_n_axis(self, n1, n2, temp_tensor, input_dtype): + g = n1 // n2 + temp_shape = temp_tensor.shape + B = temp_shape[0] + S = temp_shape[2] + D = temp_shape[3] + modify_tensor = torch.zeros([B, n1, S, D], dtype=temp_tensor.dtype) + for i in range(n1): + j = i // g + modify_tensor[:, i : i + 1, :, :] = temp_tensor[:, j : j + 1, :, :] + return modify_tensor, modify_tensor.shape + + def flatten_mx_scale_tail(self, tensor, shape): + scale_head_dim = shape[-2] * shape[-1] + flat_shape = list(shape[:-2]) + [scale_head_dim] + return tensor.reshape(flat_shape), flat_shape + + def broadcast_mx_scale_d_axis(self, tensor): + output = tensor.repeat_interleave(MX_SCALE_GROUP_SIZE, dim=-1) + return output, list(output.shape) + + def create_mask(self, m_shape, act_k, S1): + atten_masks = torch.zeros(tuple(m_shape), dtype=torch.uint8) + cmp_ratio = self.cmp_ratio + tmp_pos_orig = act_k - S1 + + for i in range(S1): + if ((tmp_pos_orig + i + 1) / cmp_ratio) < 0: + atten_masks[i, :] = 1 + else: + atten_masks[i, math.floor((tmp_pos_orig + i + 1) / cmp_ratio) :] = 1 + return atten_masks + + def create_mask_right_down( + self, m_shape, actualSeqLengthsQ, actualSeqLengthsK, batch + ): + mask_s_q = m_shape[0] + mask_s_kv = m_shape[1] + cmp_ratio = self.cmp_ratio + cmp_residual_k = self.cmp_residual_k + next_tokens_list = [] + re_mask_batch = [] + pre_tokens = 214748647 + for i in range(batch): + if len(actualSeqLengthsQ) == 0: + S1 = mask_s_q + else: + S1 = actualSeqLengthsQ[i] + + if len(actualSeqLengthsK) == 0: + S2 = mask_s_kv + else: + S2 = math.floor(actualSeqLengthsK[i]) + next_tokens = S2 - S1 + next_tokens_list.append(next_tokens) + act_k = actualSeqLengthsK[i] * cmp_ratio + cmp_residual_k[i] + atten_masks = self.create_mask(m_shape, act_k, S1) + re_mask_batch.append(np.array(atten_masks, dtype=np.bool_)) + re_mask_np = np.array(re_mask_batch, dtype=np.bool_) + cpu_mask = torch.from_numpy(re_mask_np) + return cpu_mask, next_tokens_list + + def forward( + self, + query, + key, + weights, + query_dequant_scale, + key_dequant_scale, + cu_seqlens_q, + cu_seqlens_k, + seqused_q, + seqused_k, + block_table, + output_idx_offset, + ): + print("cpu执行中...") + + # 参数的初始化 + batch_size = self.batch_size + q_seq = self.q_seq + k_seq = self.k_seq + layout_query = self.layout_query + layout_key = self.layout_key + sparse_count = self.sparse_count + sparse_mode = self.sparse_mode + out_shape = self.out_shape + q_shape = self.q_shape + head_dim = self.head_dim + q_head_num = self.q_head_num + k_head_num = self.k_head_num + q_t_size = self.q_t_size + k_t_size = self.k_t_size + block_size = self.block_size + block_num = self.block_num + q_dtype = self.qk_dtype + k_dtype = self.qk_dtype + w_shape = self.w_shape + w_dtype = self.w_dtype + actual_seq_dtype = self.actual_seq_dtype + cmp_ratio = self.cmp_ratio + return_value = self.return_value + is_mx_mode = is_mx_quant_mode(self.quant_mode) + + if layout_query == "TND": + q_scale_shape = [q_t_size, q_head_num] + self.cu_seqlens_q = self.trans_tnd_actseq(cu_seqlens_q[1:]) + actualSeqLengths_q = self.cu_seqlens_q + if seqused_q is not None: + self.seqused_q = seqused_q + self.has_seqused_q = True + actualSeqLengths_q = self.seqused_q + elif layout_query == "BSND": + q_scale_shape = [batch_size, q_seq, q_head_num] + if seqused_q is not None: + self.seqused_q = seqused_q + actual_seq_lengths_query = seqused_q + self.has_seqused_q = True + else: + actual_seq_lengths_query = torch.tensor( + np.random.uniform(q_seq, q_seq, batch_size) + ).to(torch.int32) + actualSeqLengths_q = actual_seq_lengths_query + + if layout_key == "TND": + layout_key_scale = "TN" + k_scale_shape = [k_t_size, k_head_num] + self.cu_seqlens_k = self.trans_tnd_actseq(cu_seqlens_k[1:]) + actualSeqLengths_k = self.cu_seqlens_k + k_shape = self.k_shape + if seqused_k is not None: + self.seqused_k = seqused_k + self.has_seqused_k = True + actualSeqLengths_k = self.seqused_k + elif layout_key == "BSND": + layout_key_scale = "BSN" + k_shape = self.k_shape + k_scale_shape = [batch_size, k_seq, k_head_num] + if seqused_k is not None: + self.seqused_k = seqused_k + actual_seq_lengths_key = seqused_k + self.has_seqused_k = True + else: + actual_seq_lengths_key = torch.tensor( + np.random.uniform(k_seq, k_seq, batch_size) + ).to(torch.int32) + actualSeqLengths_k = actual_seq_lengths_key + + elif layout_key == "PA_BBND": + self.actual_seq_lengths_key = seqused_k + actualSeqLengths_k = self.actual_seq_lengths_key + layout_key_scale = layout_key + k_max_s2 = math.floor(max(actualSeqLengths_k)) + k_shape = [batch_size, k_head_num, k_max_s2, head_dim] + k_scale_shape = [batch_size, k_head_num, k_max_s2] + if is_mx_mode: + mx_scale_tail_shape = [head_dim // MX_SCALE_SHAPE_ALIGN, MX_SCALE_PACK_NUM] + q_scale_shape = q_scale_shape + mx_scale_tail_shape + k_scale_shape = k_scale_shape + mx_scale_tail_shape + if layout_key in ["BSND", "TND"]: + layout_key_scale = layout_key + query = query.cpu() + key = key.cpu() + weights = weights.cpu() + query_dequant_scale = query_dequant_scale.cpu() + key_dequant_scale = key_dequant_scale.cpu() + q_scale_is_weights = True + if is_mx_mode: + query_dequant_scale, q_scale_shape = self.flatten_mx_scale_tail( + query_dequant_scale, q_scale_shape + ) + key_dequant_scale, k_scale_shape = self.flatten_mx_scale_tail( + key_dequant_scale, k_scale_shape + ) + q_scale_is_weights = False + if output_idx_offset is not None: + output_idx_offset = output_idx_offset.cpu() + + # 将输入转化为BNSD + ## BSND / TND -> BNSD + if self.layout_query == "TND": + q_bnsd_tensor, q_bnsd_shape = self.trans_shape_to_bnsd( + query, q_shape, layout_query, q_head_num, self.cu_seqlens_q + ) + else: + q_bnsd_tensor, q_bnsd_shape = self.trans_shape_to_bnsd( + query, q_shape, layout_query, q_head_num, actualSeqLengths_q + ) + + ## BSND/TND/ -> BNSD + if self.layout_key == "TND": + k_bnsd_tensor, k_bnsd_shape = self.trans_shape_to_bnsd( + key, k_shape, layout_key, k_head_num, self.cu_seqlens_k + ) + k_scale_bnsd_tensor, k_scale_bnsd_shape = self.trans_shape_to_bnsd( + key_dequant_scale, + k_scale_shape, + layout_key_scale, + k_head_num, + self.cu_seqlens_k, + ) + else: + k_bnsd_tensor, k_bnsd_shape = self.trans_shape_to_bnsd( + key, + k_shape, + layout_key, + k_head_num, + torch.floor(actualSeqLengths_k).to(actual_seq_dtype), + ) + k_scale_bnsd_tensor, k_scale_bnsd_shape = self.trans_shape_to_bnsd( + key_dequant_scale, + k_scale_shape, + layout_key_scale, + k_head_num, + torch.floor(actualSeqLengths_k).to(actual_seq_dtype), + ) + + ## BSN1 -> BNS1 TN1 -> BNS1 + is_weights = True + if self.layout_query == "TND": + wt_bnsd_tensor, wt_bnsd_shape = self.trans_shape_to_bnsd( + weights, + w_shape, + layout_query, + q_head_num, + self.cu_seqlens_q, + is_weights, + ) + q_scale_bnsd_tensor, q_scale_bnsd_shape = self.trans_shape_to_bnsd( + query_dequant_scale, + q_scale_shape, + layout_query, + q_head_num, + self.cu_seqlens_q, + q_scale_is_weights, + ) + else: + wt_bnsd_tensor, wt_bnsd_shape = self.trans_shape_to_bnsd( + weights, + w_shape, + layout_query, + q_head_num, + actualSeqLengths_q, + is_weights, + ) + # BSN1 -> BNS1 + q_scale_bnsd_tensor, q_scale_bnsd_shape = self.trans_shape_to_bnsd( + query_dequant_scale, + q_scale_shape, + layout_query, + q_head_num, + actualSeqLengths_q, + q_scale_is_weights, + ) + # 将 k n2轴 广播为 n1 + if q_head_num != k_head_num: + k_bnsd_tensor, k_bnsd_shape = self.broadcast_n_axis( + q_head_num, k_head_num, k_bnsd_tensor, k_dtype + ) + if is_mx_mode: + k_scale_bnsd_tensor, k_scale_bnsd_shape = self.broadcast_n_axis( + q_head_num, + k_head_num, + k_scale_bnsd_tensor, + k_scale_bnsd_tensor.dtype, + ) + if is_mx_mode: + q_scale_bnsd_tensor, q_scale_bnsd_shape = self.broadcast_mx_scale_d_axis( + q_scale_bnsd_tensor + ) + k_scale_bnsd_tensor, k_scale_bnsd_shape = self.broadcast_mx_scale_d_axis( + k_scale_bnsd_tensor + ) + self.q_bnsd_tensor = q_bnsd_tensor + self.q_bnsd_shape = q_bnsd_shape + self.k_bnsd_tensor = k_bnsd_tensor + self.k_bnsd_shape = k_bnsd_shape + self.wt_bnsd_tensor = wt_bnsd_tensor + self.wt_bnsd_shape = wt_bnsd_shape + self.q_scale_bnsd_tensor = q_scale_bnsd_tensor + self.q_scale_bnsd_shape = q_scale_bnsd_shape + self.k_scale_bnsd_tensor = k_scale_bnsd_tensor + self.k_scale_bnsd_shape = k_scale_bnsd_shape + # 生成mask, sparse_mode=3时使能 + m_shape_std = [q_bnsd_shape[2], k_bnsd_shape[2]] # m_shape应该是[s1,s2] + batch = q_bnsd_shape[0] + m_tensor = [] + if sparse_mode == 3: + m_tensor, next_tokens_list = self.create_mask_right_down( + m_shape_std, actualSeqLengths_q, actualSeqLengths_k, batch + ) + elif sparse_mode == 0: + pass + else: + raise ValueError("unsupported sparse_mode!") + self.m_tensor = m_tensor + y, y_value, y_value_np = self.cal_atten_bnsd(output_idx_offset) + sparse_value = torch.from_numpy(y_value_np) + + # TND & PA 需要传入out_shape为BNSD + out_shape_bnsd = copy.deepcopy(self.q_bnsd_shape) + out_shape_bnsd[1] = k_head_num + out_shape_bnsd[-1] = sparse_count + if self.layout_query == "TND": + y = self.trans_bnsd_to_layout( + y, out_shape_bnsd, layout_query, self.cu_seqlens_q + ) # TODO + if return_value: + sparse_value = self.trans_bnsd_to_layout( + sparse_value, out_shape_bnsd, layout_query, self.cu_seqlens_q + ) + else: + y = self.trans_bnsd_to_layout(y, out_shape_bnsd, layout_query, q_seq) + if return_value: + sparse_value = self.trans_bnsd_to_layout( + sparse_value, out_shape_bnsd, layout_query, q_seq + ) + return y, y_value, sparse_value + + +def trans_prefix_actseq(self, list): + list_len = len(list) + if list_len == 0: + raise ValueError("PA场景下 act_seq需要必传") + list_new = [] + list_new.append(list[0]) + for i in range(list_len - 1): + new_item = list[i + 1] - list[i] + if new_item >= 0: + list_new.append(new_item) + else: + raise ValueError(f"PA场景下act seq 为非递减数列 act_seq ={list}") + return list_new + + +def qliv2_output_single( + params, + is_batch=False, + split_s1=DEFAULT_SPLIT_S1, + s1size=DEFAULT_S1SIZE, + generate_golden=True, +): + if is_batch: + params = normalize_qliv2_params(params) + ( + batch_size, + q_seq, + k_seq, + q_t_size, + k_t_size, + q_head_num, + k_head_num, + head_dim, + block_size, + block_num, + qk_dtype, + weight_dtype, + dequant_dtype, + actual_seq_dtype, + cu_seqlens_q, + cu_seqlens_k, + seqused_q, + seqused_k, + cmp_residual_k, + max_seqlen_q, + quant_mode, + layout_query, + layout_key, + sparse_count, + sparse_mode, + query_datarange, + key_datarange, + weights_datarange, + q_scale_datarange, + k_scale_datarange, + cmp_ratio, + return_value, + output_idx_offset, + ) = params + + if is_batch: + dtype_map = { + "INT8": torch.int8, + "UINT8": torch.uint8, + "HIFLOAT8": torch.uint8, + "INT32": torch.int32, + "INT64": torch.int64, + "FP16": torch.float16, + "FLOAT16": torch.float16, + "FP32": torch.float32, + "FLOAT": torch.float32, + "FLOAT32": torch.float32, + "BF16": torch.bfloat16, + "FLOAT8_E4M3FN": torch.float8_e4m3fn, + "FLOAT8_E8M0": torch.float8_e8m0fnu, + "FLOAT8_E8M0FNU": torch.float8_e8m0fnu, + "FLOAT4_E2M1": MXFP4_TORCH_DTYPE, + "FLOAT4_E2M1FN_X2": MXFP4_TORCH_DTYPE, + } + qk_dtype = dtype_map.get(qk_dtype, qk_dtype) + weight_dtype = dtype_map.get(weight_dtype, weight_dtype) + dequant_dtype = dtype_map.get(dequant_dtype, dequant_dtype) + actual_seq_dtype = dtype_map.get(actual_seq_dtype, actual_seq_dtype) + + if cu_seqlens_q is not None and isinstance(cu_seqlens_q, str): + cu_seqlens_q = ast.literal_eval(cu_seqlens_q) + if cu_seqlens_k is not None and isinstance(cu_seqlens_k, str): + cu_seqlens_k = ast.literal_eval(cu_seqlens_k) + if seqused_q is not None and isinstance(seqused_q, str): + seqused_q = ast.literal_eval(seqused_q) + if seqused_k is not None and isinstance(seqused_k, str): + seqused_k = ast.literal_eval(seqused_k) + if cmp_residual_k is not None and isinstance(cmp_residual_k, str): + cmp_residual_k = ast.literal_eval(cmp_residual_k) + if query_datarange is not None and isinstance(query_datarange, str): + query_datarange = ast.literal_eval(query_datarange) + if key_datarange is not None and isinstance(key_datarange, str): + key_datarange = ast.literal_eval(key_datarange) + if weights_datarange is not None and isinstance(weights_datarange, str): + weights_datarange = ast.literal_eval(weights_datarange) + if output_idx_offset is not None and isinstance(output_idx_offset, str): + output_idx_offset = ast.literal_eval(output_idx_offset) + output_idx_offset = [int(x) for x in output_idx_offset] + if layout_query == "TND": + output_idx_offset_size = q_t_size * 1 + else: + output_idx_offset_size = batch_size * q_seq * 1 + output_idx_offset = [ + [ + random.randint(output_idx_offset[0], output_idx_offset[1]) + for _ in range(output_idx_offset_size) + ] + for _ in range(1) + ] + if isinstance(q_scale_datarange, str): + q_scale_datarange = ast.literal_eval(q_scale_datarange) + if isinstance(k_scale_datarange, str): + k_scale_datarange = ast.literal_eval(k_scale_datarange) + + hifp8mode = 1 if quant_mode == QUANT_MODE_HIFLOAT8 else 0 + if is_mx_quant_mode(quant_mode): + validate_mx_scale_dtype(dequant_dtype) + + # ======================== 核心推导：从 cu_seqlens / seqused 推导个体长度 ======================== + # 辅助函数：从前缀和 cu_seqlens [B+1] 推导个体长度 [B] + def _cu_seqlens_to_lengths(cu_list): + return [cu_list[i + 1] - cu_list[i] for i in range(len(cu_list) - 1)] + + # Q 侧个体长度（CPU golden 用） + if layout_query == "TND": + # TND: 必传 cu_seqlens_q，从差分推导个体长度 + assert cu_seqlens_q is not None, "TND layout requires cu_seqlens_q" + lengths_q_list = _cu_seqlens_to_lengths(cu_seqlens_q) + else: + # BSND: 从 seqused_q 获取，若 None 则用 q_seq 填满 + if seqused_q is not None: + lengths_q_list = list(seqused_q) + else: + lengths_q_list = [q_seq] * batch_size + + # K 侧个体长度（CPU golden 用） + if layout_key == "TND": + # TND: 必传 cu_seqlens_k，从差分推导个体长度 + assert cu_seqlens_k is not None, "TND layout requires cu_seqlens_k" + lengths_k_list = _cu_seqlens_to_lengths(cu_seqlens_k) + elif layout_key == "PA_BBND": + # PA_BBND: 从 seqused_k 获取 + assert seqused_k is not None, f"{layout_key} layout requires seqused_k" + lengths_k_list = list(seqused_k) + else: + # BSND: 从 seqused_k 获取，若 None 则用 q_seq 填满 + if seqused_k is not None: + lengths_k_list = list(seqused_k) + else: + lengths_k_list = [k_seq] * batch_size + + # ======================== 构造 NPU 输入 tensor ======================== + # cu_seqlens tensor（仅 TND 传入） + if layout_query == "TND": + cu_seqlens_query = torch.tensor(cu_seqlens_q).to(actual_seq_dtype) + else: + cu_seqlens_query = None + + if layout_key == "TND": + cu_seqlens_key = torch.tensor(cu_seqlens_k).to(actual_seq_dtype) + else: + cu_seqlens_key = None + + # seqused tensor + if seqused_q is not None: + seqused_q_tensor = torch.tensor(seqused_q).to(actual_seq_dtype) + else: + seqused_q_tensor = None + if seqused_k is not None: + seqused_k_tensor = torch.tensor(seqused_k).to(actual_seq_dtype) + else: + seqused_k_tensor = None + + # ======================== CPU golden forward 用的 actual_seq ======================== + # TND: actual_seq 是前缀和格式，即 cu_seqlens[1:]（去掉首位 0） + # golden.forward 内部会 trans_tnd_actseq 差分为个体长度 + # BSND/PA: actual_seq 是个体长度，即 seqused + # （actual_seq始终传入，CPU golden 也需要） + if layout_query == "TND": + actual_seq_lengths_query = torch.tensor(cu_seqlens_q[1:]).to(actual_seq_dtype) + else: + actual_seq_lengths_query = torch.tensor(lengths_q_list).to(actual_seq_dtype) + + if layout_key == "TND": + actual_seq_lengths_key = torch.tensor(cu_seqlens_k[1:]).to(actual_seq_dtype) + else: + actual_seq_lengths_key = torch.tensor(lengths_k_list).to(actual_seq_dtype) + + # PA_BBND key 构造用的 act_seq_k 列表 + act_seq_k = lengths_k_list + + # 检查 cmp_residual_k 参数 + if (sparse_mode == 0 or cmp_ratio == 1) and cmp_residual_k is not None: + print( + f"Warning: sparse_mode={sparse_mode} or cmp_ratio={cmp_ratio}, " + f"cmp_residual_k={cmp_residual_k}, should be None" + ) + print("Hint: set cmp_residual_k to None when sparse_mode==0 or cmp_ratio==1") + + # cmp_residual_k for CPU golden (always a list with zeros when cmp_ratio==1 or sparse_mode==0) + if cmp_ratio == 1 or sparse_mode == 0: + cmp_residual_k_for_cpu = [0] * batch_size + else: + cmp_residual_k_for_cpu = list(cmp_residual_k) + + # cmp_residual_k for NPU (None when cmp_ratio==1 or sparse_mode==0, tensor otherwise) + if cmp_ratio == 1 or sparse_mode == 0: + cmp_residual_k_for_npu = None + else: + cmp_residual_k_for_npu = torch.tensor(cmp_residual_k).to(actual_seq_dtype) + + if cu_seqlens_q is not None: + cu_seqlens_q = torch.tensor(cu_seqlens_q).to(torch.int32) + if cu_seqlens_k is not None: + cu_seqlens_k = torch.tensor(cu_seqlens_k).to(torch.int32) + if seqused_q is not None: + seqused_q = torch.tensor(seqused_q).to(torch.int32) + if seqused_k is not None: + seqused_k = torch.tensor(seqused_k).to(torch.int32) + # ======================== 构造 GeneralizedQLIV2 用于 CPU golden ======================== + # GeneralizedQLIV2 需要 act_seq 个体长度（用于 TND→BNSD 转换等） + test_qliv2 = GeneralizedQLIV2( + batch_size, + q_seq, + k_seq, + q_t_size, + k_t_size, + q_head_num, + k_head_num, + head_dim, + block_size, + block_num, + qk_dtype, + weight_dtype, + dequant_dtype, + actual_seq_dtype, + cu_seqlens_q, + cu_seqlens_k, + lengths_q_list, + lengths_k_list, + cmp_residual_k_for_cpu, + max_seqlen_q, + quant_mode, + layout_query, + layout_key, + sparse_count, + sparse_mode, + query_datarange, + key_datarange, + weights_datarange, + q_scale_datarange, + k_scale_datarange, + cmp_ratio, + return_value, + split_s1=split_s1, + s1size=s1size, + ) + + qk_physical_head_dim = get_qk_physical_head_dim(head_dim, quant_mode) + mx_scale_tail_shape = (head_dim // MX_SCALE_SHAPE_ALIGN, MX_SCALE_PACK_NUM) + if layout_query == "BSND": + q_logical_shape = (batch_size, q_seq, q_head_num, head_dim) + q_physical_shape = (batch_size, q_seq, q_head_num, qk_physical_head_dim) + if is_mxfp4_quant_mode(quant_mode): + query, query_cpu_ref = make_mxfp4_tensor_pair( + q_logical_shape, query_datarange, qk_dtype, to_npu=False + ) + else: + query_base = torch.tensor( + np.random.uniform( + query_datarange[0], query_datarange[1], q_logical_shape + ) + ).to(torch.float) + if hifp8mode == 1: + query = trans_float_tensor_to_hifuint8( + query_base, round_mode="hybrid", over_mode=True + ) + else: + query = query_base.to(qk_dtype) + query_cpu_ref = query + + q_scale = random.uniform(q_scale_datarange[0], q_scale_datarange[1]) + if is_mx_quant_mode(quant_mode): + query_dequant_scale, query_dequant_scale_cpu = make_mx_e8m0_scale_pair( + (batch_size, q_seq, q_head_num), + mx_scale_tail_shape, + q_scale_datarange, + dequant_dtype, + to_npu=False, + ) + elif quant_mode == QUANT_MODE_HIFLOAT8: + query_dequant_scale = torch.tensor([q_scale]).to(dequant_dtype) + query_dequant_scale_cpu = torch.tensor( + np.random.uniform(q_scale, q_scale, (batch_size, q_seq, q_head_num)) + ).to(dequant_dtype) + else: + query_dequant_scale = torch.tensor( + np.random.uniform( + q_scale_datarange[0], + q_scale_datarange[1], + (batch_size, q_seq, q_head_num), + ) + ).to(dequant_dtype) + query_dequant_scale_cpu = query_dequant_scale + + weights_cpu = torch.tensor( + np.random.uniform( + weights_datarange[0], + weights_datarange[1], + (batch_size, q_seq, q_head_num), + ) + ).to(weight_dtype) + weights = weights_cpu + if output_idx_offset is not None: + output_idx_offset = ( + torch.tensor(output_idx_offset) + .reshape(batch_size, q_seq, 1) + .to(torch.int32) + ) + elif layout_query == "TND": + q_logical_shape = (q_t_size, q_head_num, head_dim) + q_physical_shape = (q_t_size, q_head_num, qk_physical_head_dim) + if is_mxfp4_quant_mode(quant_mode): + query, query_cpu_ref = make_mxfp4_tensor_pair( + q_logical_shape, query_datarange, qk_dtype, to_npu=False + ) + else: + query_base = torch.tensor( + np.random.uniform( + query_datarange[0], query_datarange[1], q_logical_shape + ) + ).to(torch.float) + if hifp8mode == 1: + query = trans_float_tensor_to_hifuint8( + query_base, round_mode="hybrid", over_mode=True + ) + else: + query = query_base.to(qk_dtype) + query_cpu_ref = query + + q_scale = random.uniform(q_scale_datarange[0], q_scale_datarange[1]) + if is_mx_quant_mode(quant_mode): + query_dequant_scale, query_dequant_scale_cpu = make_mx_e8m0_scale_pair( + (q_t_size, q_head_num), + mx_scale_tail_shape, + q_scale_datarange, + dequant_dtype, + to_npu=False, + ) + elif quant_mode == QUANT_MODE_HIFLOAT8: + query_dequant_scale = torch.tensor([q_scale]).to(dequant_dtype) + query_dequant_scale_cpu = torch.tensor( + np.random.uniform(q_scale, q_scale, (q_t_size, q_head_num)) + ).to(dequant_dtype) + else: + query_dequant_scale = torch.tensor( + np.random.uniform( + q_scale_datarange[0], q_scale_datarange[1], (q_t_size, q_head_num) + ) + ).to(dequant_dtype) + query_dequant_scale_cpu = query_dequant_scale + + weights_cpu = torch.tensor( + np.random.uniform( + weights_datarange[0], weights_datarange[1], (q_t_size, q_head_num) + ) + ).to(weight_dtype) + weights = weights_cpu + if output_idx_offset is not None: + output_idx_offset = ( + torch.tensor(output_idx_offset).reshape(q_t_size, 1).to(torch.int32) + ) + + blockFusion = None + if layout_key == "BSND": + k_logical_shape = (batch_size, k_seq, k_head_num, head_dim) + k_physical_shape = (batch_size, k_seq, k_head_num, qk_physical_head_dim) + if is_mxfp4_quant_mode(quant_mode): + key, key_cpu_ref = make_mxfp4_tensor_pair( + k_logical_shape, key_datarange, qk_dtype, to_npu=False + ) + else: + key_base = torch.tensor( + np.random.uniform(key_datarange[0], key_datarange[1], k_logical_shape) + ).to(torch.float) + if hifp8mode == 1: + key = trans_float_tensor_to_hifuint8( + key_base, round_mode="hybrid", over_mode=True + ) + else: + key = key_base.to(qk_dtype) + key_cpu_ref = key + + k_scale = random.uniform(k_scale_datarange[0], k_scale_datarange[1]) + if is_mx_quant_mode(quant_mode): + key_dequant_scale, key_dequant_scale_cpu = make_mx_e8m0_scale_pair( + (batch_size, k_seq, k_head_num), + mx_scale_tail_shape, + k_scale_datarange, + dequant_dtype, + to_npu=False, + ) + elif quant_mode == QUANT_MODE_HIFLOAT8: + key_dequant_scale = torch.tensor([k_scale]).to(dequant_dtype) + key_dequant_scale_cpu = torch.tensor( + np.random.uniform(k_scale, k_scale, (batch_size, k_seq, k_head_num)) + ).to(dequant_dtype) + else: + key_dequant_scale = torch.tensor( + np.random.uniform( + k_scale_datarange[0], + k_scale_datarange[1], + (batch_size, k_seq, k_head_num), + ) + ).to(dequant_dtype) + key_dequant_scale_cpu = key_dequant_scale + + block_table = None + if generate_golden: + cpu_result, topk_value, cpu_topk_value = test_qliv2.forward( + query_cpu_ref, + key_cpu_ref, + weights, + query_dequant_scale_cpu, + key_dequant_scale_cpu, + cu_seqlens_q, + cu_seqlens_k, + seqused_q, + seqused_k, + block_table, + output_idx_offset, + ) + else: + cpu_result, topk_value, cpu_topk_value = None, None, None + + elif layout_key == "TND": + k_logical_shape = (k_t_size, k_head_num, head_dim) + k_physical_shape = (k_t_size, k_head_num, qk_physical_head_dim) + if is_mxfp4_quant_mode(quant_mode): + key, key_cpu_ref = make_mxfp4_tensor_pair( + k_logical_shape, key_datarange, qk_dtype, to_npu=False + ) + else: + key_base = torch.tensor( + np.random.uniform(key_datarange[0], key_datarange[1], k_logical_shape) + ).to(torch.float) + if hifp8mode == 1: + key = trans_float_tensor_to_hifuint8( + key_base, round_mode="hybrid", over_mode=True + ) + else: + key = key_base.to(qk_dtype) + key_cpu_ref = key + + k_scale = random.uniform(k_scale_datarange[0], k_scale_datarange[1]) + if is_mx_quant_mode(quant_mode): + key_dequant_scale, key_dequant_scale_cpu = make_mx_e8m0_scale_pair( + (k_t_size, k_head_num), + mx_scale_tail_shape, + k_scale_datarange, + dequant_dtype, + to_npu=False, + ) + elif quant_mode == QUANT_MODE_HIFLOAT8: + key_dequant_scale = torch.tensor([k_scale]).to(dequant_dtype) + key_dequant_scale_cpu = torch.tensor( + np.random.uniform(k_scale, k_scale, (k_t_size, k_head_num)) + ).to(dequant_dtype) + else: + key_dequant_scale = torch.tensor( + np.random.uniform( + k_scale_datarange[0], k_scale_datarange[1], (k_t_size, k_head_num) + ) + ).to(dequant_dtype) + key_dequant_scale_cpu = key_dequant_scale + + block_table = None + if generate_golden: + cpu_result, topk_value, cpu_topk_value = test_qliv2.forward( + query_cpu_ref, + key_cpu_ref, + weights, + query_dequant_scale_cpu, + key_dequant_scale_cpu, + cu_seqlens_q, + cu_seqlens_k, + seqused_q, + seqused_k, + block_table, + output_idx_offset, + ) + else: + cpu_result, topk_value, cpu_topk_value = None, None, None + + elif layout_key == "PA_BBND": + k_max_s2 = math.floor(max(act_seq_k)) + k_max_block_num_per_batch = math.ceil(k_max_s2 / block_size) + k_logical_shape = (batch_size, k_head_num, k_max_s2, head_dim) + k_physical_shape = (batch_size, k_head_num, k_max_s2, qk_physical_head_dim) + if is_mxfp4_quant_mode(quant_mode): + key_bnsd, key_bnsd_cpu_ref = make_mxfp4_tensor_pair( + k_logical_shape, key_datarange, qk_dtype, to_npu=False + ) + else: + key_bnsd_base = torch.tensor( + np.random.uniform(key_datarange[0], key_datarange[1], k_logical_shape) + ).to(torch.float) + if hifp8mode == 1: + key_bnsd = trans_float_tensor_to_hifuint8( + key_bnsd_base, round_mode="hybrid", over_mode=True + ) + else: + key_bnsd = key_bnsd_base.to(qk_dtype) + key_bnsd_cpu_ref = key_bnsd + + k_scale = random.uniform(k_scale_datarange[0], k_scale_datarange[1]) + if is_mx_quant_mode(quant_mode): + key_dequant_scale_bns_mx, key_dequant_scale_bns = make_mx_e8m0_scale_pair( + (batch_size, k_head_num, k_max_s2), + mx_scale_tail_shape, + k_scale_datarange, + dequant_dtype, + to_npu=False, + ) + elif quant_mode == QUANT_MODE_HIFLOAT8: + key_dequant_scale_bns = torch.tensor( + np.random.uniform(k_scale, k_scale, (batch_size, k_head_num, k_max_s2)) + ).to(dequant_dtype) + else: + key_dequant_scale_bns = torch.tensor( + np.random.uniform( + k_scale_datarange[0], + k_scale_datarange[1], + (batch_size, k_head_num, k_max_s2), + ) + ).to(dequant_dtype) + + key_block_num_per_batch = [] + key_block_num_sum = 0 + for cur_act_k in act_seq_k: + cur_cmp_act_k = math.floor(cur_act_k) + cur_key_block_num = math.ceil(cur_cmp_act_k / block_size) + key_block_num_per_batch.append(cur_key_block_num) + key_block_num_sum += cur_key_block_num + if block_num < key_block_num_sum: + raise ValueError("key actual block num < needed block num") + + block_id_list = np.arange(block_num) + block_id_list = np.random.permutation(block_id_list).astype(np.int32) + cur_block_id = 0 + block_table = np.full( + (batch_size, k_max_block_num_per_batch), fill_value=-1, dtype=np.int32 + ) + batch_idx = 0 + for cur_block_id_threshold in key_block_num_per_batch: + for i_block_id in range(cur_block_id_threshold): + block_table[batch_idx][i_block_id] = block_id_list[cur_block_id] + cur_block_id += 1 + batch_idx += 1 + + if is_mxfp4_quant_mode(quant_mode): + # FP4 shell dtype仅用于接口语义；PA重排按其底层packed uint8字节完成。 + key_storage_dtype = torch.uint8 + key_bnsd_storage = key_bnsd.view(torch.uint8) + else: + key_storage_dtype = qk_dtype + key_bnsd_storage = key_bnsd + key_expand = torch.zeros( + ( + batch_size, + k_head_num, + k_max_block_num_per_batch * block_size, + qk_physical_head_dim, + ), + dtype=key_storage_dtype, + ) + key_expand[:, :, :k_max_s2, :] = key_bnsd_storage + key = torch.zeros( + (block_num, block_size, k_head_num, qk_physical_head_dim), + dtype=key_storage_dtype, + ) + for i_batch in range(batch_size): + for i_block, cur_block_id in enumerate(block_table[i_batch]): + block_start_pos = i_block * block_size + if cur_block_id == -1: + continue + else: + for i_n in range(k_head_num): + key[cur_block_id, :, i_n, :] = key_expand[ + i_batch, + i_n, + block_start_pos : block_start_pos + block_size, + :, + ] + + if is_mx_quant_mode(quant_mode): + key_dequant_scale_expand = make_e8m0_zero( + ( + batch_size, + k_head_num, + k_max_block_num_per_batch * block_size, + *mx_scale_tail_shape, + ), + dequant_dtype, + ) + key_dequant_scale_expand[:, :, :k_max_s2, :, :] = key_dequant_scale_bns_mx + key_dequant_scale_block = make_e8m0_zero( + (block_num, block_size, k_head_num, *mx_scale_tail_shape), dequant_dtype + ) + for i_batch in range(batch_size): + for i_block, cur_block_id in enumerate(block_table[i_batch]): + block_start_pos = i_block * block_size + if cur_block_id == -1: + continue + else: + for i_n in range(k_head_num): + key_dequant_scale_block[cur_block_id, :, i_n, :, :] = ( + key_dequant_scale_expand[ + i_batch, + i_n, + block_start_pos : block_start_pos + block_size, + :, + :, + ] + ) + else: + key_dequant_scale_expand = torch.zeros( + (batch_size, k_head_num, k_max_block_num_per_batch * block_size), + dtype=dequant_dtype, + ) + key_dequant_scale_expand[:, :, :k_max_s2] = key_dequant_scale_bns + key_dequant_scale_block = torch.zeros( + (block_num, block_size, k_head_num), dtype=dequant_dtype + ) + for i_batch in range(batch_size): + for i_block, cur_block_id in enumerate(block_table[i_batch]): + block_start_pos = i_block * block_size + if cur_block_id == -1: + continue + else: + for i_n in range(k_head_num): + key_dequant_scale_block[cur_block_id, :, i_n] = ( + key_dequant_scale_expand[ + i_batch, + i_n, + block_start_pos : block_start_pos + block_size, + ] + ) + + if ( + not is_mx_quant_mode(quant_mode) + and quant_mode != QUANT_MODE_HIFLOAT8 + and DISCONTINUOUS_KEYS + ): + bytes_per_token = ( + head_dim + key_dequant_scale_block.element_size() // key.element_size() + ) + blockFusion = torch.zeros( + (block_num, block_size * k_head_num * bytes_per_token), dtype=qk_dtype + ) + key_flat = key.view(block_num, block_size * k_head_num * head_dim) + scale_flat = key_dequant_scale_block.view( + block_num, block_size * k_head_num + ).view(qk_dtype) + blockFusion[:, : block_size * k_head_num * head_dim] = key_flat + blockFusion[:, block_size * k_head_num * head_dim :] = scale_flat + + if is_mxfp4_quant_mode(quant_mode): + key = key.view(qk_dtype) + if quant_mode == QUANT_MODE_HIFLOAT8: + key_dequant_scale = torch.tensor([k_scale]).to(dequant_dtype) + else: + key_dequant_scale = key_dequant_scale_block + if generate_golden: + cpu_result, topk_value, cpu_topk_value = test_qliv2.forward( + query_cpu_ref, + key_bnsd_cpu_ref, + weights, + query_dequant_scale_cpu, + key_dequant_scale_bns, + cu_seqlens_q, + cu_seqlens_k, + seqused_q, + seqused_k, + block_table, + output_idx_offset, + ) + else: + cpu_result, topk_value, cpu_topk_value = None, None, None + block_table = torch.from_numpy(block_table).to(dtype=torch.int32) + # ======================== metadata 构造 ======================== + # max_seqlen 从个体长度中取 + max_seqlen_q_meta = actual_seq_lengths_query.max().item() + max_seqlen_k_meta = actual_seq_lengths_key.max().item() + + if is_batch: + if qk_dtype == torch.float8_e4m3fn: + query = query.to(dtype=torch.float16) + key = key.to(dtype=torch.float16) + if blockFusion is not None: + blockFusion = blockFusion.view(torch.uint8) + + golden_key = key_bnsd_cpu_ref if layout_key == "PA_BBND" else key_cpu_ref + golden_key_scale = ( + key_dequant_scale_bns if layout_key == "PA_BBND" else key_dequant_scale_cpu + ) + golden_block_table = block_table + if torch.is_tensor(golden_block_table): + golden_block_table = golden_block_table.detach().cpu().numpy() + golden_state = { + "model_args": ( + batch_size, + q_seq, + k_seq, + q_t_size, + k_t_size, + q_head_num, + k_head_num, + head_dim, + block_size, + block_num, + qk_dtype, + weight_dtype, + dequant_dtype, + actual_seq_dtype, + cu_seqlens_q, + cu_seqlens_k, + lengths_q_list, + lengths_k_list, + cmp_residual_k_for_cpu, + max_seqlen_q, + quant_mode, + layout_query, + layout_key, + sparse_count, + sparse_mode, + query_datarange, + key_datarange, + weights_datarange, + q_scale_datarange, + k_scale_datarange, + cmp_ratio, + return_value, + ), + "split_s1": split_s1, + "s1size": s1size, + "forward_inputs": { + "query": query_cpu_ref, + "key": golden_key, + "weights": weights, + "query_dequant_scale": query_dequant_scale_cpu, + "key_dequant_scale": golden_key_scale, + "cu_seqlens_q": cu_seqlens_q, + "cu_seqlens_k": cu_seqlens_k, + "seqused_q": seqused_q, + "seqused_k": seqused_k, + "block_table": golden_block_table, + "output_idx_offset": output_idx_offset, + }, + } + + output_tensors = { + "params": params, + "cpu_result": cpu_result, + "topk_value": topk_value, + "cpu_topk_value": cpu_topk_value, + "query": query, + "key": key, + "weights": weights, + "query_dequant_scale": query_dequant_scale, + "key_dequant_scale": key_dequant_scale, + "blockFusion": blockFusion, + "cu_seqlens_query": cu_seqlens_query, + "cu_seqlens_key": cu_seqlens_key, + "seqused_q": seqused_q_tensor, + "seqused_k": seqused_k_tensor, + "output_idx_offset": output_idx_offset, + "actual_seq_lengths_query": actual_seq_lengths_query, + "actual_seq_lengths_key": actual_seq_lengths_key, + "cmp_residual_k_for_npu": cmp_residual_k_for_npu, + "block_table": block_table, + "max_seqlen_q_meta": max_seqlen_q_meta, + "max_seqlen_k_meta": max_seqlen_k_meta, + "quant_mode": quant_mode, + "layout_query": layout_query, + "layout_key": layout_key, + "sparse_count": sparse_count, + "sparse_mode": sparse_mode, + "cmp_ratio": cmp_ratio, + "golden_state": golden_state, + } + return output_tensors + else: + metadata = torch.ops.cann_ops_transformer.quant_lightning_indexer_metadata( + cu_seqlens_q=cu_seqlens_query.npu() + if cu_seqlens_query is not None + else None, + cu_seqlens_k=cu_seqlens_key.npu() if cu_seqlens_key is not None else None, + seqused_q=seqused_q_tensor.npu() if seqused_q_tensor is not None else None, + seqused_k=seqused_k_tensor.npu() if seqused_k_tensor is not None else None, + cmp_residual_k=cmp_residual_k_for_npu.npu() + if cmp_residual_k_for_npu is not None + else None, + batch_size=batch_size, + max_seqlen_q=max_seqlen_q_meta, + max_seqlen_k=max_seqlen_k_meta, + num_heads_q=q_head_num, + num_heads_k=k_head_num, + head_dim=head_dim, + topk=sparse_count, + quant_mode=quant_mode, + mask_mode=sparse_mode, + layout_q=layout_query, + layout_k=layout_key, + cmp_ratio=cmp_ratio, + ) + metadata = metadata.npu() + if blockFusion is not None: + blockFusion = blockFusion.npu() + key = blockFusion[:, : block_size * k_head_num * head_dim].view( + block_num, block_size, k_head_num, head_dim + ) + key_dequant_scale_block = ( + blockFusion[:, block_size * k_head_num * head_dim :] + .view(dequant_dtype) + .view(block_num, block_size, k_head_num) + ) + else: + key = key.npu() + if quant_mode == QUANT_MODE_HIFLOAT8: + key_dequant_scale = torch.tensor([k_scale]).to(dequant_dtype).npu() + else: + if layout_key == "PA_BBND": + key_dequant_scale = key_dequant_scale_block.npu() + else: + key_dequant_scale = key_dequant_scale.npu() + + npu_result, npu_topk_value = ( + torch.ops.cann_ops_transformer.quant_lightning_indexer( + query.npu(), + key, + weights.npu(), + query_dequant_scale.npu(), + key_dequant_scale, + cu_seqlens_q=cu_seqlens_query.npu() + if cu_seqlens_query is not None + else None, + cu_seqlens_k=cu_seqlens_key.npu() + if cu_seqlens_key is not None + else None, + seqused_q=seqused_q_tensor.npu() + if seqused_q_tensor is not None + else None, + seqused_k=seqused_k_tensor.npu() + if seqused_k_tensor is not None + else None, + cmp_residual_k=cmp_residual_k_for_npu.npu() + if cmp_residual_k_for_npu is not None + else None, + output_idx_offset=output_idx_offset.npu() + if output_idx_offset is not None + else None, + max_seqlen_q=max_seqlen_q, + block_table=block_table.npu() if block_table is not None else None, + metadata=metadata, + quant_mode=quant_mode, + layout_q=layout_query, + layout_k=layout_key, + topk=sparse_count, + mask_mode=sparse_mode, + cmp_ratio=cmp_ratio, + return_value=return_value, + ) + ) + + torch.npu.synchronize() + if return_value: + if npu_topk_value.shape != npu_result.shape: + raise RuntimeError( + "sparse_values and sparse_indices must have the same shape when return_value=1, " + f"but got {tuple(npu_topk_value.shape)} and {tuple(npu_result.shape)}" + ) + npu_topk_value, npu_sort_order = npu_topk_value.sort( + dim=-1, descending=True + ) + npu_result = torch.gather(npu_result, dim=-1, index=npu_sort_order) + return cpu_result, npu_result, topk_value, cpu_topk_value, npu_topk_value + + +def build_qliv2_metadata_input(params, tensor_dict): + """Return the canonical metadata arguments produced by the pytest case.""" + return { + "num_heads_q": int(params[5]), + "num_heads_k": int(params[6]), + "head_dim": int(params[7]), + "topk": int(tensor_dict["sparse_count"]), + "quant_mode": int(tensor_dict["quant_mode"]), + "cu_seqlens_q": tensor_dict["cu_seqlens_query"], + "cu_seqlens_k": tensor_dict["cu_seqlens_key"], + "seqused_q": tensor_dict["seqused_q"], + "seqused_k": tensor_dict["seqused_k"], + "cmp_residual_k": tensor_dict["cmp_residual_k_for_npu"], + "batch_size": int(params[0]), + "max_seqlen_q": int(tensor_dict["max_seqlen_q_meta"]), + "max_seqlen_k": int(tensor_dict["max_seqlen_k_meta"]), + "layout_q": tensor_dict["layout_query"], + "layout_k": tensor_dict["layout_key"], + "mask_mode": int(tensor_dict["sparse_mode"]), + "cmp_ratio": int(tensor_dict["cmp_ratio"]), + } + + +def generate_qliv2_test_data( + params, + split_s1=DEFAULT_SPLIT_S1, + s1size=DEFAULT_S1SIZE, + generate_golden=True, +): + """Generate QLI_V2 inputs and optionally materialize the CPU Golden.""" + data = qliv2_output_single( + params, + is_batch=True, + split_s1=split_s1, + s1size=s1size, + generate_golden=generate_golden, + ) + data["metadata_input"] = build_qliv2_metadata_input(data["params"], data) + return data + + +def generate_cpu_golden(input_data): + """Calculate CPU Golden from input-stage data without rerunning random generation.""" + state = input_data["golden_state"] + test_qliv2 = GeneralizedQLIV2( + *state["model_args"], + split_s1=state["split_s1"], + s1size=state["s1size"], + ) + values = state["forward_inputs"] + cpu_result, topk_value, cpu_topk_value = test_qliv2.forward( + values["query"], + values["key"], + values["weights"], + values["query_dequant_scale"], + values["key_dequant_scale"], + values["cu_seqlens_q"], + values["cu_seqlens_k"], + values["seqused_q"], + values["seqused_k"], + values["block_table"], + values["output_idx_offset"], + ) + state["cpu_result"] = cpu_result + state["topk_value"] = topk_value + state["cpu_topk_value"] = cpu_topk_value + input_data["cpu_result"] = cpu_result + input_data["topk_value"] = topk_value + input_data["cpu_topk_value"] = cpu_topk_value + return cpu_result, topk_value, cpu_topk_value + + +def fp32_ta_round_to_hif8(fraction32_int, hif8_bits_num, exponent): + if exponent == HIF8_EXP_ZERO_THRESHOLD: + return True, 0 + hif8_value_tmp = fraction32_int >> (FP32_FRACTION_BITS - (hif8_bits_num + 1)) + if hif8_value_tmp == pow(2, hif8_bits_num + 1) - 1: + return True, 0 + elif hif8_value_tmp == 0: + return False, 0 + elif hif8_value_tmp % 2 == 1: + hif8_value_tmp += 1 + return False, hif8_value_tmp >> 1 + else: + return False, hif8_value_tmp >> 1 + + +def fp32_ssr_round_to_hif8(fraction32_int, hif8_bits_num, exponent): + t14_mask = SSR_T14_MASK + if exponent == HIF8_EXP_ZERO_THRESHOLD: + f14_values = (fraction32_int >> SSR_DML_SHIFT) + SSR_F14_OFFSET + t14_values = fraction32_int & t14_mask + hif8_value = 0 + else: + hif8_value = fraction32_int >> (FP32_FRACTION_BITS - hif8_bits_num) + f14_t14 = fraction32_int - (hif8_value << (FP32_FRACTION_BITS - hif8_bits_num)) + f14_values = f14_t14 >> (FP32_FRACTION_BITS - hif8_bits_num - SSR_RESERVED_BITS) + t14_values = f14_t14 & t14_mask + if f14_values >= t14_values: + if hif8_value == pow(2, hif8_bits_num) - 1: + return True, 0 + else: + hif8_value += 1 + return False, hif8_value + else: + return False, hif8_value + + +def get_hif8_fraction_bits_number(exponent): + if exponent < HIF8_EXP_DML_MIN: + return HIF8_DOT_INVALID, HIF8_EXP_BITS_DML, HIF8_FRAC_BITS_DML + if HIF8_EXP_DML_MIN <= exponent < HIF8_EXP_DML_MAX: + return HIF8_DOT_DML, HIF8_EXP_BITS_DML, HIF8_FRAC_BITS_DML + if exponent == HIF8_EXP_D0: + return HIF8_DOT_D0, HIF8_EXP_BITS_D0, HIF8_FRAC_BITS_D0 + if abs(exponent) == HIF8_EXP_D1_BOUNDARY: + return HIF8_DOT_D1, HIF8_EXP_BITS_D1, HIF8_FRAC_BITS_D1 + if HIF8_EXP_D2_MIN <= abs(exponent) <= HIF8_EXP_D2_MAX: + return HIF8_DOT_D2, HIF8_EXP_BITS_D2, HIF8_FRAC_BITS_D2 + if HIF8_EXP_D3_MIN <= abs(exponent) <= HIF8_EXP_D3_MAX: + return HIF8_DOT_D3, HIF8_EXP_BITS_D3, HIF8_FRAC_BITS_D3 + if HIF8_EXP_D4_MIN <= abs(exponent) <= HIF8_EXP_D4_MAX: + return HIF8_DOT_D4, HIF8_EXP_BITS_D4, HIF8_FRAC_BITS_D4 + if exponent > HIF8_EXP_D4_MAX: + return HIF8_DOT_D4, HIF8_EXP_BITS_D4, HIF8_DOT_INVALID + + +def cvt_float32_to_hifuint8(x, round_mode="round", over_mode=True): + sign = False + sign_int_value = 0 + x_abs = math.fabs(x) + ec = 0 + over_value = HIF8_OVERFLOW_SCALE * pow(2.0, HIF8_EXP_D4_MAX + ec) + if x < 0.0: + sign = True + sign_int_value = HIF8_SIGN_MASK + if torch.isinf(x) or x_abs >= over_value: + if sign: + if over_mode: + return HIF8_NEG_INF + else: + return HIF8_NEG_MAX + else: + if over_mode: + return HIF8_POS_INF + else: + return HIF8_POS_MAX + if torch.isnan(x): + if over_mode: + return HIF8_NAN + else: + return 0 + if x_abs == 0.0: + return 0 + exponent = math.floor(math.log2(x_abs)) + if round_mode == "hybrid": + if abs(exponent) < HYBRID_ROUND_EXP_THRESHOLD: + cut_bit_type = "TA" + else: + cut_bit_type = "SSR" + elif round_mode == "round": + cut_bit_type = "TA" + elif round_mode == "storound": + cut_bit_type = "SSR" + else: + cut_bit_type = "TA" + fraction_int = int( + x_abs * pow(2, FP32_FRACTION_BITS) * pow(2, -exponent) + - pow(2, FP32_FRACTION_BITS) + ) + dot_hif8_value, exponent_hif8_bits, fraction_hif8_bits = ( + get_hif8_fraction_bits_number(exponent) + ) + if cut_bit_type == "TA": + carry_exp_status, hif8_frac_value = fp32_ta_round_to_hif8( + fraction_int, fraction_hif8_bits, exponent + ) + elif cut_bit_type == "SSR": + carry_exp_status, hif8_frac_value = fp32_ssr_round_to_hif8( + fraction_int, fraction_hif8_bits, exponent + ) + else: + print("unknown round type") + return 0 + + if carry_exp_status: + exponent += 1 + dot_hif8_value, exponent_hif8_bits, fraction_hif8_bits_new = ( + get_hif8_fraction_bits_number(exponent) + ) + fraction_hif8_bits = fraction_hif8_bits_new + if exponent < HIF8_EXP_ZERO_THRESHOLD: + return 0 + if exponent < 0: + sig_exp = 1 + else: + sig_exp = 0 + if dot_hif8_value <= 0: + if exponent <= HIF8_EXP_ZERO_THRESHOLD: + return 0 + else: + return sign_int_value + exponent + HIF8_DML_EXP_OFFSET + elif dot_hif8_value == 1: + dot_int_value = dot_hif8_value << HIF8_DOT_BIT_SHIFT + hif8_int_value = sign_int_value + dot_int_value + hif8_frac_value + else: + abs_exponent = abs(exponent) + abs_exponent = abs_exponent - pow(2, exponent_hif8_bits - 1) + exponent_int_value = abs_exponent << fraction_hif8_bits + sig_exp = sig_exp << (exponent_hif8_bits - 1 + fraction_hif8_bits) + dot_int_value = dot_hif8_value << HIF8_DOT_BIT_SHIFT + hif8_int_value = ( + sign_int_value + + dot_int_value + + sig_exp + + exponent_int_value + + hif8_frac_value + ) + return hif8_int_value + + +def trans_float_tensor_to_hifuint8(in_tensor, round_mode="round", over_mode=True): + """ + 通过向量操作，将 float32 Tensor 批量转换为 HiF8 编码的 uint8 Tensor + """ + + shape = in_tensor.shape + x = in_tensor.reshape(-1).to(torch.float32) + + # 先用int32作为输出类型，避免出现赋值错误 + out = torch.zeros_like(x, dtype=torch.int32) + + # 1. 符号位与绝对值提取 + sign_mask = x < 0.0 + sign_int_value = torch.where(sign_mask, HIF8_SIGN_MASK, 0) + x_abs = torch.abs(x) + + # 2. 溢出与边界条件判断 (Masks) + over_value = HIF8_OVERFLOW_SCALE * (2.0**HIF8_EXP_D4_MAX) + mask_inf_or_over = torch.isinf(x) | (x_abs >= over_value) + mask_nan = torch.isnan(x) + mask_zero = x_abs == 0.0 + + # 处理特殊边界填值 + if over_mode: + out = torch.where( + mask_inf_or_over, torch.where(sign_mask, HIF8_NEG_INF, HIF8_POS_INF), out + ) + out = torch.where(mask_nan, HIF8_NAN, out) + else: + out = torch.where( + mask_inf_or_over, torch.where(sign_mask, HIF8_NEG_MAX, HIF8_POS_MAX), out + ) + out = torch.where(mask_nan, 0, out) + out = torch.where(mask_zero, 0, out) + + # 提取正常数字的 Mask + mask_normal = ~(mask_inf_or_over | mask_nan | mask_zero) + if not mask_normal.any(): + return out.reshape(shape).to(torch.uint8) + + x_norm = x_abs[mask_normal] + sign_norm = sign_int_value[mask_normal] + + # 计算基本指数 + exponent = torch.floor(torch.log2(x_norm)).to(torch.int32) + + # 确定截断模式 (TA / SSR) + if round_mode == "hybrid": + cut_bit_is_ta = torch.abs(exponent) < HYBRID_ROUND_EXP_THRESHOLD + elif round_mode == "round": + cut_bit_is_ta = torch.ones_like(exponent, dtype=torch.bool) + elif round_mode == "storound": + cut_bit_is_ta = torch.zeros_like(exponent, dtype=torch.bool) + else: + cut_bit_is_ta = torch.ones_like(exponent, dtype=torch.bool) + + # 计算 fraction_int + fraction_int = ( + x_norm * (2.0**FP32_FRACTION_BITS) * torch.pow(2.0, -exponent.float()) + - (2.0**FP32_FRACTION_BITS) + ).to(torch.int32) + + # 批量获取档位属性 (根据 exponent 映射) + abs_exp = torch.abs(exponent) + + dot = torch.full_like(exponent, HIF8_DOT_INVALID) + exp_bits = torch.zeros_like(exponent) + frac_bits = torch.zeros_like(exponent) + + # 条件区间映射 + m1 = exponent < HIF8_EXP_DML_MIN + dot = torch.where(m1, HIF8_DOT_INVALID, dot) + exp_bits = torch.where(m1, HIF8_EXP_BITS_DML, exp_bits) + frac_bits = torch.where(m1, HIF8_FRAC_BITS_DML, frac_bits) + + m2 = (~m1) & (exponent >= HIF8_EXP_DML_MIN) & (exponent < HIF8_EXP_DML_MAX) + dot = torch.where(m2, HIF8_DOT_DML, dot) + exp_bits = torch.where(m2, HIF8_EXP_BITS_DML, exp_bits) + frac_bits = torch.where(m2, HIF8_FRAC_BITS_DML, frac_bits) + + m3 = exponent == HIF8_EXP_D0 + dot = torch.where(m3, HIF8_DOT_D0, dot) + exp_bits = torch.where(m3, HIF8_EXP_BITS_D0, exp_bits) + frac_bits = torch.where(m3, HIF8_FRAC_BITS_D0, frac_bits) + + m4 = abs_exp == HIF8_EXP_D1_BOUNDARY + dot = torch.where(m4, HIF8_DOT_D1, dot) + exp_bits = torch.where(m4, HIF8_EXP_BITS_D1, exp_bits) + frac_bits = torch.where(m4, HIF8_FRAC_BITS_D1, frac_bits) + + m5 = (abs_exp >= HIF8_EXP_D2_MIN) & (abs_exp <= HIF8_EXP_D2_MAX) + dot = torch.where(m5, HIF8_DOT_D2, dot) + exp_bits = torch.where(m5, HIF8_EXP_BITS_D2, exp_bits) + frac_bits = torch.where(m5, HIF8_FRAC_BITS_D2, frac_bits) + + m6 = (abs_exp >= HIF8_EXP_D3_MIN) & (abs_exp <= HIF8_EXP_D3_MAX) + dot = torch.where(m6, HIF8_DOT_D3, dot) + exp_bits = torch.where(m6, HIF8_EXP_BITS_D3, exp_bits) + frac_bits = torch.where(m6, HIF8_FRAC_BITS_D3, frac_bits) + + m7 = (abs_exp >= HIF8_EXP_D4_MIN) & (abs_exp <= HIF8_EXP_D4_MAX) + dot = torch.where(m7, HIF8_DOT_D4, dot) + exp_bits = torch.where(m7, HIF8_EXP_BITS_D4, exp_bits) + frac_bits = torch.where(m7, HIF8_FRAC_BITS_D4, frac_bits) + + m8 = exponent > HIF8_EXP_D4_MAX + dot = torch.where(m8, HIF8_DOT_D4, dot) + exp_bits = torch.where(m8, HIF8_EXP_BITS_D4, exp_bits) + frac_bits = torch.where(m8, HIF8_DOT_INVALID, frac_bits) + + # ------------------ TA 舍入分支 ------------------ + carry_ta = torch.zeros_like(exponent, dtype=torch.bool) + frac_val_ta = torch.zeros_like(exponent) + + m_zero_thresh = exponent == HIF8_EXP_ZERO_THRESHOLD + carry_ta = torch.where(m_zero_thresh, True, carry_ta) + + m_ta_norm = ~m_zero_thresh + shift_bits = torch.clamp(FP32_FRACTION_BITS - (frac_bits + 1), min=0) + hif8_val_tmp = fraction_int >> shift_bits + + pow_frac = torch.pow(2, frac_bits + 1) - 1 + m_carry = m_ta_norm & (hif8_val_tmp == pow_frac) + carry_ta = torch.where(m_carry, True, carry_ta) + + m_odd = m_ta_norm & (~m_carry) & (hif8_val_tmp != 0) & (hif8_val_tmp % 2 == 1) + frac_val_ta = torch.where(m_odd, (hif8_val_tmp + 1) >> 1, frac_val_ta) + + m_even = m_ta_norm & (~m_carry) & (hif8_val_tmp != 0) & (hif8_val_tmp % 2 == 0) + frac_val_ta = torch.where(m_even, hif8_val_tmp >> 1, frac_val_ta) + + # ------------------ SSR 舍入分支 ------------------ + carry_ssr = torch.zeros_like(exponent, dtype=torch.bool) + frac_val_ssr = torch.zeros_like(exponent) + + f14_v1 = (fraction_int >> SSR_DML_SHIFT) + SSR_F14_OFFSET + t14_v1 = fraction_int & SSR_T14_MASK + hif8_v1 = torch.zeros_like(fraction_int) + + s_bits = torch.clamp(FP32_FRACTION_BITS - frac_bits, min=0) + hif8_v2 = fraction_int >> s_bits + f14_t14 = fraction_int - (hif8_v2 << s_bits) + s_bits_f14 = torch.clamp(FP32_FRACTION_BITS - frac_bits - SSR_RESERVED_BITS, min=0) + f14_v2 = f14_t14 >> s_bits_f14 + t14_v2 = f14_t14 & SSR_T14_MASK + + f14_values = torch.where(m_zero_thresh, f14_v1, f14_v2) + t14_values = torch.where(m_zero_thresh, t14_v1, t14_v2) + hif8_value = torch.where(m_zero_thresh, hif8_v1, hif8_v2) + + m_ge = f14_values >= t14_values + pow_frac_ssr = torch.pow(2, frac_bits) - 1 + m_ssr_carry = m_ge & (hif8_value == pow_frac_ssr) + carry_ssr = torch.where(m_ssr_carry, True, carry_ssr) + frac_val_ssr = torch.where(m_ge & (~m_ssr_carry), hif8_value + 1, frac_val_ssr) + frac_val_ssr = torch.where(~m_ge, hif8_value, frac_val_ssr) + + # ------------------ 合并舍入结果 ------------------ + carry_exp_status = torch.where(cut_bit_is_ta, carry_ta, carry_ssr) + hif8_frac_value = torch.where(cut_bit_is_ta, frac_val_ta, frac_val_ssr) + + exponent = torch.where(carry_exp_status, exponent + 1, exponent) + abs_exp = torch.abs(exponent) + + dot = torch.where( + carry_exp_status, + torch.where(exponent < HIF8_EXP_DML_MIN, HIF8_DOT_INVALID, dot), + dot, + ) + dot = torch.where( + carry_exp_status, + torch.where( + (exponent >= HIF8_EXP_DML_MIN) & (exponent < HIF8_EXP_DML_MAX), + HIF8_DOT_DML, + dot, + ), + dot, + ) + dot = torch.where( + carry_exp_status, torch.where(exponent == HIF8_EXP_D0, HIF8_DOT_D0, dot), dot + ) + dot = torch.where( + carry_exp_status, + torch.where(abs_exp == HIF8_EXP_D1_BOUNDARY, HIF8_DOT_D1, dot), + dot, + ) + dot = torch.where( + carry_exp_status, + torch.where( + (abs_exp >= HIF8_EXP_D2_MIN) & (abs_exp <= HIF8_EXP_D2_MAX), + HIF8_DOT_D2, + dot, + ), + dot, + ) + dot = torch.where( + carry_exp_status, + torch.where( + (abs_exp >= HIF8_EXP_D3_MIN) & (abs_exp <= HIF8_EXP_D3_MAX), + HIF8_DOT_D3, + dot, + ), + dot, + ) + dot = torch.where( + carry_exp_status, + torch.where( + (abs_exp >= HIF8_EXP_D4_MIN) & (abs_exp <= HIF8_EXP_D4_MAX), + HIF8_DOT_D4, + dot, + ), + dot, + ) + + frac_bits = torch.where( + carry_exp_status, + torch.where(exponent < HIF8_EXP_DML_MIN, HIF8_FRAC_BITS_DML, frac_bits), + frac_bits, + ) + frac_bits = torch.where( + carry_exp_status, + torch.where( + (exponent >= HIF8_EXP_DML_MIN) & (exponent < HIF8_EXP_DML_MAX), + HIF8_FRAC_BITS_DML, + frac_bits, + ), + frac_bits, + ) + frac_bits = torch.where( + carry_exp_status, + torch.where(exponent == HIF8_EXP_D0, HIF8_FRAC_BITS_D0, frac_bits), + frac_bits, + ) + frac_bits = torch.where( + carry_exp_status, + torch.where(abs_exp == HIF8_EXP_D1_BOUNDARY, HIF8_FRAC_BITS_D1, frac_bits), + frac_bits, + ) + frac_bits = torch.where( + carry_exp_status, + torch.where( + (abs_exp >= HIF8_EXP_D2_MIN) & (abs_exp <= HIF8_EXP_D2_MAX), + HIF8_FRAC_BITS_D2, + frac_bits, + ), + frac_bits, + ) + frac_bits = torch.where( + carry_exp_status, + torch.where( + (abs_exp >= HIF8_EXP_D3_MIN) & (abs_exp <= HIF8_EXP_D3_MAX), + HIF8_FRAC_BITS_D3, + frac_bits, + ), + frac_bits, + ) + frac_bits = torch.where( + carry_exp_status, + torch.where( + (abs_exp >= HIF8_EXP_D4_MIN) & (abs_exp <= HIF8_EXP_D4_MAX), + HIF8_FRAC_BITS_D4, + frac_bits, + ), + frac_bits, + ) + + exp_bits = torch.where( + carry_exp_status, + torch.where(exponent < HIF8_EXP_DML_MIN, HIF8_EXP_BITS_DML, exp_bits), + exp_bits, + ) + exp_bits = torch.where( + carry_exp_status, + torch.where( + (exponent >= HIF8_EXP_DML_MIN) & (exponent < HIF8_EXP_DML_MAX), + HIF8_EXP_BITS_DML, + exp_bits, + ), + exp_bits, + ) + exp_bits = torch.where( + carry_exp_status, + torch.where(exponent == HIF8_EXP_D0, HIF8_EXP_BITS_D0, exp_bits), + exp_bits, + ) + exp_bits = torch.where( + carry_exp_status, + torch.where(abs_exp == HIF8_EXP_D1_BOUNDARY, HIF8_EXP_BITS_D1, exp_bits), + exp_bits, + ) + exp_bits = torch.where( + carry_exp_status, + torch.where( + (abs_exp >= HIF8_EXP_D2_MIN) & (abs_exp <= HIF8_EXP_D2_MAX), + HIF8_EXP_BITS_D2, + exp_bits, + ), + exp_bits, + ) + exp_bits = torch.where( + carry_exp_status, + torch.where( + (abs_exp >= HIF8_EXP_D3_MIN) & (abs_exp <= HIF8_EXP_D3_MAX), + HIF8_EXP_BITS_D3, + exp_bits, + ), + exp_bits, + ) + exp_bits = torch.where( + carry_exp_status, + torch.where( + (abs_exp >= HIF8_EXP_D4_MIN) & (abs_exp <= HIF8_EXP_D4_MAX), + HIF8_EXP_BITS_D4, + exp_bits, + ), + exp_bits, + ) + + # ------------------ 组合输出编码 ------------------ + hif8_int_value = torch.zeros_like(exponent) + sig_exp = torch.where(exponent < 0, 1, 0) + + # 分支 A: dot <= 0 + m_a = dot <= 0 + val_a = torch.where( + exponent <= HIF8_EXP_ZERO_THRESHOLD, + 0, + sign_norm + exponent + HIF8_DML_EXP_OFFSET, + ) + hif8_int_value = torch.where(m_a, val_a, hif8_int_value) + + # 分支 B: dot == 1 + m_b = dot == 1 + val_b = sign_norm + (dot << HIF8_DOT_BIT_SHIFT) + hif8_frac_value + hif8_int_value = torch.where(m_b, val_b, hif8_int_value) + + # 分支 C: dot > 1 + m_c = dot > 1 + abs_exponent = torch.abs(exponent) + abs_exponent = abs_exponent - torch.pow(2, exp_bits - 1) + exponent_int_value = abs_exponent << frac_bits + sig_exp_shifted = sig_exp << (exp_bits - 1 + frac_bits) + dot_int_value = dot << HIF8_DOT_BIT_SHIFT + val_c = ( + sign_norm + + dot_int_value + + sig_exp_shifted + + exponent_int_value + + hif8_frac_value + ) + hif8_int_value = torch.where(m_c, val_c, hif8_int_value) + + hif8_int_value = torch.where(exponent < HIF8_EXP_ZERO_THRESHOLD, 0, hif8_int_value) + + out[mask_normal] = hif8_int_value + + return out.reshape(shape).to(torch.uint8) + + +def cvt_hifuint8_to_float32(x, over_mode=True): + x = int(x) + if x == HIF8_ZERO: + return float(0) + elif x == HIF8_NAN: + if over_mode: + return float("nan") + else: + return float(0) + elif x == HIF8_NEG_INF: + if over_mode: + return -torch.inf + else: + return -HIF8_MAX_FINITE_VALUE + elif x == HIF8_POS_INF: + if over_mode: + return torch.inf + else: + return HIF8_MAX_FINITE_VALUE + else: + if x >= HIF8_NAN: + sign = -1.0 + else: + sign = 1.0 + dot_4_bits = x & HIF8_DOT_MASK + dot_4_value = dot_4_bits >> 3 + if dot_4_value >= HIF8_DOT_D4: + exponent = x & HIF8_EXP_MASK_D4 + exponent_int = exponent >> 1 + if exponent_int >= 8: + exponent_value = -exponent_int + else: + exponent_value = exponent_int + 8 + + fra_int = x & HIF8_FRAC_MASK_1BIT + m_value = 1.0 + fra_int * 0.5 + elif dot_4_value >= HIF8_DOT_D3: + exponent = x & HIF8_EXP_MASK_D3 + exponent_int = exponent >> 2 + if exponent_int >= 4: + exponent_value = -exponent_int + else: + exponent_value = exponent_int + 4 + + fra_int = x & HIF8_FRAC_MASK_2BIT + m_value = 1.0 + fra_int * 0.25 + elif dot_4_value >= HIF8_DOT_D2: + exponent = x & HIF8_EXP_MASK_D2 + exponent_int = exponent >> 3 + if exponent_int >= 2: + exponent_value = -exponent_int + else: + exponent_value = exponent_int + 2 + + fra_int = x & HIF8_FRAC_MASK_3BIT + m_value = 1.0 + fra_int * 0.125 + elif dot_4_value >= HIF8_DOT_D1: + exponent = x & HIF8_EXP_SIGN_MASK_D1 + exponent_sign = exponent >> 3 + if exponent_sign >= 1: + exponent_value = -1 + else: + exponent_value = 1 + + fra_int = x & HIF8_FRAC_MASK_3BIT + m_value = 1.0 + fra_int * 0.125 + elif dot_4_value == HIF8_DOT_D0: + exponent_value = 0 + fra_int = x & HIF8_FRAC_MASK_3BIT + m_value = 1.0 + fra_int * 0.125 + elif dot_4_value == HIF8_DOT_DML: + m_value = 1 + exponent_value = (x & HIF8_EXP_MASK_DML) - HIF8_DML_EXP_OFFSET + else: + print("error, dot error") + m_value = 0.0 + exponent_value = 0 + return sign * pow(2.0, exponent_value) * m_value + + +def trans_hifuint8_tensor_to_float(in_tensor, over_mode=True): + """ + 将 HiF8 编码的 uint8 Tensor 批量转换为 float32 Tensor (支持 CPU/GPU 矢量化) + """ + shape = in_tensor.shape + x = in_tensor.reshape(-1).to(torch.int32) + out = torch.zeros_like(x, dtype=torch.float32) + + # 1. 特殊值处理 (Masks) + mask_zero = x == HIF8_ZERO + mask_nan = x == HIF8_NAN + mask_ninf = x == HIF8_NEG_INF + mask_pinf = x == HIF8_POS_INF + + if over_mode: + out = torch.where(mask_nan, torch.tensor(float("nan"), device=x.device), out) + out = torch.where(mask_ninf, torch.tensor(-torch.inf, device=x.device), out) + out = torch.where(mask_pinf, torch.tensor(torch.inf, device=x.device), out) + else: + out = torch.where(mask_nan, 0.0, out) + out = torch.where(mask_ninf, float(-HIF8_MAX_FINITE_VALUE), out) + out = torch.where(mask_pinf, float(HIF8_MAX_FINITE_VALUE), out) + + # 正常数值的 Mask (排除特殊值) + mask_normal = ~(mask_zero | mask_nan | mask_ninf | mask_pinf) + if not mask_normal.any(): + return out.reshape(shape) + + # 提取正常数值子集进行计算 + x_norm = x[mask_normal] + + # 符号位计算 + sign = torch.where(x_norm >= HIF8_NAN, -1.0, 1.0) + + # 提取 dot 档位 + dot_4_value = (x_norm & HIF8_DOT_MASK) >> 3 + + # 初始化指数和尾数乘子 + exponent_value = torch.zeros_like(x_norm, dtype=torch.float32) + m_value = torch.zeros_like(x_norm, dtype=torch.float32) + + # --- 档位 D4 --- + m_d4 = dot_4_value >= HIF8_DOT_D4 + if m_d4.any(): + exp_int = (x_norm & HIF8_EXP_MASK_D4) >> 1 + exponent_value = torch.where( + m_d4, + torch.where(exp_int >= 8, -exp_int, exp_int + 8).float(), + exponent_value, + ) + m_value = torch.where(m_d4, 1.0 + (x_norm & HIF8_FRAC_MASK_1BIT) * 0.5, m_value) + + # --- 档位 D3 --- + m_d3 = (~m_d4) & (dot_4_value >= HIF8_DOT_D3) + if m_d3.any(): + exp_int = (x_norm & HIF8_EXP_MASK_D3) >> 2 + exponent_value = torch.where( + m_d3, + torch.where(exp_int >= 4, -exp_int, exp_int + 4).float(), + exponent_value, + ) + m_value = torch.where( + m_d3, 1.0 + (x_norm & HIF8_FRAC_MASK_2BIT) * 0.25, m_value + ) + + # --- 档位 D2 --- + m_d2 = (~(m_d4 | m_d3)) & (dot_4_value >= HIF8_DOT_D2) + if m_d2.any(): + exp_int = (x_norm & HIF8_EXP_MASK_D2) >> 3 + exponent_value = torch.where( + m_d2, + torch.where(exp_int >= 2, -exp_int, exp_int + 2).float(), + exponent_value, + ) + m_value = torch.where( + m_d2, 1.0 + (x_norm & HIF8_FRAC_MASK_3BIT) * 0.125, m_value + ) + + # --- 档位 D1 --- + m_d1 = (~(m_d4 | m_d3 | m_d2)) & (dot_4_value >= HIF8_DOT_D1) + if m_d1.any(): + exp_sign = (x_norm & HIF8_EXP_SIGN_MASK_D1) >> 3 + exponent_value = torch.where( + m_d1, torch.where(exp_sign >= 1, -1.0, 1.0), exponent_value + ) + m_value = torch.where( + m_d1, 1.0 + (x_norm & HIF8_FRAC_MASK_3BIT) * 0.125, m_value + ) + + # --- 档位 D0 --- + m_d0 = dot_4_value == HIF8_DOT_D0 + if m_d0.any(): + exponent_value = torch.where(m_d0, 0.0, exponent_value) + m_value = torch.where( + m_d0, 1.0 + (x_norm & HIF8_FRAC_MASK_3BIT) * 0.125, m_value + ) + + # --- 档位 DML --- + m_dml = dot_4_value == HIF8_DOT_DML + if m_dml.any(): + exponent_value = torch.where( + m_dml, + ((x_norm & HIF8_EXP_MASK_DML) - HIF8_DML_EXP_OFFSET).float(), + exponent_value, + ) + m_value = torch.where(m_dml, 1.0, m_value) + + # 计算正常值结果并写回 + norm_res = sign * torch.pow(2.0, exponent_value) * m_value + out[mask_normal] = norm_res + + return out.reshape(shape) diff --git a/xllm_ops/attention/quant_lightning_indexer_v2/tests/pytest/result_compare_method.py b/xllm_ops/attention/quant_lightning_indexer_v2/tests/pytest/result_compare_method.py new file mode 100644 index 0000000..e8c2531 --- /dev/null +++ b/xllm_ops/attention/quant_lightning_indexer_v2/tests/pytest/result_compare_method.py @@ -0,0 +1,909 @@ +#!/usr/bin/python +# -*- coding: utf-8 -*- +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +import math +import random +import logging +import torch +import datetime +import os +import sys +import ast +import numpy as np +from time import time + +logging.basicConfig(level=logging.INFO, format="%(message)s", force=True) +logger = logging.getLogger(__name__) + + +def cal_relative_diff_np_isclose(real_data, expect_data, type_str="fp16"): + diff = abs(float(real_data) - float(expect_data)) + result = diff / (np.abs(expect_data) + 10e-10) + return result + + +def print_log(data=None, level="INFO"): + print( + "[%s] [%s]-%s:%s - %s" + % ( + datetime.datetime.now().strftime("%Y/%m/%d %H:%M:%S"), + level, + os.path.basename(sys._getframe().f_back.f_code.co_filename), + str(sys._getframe().f_back.f_lineno).zfill(4), + data, + ) + ) + + +def display_error_output(real_data, expect_data, err_idx, relative_diff): + print_log( + "Error Line-----------------------------------------------------------------------------" + ) + print_log("Loop \t ExpectOut \t RealOut \t FpDiff \t RateDiff") + print_log( + "---------------------------------------------------------------------------------------" + ) + count = 0 + len_err = len(err_idx) + for i in err_idx: + count += 1 + if count < 10 or (90 < count < 100): + print_log( + "%08d \t %.7f \t %.7f \t %.7f \t %.7f" + % ( + i, + expect_data[i], + real_data[i], + abs(np.float64(expect_data[i]) - np.float64(real_data[i])), + relative_diff[count - 1], + ) + ) + elif count == 10 or (count == 100 and len_err > 100): + dot_3 = "..." + print_log( + "%08s \t %07s \t %07s \t %07s \t %07s" + % (dot_3, dot_3, dot_3, dot_3, dot_3) + ) + elif count > 100: + break + + print_log( + "Max-RE line:---------------------------------------------------------------------------" + ) + max_error = max(relative_diff) + m_idx_list = err_idx[np.where(relative_diff == max_error)] + m_count = 0 + for m_idx in m_idx_list: + m_count += 1 + if m_count < 4: + print_log( + "%08d \t %.7f \t %.7f \t %.7f \t %.7f" + % ( + m_idx, + expect_data[m_idx], + real_data[m_idx], + abs(np.float64(expect_data[m_idx]) - np.float64(real_data[m_idx])), + max_error, + ) + ) + else: + break + print_log( + "---------------------------------------------------------------------------------------" + ) + + +def display_output_np_isclose( + real_data, expect_data, start, end, expect_fp32_data=None +): + def display_inner(idx): + j = idx + start + diff_rate = cal_relative_diff_np_isclose(real_data[j], expect_data[j]) + if "inf" in str(expect_data[j]) or "nan" in str(expect_data[j]): + diff_abs = "inf" if "inf" in str(expect_data[j]) else "nan" + if expect_fp32_data is not None: + print_log( + "%08d \t %-7s \t %-7s \t %-7s \t %-7s \t %-7s" + % ( + start + idx + 1, + expect_fp32_data[j], + expect_data[j], + real_data[j], + diff_abs, + diff_rate, + ) + ) + else: + print_log( + "%08d \t %-7s \t %-7s \t %-7s \t %-7s" + % ( + start + idx + 1, + expect_data[j], + real_data[j], + diff_abs, + diff_rate, + ) + ) + else: + diff_abs = abs(np.float64(expect_data[j]) - np.float64(real_data[j])) + if expect_fp32_data is not None: + print_log( + "%08d \t %0.7f \t %0.7f \t %0.7f \t %0.7f \t %0.7f" + % ( + start + idx + 1, + expect_fp32_data[j], + expect_data[j], + real_data[j], + diff_abs, + diff_rate, + ) + ) + else: + print_log( + "%08d \t %0.7f \t %0.7f \t %0.7f \t %0.7f" + % ( + start + idx + 1, + expect_data[j], + real_data[j], + diff_abs, + diff_rate, + ) + ) + + print_log( + "---------------------------------------------------------------------------------------" + ) + if expect_fp32_data is not None: + print_log( + "Loop \t ExpFP32Out \t ExpFP16Out \t NPUOut \tFpDiff(min) \t RateDiff" + ) + else: + print_log("Loop \t ExpectOut \t RealOut \t FpDiff \t RateDiff") + print_log( + "---------------------------------------------------------------------------------------" + ) + split_count = int(end - start) + if split_count <= 20: + for i in range(split_count + 1): + display_inner(i) + else: + for i in range(10): + display_inner(i) + print_log("... \t ... \t ... \t ... \t ...") + for i in range(split_count - 10 + 1, split_count + 1): + display_inner(i) + + +def find_batch_and_position(cu_seqlens, x): + """ + 判断x属于哪个batch以及在该batch中的位置 + + 参数: + cu_seqlens: 前缀和列表, cu_seqlens[b_idx]表示前(b_idx)个batch的总长度 + x: 需要判断的数值 + + 返回: + tuple: (batch_idx, position) + - batch_idx: 所属的batch索引(从0开始)，超出范围则为-1 + - position: 在该batch中的位置(从0开始), 超出范围则为-1 + """ + if not cu_seqlens: + return (-1, -1) + # 遍历前缀和列表查找所属批次 + for batch_idx in range(len(cu_seqlens) - 1): + # 计算当前批次的起始位置 + start = cu_seqlens[batch_idx] + # 判断是否在当前批次范围内 + if start <= x < cu_seqlens[batch_idx + 1]: + # 计算在当前批次中的位置（偏移量） + position = x - start + return (batch_idx, position) + # 超出所有批次范围 + return (-1, -1) + + +def judge_value_by_isclose(real_data, data_compe, force_bf16=False): + atol = 2.5e-05 + rtol = 0.005 + pct_thd = 0.005 + diff_thd = 0.005 + # force_bf16: QLIV2 的 returnValue 固定为 bf16，但流程中已被 .float() 转成 float32， + # 无法通过 dtype 判断，需强制按 bf16 门限对比。 + is_bfloat16 = force_bf16 or (str(real_data.dtype) in ("bfloat16", "torch.bfloat16")) + if isinstance(real_data, torch.Tensor): + real_data = real_data.detach().cpu().float().numpy() + else: + real_data = np.asarray(real_data) + if isinstance(data_compe, torch.Tensor): + data_compe = data_compe.detach().cpu().float().numpy() + else: + data_compe = np.asarray(data_compe) + start = 0 + end = real_data.size - 1 + if end < start: + end = start + split_count = int(end - start + 1) if end != start else 1 + + if is_bfloat16: + # bf16 尾数位少、舍入误差大，误差门限放宽到 1/128（约 0.0078125） + atol = 0.0001 + rtol = 1.0 / 128 + diff_thd = 1.0 / 128 + diff_result = np.isclose( + real_data.astype(np.float32), + data_compe.astype(np.float32), + rtol=rtol, + atol=atol, + equal_nan=True, + ) + else: + diff_result = np.isclose( + real_data, data_compe, rtol=rtol, atol=atol, equal_nan=True + ) + err_idx = np.where(diff_result != np.array((True,)))[0] + diff_abs = abs(data_compe - real_data) + b1 = np.maximum(np.abs(real_data), (np.abs(data_compe))) + b2 = float((1.0 / (1 << 14)) / diff_thd) + b = np.add(np.maximum(b1, b2), 10e-10) + eps = 10e-10 + err_diff = diff_abs / (b + eps) + err_diff = err_diff[err_idx] + fulfill_percent = float(split_count - err_idx.size) / float(split_count) * 100.0 + pct_thd = (1 - pct_thd) * 100.0 + result = True if (fulfill_percent >= pct_thd) else False + return result + + +def compare_topk_valid( + cur_cpu, + cur_npu, + topk_value, + bsn, + diff_npu, + diff_cpu, + cur_npu_output_value=None, + cur_cpu_output_value=None, + thres=0.001, + return_value_flag=False, + output_idx_offset=None, + layout_query=None, + cu_seqlens_q=None, + q_seq=0, +): + b_idx, s1_idx, n2_idx = bsn + max_re = 0.0 + npu_pass = True + cur_cpu = np.asarray(cur_cpu, dtype=np.int64) + cur_npu = np.asarray(cur_npu, dtype=np.int64) + + if output_idx_offset is not None: + # 统一转换后使用 + offset_data = ( + output_idx_offset.cpu().numpy() + if hasattr(output_idx_offset, "device") + and output_idx_offset.device.type != "cpu" + else np.array(output_idx_offset) + ) + offset_flat = offset_data.flatten() + if layout_query == "TND": + cur_prefix = cu_seqlens_q[b_idx] + offset = offset_flat[cur_prefix + s1_idx] + else: + offset = offset_flat[b_idx * q_seq + s1_idx] + cpu_offset_mask = cur_cpu != -1 + npu_offset_mask = cur_npu != -1 + cur_cpu = np.where(cpu_offset_mask, cur_cpu - offset, cur_cpu) + cur_npu = np.where(npu_offset_mask, cur_npu - offset, cur_npu) + + element_list = topk_value[b_idx, n2_idx, s1_idx, :] + score_size = element_list.shape[-1] + invalid_cpu = (cur_cpu < 0) | (cur_cpu >= score_size) + invalid_npu = (cur_npu < 0) | (cur_npu >= score_size) + has_duplicate_cpu = np.unique(cur_cpu).size != cur_cpu.size + has_duplicate_npu = np.unique(cur_npu).size != cur_npu.size + if ( + cur_cpu.size != cur_npu.size + or np.any(invalid_cpu) + or np.any(invalid_npu) + or has_duplicate_cpu + or has_duplicate_npu + ): + diff_cpu.append(cur_cpu.tolist()) + diff_npu.append(cur_npu.tolist()) + return False, float("inf") + + npu_set = set(cur_npu) + cpu_set = set(cur_cpu) + if npu_set != cpu_set: + value_bm = topk_value[b_idx, n2_idx, s1_idx, cur_cpu[-1]] + only_in_npu = npu_set - cpu_set + only_in_cpu = cpu_set - npu_set + only_in_npu_list = list(only_in_npu) + only_in_cpu_list = list(only_in_cpu) + for diff_idx in range(len(only_in_npu_list)): + element_npu = element_list[only_in_npu_list[diff_idx]] + element_cpu = element_list[only_in_cpu_list[diff_idx]] + npu_ae = abs(element_npu - value_bm) + cpu_ae = abs(element_cpu - value_bm) + if value_bm == 0: + if npu_ae == 0: + npu_re = 0.0 + else: + npu_re = float("inf") + if cpu_ae == 0: + cpu_re = 0.0 + else: + cpu_re = float("inf") + else: + npu_re = abs(npu_ae / value_bm) + cpu_re = abs(cpu_ae / value_bm) + if npu_re > thres or cpu_re > thres: + if return_value_flag: + # 将 value 输出统一转为 numpy array，bfloat16 需先转 float32 再转 numpy + if torch.is_tensor(cur_npu_output_value): + npuValueArr = ( + cur_npu_output_value.float().cpu().numpy() + if cur_npu_output_value.dtype == torch.bfloat16 + else cur_npu_output_value.cpu().numpy() + ) + else: + npuValueArr = np.asarray(cur_npu_output_value) + if torch.is_tensor(cur_cpu_output_value): + cpuValueArr = ( + cur_cpu_output_value.float().cpu().numpy() + if cur_cpu_output_value.dtype == torch.bfloat16 + else cur_cpu_output_value.cpu().numpy() + ) + else: + cpuValueArr = np.asarray(cur_cpu_output_value) + if not judge_value_by_isclose(npuValueArr, cpuValueArr): + npu_pass = False + diff_npu.append(element_npu) + diff_cpu.append(element_cpu) + max_re = max(max_re, npu_re, cpu_re) + else: + npu_pass = False + diff_npu.append(element_npu) + diff_cpu.append(element_cpu) + max_re = max(max_re, npu_re) + return npu_pass, max_re + + +def compare_return_value(cur_npu_output_value=None, cur_cpu_output_value=None): + max_re = 0.0 + npu_pass = True + npu_pass = judge_value_by_isclose(cur_npu_output_value, cur_cpu_output_value) + return npu_pass, max_re + + +def trans_tnd_actseq(list): + list_len = len(list) + if list_len == 0: + raise ValueError("TND情况下 act_seq需要必传") + list_new = [] + list_new.append(list[0]) + for i in range(list_len - 1): + new_item = list[i + 1] - list[i] + if new_item >= 0: + list_new.append(new_item) + else: + raise ValueError(f"TND情况下 act_seq_len 为非递减数列 act_seq_len={list}") + return list_new + + +def _reshape_topk_value(topk_value, total_rows, sparse_count, params): + if isinstance(topk_value, torch.Tensor): + topk_value = topk_value.detach().cpu().float().numpy() + else: + topk_value = np.asarray(topk_value) + if topk_value.size == total_rows * sparse_count: + return topk_value.reshape(total_rows, sparse_count) + + batch_size, cu_seqlens_q, layout_query = params[0], params[14], params[21] + if layout_query != "TND" or topk_value.ndim != 4: + raise ValueError( + f"topk value shape {topk_value.shape} cannot reshape to " + f"({total_rows}, {sparse_count})" + ) + cu_seqlens_q = _get_tnd_query_prefix(cu_seqlens_q, batch_size) + topk_value = np.concatenate( + [ + topk_value[ + batch_idx, :, : cu_seqlens_q[batch_idx + 1] - cu_seqlens_q[batch_idx], : + ] + .transpose(1, 0, 2) + .reshape(-1, sparse_count) + for batch_idx in range(batch_size) + ] + ) + return topk_value.reshape(total_rows, sparse_count) + + +def check_result( + expect, + result, + topk_value, + output_idx_offset, + params, + cpu_topk_value, + npu_topk_value, +): + ( + batch_size, + q_seq, + k_seq, + q_t_size, + k_t_size, + q_head_num, + k_head_num, + head_dim, + block_size, + block_num, + qk_dtype, + weight_dtype, + dequant_dtype, + actual_seq_dtype, + cu_seqlens_q, + cu_seqlens_k, + seqused_q, + seqused_k, + cmp_residual_k, + max_seqlen_q, + quant_mode, + layout_query, + layout_key, + sparse_count, + sparse_mode, + query_datarange, + key_datarange, + weights_datarange, + q_scale_datarange, + k_scale_datarange, + cmp_ratio, + return_value, + _, + ) = params + + # Q 侧个体长度 + if layout_query == "TND": + # TND: 必传 cu_seqlens_q，从差分推导个体长度 + if isinstance(cu_seqlens_q, str): + lengths_q_list = ast.literal_eval(cu_seqlens_q) + else: + lengths_q_list = cu_seqlens_q[1:] + else: + # BSND: 从 seqused_q 获取，若 None 则用 q_seq 填满 + if seqused_q is not None: + if isinstance(seqused_q, str): + lengths_q_list = ast.literal_eval(seqused_q) + else: + lengths_q_list = list(seqused_q) + else: + lengths_q_list = [q_seq] * batch_size + + # K 侧个体长度 + if layout_key == "TND": + # TND: 必传 cu_seqlens_k，从差分推导个体长度 + if isinstance(cu_seqlens_k, str): + lengths_k_list = ast.literal_eval(cu_seqlens_k) + else: + lengths_k_list = cu_seqlens_k[1:] + elif layout_key == "PA_BBND": + # PA_BBND: 从 seqused_k 获取 + assert seqused_k is not None, f"{layout_key} layout requires seqused_k" + if isinstance(seqused_k, str): + lengths_k_list = ast.literal_eval(seqused_k) + else: + lengths_k_list = list(seqused_k) + else: + # BSND: 从 seqused_k 获取，若 None 则用 q_seq 填满 + if seqused_k is not None: + if isinstance(seqused_k, str): + lengths_k_list = ast.literal_eval(seqused_k) + else: + lengths_k_list = list(seqused_k) + else: + lengths_k_list = [k_seq] * batch_size + + act_seq_q = lengths_q_list + act_seq_k = lengths_k_list + + if isinstance(act_seq_q, int): + act_seq_q = [act_seq_q] + elif isinstance(act_seq_q, list): + act_seq_q = act_seq_q + else: + act_seq_q = ast.literal_eval(act_seq_q) + if isinstance(act_seq_k, int): + act_seq_k = [act_seq_k] + elif isinstance(act_seq_k, list): + act_seq_k = act_seq_k + else: + act_seq_k = ast.literal_eval(act_seq_k) + + if isinstance(cu_seqlens_q, int): + cu_seqlens_q = [cu_seqlens_q] + elif isinstance(cu_seqlens_q, list): + cu_seqlens_q = cu_seqlens_q + elif cu_seqlens_q is not None: + cu_seqlens_q = ast.literal_eval(cu_seqlens_q) + + if isinstance(cu_seqlens_k, int): + cu_seqlens_k = [cu_seqlens_k] + elif isinstance(cu_seqlens_k, list): + cu_seqlens_k = cu_seqlens_k + elif cu_seqlens_k is not None: + cu_seqlens_k = ast.literal_eval(cu_seqlens_k) + + if isinstance(seqused_q, int): + seqused_q = [seqused_q] + elif isinstance(seqused_q, list): + seqused_q = seqused_q + elif seqused_q is not None: + seqused_q = ast.literal_eval(seqused_q) + + if isinstance(seqused_k, int): + seqused_k = [seqused_k] + elif isinstance(seqused_k, list): + seqused_k = seqused_k + elif seqused_k is not None: + seqused_k = ast.literal_eval(seqused_k) + npu_pass = True + max_error = 0 + max_re = 0 + thres = 0.001 + diff_thd = 0.01 + pct_thd = 0.005 + max_diff_hd = 0.1 + rtol = 0.005 + atol = 0.000025 + max_error_idx = 10000000 + cpu_output = expect.cpu().numpy() + npu_output = result.cpu().numpy() + real_data = result.cpu().numpy() + data_compe = expect.cpu().numpy() + real_data = npu_output.flatten() + data_compe = cpu_output.flatten() + diff_cpu = [] + diff_npu = [] + + if layout_query in ["BSND"]: + sp = (batch_size, q_seq, k_head_num) + total_rows = batch_size * q_seq * k_head_num + elif layout_query in ["TND"]: + sp = (q_t_size, k_head_num) + total_rows = q_t_size * k_head_num + else: + total_rows = 0 + sp = (0, 0) + print(f"total_line is {total_rows}") + npu_reshape = npu_output.reshape([total_rows, sparse_count]) + cpu_reshape = cpu_output.reshape([total_rows, sparse_count]) + if return_value: + cpu_topk_value = _reshape_topk_value( + cpu_topk_value, total_rows, sparse_count, params + ) + npu_topk_value = _reshape_topk_value( + npu_topk_value, total_rows, sparse_count, params + ) + start_time = time() + invalid_data = cpu_reshape != -1 + valid_lens = invalid_data.sum(axis=-1) # (total_rows,) + # 判断有效值部分集合是否相同 + cpu_output_sorted = np.sort(cpu_reshape, axis=1) + npu_output_sorted = np.sort(npu_reshape, axis=1) + diff_rows = np.zeros(total_rows, dtype=bool) + diff_rows |= np.any( + cpu_output_sorted != npu_output_sorted, axis=1 + ) # 标记存在差异的行 + test_id = [] + rows = [] + if np.any(diff_rows): + rows = np.where(diff_rows)[0] + num_rows = len(rows) + if num_rows: + print(f"需要进行第二步比较的batch有{num_rows}") + else: + print("有效值集合相同，无需进行比较") + for t_id in rows: + bsn = np.unravel_index(t_id, sp) + npu_topk_output_value = None + cpu_topk_output_value = None + if layout_query == "TND": + b_idx, s1_idx = find_batch_and_position(cu_seqlens_q, bsn[0]) + bsn = (b_idx, s1_idx, bsn[-1]) + if return_value: + cpu_topk_output_value = cpu_topk_value[t_id, :] + npu_topk_output_value = npu_topk_value[t_id, :] + npu_pass_t = True + max_re_t = 0 + valid_len = valid_lens[t_id] + npu_pass_t, max_re_t = compare_topk_valid( + cpu_reshape[t_id, :valid_len], + npu_reshape[t_id, :valid_len], + topk_value, + bsn, + diff_npu, + diff_cpu, + npu_topk_output_value, + cpu_topk_output_value, + thres, + return_value, + output_idx_offset, + layout_query, + cu_seqlens_q, + q_seq, + ) + if not npu_pass_t: + npu_pass = False + end_time = time() + print(f"耗时：{end_time - start_time:.6f} 秒") + topk_precision = not diff_npu and not diff_cpu + if topk_precision: + print("[success]TopK精度通过, idx不同的地方的value误差在阈值之内") + else: + print("[fail]TopK精度失败") + print(f"npu_pass is {npu_pass}") + if real_data.size == 0 and real_data.size == data_compe.size: + print_log( + 'The npu_output is [],and it is same as bm_output, the result of data_compare is "Pass"' + ) + return "Pass", 100.0, 0 + start = 0 + end = real_data.size - 1 + if end < start: + end = start + diff_result = np.isclose( + real_data, data_compe, rtol=rtol, atol=atol, equal_nan=True + ) + err_idx = np.where(diff_result != np.array((True,)))[0] + diff_abs = abs(data_compe - real_data) + b1 = np.maximum(np.abs(real_data), (np.abs(data_compe))) + b2 = float((1.0 / (1 << 14)) / diff_thd) + b = np.add(np.maximum(b1, b2), 10e-10) + eps = 10e-10 + err_diff = diff_abs / (b + eps) + err_diff = err_diff[err_idx] + split_count = int(end - start + 1) if end != start else 1 + print_log("split_count:%s; max_diff_hd:%s;" % (float(split_count), max_diff_hd)) + fulfill_percent = float(split_count - err_idx.size) / float(split_count) * 100.0 + display_output_np_isclose(real_data, data_compe, start, end) + pct_thd = (1 - pct_thd) * 100.0 + result = "Pass" if (npu_pass or topk_precision) else "Failed" + print_log( + "---------------------------------------------------------------------------------------" + ) + print_log("Rtol \t Atol \t PctThd \t PctRlt \t Result") + print_log( + "---------------------------------------------------------------------------------------" + ) + print_log( + "%.4f \t %.6f \t %.2f%% \t %.6f%% \t %s" + % (rtol, atol, pct_thd, fulfill_percent, result) + ) + if len(err_diff) > 0: + print_log( + "Max-RelativeError is: %s. Threshold is: %s." % (max_error, max_diff_hd) + ) + if result == "Failed": + display_error_output(real_data, data_compe, err_idx, err_diff[0:max_error_idx]) + return result, fulfill_percent + + +def _to_flat_numpy(value, dtype=np.int64): + if value is None: + return None + if isinstance(value, str): + value = ast.literal_eval(value) + if isinstance(value, torch.Tensor): + value = value.detach().cpu().numpy() + return np.asarray(value, dtype=dtype).reshape(-1) + + +def _get_tnd_query_prefix(cu_seqlens_q, batch_size): + prefix = _to_flat_numpy(cu_seqlens_q) + if prefix is None: + raise ValueError("cu_seqlens_q is required") + if prefix.size == batch_size + 1 and prefix[0] == 0: + return prefix + if prefix.size == batch_size: + return np.concatenate((np.array([0], dtype=prefix.dtype), prefix)) + raise ValueError(f"invalid TND cu_seqlens_q length: {prefix.size}") + + +def _gather_return_values_by_index( + topk_value, result_indices, params, output_idx_offset +): + batch_size = params[0] + q_seq = params[1] + q_t_size = params[3] + k_head_num = params[6] + cu_seqlens_q = params[14] + layout_query = params[21] + sparse_count = params[23] + + full_score = topk_value.detach().cpu().float().numpy() + npu_indices = result_indices.detach().cpu().numpy().reshape(-1, sparse_count) + expected = np.full(npu_indices.shape, -np.inf, dtype=np.float32) + invalid_index = np.zeros(npu_indices.shape, dtype=bool) + offsets = _to_flat_numpy(output_idx_offset) + query_prefix = ( + _get_tnd_query_prefix(cu_seqlens_q, batch_size) + if layout_query == "TND" + else None + ) + + for row_idx in range(npu_indices.shape[0]): + if layout_query == "BSND": + b_idx, s1_idx, n2_idx = np.unravel_index( + row_idx, (batch_size, q_seq, k_head_num) + ) + offset_pos = b_idx * q_seq + s1_idx + elif layout_query == "TND": + t_idx, n2_idx = np.unravel_index(row_idx, (q_t_size, k_head_num)) + b_idx = int(np.searchsorted(query_prefix[1:], t_idx, side="right")) + s1_idx = int(t_idx - query_prefix[b_idx]) + offset_pos = t_idx + else: + raise ValueError(f"unsupported query layout: {layout_query}") + + row_indices = npu_indices[row_idx] + logical_indices = row_indices.astype(np.int64, copy=True) + if offsets is not None: + logical_indices[row_indices >= 0] -= int(offsets[offset_pos]) + row_score = full_score[b_idx, n2_idx, s1_idx] + valid = ( + (row_indices >= 0) + & (logical_indices >= 0) + & (logical_indices < row_score.shape[-1]) + ) + expected[row_idx, valid] = row_score[logical_indices[valid]] + invalid_index[row_idx] = (row_indices < -1) | ((row_indices >= 0) & ~valid) + + return expected, invalid_index + + +def check_result_return_value( + expect, + result, + params, + expect_indices=None, + result_indices=None, + topk_value=None, + output_idx_offset=None, +): + ( + batch_size, + q_seq, + k_seq, + q_t_size, + k_t_size, + q_head_num, + k_head_num, + head_dim, + block_size, + block_num, + qk_dtype, + weight_dtype, + dequant_dtype, + actual_seq_dtype, + cu_seqlens_q, + cu_seqlens_k, + seqused_q, + seqused_k, + cmp_residual_k, + max_seqlen_q, + quant_mode, + layout_query, + layout_key, + sparse_count, + sparse_mode, + query_datarange, + key_datarange, + weights_datarange, + q_scale_datarange, + k_scale_datarange, + cmp_ratio, + return_value, + _, + ) = params + + npu_pass = True + max_error = 0 + max_re = 0 + thres = 0.0001 + diff_thd = 0.01 + pct_thd = 0.005 + max_diff_hd = 0.1 + rtol = 0.005 + atol = 0.000025 + max_error_idx = 10000000 + npu_output = result.cpu().float().numpy() + if topk_value is not None and result_indices is not None: + if output_idx_offset is None: + output_idx_offset = params[32] + cpu_output, invalid_index = _gather_return_values_by_index( + topk_value, result_indices, params, output_idx_offset + ) + else: + cpu_output = expect.cpu().float().numpy() + invalid_index = np.zeros(cpu_output.shape, dtype=bool) + real_data = npu_output.flatten() + data_compe = cpu_output.flatten() + + if layout_query in ["BSND"]: + sp = (batch_size, q_seq, k_head_num) + total_rows = batch_size * q_seq * k_head_num + elif layout_query in ["TND"]: + sp = (q_t_size, k_head_num) + total_rows = q_t_size * k_head_num + else: + total_rows = 0 + sp = (0, 0) + print(f"total_line is {total_rows}") + npu_reshape = npu_output.reshape([total_rows, sparse_count]) + cpu_reshape = cpu_output.reshape([total_rows, sparse_count]) + invalid_index_reshape = invalid_index.reshape([total_rows, sparse_count]) + + start_time = time() + + # QLIV2 returnValue 为 bf16，强制使用 bf16 门限（误差阈值 1/128） + npu_pass = judge_value_by_isclose(npu_reshape, cpu_reshape, force_bf16=True) + if np.any(invalid_index_reshape): + npu_pass = False + end_time = time() + print(f"耗时：{end_time - start_time:.6f} 秒") + print(f"npu_pass is {npu_pass}") + if real_data.size == 0 and real_data.size == data_compe.size: + print_log( + 'The npu_output is [],and it is same as bm_output, the result of data_compare is "Pass"' + ) + return "Pass", 100.0, 0 + start = 0 + end = real_data.size - 1 + if end < start: + end = start + diff_result = np.isclose( + real_data, data_compe, rtol=rtol, atol=atol, equal_nan=True + ) + err_idx = np.where(diff_result != np.array((True,)))[0] + diff_abs = abs(data_compe - real_data) + b1 = np.maximum(np.abs(real_data), (np.abs(data_compe))) + b2 = float((1.0 / (1 << 14)) / diff_thd) + b = np.add(np.maximum(b1, b2), 10e-10) + eps = 10e-10 + err_diff = diff_abs / (b + eps) + err_diff = err_diff[err_idx] + split_count = int(end - start + 1) if end != start else 1 + print_log("split_count:%s; max_diff_hd:%s;" % (float(split_count), max_diff_hd)) + fulfill_percent = float(split_count - err_idx.size) / float(split_count) * 100.0 + display_output_np_isclose(real_data, data_compe, start, end) + pct_thd = (1 - pct_thd) * 100.0 + result = "Pass" if npu_pass else "Failed" + print_log( + "---------------------------------------------------------------------------------------" + ) + print_log("Rtol \t Atol \t PctThd \t PctRlt \t Result") + print_log( + "---------------------------------------------------------------------------------------" + ) + print_log( + "%.4f \t %.6f \t %.2f%% \t %.6f%% \t %s" + % (rtol, atol, pct_thd, fulfill_percent, result) + ) + if len(err_diff) > 0: + print_log( + "Max-RelativeError is: %s. Threshold is: %s." % (max_error, max_diff_hd) + ) + if result == "Failed": + display_error_output(real_data, data_compe, err_idx, err_diff[0:max_error_idx]) + return result, fulfill_percent diff --git a/xllm_ops/attention/quant_lightning_indexer_v2/tests/pytest/test_qliv2_test_utils.py b/xllm_ops/attention/quant_lightning_indexer_v2/tests/pytest/test_qliv2_test_utils.py new file mode 100644 index 0000000..bbb59d2 --- /dev/null +++ b/xllm_ops/attention/quant_lightning_indexer_v2/tests/pytest/test_qliv2_test_utils.py @@ -0,0 +1,100 @@ +#!/usr/bin/python +# -*- coding: utf-8 -*- +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +import pickle +from pathlib import Path + +import pandas as pd +import pytest + +from qliv2_parameter_normalization import normalize_qliv2_params +from qliv2_test_utils import ( + PARAM_NAMES, + QliV2CaseSelector, + QliV2ResultWriter, + ensure_comparison_passed, +) + + +def test_case_selector_preserves_requested_order(tmp_path): + for name in ("case_10.pt", "case_2.pt", "case_1.pt"): + (tmp_path / name).touch() + + natural = QliV2CaseSelector.resolve(tmp_path) + assert [Path(path).name for path in natural] == ["case_1.pt", "case_2.pt", "case_10.pt"] + + indexed = QliV2CaseSelector.resolve(tmp_path, case_indexes="3,1,2-2") + assert [Path(path).name for path in indexed] == ["case_10.pt", "case_1.pt", "case_2.pt"] + + named = QliV2CaseSelector.resolve(tmp_path, case_names="case_2,case_1.pt") + assert [Path(path).name for path in named] == ["case_2.pt", "case_1.pt"] + + +def test_case_selector_rejects_ambiguous_or_invalid_selection(tmp_path): + (tmp_path / "case_1.pt").touch() + with pytest.raises(ValueError, match="cannot be specified together"): + QliV2CaseSelector.resolve(tmp_path, case_names="case_1", case_indexes="1") + with pytest.raises(ValueError, match="out of range"): + QliV2CaseSelector.resolve(tmp_path, case_indexes="2") + + +def test_result_writer_uses_readable_name_and_migrates_legacy_result(tmp_path): + params = list(range(len(PARAM_NAMES))) + params[PARAM_NAMES.index("qk_dtype")] = "INT8" + params[PARAM_NAMES.index("layout_query")] = "BSND" + params[PARAM_NAMES.index("layout_key")] = "PA_BBND" + name = QliV2ResultWriter.case_name(params) + assert name == QliV2ResultWriter.case_name(params) + assert name == ( + "QLI_B0_S11_S22_N15_N26_D7_BSND_PA_BBND_INT8_" + "QM20_SM24_CR30_K23_RV31" + ) + assert QliV2ResultWriter.case_name( + params, + explicit_name="quant li/default:a5 v2", + ) == "quant_li_default_a5_v2" + + row = QliV2ResultWriter.row(name, params, "Pass", 100.0) + output = tmp_path / "result.xlsx" + legacy = pd.DataFrame([{key: value for key, value in row.items() if key != "return_value"}]) + legacy.to_excel(output, index=False) + + QliV2ResultWriter.append(output, row) + result = pd.read_excel(output) + assert list(result.columns) == list(row.keys()) + assert len(result) == 2 + assert result.iloc[1]["return_value"] == params[PARAM_NAMES.index("return_value")] + + +def test_comparison_failure_raises_serializable_assertion(): + ensure_comparison_passed("case_pass", "Pass", 100.0) + + with pytest.raises(AssertionError, match="case_index_fail.*index result=Failed") as caught: + ensure_comparison_passed("case_index_fail", "Failed", 97.5) + restored = pickle.loads(pickle.dumps(caught.value)) + assert str(restored) == str(caught.value) + + with pytest.raises(AssertionError, match="case_value_fail.*value result=Failed"): + ensure_comparison_passed("case_value_fail", "Pass", 100.0, "Failed", 90.0) + + +def test_normalize_legacy_params_adds_weight_dtype(): + params = list(range(32)) + params[10] = "INT8" + params[11] = "FP16" + + normalized = normalize_qliv2_params(params) + + assert len(normalized) == 33 + assert normalized[:11] == tuple(params[:11]) + assert normalized[11] == params[11] + assert normalized[12:] == tuple(params[11:]) diff --git a/xllm_ops/attention/quant_lightning_indexer_v2/tests/pytest/test_quant_lightning_indexer_v2_batch.py b/xllm_ops/attention/quant_lightning_indexer_v2/tests/pytest/test_quant_lightning_indexer_v2_batch.py new file mode 100644 index 0000000..85bf391 --- /dev/null +++ b/xllm_ops/attention/quant_lightning_indexer_v2/tests/pytest/test_quant_lightning_indexer_v2_batch.py @@ -0,0 +1,192 @@ +#!/usr/bin/python +# -*- coding: utf-8 -*- +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +import concurrent.futures +import os +from pathlib import Path + +import pytest + +import result_compare_method +from batch import quant_lightning_indexer_v2_pt_loadprocess +from qliv2_test_utils import ( + QliV2CaseSelector, + QliV2ResultWriter, + ensure_comparison_passed, +) + +TEST_INPUT_PATH_ENV = os.environ.get("QLIV2_TESTCASE_DIR", "").strip() +TEST_INPUT_PATH = TEST_INPUT_PATH_ENV or "pt_path" +RESULT_PATH = Path(os.environ.get("QLIV2_RESULT_PATH", "result.xlsx").strip()) +DEVICE_ID = int(os.environ.get("QLIV2_DEVICE_ID", "0")) + +# 支持通过环境变量 QLIV2_TESTCASE_PATH 指定单条用例文件，实现进程级隔离执行： +# - 设置时：仅运行该条用例（配合 batch_isolated_run.sh 每条用例拉起独立进程） +# - 未设置：回退为原有行为，一次性加载目录下全部用例 +SINGLE_CASE_PATH = os.environ.get("QLIV2_TESTCASE_PATH", "").strip() +# flag：是否处于批量隔离模式（由 batch_isolated_run.sh 设置 QLIV2_TESTCASE_PATH 触发） +IS_ISOLATED_MODE = bool(SINGLE_CASE_PATH) +# flag：运行模式 eager / graph（通过环境变量 QLIV2_RUN_MODE 或命令行参数设置，默认 eager） +RUN_MODE = os.environ.get("QLIV2_RUN_MODE", "eager").strip().lower() +# 支持通过环境变量 QLIV2_PT_FILE_LIST 指定用例文件列表（逗号分隔），用于从 Excel 筛选的 batch_exec 模式 +PT_FILE_LIST = os.environ.get("QLIV2_PT_FILE_LIST", "").strip() +CASE_NAMES = os.environ.get("QLIV2_CASE_NAMES", "").strip() +CASE_INDEXES = os.environ.get("QLIV2_CASE_INDEXES", "").strip() + +try: + if SINGLE_CASE_PATH: + TESTCASE_FILES = QliV2CaseSelector.resolve( + TEST_INPUT_PATH, explicit_files=SINGLE_CASE_PATH + ) + print(f"单用例隔离模式, 仅执行: {SINGLE_CASE_PATH}") + else: + TESTCASE_FILES = QliV2CaseSelector.resolve( + TEST_INPUT_PATH, + explicit_files=PT_FILE_LIST, + case_names=CASE_NAMES, + case_indexes=CASE_INDEXES, + ) + print(f"找到 {len(TESTCASE_FILES)} 个测试用例文件") +except ValueError as error: + has_explicit_selection = any( + ( + TEST_INPUT_PATH_ENV, + SINGLE_CASE_PATH, + PT_FILE_LIST, + CASE_NAMES, + CASE_INDEXES, + ) + ) + if has_explicit_selection: + raise + print(f"未配置 batch PT 用例，跳过收集: {error}") + TESTCASE_FILES = [] + + +def qliv2(testcase_file): + try: + if RUN_MODE == "graph": + ( + cpu_result, + npu_result, + topk_value, + cpu_topk_value, + npu_topk_value, + output_idx_offset, + params, + ) = quant_lightning_indexer_v2_pt_loadprocess.test_qliv2_process_graph( + testcase_file, device_id=DEVICE_ID + ) + else: + ( + cpu_result, + npu_result, + topk_value, + cpu_topk_value, + npu_topk_value, + output_idx_offset, + params, + ) = quant_lightning_indexer_v2_pt_loadprocess.test_qliv2_process( + testcase_file, device_id=DEVICE_ID + ) + if npu_result is not None: + result, fulfill_percent = result_compare_method.check_result( + cpu_result, + npu_result, + topk_value, + output_idx_offset, + params, + cpu_topk_value, + npu_topk_value, + ) + else: + result = "Failed" + fulfill_percent = 0 + return_value = params[31] + if return_value: + result_return_value, fulfill_precent_return_value = ( + result_compare_method.check_result_return_value( + cpu_topk_value, + npu_topk_value, + params, + cpu_result, + npu_result, + topk_value, + output_idx_offset, + ) + ) + print(f"result_return_value: {result_return_value}") + print(f"fulfill_precent_return_value: {fulfill_precent_return_value}") + else: + result_return_value = "N/A" + fulfill_precent_return_value = 0 + except Exception as error: + print("NPU ERROR：", error) + result = "NPU ERROR" + fulfill_percent = 0 + result_return_value = "N/A" + fulfill_precent_return_value = 0 + params = [None] * 33 + + row_data = QliV2ResultWriter.row( + Path(testcase_file).stem, + params, + result, + fulfill_percent, + result_return_value, + fulfill_precent_return_value, + ) + QliV2ResultWriter.append(RESULT_PATH, row_data) + + case_name = Path(testcase_file).stem + if result != "NPU ERROR": + try: + ensure_comparison_passed( + case_name, + result, + fulfill_percent, + result_return_value, + fulfill_precent_return_value, + ) + except AssertionError as error: + return str(error) + + if result == "NPU ERROR": + return f"用例执行失败:{Path(testcase_file).stem}" + return None + + +@pytest.mark.ci +@pytest.mark.parametrize("testcase_file", TESTCASE_FILES) +def test_qliv2(testcase_file): + if IS_ISOLATED_MODE: + # 批量隔离模式：shell 层已通过独立 pytest 进程提供进程隔离，内部使用线程池即可 + with concurrent.futures.ThreadPoolExecutor(max_workers=1) as executor: + futures = executor.submit(qliv2, testcase_file) + for future in concurrent.futures.as_completed([futures]): + try: + result = future.result() + if result is not None: + pytest.fail(str(result)) + except Exception as e: + pytest.fail(f"当前用例线程执行失败：{e}") + else: + # 非隔离模式（直接 pytest）：使用子进程隔离，防止单条用例崩溃影响整体 + with concurrent.futures.ProcessPoolExecutor(max_workers=1) as executor: + future1 = executor.submit(qliv2, testcase_file) + for future in concurrent.futures.as_completed([future1]): + try: + result = future.result() + if result is not None: + pytest.fail(str(result)) + except Exception as e: + pytest.fail(f"当前用例子进程执行失败：{e}") diff --git a/xllm_ops/attention/quant_lightning_indexer_v2/tests/pytest/test_quant_lightning_indexer_v2_paramset.py b/xllm_ops/attention/quant_lightning_indexer_v2/tests/pytest/test_quant_lightning_indexer_v2_paramset.py new file mode 100644 index 0000000..690590a --- /dev/null +++ b/xllm_ops/attention/quant_lightning_indexer_v2/tests/pytest/test_quant_lightning_indexer_v2_paramset.py @@ -0,0 +1,947 @@ +#!/usr/bin/python +# -*- coding: utf-8 -*- +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +import torch +import random + +# 定义测试参数组合 +# 参数规则： +# cu_seqlens_q/k: TND 时必传 [B+1] 前缀和（首元素=0），非 TND 时为 None +# seqused_q/k: 每个 batch 的实际有效元素数 [B] +# TND 时可选（golden 可从 cu_seqlens 推导） +# BSND 时可选（None 则用 q_seq/k_seq 填满） +# PA_BBND 时 seqused_k 必传 +TEST_PARAMS = { + # Ascend950 基础场景: BSND query + PA_BBND key + "quant_li_default_a5": { + "batch_size": [8], + "q_seq": [15], + "k_seq": [111], + "q_t_size": [8], + "k_t_size": [15], # 压缩后的值 + "q_head_num": [64], + "k_head_num": [1], + "head_dim": [128], + "block_size": [512], # 取16的整数倍，最多支持到1024 + "block_num": [8], + "qk_dtype": [torch.float8_e4m3fn], + "dequant_dtype": [torch.float32], + "actual_seq_dtype": [torch.int32], + "cu_seqlens_q": [None], # BSND: cu_seqlens_q 不传 + "cu_seqlens_k": [None], # PA_BBND: cu_seqlens_k 不传 + "seqused_q": [[3, 3, 3, 3, 3, 3, 3, 3]], + "seqused_k": [[28, 24, 80, 96, 47, 76, 0, 111]], # PA场景每个batch的实际token数 + "cmp_residual_k": [None], + "max_seqlen_q": [-1], + "quant_mode": [1], + "layout_query": ["BSND"], + "layout_key": ["PA_BBND"], + "sparse_count": [512], + "sparse_mode": [3], + "query_datarange": [[-448, 448]], + "key_datarange": [[-20, 20]], + "weights_datarange": [[-123, 123]], + "q_scale_datarange": [[0, 255]], + "k_scale_datarange": [[0, 65504]], + "cmp_ratio": [1], # 1/2/4/8/16/32/64/128 + "return_value": [0], + "output_idx_offset": [None], + "run_mode": ["eager"], + }, + # Ascend950 基础场景v2: BSND query + PA_BBND key + "quant_li_default_a5_v2": { + "batch_size": [104], + "q_seq": [4], + "k_seq": [32768], + "q_t_size": [8], + "k_t_size": [15], # 压缩后的值 + "q_head_num": [64], + "k_head_num": [1], + "head_dim": [128], + "block_size": [64], # 取16的整数倍，最多支持到1024 + "block_num": [53256], + "qk_dtype": [torch.float8_e4m3fn], + "dequant_dtype": [torch.float32], + "actual_seq_dtype": [torch.int32], + "cu_seqlens_q": [None], # BSND: cu_seqlens_q 不传 + "cu_seqlens_k": [None], # PA_BBND: cu_seqlens_k 不传 + "seqused_q": [None], + "seqused_k": [[14224] * 103 + [32768]], # PA场景每个batch的实际token数 + "cmp_residual_k": [None], + "max_seqlen_q": [4], + "quant_mode": [1], + "layout_query": ["BSND"], + "layout_key": ["PA_BBND"], + "sparse_count": [512], + "sparse_mode": [0], + "query_datarange": [[2, 10]], + "key_datarange": [[-1, 1]], + "weights_datarange": [[-1, 1]], + "q_scale_datarange": [[0, 1]], + "k_scale_datarange": [[0, 1]], + "cmp_ratio": [4], # 1/2/4/8/16/32/64/128 + "return_value": [1], + "output_idx_offset": [ + [random.randint(2, 10) for _ in range(416)] for _ in range(1) + ], + }, + # Ascend950 基础场景v3: BSND query + PA_BBND key + "quant_li_default_a5_v3": { + "batch_size": [26], + "q_seq": [4], + "k_seq": [262144], + "q_t_size": [8], + "k_t_size": [15], # 压缩后的值 + "q_head_num": [64], + "k_head_num": [1], + "head_dim": [128], + "block_size": [64], # 取16的整数倍，最多支持到1024 + "block_num": [106528], + "qk_dtype": [torch.float8_e4m3fn], + "dequant_dtype": [torch.float32], + "actual_seq_dtype": [torch.int32], + "cu_seqlens_q": [None], # BSND: cu_seqlens_q 不传 + "cu_seqlens_k": [None], # PA_BBND: cu_seqlens_k 不传 + "seqused_q": [None], + "seqused_k": [[18244] * 25 + [262144]], # PA场景每个batch的实际token数 + "cmp_residual_k": [ + [ + 0, + 2, + 1, + 2, + 1, + 0, + 0, + 1, + 0, + 3, + 3, + 3, + 3, + 3, + 2, + 1, + 3, + 0, + 3, + 0, + 3, + 0, + 1, + 0, + 2, + 2, + ] + ], + "max_seqlen_q": [4], + "quant_mode": [1], + "layout_query": ["BSND"], + "layout_key": ["PA_BBND"], + "sparse_count": [512], + "sparse_mode": [3], + "query_datarange": [[-1, 1]], + "key_datarange": [[-1, 1]], + "weights_datarange": [[-1, 1]], + "q_scale_datarange": [[0, 0.5]], + "k_scale_datarange": [[0, 0.5]], + "cmp_ratio": [4], # 1/2/4/8/16/32/64/128 + "return_value": [1], + "output_idx_offset": [ + [random.randint(0, 1000) for _ in range(104)] for _ in range(1) + ], + }, + # Ascend950 基础场景v4: BSND query + PA_BBND key + "quant_li_default_a5_v4": { + "batch_size": [56], + "q_seq": [4], + "k_seq": [2048], + "q_t_size": [8], + "k_t_size": [15], # 压缩后的值 + "q_head_num": [64], + "k_head_num": [1], + "head_dim": [128], + "block_size": [64], # 取16的整数倍，最多支持到1024 + "block_num": [1792], + "qk_dtype": [torch.float8_e4m3fn], + "dequant_dtype": [torch.float32], + "actual_seq_dtype": [torch.int32], + "cu_seqlens_q": [None], # BSND: cu_seqlens_q 不传 + "cu_seqlens_k": [None], # PA_BBND: cu_seqlens_k 不传 + "seqused_q": [None], + "seqused_k": [[1804] * 55 + [2048]], # PA场景每个batch的实际token数 + "cmp_residual_k": [None], + "max_seqlen_q": [4], + "quant_mode": [1], + "layout_query": ["BSND"], + "layout_key": ["PA_BBND"], + "sparse_count": [1024], + "sparse_mode": [0], + "query_datarange": [[-1, 1]], + "key_datarange": [[0, 0.001]], + "weights_datarange": [[-20, 20]], + "q_scale_datarange": [[0, 5]], + "k_scale_datarange": [[0, 5]], + "cmp_ratio": [4], # 1/2/4/8/16/32/64/128 + "return_value": [1], + "output_idx_offset": [ + [random.randint(0, 100) for _ in range(224)] for _ in range(1) + ], + }, + # Ascend950 基础场景v5: BSND query + PA_BBND key + "quant_li_default_a5_v5": { + "batch_size": [42], + "q_seq": [4], + "k_seq": [32768], + "q_t_size": [8], + "k_t_size": [15], # 压缩后的值 + "q_head_num": [64], + "k_head_num": [1], + "head_dim": [128], + "block_size": [64], # 取16的整数倍，最多支持到1024 + "block_num": [21523], + "qk_dtype": [torch.float8_e4m3fn], + "dequant_dtype": [torch.float32], + "actual_seq_dtype": [torch.int32], + "cu_seqlens_q": [None], # BSND: cu_seqlens_q 不传 + "cu_seqlens_k": [None], # PA_BBND: cu_seqlens_k 不传 + "seqused_q": [None], + "seqused_k": [[20912] * 41 + [32768]], # PA场景每个batch的实际token数 + "cmp_residual_k": [ + [ + 2, + 1, + 3, + 0, + 3, + 0, + 0, + 0, + 3, + 1, + 1, + 2, + 2, + 2, + 3, + 2, + 1, + 2, + 1, + 2, + 1, + 2, + 2, + 0, + 2, + 1, + 2, + 3, + 1, + 0, + 0, + 0, + 3, + 2, + 0, + 3, + 0, + 2, + 2, + 1, + 2, + 0, + ] + ], + "max_seqlen_q": [4], + "quant_mode": [1], + "layout_query": ["BSND"], + "layout_key": ["PA_BBND"], + "sparse_count": [1024], + "sparse_mode": [3], + "query_datarange": [[0, 0.001]], + "key_datarange": [[0.001, 0.01]], + "weights_datarange": [[-0.5, 0.5]], + "q_scale_datarange": [[0, 1]], + "k_scale_datarange": [[0, 1]], + "cmp_ratio": [4], # 1/2/4/8/16/32/64/128 + "return_value": [1], + "output_idx_offset": [ + [random.randint(2, 10) for _ in range(168)] for _ in range(1) + ], + }, + # Ascend950 基础场景v6: BSND query + PA_BBND key + "quant_li_default_a5_v6": { + "batch_size": [13], + "q_seq": [4], + "k_seq": [262144], + "q_t_size": [8], + "k_t_size": [15], # 压缩后的值 + "q_head_num": [64], + "k_head_num": [1], + "head_dim": [128], + "block_size": [64], # 取16的整数倍，最多支持到1024 + "block_num": [53290], + "qk_dtype": [torch.float8_e4m3fn], + "dequant_dtype": [torch.float32], + "actual_seq_dtype": [torch.int32], + "cu_seqlens_q": [None], # BSND: cu_seqlens_q 不传 + "cu_seqlens_k": [None], # PA_BBND: cu_seqlens_k 不传 + "seqused_q": [None], + "seqused_k": [[201988] * 12 + [262144]], # PA场景每个batch的实际token数 + "cmp_residual_k": [[0, 2, 3, 3, 0, 1, 2, 0, 3, 2, 0, 0, 3]], + "max_seqlen_q": [4], + "quant_mode": [1], + "layout_query": ["BSND"], + "layout_key": ["PA_BBND"], + "sparse_count": [1024], + "sparse_mode": [3], + "query_datarange": [[0.001, 0.01]], + "key_datarange": [[-5, 5]], + "weights_datarange": [[-2, -1]], + "q_scale_datarange": [[10, 255]], + "k_scale_datarange": [[10, 255]], + "cmp_ratio": [4], # 1/2/4/8/16/32/64/128 + "return_value": [1], + "output_idx_offset": [ + [random.randint(0, 1000000) for _ in range(52)] for _ in range(1) + ], + }, + # Ascend950 基础场景v7: BSND query + PA_BBND key + "quant_li_default_a5_v7": { + "batch_size": [4], + "q_seq": [4096], + "k_seq": [1024], + "q_t_size": [8], + "k_t_size": [15], # 压缩后的值 + "q_head_num": [64], + "k_head_num": [1], + "head_dim": [128], + "block_size": [512], # 取16的整数倍，最多支持到1024 + "block_num": [45], + "qk_dtype": [torch.float8_e4m3fn], + "dequant_dtype": [torch.float32], + "actual_seq_dtype": [torch.int32], + "cu_seqlens_q": [None], # BSND: cu_seqlens_q 不传 + "cu_seqlens_k": [None], # PA_BBND: cu_seqlens_k 不传 + "seqused_q": [None], + "seqused_k": [[1000] * 3 + [1024]], # PA场景每个batch的实际token数 + "cmp_residual_k": [None], + "max_seqlen_q": [4096], + "quant_mode": [1], + "layout_query": ["BSND"], + "layout_key": ["PA_BBND"], + "sparse_count": [512], + "sparse_mode": [0], + "query_datarange": [[-5, 5]], + "key_datarange": [[-100, 100]], + "weights_datarange": [[-1, 1]], + "q_scale_datarange": [[0, 1]], + "k_scale_datarange": [[0, 1]], + "cmp_ratio": [4], # 1/2/4/8/16/32/64/128 + "return_value": [1], + "output_idx_offset": [ + [random.randint(0, 1000000) for _ in range(16384)] for _ in range(1) + ], + }, + # Ascend950 基础场景v8: BSND query + PA_BBND key + "quant_li_default_a5_v8": { + "batch_size": [4], + "q_seq": [4096], + "k_seq": [1024], + "q_t_size": [8], + "k_t_size": [15], # 压缩后的值 + "q_head_num": [64], + "k_head_num": [1], + "head_dim": [128], + "block_size": [512], # 取16的整数倍，最多支持到1024 + "block_num": [41], + "qk_dtype": [torch.float8_e4m3fn], + "dequant_dtype": [torch.float32], + "actual_seq_dtype": [torch.int32], + "cu_seqlens_q": [None], # BSND: cu_seqlens_q 不传 + "cu_seqlens_k": [None], # PA_BBND: cu_seqlens_k 不传 + "seqused_q": [[2736] * 3 + [4096]], + "seqused_k": [[92] * 3 + [1024]], # PA场景每个batch的实际token数 + "cmp_residual_k": [[3, 3, 2, 1]], + "max_seqlen_q": [4096], + "quant_mode": [1], + "layout_query": ["BSND"], + "layout_key": ["PA_BBND"], + "sparse_count": [1024], + "sparse_mode": [3], + "query_datarange": [[-1, 1]], + "key_datarange": [[-0.5, 0.5]], + "weights_datarange": [[-1, 1]], + "q_scale_datarange": [[0, 1]], + "k_scale_datarange": [[0, 1]], + "cmp_ratio": [4], # 1/2/4/8/16/32/64/128 + "return_value": [0], + "output_idx_offset": [ + [random.randint(0, 1) for _ in range(16384)] for _ in range(1) + ], + }, + # Ascend950 hifp8 场景: BSND query + PA_BBND key + "quant_li_default_hifp8_a5": { + "batch_size": [3], + "q_seq": [13], + "k_seq": [111], + "q_t_size": [8], + "k_t_size": [15], # 压缩后的值 + "q_head_num": [64], + "k_head_num": [1], + "head_dim": [128], + "block_size": [1024], # 取16的整数倍，最多支持到1024 + "block_num": [100], + "qk_dtype": [torch.uint8], + "dequant_dtype": [torch.float32], + "actual_seq_dtype": [torch.int32], + "cu_seqlens_q": [None], # BSND: cu_seqlens_q 不传 + "cu_seqlens_k": [None], # PA_BBND: cu_seqlens_k 不传 + "seqused_q": [[2, 5, 13]], + "seqused_k": [[2080, 2114, 1180]], + "max_seqlen_q": [-1], + "cmp_residual_k": [[3, 1, 3]], + "quant_mode": [4], + "layout_query": ["BSND"], + "layout_key": ["PA_BBND"], + "sparse_count": [512], + "sparse_mode": [3], + "query_datarange": [[-448, 448]], + "key_datarange": [[-20, 20]], + "weights_datarange": [[-123, 123]], + "q_scale_datarange": [[0, 255]], + "k_scale_datarange": [[0, 65504]], + "cmp_ratio": [4], # 1/2/4/8/16/32/64/128 + "return_value": [0], + "output_idx_offset": [None], + }, + # Ascend910_93 场景: TND query + PA_BBND key + "quant_li_default_a3": { + "batch_size": [1], + "q_seq": [1], + "k_seq": [8192], + "q_t_size": [1], + "k_t_size": [8192], # 压缩后的值 + "q_head_num": [64], + "k_head_num": [1], + "head_dim": [128], + "block_size": [1024], # 取16的整数倍，最多支持到1024 + "block_num": [17], + "qk_dtype": [torch.int8], + "dequant_dtype": [torch.float16], + "actual_seq_dtype": [torch.int32], + "cu_seqlens_q": [[0, 1]], # TND: cu_seqlens_q 必传 [B+1] + "cu_seqlens_k": [None], # PA_BBND: cu_seqlens_k 不传 + "seqused_q": [None], # TND: seqused_q 可选，None 时从 cu_seqlens 推导 + "seqused_k": [[8196]], # PA_BBND: seqused_k 必传 + "cmp_residual_k": [[1]], # cmp_ratio=4 时需要 + "quant_mode": [2], # 910_93 tiling 要求 quant_mode=2 + "layout_query": ["TND"], + "layout_key": ["PA_BBND"], + "sparse_count": [512], + "sparse_mode": [3], + "query_datarange": [[-100, 100]], + "key_datarange": [[-100, 100]], + "weights_datarange": [[-25, 25]], + "q_scale_datarange": [[0, 255]], + "k_scale_datarange": [[0, 65504]], + "cmp_ratio": [4], # 1/2/4/8/16/32/64/128 + "max_seqlen_q": [-1], + "return_value": [0], + "output_idx_offset": [None], + }, + # ==================== 白盒测试用例（针对LD+returnValue修改）==================== + # WB1: LD + return_value=0 — 验证合并分支仅 isNeedLD=true 时不输出 value + "wb_ld_rv0_bsnd_pa": { + "batch_size": [8], + "q_seq": [4], + "k_seq": [32768], + "q_t_size": [8], + "k_t_size": [15], + "q_head_num": [64], + "k_head_num": [1], + "head_dim": [128], + "block_size": [64], + "block_num": [4096], + "qk_dtype": [torch.float8_e4m3fn], + "dequant_dtype": [torch.float32], + "actual_seq_dtype": [torch.int32], + "cu_seqlens_q": [None], + "cu_seqlens_k": [None], + "seqused_q": [None], + "seqused_k": [[4096] * 7 + [32768]], + "cmp_residual_k": [None], + "max_seqlen_q": [4], + "quant_mode": [1], + "layout_query": ["BSND"], + "layout_key": ["PA_BBND"], + "sparse_count": [512], + "sparse_mode": [3], + "query_datarange": [[-1, 1]], + "key_datarange": [[-1, 1]], + "weights_datarange": [[-1, 1]], + "q_scale_datarange": [[0, 1]], + "k_scale_datarange": [[0, 1]], + "cmp_ratio": [1], + "return_value": [0], + "output_idx_offset": [None], + "run_mode": ["eager"], + }, + # WB2: non-LD + return_value=1 — k_seq=128, block_size=128 → s2BlockNum=1, 不触发LD + # 验证合并分支仅 returnValueFlag=true 时正确输出 value + "wb_nold_rv1_bsnd_pa": { + "batch_size": [2], + "q_seq": [4], + "k_seq": [128], + "q_t_size": [8], + "k_t_size": [15], + "q_head_num": [64], + "k_head_num": [1], + "head_dim": [128], + "block_size": [128], + "block_num": [2], + "qk_dtype": [torch.float8_e4m3fn], + "dequant_dtype": [torch.float32], + "actual_seq_dtype": [torch.int32], + "cu_seqlens_q": [None], + "cu_seqlens_k": [None], + "seqused_q": [None], + "seqused_k": [[128, 128]], + "cmp_residual_k": [None], + "max_seqlen_q": [4], + "quant_mode": [1], + "layout_query": ["BSND"], + "layout_key": ["PA_BBND"], + "sparse_count": [128], + "sparse_mode": [0], + "query_datarange": [[-1, 1]], + "key_datarange": [[-1, 1]], + "weights_datarange": [[-1, 1]], + "q_scale_datarange": [[0, 1]], + "k_scale_datarange": [[0, 1]], + "cmp_ratio": [1], + "return_value": [1], + "output_idx_offset": [[0] * 8], + "run_mode": ["eager"], + }, + # WB3: TND query + PA_BBND key + LD + return_value=1 + # 验证 infershape TND 分支 + ProcessLD value 输出 + "wb_ld_rv1_tnd_pa": { + "batch_size": [4], + "q_seq": [4], + "k_seq": [8192], + "q_t_size": [16], + "k_t_size": [15], + "q_head_num": [64], + "k_head_num": [1], + "head_dim": [128], + "block_size": [64], + "block_num": [512], + "qk_dtype": [torch.float8_e4m3fn], + "dequant_dtype": [torch.float32], + "actual_seq_dtype": [torch.int32], + "cu_seqlens_q": [[0, 4, 8, 12, 16]], + "cu_seqlens_k": [None], + "seqused_q": [None], + "seqused_k": [[2048, 2048, 2048, 8192]], + "cmp_residual_k": [None], + "max_seqlen_q": [4], + "quant_mode": [1], + "layout_query": ["TND"], + "layout_key": ["PA_BBND"], + "sparse_count": [512], + "sparse_mode": [3], + "query_datarange": [[-1, 1]], + "key_datarange": [[-1, 1]], + "weights_datarange": [[-1, 1]], + "q_scale_datarange": [[0, 1]], + "k_scale_datarange": [[0, 1]], + "cmp_ratio": [1], + "return_value": [1], + "output_idx_offset": [[0] * 16], + "run_mode": ["eager"], + }, + # WB4: TND query + TND key + LD + return_value=1 + # 验证 TND+TND layout 路径 + ProcessLD + "wb_ld_rv1_tnd_tnd": { + "batch_size": [3], + "q_seq": [4], + "k_seq": [3072], + "q_t_size": [12], + "k_t_size": [3072], + "q_head_num": [64], + "k_head_num": [1], + "head_dim": [128], + "block_size": [64], + "block_num": [48], + "qk_dtype": [torch.float8_e4m3fn], + "dequant_dtype": [torch.float32], + "actual_seq_dtype": [torch.int32], + "cu_seqlens_q": [[0, 4, 8, 12]], + "cu_seqlens_k": [[0, 1024, 2048, 3072]], + "seqused_q": [None], + "seqused_k": [None], + "cmp_residual_k": [None], + "max_seqlen_q": [4], + "quant_mode": [1], + "layout_query": ["TND"], + "layout_key": ["TND"], + "sparse_count": [256], + "sparse_mode": [0], + "query_datarange": [[-1, 1]], + "key_datarange": [[-1, 1]], + "weights_datarange": [[-1, 1]], + "q_scale_datarange": [[0, 1]], + "k_scale_datarange": [[0, 1]], + "cmp_ratio": [1], + "return_value": [1], + "output_idx_offset": [[0] * 12], + "run_mode": ["eager"], + }, + # WB5: BSND query + BSND key + LD + return_value=1 + # 验证非 PA key 路径 + ProcessLD + "wb_ld_rv1_bsnd_bsnd": { + "batch_size": [4], + "q_seq": [4], + "k_seq": [4096], + "q_t_size": [8], + "k_t_size": [15], + "q_head_num": [64], + "k_head_num": [1], + "head_dim": [128], + "block_size": [64], + "block_num": [256], + "qk_dtype": [torch.float8_e4m3fn], + "dequant_dtype": [torch.float32], + "actual_seq_dtype": [torch.int32], + "cu_seqlens_q": [None], + "cu_seqlens_k": [None], + "seqused_q": [None], + "seqused_k": [[4096, 4096, 4096, 4096]], + "cmp_residual_k": [None], + "max_seqlen_q": [4], + "quant_mode": [1], + "layout_query": ["BSND"], + "layout_key": ["BSND"], + "sparse_count": [512], + "sparse_mode": [0], + "query_datarange": [[-1, 1]], + "key_datarange": [[-1, 1]], + "weights_datarange": [[-1, 1]], + "q_scale_datarange": [[0, 1]], + "k_scale_datarange": [[0, 1]], + "cmp_ratio": [1], + "return_value": [1], + "output_idx_offset": [[0] * 16], + "run_mode": ["eager"], + }, + # WB6: quant_mode=4 (HIFLOAT8) + LD + return_value=1 + # 验证 HIFLOAT8 dtype + ProcessLD value 输出 + "wb_ld_rv1_hifp8": { + "batch_size": [4], + "q_seq": [4], + "k_seq": [8192], + "q_t_size": [8], + "k_t_size": [15], + "q_head_num": [64], + "k_head_num": [1], + "head_dim": [128], + "block_size": [64], + "block_num": [512], + "qk_dtype": [torch.uint8], + "dequant_dtype": [torch.float32], + "actual_seq_dtype": [torch.int32], + "cu_seqlens_q": [None], + "cu_seqlens_k": [None], + "seqused_q": [None], + "seqused_k": [[2048, 2048, 2048, 8192]], + "cmp_residual_k": [[0, 1, 2, 3]], + "max_seqlen_q": [4], + "quant_mode": [4], + "layout_query": ["BSND"], + "layout_key": ["PA_BBND"], + "sparse_count": [512], + "sparse_mode": [3], + "query_datarange": [[-448, 448]], + "key_datarange": [[-20, 20]], + "weights_datarange": [[-123, 123]], + "q_scale_datarange": [[0, 255]], + "k_scale_datarange": [[0, 65504]], + "cmp_ratio": [4], + "return_value": [1], + "output_idx_offset": [[0] * 16], + "run_mode": ["eager"], + }, + # WB7: group_size=32 (q_head_num=32) + LD + return_value=1 + # 验证 gSize=32 路径 + ProcessLD + "wb_ld_rv1_g32": { + "batch_size": [4], + "q_seq": [4], + "k_seq": [8192], + "q_t_size": [8], + "k_t_size": [15], + "q_head_num": [32], + "k_head_num": [1], + "head_dim": [128], + "block_size": [64], + "block_num": [512], + "qk_dtype": [torch.float8_e4m3fn], + "dequant_dtype": [torch.float32], + "actual_seq_dtype": [torch.int32], + "cu_seqlens_q": [None], + "cu_seqlens_k": [None], + "seqused_q": [None], + "seqused_k": [[2048, 2048, 2048, 8192]], + "cmp_residual_k": [None], + "max_seqlen_q": [4], + "quant_mode": [1], + "layout_query": ["BSND"], + "layout_key": ["PA_BBND"], + "sparse_count": [512], + "sparse_mode": [3], + "query_datarange": [[-1, 1]], + "key_datarange": [[-1, 1]], + "weights_datarange": [[-1, 1]], + "q_scale_datarange": [[0, 1]], + "k_scale_datarange": [[0, 1]], + "cmp_ratio": [1], + "return_value": [1], + "output_idx_offset": [[0] * 16], + "run_mode": ["eager"], + }, + # WB8: sparse_count=1 (最小边界) + LD + return_value=1 + # 验证 ProcessLD topk 边界 + "wb_ld_rv1_sparse1": { + "batch_size": [4], + "q_seq": [4], + "k_seq": [8192], + "q_t_size": [8], + "k_t_size": [15], + "q_head_num": [64], + "k_head_num": [1], + "head_dim": [128], + "block_size": [64], + "block_num": [512], + "qk_dtype": [torch.float8_e4m3fn], + "dequant_dtype": [torch.float32], + "actual_seq_dtype": [torch.int32], + "cu_seqlens_q": [None], + "cu_seqlens_k": [None], + "seqused_q": [None], + "seqused_k": [[2048, 2048, 2048, 8192]], + "cmp_residual_k": [None], + "max_seqlen_q": [4], + "quant_mode": [1], + "layout_query": ["BSND"], + "layout_key": ["PA_BBND"], + "sparse_count": [1], + "sparse_mode": [0], + "query_datarange": [[-1, 1]], + "key_datarange": [[-1, 1]], + "weights_datarange": [[-1, 1]], + "q_scale_datarange": [[0, 1]], + "k_scale_datarange": [[0, 1]], + "cmp_ratio": [1], + "return_value": [1], + "output_idx_offset": [[0] * 16], + "run_mode": ["eager"], + }, + # WB9: block_size=16 (最小边界) + LD + return_value=1 + # 验证 ProcessLD 对齐边界 + "wb_ld_rv1_block16": { + "batch_size": [4], + "q_seq": [4], + "k_seq": [4096], + "q_t_size": [8], + "k_t_size": [15], + "q_head_num": [64], + "k_head_num": [1], + "head_dim": [128], + "block_size": [16], + "block_num": [1024], + "qk_dtype": [torch.float8_e4m3fn], + "dequant_dtype": [torch.float32], + "actual_seq_dtype": [torch.int32], + "cu_seqlens_q": [None], + "cu_seqlens_k": [None], + "seqused_q": [None], + "seqused_k": [[1024, 1024, 1024, 4096]], + "cmp_residual_k": [None], + "max_seqlen_q": [4], + "quant_mode": [1], + "layout_query": ["BSND"], + "layout_key": ["PA_BBND"], + "sparse_count": [512], + "sparse_mode": [0], + "query_datarange": [[-1, 1]], + "key_datarange": [[-1, 1]], + "weights_datarange": [[-1, 1]], + "q_scale_datarange": [[0, 1]], + "k_scale_datarange": [[0, 1]], + "cmp_ratio": [1], + "return_value": [1], + "output_idx_offset": [[0] * 16], + "run_mode": ["eager"], + }, + # WB10: block_size=1024 (最大边界) + LD + return_value=1 + # 验证大 block_size 下 ProcessLD + "wb_ld_rv1_block1024": { + "batch_size": [4], + "q_seq": [4], + "k_seq": [32768], + "q_t_size": [8], + "k_t_size": [15], + "q_head_num": [64], + "k_head_num": [1], + "head_dim": [128], + "block_size": [1024], + "block_num": [128], + "qk_dtype": [torch.float8_e4m3fn], + "dequant_dtype": [torch.float32], + "actual_seq_dtype": [torch.int32], + "cu_seqlens_q": [None], + "cu_seqlens_k": [None], + "seqused_q": [None], + "seqused_k": [[8192, 8192, 8192, 32768]], + "cmp_residual_k": [None], + "max_seqlen_q": [4], + "quant_mode": [1], + "layout_query": ["BSND"], + "layout_key": ["PA_BBND"], + "sparse_count": [512], + "sparse_mode": [3], + "query_datarange": [[-1, 1]], + "key_datarange": [[-1, 1]], + "weights_datarange": [[-1, 1]], + "q_scale_datarange": [[0, 1]], + "k_scale_datarange": [[0, 1]], + "cmp_ratio": [1], + "return_value": [1], + "output_idx_offset": [[0] * 16], + "run_mode": ["eager"], + }, + # WB11: sparse_count=2048 (最大边界) + LD + return_value=1 + # 验证 ProcessLD topk 最大值 + "wb_ld_rv1_sparse2048": { + "batch_size": [4], + "q_seq": [4], + "k_seq": [32768], + "q_t_size": [8], + "k_t_size": [15], + "q_head_num": [64], + "k_head_num": [1], + "head_dim": [128], + "block_size": [64], + "block_num": [4096], + "qk_dtype": [torch.float8_e4m3fn], + "dequant_dtype": [torch.float32], + "actual_seq_dtype": [torch.int32], + "cu_seqlens_q": [None], + "cu_seqlens_k": [None], + "seqused_q": [None], + "seqused_k": [[8192, 8192, 8192, 32768]], + "cmp_residual_k": [None], + "max_seqlen_q": [4], + "quant_mode": [1], + "layout_query": ["BSND"], + "layout_key": ["PA_BBND"], + "sparse_count": [2048], + "sparse_mode": [0], + "query_datarange": [[-1, 1]], + "key_datarange": [[-1, 1]], + "weights_datarange": [[-1, 1]], + "q_scale_datarange": [[0, 1]], + "k_scale_datarange": [[0, 1]], + "cmp_ratio": [1], + "return_value": [1], + "output_idx_offset": [[0] * 16], + "run_mode": ["eager"], + }, + # WB12: non-LD + return_value=0 — 基线对照（两条件均为false） + "wb_nold_rv0_bsnd_pa": { + "batch_size": [2], + "q_seq": [4], + "k_seq": [128], + "q_t_size": [8], + "k_t_size": [15], + "q_head_num": [64], + "k_head_num": [1], + "head_dim": [128], + "block_size": [128], + "block_num": [2], + "qk_dtype": [torch.float8_e4m3fn], + "dequant_dtype": [torch.float32], + "actual_seq_dtype": [torch.int32], + "cu_seqlens_q": [None], + "cu_seqlens_k": [None], + "seqused_q": [None], + "seqused_k": [[128, 128]], + "cmp_residual_k": [None], + "max_seqlen_q": [4], + "quant_mode": [1], + "layout_query": ["BSND"], + "layout_key": ["PA_BBND"], + "sparse_count": [128], + "sparse_mode": [0], + "query_datarange": [[-1, 1]], + "key_datarange": [[-1, 1]], + "weights_datarange": [[-1, 1]], + "q_scale_datarange": [[0, 1]], + "k_scale_datarange": [[0, 1]], + "cmp_ratio": [1], + "return_value": [0], + "output_idx_offset": [None], + "run_mode": ["eager"], + }, +} + +# 按需选择要启用的测试参数（例如默认启用所有） +properties = torch.npu.get_device_properties() +if "Ascend910_93" in properties.name: + ENABLED_PARAMSETS = [ + ("quant_li_default_a3", TEST_PARAMS["quant_li_default_a3"]), + ] +elif "Ascend950" in properties.name: + ENABLED_PARAMSETS = [ + (name, TEST_PARAMS[name]) + for name in ( + "quant_li_default_a5_v2", + "quant_li_default_a5_v3", + "quant_li_default_a5_v4", + "quant_li_default_a5_v5", + "quant_li_default_a5_v6", + "quant_li_default_a5_v7", + "quant_li_default_a5_v8", + "quant_li_default_a5_mxfp8", + "quant_li_default_a5_mxfp4", + "quant_li_default_a5_mxfp8_bsnd", + "quant_li_default_a5_mxfp4_bsnd", + "quant_li_default_a5_mxfp8_tnd", + "quant_li_default_a5_mxfp4_tnd", + # 白盒测试用例 + "wb_ld_rv0_bsnd_pa", + "wb_nold_rv1_bsnd_pa", + "wb_ld_rv1_tnd_pa", + "wb_ld_rv1_tnd_tnd", + "wb_ld_rv1_bsnd_bsnd", + "wb_ld_rv1_hifp8", + "wb_ld_rv1_g32", + "wb_ld_rv1_sparse1", + "wb_ld_rv1_block16", + "wb_ld_rv1_block1024", + "wb_ld_rv1_sparse2048", + "wb_nold_rv0_bsnd_pa", + ) + ] + +ENABLED_PARAMS = [params for _, params in ENABLED_PARAMSETS] diff --git a/xllm_ops/attention/quant_lightning_indexer_v2/tests/pytest/test_quant_lightning_indexer_v2_single.py b/xllm_ops/attention/quant_lightning_indexer_v2/tests/pytest/test_quant_lightning_indexer_v2_single.py new file mode 100644 index 0000000..4cba557 --- /dev/null +++ b/xllm_ops/attention/quant_lightning_indexer_v2/tests/pytest/test_quant_lightning_indexer_v2_single.py @@ -0,0 +1,261 @@ +#!/usr/bin/python +# -*- coding: utf-8 -*- +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +import itertools +import os +from pathlib import Path + +import torch +import torch_npu +from test_quant_lightning_indexer_v2_paramset import ENABLED_PARAMSETS +import result_compare_method +import quant_lightning_indexer_v2_golden +import pytest +from batch import quant_lightning_indexer_v2_pt_loadprocess +from qliv2_test_utils import QliV2ResultWriter, ensure_comparison_passed + + +SAVE_PT_DIR = os.environ.get("QLIV2_SINGLE_SAVE_PT_DIR", "").strip() +RESULT_PATH = os.environ.get("QLIV2_SINGLE_RESULT_PATH", "").strip() + +param_names = [ + "batch_size", + "q_seq", + "k_seq", + "q_t_size", + "k_t_size", + "q_head_num", + "k_head_num", + "head_dim", + "block_size", + "block_num", + "qk_dtype", + "weight_dtype", + "dequant_dtype", + "actual_seq_dtype", + "cu_seqlens_q", + "cu_seqlens_k", + "seqused_q", + "seqused_k", + "cmp_residual_k", + "max_seqlen_q", + "quant_mode", + "layout_query", + "layout_key", + "sparse_count", + "sparse_mode", + "query_datarange", + "key_datarange", + "weights_datarange", + "q_scale_datarange", + "k_scale_datarange", + "cmp_ratio", + "return_value", + "output_idx_offset", + "run_mode", +] + +param_combinations = [] +for paramset_name, params in ENABLED_PARAMSETS: + param_values = [ + params.get( + name, params["dequant_dtype"] if name == "weight_dtype" else ["eager"] + ) + for name in param_names + ] + combinations = list(itertools.product(*param_values)) + for combo_index, combo in enumerate(combinations, start=1): + param_dict = dict(zip(param_names, combo)) + param_dict["case_name"] = ( + paramset_name + if len(combinations) == 1 + else f"{paramset_name}_{combo_index:03d}" + ) + param_combinations.append(param_dict) + + +@pytest.mark.ci +@pytest.mark.parametrize("param_combinations", param_combinations) +def test_qliv2(param_combinations): # Init params and tensors + batch_size = param_combinations["batch_size"] + q_seq = param_combinations["q_seq"] + k_seq = param_combinations["k_seq"] + q_t_size = param_combinations["q_t_size"] + k_t_size = param_combinations["k_t_size"] + q_head_num = param_combinations["q_head_num"] + k_head_num = param_combinations["k_head_num"] + head_dim = param_combinations["head_dim"] + block_size = param_combinations["block_size"] + block_num = param_combinations["block_num"] + qk_dtype = param_combinations["qk_dtype"] + weight_dtype = param_combinations["weight_dtype"] + dequant_dtype = param_combinations["dequant_dtype"] + actual_seq_dtype = param_combinations["actual_seq_dtype"] + cu_seqlens_q = param_combinations["cu_seqlens_q"] + cu_seqlens_k = param_combinations["cu_seqlens_k"] + seqused_q = param_combinations["seqused_q"] + seqused_k = param_combinations["seqused_k"] + cmp_residual_k = param_combinations["cmp_residual_k"] + max_seqlen_q = param_combinations["max_seqlen_q"] + quant_mode = param_combinations["quant_mode"] + layout_query = param_combinations["layout_query"] + layout_key = param_combinations["layout_key"] + sparse_count = param_combinations["sparse_count"] + sparse_mode = param_combinations["sparse_mode"] + query_datarange = param_combinations["query_datarange"] + key_datarange = param_combinations["key_datarange"] + weights_datarange = param_combinations["weights_datarange"] + q_scale_datarange = param_combinations["q_scale_datarange"] + k_scale_datarange = param_combinations["k_scale_datarange"] + cmp_ratio = param_combinations["cmp_ratio"] + return_value = param_combinations["return_value"] + output_idx_offset = param_combinations["output_idx_offset"] + run_mode = ( + os.environ.get("QLIV2_RUN_MODE", param_combinations["run_mode"]).strip().lower() + ) + torch_npu.npu.set_device(0) + test_data = ( + batch_size, + q_seq, + k_seq, + q_t_size, + k_t_size, + q_head_num, + k_head_num, + head_dim, + block_size, + block_num, + qk_dtype, + weight_dtype, + dequant_dtype, + actual_seq_dtype, + cu_seqlens_q, + cu_seqlens_k, + seqused_q, + seqused_k, + cmp_residual_k, + max_seqlen_q, + quant_mode, + layout_query, + layout_key, + sparse_count, + sparse_mode, + query_datarange, + key_datarange, + weights_datarange, + q_scale_datarange, + k_scale_datarange, + cmp_ratio, + return_value, + output_idx_offset, + ) + + case_name = QliV2ResultWriter.case_name( + test_data, + explicit_name=param_combinations["case_name"], + ) + if SAVE_PT_DIR: + case_data = quant_lightning_indexer_v2_golden.generate_qliv2_test_data( + test_data + ) + case_path = Path(SAVE_PT_DIR) / f"{case_name}.pt" + case_path.parent.mkdir(parents=True, exist_ok=True) + torch.save(case_data, case_path) + print(f"当前用例 PT 已保存: {case_path}") + if run_mode == "eager": + ( + cpu_result, + npu_result, + topk_value, + cpu_topk_value, + npu_topk_value, + output_idx_offset, + _, + ) = quant_lightning_indexer_v2_pt_loadprocess.test_qliv2_process( + case_path, device_id=0 + ) + elif run_mode == "graph": + ( + cpu_result, + npu_result, + topk_value, + cpu_topk_value, + npu_topk_value, + output_idx_offset, + _, + ) = quant_lightning_indexer_v2_pt_loadprocess.test_qliv2_process_graph( + case_path, device_id=0 + ) + else: + raise ValueError(f"unsupported run_mode: {run_mode}") + elif run_mode == "eager": + cpu_result, npu_result, topk_value, cpu_topk_value, npu_topk_value = ( + quant_lightning_indexer_v2_golden.qliv2_output_single(test_data) + ) + elif run_mode == "graph": + import quant_lightning_indexer_v2_acl_graph + + cpu_result, npu_result, topk_value, cpu_topk_value, npu_topk_value = ( + quant_lightning_indexer_v2_acl_graph.qliv2_output_acl_graph(test_data) + ) + else: + raise ValueError(f"unsupported run_mode: {run_mode}") + # print("npu_result", npu_result) + # print("cpu_result:", cpu_result) + # Compare result accuracy + result, fulfill_percent = result_compare_method.check_result( + cpu_result, + npu_result, + topk_value, + output_idx_offset, + test_data, + cpu_topk_value, + npu_topk_value, + ) + print("result", result) + print("result", fulfill_percent) + result_return_value = "N/A" + fulfill_precent_return_value = 0 + if return_value: + result_return_value, fulfill_precent_return_value = ( + result_compare_method.check_result_return_value( + cpu_topk_value, + npu_topk_value, + test_data, + cpu_result, + npu_result, + topk_value, + output_idx_offset, + ) + ) + print(f"result_return_value: {result_return_value}") + print(f"result_return_value: {fulfill_precent_return_value}") + + if RESULT_PATH: + row = QliV2ResultWriter.row( + case_name, + test_data, + result, + fulfill_percent, + result_return_value, + fulfill_precent_return_value, + ) + QliV2ResultWriter.append(RESULT_PATH, row) + print(f"当前用例结果已写入: {RESULT_PATH}") + + ensure_comparison_passed( + case_name, + result, + fulfill_percent, + result_return_value, + fulfill_precent_return_value, + ) diff --git a/xllm_ops/attention/quant_lightning_indexer_v2/tests/pytest/test_run.sh b/xllm_ops/attention/quant_lightning_indexer_v2/tests/pytest/test_run.sh new file mode 100644 index 0000000..ff7f047 --- /dev/null +++ b/xllm_ops/attention/quant_lightning_indexer_v2/tests/pytest/test_run.sh @@ -0,0 +1,226 @@ +#!/bin/bash +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +set -o pipefail + +SCRIPT_DIR=$(cd "$(dirname "$0")" && pwd) +DEFAULT_PT_PATH="$SCRIPT_DIR/pt_path" +PT_SAVE_SCRIPT="$SCRIPT_DIR/batch/quant_lightning_indexer_v2_pt_save.py" +LIST_PT_SCRIPT="$SCRIPT_DIR/batch/list_pt_from_excel.py" +BATCH_TEST_SCRIPT="$SCRIPT_DIR/test_quant_lightning_indexer_v2_batch.py" +SINGLE_TEST_SCRIPT="$SCRIPT_DIR/test_quant_lightning_indexer_v2_single.py" + +show_help() { + cat < [选项] + +命令: + single 执行 paramset 中的单用例，可保存本次实际输入 PT 和结果表 + batch 有 -E 时生成 PT 后执行；无 -E 时直接执行已有 PT + batch_exec 按 Excel 中的 Testcase_Name 筛选已有 PT，仅运行 NPU 和 compare + help 显示帮助 + +通用选项: + -M, --run-mode MODE eager|graph，默认 eager + -O, --output FILE 结果 Excel 路径；single 默认 single_result.xlsx，batch 默认 result.xlsx + +batch/batch_exec 选项: + -C, --cases NAMES 按给定顺序执行 case 名，逗号分隔，可省略 .pt + -I, --indexes INDEXES 按自然排序后的 1-based 序号执行，如 3,1,5-7 + -E, --excel FILE Excel 路径；batch 不传时跳过 PT 生成 + -S, --sheet NAME Sheet 名，默认 Sheet1 + -P, --pt-path DIR PT 生成和读取目录，默认 $DEFAULT_PT_PATH + +single 选项: + --save-pt DIR single 保存本次实际输入和 CPU golden 的目录 + +示例: + $0 single --save-pt ./single_pt -O ./result/single.xlsx + $0 batch -P ./pt_path + $0 batch -E ./excel/test_cases.xlsx -P ./pt_path -O ./result/batch.xlsx + $0 batch -P ./pt_path -I 3,1,5-7 + $0 batch_exec -E ./excel/test_cases.xlsx -P ./pt_path -M graph +EOF +} + +require_value() { + if [ -z "$2" ]; then + echo "错误: $1 缺少参数值" >&2 + exit 2 + fi +} + +validate_run_mode() { + if [ "$RUN_MODE" != "eager" ] && [ "$RUN_MODE" != "graph" ]; then + echo "错误: run mode 仅支持 eager/graph，当前值: $RUN_MODE" >&2 + exit 2 + fi +} + +run_batch_pytest() { + local explicit_files="$1" + QLIV2_TESTCASE_DIR="$PT_PATH" \ + QLIV2_PT_FILE_LIST="$explicit_files" \ + QLIV2_CASE_NAMES="$CASE_NAMES" \ + QLIV2_CASE_INDEXES="$CASE_INDEXES" \ + QLIV2_RESULT_PATH="$RESULT_PATH" \ + QLIV2_RUN_MODE="$RUN_MODE" \ + python3 -m pytest -rA -s "$BATCH_TEST_SCRIPT" -v -m ci \ + -W ignore::UserWarning -W ignore::DeprecationWarning +} + +run_single() { + echo "===== QLI_V2 single: mode=$RUN_MODE result=$RESULT_PATH =====" + QLIV2_SINGLE_SAVE_PT_DIR="$SAVE_PT_DIR" \ + QLIV2_SINGLE_RESULT_PATH="$RESULT_PATH" \ + QLIV2_RUN_MODE="$RUN_MODE" \ + python3 -m pytest -rA -s "$SINGLE_TEST_SCRIPT" -v -m ci \ + -W ignore::UserWarning -W ignore::DeprecationWarning +} + +run_batch() { + if [ -n "$EXCEL_PATH" ]; then + if [ ! -f "$EXCEL_PATH" ]; then + echo "错误: Excel 文件不存在: $EXCEL_PATH" >&2 + exit 1 + fi + echo "===== 生成 PT: excel=$EXCEL_PATH sheet=$EXCEL_SHEET output=$PT_PATH =====" + python3 "$PT_SAVE_SCRIPT" "$EXCEL_PATH" "$PT_PATH" --sheet "$EXCEL_SHEET" || exit 1 + elif [ ! -d "$PT_PATH" ]; then + echo "错误: PT 目录不存在: $PT_PATH" >&2 + exit 1 + fi + echo "===== 执行 PT: input=$PT_PATH mode=$RUN_MODE result=$RESULT_PATH =====" + run_batch_pytest "" +} + +run_batch_from_excel() { + if [ -z "$EXCEL_PATH" ]; then + echo "错误: batch_exec 必须指定 -E/--excel" >&2 + exit 2 + fi + if [ ! -f "$EXCEL_PATH" ]; then + echo "错误: Excel 文件不存在: $EXCEL_PATH" >&2 + exit 1 + fi + if [ ! -d "$PT_PATH" ]; then + echo "错误: PT 目录不存在: $PT_PATH" >&2 + exit 1 + fi + local file_list + file_list=$(python3 "$LIST_PT_SCRIPT" "$EXCEL_PATH" "$PT_PATH" --sheet "$EXCEL_SHEET") || exit 1 + echo "===== Excel 筛选后仅执行 NPU + compare: mode=$RUN_MODE result=$RESULT_PATH =====" + run_batch_pytest "$file_list" +} + +if [ $# -lt 1 ]; then + show_help + exit 2 +fi + +COMMAND="$1" +shift + +EXCEL_PATH="" +EXCEL_SHEET="Sheet1" +PT_PATH="$DEFAULT_PT_PATH" +RUN_MODE="eager" +RESULT_PATH="" +CASE_NAMES="" +CASE_INDEXES="" +SAVE_PT_DIR="" + +while [ $# -gt 0 ]; do + case "$1" in + -E|--excel) + require_value "$1" "$2" + EXCEL_PATH="$2" + shift 2 + ;; + -S|--sheet) + require_value "$1" "$2" + EXCEL_SHEET="$2" + shift 2 + ;; + -P|--pt-path) + require_value "$1" "$2" + PT_PATH="$2" + shift 2 + ;; + -M|--run-mode) + require_value "$1" "$2" + RUN_MODE="$2" + shift 2 + ;; + -O|--output) + require_value "$1" "$2" + RESULT_PATH="$2" + shift 2 + ;; + -C|--cases) + require_value "$1" "$2" + CASE_NAMES="$2" + shift 2 + ;; + -I|--indexes) + require_value "$1" "$2" + CASE_INDEXES="$2" + shift 2 + ;; + --save-pt) + require_value "$1" "$2" + SAVE_PT_DIR="$2" + shift 2 + ;; + -h|--help) + show_help + exit 0 + ;; + *) + echo "错误: 未知选项 $1" >&2 + show_help + exit 2 + ;; + esac +done + +if [ -n "$CASE_NAMES" ] && [ -n "$CASE_INDEXES" ]; then + echo "错误: --cases 和 --indexes 不能同时使用" >&2 + exit 2 +fi +if [ -z "$RESULT_PATH" ]; then + if [ "$COMMAND" = "single" ]; then + RESULT_PATH="$SCRIPT_DIR/single_result.xlsx" + else + RESULT_PATH="$SCRIPT_DIR/result.xlsx" + fi +fi +validate_run_mode + +case "$COMMAND" in + single) + run_single + ;; + batch) + run_batch + ;; + batch_exec) + run_batch_from_excel + ;; + help) + show_help + ;; + *) + echo "错误: 未知命令 $COMMAND" >&2 + show_help + exit 2 + ;; +esac diff --git a/xllm_ops/attention/quant_lightning_indexer_v2/torch_extension/__init__.py b/xllm_ops/attention/quant_lightning_indexer_v2/torch_extension/__init__.py new file mode 100644 index 0000000..764be8a --- /dev/null +++ b/xllm_ops/attention/quant_lightning_indexer_v2/torch_extension/__init__.py @@ -0,0 +1,7 @@ +__all__ = ["quant_lightning_indexer", "quant_lightning_indexer_metadata"] + +from .quant_lightning_indexer import ( + quant_lightning_indexer, + quant_lightning_indexer_metadata, +) +from . import graph_convert_quant_lightning_indexer diff --git a/xllm_ops/attention/quant_lightning_indexer_v2/torch_extension/csrc/quant_lightning_indexer.cpp b/xllm_ops/attention/quant_lightning_indexer_v2/torch_extension/csrc/quant_lightning_indexer.cpp new file mode 100644 index 0000000..2a0828e --- /dev/null +++ b/xllm_ops/attention/quant_lightning_indexer_v2/torch_extension/csrc/quant_lightning_indexer.cpp @@ -0,0 +1,207 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file quant_lightning_indexer.cpp + * \brief + */ + +#include +#include "aclnn_common.h" + +namespace op_api { +using namespace at_npu::native; + +inline TensorWrapper MakeWrapper(const at::Tensor &tensor) +{ + return {tensor, ConvertToAclDataType(tensor.scalar_type())}; +} + +inline bool IsMxQuantMode(int64_t quantMode) { return quantMode == 3 || quantMode == 5; } + +inline bool IsE8M0Tensor(const at::Tensor &tensor) { return tensor.scalar_type() == at::kFloat8_e8m0fnu; } + +inline bool IsFp4CompatibleTensor(const at::Tensor &tensor) +{ + return tensor.scalar_type() == at::kFloat4_e2m1fn_x2 || tensor.scalar_type() == at::kByte; +} + +constexpr int64_t E8M0_SCALE_PACK_NUM = 2; + +inline void FixQLIV2AclDtypes(int64_t quantMode, TensorWrapper &queryWrapper, TensorWrapper &keyWrapper, + TensorWrapper &queryScaleWrapper, TensorWrapper &keyScaleWrapper) +{ + if (quantMode == 4) { + TORCH_CHECK(queryWrapper.tensor_.scalar_type() == at::kByte, "When quant_mode is 4, query must be hifp8 type"); + TORCH_CHECK(keyWrapper.tensor_.scalar_type() == at::kByte, "When quant_mode is 4, key must be hifp8 type"); + queryWrapper.dtype = ACL_HIFLOAT8; + keyWrapper.dtype = ACL_HIFLOAT8; + return; + } + + if (quantMode == 5) { + TORCH_CHECK(IsFp4CompatibleTensor(queryWrapper.tensor_), + "When quant_mode is 5, query must be torch.float4_e2m1fn_x2 or packed torch.uint8"); + TORCH_CHECK(IsFp4CompatibleTensor(keyWrapper.tensor_), + "When quant_mode is 5, key must be torch.float4_e2m1fn_x2 or packed torch.uint8"); + queryWrapper.dtype = ACL_FLOAT4_E2M1; + keyWrapper.dtype = ACL_FLOAT4_E2M1; + } + + if (IsMxQuantMode(quantMode)) { + TORCH_CHECK(IsE8M0Tensor(queryScaleWrapper.tensor_), + "When quant_mode is 3 or 5, query_dequant_scale must be torch.float8_e8m0fnu"); + TORCH_CHECK(IsE8M0Tensor(keyScaleWrapper.tensor_), + "When quant_mode is 3 or 5, key_dequant_scale must be torch.float8_e8m0fnu"); + // Cube loads E8M0 scales in two-byte groups through a bfloat16_t view. + TORCH_CHECK(queryScaleWrapper.tensor_.storage_offset() % E8M0_SCALE_PACK_NUM == 0, + "When quant_mode is 3 or 5, query_dequant_scale storage offset must satisfy 2-element E8M0 packing " + "alignment, but got ", + queryScaleWrapper.tensor_.storage_offset()); + TORCH_CHECK(keyScaleWrapper.tensor_.storage_offset() % E8M0_SCALE_PACK_NUM == 0, + "When quant_mode is 3 or 5, key_dequant_scale storage offset must satisfy 2-element E8M0 packing " + "alignment, but got ", + keyScaleWrapper.tensor_.storage_offset()); + queryScaleWrapper.dtype = ACL_FLOAT8_E8M0; + keyScaleWrapper.dtype = ACL_FLOAT8_E8M0; + } +} + +// npu tensor max size +const int SIZE = 8; +const int DIM_0 = 0; +const int DIM_1 = 1; +const int DIM_2 = 2; +const int DIM_3 = 3; + +constexpr int64_t QLI_V2_METADATA_SIZE = 1024; + +at::Tensor QuantLightningIndexerMetadata(int64_t numHeadsQ, int64_t numHeadsK, int64_t headDim, int64_t topk, + int64_t quantMode, const c10::optional &cuSeqlensQ, + const c10::optional &cuSeqlensK, + const c10::optional &sequsedQ, + const c10::optional &sequsedK, + const c10::optional &cmpResidualK, int64_t batchSize, + int64_t maxSeqlenQ, int64_t maxSeqlenK, c10::string_view layoutQ, + c10::string_view layoutK, int64_t maskMode, int64_t cmpRatio) +{ + at::Device outputDevice = at::Device(std::string("npu")); + if (cuSeqlensQ.has_value()) { + outputDevice = cuSeqlensQ.value().device(); + } else if (cuSeqlensK.has_value()) { + outputDevice = cuSeqlensK.value().device(); + } else if (sequsedQ.has_value()) { + outputDevice = sequsedQ.value().device(); + } else if (sequsedK.has_value()) { + outputDevice = sequsedK.value().device(); + } else if (cmpResidualK.has_value()) { + outputDevice = cmpResidualK.value().device(); + } + + at::Tensor output = torch::empty({QLI_V2_METADATA_SIZE}, torch::dtype(torch::kInt32).device(outputDevice)); + auto cuSeqlensQVal = get_valid_tensor(cuSeqlensQ, outputDevice); + auto cuSeqlensKVal = get_valid_tensor(cuSeqlensK, outputDevice); + auto sequsedQVal = get_valid_tensor(sequsedQ, outputDevice); + auto sequsedKVal = get_valid_tensor(sequsedK, outputDevice); + auto cmpResidualKVal = get_valid_tensor(cmpResidualK, outputDevice); + + std::string layoutQStr = std::string(layoutQ); + std::string layoutKStr = std::string(layoutK); + char *layoutQPtr = const_cast(layoutQStr.c_str()); + char *layoutKPtr = const_cast(layoutKStr.c_str()); + + ACLNN_CMD(aclnnQuantLightningIndexerV2Metadata, cuSeqlensQVal, cuSeqlensKVal, sequsedQVal, sequsedKVal, + cmpResidualKVal, numHeadsQ, numHeadsK, headDim, topk, quantMode, batchSize, maxSeqlenQ, maxSeqlenK, + layoutQPtr, layoutKPtr, maskMode, cmpRatio, output); + return output; +} + +// 工具函数，推导输出shape +std::tuple ConstructQuantLightningIndexerOutputTensor( + const at::Tensor &query, const at::Tensor &key, int64_t sparseCount, std::string queryLayoutStr, + std::string keyLayoutStr, int64_t returnValue) +{ + at::SmallVector outputSize; + for (size_t i = 0; i < query.sizes().size(); i++) { + TORCH_CHECK(query.size(i) > 0, + "All values within query's shape should be greater " + "than 0, but shape[", + i, "] is ", query.size(i)); + } + for (size_t i = 0; i < key.sizes().size(); i++) { + TORCH_CHECK(key.size(i) > 0, + "All values within key's shape should be greater " + "than 0, but shape[", + i, "] is ", key.size(i)); + } + TORCH_CHECK(sparseCount > 0, "sparse count should be greater than 0, but now is ", sparseCount); + int64_t keyHeadNum = (keyLayoutStr == "TND") ? key.size(DIM_1) : key.size(DIM_2); + if (queryLayoutStr == "BSND") { + outputSize = {query.size(DIM_0), query.size(DIM_1), keyHeadNum, sparseCount}; + } else { + int nDimIndex = 0; + nDimIndex = (keyLayoutStr == "TND") ? DIM_1 : DIM_2; + outputSize = {query.size(DIM_0), key.size(nDimIndex), sparseCount}; + } + at::Tensor sparseIndicesOut = at::empty(outputSize, query.options().dtype(at::kInt)); + at::Tensor sparseValuesOut; + if (returnValue) { + sparseValuesOut = at::empty(outputSize, query.options().dtype(at::kBFloat16)); + } else { + sparseValuesOut = at::empty({0}, query.options().dtype(at::kBFloat16)); + } + + return std::tuple(sparseIndicesOut, sparseValuesOut); +} + +std::tuple QuantLightningIndexer( + const at::Tensor &query, const at::Tensor &key, const at::Tensor &weights, const at::Tensor &queryDequantScale, + const at::Tensor &keyDequantScale, int64_t topk, int64_t quantMode, const c10::optional &cuSeqlensQ, + const c10::optional &cuSeqlensK, const c10::optional &sequsedQ, + const c10::optional &sequsedK, const c10::optional &cmpResidualK, + const c10::optional &blockTable, const c10::optional &outputIdxOffset, + const c10::optional &metadata, int64_t maxSeqlenQ, c10::string_view layoutQ, c10::string_view layoutK, + int64_t maskMode, int64_t cmpRatio, int64_t returnValue) +{ + TORCH_CHECK(query.numel() > 0, "Tensor query is empty.") + TORCH_CHECK(key.numel() > 0, "Tensor key is empty.") + + std::string queryLayoutStr = std::string(layoutQ); + std::string keyLayoutStr = std::string(layoutK); + + // construct the output tensor + std::tuple quantLightningIndexerOutput = + ConstructQuantLightningIndexerOutputTensor(query, key, topk, queryLayoutStr, keyLayoutStr, returnValue); + at::Tensor sparseIndicesOut = std::get<0>(quantLightningIndexerOutput); + at::Tensor sparseValuesOut = std::get<1>(quantLightningIndexerOutput); + // convert str + char *queryLayoutPtr = const_cast(queryLayoutStr.c_str()); + char *keyLayoutPtr = const_cast(keyLayoutStr.c_str()); + + auto queryWrapper = MakeWrapper(query); + auto keyWrapper = MakeWrapper(key); + auto queryScaleWrapper = MakeWrapper(queryDequantScale); + auto keyScaleWrapper = MakeWrapper(keyDequantScale); + FixQLIV2AclDtypes(quantMode, queryWrapper, keyWrapper, queryScaleWrapper, keyScaleWrapper); + + ACLNN_CMD(aclnnQuantLightningIndexerV2, queryWrapper, keyWrapper, weights, queryScaleWrapper, keyScaleWrapper, + cuSeqlensQ, cuSeqlensK, sequsedQ, sequsedK, cmpResidualK, blockTable, outputIdxOffset, metadata, topk, + quantMode, maxSeqlenQ, queryLayoutPtr, keyLayoutPtr, maskMode, cmpRatio, returnValue, sparseIndicesOut, + sparseValuesOut); + + return std::tuple(sparseIndicesOut, sparseValuesOut); +} +// Bind the C++ function to Python module +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) +{ + m.def("quant_lightning_indexer_metadata", &QuantLightningIndexerMetadata, "quant_lightning_indexer_metadata"); + m.def("quant_lightning_indexer", &QuantLightningIndexer, "quant_lightning_indexer"); +} +} // namespace op_api diff --git a/xllm_ops/attention/quant_lightning_indexer_v2/torch_extension/graph_convert_quant_lightning_indexer.py b/xllm_ops/attention/quant_lightning_indexer_v2/torch_extension/graph_convert_quant_lightning_indexer.py new file mode 100644 index 0000000..e654b07 --- /dev/null +++ b/xllm_ops/attention/quant_lightning_indexer_v2/torch_extension/graph_convert_quant_lightning_indexer.py @@ -0,0 +1,72 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- +# GE Converter for Graph Mode + +try: + import torch + import torch_npu + import torchair + from torch.library import impl + from torchair._ge_concrete_graph import ge_apis as ge + from torchair.ge._ge_graph import Tensor, TensorSpec + from torchair._ge_concrete_graph.fx2ge_converter import ( + declare_supported, + register_fx_node_ge_converter, + ) + from torchair._ge_concrete_graph.supported_declaration import Support + from typing import Any, Dict, List, Tuple, Union, Callable, Optional + from torchair._ge_concrete_graph.ge_ir_pb2 import ( + GraphDef, + OpDef, + TensorDescriptor, + TensorDef, + ) + from torchair.ge._ge_graph import get_default_ge_graph, next_unique_name + from torchair.ge._ge_graph import auto_convert_to_tensor + from torchair.ge._ge_graph import DataType, TensorType + from torchair.ge._ge_graph import compat_as_bytes, compat_as_bytes_list + from torchair.ge._ge_graph import trans_to_list_list_int, trans_to_list_list_float + from torchair.ge._ge_graph import get_invalid_desc + from torchair._ge_concrete_graph.compat_ir import ge_op, IrDef + from torchair.ge import attr + + _TORCHAIR_AVAILABLE = True +except ImportError: + _TORCHAIR_AVAILABLE = False + +if _TORCHAIR_AVAILABLE: + + @register_fx_node_ge_converter( + torch.ops.cann_ops_transformer.quant_lightning_indexer_metadata.default + ) + def convert_quant_lightning_indexer_metadata( + num_heads_q: int, + num_heads_kv: int, + head_dim: int, + topk: int, + quant_mode: int, + *, + cu_seqlens_q: Optional[Tensor] = None, + cu_seqlens_k: Optional[Tensor] = None, + seqused_q: Optional[Tensor] = None, + seqused_k: Optional[Tensor] = None, + cmp_residual_k: Optional[Tensor] = None, + batch_size: Optional[int] = None, + max_seqlen_q: Optional[int] = None, + max_seqlen_k: Optional[int] = None, + layout_q: Optional[str] = None, + layout_k: Optional[str] = None, + mask_mode: Optional[int] = None, + cmp_ratio: Optional[int] = None, + meta_outputs: TensorSpec = None, + ): + raise RuntimeError( + "GE converter doesn't support op: 'quant_lightning_indexer_metadata'" + ) diff --git a/xllm_ops/attention/quant_lightning_indexer_v2/torch_extension/quant_lightning_indexer.py b/xllm_ops/attention/quant_lightning_indexer_v2/torch_extension/quant_lightning_indexer.py new file mode 100644 index 0000000..19b12ea --- /dev/null +++ b/xllm_ops/attention/quant_lightning_indexer_v2/torch_extension/quant_lightning_indexer.py @@ -0,0 +1,289 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- +from typing import Optional +import torch +from torch.library import impl +from cann_ops_transformer.op_builder import OpBuilder, get_as_library + +QLI_METADATA_SIZE = 1024 +QLI_METADATA_OP_NAME = "quant_lightning_indexer_metadata" + + +class QuantLightningIndexerOpBuilder(OpBuilder): + def __init__(self): + super(QuantLightningIndexerOpBuilder, self).__init__( + "quant_lightning_indexer", category="attention" + ) + + def sources(self): + """Path to C++ source code.""" + return ["csrc/attention/quant_lightning_indexer.cpp"] + + def schema(self) -> str: + """PyTorch operator signature.""" + return [ + "quant_lightning_indexer_metadata(int num_heads_q, int num_heads_k, int head_dim, int topk, " + "int quant_mode, *, Tensor? cu_seqlens_q=None, Tensor? cu_seqlens_k=None, Tensor? seqused_q=None," + "Tensor? seqused_k=None, Tensor? cmp_residual_k=None, int? batch_size=None, int? max_seqlen_q=None," + "int? max_seqlen_k=None, str? layout_q=None, str? layout_k=None, int? mask_mode=None, " + "int? cmp_ratio=None) -> Tensor", + "quant_lightning_indexer(Tensor query, Tensor key, Tensor weights, Tensor query_dequant_scale, " + "Tensor key_dequant_scale, int topk, int quant_mode, *, Tensor? cu_seqlens_q=None, " + "Tensor? cu_seqlens_k=None, Tensor? seqused_q=None, Tensor? seqused_k=None, Tensor? " + "cmp_residual_k = None, Tensor? block_table=None, Tensor? output_idx_offset=None, Tensor? metadata=None, " + 'int max_seqlen_q=-1, str layout_q="BSND", str layout_k="BSND", int mask_mode=0, ' + "int cmp_ratio=1, int return_value=0) -> (Tensor, Tensor)", + ] + + def register_meta(self): + """ + Registers the Meta implementation (Shape/Dtype inference). + Essential for Autograd and FakeTensor support. + """ + + @torch.library.register_fake("cann_ops_transformer::" + QLI_METADATA_OP_NAME) + def quant_lightning_indexer_metadata_meta( + num_heads_q: int, + num_heads_k: int, + head_dim: int, + topk: int, + quant_mode: int, + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_k: Optional[torch.Tensor] = None, + seqused_q: Optional[torch.Tensor] = None, + seqused_k: Optional[torch.Tensor] = None, + cmp_residual_k: Optional[torch.Tensor] = None, + batch_size: Optional[int] = None, + max_seqlen_q: Optional[int] = None, + max_seqlen_k: Optional[int] = None, + layout_q: Optional[str] = None, + layout_k: Optional[str] = None, + mask_mode: Optional[int] = None, + cmp_ratio: Optional[int] = None, + ): + return torch.empty((QLI_METADATA_SIZE), dtype=torch.int32, device="npu") + + @impl(get_as_library(), self.name, "Meta") + def quant_lightning_indexer_meta( + query, + key, + weights, + query_dequant_scale, + key_dequant_scale, + topk, + quant_mode, + *, + cu_seqlens_q=None, + cu_seqlens_k=None, + seqused_q=None, + seqused_k=None, + cmp_residual_k=None, + block_table=None, + output_idx_offset=None, + metadata=None, + max_seqlen_q=-1, + layout_q="BSND", + layout_k="BSND", + mask_mode=0, + cmp_ratio=1, + return_value=0, + ): + key_head_num = key.shape[1] if layout_k == "TND" else key.shape[2] + + if layout_q == "BSND": + sparse_indices_out = torch.empty( + [query.shape[0], query.shape[1], key_head_num, topk], + dtype=torch.int32, + device="meta", + ) + else: + sparse_indices_out = torch.empty( + [query.shape[0], key_head_num, topk], + dtype=torch.int32, + device="meta", + ) + if return_value: + if layout_q == "BSND": + sparse_values_out = torch.empty( + [query.shape[0], query.shape[1], key_head_num, topk], + dtype=torch.bfloat16, + device="meta", + ) + else: + sparse_values_out = torch.empty( + [query.shape[0], key_head_num, topk], + dtype=torch.bfloat16, + device="meta", + ) + else: + sparse_values_out = torch.empty( + [0], dtype=torch.bfloat16, device="meta" + ) + return (sparse_indices_out, sparse_values_out) + + +# Instantiate the builder +quant_lightning_indexer_op_builder = QuantLightningIndexerOpBuilder() +quant_lightning_indexer_op_builder._ensure_initialized() + + +@impl(get_as_library(), QLI_METADATA_OP_NAME, "PrivateUse1") +def quant_lightning_indexer_metadata( + num_heads_q: int, + num_heads_k: int, + head_dim: int, + topk: int, + quant_mode: int, + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_k: Optional[torch.Tensor] = None, + seqused_q: Optional[torch.Tensor] = None, + seqused_k: Optional[torch.Tensor] = None, + cmp_residual_k: Optional[torch.Tensor] = None, + batch_size: Optional[int] = None, + max_seqlen_q: Optional[int] = None, + max_seqlen_k: Optional[int] = None, + layout_q: Optional[str] = None, + layout_k: Optional[str] = None, + mask_mode: Optional[int] = None, + cmp_ratio: Optional[int] = None, +): + """ + dispatcher implementation for NPU.zhe + 'PrivateUse1' is the combine key for custom NPU backends. + """ + batch_size = 0 if batch_size is None else batch_size + max_seqlen_q = -1 if max_seqlen_q is None else max_seqlen_q + max_seqlen_k = -1 if max_seqlen_k is None else max_seqlen_k + layout_q = "BSND" if layout_q is None else layout_q + layout_k = "BSND" if layout_k is None else layout_k + mask_mode = 0 if mask_mode is None else mask_mode + cmp_ratio = 1 if cmp_ratio is None else cmp_ratio + + op_module = quant_lightning_indexer_op_builder.load() + return op_module.quant_lightning_indexer_metadata( + num_heads_q, + num_heads_k, + head_dim, + topk, + quant_mode, + cu_seqlens_q, + cu_seqlens_k, + seqused_q, + seqused_k, + cmp_residual_k, + batch_size, + max_seqlen_q, + max_seqlen_k, + layout_q, + layout_k, + mask_mode, + cmp_ratio, + ) + + +@torch.library.register_kernel("cann_ops_transformer::" + QLI_METADATA_OP_NAME, None) +def quant_lightning_indexer_metadata_fallback( + num_heads_q: int, + num_heads_k: int, + head_dim: int, + topk: int, + quant_mode: int, + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_k: Optional[torch.Tensor] = None, + seqused_q: Optional[torch.Tensor] = None, + seqused_k: Optional[torch.Tensor] = None, + cmp_residual_k: Optional[torch.Tensor] = None, + batch_size: Optional[int] = None, + max_seqlen_q: Optional[int] = None, + max_seqlen_k: Optional[int] = None, + layout_q: Optional[str] = None, + layout_k: Optional[str] = None, + mask_mode: Optional[int] = None, + cmp_ratio: Optional[int] = None, +): + # 处理所有 tensor 都为 None 的情况 + # 调用 NPU 实现 + return quant_lightning_indexer_metadata( + num_heads_q, + num_heads_k, + head_dim, + topk, + quant_mode, + cu_seqlens_q, + cu_seqlens_k, + seqused_q, + seqused_k, + cmp_residual_k, + batch_size, + max_seqlen_q, + max_seqlen_k, + layout_q, + layout_k, + mask_mode, + cmp_ratio, + ) + + +torch.compiler.allow_in_graph(quant_lightning_indexer_metadata) + + +@impl(get_as_library(), quant_lightning_indexer_op_builder.name, "PrivateUse1") +def quant_lightning_indexer( + query, + key, + weights, + query_dequant_scale, + key_dequant_scale, + topk, + quant_mode, + *, + cu_seqlens_q=None, + cu_seqlens_k=None, + seqused_q=None, + seqused_k=None, + cmp_residual_k=None, + block_table=None, + output_idx_offset=None, + metadata=None, + max_seqlen_q=-1, + layout_q="BSND", + layout_k="BSND", + mask_mode=0, + cmp_ratio=1, + return_value=0, +): + """ + dispatcher implementation for NPU.zhe + 'PrivateUse1' is the combine key for custom NPU backends. + """ + op_module = quant_lightning_indexer_op_builder.load() + return op_module.quant_lightning_indexer( + query, + key, + weights, + query_dequant_scale, + key_dequant_scale, + topk, + quant_mode, + cu_seqlens_q, + cu_seqlens_k, + seqused_q, + seqused_k, + cmp_residual_k, + block_table, + output_idx_offset, + metadata, + max_seqlen_q, + layout_q, + layout_k, + mask_mode, + cmp_ratio, + return_value, + ) diff --git a/xllm_ops/attention/quant_lightning_indexer_v2_metadata/CMakeLists.txt b/xllm_ops/attention/quant_lightning_indexer_v2_metadata/CMakeLists.txt new file mode 100644 index 0000000..fb93e0d --- /dev/null +++ b/xllm_ops/attention/quant_lightning_indexer_v2_metadata/CMakeLists.txt @@ -0,0 +1,17 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +file(GLOB CURRENT_DIRS RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*) +list(REMOVE_ITEM CURRENT_DIRS tests) +foreach(SUB_DIR ${CURRENT_DIRS}) + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/${SUB_DIR}/CMakeLists.txt") + add_subdirectory(${SUB_DIR}) + endif() +endforeach() \ No newline at end of file diff --git a/xllm_ops/attention/quant_lightning_indexer_v2_metadata/README.md b/xllm_ops/attention/quant_lightning_indexer_v2_metadata/README.md new file mode 100644 index 0000000..d5f5606 --- /dev/null +++ b/xllm_ops/attention/quant_lightning_indexer_v2_metadata/README.md @@ -0,0 +1,205 @@ +# QuantLightningIndexerV2Metadata + +## 产品支持情况 + +| 产品 | 是否支持 | +| :----------------------------------------------------------- | :------: | +|Ascend 950PR/Ascend 950DT| √ | +|Atlas A3 训练系列产品/Atlas A3 推理系列产品| √ | +|Atlas A2 训练系列产品/Atlas A2 推理系列产品| √ | +|Atlas 200I/500 A2 推理产品| × | +|Atlas 推理系列产品| × | +|Atlas 训练系列产品| × | + +## 功能说明 + +- 算子功能：`QuantLightningIndexerV2Metadata`是`QuantLightningIndexerV2`算子的前置算子，用于生成负载均衡的任务划分方案。本算子不执行实际的LightningIndexer计算，而是根据输入参数在AI CPU计算出每个AI Core应处理的计算起止范围，从而最大化计算资源的利用率，避免各Core间负载不均衡的问题。 + +## 参数说明 + ++ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +
参数名输入/输出/属性描述数据类型数据格式
cu_seqlens_q可选输入表示不同Batch中q的有效Sequence Length，shape为(B+1, )，仅layout_q为TND场景下必传，第一个值固定为0。INT32ND
cu_seqlens_k可选输入表示不同Batch中k的有效Sequence Length，shape为(B+1, )，仅layout_k为TND场景下必传，第一个值固定为0。INT32ND
seqused_q可选输入表示不同Batch中q实际参与运算的Sequence Length，shape为(B, )。INT32ND
seqused_k可选输入表示不同Batch中k实际参与运算的Sequence Length，shape为(B, )。INT32ND
cmp_residual_k可选输入表示不同Batch中k压缩后Sequence Length的余数，配合cmp_ratio实现mask和负载计算，shape为(B, )。cmp_ratio不为1且mask_mode为3场景下必传。INT32ND
num_heads_q属性表示q的head个数，当前支持[1, 64]。INT32-
num_heads_k属性表示k的head个数，当前仅支持1。INT32-
head_dim属性表示注意力头的维度，当前仅支持128。INT32-
topk属性表示从q中筛选出的关键稀疏token的个数，当前仅支持[1, 8192]。INT32-
quant_mode属性表示量化模式，当前支持1/2/3/4/5。1表示qk: fp8(e4m3) per-token-head, scale: fp32；2表示qk: int8 per-token-head, scale: fp16, w: fp16；3表示qk: mxfp8(e4m3), scale: fp8(e8m0)；4表示qk: hif8 per-tensor, scale: fp32；5表示qk: mxfp4(e2m1), scale: fp8(e8m0)。INT32-
batch_size可选属性表示Batch数量，默认值为0。INT32-
max_seqlen_q可选属性表示q的最长Sequence Length，-1表示任意可能长度，默认值为-1。INT32-
max_seqlen_k可选属性表示k的最长Sequence Length，-1表示任意可能长度，默认值为-1。INT32-
layout_q可选属性表示q的排列格式，支持BSND、TND，默认值为BSND。STRING-
layout_k可选属性表示k的排列格式，支持BSND、TND、PA_BBND，默认值为BSND。STRING-
mask_mode可选属性表示sparse模式，0表示No mask，3表示rightDownCausal模式，默认值为0。INT32-
cmp_ratio可选属性表示k的压缩率，取值范围[1, 128]，默认值为1，表示无压缩。INT32-
metadata输出表示负载均衡结果输出，shape固定为(1024, )INT32ND
+ +
    +
  • Atlas A3 训练系列产品/Atlas A3 推理系列产品 ：num_heads_q仅支持64，不支持quant_mode = 1/3/4/5，topk仅支持[1, 2048]，不支持layout_k = BSND/TND，不支持cmp_ratio在[1，128]任意取值，仅支持cmp_ratio = 1/2/4/8/16/32/64/128。
    • +
  • Atlas A2 训练系列产品/Atlas A2 推理系列产品 ：num_heads_q仅支持64，不支持quant_mode = 1/3/4/5，topk仅支持[1, 2048]，不支持layout_k = BSND/TND，不支持cmp_ratio在[1，128]任意取值，仅支持cmp_ratio = 1/2/4/8/16/32/64/128。
    • +
+ +## 约束说明 + +- QuantLightningIndexerV2Metadata算子需要与QuantLightningIndexerV2算子配套使用。 +- B（Batch）表示输入样本批量大小，q、k为配套的QuantLightningIndexerV2算子的入参，S1表示layout_q=BSND时，q shape中的S轴的大小，S2表示layout_k=BSND时，k shape中的S轴的大小。 +- 参数cu_seqlens_q、cu_seqlens_k要求其值为当前Batch与前序Batch有效token数的累加值，第一个元素固定为0，后一个元素的值必须大于等于前一个元素的值。 +- 参数seqused_q、seqused_k要求其值表示每个Batch中的有效token数。 +- 参数cmp_residual_k需满足cmp_residual_k[i] < cmp_ratio。 +- mask_mode所表示的mask模式的详细介绍见[sparse_mode参数说明](../../docs/zh/context/sparse_mode_introduction.md)。 +- 非PA场景layout_q、layout_k须相同。 +- layout_q=BSND场景 + - max_seqlen_q必须传入S1的值。 +- layout_k=BSND场景 + - max_seqlen_k必须传入S2的值。 +- layout_q=TND场景 + - cu_seqlens_q必须传入。 +- layout_k=TND场景 + - cu_seqlens_k必须传入。 +- layout_k=PA_BBND场景 + - seqused_k必须传入。 +- Batch取值规则 + - layout_q为BSND时，优先通过seqused_q的shape推导batch，seqused_q未传入则通过batch_size获取batch数。 + - layout_q为TND时，优先通过seqused_q的shape推导batch，seqused_q未传入则通过cu_seqlens_q的shape推导batch。 +- q Seqlen取值规则 + - layout_q为BSND时，优先通过seqused_q中的元素获取seqlen，seqused_q未传入则通过max_seqlen_q获取seqlen。 + - layout_q为TND时，优先通过seqused_q中的元素获取seqlen，seqused_q未传入则通过cu_seqlens_q中的元素获取seqlen。 +- k Seqlen取值规则 + - layout_k为BSND时，优先通过seqused_k中的元素获取seqlen，seqused_k未传入则通过max_seqlen_k获取seqlen。 + - layout_k为TND时，优先通过seqused_k中的元素获取seqlen，seqused_k未传入则通过cu_seqlens_k中的元素获取seqlen。 + +## 调用说明 + +| 调用方式 | 样例代码 | 说明 | +| --------- | ------------------------------------------------------------ | ------------------------------------------------------------ | +| aclnn API | [test_aclnn_quant_lightning_indexer_v2_metadata](./examples/test_aclnn_quant_lightning_indexer_v2_metadata.cpp) | 通过[aclnnQuantLightningIndexerV2Metadata](./docs/aclnnQuantLightningIndexerV2Metadata.md)接口调用QuantLightningIndexerV2Metadata算子 | +| PyTorch API | [test_torch_quant_lightning_indexer_v2_metadata](./examples/test_torch_quant_lightning_indexer_v2_metadata.py) | 通过[quant_lightning_indexer_metadata](../../torch_extension/cann_ops_transformer/docs/zh/quant_lightning_indexer.md)接口调用QuantLightningIndexerV2Metadata算子 | diff --git a/xllm_ops/attention/quant_lightning_indexer_v2_metadata/docs/aclnnQuantLightningIndexerV2Metadata.md b/xllm_ops/attention/quant_lightning_indexer_v2_metadata/docs/aclnnQuantLightningIndexerV2Metadata.md new file mode 100644 index 0000000..a9ef7b1 --- /dev/null +++ b/xllm_ops/attention/quant_lightning_indexer_v2_metadata/docs/aclnnQuantLightningIndexerV2Metadata.md @@ -0,0 +1,743 @@ +# aclnnQuantLightningIndexerV2Metadata + +[📄 查看源码](https://gitcode.com/cann/ops-transformer/tree/master/attention/quant_lightning_indexer_v2_metadata) + +## 产品支持情况 + + +- Ascend 950PR/Ascend 950DT：支持 + + +- Atlas A3 训练系列产品/Atlas A3 推理系列产品：支持 + + +- Atlas A2 训练系列产品/Atlas A2 推理系列产品：支持 + + +- Atlas 200I/500 A2 推理产品：x + + +- Atlas 推理系列产品：x + + +- Atlas 训练系列产品：x + + +## 功能说明 + +- 算子功能：`aclnnQuantLightningIndexerV2Metadata`是`aclnnQuantLightningIndexerV2`算子的前置算子，用于生成负载均衡的任务划分方案。本算子不执行实际的LightningIndexer计算，而是根据输入参数在AI CPU计算出每个AI Core应处理的计算起止范围，从而最大化计算资源的利用率，避免各Core间负载不均衡的问题。 + + **该算子不建议单独使用，建议与aclnnQuantLightningIndexerV2算子配合使用，形成完整的工作流。** + 1. 接受aclnnQuantLightningIndexerV2算子接口输入数据shape信息，包含batchSize、qSeqlen、kSeqlen、mask。通过对输入分块并模拟计算耗时，均匀分配分块到可用核上，以降低aclnnQuantLightningIndexerV2算子的整体计算耗时，并提高硬件利用率。 + 2. 分配结果输出后，后续作为输入供aclnnQuantLightningIndexerV2算子使用。 + 3. 分配结果包含每个AIC核基本块的起始点和终止点，已经每个AIV核的FD任务信息。详细内容可以参考[调用示例](#调用示例)。 + +## 函数原型 + +每个算子分为[两段式接口](../../../docs/zh/context/two_phase_api.md)，必须先调用"aclnnQuantLightningIndexerV2MetadataGetWorkspaceSize"获取workspace大小，在调用"aclnnQuantLightningIndexerV2Metadata"执行计算 + +``` cpp +aclnnStatus aclnnQuantLightningIndexerV2MetadataGetWorkspaceSize( + const aclTensor *cuSeqlensQOptional, + const aclTensor *cuSeqlensKOptional, + const aclTensor *sequsedQOptional, + const aclTensor *sequsedKOptional, + const aclTensor *cmpResidualKOptional, + int64_t numHeadsQ, + int64_t numHeadsK, + int64_t headDim, + int64_t topk, + int64_t quantMode, + int64_t batchSize, + int64_t maxSeqlenQ, + int64_t maxSeqlenK, + char *layoutQOptional, + char *layoutKOptional, + int64_t maskMode, + int64_t cmpRatio, + const aclTensor *metadata, + uint64_t *workspaceSize, + aclOpExecutor **executor) +``` + +``` cpp +aclnnStatus aclnnQuantLightningIndexerV2Metadata( + void *workspace, + uint64_t workspaceSize, + aclOpExecutor *executor, + aclrtStream stream) +``` + +## aclnnQuantLightningIndexerV2MetadataGetWorkspaceSize + +- **参数说明** + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +
参数名输入/输出描述使用说明数据类型数据格式维度(shape)非连续Tensor
cuSeqlensQOptional（aclTensor*）输入表示不同Batch中q的有效Sequence Length。
  • 支持空Tensor
  • shape固定为(B+1, )。
INT32ND1维√
cuSeqlensKOptional（aclTensor*）输入表示不同Batch中k的有效Sequence Length。
  • 支持空Tensor。
  • shape固定为(B+1, )。
INT32ND1维√
sequsedQOptional（aclTensor*）输入表示不同Batch中q实际参与运算的Sequence Length。
  • 支持空Tensor。
  • shape固定为(B, )。
INT32ND1维√
sequsedKOptional（aclTensor*）输入表示不同Batch中k实际参与运算的Sequence Length。
  • 支持空Tensor。
  • shape固定为(B, )。
INT32ND1维√
cmpResidualKOptional（aclTensor*）输入表示不同Batch中k压缩后Sequence Length的余数，配合cmpRatio实现mask和负载计算。
  • 支持空Tensor。
  • cmpRatio不为1，且mask为3场景下必传。
  • shape固定为(B, )。
INT32ND1维√
numHeadsQ（int64_t）输入表示q的head个数。当前支持[1, 64]。----
numHeadsK（int64_t）输入表示k的head个数。当前仅支持1。----
headDim（int64_t）输入注意力头的维度。当前仅支持128。----
topk（int64_t）输入表示从q中筛选出的关键稀疏token的个数。当前仅支持[1, 8192]。----
quantMode（int64_t）输入表示量化模式。
  • 当前支持1/2/3/4/5。
  • 1: qk: fp8(e4m3) per-token-head; scale: fp32。
  • 2: qk: int8 per-token-head; scale: fp16 w: fp16。
  • 3: qk: mxfp8(e4m3), scale: fp8(e8m0)。
  • 4: qk: hif8 per-tensor; scale: fp32。
  • 5: mxfp4(e2m1), scale: fp8(e8m0)。
----
batchSize（int64_t）输入表示Batch数量。
  • 支持非负数。
  • 建议值为0。
----
maxSeqlenQ（int64_t）输入表示q的最长Sequence Length。
  • 取值范围≥-1，-1表示任意可能长度。
  • 建议值为-1。
----
maxSeqlenK（int64_t）输入表示k的最长Sequence Length。
  • 取值范围≥-1，-1表示任意可能长度。
  • 建议值为-1。
----
layoutQOptional（char*）输入表示q的排列格式。
  • 支持 BSND、TND。
  • 建议值为BSND。
----
layoutKOptional（char*）输入表示k的排列格式。
  • 支持 BSND、TND、PA_BBND。
  • 建议值为BSND。
----
maskMode（int64_t）输入表示sparse模式。
  • 0: No mask。
  • 3: rightDownCausal模式的mask，对应以右顶点为划分的下三角场景。
  • 建议值为0。
----
cmpRatio（int64_t）输入表示k的压缩率。
  • 取值范围[1，128]。
  • 建议值1，表示无压缩。
----
metadata（aclTensor*）输出表示负载均衡结果输出。shape固定为(1024, )。INT32ND1维×
workspaceSize（uint64_t*）输出返回需要在Device侧申请的workspace大小。-----
executor（aclOpExecutor**）输出返回op执行器，包含了算子计算流程。-----
+ +
    + +
  • Atlas A3 训练系列产品/Atlas A3 推理系列产品 ：numHeadsQ仅支持64，不支持quantMode = 1/3/4/5，topk仅支持[1, 2048]，不支持layoutKOptional = BSND/TND，不支持cmpRatio在[1，128]任意取值，仅支持cmpRatio = 1/2/4/8/16/32/64/128。
    • + + +
  • Atlas A2 训练系列产品/Atlas A2 推理系列产品 ：numHeadsQ仅支持64，不支持quantMode = 1/3/4/5，topk仅支持[1, 2048]，不支持layoutKOptional = BSND/TND，不支持cmpRatio在[1，128]任意取值，仅支持cmpRatio = 1/2/4/8/16/32/64/128。
    • + +
+ +- **返回值：** + + 返回aclnnStatus状态码，具体参见[aclnn返回码](../../../docs/zh/context/aclnn_return_code.md)。 + + 第一段接口完成入参校验，出现以下场景时报错： + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +
返回值错误码描述
ACLNN_ERR_INNER_CREATE_EXECUTOR561101创建aclOpExecutor失败。
ACLNN_ERR_INNER_NULLPTR561103参数workspaceSize、executor是空指针，或参数cuSeqlensQOptional、cuSeqlensKOptional、sequsedQOptional、sequsedKOptional、cmpResidualKOptional进行Contiguous处理后为空指针。
ACLNN_ERR_PARAM_INVALID161002参数cuSeqlensQOptional、cuSeqlensKOptional、sequsedQOptional、sequsedKOptional、cmpResidualKOptional、numHeadsQ、numHeadsK、headDim、topk、quantMode、batchSize、maxSeqlenQ、maxSeqlenK、layoutQOptional、layoutKOptional、maskMode、cmpRatio的规格不在支持范围内。
+ +## aclnnQuantLightningIndexerV2Metadata + +- **参数说明：** + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +
参数名输入/输出描述
workspace输入在Device侧申请的workspace内存地址
workspaceSize输入在Device侧申请的workspace大小，由第一段接口aclnnQuantLightningIndexerV2MetadataGetWorkspaceSize获取
executor输入op执行器，包含了算子计算流程
stream输入指定执行任务的Stream
+ +- **返回值：** + + 返回aclnnStatus状态码，具体参见[aclnn返回码](../../../docs/zh/context/aclnn_return_code.md)。 + +## 约束说明 + + - aclnnQuantLightningIndexerV2Metadata默认确定性实现。 + - B（Batch）表示输入样本批量大小，q、k为配套的aclnnQuantLightningIndexerV2算子的入参，S1表示layoutQOptional=BSND时，q shape中的S轴的大小，S2表示layoutKOptional=BSND时，k shape中的S轴的大小。 + - 参数cuSeqlensQOptional、cuSeqlensKOptional要求其值为当前Batch与前序Batch有效token数的累加值，第一个元素固定为0，后一个元素的值必须大于等于前一个元素的值。 + - 参数sequsedQOptional、sequsedKOptional要求其值表示每个Batch中的有效token数。 + - 非PA场景layoutQOptional、layoutKOptional须相同。 + - 参数cmpResidualKOptional需满足cmpResidualKOptional[i] < cmpRatio。 + - layoutQOptional=BSND场景 + - maxSeqlenQ必须传入S1的值。 + - layoutKOptional=BSND场景 + - maxSeqlenK必须传入S2的值。 + - layoutQOptional=TND场景 + - cuSeqlensQOptional必须传入。 + - layoutKOptional=TND场景 + - cuSeqlensKOptional必须传入。 + - layoutKOptional=PA_BBND场景 + - sequsedKOptional必须传入。 + - Batch取值规则 + - layoutQOptional为BSND时，优先通过sequsedQOptional的shape推导batch，sequsedQOptional未传入则通过batchSize获取batch数。 + - layoutQOptional为TND时，优先通过sequsedQOptional的shape推导batch，sequsedQOptional未传入则通过cuSeqlensQOptional的shape推导batch。 + - q Seqlen取值规则 + - layoutQOptional为BSND时，优先通过sequsedQOptional中的元素获取seqlen，sequsedQOptional未传入则通过maxSeqlenQ获取seqlen。 + - layoutQOptional为TND时，优先通过sequsedQOptional中的元素获取seqlen，sequsedQOptional未传入则通过cuSeqlensQOptional中的元素获取seqlen。 + - k Seqlen取值规则 + - layoutKOptional为BSND时，优先通过sequsedKOptional中的元素获取seqlen，sequsedKOptional未传入则通过maxSeqlenK获取seqlen。 + - layoutKOptional为TND时，优先通过sequsedKOptional中的元素获取seqlen，sequsedKOptional未传入则通过cuSeqlensKOptional中的元素获取seqlen。 + +## 调用示例 + +示例代码如下，仅供参考，具体编译和执行过程请参考[编译与运行样例](../../../docs/zh/context/compile_and_run_sample.md)。 + +``` cpp +#include +#include +#include +#include +#include +#include +#include +#include "acl/acl.h" +#include "aclnnop/aclnn_quant_lightning_indexer_v2_metadata.h" + +#define CHECK_LOG_RET(cond, ret_val, fmt, ...) \ + do { \ + if (!(cond)) { \ + printf(fmt "\n", ##__VA_ARGS__); \ + return (ret_val); \ + } \ + } while (0) + +// 参考 quant_lightning_indexer_v2_metadata.h +constexpr uint32_t AIC_CORE_MAX_NUM = 36; +constexpr uint32_t AIV_CORE_MAX_NUM = 72; +constexpr uint32_t QLI_V2_METADATA_TOTAL_SIZE = 1024; +constexpr uint32_t QLI_V2_METADATA_SIZE = 8; +constexpr uint32_t QLD_V2_METADATA_SIZE = 8; + +// QLI Metadata Index Definitions +constexpr uint32_t QLI_V2_CORE_ENABLE_INDEX = 0; +constexpr uint32_t QLI_V2_BN2_START_INDEX = 1; +constexpr uint32_t QLI_V2_M_START_INDEX = 2; +constexpr uint32_t QLI_V2_S2_START_INDEX = 3; +constexpr uint32_t QLI_V2_BN2_END_INDEX = 4; +constexpr uint32_t QLI_V2_M_END_INDEX = 5; +constexpr uint32_t QLI_V2_S2_END_INDEX = 6; +constexpr uint32_t QLI_V2_FIRST_QLD_V2_DATA_WORKSPACE_IDX_INDEX = 7; + +// QLD Metadata Index Definitions +constexpr uint32_t QLD_V2_CORE_ENABLE_INDEX = 0; +constexpr uint32_t QLD_V2_BN2_IDX_INDEX = 1; +constexpr uint32_t QLD_V2_M_IDX_INDEX = 2; +constexpr uint32_t QLD_V2_WORKSPACE_IDX_INDEX = 3; +constexpr uint32_t QLD_V2_WORKSPACE_NUM_INDEX = 4; +constexpr uint32_t QLD_V2_M_START_INDEX = 5; +constexpr uint32_t QLD_V2_M_NUM_INDEX = 6; + +struct QliV2Metadata { + uint32_t faData[AIC_CORE_MAX_NUM][QLI_V2_METADATA_SIZE]; + uint32_t fdData[AIV_CORE_MAX_NUM][QLD_V2_METADATA_SIZE]; +}; + +struct ScopeGuard +{ + explicit ScopeGuard(std::function onExitScope) : m_exitFunc(std::move(onExitScope)), + m_isDismissed(false) {} + // 禁止拷贝 + ScopeGuard(const ScopeGuard&) = delete; + ScopeGuard& operator=(const ScopeGuard&) = delete; + + ~ScopeGuard() + { + if (!m_isDismissed) { + m_exitFunc(); + } + } + + void Dismiss() + { + m_isDismissed = true; + } + + std::function m_exitFunc; + bool m_isDismissed; +}; + +struct Tensor { + void *hostAddr { nullptr }; + void *deviceAddr { nullptr }; + aclTensor *data { nullptr }; +}; + +struct ArgScenario { + bool hasCuSeq { false }; + bool hasSeqused { false }; +}; + +struct ArgContext { + // required input + int64_t numHeadsQ { 0 }; + int64_t numHeadsK { 0 }; + int64_t headDim { 0 }; + int64_t topk { 0 }; + int64_t quantMode { 2 }; + // optional input + Tensor cuSeqlensQOptional {}; + Tensor cuSeqlensKOptional {}; + Tensor sequsedQOptional {}; + Tensor sequsedKOptional {}; + Tensor cmpResidualKOptional {}; + int64_t batchSize { 0 }; + int64_t maxSeqlenQ { 0 }; + int64_t maxSeqlenK { 0 }; + char *layoutQOptional { nullptr }; + char *layoutKOptional { nullptr }; + int64_t maskMode { 0 }; + int64_t cmpRatio { 0 }; + // output + Tensor metadata {}; +}; + +int64_t GetShapeSize(const std::vector& shape) +{ + int64_t shapeSize = 1; + for (auto i : shape) { + shapeSize *= i; + } + return shapeSize; +} + +aclnnStatus Init(int32_t deviceId, aclrtStream* stream) +{ + // 固定写法，初始化 + auto ret = aclInit(nullptr); + CHECK_LOG_RET(ret == ACL_SUCCESS, ret, "aclInit failed. ERROR: %d", ret); + ret = aclrtSetDevice(deviceId); + CHECK_LOG_RET(ret == ACL_SUCCESS, ret, "aclrtSetDevice failed. ERROR: %d", ret); + ret = aclrtCreateStream(stream); + CHECK_LOG_RET(ret == ACL_SUCCESS, ret, "aclrtCreateStream failed. ERROR: %d", ret); + return ACL_SUCCESS; +} + +void Finalize(int32_t deviceId, aclrtStream stream) +{ + aclrtDestroyStream(stream); + aclrtResetDevice(deviceId); + aclFinalize(); +} + +aclnnStatus CreateTensor(aclDataType dataType, const std::vector &shape, Tensor &tensor) +{ + auto size = GetShapeSize(shape) * aclDataTypeSize(dataType); + // 调用aclrtMallocHost申请host侧内存 + auto ret = aclrtMallocHost(&(tensor.hostAddr), size); + CHECK_LOG_RET(ret == ACL_SUCCESS, ret, "aclrtMallocHost failed. ERROR: %d", ret); + memset(tensor.hostAddr, 0, size); + // 调用aclrtMalloc申请device侧内存 + ret = aclrtMalloc(&(tensor.deviceAddr), size, ACL_MEM_MALLOC_HUGE_FIRST); + CHECK_LOG_RET(ret == ACL_SUCCESS, ret, "aclrtMalloc failed. ERROR: %d", ret); + // 调用aclCreateTensor接口创建aclTensor + tensor.data = aclCreateTensor(shape.data(), shape.size(), dataType, nullptr, 0, aclFormat::ACL_FORMAT_ND, + shape.data(), shape.size(), tensor.deviceAddr); + + // 调用aclrtMemcpy将host侧数据拷贝到device侧内存上 + ret = aclrtMemcpy(tensor.deviceAddr, size, tensor.hostAddr, size, ACL_MEMCPY_HOST_TO_DEVICE); + CHECK_LOG_RET(ret == ACL_SUCCESS, ret, "aclrtMemcpy failed. ERROR: %d", ret); + return ACL_SUCCESS; +} + +void DestroyTensor(Tensor &tensor) +{ + if (tensor.data != nullptr) { + aclDestroyTensor(tensor.data); + tensor.data = nullptr; + } + if (tensor.deviceAddr != nullptr) { + aclrtFree(tensor.deviceAddr); + tensor.deviceAddr = nullptr; + } + if (tensor.hostAddr != nullptr) { + aclrtFreeHost(tensor.hostAddr); + tensor.hostAddr = nullptr; + } +} + +void DestroyArgs(ArgContext &context) +{ + DestroyTensor(context.metadata); + DestroyTensor(context.cuSeqlensQOptional); + DestroyTensor(context.cuSeqlensKOptional); + DestroyTensor(context.sequsedQOptional); + DestroyTensor(context.sequsedKOptional); + DestroyTensor(context.cmpResidualKOptional); + + if (context.layoutQOptional != nullptr) { + free(context.layoutQOptional); + context.layoutQOptional = nullptr; + } + if (context.layoutKOptional != nullptr) { + free(context.layoutKOptional); + context.layoutKOptional = nullptr; + } +} + +aclnnStatus CreateArgs(const ArgScenario &scenario, ArgContext &context) +{ + ScopeGuard argsGuard([&] { DestroyArgs(context); }); + aclnnStatus ret; + int64_t batchSize = 4; + + context.numHeadsQ = 1; + context.numHeadsK = 1; + context.headDim = 128; + context.topk = 0; + context.quantMode = 2; // 2: per-token-head / 3: group-scaling + ret = CreateTensor(aclDataType::ACL_INT32, { QLI_V2_METADATA_TOTAL_SIZE }, context.metadata); // 1024: Fix size + CHECK_LOG_RET(ret == ACL_SUCCESS, ret, "Create meta failed. Error: %d", ret); + + context.maskMode = 0; // 0: no mask, 3: causal + context.cmpRatio = 1; // [1, 128], 1: no compress + context.layoutQOptional = (char *)malloc(sizeof(char) * 16); + context.layoutKOptional = (char *)malloc(sizeof(char) * 16); + strcpy(context.layoutQOptional, "BSND"); // BSND,TND + strcpy(context.layoutKOptional, "BSND"); // BSND,TND,PA_BBND + + if (!scenario.hasCuSeq && !scenario.hasSeqused) { + context.batchSize = batchSize; + context.maxSeqlenK = 1024; + context.maxSeqlenQ = 1024; + argsGuard.Dismiss(); + return ACL_SUCCESS; + } + + if (scenario.hasCuSeq) { + // (B+1,), first element is always 0 + ret = CreateTensor(aclDataType::ACL_INT32, { batchSize + 1 }, context.cuSeqlensQOptional); + CHECK_LOG_RET(ret == ACL_SUCCESS, ret, "Create cuSeqlensQOptional failed. Error: %d", ret); + ret = CreateTensor(aclDataType::ACL_INT32, { batchSize + 1 }, context.cuSeqlensKOptional); + CHECK_LOG_RET(ret == ACL_SUCCESS, ret, "Create cuSeqlensKOptional failed. Error: %d", ret); + } + + if (scenario.hasSeqused) { + // (B,) + ret = CreateTensor(aclDataType::ACL_INT32, { batchSize }, context.sequsedQOptional); + CHECK_LOG_RET(ret == ACL_SUCCESS, ret, "Create sequsedQOptional failed. Error: %d", ret); + ret = CreateTensor(aclDataType::ACL_INT32, { batchSize }, context.sequsedKOptional); + CHECK_LOG_RET(ret == ACL_SUCCESS, ret, "Create sequsedKOptional failed. Error: %d", ret); + } + + argsGuard.Dismiss(); + return ACL_SUCCESS; +} + +int main() { + // 1.（固定写法）device/stream初始化，参考对外接口列表 + // 根据自己的实际device填写deviceId + int32_t deviceId = 0; + aclrtStream stream; + auto ret = Init(deviceId, &stream); + CHECK_LOG_RET(ret == ACL_SUCCESS, ret, "Init acl failed. ERROR: %d", ret); + ScopeGuard sysGuard([&] { Finalize(deviceId, stream); }); + + // 2. 构造输入与输出，需要根据API的接口定义构造 + ArgScenario scenario {}; + scenario.hasCuSeq = true; + scenario.hasSeqused = true; + ArgContext context {}; + ret = CreateArgs(scenario, context); + CHECK_LOG_RET(ret == ACL_SUCCESS, ret, "Create input arguments failed. ERROR: %d", ret); + ScopeGuard argsGuard([&] { DestroyArgs(context); }); + + // 3. 调用CANN算子库API，需要修改为具体的API + // 调用aclnnQuantLightningIndexerV2Metadata第一段接口 + uint64_t workspaceSize = 0; + aclOpExecutor *executor = nullptr; + void *workspaceAddr = nullptr; + ret = aclnnQuantLightningIndexerV2MetadataGetWorkspaceSize( + context.cuSeqlensQOptional.data, context.cuSeqlensKOptional.data, context.sequsedQOptional.data, + context.sequsedKOptional.data, context.cmpResidualKOptional.data, + context.numHeadsQ, context.numHeadsK, context.headDim, context.topk, context.quantMode, + context.batchSize, context.maxSeqlenQ, context.maxSeqlenK, context.layoutQOptional, + context.layoutKOptional, context.maskMode, context.cmpRatio, + context.metadata.data, &workspaceSize, &executor); + CHECK_LOG_RET(ret == ACL_SUCCESS, ret, + "aclnnQuantLightningIndexerV2MetadataGetWorkspaceSize failed. ERROR: %d\n", ret); + + if (workspaceSize > static_cast(0)) { + ret = aclrtMalloc(&workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST); + CHECK_LOG_RET(ret == ACL_SUCCESS, ret, "allocate workspace failed. ERROR: %d\n", ret); + } + ScopeGuard workspaceGuard([&] { + if (workspaceAddr != nullptr) { + aclrtFree(workspaceAddr); + workspaceAddr = nullptr; + } + }); + + // 调用aclnnQuantLightningIndexerV2Metadata第二段接口 + ret = aclnnQuantLightningIndexerV2Metadata(workspaceAddr, workspaceSize, executor, stream); + CHECK_LOG_RET(ret == ACL_SUCCESS, ret, "aclnnQuantLightningIndexerV2Metadata failed. ERROR: %d\n", ret); + + // 4.（固定写法）同步等待任务执行结束 + ret = aclrtSynchronizeStream(stream); + CHECK_LOG_RET(ret == ACL_SUCCESS, ret, "aclrtSynchronizeStream failed. ERROR: %d\n", ret); + + // 5. 打印输出 + QliV2Metadata result {}; + ret = aclrtMemcpy(&result, sizeof(result), context.metadata.deviceAddr, sizeof(result), ACL_MEMCPY_DEVICE_TO_HOST); + CHECK_LOG_RET(ret == ACL_SUCCESS, ret, "aclrtMemcpy failed. ERROR: %d\n", ret); + + for (uint32_t i = 0; i < AIC_CORE_MAX_NUM; ++i) { + printf("AIC Core%u\n", i); + printf(" Core Enable : %u\n", result.faData[i][QLI_V2_CORE_ENABLE_INDEX]); + printf(" Start BN2 : %u\n", result.faData[i][QLI_V2_BN2_START_INDEX]); + printf(" Start M : %u\n", result.faData[i][QLI_V2_M_START_INDEX]); + printf(" Start S2 : %u\n", result.faData[i][QLI_V2_S2_START_INDEX]); + printf(" End BN2 : %u\n", result.faData[i][QLI_V2_BN2_END_INDEX]); + printf(" End M : %u\n", result.faData[i][QLI_V2_M_END_INDEX]); + printf(" End S2 : %u\n", result.faData[i][QLI_V2_S2_END_INDEX]); + printf(" First Worksapce Index : %u\n", result.faData[i][QLI_V2_FIRST_QLD_V2_DATA_WORKSPACE_IDX_INDEX]); + } + for (uint32_t i = 0; i < AIV_CORE_MAX_NUM; ++i) { + printf("AIV Core%u\n", i); + printf(" Core Enable : %u\n", result.fdData[i][QLD_V2_CORE_ENABLE_INDEX]); + printf(" FD Task BN2 Idx : %u\n", result.fdData[i][QLD_V2_BN2_IDX_INDEX]); + printf(" FD Task M Idx : %u\n", result.fdData[i][QLD_V2_M_IDX_INDEX]); + printf(" FD Task S2 Idx : %u\n", result.fdData[i][QLD_V2_WORKSPACE_IDX_INDEX]); + printf(" FD Task Workspace Num : %u\n", result.fdData[i][QLD_V2_WORKSPACE_NUM_INDEX]); + printf(" FD Subtask M Start : %u\n", result.fdData[i][QLD_V2_M_START_INDEX]); + printf(" FD Subtask M Num : %u\n", result.fdData[i][QLD_V2_M_NUM_INDEX]); + } + + return 0; +} +``` diff --git a/xllm_ops/attention/quant_lightning_indexer_v2_metadata/examples/test_aclnn_quant_lightning_indexer_v2_metadata.cpp b/xllm_ops/attention/quant_lightning_indexer_v2_metadata/examples/test_aclnn_quant_lightning_indexer_v2_metadata.cpp new file mode 100644 index 0000000..8836a71 --- /dev/null +++ b/xllm_ops/attention/quant_lightning_indexer_v2_metadata/examples/test_aclnn_quant_lightning_indexer_v2_metadata.cpp @@ -0,0 +1,334 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/** + * @file test_aclnn_quant_lightning_indexer_v2_metadata.cpp + */ +#include +#include +#include +#include +#include +#include +#include +#include "acl/acl.h" +#include "aclnnop/aclnn_quant_lightning_indexer_v2_metadata.h" + +#define CHECK_LOG_RET(cond, ret_val, fmt, ...) \ + do { \ + if (!(cond)) { \ + printf(fmt "\n", ##__VA_ARGS__); \ + return (ret_val); \ + } \ + } while (0) + +// 参考 quant_lightning_indexer_v2_metadata.h +constexpr uint32_t AIC_CORE_MAX_NUM = 36; +constexpr uint32_t AIV_CORE_MAX_NUM = 72; +constexpr uint32_t QLI_V2_METADATA_TOTAL_SIZE = 1024; +constexpr uint32_t QLI_V2_METADATA_SIZE = 8; +constexpr uint32_t QLD_V2_METADATA_SIZE = 8; + +// QLI Metadata Index Definitions +constexpr uint32_t QLI_V2_CORE_ENABLE_INDEX = 0; +constexpr uint32_t QLI_V2_BN2_START_INDEX = 1; +constexpr uint32_t QLI_V2_M_START_INDEX = 2; +constexpr uint32_t QLI_V2_S2_START_INDEX = 3; +constexpr uint32_t QLI_V2_BN2_END_INDEX = 4; +constexpr uint32_t QLI_V2_M_END_INDEX = 5; +constexpr uint32_t QLI_V2_S2_END_INDEX = 6; +constexpr uint32_t QLI_V2_FIRST_QLD_V2_DATA_WORKSPACE_IDX_INDEX = 7; + +// QLD Metadata Index Definitions +constexpr uint32_t QLD_V2_CORE_ENABLE_INDEX = 0; +constexpr uint32_t QLD_V2_BN2_IDX_INDEX = 1; +constexpr uint32_t QLD_V2_M_IDX_INDEX = 2; +constexpr uint32_t QLD_V2_WORKSPACE_IDX_INDEX = 3; +constexpr uint32_t QLD_V2_WORKSPACE_NUM_INDEX = 4; +constexpr uint32_t QLD_V2_M_START_INDEX = 5; +constexpr uint32_t QLD_V2_M_NUM_INDEX = 6; + +struct QliV2Metadata { + uint32_t faData[AIC_CORE_MAX_NUM][QLI_V2_METADATA_SIZE]; + uint32_t fdData[AIV_CORE_MAX_NUM][QLD_V2_METADATA_SIZE]; +}; + +struct ScopeGuard +{ + explicit ScopeGuard(std::function onExitScope) : m_exitFunc(std::move(onExitScope)), + m_isDismissed(false) {} + // 禁止拷贝 + ScopeGuard(const ScopeGuard&) = delete; + ScopeGuard& operator=(const ScopeGuard&) = delete; + + ~ScopeGuard() + { + if (!m_isDismissed) { + m_exitFunc(); + } + } + + void Dismiss() + { + m_isDismissed = true; + } + + std::function m_exitFunc; + bool m_isDismissed; +}; + +struct Tensor { + void *hostAddr { nullptr }; + void *deviceAddr { nullptr }; + aclTensor *data { nullptr }; +}; + +struct ArgScenario { + bool hasCuSeq { false }; + bool hasSeqused { false }; +}; + +struct ArgContext { + // required input + int64_t numHeadsQ { 0 }; + int64_t numHeadsK { 0 }; + int64_t headDim { 0 }; + int64_t topk { 0 }; + int64_t quantMode { 2 }; + // optional input + Tensor cuSeqlensQOptional {}; + Tensor cuSeqlensKOptional {}; + Tensor sequsedQOptional {}; + Tensor sequsedKOptional {}; + Tensor cmpResidualKOptional {}; + int64_t batchSize { 0 }; + int64_t maxSeqlenQ { 0 }; + int64_t maxSeqlenK { 0 }; + char *layoutQOptional { nullptr }; + char *layoutKOptional { nullptr }; + int64_t maskMode { 0 }; + int64_t cmpRatio { 0 }; + // output + Tensor metadata {}; +}; + +int64_t GetShapeSize(const std::vector& shape) +{ + int64_t shapeSize = 1; + for (auto i : shape) { + shapeSize *= i; + } + return shapeSize; +} + +aclnnStatus Init(int32_t deviceId, aclrtStream* stream) +{ + // 固定写法，初始化 + auto ret = aclInit(nullptr); + CHECK_LOG_RET(ret == ACL_SUCCESS, ret, "aclInit failed. ERROR: %d", ret); + ret = aclrtSetDevice(deviceId); + CHECK_LOG_RET(ret == ACL_SUCCESS, ret, "aclrtSetDevice failed. ERROR: %d", ret); + ret = aclrtCreateStream(stream); + CHECK_LOG_RET(ret == ACL_SUCCESS, ret, "aclrtCreateStream failed. ERROR: %d", ret); + return ACL_SUCCESS; +} + +void Finalize(int32_t deviceId, aclrtStream stream) +{ + aclrtDestroyStream(stream); + aclrtResetDevice(deviceId); + aclFinalize(); +} + +aclnnStatus CreateTensor(aclDataType dataType, const std::vector &shape, Tensor &tensor) +{ + auto size = GetShapeSize(shape) * aclDataTypeSize(dataType); + // 调用aclrtMallocHost申请host侧内存 + auto ret = aclrtMallocHost(&(tensor.hostAddr), size); + CHECK_LOG_RET(ret == ACL_SUCCESS, ret, "aclrtMallocHost failed. ERROR: %d", ret); + memset(tensor.hostAddr, 0, size); + // 调用aclrtMalloc申请device侧内存 + ret = aclrtMalloc(&(tensor.deviceAddr), size, ACL_MEM_MALLOC_HUGE_FIRST); + CHECK_LOG_RET(ret == ACL_SUCCESS, ret, "aclrtMalloc failed. ERROR: %d", ret); + // 调用aclCreateTensor接口创建aclTensor + tensor.data = aclCreateTensor(shape.data(), shape.size(), dataType, nullptr, 0, aclFormat::ACL_FORMAT_ND, + shape.data(), shape.size(), tensor.deviceAddr); + + // 调用aclrtMemcpy将host侧数据拷贝到device侧内存上 + ret = aclrtMemcpy(tensor.deviceAddr, size, tensor.hostAddr, size, ACL_MEMCPY_HOST_TO_DEVICE); + CHECK_LOG_RET(ret == ACL_SUCCESS, ret, "aclrtMemcpy failed. ERROR: %d", ret); + return ACL_SUCCESS; +} + +void DestroyTensor(Tensor &tensor) +{ + if (tensor.data != nullptr) { + aclDestroyTensor(tensor.data); + tensor.data = nullptr; + } + if (tensor.deviceAddr != nullptr) { + aclrtFree(tensor.deviceAddr); + tensor.deviceAddr = nullptr; + } + if (tensor.hostAddr != nullptr) { + aclrtFreeHost(tensor.hostAddr); + tensor.hostAddr = nullptr; + } +} + +void DestroyArgs(ArgContext &context) +{ + DestroyTensor(context.metadata); + DestroyTensor(context.cuSeqlensQOptional); + DestroyTensor(context.cuSeqlensKOptional); + DestroyTensor(context.sequsedQOptional); + DestroyTensor(context.sequsedKOptional); + DestroyTensor(context.cmpResidualKOptional); + + if (context.layoutQOptional != nullptr) { + free(context.layoutQOptional); + context.layoutQOptional = nullptr; + } + if (context.layoutKOptional != nullptr) { + free(context.layoutKOptional); + context.layoutKOptional = nullptr; + } +} + +aclnnStatus CreateArgs(const ArgScenario &scenario, ArgContext &context) +{ + ScopeGuard argsGuard([&] { DestroyArgs(context); }); + aclnnStatus ret; + int64_t batchSize = 4; + + context.numHeadsQ = 1; + context.numHeadsK = 1; + context.headDim = 128; + context.topk = 0; + context.quantMode = 2; // 2: per-token-head / 3: group-scaling + ret = CreateTensor(aclDataType::ACL_INT32, { QLI_V2_METADATA_TOTAL_SIZE }, context.metadata); // 1024: Fix size + CHECK_LOG_RET(ret == ACL_SUCCESS, ret, "Create meta failed. Error: %d", ret); + + context.maskMode = 0; // 0: no mask, 3: causal + context.cmpRatio = 1; // [1, 128], 1: no compress + context.layoutQOptional = (char *)malloc(sizeof(char) * 16); + context.layoutKOptional = (char *)malloc(sizeof(char) * 16); + strcpy(context.layoutQOptional, "BSND"); // BSND,TND + strcpy(context.layoutKOptional, "BSND"); // BSND,TND,PA_BBND + + if (!scenario.hasCuSeq && !scenario.hasSeqused) { + context.batchSize = batchSize; + context.maxSeqlenK = 1024; + context.maxSeqlenQ = 1024; + argsGuard.Dismiss(); + return ACL_SUCCESS; + } + + if (scenario.hasCuSeq) { + // (B+1,), first element is always 0 + ret = CreateTensor(aclDataType::ACL_INT32, { batchSize + 1 }, context.cuSeqlensQOptional); + CHECK_LOG_RET(ret == ACL_SUCCESS, ret, "Create cuSeqlensQOptional failed. Error: %d", ret); + ret = CreateTensor(aclDataType::ACL_INT32, { batchSize + 1 }, context.cuSeqlensKOptional); + CHECK_LOG_RET(ret == ACL_SUCCESS, ret, "Create cuSeqlensKOptional failed. Error: %d", ret); + } + + if (scenario.hasSeqused) { + // (B,) + ret = CreateTensor(aclDataType::ACL_INT32, { batchSize }, context.sequsedQOptional); + CHECK_LOG_RET(ret == ACL_SUCCESS, ret, "Create sequsedQOptional failed. Error: %d", ret); + ret = CreateTensor(aclDataType::ACL_INT32, { batchSize }, context.sequsedKOptional); + CHECK_LOG_RET(ret == ACL_SUCCESS, ret, "Create sequsedKOptional failed. Error: %d", ret); + } + + argsGuard.Dismiss(); + return ACL_SUCCESS; +} + +int main() { + // 1. （固定写法）device/stream初始化，参考对外接口列表 + // 根据自己的实际device填写deviceId + int32_t deviceId = 0; + aclrtStream stream; + auto ret = Init(deviceId, &stream); + CHECK_LOG_RET(ret == ACL_SUCCESS, ret, "Init acl failed. ERROR: %d", ret); + ScopeGuard sysGuard([&] { Finalize(deviceId, stream); }); + + // 2. 构造输入与输出，需要根据API的接口定义构造 + ArgScenario scenario {}; + scenario.hasCuSeq = true; + scenario.hasSeqused = true; + ArgContext context {}; + ret = CreateArgs(scenario, context); + CHECK_LOG_RET(ret == ACL_SUCCESS, ret, "Create input arguments failed. ERROR: %d", ret); + ScopeGuard argsGuard([&] { DestroyArgs(context); }); + + // 3. 调用CANN算子库API，需要修改为具体的API + // 调用aclnnLightningIndexerV2Metadata第一段接口 + uint64_t workspaceSize = 0; + aclOpExecutor *executor = nullptr; + void *workspaceAddr = nullptr; + ret = aclnnQuantLightningIndexerV2MetadataGetWorkspaceSize( + context.cuSeqlensQOptional.data, context.cuSeqlensKOptional.data, context.sequsedQOptional.data, + context.sequsedKOptional.data, context.cmpResidualKOptional.data, + context.numHeadsQ, context.numHeadsK, context.headDim, context.topk, context.quantMode, + context.batchSize, context.maxSeqlenQ, context.maxSeqlenK, context.layoutQOptional, + context.layoutKOptional, context.maskMode, context.cmpRatio, + context.metadata.data, &workspaceSize, &executor); + CHECK_LOG_RET(ret == ACL_SUCCESS, ret, + "aclnnQuantLightningIndexerV2MetadataGetWorkspaceSize failed. ERROR: %d\n", ret); + + if (workspaceSize > static_cast(0)) { + ret = aclrtMalloc(&workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST); + CHECK_LOG_RET(ret == ACL_SUCCESS, ret, "allocate workspace failed. ERROR: %d\n", ret); + } + ScopeGuard workspaceGuard([&] { + if (workspaceAddr != nullptr) { + aclrtFree(workspaceAddr); + workspaceAddr = nullptr; + } + }); + + // 调用aclnnLightningIndexerV2Metadata第二段接口 + ret = aclnnQuantLightningIndexerV2Metadata(workspaceAddr, workspaceSize, executor, stream); + CHECK_LOG_RET(ret == ACL_SUCCESS, ret, "aclnnQuantLightningIndexerV2Metadata failed. ERROR: %d\n", ret); + + // 4. （固定写法）同步等待任务执行结束 + ret = aclrtSynchronizeStream(stream); + CHECK_LOG_RET(ret == ACL_SUCCESS, ret, "aclrtSynchronizeStream failed. ERROR: %d\n", ret); + + // 5. 打印输出 + QliV2Metadata result {}; + ret = aclrtMemcpy(&result, sizeof(result), context.metadata.deviceAddr, sizeof(result), ACL_MEMCPY_DEVICE_TO_HOST); + CHECK_LOG_RET(ret == ACL_SUCCESS, ret, "aclrtMemcpy failed. ERROR: %d\n", ret); + + for (uint32_t i = 0; i < AIC_CORE_MAX_NUM; ++i) { + printf("AIC Core%u\n", i); + printf(" Core Enable : %u\n", result.faData[i][QLI_V2_CORE_ENABLE_INDEX]); + printf(" Start BN2 : %u\n", result.faData[i][QLI_V2_BN2_START_INDEX]); + printf(" Start M : %u\n", result.faData[i][QLI_V2_M_START_INDEX]); + printf(" Start S2 : %u\n", result.faData[i][QLI_V2_S2_START_INDEX]); + printf(" End BN2 : %u\n", result.faData[i][QLI_V2_BN2_END_INDEX]); + printf(" End M : %u\n", result.faData[i][QLI_V2_M_END_INDEX]); + printf(" End S2 : %u\n", result.faData[i][QLI_V2_S2_END_INDEX]); + printf(" First Worksapce Index : %u\n", result.faData[i][QLI_V2_FIRST_QLD_V2_DATA_WORKSPACE_IDX_INDEX]); + } + for (uint32_t i = 0; i < AIV_CORE_MAX_NUM; ++i) { + printf("AIV Core%u\n", i); + printf(" Core Enable : %u\n", result.fdData[i][QLD_V2_CORE_ENABLE_INDEX]); + printf(" FD Task BN2 Idx : %u\n", result.fdData[i][QLD_V2_BN2_IDX_INDEX]); + printf(" FD Task M Idx : %u\n", result.fdData[i][QLD_V2_M_IDX_INDEX]); + printf(" FD Task S2 Idx : %u\n", result.fdData[i][QLD_V2_WORKSPACE_IDX_INDEX]); + printf(" FD Task Workspace Num : %u\n", result.fdData[i][QLD_V2_WORKSPACE_NUM_INDEX]); + printf(" FD Subtask M Start : %u\n", result.fdData[i][QLD_V2_M_START_INDEX]); + printf(" FD Subtask M Num : %u\n", result.fdData[i][QLD_V2_M_NUM_INDEX]); + } + + return 0; +} \ No newline at end of file diff --git a/xllm_ops/attention/quant_lightning_indexer_v2_metadata/examples/test_torch_quant_lightning_indexer_v2_metadata.py b/xllm_ops/attention/quant_lightning_indexer_v2_metadata/examples/test_torch_quant_lightning_indexer_v2_metadata.py new file mode 100644 index 0000000..e88d2b6 --- /dev/null +++ b/xllm_ops/attention/quant_lightning_indexer_v2_metadata/examples/test_torch_quant_lightning_indexer_v2_metadata.py @@ -0,0 +1,36 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +import torch +import torch_npu +import torchair +import numpy as np +import torch.nn as nn +import cann_ops_transformer + +metadata = torch.ops.cann_ops_transformer.quant_lightning_indexer_metadata( + cu_seqlens_q = None, + cu_seqlens_k = None, + seqused_q = None, + seqused_k = None, + cmp_residual_k = torch.tensor([3,3,3,3,3,3,3,3], dtype=torch.int32).npu(), + batch_size = 8, + max_seqlen_q = 10, + max_seqlen_k = 10, + num_heads_q = 64, + num_heads_k = 1, + head_dim = 128, + topk = 2048, + quant_mode = 1, + mask_mode = 3, + layout_q = "BSND", + layout_k = "BSND", + cmp_ratio = 128 +) \ No newline at end of file diff --git a/xllm_ops/attention/quant_lightning_indexer_v2_metadata/op_host/CMakeLists.txt b/xllm_ops/attention/quant_lightning_indexer_v2_metadata/op_host/CMakeLists.txt new file mode 100644 index 0000000..5e263c3 --- /dev/null +++ b/xllm_ops/attention/quant_lightning_indexer_v2_metadata/op_host/CMakeLists.txt @@ -0,0 +1,12 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +add_op_to_compiled_list() +add_modules_sources(OPTYPE quant_lightning_indexer_v2_metadata ACLNNTYPE aclnn) \ No newline at end of file diff --git a/xllm_ops/attention/quant_lightning_indexer_v2_metadata/op_host/op_api/aclnn_quant_lightning_indexer_v2_metadata.cpp b/xllm_ops/attention/quant_lightning_indexer_v2_metadata/op_host/op_api/aclnn_quant_lightning_indexer_v2_metadata.cpp new file mode 100644 index 0000000..a737cc4 --- /dev/null +++ b/xllm_ops/attention/quant_lightning_indexer_v2_metadata/op_host/op_api/aclnn_quant_lightning_indexer_v2_metadata.cpp @@ -0,0 +1,134 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file aclnn_quant_lightning_indexer_v2_metadata.cpp + * \brief + */ + +#include "aclnn_quant_lightning_indexer_v2_metadata.h" +#include "aclnn/aclnn_base.h" +#include "aclnn_kernels/common/op_error_check.h" +#include "aclnn_kernels/contiguous.h" +#include "aclnn_kernels/reshape.h" +#include "quant_lightning_indexer_v2_metadata.h" +#include "opdev/common_types.h" +#include "opdev/data_type_utils.h" +#include "opdev/format_utils.h" +#include "opdev/make_op_executor.h" +#include "opdev/op_dfx.h" +#include "opdev/op_executor.h" +#include "opdev/op_log.h" +#include "opdev/tensor_view_utils.h" +#include "opdev/platform.h" +#include "../quant_lightning_indexer_v2_metadata_check.h" + +#ifdef __cplusplus +extern "C" { +#endif + +aclnnStatus aclnnQuantLightningIndexerV2MetadataGetWorkspaceSize( + const aclTensor *cuSeqlensQOptional, const aclTensor *cuSeqlensKOptional, const aclTensor *sequsedQOptional, + const aclTensor *sequsedKOptional, const aclTensor *cmpResidualKOptional, int64_t numHeadsQ, int64_t numHeadsK, + int64_t headDim, int64_t topk, int64_t quantMode, int64_t batchSize, int64_t maxSeqlenQ, int64_t maxSeqlenK, + char *layoutQOptional, char *layoutKOptional, int64_t maskMode, int64_t cmpRatio, const aclTensor *metadata, + uint64_t *workspaceSize, aclOpExecutor **executor) +{ + if (workspaceSize == nullptr) { + OP_LOGE(ACLNN_ERR_INNER_NULLPTR, "workspaceSize is nullptr"); + return ACLNN_ERR_INNER_NULLPTR; + } + if (executor == nullptr) { + OP_LOGE(ACLNN_ERR_INNER_NULLPTR, "executor is nullptr"); + return ACLNN_ERR_INNER_NULLPTR; + } + L2_DFX_PHASE_1(aclnnQuantLightningIndexerV2Metadata, + DFX_IN(cuSeqlensQOptional, cuSeqlensKOptional, sequsedQOptional, sequsedKOptional, + cmpResidualKOptional, numHeadsQ, numHeadsK, headDim, topk, quantMode, batchSize, maxSeqlenQ, + maxSeqlenK, layoutQOptional, layoutKOptional, maskMode, cmpRatio), + DFX_OUT(metadata)); + + auto uniqueExecutor = CREATE_EXECUTOR(); + CHECK_RET(uniqueExecutor.get() != nullptr, ACLNN_ERR_INNER_CREATE_EXECUTOR); + + const op::PlatformInfo &npuInfo = op::GetCurrentPlatformInfo(); + uint32_t aicCoreNum = npuInfo.GetCubeCoreNum(); + uint32_t aivCoreNum = npuInfo.GetVectorCoreNum(); + const std::string socVersion = npuInfo.GetSocLongVersion(); + + auto ret = ParamsCheckQliV2(cuSeqlensQOptional, cuSeqlensKOptional, sequsedQOptional, sequsedKOptional, + cmpResidualKOptional, numHeadsQ, numHeadsK, headDim, topk, quantMode, batchSize, + maxSeqlenQ, maxSeqlenK, layoutQOptional, layoutKOptional, maskMode, cmpRatio, metadata, + aicCoreNum, aivCoreNum, socVersion); + CHECK_RET(ret == ACLNN_SUCCESS, ret); + + const aclTensor *cuSeqlensQOptionalContiguous = nullptr; + if (cuSeqlensQOptional != nullptr) { + cuSeqlensQOptionalContiguous = l0op::Contiguous(cuSeqlensQOptional, uniqueExecutor.get()); + if (cuSeqlensQOptionalContiguous == nullptr) { + OP_LOGE(ACLNN_ERR_INNER_NULLPTR, "cu_seqlens_q contiguous is null"); + return ACLNN_ERR_INNER_NULLPTR; + } + } + const aclTensor *cuSeqlensKOptionalContiguous = nullptr; + if (cuSeqlensKOptional != nullptr) { + cuSeqlensKOptionalContiguous = l0op::Contiguous(cuSeqlensKOptional, uniqueExecutor.get()); + if (cuSeqlensKOptionalContiguous == nullptr) { + OP_LOGE(ACLNN_ERR_INNER_NULLPTR, "cu_seqlens_k contiguous is null"); + return ACLNN_ERR_INNER_NULLPTR; + } + } + const aclTensor *sequsedQOptionalContiguous = nullptr; + if (sequsedQOptional != nullptr) { + sequsedQOptionalContiguous = l0op::Contiguous(sequsedQOptional, uniqueExecutor.get()); + if (sequsedQOptionalContiguous == nullptr) { + OP_LOGE(ACLNN_ERR_INNER_NULLPTR, "seqused_q contiguous is null"); + return ACLNN_ERR_INNER_NULLPTR; + } + } + const aclTensor *sequsedKOptionalContiguous = nullptr; + if (sequsedKOptional != nullptr) { + sequsedKOptionalContiguous = l0op::Contiguous(sequsedKOptional, uniqueExecutor.get()); + if (sequsedKOptionalContiguous == nullptr) { + OP_LOGE(ACLNN_ERR_INNER_NULLPTR, "seqused_k contiguous is null"); + return ACLNN_ERR_INNER_NULLPTR; + } + } + const aclTensor *cmpResidualLKOptionalContiguous = nullptr; + if (cmpResidualKOptional != nullptr) { + cmpResidualLKOptionalContiguous = l0op::Contiguous(cmpResidualKOptional, uniqueExecutor.get()); + if (cmpResidualLKOptionalContiguous == nullptr) { + OP_LOGE(ACLNN_ERR_INNER_NULLPTR, "cmp_residual_k contiguous is null"); + return ACLNN_ERR_INNER_NULLPTR; + } + } + + auto output = l0op::QuantLightningIndexerV2Metadata( + cuSeqlensQOptionalContiguous, cuSeqlensKOptionalContiguous, sequsedQOptionalContiguous, + sequsedKOptionalContiguous, cmpResidualLKOptionalContiguous, numHeadsQ, numHeadsK, headDim, topk, quantMode, + batchSize, maxSeqlenQ, maxSeqlenK, layoutQOptional, layoutKOptional, maskMode, cmpRatio, aicCoreNum, aivCoreNum, + socVersion.c_str(), metadata, uniqueExecutor.get()); + CHECK_RET(output != nullptr, ACLNN_ERR_INNER_NULLPTR); + + *workspaceSize = uniqueExecutor->GetWorkspaceSize(); + uniqueExecutor.ReleaseTo(executor); + return ACLNN_SUCCESS; +} + +aclnnStatus aclnnQuantLightningIndexerV2Metadata(void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, + aclrtStream stream) +{ + L2_DFX_PHASE_2(aclnnQuantLightningIndexerV2Metadata); + return CommonOpExecutorRun(workspace, workspaceSize, executor, stream); +} + +#ifdef __cplusplus +} +#endif diff --git a/xllm_ops/attention/quant_lightning_indexer_v2_metadata/op_host/op_api/aclnn_quant_lightning_indexer_v2_metadata.h b/xllm_ops/attention/quant_lightning_indexer_v2_metadata/op_host/op_api/aclnn_quant_lightning_indexer_v2_metadata.h new file mode 100644 index 0000000..10f8215 --- /dev/null +++ b/xllm_ops/attention/quant_lightning_indexer_v2_metadata/op_host/op_api/aclnn_quant_lightning_indexer_v2_metadata.h @@ -0,0 +1,37 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef ACLNN_QUANT_LIGHTNING_INDEXER_V2_METADATA_H +#define ACLNN_QUANT_LIGHTNING_INDEXER_V2_METADATA_H + +#include +#include "aclnn/aclnn_base.h" + +#ifdef __cplusplus +extern "C" { +#endif + +__attribute__((visibility("default"))) +aclnnStatus aclnnQuantLightningIndexerV2MetadataGetWorkspaceSize( + const aclTensor *cuSeqlensQOptional, const aclTensor *cuSeqlensKOptional, const aclTensor *sequsedQOptional, + const aclTensor *sequsedKOptional, const aclTensor *cmpResidualKOptional, int64_t numHeadsQ, int64_t numHeadsK, + int64_t headDim, int64_t topk, int64_t quantMode, int64_t batchSize, int64_t maxSeqlenQ, int64_t maxSeqlenK, + char *layoutQOptional, char *layoutKOptional, int64_t maskMode, int64_t cmpRatio, const aclTensor *metadata, + uint64_t *workspaceSize, aclOpExecutor **executor); + +__attribute__((visibility("default"))) +aclnnStatus aclnnQuantLightningIndexerV2Metadata(void* workspace, uint64_t workspaceSize, aclOpExecutor *executor, + aclrtStream stream); + +#ifdef __cplusplus +} +#endif + +#endif // ACLNN_QUANT_LIGHTNING_INDEXER_V2_METADATA_AICPU_H \ No newline at end of file diff --git a/xllm_ops/attention/quant_lightning_indexer_v2_metadata/op_host/op_api/quant_lightning_indexer_v2_metadata.cpp b/xllm_ops/attention/quant_lightning_indexer_v2_metadata/op_host/op_api/quant_lightning_indexer_v2_metadata.cpp new file mode 100644 index 0000000..dc4d283 --- /dev/null +++ b/xllm_ops/attention/quant_lightning_indexer_v2_metadata/op_host/op_api/quant_lightning_indexer_v2_metadata.cpp @@ -0,0 +1,57 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file quant_lightning_indexer_v2_metadata.cpp + * \brief + */ + +#include "quant_lightning_indexer_v2_metadata.h" +#include "opdev/aicpu/aicpu_task.h" +#include "opdev/make_op_executor.h" +#include "opdev/op_def.h" +#include "opdev/op_dfx.h" +#include "opdev/op_executor.h" +#include "opdev/op_log.h" +#include "opdev/shape_utils.h" + +using namespace op; +namespace l0op { +OP_TYPE_REGISTER(QuantLightningIndexerV2Metadata); + +const aclTensor *QuantLightningIndexerV2Metadata( + const aclTensor *cuSeqlensQOptional, const aclTensor *cuSeqlensKOptional, const aclTensor *sequsedQOptional, + const aclTensor *sequsedKOptional, const aclTensor *cmpResidualKOptional, int64_t numHeadsQ, int64_t numHeadsK, + int64_t headDim, int64_t topk, int64_t quantMode, int64_t batchSize, int64_t maxSeqlenQ, int64_t maxSeqlenK, + char *layoutQOptional, char *layoutKOptional, int64_t maskMode, int64_t cmpRatio, int64_t aicCoreNum, + int64_t aivCoreNum, const char *socVersion, const aclTensor *metadata, aclOpExecutor *executor) +{ + L0_DFX(QuantLightningIndexerV2Metadata, cuSeqlensQOptional, cuSeqlensKOptional, sequsedQOptional, sequsedKOptional, + cmpResidualKOptional, numHeadsQ, numHeadsK, headDim, topk, quantMode, batchSize, maxSeqlenQ, maxSeqlenK, + layoutQOptional, layoutKOptional, maskMode, cmpRatio, aicCoreNum, aivCoreNum, socVersion, metadata); + + static internal::AicpuTaskSpace space("QuantLightningIndexerV2Metadata"); + + auto ret = ADD_TO_LAUNCHER_LIST_AICPU( + QuantLightningIndexerV2Metadata, + OP_ATTR_NAMES({ "num_heads_q", "num_heads_k", "head_dim", "topk", "quant_mode", "batch_size", "max_seqlen_q", + "max_seqlen_k", "layout_q", "layout_k", "mask_mode", "cmp_ratio", "aic_core_num", + "aiv_core_num", "soc_version" }), + OP_INPUT(cuSeqlensQOptional, cuSeqlensKOptional, sequsedQOptional, sequsedKOptional, cmpResidualKOptional), + OP_OUTPUT(metadata), + OP_ATTR(numHeadsQ, numHeadsK, headDim, topk, quantMode, batchSize, maxSeqlenQ, maxSeqlenK, layoutQOptional, + layoutKOptional, maskMode, cmpRatio, aicCoreNum, aivCoreNum, socVersion)); + + OP_CHECK(ret == ACL_SUCCESS, + OP_LOGE(ACLNN_ERR_INNER_NULLPTR, "QuantLightningIndexerV2Metadata ADD_TO_LAUNCHER_LIST_AICPU failed."), + return nullptr); + return metadata; +} +} // namespace l0op diff --git a/xllm_ops/attention/quant_lightning_indexer_v2_metadata/op_host/op_api/quant_lightning_indexer_v2_metadata.h b/xllm_ops/attention/quant_lightning_indexer_v2_metadata/op_host/op_api/quant_lightning_indexer_v2_metadata.h new file mode 100644 index 0000000..7eb4458 --- /dev/null +++ b/xllm_ops/attention/quant_lightning_indexer_v2_metadata/op_host/op_api/quant_lightning_indexer_v2_metadata.h @@ -0,0 +1,25 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef L0_QUANT_LIGHTNING_INDEXER_V2_METADATA_H +#define L0_QUANT_LIGHTNING_INDEXER_V2_METADATA_H + +#include "opdev/op_executor.h" + +namespace l0op { +const aclTensor *QuantLightningIndexerV2Metadata( + const aclTensor *cuSeqlensQOptional, const aclTensor *cuSeqlensKOptional, const aclTensor *sequsedQOptional, + const aclTensor *sequsedKOptional, const aclTensor *cmpResidualKOptional, int64_t numHeadsQ, int64_t numHeadsK, + int64_t headDim, int64_t topk, int64_t quantMode, int64_t batchSize, int64_t maxSeqlenQ, int64_t maxSeqlenK, + char *layoutQOptional, char *layoutKOptional, int64_t maskMode, int64_t cmpRatio, int64_t aicCoreNum, + int64_t aivCoreNum, const char *socVersion, const aclTensor *metadata, aclOpExecutor *executor); +} // namespace l0op + +#endif \ No newline at end of file diff --git a/xllm_ops/attention/quant_lightning_indexer_v2_metadata/op_host/quant_lightning_indexer_v2_metadata_check.h b/xllm_ops/attention/quant_lightning_indexer_v2_metadata/op_host/quant_lightning_indexer_v2_metadata_check.h new file mode 100644 index 0000000..feecdef --- /dev/null +++ b/xllm_ops/attention/quant_lightning_indexer_v2_metadata/op_host/quant_lightning_indexer_v2_metadata_check.h @@ -0,0 +1,533 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file quant_lightning_indexer_v2_metadata_check.h + * \brief + */ + +#include "log/log.h" +#include "../../quant_lightning_indexer_v2/op_host/quant_lightning_indexer_v2_error_log.h" +#include "opdev/format_utils.h" +#include "opdev/op_log.h" +#include "opdev/data_type_utils.h" +#include "opdev/tensor_view_utils.h" +#include "../../quant_lightning_indexer_v2/op_kernel/quant_lightning_indexer_v2_metadata.h" + +#ifdef __cplusplus +extern "C" { +#endif + +namespace { + +static constexpr const char *QLI_V2_ACLNN_OP_NAME = "QuantLightningIndexerV2Metadata"; + +inline constexpr int64_t QLI_V2_QUANT_MODE_1 = 1; +inline constexpr int64_t QLI_V2_QUANT_MODE_2 = 2; +inline constexpr int64_t QLI_V2_QUANT_MODE_3 = 3; +inline constexpr int64_t QLI_V2_QUANT_MODE_4 = 4; +inline constexpr int64_t QLI_V2_QUANT_MODE_5 = 5; +inline constexpr int64_t QLI_V2_NO_MASK_MODE = 0; +inline constexpr int64_t QLI_V2_CAUSAL_MASK_MODE = 3; +inline constexpr int64_t QLI_V2_CMP_RATIO_LOWER_BOUND = 1; +inline constexpr int64_t QLI_V2_CMP_RATIO_UPPER_BOUND = 128; +inline constexpr int64_t QLI_V2_NUM_HEADS_Q_LOWER_BOUND = 1; +inline constexpr int64_t QLI_V2_NUM_HEADS_Q_UPPER_BOUND = 64; +inline constexpr int64_t QLI_V2_TOPK_LOWER_BOUND = 1; +inline constexpr int64_t QLI_V2_A5_TOPK_UPPER_BOUND = 8192; +inline constexpr int64_t QLI_V2_A3_TOPK_UPPER_BOUND = 2048; + +inline bool IsTensorExistQliV2(const aclTensor *tensor) +{ + return (tensor != nullptr) && (tensor->GetViewShape().GetDimNum() > 0) && (tensor->GetViewShape().GetDim(0) > 0); +} + +int64_t GetDimNumQliV2(const aclTensor *tensor) +{ + if (tensor == nullptr) { + return -1; + } + return tensor->GetViewShape().GetDimNum(); +} + +aclDataType GetDataTypeQliV2(const aclTensor *tensor) +{ + aclDataType dataType = aclDataType::ACL_DT_UNDEFINED; + if (tensor == nullptr) { + return dataType; + } + aclGetDataType(tensor, &dataType); + return dataType; +} + +inline bool IsTensorSourceQLiV2(const std::string &source) { return source != "batch_size"; } + +inline int64_t GetRawShapeSizeQLiV2(const std::string &source, int64_t batchValue) +{ + if (source.find("cu_seqlens") != std::string::npos) { + return batchValue + 1; + } + return batchValue; +} + +int64_t GetQueryBatchSizeQliV2(int64_t batchSize, const aclTensor *cuSeqlensQOptional, + const aclTensor *sequsedQOptional, const char *layoutQOptional, std::string &source) +{ + // 1. 如果sequsedQOptional 传了，使用sequsedQOptional获取BatchSize + if (IsTensorExistQliV2(sequsedQOptional)) { + source = "seqused_q"; + return sequsedQOptional->GetViewShape().GetDim(0); + } + // 2. 如果sequsedQOptional 没传，使用cuSeqlensQOptional获取BatchSize + if (strcmp(layoutQOptional, "TND") == 0) { + if (IsTensorExistQliV2( + cuSeqlensQOptional)) { // 前序校验已保证layout_q = TND时，cu_seqlens_q必须传入，此通路必达 + source = "cu_seqlens_q"; + return cuSeqlensQOptional->GetViewShape().GetDim(0) - 1; + } + } + source = "batch_size"; + // 3. 使用batchSize + return batchSize; +} + +int64_t GetKeyBatchSizeQliV2(int64_t batchSize, const aclTensor *cuSeqlensKOptional, const aclTensor *sequsedKOptional, + const char *layoutKOptional, std::string &source) +{ + // 1. 如果sequsedKOptional 传了，使用sequsedKOptional获取BatchSize + if (IsTensorExistQliV2(sequsedKOptional)) { + source = "seqused_q"; + return sequsedKOptional->GetViewShape().GetDim(0); + } + // 如果是 TND，必须使用 cuSeqlensKOptional获取BatchSize + if (strcmp(layoutKOptional, "TND") == 0) { + if (IsTensorExistQliV2( + cuSeqlensKOptional)) { // 前序校验已保证layout_k = TND时，cu_seqlens_k必须传入，此通路必达 + source = "cu_seqlens_q"; + return cuSeqlensKOptional->GetViewShape().GetDim(0) - 1; + } + } + source = "batch_size"; + // 3. 使用batchSize + return batchSize; +} + +aclnnStatus CheckSingleParamQliV2(int64_t numHeadsQ, int64_t numHeadsK, int64_t headDim, int64_t topk, + int64_t quantMode, int64_t batchSize, int64_t maxSeqlenQ, int64_t maxSeqlenK, + const char *layoutQOptional, const char *layoutKOptional, int64_t maskMode, + int64_t cmpRatio, uint32_t aicCoreNum, uint32_t aivCoreNum, + const std::string &socVersion) +{ + // num_heads_k 校验 + if (numHeadsK != 1) { + OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(QLI_V2_ACLNN_OP_NAME, "num_heads_kv", std::to_string(numHeadsK), + "The value of num_heads_kv should be 1"); + return ACLNN_ERR_PARAM_INVALID; + } + // head_dim 校验 + if (headDim != 128) { + OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(QLI_V2_ACLNN_OP_NAME, "head_dim", std::to_string(numHeadsK), + "The value of head_dim should be 128"); + return ACLNN_ERR_PARAM_INVALID; + } + // batch_size 非负校验 + if (batchSize < 0) { + OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(QLI_V2_ACLNN_OP_NAME, "batch_size", std::to_string(batchSize), + "The value of batch_size should not be negative"); + return ACLNN_ERR_PARAM_INVALID; + } + // max_seqlen_q 校验 + if (maxSeqlenQ < -1) { + OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(QLI_V2_ACLNN_OP_NAME, "max_seqlen_q", std::to_string(maxSeqlenQ), + "The value of max_seqlen_q should be >= -1"); + return ACLNN_ERR_PARAM_INVALID; + } + // max_seqlen_k 校验 + if (maxSeqlenK < -1) { + OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(QLI_V2_ACLNN_OP_NAME, "max_seqlen_k", std::to_string(maxSeqlenK), + "The value of max_seqlen_k should be >= -1"); + return ACLNN_ERR_PARAM_INVALID; + } + // mask_mode 校验 + if ((maskMode != QLI_V2_NO_MASK_MODE) && (maskMode != QLI_V2_CAUSAL_MASK_MODE)) { + OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(QLI_V2_ACLNN_OP_NAME, "mask_mode", std::to_string(maskMode), + "The value of mask_mode should be " + + std::to_string(QLI_V2_NO_MASK_MODE) + " or " + + std::to_string(QLI_V2_CAUSAL_MASK_MODE)); + return ACLNN_ERR_PARAM_INVALID; + } + // layout_q 校验 + if (layoutQOptional == nullptr) { + OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON(QLI_V2_ACLNN_OP_NAME, "layout_q", "Layout_q is null"); + return ACLNN_ERR_PARAM_INVALID; + } + if ((strcmp(layoutQOptional, "TND") != 0) && (strcmp(layoutQOptional, "BSND") != 0)) { + OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(QLI_V2_ACLNN_OP_NAME, "layout_q", layoutQOptional, + "The value of layout_q must be TND or BSND"); + return ACLNN_ERR_PARAM_INVALID; + } + // layout_k 校验 + if (layoutKOptional == nullptr) { + OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON(QLI_V2_ACLNN_OP_NAME, "layout_k", "Layout_k is null"); + return ACLNN_ERR_PARAM_INVALID; + } + if ((strcmp(layoutKOptional, "PA_BBND") != 0) && (strcmp(layoutQOptional, layoutKOptional) != 0)) { + OP_LOGE_FOR_INVALID_VALUES_WITH_REASON(QLI_V2_ACLNN_OP_NAME, "layout_q and layout_k", + std::string(layoutQOptional) + " and " + std::string(layoutKOptional), + "For layout_k != PA_BBND, layout_q and layout_k must be the same"); + return ACLNN_ERR_PARAM_INVALID; + } + // 校验 A2/A3 参数 + if (socVersion.find("Ascend950") == std::string::npos) { + // num_heads_q 校验 + CHECK_COND(numHeadsQ == QLI_V2_NUM_HEADS_Q_UPPER_BOUND, ACLNN_ERR_PARAM_INVALID, + "num_heads_q should be %lld, but got %lld", QLI_V2_NUM_HEADS_Q_UPPER_BOUND, numHeadsQ); + // topk 校验 + CHECK_COND(topk >= QLI_V2_TOPK_LOWER_BOUND && topk <= QLI_V2_A3_TOPK_UPPER_BOUND, ACLNN_ERR_PARAM_INVALID, + "topk should be [%lld, %lld], but got %lld", QLI_V2_TOPK_LOWER_BOUND, QLI_V2_A3_TOPK_UPPER_BOUND, + topk); + // quant_mode 校验 + CHECK_COND(quantMode == QLI_V2_QUANT_MODE_2, ACLNN_ERR_PARAM_INVALID, "quant_mode should be 2, but got %lld", + quantMode); + // cmp_ratio 校验 + CHECK_COND((cmpRatio >= QLI_V2_CMP_RATIO_LOWER_BOUND) && (cmpRatio <= QLI_V2_CMP_RATIO_UPPER_BOUND) && + ((cmpRatio & (cmpRatio - 1)) == 0), + ACLNN_ERR_PARAM_INVALID, "cmp_ratio should be 1/2/4/8/16/32/64/128, but got %lld", cmpRatio); + CHECK_COND(strcmp(layoutKOptional, "PA_BBND") == 0, ACLNN_ERR_PARAM_INVALID, + "layout_k must be PA_BBND, but got %s", layoutKOptional); + } else { // 校验 A5参数 + // num_heads_q 校验 + if (numHeadsQ < QLI_V2_NUM_HEADS_Q_LOWER_BOUND || numHeadsQ > QLI_V2_NUM_HEADS_Q_UPPER_BOUND) { + OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(QLI_V2_ACLNN_OP_NAME, "num_heads_q", std::to_string(numHeadsQ), + "The value of num_heads_q should be in range [" + + std::to_string(QLI_V2_NUM_HEADS_Q_LOWER_BOUND) + ", " + + std::to_string(QLI_V2_NUM_HEADS_Q_UPPER_BOUND) + "]"); + return ACLNN_ERR_PARAM_INVALID; + } + // topk 校验 + if (topk < QLI_V2_TOPK_LOWER_BOUND || topk > QLI_V2_A5_TOPK_UPPER_BOUND) { + OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(QLI_V2_ACLNN_OP_NAME, "topk", std::to_string(topk), + "The value of topk should be in range [" + + std::to_string(QLI_V2_TOPK_LOWER_BOUND) + ", " + + std::to_string(QLI_V2_A5_TOPK_UPPER_BOUND) + "]"); + return ACLNN_ERR_PARAM_INVALID; + } + // quant_mode 校验 + if ((quantMode != QLI_V2_QUANT_MODE_1) && (quantMode != QLI_V2_QUANT_MODE_2) && + (quantMode != QLI_V2_QUANT_MODE_3) && (quantMode != QLI_V2_QUANT_MODE_4) && + (quantMode != QLI_V2_QUANT_MODE_5)) { + OP_LOGE_FOR_INVALID_VALUE_WITH_REASON( + QLI_V2_ACLNN_OP_NAME, "quant_mode", std::to_string(quantMode), + "The value of quant_mode should be in [" + std::to_string(QLI_V2_QUANT_MODE_1) + ", " + + std::to_string(QLI_V2_QUANT_MODE_2) + ", " + std::to_string(QLI_V2_QUANT_MODE_3) + ", " + + std::to_string(QLI_V2_QUANT_MODE_4) + ", " + std::to_string(QLI_V2_QUANT_MODE_5) + "]"); + return ACLNN_ERR_PARAM_INVALID; + } + // cmp_ratio 校验 + if ((cmpRatio < QLI_V2_CMP_RATIO_LOWER_BOUND) || (cmpRatio > QLI_V2_CMP_RATIO_UPPER_BOUND)) { + OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(QLI_V2_ACLNN_OP_NAME, "cmp_ratio", std::to_string(cmpRatio), + "The value of cmp_ratio should be in range [" + + std::to_string(QLI_V2_CMP_RATIO_LOWER_BOUND) + ", " + + std::to_string(QLI_V2_CMP_RATIO_UPPER_BOUND) + "]"); + return ACLNN_ERR_PARAM_INVALID; + } + if ((strcmp(layoutKOptional, "TND") != 0) && (strcmp(layoutKOptional, "BSND") != 0) && + (strcmp(layoutKOptional, "PA_BBND") != 0)) { + OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(QLI_V2_ACLNN_OP_NAME, "layout_k", layoutKOptional, + "The value of layout_k must be in [TND, BSND, PA_BBND]"); + return ACLNN_ERR_PARAM_INVALID; + } + } + // 校验 layout_q 为 BSND 时，max_seqlen_q 必须大于 0 + if ((strcmp(layoutQOptional, "BSND") == 0) && (maxSeqlenQ <= 0)) { + OP_LOGE_FOR_INVALID_VALUES_WITH_REASON(QLI_V2_ACLNN_OP_NAME, "max_seqlen_q", std::to_string(maxSeqlenQ), + "When layout_q is BSND, the value of max_seqlen_q " + "must be equal to the size of the second axis of q"); + return ACLNN_ERR_PARAM_INVALID; + } + // 校验 layout_k 为 BSND 时，max_seqlen_k 必须大于 0 + if ((strcmp(layoutKOptional, "BSND") == 0) && (maxSeqlenK <= 0)) { + OP_LOGE_FOR_INVALID_VALUES_WITH_REASON(QLI_V2_ACLNN_OP_NAME, "max_seqlen_k", std::to_string(maxSeqlenK), + "When layout_k is BSND, the value of max_seqlen_k " + "must be equal to the size of the second axis of k"); + return ACLNN_ERR_PARAM_INVALID; + } + // 核心数校验 + CHECK_COND(aicCoreNum > 0, ACLNN_ERR_PARAM_INVALID, "AIC num should be larger than 0, but got %u", aicCoreNum); + CHECK_COND(aicCoreNum <= optiling::AIC_CORE_MAX_NUM, ACLNN_ERR_PARAM_INVALID, + "The maximum supported AIC num is %u, but got %u", optiling::AIC_CORE_MAX_NUM, aicCoreNum); + CHECK_COND(aivCoreNum > 0, ACLNN_ERR_PARAM_INVALID, "AIV num should be larger than 0, but got %u", aivCoreNum); + CHECK_COND(aivCoreNum <= optiling::AIV_CORE_MAX_NUM, ACLNN_ERR_PARAM_INVALID, + "The maximum supported AIV num is %u, but got %u", optiling::AIV_CORE_MAX_NUM, aivCoreNum); + return ACLNN_SUCCESS; +} + +aclnnStatus CheckExistenceQliV2(int64_t maskMode, int64_t cmpRatio, const aclTensor *cuSeqlensQOptional, + const aclTensor *cuSeqlensKOptional, const aclTensor *sequsedQOptional, + const aclTensor *sequsedKOptional, const aclTensor *cmpResidualKOptional, + int64_t maxSeqlenQ, int64_t maxSeqlenK, const char *layoutQOptional, + const char *layoutKOptional, const aclTensor *metadata) +{ + // cu_seqlens_q 存在性校验 + if ((strcmp(layoutQOptional, "TND") == 0) && !IsTensorExistQliV2(cuSeqlensQOptional)) { + OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON(QLI_V2_ACLNN_OP_NAME, "cu_seqlens_q", + "When layout_q is TND, cu_seqlens_q must be provided"); + return ACLNN_ERR_PARAM_INVALID; + } + // layout_q BSND, seqused_q 不存在时，max_seqlen_q 不能为-1 + if (strcmp(layoutQOptional, "BSND") == 0 && !IsTensorExistQliV2(sequsedQOptional) && (maxSeqlenQ <= -1)) { + OP_LOGE_FOR_INVALID_VALUE_WITH_REASON( + QLI_V2_ACLNN_OP_NAME, "max_seqlen_q", std::to_string(maxSeqlenQ), + "When layout_q is BSND and seqused_q is not provided, max_seqlen_q can not be -1"); + return ACLNN_ERR_PARAM_INVALID; + } + // cu_seqlens_k 存在性校验 + if ((strcmp(layoutKOptional, "TND") == 0) && !IsTensorExistQliV2(cuSeqlensKOptional)) { + OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON(QLI_V2_ACLNN_OP_NAME, "cu_seqlens_k", + "When layout_k is TND, cu_seqlens_k must be provided"); + return ACLNN_ERR_PARAM_INVALID; + } + // seqused_k 存在性校验 + if ((strcmp(layoutKOptional, "PA_BBND") == 0) && !IsTensorExistQliV2(sequsedKOptional)) { + OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON(QLI_V2_ACLNN_OP_NAME, "seqused_k", + "When layout_k is PA_BBND, seqused_k must be provided"); + return ACLNN_ERR_PARAM_INVALID; + } + // layout_k BSND, seqused_k 不存在时，max_seqlen_k 不能为-1 + if (strcmp(layoutKOptional, "BSND") == 0 && !IsTensorExistQliV2(sequsedKOptional) && maxSeqlenK <= -1) { + OP_LOGE_FOR_INVALID_VALUE_WITH_REASON( + QLI_V2_ACLNN_OP_NAME, "max_seqlen_k", std::to_string(maxSeqlenK), + "When layout_k is BSND and seqused_k is not provided, max_seqlen_k can not be -1"); + return ACLNN_ERR_PARAM_INVALID; + } + // cmp_residual_k 存在性校验 + if ((cmpRatio != QLI_V2_CMP_RATIO_LOWER_BOUND) && (maskMode == QLI_V2_CAUSAL_MASK_MODE) && + !IsTensorExistQliV2(cmpResidualKOptional)) { + OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON( + QLI_V2_ACLNN_OP_NAME, "cmp_residual_k", + "When cmp_ratio is not 1 and mask_mode is CAUSAL, cmp_residual_k must be provided"); + return ACLNN_ERR_PARAM_INVALID; + } + // metadata 存在性校验 + if (!IsTensorExistQliV2(metadata)) { + OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON(QLI_V2_ACLNN_OP_NAME, "metadata", "Metadata is nullptr"); + return ACLNN_ERR_PARAM_INVALID; + } + return ACLNN_SUCCESS; +} + +aclnnStatus CheckConsistencyQliV2(int64_t batchSize, const aclTensor *cuSeqlensQOptional, + const aclTensor *cuSeqlensKOptional, const aclTensor *sequsedQOptional, + const aclTensor *sequsedKOptional, const aclTensor *cmpResidualKOptional, + const char *layoutQOptional, const char *layoutKOptional, const aclTensor *metadata) +{ + int64_t dimNum = -1; + aclDataType dataType = aclDataType::ACL_DT_UNDEFINED; + + // 校验 cu_seqlens_q + if (IsTensorExistQliV2(cuSeqlensQOptional)) { + dimNum = GetDimNumQliV2(cuSeqlensQOptional); + if (dimNum != 1) { + OP_LOGE_FOR_INVALID_SHAPEDIM(QLI_V2_ACLNN_OP_NAME, "cu_seqlens_q", std::to_string(dimNum), "1"); + return ACLNN_ERR_PARAM_INVALID; + } + dataType = GetDataTypeQliV2(cuSeqlensQOptional); + if (dataType != aclDataType::ACL_INT32) { + OP_LOGE_FOR_INVALID_DTYPE_WITH_REASON(QLI_V2_ACLNN_OP_NAME, "cu_seqlens_q", ToString(dataType).GetString(), + "The dtype of cu_seqlens_q must be int32"); + return ACLNN_ERR_PARAM_INVALID; + } + } + // 校验 cu_seqlens_k + if (IsTensorExistQliV2(cuSeqlensKOptional)) { + dimNum = GetDimNumQliV2(cuSeqlensKOptional); + if (dimNum != 1) { + OP_LOGE_FOR_INVALID_SHAPEDIM(QLI_V2_ACLNN_OP_NAME, "cu_seqlens_k", std::to_string(dimNum), "1"); + return ACLNN_ERR_PARAM_INVALID; + } + dataType = GetDataTypeQliV2(cuSeqlensKOptional); + if (dataType != aclDataType::ACL_INT32) { + OP_LOGE_FOR_INVALID_DTYPE_WITH_REASON(QLI_V2_ACLNN_OP_NAME, "cu_seqlens_k", ToString(dataType).GetString(), + "The dtype of cu_seqlens_k must be int32"); + return ACLNN_ERR_PARAM_INVALID; + } + } + // 校验 seqused_q + if (IsTensorExistQliV2(sequsedQOptional)) { + dimNum = GetDimNumQliV2(sequsedQOptional); + if (dimNum != 1) { + OP_LOGE_FOR_INVALID_SHAPEDIM(QLI_V2_ACLNN_OP_NAME, "seqused_q", std::to_string(dimNum), "1"); + return ACLNN_ERR_PARAM_INVALID; + } + dataType = GetDataTypeQliV2(sequsedQOptional); + if (dataType != aclDataType::ACL_INT32) { + OP_LOGE_FOR_INVALID_DTYPE_WITH_REASON(QLI_V2_ACLNN_OP_NAME, "seqused_q", ToString(dataType).GetString(), + "The dtype of seqused_q must be int32"); + return ACLNN_ERR_PARAM_INVALID; + } + } + // 校验 seqused_k + if (IsTensorExistQliV2(sequsedKOptional)) { + dimNum = GetDimNumQliV2(sequsedKOptional); + if (dimNum != 1) { + OP_LOGE_FOR_INVALID_SHAPEDIM(QLI_V2_ACLNN_OP_NAME, "seqused_k", std::to_string(dimNum), "1"); + return ACLNN_ERR_PARAM_INVALID; + } + dataType = GetDataTypeQliV2(sequsedKOptional); + if (dataType != aclDataType::ACL_INT32) { + OP_LOGE_FOR_INVALID_DTYPE_WITH_REASON(QLI_V2_ACLNN_OP_NAME, "seqused_k", ToString(dataType).GetString(), + "The dtype of seqused_k must be int32"); + return ACLNN_ERR_PARAM_INVALID; + } + } + // 校验 cmp_residual_k + if (IsTensorExistQliV2(cmpResidualKOptional)) { + dimNum = GetDimNumQliV2(cmpResidualKOptional); + if (dimNum != 1) { + OP_LOGE_FOR_INVALID_SHAPEDIM(QLI_V2_ACLNN_OP_NAME, "cmp_residual_k", std::to_string(dimNum), "1"); + return ACLNN_ERR_PARAM_INVALID; + } + dataType = GetDataTypeQliV2(cmpResidualKOptional); + if (dataType != aclDataType::ACL_INT32) { + OP_LOGE_FOR_INVALID_DTYPE_WITH_REASON(QLI_V2_ACLNN_OP_NAME, "cmp_residual_k", + ToString(dataType).GetString(), + "The dtype of cmp_residual_k must be int32"); + return ACLNN_ERR_PARAM_INVALID; + } + } + // 校验 metadata + if (IsTensorExistQliV2(metadata)) { + dimNum = GetDimNumQliV2(metadata); + if (dimNum != 1) { + OP_LOGE_FOR_INVALID_SHAPEDIM(QLI_V2_ACLNN_OP_NAME, "metadata", std::to_string(dimNum), "1"); + return ACLNN_ERR_PARAM_INVALID; + } + dataType = GetDataTypeQliV2(metadata); + if (dataType != aclDataType::ACL_INT32) { + OP_LOGE_FOR_INVALID_DTYPE_WITH_REASON(QLI_V2_ACLNN_OP_NAME, "metadata", ToString(dataType).GetString(), + "The dtype of metadata must be int32"); + return ACLNN_ERR_PARAM_INVALID; + } + // 校验 metadata 元素数 + if (metadata->GetViewShape().GetDim(0) != optiling::QLI_V2_METADATA_TOTAL_SIZE) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "The element num of metadata must be %u, but got %lld", + optiling::QLI_V2_METADATA_TOTAL_SIZE, metadata->GetViewShape().GetDim(0)); + return ACLNN_ERR_PARAM_INVALID; + } + } + // 校验batch + std::string querySource; + std::string keySource; + int64_t queryBatchSize = + GetQueryBatchSizeQliV2(batchSize, cuSeqlensQOptional, sequsedQOptional, layoutQOptional, querySource); + int64_t keyBatchSize = + GetKeyBatchSizeQliV2(batchSize, cuSeqlensKOptional, sequsedKOptional, layoutKOptional, keySource); + if (queryBatchSize != keyBatchSize) { + OP_LOGE_FOR_INVALID_SHAPESIZES_WITH_REASON( + QLI_V2_ACLNN_OP_NAME, querySource + " and " + keySource, + std::to_string(GetRawShapeSizeQLiV2(querySource, queryBatchSize)) + " and " + + std::to_string(GetRawShapeSizeQLiV2(keySource, keyBatchSize)), + "The batch_size obtained from query should be the same as that obtained from key"); + if (IsTensorSourceQLiV2(querySource) && IsTensorSourceQLiV2(keySource)) { + OP_LOGE_FOR_INVALID_SHAPESIZES_WITH_REASON( + QLI_V2_ACLNN_OP_NAME, querySource + " and " + keySource, + std::to_string(GetRawShapeSizeQLiV2(querySource, queryBatchSize)) + " and " + + std::to_string(GetRawShapeSizeQLiV2(keySource, keyBatchSize)), + "The batch_size obtained from query should be the same as that obtained from key"); + } else if (IsTensorSourceQLiV2(querySource)) { + OP_LOGE_FOR_INVALID_SHAPESIZE_WITH_REASON( + QLI_V2_ACLNN_OP_NAME, querySource, std::to_string(GetRawShapeSizeQLiV2(querySource, queryBatchSize)), + "The batch_size obtained from query should be the same as that obtained from key"); + } else { + OP_LOGE_FOR_INVALID_SHAPESIZE_WITH_REASON( + QLI_V2_ACLNN_OP_NAME, keySource, std::to_string(GetRawShapeSizeQLiV2(keySource, keyBatchSize)), + "The batch_size obtained from query should be the same as that obtained from key"); + } + return ACLNN_ERR_PARAM_INVALID; + } + // 校验TND场景q维度一致性 + if (strcmp(layoutQOptional, "TND") == 0 && IsTensorExistQliV2(sequsedQOptional)) { + int64_t cuSeqlensQBatchSize = cuSeqlensQOptional->GetViewShape().GetDim(0) - 1; + CHECK_COND( + cuSeqlensQBatchSize == queryBatchSize, ACLNN_ERR_PARAM_INVALID, + "When layout_q is TND and seqused_q is passed, The batch_size obtained from cu_seqlens_q should be the " + "same as that obtained from seqused_q, but got %lld and %lld", + cuSeqlensQBatchSize, queryBatchSize); + if (cuSeqlensQBatchSize != queryBatchSize) { + OP_LOGE_FOR_INVALID_SHAPESIZES_WITH_REASON( + QLI_V2_ACLNN_OP_NAME, "cu_seqlens_q and seqused_q", + std::to_string(cuSeqlensQBatchSize) + " and " + std::to_string(queryBatchSize), + "When layout_q is TND and seqused_q is passed, " + "the shape size of cu_seqlens_q minus 1 must be equal to " + "the shape size of seqused_q"); + return ACLNN_ERR_PARAM_INVALID; + } + } + // 校验TND场景k维度一致性 + if (strcmp(layoutKOptional, "TND") == 0 && IsTensorExistQliV2(sequsedKOptional)) { + int64_t cuSeqlensKBatchSize = cuSeqlensKOptional->GetViewShape().GetDim(0) - 1; + if (cuSeqlensKBatchSize != keyBatchSize) { + OP_LOGE_FOR_INVALID_SHAPESIZES_WITH_REASON( + QLI_V2_ACLNN_OP_NAME, "cu_seqlens_k and seqused_k", + std::to_string(cuSeqlensKBatchSize) + " and " + std::to_string(keyBatchSize), + "When layout_k is TND and seqused_k is passed, " + "the shape size of cu_seqlens_k minus 1 must be equal to " + "the shape size of seqused_k"); + return ACLNN_ERR_PARAM_INVALID; + } + } + // 校验 cmp_residual_k 元素数 + auto cmpResidualKBatch = cmpResidualKOptional->GetViewShape().GetDim(0); + if (IsTensorExistQliV2(cmpResidualKOptional) && (cmpResidualKBatch != queryBatchSize)) { + if (IsTensorSourceQLiV2(querySource)) { + OP_LOGE_FOR_INVALID_SHAPESIZES_WITH_REASON( + QLI_V2_ACLNN_OP_NAME, "cmp_residual_k and " + querySource, + std::to_string(cmpResidualKBatch) + " and " + + std::to_string(GetRawShapeSizeQLiV2(querySource, queryBatchSize)), + "The batch_size of cmp_residual_k should match the valid batch size"); + } else { + OP_LOGE_FOR_INVALID_SHAPESIZE_WITH_REASON( + QLI_V2_ACLNN_OP_NAME, "cmp_residual_k", std::to_string(cmpResidualKBatch), + "The batch_size of cmp_residual_k should match the valid batch size"); + } + return ACLNN_ERR_PARAM_INVALID; + } + return ACLNN_SUCCESS; +} + +aclnnStatus ParamsCheckQliV2(const aclTensor *cuSeqlensQOptional, const aclTensor *cuSeqlensKOptional, + const aclTensor *sequsedQOptional, const aclTensor *sequsedKOptional, + const aclTensor *cmpResidualKOptional, int64_t numHeadsQ, int64_t numHeadsK, + int64_t headDim, int64_t topk, int64_t quantMode, int64_t batchSize, int64_t maxSeqlenQ, + int64_t maxSeqlenK, char *layoutQOptional, char *layoutKOptional, int64_t maskMode, + int64_t cmpRatio, const aclTensor *metadata, uint32_t aicCoreNum, uint32_t aivCoreNum, + const std::string &socVersion) +{ + auto ret = + CheckSingleParamQliV2(numHeadsQ, numHeadsK, headDim, topk, quantMode, batchSize, maxSeqlenQ, maxSeqlenK, + layoutQOptional, layoutKOptional, maskMode, cmpRatio, aicCoreNum, aivCoreNum, socVersion); + CHECK_RET(ret == ACLNN_SUCCESS, ACLNN_ERR_PARAM_INVALID); + + ret = CheckExistenceQliV2(maskMode, cmpRatio, cuSeqlensQOptional, cuSeqlensKOptional, sequsedQOptional, + sequsedKOptional, cmpResidualKOptional, maxSeqlenQ, maxSeqlenK, layoutQOptional, + layoutKOptional, metadata); + CHECK_RET(ret == ACLNN_SUCCESS, ACLNN_ERR_PARAM_INVALID); + + ret = CheckConsistencyQliV2(batchSize, cuSeqlensQOptional, cuSeqlensKOptional, sequsedQOptional, sequsedKOptional, + cmpResidualKOptional, layoutQOptional, layoutKOptional, metadata); + CHECK_RET(ret == ACLNN_SUCCESS, ACLNN_ERR_PARAM_INVALID); + + return ACLNN_SUCCESS; +} +} // namespace + +#ifdef __cplusplus +} +#endif diff --git a/xllm_ops/attention/quant_lightning_indexer_v2_metadata/op_kernel_aicpu/quant_lightning_indexer_v2_metadata_aicpu.cpp b/xllm_ops/attention/quant_lightning_indexer_v2_metadata/op_kernel_aicpu/quant_lightning_indexer_v2_metadata_aicpu.cpp new file mode 100644 index 0000000..d6d5638 --- /dev/null +++ b/xllm_ops/attention/quant_lightning_indexer_v2_metadata/op_kernel_aicpu/quant_lightning_indexer_v2_metadata_aicpu.cpp @@ -0,0 +1,933 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#include "quant_lightning_indexer_v2_metadata_aicpu.h" +#include +#include +#include "status.h" + +using namespace optiling; + +namespace aicpu { +uint32_t QuantLightningIndexerV2MetadataCpuKernel::Compute(CpuKernelContext &ctx) +{ + bool success = Prepare(ctx); + if (!success) { + return KERNEL_STATUS_PARAM_INVALID; + } + SplitResult splitRes{aicCoreNum_, aivCoreNum_}; + success = BalanceSchedule(splitRes) && GenMetadata(splitRes); + return success ? KERNEL_STATUS_OK : KERNEL_STATUS_PARAM_INVALID; +} + +bool QuantLightningIndexerV2MetadataCpuKernel::Prepare(CpuKernelContext &ctx) +{ + // input + cuSeqlensQ_ = ctx.Input(static_cast(ParamId::actSeqLenQ)); + cuSeqlensK_ = ctx.Input(static_cast(ParamId::actSeqLenK)); + sequsedQ_ = ctx.Input(static_cast(ParamId::seqUsedQ)); + sequsedK_ = ctx.Input(static_cast(ParamId::seqUsedK)); + cmpResidualK_ = ctx.Input(static_cast(ParamId::cmpResidualK)); + // output + metadata_ = ctx.Output(static_cast(ParamId::metadata)); + + bool requiredAttrs = GetAttrValue(ctx, "aic_core_num", aicCoreNum_) && + GetAttrValue(ctx, "aiv_core_num", aivCoreNum_) && + GetAttrValue(ctx, "soc_version", socVersion_) && + GetAttrValue(ctx, "num_heads_q", numHeadsQ_) && GetAttrValue(ctx, "num_heads_k", numHeadsK_) && + GetAttrValue(ctx, "head_dim", headDim_) && GetAttrValue(ctx, "topk", topk_); + if (!requiredAttrs) { + return false; + } + + // attributes optional + GetAttrValueOpt(ctx, "batch_size", batchSize_); + GetAttrValueOpt(ctx, "max_seqlen_q", maxSeqlenQ_); + GetAttrValueOpt(ctx, "max_seqlen_k", maxSeqlenK_); + GetAttrValueOpt(ctx, "layout_q", layoutQ_); + GetAttrValueOpt(ctx, "layout_k", layoutK_); + GetAttrValueOpt(ctx, "mask_mode", maskMode_); + GetAttrValueOpt(ctx, "cmp_ratio", cmpRatio_); + + return (ParamsCheck() && ParamsInit()); +} + +bool QuantLightningIndexerV2MetadataCpuKernel::ParamsCheck() +{ + // 校验输出 metadata 是否为空 + if (metadata_ == nullptr) { + KERNEL_LOG_ERROR("Output metadata is nullptr"); + return false; + } else if (metadata_->GetData() == nullptr) { + KERNEL_LOG_ERROR("Output metadata data is nullptr"); + return false; + } + int32_t batchSize = GetQueryBatchSize(); + // 校验 cu_seqlens_q 元素 + if (layoutQ_ == "TND") { + if (cuSeqlensQ_ != nullptr && cuSeqlensQ_->GetData() != nullptr) { + const int32_t *cuSeqlensQPtr = static_cast(cuSeqlensQ_->GetData()); + // 校验 cu_seqlens_q 首元素为 0 + if (cuSeqlensQPtr[0] != 0) { + KERNEL_LOG_ERROR("The first element of cu_seqlens_q should be 0, but got %d", cuSeqlensQPtr[0]); + return false; + } + for (int i = 0; i < batchSize + 1; i++) { + // 校验 cu_seqlens_q 元素递增 + if (i > 0 && cuSeqlensQPtr[i - 1] > cuSeqlensQPtr[i]) { + KERNEL_LOG_ERROR("The elements in cu_seqlens_q must be in ascending order, " + "but got cu_seqlens_q[%d] = %d, cu_seqlens_q[%d] = %d", + i - 1, cuSeqlensQPtr[i - 1], i, cuSeqlensQPtr[i]); + return false; + } + } + } + } + // 校验 cu_seqlens_k 元素 + if (layoutK_ == "TND") { + if (cuSeqlensK_ != nullptr && cuSeqlensK_->GetData() != nullptr) { + const int32_t *cuSeqlensKPtr = static_cast(cuSeqlensK_->GetData()); + // 校验 cu_seqlens_k 首元素为 0 + if (cuSeqlensKPtr[0] != 0) { + KERNEL_LOG_ERROR("The first element of cu_seqlens_k should be 0, but got %d", cuSeqlensKPtr[0]); + return false; + } + for (int i = 0; i < batchSize + 1; i++) { + // 校验 cu_seqlens_k 元素递增 + if (i > 0 && cuSeqlensKPtr[i - 1] > cuSeqlensKPtr[i]) { + KERNEL_LOG_ERROR("The elements in cu_seqlens_k must be in ascending order, " + "but got cu_seqlens_k[%d] = %d, cu_seqlens_k[%d] = %d", + i - 1, cuSeqlensKPtr[i - 1], i, cuSeqlensKPtr[i]); + return false; + } + } + } + } + // 校验 seqused_q 元素非负 + if (sequsedQ_ != nullptr && sequsedQ_->GetData() != nullptr) { + const int32_t *sequsedQPtr = static_cast(sequsedQ_->GetData()); + const int32_t *cuSeqlensQPtr = (layoutQ_ == "TND" && cuSeqlensQ_ != nullptr && + cuSeqlensQ_->GetData() != nullptr) ? + static_cast(cuSeqlensQ_->GetData()) : nullptr; + for (int i = 0; i < batchSize; i++) { + if (sequsedQPtr[i] < 0) { + KERNEL_LOG_ERROR("The elements in seqused_q should be >= 0, but got seqused_q[%d] = %d", i, + sequsedQPtr[i]); + return false; + } + // 校验 seqused_q 元素不大于 max_seqlen_q (BSND) 或 cu_seqlens_q 序列长度 (TND) + if (layoutQ_ == "BSND" && sequsedQPtr[i] > maxSeqlenQ_) { + KERNEL_LOG_ERROR("The elements in seqused_q should not be greater than max_seqlen_q %d, " + "but got seqused_q[%d] = %d", maxSeqlenQ_, i, sequsedQPtr[i]); + return false; + } + if (cuSeqlensQPtr != nullptr) { + int32_t seqLen = cuSeqlensQPtr[i + 1] - cuSeqlensQPtr[i]; + if (sequsedQPtr[i] > seqLen) { + KERNEL_LOG_ERROR("The elements in seqused_q should not be greater than the sequence length " + "from cu_seqlens_q %d, but got seqused_q[%d] = %d", seqLen, i, sequsedQPtr[i]); + return false; + } + } + } + } + // 校验 seqused_k 元素非负 + if (sequsedK_ != nullptr && sequsedK_->GetData() != nullptr) { + const int32_t *sequsedKPtr = static_cast(sequsedK_->GetData()); + const int32_t *cuSeqlensKPtr = (layoutK_ == "TND" && cuSeqlensK_ != nullptr && + cuSeqlensK_->GetData() != nullptr) ? + static_cast(cuSeqlensK_->GetData()) : nullptr; + for (int i = 0; i < batchSize; i++) { + if (sequsedKPtr[i] < 0) { + KERNEL_LOG_ERROR("The elements in seqused_k should be >= 0, but got seqused_k[%d] = %d", i, + sequsedKPtr[i]); + return false; + } + // 校验 seqused_k 元素不大于 max_seqlen_k (BSND) 或 cu_seqlens_k 序列长度 (TND) + if (layoutK_ == "BSND" && sequsedKPtr[i] > maxSeqlenK_) { + KERNEL_LOG_ERROR("The elements in seqused_k should not be greater than max_seqlen_k %d, " + "but got seqused_k[%d] = %d", maxSeqlenK_, i, sequsedKPtr[i]); + return false; + } + if (cuSeqlensKPtr != nullptr) { + int32_t seqLen = cuSeqlensKPtr[i + 1] - cuSeqlensKPtr[i]; + if (sequsedKPtr[i] > seqLen) { + KERNEL_LOG_ERROR("The elements in seqused_k should not be greater than the sequence length " + "from cu_seqlens_k %d, but got seqused_k[%d] = %d", seqLen, i, sequsedKPtr[i]); + return false; + } + } + } + } + // 校验 cmp_residual_k 元素 + if (cmpResidualK_ != nullptr && cmpResidualK_->GetData() != nullptr) { + const int32_t *cmpResidualKPtr = static_cast(cmpResidualK_->GetData()); + for (int i = 0; i < batchSize; i++) { + if (cmpResidualKPtr[i] < 0 || cmpResidualKPtr[i] >= cmpRatio_) { + KERNEL_LOG_ERROR("The elements in cmp_residual_k should be in [0, cmpRatio_(%d)), but got " + "cmp_residual_k[%d] = %d", cmpRatio_, + i, cmpResidualKPtr[i]); + return false; + } + } + } + return true; +} + +int32_t QuantLightningIndexerV2MetadataCpuKernel::GetQueryBatchSize() +{ + // 1. 如果sequsedQ_传了，使用sequsedQ_获取BatchSize + if (sequsedQ_ != nullptr && sequsedQ_->GetData() != nullptr) { + if (sequsedQ_->GetTensorShape() != nullptr) { + return sequsedQ_->GetTensorShape()->GetDimSize(0); + } + } + // 2. sequsedQ_ 没传，判断 Layout + if (layoutQ_ == "TND") { + // 如果是 TND，尝试使用 cuSeqlensQ_获取BatchSize + if (cuSeqlensQ_ != nullptr && cuSeqlensQ_->GetData() != nullptr) { + if (cuSeqlensQ_->GetTensorShape() != nullptr) { + return cuSeqlensQ_->GetTensorShape()->GetDimSize(0) - 1; + } + } + } + // 3. 如果不是 TND，或者 cuSeqlensQ_ 为空，使用batchSize_ + return batchSize_; +} + +ValidSocVersion QuantLightningIndexerV2MetadataCpuKernel::ProcessSocVersion() +{ + const std::string ascend950 = "Ascend950"; + if (socVersion_.find(ascend950) != std::string::npos) { + return ValidSocVersion::ASCEND950; + } else { + return ValidSocVersion::ASCEND910B; + } +} + +bool QuantLightningIndexerV2MetadataCpuKernel::ParamsInit() +{ + batchSize_ = GetQueryBatchSize(); + auto mode = static_cast(maskMode_); + if (mode == SparseMode::RIGHT_DOWN_CAUSAL) { + attentionMode_ = 1; + preToken_ = INT64_MAX; + } else if (mode == SparseMode::DEFAULT_MASK) { + attentionMode_ = 0; + } else if (mode == SparseMode::BAND) { + attentionMode_ = 1; + } + groupSize_ = numHeadsQ_ / numHeadsK_; + ValidSocVersion validSocVersion = ProcessSocVersion(); + if (validSocVersion == ValidSocVersion::ASCEND910B) { + mBaseSize_ = s1BaseSize_ * groupSize_; + s2BaseSize_ = 2048U; + } else if (validSocVersion == ValidSocVersion::ASCEND950) { + if (topk_ > TOPK_6K) { + s1BaseSize_ = S1_BASE_SIZE_SMALL; + } + mBaseSize_ = s1BaseSize_ * groupSize_; + s2BaseSize_ = 128U; + } else { + mBaseSize_ = s1BaseSize_ * groupSize_; + s2BaseSize_ = 128U; + } + return true; +} + +uint32_t QuantLightningIndexerV2MetadataCpuKernel::GetS1SeqSize(uint32_t bIdx) +{ + // 1. 如果 sequsedQ_ 传了，直接使用 + if (sequsedQ_ != nullptr && sequsedQ_->GetData() != nullptr) { + const int32_t *seqUsedPtr = static_cast(sequsedQ_->GetData()); + return static_cast(seqUsedPtr[bIdx]); + } + // 2. sequsedQ_ 没传，判断 Layout + if (layoutQ_ == "TND") { + // 如果是 TND，尝试使用 cuSeqlensQ_ + if (cuSeqlensQ_ != nullptr && cuSeqlensQ_->GetData() != nullptr) { + const int32_t *s1Ptr = static_cast(cuSeqlensQ_->GetData()); + return static_cast(s1Ptr[bIdx + 1U] - s1Ptr[bIdx]); + } + } + // 3. 如果不是 TND，或者 cuSeqlensQ_ 为空，使用 querySeqSize_ + return static_cast(maxSeqlenQ_); +} + +uint32_t QuantLightningIndexerV2MetadataCpuKernel::GetS2SeqSize(uint32_t bIdx) +{ + // 如果 seqUsedKv_ 传了，直接使用 + if (sequsedK_ != nullptr && sequsedK_->GetData() != nullptr) { + const int32_t *seqUsedPtr = static_cast(sequsedK_->GetData()); + return static_cast(seqUsedPtr[bIdx]); + } + // seqUsedKv_ 没传，判断 Layout + if (layoutK_ == "TND") { + // 如果是 TND，尝试使用 actSeqLenOriKv_ + if (cuSeqlensK_ != nullptr && cuSeqlensK_->GetData() != nullptr) { + const int32_t *s2Ptr = static_cast(cuSeqlensK_->GetData()); + return static_cast(s2Ptr[bIdx + 1U] - s2Ptr[bIdx]); + } + } + // 使用 max_seqlen_k + return static_cast(maxSeqlenK_); +} + +uint64_t QuantLightningIndexerV2MetadataCpuKernel::GetRevertS2Size(uint32_t bIdx) +{ + uint32_t cmpS2Size = GetS2SeqSize(bIdx); + if (cmpResidualK_ != nullptr && cmpResidualK_->GetData() != nullptr) { + const int32_t *residualPtr = static_cast(cmpResidualK_->GetData()); + return static_cast(cmpS2Size) * static_cast(cmpRatio_) + residualPtr[bIdx]; + } else { + return static_cast(cmpS2Size) * static_cast(cmpRatio_); + } +} + +void QuantLightningIndexerV2MetadataCpuKernel::CalcSplitInfo(SplitContext &splitContext) +{ + // 计算每个batch的切分，统计是否为空batch，记录最后有效batch（每个batch的每个N2切分是一样的） + SplitInfo &splitInfo = splitContext.splitInfo; + for (uint32_t bIdx = 0; bIdx < batchSize_; bIdx++) { + uint32_t s1Size = GetS1SeqSize(bIdx); + uint32_t s2Size = GetS2SeqSize(bIdx); + maxS2Size_ = std::max(maxS2Size_, s2Size); + splitInfo.s1GBaseNum[bIdx] = (static_cast(s1Size) * groupSize_ + (mBaseSize_ - 1U)) / mBaseSize_; + splitInfo.s1GTailSize[bIdx] = (static_cast(s1Size) * groupSize_) % mBaseSize_; + splitInfo.s2BaseNum[bIdx] = (s2Size + s2BaseSize_ - 1U) / s2BaseSize_; + splitInfo.s2TailSize[bIdx] = s2Size % s2BaseSize_; + if (splitInfo.s1GBaseNum[bIdx] != 0U && splitInfo.s2BaseNum[bIdx] != 0U) { + splitInfo.isKvSeqAllZero = false; + } + } + ValidSocVersion validSocVersion = ProcessSocVersion(); + if (validSocVersion == ValidSocVersion::ASCEND950) { + if (maxS2Size_ < topk_) { + supportFd_ = false; + return; + } + if (maxS2Size_ > fdToleranceRatio * s2BaseSize_) { + supportFd_ = true; + return; + } + } +} + +int64_t QuantLightningIndexerV2MetadataCpuKernel::CalcPreTokenLeftUp(uint32_t s1Size, uint64_t s2Size) +{ + auto mode = static_cast(maskMode_); + if (mode == SparseMode::BAND) { + return static_cast(s1Size) - static_cast(s2Size) + preToken_; + } + return preToken_; +} + +int64_t QuantLightningIndexerV2MetadataCpuKernel::CalcNextTokenLeftUp(uint32_t s1Size, uint64_t s2Size) +{ + auto mode = static_cast(maskMode_); + switch (mode) { + case SparseMode::DEFAULT_MASK: + case SparseMode::ALL_MASK: + case SparseMode::LEFT_UP_CAUSAL: + return nextToken_; + case SparseMode::RIGHT_DOWN_CAUSAL: + return static_cast(s2Size) - static_cast(s1Size); + case SparseMode::BAND: + return static_cast(s2Size) - static_cast(s1Size) + nextToken_; + default: + return nextToken_; + } +} + +int64_t QuantLightningIndexerV2MetadataCpuKernel::CalcCost(uint32_t basicM, uint32_t basicS2) +{ + uint32_t alignCoefM = 16U; + uint32_t alignCoefS2 = 64U; + uint32_t alignBasicM = (basicM + alignCoefM - 1U) >> 4U; // 按alignCoefM对齐，向上取整，4：移位操作实现除16 + uint32_t alignBasicS2 = (basicS2 + alignCoefS2 - 1U) >> 6U; // 按alignCoefS2对齐，向上取整，6：移位操作实现除64 + return static_cast(COST_WEIGHT_M * alignBasicM + COST_WEIGHT_S2 * alignBasicS2); +} + +BlockCost QuantLightningIndexerV2MetadataCpuKernel::CalcCostTable(uint32_t s1NormalSize, uint32_t s2NormalSize, + uint32_t s1GTailSize, uint32_t s2TailSize) +{ + BlockCost typeCost{}; + typeCost[NORMAL_BLOCK][NORMAL_BLOCK] = CalcCost(s1NormalSize, s2NormalSize); + typeCost[TAIL_BLOCK][NORMAL_BLOCK] = (s1GTailSize == 0U) ? 0U : CalcCost(s1GTailSize, s2NormalSize); + typeCost[NORMAL_BLOCK][TAIL_BLOCK] = (s2TailSize == 0U) ? 0U : CalcCost(s1NormalSize, s2TailSize); + typeCost[TAIL_BLOCK][TAIL_BLOCK] = (s1GTailSize == 0U || s2TailSize == 0U) ? 0U : CalcCost(s1GTailSize, s2TailSize); + return typeCost; +} + +Range QuantLightningIndexerV2MetadataCpuKernel::CalcS2TokenRange(uint32_t s1GIdx, const BatchCache &batchCache) +{ + // no mask + if (!attentionMode_) { + return std::make_pair(0, static_cast(batchCache.revertS2Size)); + } + + // 1. calc index of s2FirstToken, s2LastToken by index of s1GFirstToken, s1GLastToken + int64_t s1GFirstToken = static_cast(s1GIdx) * static_cast(mBaseSize_); + int64_t s1GLastToken = std::min(s1GFirstToken + static_cast(mBaseSize_), + static_cast(batchCache.s1Size) * static_cast(groupSize_)) - + 1; + int64_t s1FirstToken = 0; + int64_t s1LastToken = 0; + if (isS1G_) { + s1FirstToken = s1GFirstToken / static_cast(groupSize_); + s1LastToken = s1GLastToken / static_cast(groupSize_); + } else { + if (s1GFirstToken / batchCache.s1Size == s1GLastToken / batchCache.s1Size) { + // start and end locate in one G + s1FirstToken = s1GFirstToken % static_cast(batchCache.s1Size); + s1LastToken = s1GLastToken % static_cast(batchCache.s1Size); + } else { + // start and end locate in two or more G, but working same as crossing a complete block + s1FirstToken = 0; + s1LastToken = batchCache.s1Size; + } + } + + int64_t s2FirstToken = s1FirstToken - batchCache.preTokenLeftUp; + int64_t s2LastToken = s1LastToken + batchCache.nextTokenLeftUp; + + return std::make_pair(s2FirstToken, s2LastToken); +} + +void QuantLightningIndexerV2MetadataCpuKernel::CalcBatchCache(uint32_t bIdx, const SplitContext &splitContext, + BatchCache &batchCache) +{ + const SplitInfo &splitInfo = splitContext.splitInfo; + + batchCache.bIdx = bIdx; + batchCache.s1Size = GetS1SeqSize(bIdx); + batchCache.revertS2Size = GetRevertS2Size(bIdx); + batchCache.preTokenLeftUp = CalcPreTokenLeftUp(batchCache.s1Size, batchCache.revertS2Size); + batchCache.nextTokenLeftUp = CalcNextTokenLeftUp(batchCache.s1Size, batchCache.revertS2Size); +} + +void QuantLightningIndexerV2MetadataCpuKernel::CalcS1GCache(uint32_t s1GIdx, const SplitContext &splitContext, + const BatchCache &batchCache, S1GCache &s1GCache) +{ + const SplitInfo &splitInfo = splitContext.splitInfo; + + s1GCache.bIdx = batchCache.bIdx; + s1GCache.s1GIdx = s1GIdx; + + if (splitInfo.s1GBaseNum[batchCache.bIdx] == 0 || splitInfo.s2BaseNum[batchCache.bIdx] == 0) { + s1GCache.s1GBlock = 0; + s1GCache.s1GCost = 0; + s1GCache.s1GLastBlockCost = 0; + s1GCache.s1GNormalBlockCost = 0; + return; + } + + auto s2TokenRange = CalcS2TokenRange(s1GIdx, batchCache); + int64_t s2FirstToken = s2TokenRange.first; + int64_t s2LastToken = s2TokenRange.second; + // get valid range + s2FirstToken = Clip(s2FirstToken, static_cast(0), static_cast(batchCache.revertS2Size - 1U)); + s2LastToken = Clip(s2LastToken, static_cast(0), static_cast(batchCache.revertS2Size - 1U)); + + // get block start & end + int64_t s2CmpLength = (s2LastToken - s2FirstToken + 1) / cmpRatio_; + if (s2CmpLength <= 0) { + s1GCache.s2Start = 0U; + s1GCache.s2End = 0U; + } else { + s1GCache.s2Start = 0U; + // end of block index, Right-open interval + s1GCache.s2End = (static_cast(s2CmpLength) + s2BaseSize_ - 1) / s2BaseSize_; + } + + if (s1GCache.s2Start >= s1GCache.s2End) { + s1GCache.s1GBlock = 0; + s1GCache.s1GCost = 0; + s1GCache.s1GLastBlockCost = 0; + s1GCache.s1GNormalBlockCost = 0; + return; + } + + uint32_t s2TailSize = static_cast(s2CmpLength) % s2BaseSize_; + + // 计算S2方向满块、尾块数量 + s1GCache.s1GBlock = s1GCache.s2End - s1GCache.s2Start; + uint32_t curTailS2Num = s2TailSize != 0 ? 1U : 0U; + uint32_t curNormalS2Num = s1GCache.s1GBlock - curTailS2Num; + + BlockCost typeCost = + CalcCostTable(mBaseSize_, s2BaseSize_, splitInfo.s1GTailSize[batchCache.bIdx], s2TailSize); + + if (s1GIdx == (splitInfo.s1GBaseNum[batchCache.bIdx] - 1U) && splitInfo.s1GTailSize[batchCache.bIdx] != 0U) { + s1GCache.s1GCost = + typeCost[TAIL_BLOCK][NORMAL_BLOCK] * curNormalS2Num + typeCost[TAIL_BLOCK][TAIL_BLOCK] * curTailS2Num; + s1GCache.s1GLastBlockCost = + curTailS2Num > 0U ? typeCost[TAIL_BLOCK][TAIL_BLOCK] : typeCost[TAIL_BLOCK][NORMAL_BLOCK]; + s1GCache.s1GNormalBlockCost = typeCost[TAIL_BLOCK][NORMAL_BLOCK]; + } else { + s1GCache.s1GCost = + typeCost[NORMAL_BLOCK][NORMAL_BLOCK] * curNormalS2Num + typeCost[NORMAL_BLOCK][TAIL_BLOCK] * curTailS2Num; + s1GCache.s1GLastBlockCost = + curTailS2Num > 0U ? typeCost[NORMAL_BLOCK][TAIL_BLOCK] : typeCost[NORMAL_BLOCK][NORMAL_BLOCK]; + s1GCache.s1GNormalBlockCost = typeCost[NORMAL_BLOCK][NORMAL_BLOCK]; + } +} + +void QuantLightningIndexerV2MetadataCpuKernel::CalcBatchCost(uint32_t bIdx, const SplitContext &splitContext, + CostInfo &costInfo) +{ + const SplitInfo &splitInfo = splitContext.splitInfo; + + costInfo.bN2CostOfEachBatch[bIdx] = 0; + costInfo.bN2BlockOfEachBatch[bIdx] = 0U; + costInfo.bN2LastBlockCostOfEachBatch[bIdx] = 0U; + + if (GetS1SeqSize(bIdx) == 0U || GetS2SeqSize(bIdx) == 0U) { + return; + } + + BatchCache bCache; + S1GCache s1GCache; + CalcBatchCache(bIdx, splitContext, bCache); + for (uint32_t s1GIdx = 0; s1GIdx < splitInfo.s1GBaseNum[bIdx]; s1GIdx++) { + CalcS1GCache(s1GIdx, splitContext, bCache, s1GCache); + costInfo.bN2CostOfEachBatch[bIdx] += s1GCache.s1GCost; + costInfo.bN2BlockOfEachBatch[bIdx] += s1GCache.s1GBlock; + // 更新最大S1G行开销 + if (s1GCache.s1GCost > costInfo.maxS1GCost) { + costInfo.maxS1GCost = s1GCache.s1GCost; + } + + if (s1GCache.s1GBlock > 0) { + costInfo.bN2LastBlockCostOfEachBatch[bIdx] = s1GCache.s1GLastBlockCost; + } + } +} + +void QuantLightningIndexerV2MetadataCpuKernel::CalcCostInfo(SplitContext &splitContext) +{ + const SplitInfo &splitInfo = splitContext.splitInfo; + CostInfo &costInfo = splitContext.costInfo; + + if (splitInfo.isKvSeqAllZero) { + costInfo.totalCost = 0; + costInfo.totalBlockNum = 0U; + return; + } + + // 计算batch的负载并记录，用于按batch分配，需要按行计算起止点，统计块数、负载 + for (uint32_t bIdx = 0; bIdx < batchSize_; bIdx++) { + CalcBatchCost(bIdx, splitContext, costInfo); + costInfo.totalCost += costInfo.bN2CostOfEachBatch[bIdx] * numHeadsK_; + costInfo.totalBlockNum += costInfo.bN2BlockOfEachBatch[bIdx] * numHeadsK_; + } +} + +void QuantLightningIndexerV2MetadataCpuKernel::UpdateCursor(const SplitContext &splitContext, + AssignContext &assignContext) +{ + const SplitInfo &splitInfo = splitContext.splitInfo; + const CostInfo &costInfo = splitContext.costInfo; + + bool UpdateS1G = false; + bool UpdateBatch = false; + + // Update S2 + if (assignContext.curS2Idx >= assignContext.s1GCache.s2End) { // 边界assignInfo.s2End是取不到的开区间 + assignContext.curS2Idx = 0U; + assignContext.curS1GIdx++; + UpdateS1G = true; + } + + // Update S1G + if (assignContext.curS1GIdx >= splitInfo.s1GBaseNum[assignContext.curBIdx]) { + assignContext.curS1GIdx = 0U; + assignContext.curBN2Idx++; + } + + // Update Batch + if (assignContext.curBN2Idx == batchSize_ * numHeadsK_) { // 所有负载全部分配完，设置最后一个核的右开区间，返回 + assignContext.curS1GIdx = 0U; + assignContext.curS2Idx = 0U; + assignContext.isFinished = true; + return; + } + + if (assignContext.curBN2Idx / numHeadsK_ != assignContext.curBIdx) { + assignContext.curBIdx = assignContext.curBN2Idx / numHeadsK_; + assignContext.curS1GIdx = 0U; + UpdateBatch = true; + UpdateS1G = true; + } + + // Update Cache + if (UpdateBatch) { + CalcBatchCache(assignContext.curBIdx, splitContext, assignContext.batchCache); + assignContext.bN2Cost = costInfo.bN2CostOfEachBatch[assignContext.curBIdx]; + assignContext.bN2Block = costInfo.bN2BlockOfEachBatch[assignContext.curBIdx]; + } + if (UpdateS1G) { + CalcS1GCache(assignContext.curS1GIdx, splitContext, assignContext.batchCache, assignContext.s1GCache); + assignContext.curS2Idx = assignContext.s1GCache.s2Start; + } +} + +void QuantLightningIndexerV2MetadataCpuKernel::AssignByBatch(const SplitContext &splitContext, + AssignContext &assignContext) +{ + if (assignContext.isFinished) { + return; + } + const CostInfo &costInfo = splitContext.costInfo; + while (assignContext.bN2Cost == 0 || + IsWithinTolerance(assignContext.coreCache.costLimit, + costInfo.bN2LastBlockCostOfEachBatch[assignContext.curBIdx] / FA_TOLERANCE_RATIO, + assignContext.coreCache.cost + assignContext.bN2Cost)) { + assignContext.coreCache.cost += assignContext.bN2Cost; + assignContext.coreCache.block += assignContext.bN2Block; + assignContext.curBN2Idx++; + + // to the end + if (assignContext.curBN2Idx == batchSize_ * numHeadsK_) { + assignContext.curS1GIdx = 0U; + assignContext.curS2Idx = 0U; + assignContext.isFinished = true; + return; + } + + // next batch + if (assignContext.curBN2Idx / numHeadsK_ != assignContext.curBIdx) { + assignContext.curBIdx = assignContext.curBN2Idx / numHeadsK_; + CalcBatchCache(assignContext.curBIdx, splitContext, assignContext.batchCache); + } + + assignContext.bN2Cost = costInfo.bN2CostOfEachBatch[assignContext.curBIdx]; + assignContext.bN2Block = costInfo.bN2BlockOfEachBatch[assignContext.curBIdx]; + assignContext.curS1GIdx = 0U; + CalcS1GCache(assignContext.curS1GIdx, splitContext, assignContext.batchCache, assignContext.s1GCache); + assignContext.curS2Idx = assignContext.s1GCache.s2Start; + } +} + +void QuantLightningIndexerV2MetadataCpuKernel::AssignByRow(const SplitContext &splitContext, + AssignContext &assignContext) +{ + if (assignContext.isFinished) { + return; + } + + while (IsWithinTolerance(assignContext.coreCache.costLimit, + assignContext.s1GCache.s1GLastBlockCost / FA_TOLERANCE_RATIO, + assignContext.coreCache.cost + assignContext.s1GCache.s1GCost)) { + assignContext.coreCache.cost += assignContext.s1GCache.s1GCost; + assignContext.coreCache.block += assignContext.s1GCache.s1GBlock; + + // 当前batch被分配一行出去，更新剩余负载 + assignContext.bN2Cost = assignContext.bN2Cost > assignContext.s1GCache.s1GCost ? + assignContext.bN2Cost - assignContext.s1GCache.s1GCost : + 0; + assignContext.bN2Block = assignContext.bN2Block > assignContext.s1GCache.s1GBlock ? + assignContext.bN2Block - assignContext.s1GCache.s1GBlock : + 0U; + // 计算新一行的信息 + do { + assignContext.curS1GIdx++; + CalcS1GCache(assignContext.curS1GIdx, splitContext, assignContext.batchCache, assignContext.s1GCache); + } while (assignContext.s1GCache.s1GBlock == 0); + assignContext.curS2Idx = assignContext.s1GCache.s2Start; + } +} + +void QuantLightningIndexerV2MetadataCpuKernel::AssignByBlock(const SplitContext &splitContext, + AssignContext &assignContext) +{ + if (assignContext.isFinished) { + return; + } + + int64_t curCost = assignContext.s1GCache.s1GNormalBlockCost; + if (assignContext.curS2Idx == (assignContext.s1GCache.s2End - 1U)) { + curCost = assignContext.s1GCache.s1GLastBlockCost; + } + + // (costLimit - curCostOnCore) * FA_TOLERANCE_RATIO > curCost；至少分配1块 + while (IsWithinTolerance(assignContext.coreCache.costLimit, curCost / FA_TOLERANCE_RATIO, + assignContext.coreCache.cost + curCost)) { + assignContext.coreCache.cost += curCost; + assignContext.coreCache.block++; + assignContext.curS2Idx++; + // 当前batch被分配一块出去，更新剩余负载 + assignContext.bN2Cost = assignContext.bN2Cost - curCost; + // 当前行被分配一块出去，更新剩余负载 + assignContext.s1GCache.s1GCost = assignContext.s1GCache.s1GCost - curCost; + assignContext.bN2Block--; + assignContext.s1GCache.s1GBlock--; + } +} + +void QuantLightningIndexerV2MetadataCpuKernel::ForceAssign(const SplitContext &splitContext, + AssignContext &assignContext) +{ + if (assignContext.isFinished) { + return; + } + + int64_t curCost = assignContext.s1GCache.s1GNormalBlockCost; + if (assignContext.curS2Idx == (assignContext.s1GCache.s2End - 1U)) { + curCost = assignContext.s1GCache.s1GLastBlockCost; + } + + assignContext.coreCache.cost += curCost; + assignContext.coreCache.block++; + assignContext.curS2Idx++; + // 当前batch被分配一块出去，更新剩余负载 + assignContext.bN2Cost = assignContext.bN2Cost - curCost; + assignContext.bN2Block--; + // 当前行被分配一块出去，更新剩余负载 + assignContext.s1GCache.s1GCost = assignContext.s1GCache.s1GCost - curCost; + assignContext.s1GCache.s1GBlock--; + UpdateCursor(splitContext, assignContext); +} + +bool QuantLightningIndexerV2MetadataCpuKernel::IsNeedRecordFDInfo(const AssignContext &assignContext, + const SplitResult &splitRes) +{ + // 切分点大概率不会刚好在行尾，因此滞后处理归约信息的统计，到下一个切分点再判断是否需要归约 + // 核0无需处理 + if (assignContext.curCoreIdx == 0U) { + return false; + } + // 无跨核行，无需处理 + if (assignContext.curKvSplitPart <= 1U) { + return false; + } + // 需要归约的行还未处理完 + if (assignContext.curBN2Idx == splitRes.bN2End[assignContext.curCoreIdx - 1U] && + assignContext.curS1GIdx == splitRes.gS1End[assignContext.curCoreIdx - 1U]) { + return false; + } + return true; +} + +void QuantLightningIndexerV2MetadataCpuKernel::RecordFDInfo(const SplitContext &splitContext, + const AssignContext &assignContext, SplitResult &result) +{ + const SplitInfo &splitInfo = splitContext.splitInfo; + // 需要规约的行是上一个核的切分点所在位置 + uint32_t splitBIdx = result.bN2End[assignContext.curCoreIdx - 1U] / numHeadsK_; + uint32_t splitS1GIdx = result.gS1End[assignContext.curCoreIdx - 1U]; + uint32_t s1Size = GetS1SeqSize(splitBIdx); + + // 计算归约数据的FD均衡划分信息 + uint32_t curFdS1gSize = + (splitS1GIdx == splitInfo.s1GBaseNum[splitBIdx] - 1U) ? + (static_cast(s1Size) * groupSize_ - static_cast(splitS1GIdx) * mBaseSize_) : + mBaseSize_; + // 记录 + result.maxS2SplitNum = std::max(result.maxS2SplitNum, assignContext.curKvSplitPart); + // 若存在头归约，则切分点一定为上一个核结束的位置 + result.fdRes.fdBN2Idx[result.numOfFdHead] = result.bN2End[assignContext.curCoreIdx - 1U]; + result.fdRes.fdMIdx[result.numOfFdHead] = result.gS1End[assignContext.curCoreIdx - 1U]; + result.fdRes.fdWorkspaceIdx[result.numOfFdHead] = assignContext.preFdDataNum; + result.fdRes.fdS2SplitNum[result.numOfFdHead] = assignContext.curKvSplitPart; + result.fdRes.fdMSize[result.numOfFdHead] = curFdS1gSize / groupSize_; + result.numOfFdHead++; +} + +void QuantLightningIndexerV2MetadataCpuKernel::AssignBlockToCore(const SplitContext &splitContext, + AssignContext &assignContext, SplitResult &result) +{ + const CostInfo &costInfo = splitContext.costInfo; + result.firstFdDataWorkspaceIdx[assignContext.curCoreIdx] = + assignContext.preFdDataNum + assignContext.curKvSplitPart - 1U; + assignContext.coreCache = {}; + assignContext.coreCache.costLimit = assignContext.unassignedCost / (aicCoreNum_ - assignContext.curCoreIdx); + if (!supportFd_) { + assignContext.coreCache.costLimit = costInfo.maxS1GCost > assignContext.coreCache.costLimit ? + costInfo.maxS1GCost : + assignContext.coreCache.costLimit; + } + // 1、按整batch分配 + AssignByBatch(splitContext, assignContext); + // 2、按行分配 + AssignByRow(splitContext, assignContext); + // 3、按块分配 + if (supportFd_) { + AssignByBlock(splitContext, assignContext); + // 4、强制分配 + if (assignContext.coreCache.block == 0) { + ForceAssign(splitContext, assignContext); + } + } + result.bN2End[assignContext.curCoreIdx] = assignContext.curBN2Idx; + result.gS1End[assignContext.curCoreIdx] = assignContext.curS1GIdx; + result.s2End[assignContext.curCoreIdx] = assignContext.curS2Idx; + result.maxCost = std::max(result.maxCost, assignContext.coreCache.cost); + assignContext.unassignedCost -= assignContext.coreCache.cost; + // 对之前的归约信息进行记录并清理 + if (supportFd_ && IsNeedRecordFDInfo(assignContext, result)) { + RecordFDInfo(splitContext, assignContext, result); + assignContext.preFdDataNum += assignContext.curKvSplitPart; + assignContext.curKvSplitPart = 1U; + } + // 更新S2切分信息 + if (supportFd_ && assignContext.curS2Idx > assignContext.s1GCache.s2Start && + assignContext.curS2Idx <= assignContext.s1GCache.s2End) { + assignContext.curKvSplitPart++; + } +} + +void QuantLightningIndexerV2MetadataCpuKernel::CalcSplitPlan(int64_t costLimit, const SplitContext &splitContext, + SplitResult &result) +{ + const CostInfo &costInfo = splitContext.costInfo; + if (aicCoreNum_ == 0U) { + return; + } + result.maxCost = 0U; + result.usedCoreNum = 0U; + AssignContext assignContext{}; + assignContext.curBIdx = 0U; + assignContext.curS1GIdx = 0U; + assignContext.unassignedCost = costInfo.totalCost; + assignContext.bN2Cost = costInfo.bN2CostOfEachBatch[assignContext.curBIdx]; + assignContext.bN2Block = costInfo.bN2BlockOfEachBatch[assignContext.curBIdx]; + CalcBatchCache(assignContext.curBIdx, splitContext, assignContext.batchCache); + CalcS1GCache(assignContext.curS1GIdx, splitContext, assignContext.batchCache, assignContext.s1GCache); + assignContext.curS2Idx = assignContext.s1GCache.s2Start; + for (uint32_t i = 0; i < aicCoreNum_; ++i) { + if (result.maxCost > costLimit) { + return; + } + if (assignContext.isFinished || assignContext.unassignedCost <= 0) { + break; + } + assignContext.curCoreIdx = i; + AssignBlockToCore(splitContext, assignContext, result); + } + result.usedCoreNum = assignContext.curCoreIdx + 1; +} + +void QuantLightningIndexerV2MetadataCpuKernel::SplitFD(SplitResult &splitRes) +{ + // 计算FD的总数据量 + uint64_t totalFDLoad = 0; + for (uint32_t i = 0; i < splitRes.numOfFdHead; i++) { + totalFDLoad += splitRes.fdRes.fdS2SplitNum[i] * splitRes.fdRes.fdMSize[i]; + } + // 计算当前最大冗余vec核数 + uint32_t emptyVectorNum = aivCoreNum_ - splitRes.numOfFdHead; + // 计算每个核处理的load + // 向上取整，避免核负载为0 + uint64_t averageLoad = (totalFDLoad + aivCoreNum_ - 1U) / aivCoreNum_; + uint32_t curCoreIndex = 0; + for (uint32_t i = 0; i < splitRes.numOfFdHead; i++) { + // 冗余vec核数为0，此时规约任务无法进行更小的切分，只能1个vec核计算1个规约任务 + if (emptyVectorNum == 0U) { + splitRes.fdRes.fdIdx[curCoreIndex] = i; + splitRes.fdRes.fdMStart[curCoreIndex] = 0U; + splitRes.fdRes.fdMNum[curCoreIndex] = splitRes.fdRes.fdMSize[i]; + curCoreIndex++; + continue; + } + // 计算当前归约任务所用核数，向下取整，避免使用核数超出总核数 + uint32_t curFDVectorNum = splitRes.fdRes.fdS2SplitNum[i] * splitRes.fdRes.fdMSize[i] / averageLoad; + curFDVectorNum = std::max(1U, curFDVectorNum); + // 计算当前归约任务每个核的行数，向上取整，避免行数为0 + uint32_t curAveMSize = (splitRes.fdRes.fdMSize[i] + curFDVectorNum - 1U) / curFDVectorNum; + curFDVectorNum = (splitRes.fdRes.fdMSize[i] + curAveMSize - 1U) / curAveMSize; + // 需要使用的vec核数与当前剩余可用vec核数取最小 + curFDVectorNum = std::min(curFDVectorNum, emptyVectorNum + 1U); // 1: Fd任务自身带一个核 + // FD负载分配 + for (uint32_t vid = 0; vid < curFDVectorNum; vid++) { + splitRes.fdRes.fdIdx[curCoreIndex] = i; + splitRes.fdRes.fdMStart[curCoreIndex] = vid * curAveMSize; + splitRes.fdRes.fdMNum[curCoreIndex] = + (vid < curFDVectorNum - 1) ? curAveMSize : (splitRes.fdRes.fdMSize[i] - vid * curAveMSize); + curCoreIndex++; + } + // 更新冗余vec核数 + emptyVectorNum -= (curFDVectorNum - 1U); // 1: 空余核不包含FD自身的核，要-1 + } + splitRes.fdRes.fdUsedVecNum = curCoreIndex; +} + +bool QuantLightningIndexerV2MetadataCpuKernel::BalanceSchedule(SplitResult &splitRes) +{ + SplitContext splitContext(batchSize_); + // 1、划分基本块，统计信息 + CalcSplitInfo(splitContext); + // 全空case + if (splitContext.splitInfo.isKvSeqAllZero) { + splitRes.usedCoreNum = 1U; + splitRes.bN2End[0] = batchSize_ * numHeadsK_; + splitRes.gS1End[0] = 0U; + splitRes.s2End[0] = 0U; + return true; + } + CalcCostInfo(splitContext); + + splitRes.maxCost = INT64_MAX; + splitRes.usedCoreNum = 1U; + CalcSplitPlan(splitRes.maxCost, splitContext, splitRes); + // 3、存在FD任务，对FD进行负载均衡分配 + if (supportFd_ && splitRes.numOfFdHead > 0U) { + SplitFD(splitRes); + } + splitRes.usedCoreNum = std::max(splitRes.usedCoreNum, 1U); // 至少使用1个core + return true; +} + +bool QuantLightningIndexerV2MetadataCpuKernel::GenMetadata(SplitResult &splitRes) +{ + optiling::detail::QliV2Metadata *metadataPtr = static_cast(metadata_->GetData()); + *metadataPtr = {}; + // LI Metadata Generate + for (size_t i = 0; i < aicCoreNum_; ++i) { + if (i >= splitRes.usedCoreNum) { + metadataPtr->qliV2Metadata[i][QLI_V2_CORE_ENABLE_INDEX] = 0; // AIC disenable + continue; + } + metadataPtr->qliV2Metadata[i][QLI_V2_CORE_ENABLE_INDEX] = 1; // AIC enable + // FA START + metadataPtr->qliV2Metadata[i][QLI_V2_BN2_START_INDEX] = i == 0 ? 0 : splitRes.bN2End[i - 1]; + metadataPtr->qliV2Metadata[i][QLI_V2_M_START_INDEX] = i == 0 ? 0 : splitRes.gS1End[i - 1]; + metadataPtr->qliV2Metadata[i][QLI_V2_S2_START_INDEX] = i == 0 ? 0 : splitRes.s2End[i - 1]; + // FA END + metadataPtr->qliV2Metadata[i][QLI_V2_BN2_END_INDEX] = splitRes.bN2End[i]; + metadataPtr->qliV2Metadata[i][QLI_V2_M_END_INDEX] = splitRes.gS1End[i]; + metadataPtr->qliV2Metadata[i][QLI_V2_S2_END_INDEX] = splitRes.s2End[i]; + metadataPtr->qliV2Metadata[i][QLI_V2_FIRST_QLD_V2_DATA_WORKSPACE_IDX_INDEX] = + splitRes.firstFdDataWorkspaceIdx[i]; + } + + // LD Metadata Generate + for (size_t i = 0; i < aivCoreNum_; ++i) { + if (i >= splitRes.fdRes.fdUsedVecNum) { + metadataPtr->qldV2Metadata[i][QLD_V2_CORE_ENABLE_INDEX] = 0; // AIV disenable + continue; + } + metadataPtr->qldV2Metadata[i][QLD_V2_CORE_ENABLE_INDEX] = 1; // AIV enable + uint32_t curFdIdx = splitRes.fdRes.fdIdx[i]; + metadataPtr->qldV2Metadata[i][QLD_V2_BN2_IDX_INDEX] = splitRes.fdRes.fdBN2Idx[curFdIdx]; + metadataPtr->qldV2Metadata[i][QLD_V2_M_IDX_INDEX] = splitRes.fdRes.fdMIdx[curFdIdx]; + metadataPtr->qldV2Metadata[i][QLD_V2_WORKSPACE_IDX_INDEX] = splitRes.fdRes.fdWorkspaceIdx[curFdIdx]; + metadataPtr->qldV2Metadata[i][QLD_V2_WORKSPACE_NUM_INDEX] = splitRes.fdRes.fdS2SplitNum[curFdIdx]; + metadataPtr->qldV2Metadata[i][QLD_V2_M_START_INDEX] = splitRes.fdRes.fdMStart[i]; + metadataPtr->qldV2Metadata[i][QLD_V2_M_NUM_INDEX] = splitRes.fdRes.fdMNum[i]; + } + return true; +} + +namespace { +static const char *kernelType = "QuantLightningIndexerV2Metadata"; +REGISTER_CPU_KERNEL(kernelType, QuantLightningIndexerV2MetadataCpuKernel); +} // namespace +}; // namespace aicpu diff --git a/xllm_ops/attention/quant_lightning_indexer_v2_metadata/op_kernel_aicpu/quant_lightning_indexer_v2_metadata_aicpu.h b/xllm_ops/attention/quant_lightning_indexer_v2_metadata/op_kernel_aicpu/quant_lightning_indexer_v2_metadata_aicpu.h new file mode 100644 index 0000000..6cbc92d --- /dev/null +++ b/xllm_ops/attention/quant_lightning_indexer_v2_metadata/op_kernel_aicpu/quant_lightning_indexer_v2_metadata_aicpu.h @@ -0,0 +1,302 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file quant_lightning_indexer_v2_metadata_aicpu.h + * \brief + */ + +#ifndef QUANT_LIGHTNING_INDEXER_V2_METADATA_AICPU_H +#define QUANT_LIGHTNING_INDEXER_V2_METADATA_AICPU_H + +#include +#include +#include +#include "log.h" +#include "cpu_context.h" +#include "cpu_kernel.h" +#include "cpu_tensor.h" +#include "../../quant_lightning_indexer_v2/op_kernel/quant_lightning_indexer_v2_metadata.h" +#include "../../common/op_kernel/aicpu_common.h" + +namespace aicpu { +constexpr int64_t FA_TOLERANCE_RATIO = 2; +constexpr uint32_t COST_WEIGHT_M = 6U; +constexpr uint32_t COST_WEIGHT_S2 = 10U; +constexpr uint32_t S1_BASE_SIZE_SMALL = 2; +constexpr uint32_t TOPK_6K = 6144; + +enum BlockType : uint32_t { NORMAL_BLOCK = 0, TAIL_BLOCK, BLOCK_MAX_TYPE }; + +enum class SparseMode : uint8_t { + DEFAULT_MASK = 0, + ALL_MASK, + LEFT_UP_CAUSAL, + RIGHT_DOWN_CAUSAL, + BAND, + SPARSE_BUTT, +}; + +enum class ValidSocVersion { ASCEND910B = 0, ASCEND950, RESERVED_VERSION = 99999 }; + +template +using Range = std::pair; + +template +using BlockCost = std::array(BLOCK_MAX_TYPE)>, static_cast(BLOCK_MAX_TYPE)>; + +template +T Clip(T value, T minValue, T maxValue) +{ + if (value < minValue) { + return minValue; + } + if (value > maxValue) { + return maxValue; + } + return value; +} + +template +inline bool IsWithinTolerance(T limit, T tolerance, T value) +{ + return limit + tolerance >= value; +} + +// 分核功能模块输出：FD信息，包含需要归约的数据索引及其分核信息 +struct FlashDecodeResult { + uint32_t fdUsedVecNum{0U}; // 归约过程使用的vector数量 + // 1、归约任务的索引信息 + std::vector fdBN2Idx{}; // 每个归约任务的BN2索引，脚标为归约任务的序号，最大为核数-1 + std::vector fdMIdx{}; // 每个归约任务的GS1索引，脚标为归约任务的序号 + std::vector fdWorkspaceIdx{}; // 每个归约任务在workspace中的存放位置 + std::vector fdS2SplitNum{}; // 每个归约任务的S2核间切分份数，脚标为归约任务的序号 + std::vector fdMSize{}; // 每个归约任务m轴大小，脚标为归约任务的序号 + // 2、FD负载均衡阶段，归约任务的分核（vec）信息 + std::vector fdIdx{}; // FD负载均衡阶段，每个vector处理的归约任务对应ID + std::vector fdMStart{}; // FD负载均衡阶段，每个vector处理的归约任务的m轴起点 + std::vector fdMNum{}; // FD负载均衡阶段，每个vector处理的归约任务的m轴行数 + + FlashDecodeResult(uint32_t aicNum, uint32_t aivNum) + : fdBN2Idx(aicNum), + fdMIdx(aicNum), + fdWorkspaceIdx(aicNum), + fdS2SplitNum(aicNum), + fdMSize(aicNum), + fdIdx(aivNum), + fdMStart(aivNum), + fdMNum(aivNum) + {} +}; + +// 分核功能模块输出：FA阶段的核间分核信息 +struct SplitResult { + uint32_t usedCoreNum{0U}; // 使用的核数量 + std::vector bN2End{}; // 每个核处理数据的BN2结束点 + std::vector gS1End{}; // 每个核处理数据的GS1结束点 + std::vector s2End{}; // 每个核处理数据的S2结束点 + std::vector firstFdDataWorkspaceIdx{}; // 每个核第一份归约任务的存放位置 + int64_t maxCost{0}; // 慢核开销 + uint32_t numOfFdHead{0U}; // 归约任务数量 + uint32_t maxS2SplitNum{0U}; // 单个归约任务最大分核数量 + FlashDecodeResult fdRes{0U, 0U}; // FD信息 + + SplitResult(uint32_t aicNum, uint32_t aivNum) + : bN2End(aicNum), gS1End(aicNum), s2End(aicNum), firstFdDataWorkspaceIdx(aicNum), fdRes(aicNum, aivNum) {}; +}; + +// 分核功能模块内部使用：记录切分信息 +struct SplitInfo { + std::vector s1GBaseNum{}; // S1G方向，切了多少个基本块 + std::vector s2BaseNum{}; // S2方向，切了多少个基本块 + std::vector s1GTailSize{}; // S1G方向，尾块size + std::vector s2TailSize{}; // S2方向，尾块size + bool isKvSeqAllZero{true}; + + explicit SplitInfo(uint32_t batchSize) + : s1GBaseNum(batchSize), s2BaseNum(batchSize), s1GTailSize(batchSize), s2TailSize(batchSize) + {} +}; + +// 分核功能模块内部使用：记录batch的开销信息 +struct CostInfo { + std::vector bN2CostOfEachBatch{}; // 整个batch的开销 + std::vector bN2BlockOfEachBatch{}; // 整个batch的开销 + std::vector bN2LastBlockCostOfEachBatch{}; // batch最后一块的开销 + uint32_t totalBlockNum{0U}; + int64_t totalCost{0}; + uint64_t maxS1GCost{0}; // 新增 + + explicit CostInfo(uint32_t batchSize) + : bN2CostOfEachBatch(batchSize), bN2BlockOfEachBatch(batchSize), bN2LastBlockCostOfEachBatch(batchSize) + {} +}; + +// 分核功能模块内部使用：分核过程中，case基本信息的上下文信息，组合以减少接口传参数量 +struct SplitContext { + SplitInfo splitInfo{0U}; + CostInfo costInfo{0U}; + + explicit SplitContext(uint32_t batchSize) : splitInfo(batchSize), costInfo(batchSize) {} +}; + +// 分核功能模块内部使用：记录batch相关的临时信息 +struct BatchCache { + uint32_t bIdx{0U}; + uint32_t s1Size{0U}; + uint64_t revertS2Size{0U}; + int64_t preTokenLeftUp{0}; + int64_t nextTokenLeftUp{0}; + BlockCost typeCost{}; +}; + +// 分核功能模块内部使用：记录当前行（S1G）的临时信息 +struct S1GCache { + uint32_t bIdx{0U}; + uint32_t s1GIdx{0U}; + uint32_t s2Start{0U}; + uint32_t s2End{0U}; + int64_t s1GCost{0}; + int64_t s1GLastBlockCost{0}; + uint32_t s1GBlock{0U}; + int64_t s1GNormalBlockCost{0}; +}; + +// 分核功能模块内部使用：记录分配过程中，当前核的负载信息 +struct CoreCache { + int64_t costLimit{0}; // 负载上限 + int64_t cost{0}; // 已分配负载 + uint32_t block{0U}; // 已分配块数 +}; + +// 分核功能模块内部使用：记录分配过程中的上下文信息 +struct AssignContext { + uint32_t curBIdx{0U}; + uint32_t curBN2Idx{0U}; + uint32_t curS1GIdx{0U}; + uint32_t curS2Idx{0U}; + uint32_t curCoreIdx{0U}; + int64_t unassignedCost{0}; + uint32_t curKvSplitPart{1U}; + uint32_t preFdDataNum{0U}; + + int64_t bN2Cost{0}; + uint32_t bN2Block{0U}; + bool isFinished{false}; + BatchCache batchCache{}; + S1GCache s1GCache{}; + CoreCache coreCache{}; +}; +class QuantLightningIndexerV2MetadataCpuKernel : public CpuKernel { +public: + QuantLightningIndexerV2MetadataCpuKernel() = default; + ~QuantLightningIndexerV2MetadataCpuKernel() = default; + uint32_t Compute(CpuKernelContext &ctx) override; + +private: + bool Prepare(CpuKernelContext &ctx); + bool ParamsCheck(); + int32_t GetQueryBatchSize(); + ValidSocVersion ProcessSocVersion(); + bool ParamsInit(); + bool BalanceSchedule(SplitResult &splitRes); + bool GenMetadata(SplitResult &splitRes); + + // util + uint32_t GetS1SeqSize(uint32_t bIdx); + uint32_t GetS2SeqSize(uint32_t bIdx); + uint64_t GetRevertS2Size(uint32_t bIdx); + int64_t CalcPreTokenLeftUp(uint32_t s1Size, uint64_t s2Size); + int64_t CalcNextTokenLeftUp(uint32_t s1Size, uint64_t s2Size); + Range CalcS2TokenRange(uint32_t s1GIdx, const BatchCache &batchCache); + int64_t CalcCost(uint32_t basicM, uint32_t basicS2); + BlockCost CalcCostTable(uint32_t s1NormalSize, uint32_t s2NormalSize, uint32_t s1GTailSize, + uint32_t s2TailSize); + + // cache calculation + void CalcBatchCache(uint32_t bIdx, const SplitContext &splitContext, BatchCache &batchCache); + void CalcS1GCache(uint32_t s1GIdx, const SplitContext &splitContext, const BatchCache &batchCache, + S1GCache &s1GCache); + + // preprocess + void CalcSplitInfo(SplitContext &splitContext); + void CalcBatchCost(uint32_t bIdx, const SplitContext &splitContext, CostInfo &costInfo); + void CalcCostInfo(SplitContext &splitContext); + + // assign + void UpdateCursor(const SplitContext &splitContext, AssignContext &assignContext); + void AssignByBatch(const SplitContext &splitContext, AssignContext &assignContext); + void AssignByRow(const SplitContext &splitContext, AssignContext &assignContext); + void AssignByBlock(const SplitContext &splitContext, AssignContext &assignContext); + void ForceAssign(const SplitContext &splitContext, AssignContext &assignContext); + void AssignBlockToCore(const SplitContext &splitContext, AssignContext &assignContext, SplitResult &result); + + // FD + bool IsNeedRecordFDInfo(const AssignContext &assignContext, const SplitResult &splitRes); + void RecordFDInfo(const SplitContext &splitContext, const AssignContext &assignContext, SplitResult &result); + + // main + void SplitFD(SplitResult &splitRes); + void CalcSplitPlan(int64_t costLimit, const SplitContext &splitContext, SplitResult &result); + +private: + CpuKernelContext *context_ = nullptr; + // input + Tensor *cuSeqlensQ_ = nullptr; + Tensor *cuSeqlensK_ = nullptr; + Tensor *sequsedQ_ = nullptr; + Tensor *sequsedK_ = nullptr; + Tensor *cmpResidualK_ = nullptr; + // output + Tensor *metadata_ = nullptr; + // attributes + std::string socVersion_ = ""; + bool supportFd_ = false; + uint32_t aicCoreNum_ = optiling::AIC_CORE_MAX_NUM; + uint32_t aivCoreNum_ = optiling::AIV_CORE_MAX_NUM; + int32_t batchSize_ = 0; + int32_t maxSeqlenQ_ = 0; + int32_t maxSeqlenK_ = 0; + int32_t numHeadsQ_ = 0; + int32_t numHeadsK_ = 0; + int32_t headDim_ = 0; + int32_t quantMode_ = 0; + int32_t topk_ = 0; + std::string layoutQ_ = "BSND"; + std::string layoutK_ = "BSND"; + int32_t maskMode_ = 0; + int32_t cmpRatio_ = 1; + uint32_t attentionMode_ = 0; + + // SplitParams + int64_t preToken_ = INT64_MAX; + int64_t nextToken_ = INT64_MAX; + uint32_t groupSize_ = 0; + uint32_t mBaseSize_ = 256; + uint32_t s1BaseSize_ = 4; + uint32_t s2BaseSize_ = 128; + bool isS1G_ = true; + uint32_t maxS2Size_ = 0; + uint32_t fdToleranceRatio = 5U; + +private: + enum class ParamId : uint32_t { + // input + actSeqLenQ = 0, + actSeqLenK = 1, + seqUsedQ = 2, + seqUsedK = 3, + cmpResidualK = 4, + // output + metadata = 0, + }; +}; +} // namespace aicpu +#endif diff --git a/xllm_ops/attention/quant_lightning_indexer_v2_metadata/op_kernel_aicpu/quant_lightning_indexer_v2_metadata_aicpu.json b/xllm_ops/attention/quant_lightning_indexer_v2_metadata/op_kernel_aicpu/quant_lightning_indexer_v2_metadata_aicpu.json new file mode 100644 index 0000000..a56f16c --- /dev/null +++ b/xllm_ops/attention/quant_lightning_indexer_v2_metadata/op_kernel_aicpu/quant_lightning_indexer_v2_metadata_aicpu.json @@ -0,0 +1,15 @@ +{ + "QuantLightningIndexerV2Metadata":{ + "opInfo":{ + "computeCost":"100", + "engine":"DNN_VM_AICPU", + "flagAsync":"False", + "flagPartial":"False", + "functionName":"RunCpuKernel", + "kernelSo":"libtransformer_aicpu_kernels.so", + "opKernelLib":"CUSTAICPUKernel", + "userDefined":"True", + "workspaceSize":"100" + } + } +} \ No newline at end of file diff --git a/xllm_ops/build_aclnn.sh b/xllm_ops/build_aclnn.sh index 935e8e7..a4a6fd8 100644 --- a/xllm_ops/build_aclnn.sh +++ b/xllm_ops/build_aclnn.sh @@ -193,6 +193,8 @@ elif [[ "$SOC_VERSION" =~ ^ascend910_93 ]]; then "compressor" "quant_lightning_indexer" ## 已在 CANN 中内置，见 opp/built-in/op_impl/ai_core/tbe/impl/ops_transformer/ascendc/quant_lightning_indexer "quant_lightning_indexer_metadata" + "quant_lightning_indexer_v2" + "quant_lightning_indexer_v2_metadata" "sparse_attn_sharedkv" "sparse_attn_sharedkv_metadata" "hc_pre_sinkhorn"