mixed_quant_sparse_flash_mla算子测试框架
功能说明
基于pytest测试框架，实现mixed_quant_sparse_flash_mla算子的功能验证：
- CPU侧：复现算子功能用以生成golden数据
- NPU侧：通过torch_npu进行算子直调、图模式调用获取实际数据
- 精度对比：进行CPU与NPU结果的精度对比验证算子功能
当前实现范围
参数限制
- 数据格式:
- layout_q：BSND、TND
- layout_kv：PA_BBND
环境配置
前置要求
1、 确认torch_npu为最新版本
custom包调用
支持custom包调用
文件结构
pytest文件结构说明
- test_run.sh # 用例执行脚本
- mixed_quant_sparse_flash_mla_golden.py # 算子入参处理及cpu侧golden实现
- result_compare_method.py # cpu golden与npu结果精度对比脚本
- pytest.ini # 创建测试标记
单用例测试：
- test_mixed_quant_sparse_flash_mla_single.py # pytest测试单用例运行主程序
- mixed_quant_sparse_flash_mla_paramset.py # 单用例入参数配置
批量测试：
- test_mixed_quant_sparse_flash_mla_batch.py # 用例批量测试主程序，读取pt用例并获取npu输出，结果保存至excel文件
- ./batch/test_mixed_quant_sparse_flash_mla_pt_save.py # 读取excel表格批量生成用例pt文件
- ./batch/mixed_quant_sparse_flash_mla_process.py # 调用算子接口获取npu输出
使用方法
在pytest文件夹路径下执行：
运行测试用例
单用例测试
1、手动配置mixed_quant_sparse_flash_mla_paramset.py的参数
2、执行指令：
bash test_run.sh single
批量测试
支持三种批量测试模式，均可通过选项指定excel文件、sheet名、pt保存目录：
| 命令 | 功能 |
|---|---|
batch_save |
从excel读取用例 → golden计算 → 保存pt文件 |
batch_exec |
读取pt文件 → NPU测试（需先 batch_save） |
batch |
全流程：batch_save + batch_exec，默认完成后清理pt文件 |
选项参数
| 选项 | 说明 | 默认值 | 适用命令 |
|---|---|---|---|
--excel <路径> |
指定excel文件路径 | ./excel/example.xlsx |
batch_save / batch_exec / batch |
--sheet <名称> |
指定excel sheet名 | decode |
batch_save / batch |
--pt-dir <目录> |
指定pt文件保存/读取目录 | qsmla_testcase |
batch_save / batch_exec / batch |
--keep-pt |
执行完成后保留pt文件 | 默认清理 | 仅 batch |
使用示例
# batch_save：只生成pt文件
bash test_run.sh batch_save
bash test_run.sh batch_save --excel my.xlsx --sheet prefill --pt-dir my_pt
# batch_exec：只执行NPU测试（需先batch_save生成pt文件）
bash test_run.sh batch_exec
bash test_run.sh batch_exec --pt-dir my_pt
# batch：全流程（默认清理pt文件）
bash test_run.sh batch
# batch：全流程，保留pt文件
bash test_run.sh batch --keep-pt
# batch：全流程，自定义所有参数
bash test_run.sh batch --excel my.xlsx --sheet prefill --pt-dir my_pt --keep-pt