已合并
bugfix: golden行无效逻辑修复+ lse适配 #10157
Ccccccc创建于 9 天前
bugfix: golden行无效逻辑修复+ lse适配 #10157
已合并
共 9 个文件变更+211-138
| @@ -12,9 +12,6 @@ | |||
| 12 | 12 | ||
| 13 | import os | 13 | import os |
| 14 | import torch | 14 | import torch |
| 15 | -import torch_npu | ||
| 16 | -import check_valid_param | ||
| 17 | -import pytest | ||
| 18 | import random | 15 | import random |
| 19 | import numpy as np | 16 | import numpy as np |
| 20 | import math | 17 | import math |
| @@ -233,13 +230,6 @@ class GeneralizedSFAQuant: | |||
| 233 | print( | 230 | print( |
| 234 | f" 进度：{current_pct:.1f}% | 步数：{i_S1:>{len(str(cur_act_q))}}/{cur_act_q}" | 231 | f" 进度：{current_pct:.1f}% | 步数：{i_S1:>{len(str(cur_act_q))}}/{cur_act_q}" |
| 235 | ) | 232 | ) |
| 236 | - if ( | ||
| 237 | - self.ori_mask_mode == 3 and i_S1 < cur_act_q - cur_ori_act_kv | ||
| 238 | - ): # 根据 win_kv 判断行无效 | ||
| 239 | - attn_out[i_B, i_N2 * G : (i_N2 + 1) * G, i_S1, :] = torch.zeros( | ||
| 240 | - [G, self.D], dtype=torch.float | ||
| 241 | - ) | ||
| 242 | - continue | ||
| 243 | ori_threshold = cur_ori_act_kv - cur_act_q + i_S1 + 1 | 233 | ori_threshold = cur_ori_act_kv - cur_act_q + i_S1 + 1 |
| 244 | if self.ori_mask_mode == 3: | 234 | if self.ori_mask_mode == 3: |
| 245 | ori_win_start = 0 | 235 | ori_win_start = 0 |
| @@ -263,13 +253,14 @@ class GeneralizedSFAQuant: | |||
| 263 | ori_win_end = cur_ori_act_kv | 253 | ori_win_end = cur_ori_act_kv |
| 264 | ori_win_start = min(ori_win_start, cur_ori_act_kv) | 254 | ori_win_start = min(ori_win_start, cur_ori_act_kv) |
| 265 | 255 | ||
| 266 | - if ori_win_start >= ori_win_end: # 根据 win_kv 判断行无效 | 256 | + if ori_win_start >= ori_win_end: |
| 267 | - attn_out[i_B, i_N2 * G : (i_N2 + 1) * G, i_S1, :] = torch.zeros( | 257 | + cur_ori_k_bnsd = torch.zeros( |
| 268 | - [G, self.D], dtype=torch.float | 258 | + [0, self.D], dtype=ori_k_bnsd.dtype |
| 269 | ) | 259 | ) |
| 270 | - continue | 260 | + else: |
| 271 | - | 261 | + cur_ori_k_bnsd = ori_k_bnsd[ |
| 272 | - cur_ori_k_bnsd = ori_k_bnsd[i_B, i_N2, ori_win_start:ori_win_end, :] | 262 | + i_B, i_N2, ori_win_start:ori_win_end, : |
| 263 | + ] | ||
| 273 | 264 | ||
| 274 | if ( | 265 | if ( |
| 275 | self.template_run_mode == "CSA" | 266 | self.template_run_mode == "CSA" |
| @@ -351,6 +342,12 @@ class GeneralizedSFAQuant: | |||
| 351 | cmp_s2_loop_time = math.ceil(cur_cmp_k.size(0) / s2_base_size) | 342 | cmp_s2_loop_time = math.ceil(cur_cmp_k.size(0) / s2_base_size) |
| 352 | cur_cmp_k_fp32 = cur_cmp_k.to(dtype=torch.float32) | 343 | cur_cmp_k_fp32 = cur_cmp_k.to(dtype=torch.float32) |
| 353 | 344 | ||
| 345 | + if cur_ori_k_bnsd.size(0) == 0 and cur_cmp_k == []: | ||
| 346 | + attn_out[i_B, i_N2 * G : (i_N2 + 1) * G, i_S1, :] = torch.zeros( | ||
| 347 | + [G, self.D], dtype=torch.float | ||
| 348 | + ) | ||
| 349 | + continue | ||
| 350 | + | ||
| 354 | cur_attn_out = attn_out[i_B, i_N2 * G : (i_N2 + 1) * G, i_S1, :] | 351 | cur_attn_out = attn_out[i_B, i_N2 * G : (i_N2 + 1) * G, i_S1, :] |
| 355 | q_curr = q_bnsd[i_B, i_N2 * G : (i_N2 + 1) * G, i_S1, :] | 352 | q_curr = q_bnsd[i_B, i_N2 * G : (i_N2 + 1) * G, i_S1, :] |
| 356 | q_curr_fp32 = q_curr.to(dtype=torch.float32) | 353 | q_curr_fp32 = q_curr.to(dtype=torch.float32) |
| @@ -364,13 +361,17 @@ class GeneralizedSFAQuant: | |||
| 364 | 361 | ||
| 365 | mm1_res = torch.matmul(q_curr_fp32, k_concat_fp32.T) | 362 | mm1_res = torch.matmul(q_curr_fp32, k_concat_fp32.T) |
| 366 | scale_res = mm1_res * self.softmax_scale | 363 | scale_res = mm1_res * self.softmax_scale |
| 367 | - softmax_res, softmax_sum = self.sinks_softmax( | 364 | + softmax_res, x_max, softmax_sum = self.sinks_softmax( |
| 368 | scale_res, cur_sinks_expand | 365 | scale_res, cur_sinks_expand |
| 369 | ) | 366 | ) |
| 370 | softmax_res = softmax_res.to(q_bnsd.dtype).to(torch.float32) | 367 | softmax_res = softmax_res.to(q_bnsd.dtype).to(torch.float32) |
| 371 | mm2_res = torch.matmul(softmax_res, v_concat_fp32) | 368 | mm2_res = torch.matmul(softmax_res, v_concat_fp32) |
| 372 | v2_res = torch.div(mm2_res, softmax_sum) | 369 | v2_res = torch.div(mm2_res, softmax_sum) |
| 373 | attn_out[i_B, i_N2 * G : (i_N2 + 1) * G, i_S1, :] = v2_res | 370 | attn_out[i_B, i_N2 * G : (i_N2 + 1) * G, i_S1, :] = v2_res |
| 371 | + if return_softmax_lse: | ||
| 372 | + softmax_lse[i_B, i_N2, i_S1, :] = x_max[:, 0] + torch.log( | ||
| 373 | + softmax_sum[:, 0] + 1e-10 | ||
| 374 | + ) | ||
| 374 | elif RUN_MODE == 1: | 375 | elif RUN_MODE == 1: |
| 375 | ori_s2_loop_time = math.ceil( | 376 | ori_s2_loop_time = math.ceil( |
| 376 | cur_ori_k_bnsd.size(0) / s2_base_size | 377 | cur_ori_k_bnsd.size(0) / s2_base_size |
| @@ -559,7 +560,7 @@ class GeneralizedSFAQuant: | |||
| 559 | x_sub = x - x_max | 560 | x_sub = x - x_max |
| 560 | y = torch.exp(x_sub) | 561 | y = torch.exp(x_sub) |
| 561 | x_sum = y.sum(dim=-1, keepdims=True) + torch.exp(sinks - x_max) | 562 | x_sum = y.sum(dim=-1, keepdims=True) + torch.exp(sinks - x_max) |
| 562 | - return y, x_sum | 563 | + return y, x_max, x_sum |
| 563 | 564 | ||
| 564 | def trans_shape_to_bnsd( | 565 | def trans_shape_to_bnsd( |
| 565 | self, tensor, shape, layout, cu_seqlens_q=None, seqused_q=None | 566 | self, tensor, shape, layout, cu_seqlens_q=None, seqused_q=None |
Mattention/mixed_quant_sparse_flash_mla/tests/pytest/test_mixed_quant_sparse_flash_mla_batch.py+39-24
| @@ -10,38 +10,41 @@ | |||
| 10 | # See LICENSE in the root of the software repository for the full text of the License. | 10 | # See LICENSE in the root of the software repository for the full text of the License. |
| 11 | # ----------------------------------------------------------------------------------------------------------- | 11 | # ----------------------------------------------------------------------------------------------------------- |
| 12 | 12 | ||
| 13 | -import itertools | ||
| 14 | import torch | 13 | import torch |
| 15 | import pytest | 14 | import pytest |
| 16 | -import random | ||
| 17 | import pandas as pd | 15 | import pandas as pd |
| 18 | from pathlib import Path | 16 | from pathlib import Path |
| 19 | -import numpy as np | ||
| 20 | -import math | ||
| 21 | import os | 17 | import os |
| 22 | -import multiprocessing as mp | ||
| 23 | import concurrent.futures | 18 | import concurrent.futures |
| 24 | import result_compare_method | 19 | import result_compare_method |
| 25 | -import check_valid_param | ||
| 26 | from batch import mixed_quant_sparse_flash_mla_process | 20 | from batch import mixed_quant_sparse_flash_mla_process |
| 27 | import utils | 21 | import utils |
| 28 | 22 | ||
| 29 | testcase_path = os.environ.get("MQSMLA_PT_DIR", "mqsmla_testcase") | 23 | testcase_path = os.environ.get("MQSMLA_PT_DIR", "mqsmla_testcase") |
| 30 | -batch_test_mode = int(os.environ.get("MQSMLA_BATCH_TEST_MODE", 0)) # 0:路径下全量批跑，1:按表格中case批跑 | 24 | +batch_test_mode = int( |
| 31 | -excel_path = os.environ.get("MQSMLA_EXCEL_PATH", os.path.join(os.path.dirname(__file__), "excel", "testcase.xlsx")) | 25 | + os.environ.get("MQSMLA_BATCH_TEST_MODE", 0) |
| 32 | -result_path = os.environ.get("MQSMLA_RESULT_SAVE_PATH", './mqsmla_result.xlsx') | 26 | +) # 0:路径下全量批跑，1:按表格中case批跑 |
| 27 | +excel_path = os.environ.get( | ||
| 28 | + "MQSMLA_EXCEL_PATH", | ||
| 29 | + os.path.join(os.path.dirname(__file__), "excel", "testcase.xlsx"), | ||
| 30 | +) | ||
| 31 | +result_path = os.environ.get("MQSMLA_RESULT_SAVE_PATH", "./mqsmla_result.xlsx") | ||
| 33 | device_id = int(os.environ.get("MQSMLA_DEVICE_ID", 0)) | 32 | device_id = int(os.environ.get("MQSMLA_DEVICE_ID", 0)) |
| 34 | 33 | ||
| 35 | testcase_files = [] | 34 | testcase_files = [] |
| 36 | if os.path.isdir(testcase_path): | 35 | if os.path.isdir(testcase_path): |
| 37 | if batch_test_mode == 1: | 36 | if batch_test_mode == 1: |
| 38 | df = pd.read_excel(excel_path) | 37 | df = pd.read_excel(excel_path) |
| 39 | - target_names = [str(name) for name in df['Testcase_Name'].dropna().tolist() if str(name) != 'None'] | 38 | + target_names = [ |
| 39 | + str(name) | ||
| 40 | + for name in df["Testcase_Name"].dropna().tolist() | ||
| 41 | + if str(name) != "None" | ||
| 42 | + ] | ||
| 40 | if not target_names: | 43 | if not target_names: |
| 41 | print(f"错误: 表格中没有有效的Testcase_Name: {excel_path}") | 44 | print(f"错误: 表格中没有有效的Testcase_Name: {excel_path}") |
| 42 | else: | 45 | else: |
| 43 | print(f"从表格中读取到 {len(target_names)} 个目标用例名") | 46 | print(f"从表格中读取到 {len(target_names)} 个目标用例名") |
| 44 | - pt_files = [f for f in os.listdir(testcase_path) if f.endswith('.pt')] | 47 | + pt_files = [f for f in os.listdir(testcase_path) if f.endswith(".pt")] |
W | |||
| 45 | for target_name in target_names: | 48 | for target_name in target_names: |
| 46 | matched = [f for f in pt_files if target_name in f] | 49 | matched = [f for f in pt_files if target_name in f] |
| 47 | if matched: | 50 | if matched: |
| @@ -53,7 +56,7 @@ if os.path.isdir(testcase_path): | |||
| 53 | print(f"警告: 用例名 '{target_name}' 未匹配到任何.pt文件") | 56 | print(f"警告: 用例名 '{target_name}' 未匹配到任何.pt文件") |
| 54 | print(f"按表格筛选后共 {len(testcase_files)} 个测试用例文件") | 57 | print(f"按表格筛选后共 {len(testcase_files)} 个测试用例文件") |
| 55 | else: | 58 | else: |
| 56 | - pt_files = [f for f in os.listdir(testcase_path) if f.endswith('.pt')] | 59 | + pt_files = [f for f in os.listdir(testcase_path) if f.endswith(".pt")] |
| 57 | if not pt_files: | 60 | if not pt_files: |
| 58 | print(f"错误: 目录中没有找到.pt文件: {testcase_path}") | 61 | print(f"错误: 目录中没有找到.pt文件: {testcase_path}") |
| 59 | else: | 62 | else: |
| @@ -64,22 +67,28 @@ if os.path.isdir(testcase_path): | |||
| 64 | else: | 67 | else: |
| 65 | print(f"错误: 输出目录不存在: {testcase_path}") | 68 | print(f"错误: 输出目录不存在: {testcase_path}") |
| 66 | 69 | ||
| 70 | + | ||
| 67 | def mqsmla(testcase_files): | 71 | def mqsmla(testcase_files): |
| 68 | test_data = torch.load(testcase_files, map_location="cpu", weights_only=False) | 72 | test_data = torch.load(testcase_files, map_location="cpu", weights_only=False) |
| 69 | npu_error_msg = None | 73 | npu_error_msg = None |
| 70 | try: | 74 | try: |
| 71 | - npu_result, cpu_quant_result, cpu_lse, npu_lse = mixed_quant_sparse_flash_mla_process.test_mqsmla_quant_process_ci( | 75 | + npu_result, cpu_quant_result, cpu_lse, npu_lse = ( |
| 72 | - test_data, device_id=device_id) | 76 | + mixed_quant_sparse_flash_mla_process.test_mqsmla_quant_process_ci( |
| 73 | - attn_result, attn_percent = result_compare_method.check_result(cpu_quant_result, npu_result) | 77 | + test_data, device_id=device_id |
| 78 | + ) | ||
| 79 | + ) | ||
| 80 | + attn_result, attn_percent = result_compare_method.check_result( | ||
| 81 | + cpu_quant_result, npu_result | ||
| 82 | + ) | ||
| 74 | lse_res, lse_pct = None, None | 83 | lse_res, lse_pct = None, None |
| 75 | fail_info = [] | 84 | fail_info = [] |
| 76 | min_fulfill = attn_percent # 修复：main_pct → attn_percent | 85 | min_fulfill = attn_percent # 修复：main_pct → attn_percent |
| 77 | 86 | ||
| 78 | - if test_data['params'].get('return_softmax_lse'): | 87 | + if test_data["params"].get("return_softmax_lse"): |
| 79 | print("return_softmax_lse is true!!!") | 88 | print("return_softmax_lse is true!!!") |
| 80 | lse_res, lse_pct = result_compare_method.check_result(cpu_lse, npu_lse) | 89 | lse_res, lse_pct = result_compare_method.check_result(cpu_lse, npu_lse) |
| 81 | min_fulfill = min(min_fulfill, lse_pct) | 90 | min_fulfill = min(min_fulfill, lse_pct) |
| 82 | - | 91 | + |
| 83 | if attn_result != "Pass": | 92 | if attn_result != "Pass": |
| 84 | fail_info.append(f"MAIN_FAILED:{attn_result}") | 93 | fail_info.append(f"MAIN_FAILED:{attn_result}") |
| 85 | # LSE失败记录 | 94 | # LSE失败记录 |
| @@ -98,15 +107,21 @@ def mqsmla(testcase_files): | |||
| 98 | result = "NPU ERROR" | 107 | result = "NPU ERROR" |
| 99 | fulfill_percent = 0 | 108 | fulfill_percent = 0 |
| 100 | 109 | ||
| 101 | - utils.save_result(test_data['params'], result, fulfill_percent, Path(result_path)) | 110 | + utils.save_result(test_data["params"], result, fulfill_percent, Path(result_path)) |
| 102 | 111 | ||
| 103 | - if result == "Failed": | ||
| 104 | - pytest.fail(f"用例精度失败:{test_data['Testcase_Name']} 精度:{fulfill_percent:.2f}%") | ||
| 105 | if result == "NPU ERROR": | 112 | if result == "NPU ERROR": |
| 106 | - pytest.fail(f"用例执行失败:{test_data['Testcase_Name']} NPU ERROR: {npu_error_msg}") | 113 | + pytest.fail( |
| 114 | + f"用例执行失败:{test_data['Testcase_Name']} NPU ERROR: {npu_error_msg}" | ||
| 115 | + ) | ||
| 116 | + elif result != "Pass": | ||
| 117 | + pytest.fail( | ||
| 118 | + f"用例精度失败:{test_data['Testcase_Name']} 精度:{fulfill_percent:.2f}%" | ||
| 119 | + ) | ||
| 120 | + | ||
| 107 | 121 | ||
| 108 | testcase_ids = [os.path.splitext(os.path.basename(f))[0] for f in testcase_files] | 122 | testcase_ids = [os.path.splitext(os.path.basename(f))[0] for f in testcase_files] |
| 109 | 123 | ||
| 124 | + | ||
| 110 | 125 | ||
| 111 | 126 | ||
| 112 | def test_mixed_quant_sparse_flash_mla(testcase_files): | 127 | def test_mixed_quant_sparse_flash_mla(testcase_files): |
| @@ -116,6 +131,6 @@ def test_mixed_quant_sparse_flash_mla(testcase_files): | |||
| 116 | # 等待并获取结果 | 131 | # 等待并获取结果 |
| 117 | for future in concurrent.futures.as_completed([futures]): | 132 | for future in concurrent.futures.as_completed([futures]): |
| 118 | try: | 133 | try: |
| 119 | - result = future.result() | 134 | + future.result() |
| 120 | - except Exception as e: | 135 | + except Exception: |
| 121 | - pytest.fail(f"当前用例线程执行失败") | 136 | + pytest.fail("当前用例线程执行失败") |
Mattention/mixed_quant_sparse_flash_mla/tests/pytest/test_mixed_quant_sparse_flash_mla_batch_graph.py+42-25
| @@ -10,38 +10,41 @@ | |||
| 10 | # See LICENSE in the root of the software repository for the full text of the License. | 10 | # See LICENSE in the root of the software repository for the full text of the License. |
| 11 | # ----------------------------------------------------------------------------------------------------------- | 11 | # ----------------------------------------------------------------------------------------------------------- |
| 12 | 12 | ||
| 13 | -import itertools | ||
| 14 | import torch | 13 | import torch |
| 15 | import pytest | 14 | import pytest |
| 16 | -import random | ||
| 17 | import pandas as pd | 15 | import pandas as pd |
| 18 | from pathlib import Path | 16 | from pathlib import Path |
| 19 | -import numpy as np | ||
| 20 | -import math | ||
| 21 | import os | 17 | import os |
| 22 | -import multiprocessing as mp | ||
| 23 | import concurrent.futures | 18 | import concurrent.futures |
| 24 | import result_compare_method | 19 | import result_compare_method |
| 25 | -import check_valid_param | ||
| 26 | from batch import mixed_quant_sparse_flash_mla_process | 20 | from batch import mixed_quant_sparse_flash_mla_process |
| 27 | import utils | 21 | import utils |
| 28 | 22 | ||
| 29 | testcase_path = os.environ.get("MQSMLA_PT_DIR", "mqsmla_testcase") | 23 | testcase_path = os.environ.get("MQSMLA_PT_DIR", "mqsmla_testcase") |
| 30 | -batch_test_mode = int(os.environ.get("MQSMLA_BATCH_TEST_MODE", 0)) # 0:路径下全量批跑，1:按表格中case批跑 | 24 | +batch_test_mode = int( |
| 31 | -excel_path = os.environ.get("MQSMLA_EXCEL_PATH", os.path.join(os.path.dirname(__file__), "excel", "testcase.xlsx")) | 25 | + os.environ.get("MQSMLA_BATCH_TEST_MODE", 0) |
| 32 | -result_path = os.environ.get("MQSMLA_RESULT_SAVE_PATH", './mqsmla_result.xlsx') | 26 | +) # 0:路径下全量批跑，1:按表格中case批跑 |
| 27 | +excel_path = os.environ.get( | ||
| 28 | + "MQSMLA_EXCEL_PATH", | ||
| 29 | + os.path.join(os.path.dirname(__file__), "excel", "testcase.xlsx"), | ||
| 30 | +) | ||
| 31 | +result_path = os.environ.get("MQSMLA_RESULT_SAVE_PATH", "./mqsmla_result.xlsx") | ||
| 33 | device_id = int(os.environ.get("MQSMLA_DEVICE_ID", 0)) | 32 | device_id = int(os.environ.get("MQSMLA_DEVICE_ID", 0)) |
| 34 | 33 | ||
| 35 | locals()["testcase_files"] = [] | 34 | locals()["testcase_files"] = [] |
| 36 | if os.path.isdir(testcase_path): | 35 | if os.path.isdir(testcase_path): |
| 37 | if batch_test_mode == 1: | 36 | if batch_test_mode == 1: |
| 38 | df = pd.read_excel(excel_path) | 37 | df = pd.read_excel(excel_path) |
| 39 | - target_names = [str(name) for name in df['Testcase_Name'].dropna().tolist() if str(name) != 'None'] | 38 | + target_names = [ |
| 39 | + str(name) | ||
| 40 | + for name in df["Testcase_Name"].dropna().tolist() | ||
| 41 | + if str(name) != "None" | ||
| 42 | + ] | ||
| 40 | if not target_names: | 43 | if not target_names: |
| 41 | print(f"错误: 表格中没有有效的Testcase_Name: {excel_path}") | 44 | print(f"错误: 表格中没有有效的Testcase_Name: {excel_path}") |
| 42 | else: | 45 | else: |
| 43 | print(f"从表格中读取到 {len(target_names)} 个目标用例名") | 46 | print(f"从表格中读取到 {len(target_names)} 个目标用例名") |
| 44 | - pt_files = [f for f in os.listdir(testcase_path) if f.endswith('.pt')] | 47 | + pt_files = [f for f in os.listdir(testcase_path) if f.endswith(".pt")] |
| 45 | for target_name in target_names: | 48 | for target_name in target_names: |
| 46 | matched = [f for f in pt_files if target_name in f] | 49 | matched = [f for f in pt_files if target_name in f] |
| 47 | if matched: | 50 | if matched: |
| @@ -53,7 +56,7 @@ if os.path.isdir(testcase_path): | |||
| 53 | print(f"警告: 用例名 '{target_name}' 未匹配到任何.pt文件") | 56 | print(f"警告: 用例名 '{target_name}' 未匹配到任何.pt文件") |
| 54 | print(f"按表格筛选后共 {len(locals()['testcase_files'])} 个测试用例文件") | 57 | print(f"按表格筛选后共 {len(locals()['testcase_files'])} 个测试用例文件") |
| 55 | else: | 58 | else: |
| 56 | - pt_files = [f for f in os.listdir(testcase_path) if f.endswith('.pt')] | 59 | + pt_files = [f for f in os.listdir(testcase_path) if f.endswith(".pt")] |
| 57 | if not pt_files: | 60 | if not pt_files: |
| 58 | print(f"错误: 目录中没有找到.pt文件: {testcase_path}") | 61 | print(f"错误: 目录中没有找到.pt文件: {testcase_path}") |
| 59 | else: | 62 | else: |
| @@ -64,14 +67,20 @@ if os.path.isdir(testcase_path): | |||
| 64 | else: | 67 | else: |
| 65 | print(f"错误: 输出目录不存在: {testcase_path}") | 68 | print(f"错误: 输出目录不存在: {testcase_path}") |
| 66 | 69 | ||
| 70 | + | ||
| 67 | def mqsmla_aclgraph(testcase_files): | 71 | def mqsmla_aclgraph(testcase_files): |
| 68 | test_data = torch.load(testcase_files, map_location="cpu") | 72 | test_data = torch.load(testcase_files, map_location="cpu") |
| 69 | npu_error_msg = None | 73 | npu_error_msg = None |
| 70 | try: | 74 | try: |
| 71 | - npu_result, cpu_quant_result, npu_lse, cpu_lse = mixed_quant_sparse_flash_mla_process.test_mqsmla_quant_process_graph( | 75 | + npu_result, cpu_quant_result, npu_lse, cpu_lse = ( |
| 72 | - test_data, device_id=device_id) | 76 | + mixed_quant_sparse_flash_mla_process.test_mqsmla_quant_process_graph( |
| 73 | - | 77 | + test_data, device_id=device_id |
| 74 | - attn_result, attn_pct = result_compare_method.check_result(cpu_quant_result, npu_result) | 78 | + ) |
| 79 | + ) | ||
| 80 | + | ||
| 81 | + attn_result, attn_pct = result_compare_method.check_result( | ||
| 82 | + cpu_quant_result, npu_result | ||
| 83 | + ) | ||
| 75 | lse_result, lse_pct = None, 0.0 | 84 | lse_result, lse_pct = None, 0.0 |
| 76 | fail_info = [] | 85 | fail_info = [] |
| 77 | min_fulfill = attn_pct | 86 | min_fulfill = attn_pct |
| @@ -79,7 +88,7 @@ def mqsmla_aclgraph(testcase_files): | |||
| 79 | if attn_result != "Pass": | 88 | if attn_result != "Pass": |
| 80 | fail_info.append(f"MAIN_FAILED:{attn_result}") | 89 | fail_info.append(f"MAIN_FAILED:{attn_result}") |
| 81 | 90 | ||
| 82 | - if test_data['params'].get('return_softmax_lse'): | 91 | + if test_data["params"].get("return_softmax_lse"): |
| 83 | print("return_softmax_lse is true!!!") | 92 | print("return_softmax_lse is true!!!") |
| 84 | lse_result, lse_pct = result_compare_method.check_result(cpu_lse, npu_lse) | 93 | lse_result, lse_pct = result_compare_method.check_result(cpu_lse, npu_lse) |
| 85 | min_fulfill = min(min_fulfill, lse_pct) | 94 | min_fulfill = min(min_fulfill, lse_pct) |
| @@ -95,14 +104,22 @@ def mqsmla_aclgraph(testcase_files): | |||
| 95 | result = "NPU ERROR" | 104 | result = "NPU ERROR" |
| 96 | fulfill_percent = 0 | 105 | fulfill_percent = 0 |
| 97 | 106 | ||
| 98 | - utils.save_result(test_data['params'], result, fulfill_percent, Path(result_path)) | 107 | + utils.save_result(test_data["params"], result, fulfill_percent, Path(result_path)) |
| 99 | 108 | ||
| 100 | - if result == "Failed": | ||
| 101 | - pytest.fail(f"用例精度失败:{test_data['Testcase_Name']} 精度:{fulfill_percent:.2f}%") | ||
| 102 | if result == "NPU ERROR": | 109 | if result == "NPU ERROR": |
| 103 | - pytest.fail(f"用例执行失败:{test_data['Testcase_Name']} NPU ERROR: {npu_error_msg}") | 110 | + pytest.fail( |
| 111 | + f"用例执行失败:{test_data['Testcase_Name']} NPU ERROR: {npu_error_msg}" | ||
| 112 | + ) | ||
| 113 | + elif result != "Pass": | ||
| 114 | + pytest.fail( | ||
| 115 | + f"用例精度失败:{test_data['Testcase_Name']} 精度:{fulfill_percent:.2f}%" | ||
| 116 | + ) | ||
| 117 | + | ||
| 118 | + | ||
| 119 | +testcase_ids = [ | ||
| 120 | + os.path.splitext(os.path.basename(f))[0] for f in locals()["testcase_files"] | ||
| 121 | +] | ||
| 104 | 122 | ||
| 105 | -testcase_ids = [os.path.splitext(os.path.basename(f))[0] for f in locals()["testcase_files"]] | ||
| 106 | 123 | ||
| 107 | 124 | ||
| 108 | 125 | ||
| @@ -113,6 +130,6 @@ def test_mixed_quant_sparse_flash_mla(testcase_files): | |||
| 113 | # 等待并获取结果 | 130 | # 等待并获取结果 |
| 114 | for future in concurrent.futures.as_completed([futures]): | 131 | for future in concurrent.futures.as_completed([futures]): |
| 115 | try: | 132 | try: |
| 116 | - result = future.result() | 133 | + future.result() |
| 117 | - except Exception as e: | 134 | + except Exception: |
| 118 | - pytest.fail(f"当前用例线程执行失败") | 135 | + pytest.fail("当前用例线程执行失败") |
Mattention/mixed_quant_sparse_flash_mla/tests/pytest/test_mixed_quant_sparse_flash_mla_single.py+4-4
| @@ -163,14 +163,14 @@ def mqsmla(param_combinations): | |||
| 163 | 163 | ||
| 164 | utils.save_result(test_data["params"], result, fulfill_percent, result_path) | 164 | utils.save_result(test_data["params"], result, fulfill_percent, result_path) |
| 165 | 165 | ||
| 166 | - if result == "Failed": | ||
| 167 | - pytest.fail( | ||
| 168 | - f"用例精度失败:{test_data['Testcase_Name']} 精度:{fulfill_percent:.2f}%" | ||
| 169 | - ) | ||
| 170 | if result == "NPU ERROR": | 166 | if result == "NPU ERROR": |
| 171 | pytest.fail( | 167 | pytest.fail( |
| 172 | f"用例执行失败:{test_data['Testcase_Name']} NPU ERROR: {npu_error_msg}" | 168 | f"用例执行失败:{test_data['Testcase_Name']} NPU ERROR: {npu_error_msg}" |
| 173 | ) | 169 | ) |
| 170 | + elif result != "Pass": | ||
| 171 | + pytest.fail( | ||
| 172 | + f"用例精度失败:{test_data['Testcase_Name']} 精度:{fulfill_percent:.2f}%" | ||
| 173 | + ) | ||
| 174 | 174 | ||
| 175 | 175 | ||
| 176 | def _gen_testcase_id(params, idx): | 176 | def _gen_testcase_id(params, idx): |
| @@ -161,12 +161,6 @@ class GeneralizedSFAQuant: | |||
| 161 | logging.info( | 161 | logging.info( |
| 162 | f" 进度：{current_pct:.1f}% | 步数：{i_S1:>{len(str(cur_act_q))}}/{cur_act_q}" | 162 | f" 进度：{current_pct:.1f}% | 步数：{i_S1:>{len(str(cur_act_q))}}/{cur_act_q}" |
| 163 | ) | 163 | ) |
| 164 | - if self.ori_mask_mode == 3 and i_S1 < cur_act_q - cur_ori_act_kv: | ||
| 165 | - attn_out[i_B, i_N2 * G : (i_N2 + 1) * G, i_S1, :] = torch.zeros( | ||
| 166 | - [G, self.D], dtype=torch.float | ||
| 167 | - ) | ||
| 168 | - continue | ||
| 169 | - | ||
| 170 | if self.ori_mask_mode == 0: | 164 | if self.ori_mask_mode == 0: |
| 171 | ori_win_start = 0 | 165 | ori_win_start = 0 |
| 172 | ori_win_end = cur_ori_act_kv | 166 | ori_win_end = cur_ori_act_kv |
| @@ -192,7 +186,14 @@ class GeneralizedSFAQuant: | |||
| 192 | cur_ori_act_kv, | 186 | cur_ori_act_kv, |
| 193 | ) | 187 | ) |
| 194 | 188 | ||
| 195 | - cur_ori_k_bnsd = ori_k_bnsd[i_B, i_N2, ori_win_start:ori_win_end, :] | 189 | + if ori_win_start >= ori_win_end: |
| 190 | + cur_ori_k_bnsd = torch.zeros( | ||
| 191 | + [0, self.D], dtype=ori_k_bnsd.dtype | ||
| 192 | + ) | ||
| 193 | + else: | ||
| 194 | + cur_ori_k_bnsd = ori_k_bnsd[ | ||
| 195 | + i_B, i_N2, ori_win_start:ori_win_end, : | ||
| 196 | + ] | ||
| 196 | 197 | ||
| 197 | if ( | 198 | if ( |
| 198 | self.template_run_mode == "CSA" | 199 | self.template_run_mode == "CSA" |
| @@ -274,6 +275,12 @@ class GeneralizedSFAQuant: | |||
| 274 | cmp_s2_loop_time = math.ceil(cur_cmp_k.size(0) / s2_base_size) | 275 | cmp_s2_loop_time = math.ceil(cur_cmp_k.size(0) / s2_base_size) |
| 275 | cur_cmp_k_fp32 = cur_cmp_k.to(dtype=torch.float32) | 276 | cur_cmp_k_fp32 = cur_cmp_k.to(dtype=torch.float32) |
| 276 | 277 | ||
| 278 | + if cur_ori_k_bnsd.size(0) == 0 and cur_cmp_k == []: | ||
| 279 | + attn_out[i_B, i_N2 * G : (i_N2 + 1) * G, i_S1, :] = torch.zeros( | ||
| 280 | + [G, self.D], dtype=torch.float | ||
| 281 | + ) | ||
| 282 | + continue | ||
| 283 | + | ||
| 277 | q_curr = q_bnsd[i_B, i_N2 * G : (i_N2 + 1) * G, i_S1, :] | 284 | q_curr = q_bnsd[i_B, i_N2 * G : (i_N2 + 1) * G, i_S1, :] |
| 278 | q_curr_fp32 = q_curr.to(dtype=torch.float32) | 285 | q_curr_fp32 = q_curr.to(dtype=torch.float32) |
| 279 | ori_s2_loop_time = math.ceil(cur_ori_k_bnsd.size(0) / s2_base_size) | 286 | ori_s2_loop_time = math.ceil(cur_ori_k_bnsd.size(0) / s2_base_size) |
| @@ -97,7 +97,7 @@ def qsmla(testcase_files): | |||
| 97 | 97 | ||
| 98 | utils.save_result(test_data["params"], result, fulfill_percent, Path(result_path)) | 98 | utils.save_result(test_data["params"], result, fulfill_percent, Path(result_path)) |
| 99 | 99 | ||
| 100 | - if result in ("NPU ERROR", "Failed"): | 100 | + if result != "PASS": |
| 101 | pytest.fail( | 101 | pytest.fail( |
| 102 | f"用例执行失败:{test_data['Testcase_Name']} 精度:{fulfill_percent:.2f}%" | 102 | f"用例执行失败:{test_data['Testcase_Name']} 精度:{fulfill_percent:.2f}%" |
| 103 | ) | 103 | ) |
| @@ -176,7 +176,7 @@ def qsmla(param_combinations): | |||
| 176 | 176 | ||
| 177 | utils.save_result(test_data["params"], result, fulfill_percent, result_path) | 177 | utils.save_result(test_data["params"], result, fulfill_percent, result_path) |
| 178 | 178 | ||
| 179 | - if result in ("NPU ERROR", "Failed"): | 179 | + if result != "PASS": |
| 180 | pytest.fail( | 180 | pytest.fail( |
| 181 | f"用例执行失败:{test_data['Testcase_Name']} 精度:{fulfill_percent:.2f}%" | 181 | f"用例执行失败:{test_data['Testcase_Name']} 精度:{fulfill_percent:.2f}%" |
| 182 | ) | 182 | ) |
| @@ -25,6 +25,15 @@ DATA_RANGE_RIGHT = 10 | |||
| 25 | 25 | ||
| 26 | RUN_MODE = 1 | 26 | RUN_MODE = 1 |
| 27 | 27 | ||
| 28 | + | ||
| 29 | +def is_empty(obj): | ||
| 30 | + if isinstance(obj, list): | ||
| 31 | + return len(obj) == 0 | ||
| 32 | + if isinstance(obj, torch.Tensor): | ||
| 33 | + return obj.nelement() == 0 | ||
| 34 | + return obj is None | ||
| 35 | + | ||
| 36 | + | ||
| 28 | # np.random.seed(42) | 37 | # np.random.seed(42) |
| 29 | # torch.manual_seed(42) | 38 | # torch.manual_seed(42) |
| 30 | 39 | ||
| @@ -118,24 +127,6 @@ class GeneralizedSFA: | |||
| 118 | attn_out = torch.zeros(q_bnsd.shape, dtype=q_bnsd.dtype) | 127 | attn_out = torch.zeros(q_bnsd.shape, dtype=q_bnsd.dtype) |
| 119 | softmax_lse = None | 128 | softmax_lse = None |
| 120 | if return_softmax_lse: | 129 | if return_softmax_lse: |
| 121 | - softmax_max = torch.zeros( | ||
| 122 | - ( | ||
| 123 | - q_bnsd.shape[0], | ||
| 124 | - ori_k_bnsd.shape[1], | ||
| 125 | - q_bnsd.shape[2], | ||
| 126 | - q_bnsd.shape[1] // ori_k_bnsd.shape[1], | ||
| 127 | - ), | ||
| 128 | - dtype=torch.float32, | ||
| 129 | - ) | ||
| 130 | - softmax_sum = torch.zeros( | ||
| 131 | - ( | ||
| 132 | - q_bnsd.shape[0], | ||
| 133 | - ori_k_bnsd.shape[1], | ||
| 134 | - q_bnsd.shape[2], | ||
| 135 | - q_bnsd.shape[1] // ori_k_bnsd.shape[1], | ||
| 136 | - ), | ||
| 137 | - dtype=torch.float32, | ||
| 138 | - ) | ||
| 139 | softmax_lse = torch.zeros( | 130 | softmax_lse = torch.zeros( |
| 140 | ( | 131 | ( |
| 141 | q_bnsd.shape[0], | 132 | q_bnsd.shape[0], |
| @@ -184,14 +175,6 @@ class GeneralizedSFA: | |||
| 184 | print( | 175 | print( |
| 185 | f" 进度：{current_pct:.1f}% | 步数：{i_S1:>{len(str(cur_act_q))}}/{cur_act_q}" | 176 | f" 进度：{current_pct:.1f}% | 步数：{i_S1:>{len(str(cur_act_q))}}/{cur_act_q}" |
| 186 | ) | 177 | ) |
| 187 | - if ( | ||
| 188 | - self.ori_mask_mode == 3 | ||
| 189 | - ) and i_S1 < cur_act_q - cur_ori_act_kv: # 根据 ori_kv 判断行无效 | ||
| 190 | - attn_out[i_B, i_N2 * G : (i_N2 + 1) * G, i_S1, :] = torch.zeros( | ||
| 191 | - [G, self.D], dtype=torch.float | ||
| 192 | - ) | ||
| 193 | - continue | ||
| 194 | - | ||
| 195 | q_curr = q_bnsd[i_B, i_N2 * G : (i_N2 + 1) * G, i_S1, :] | 178 | q_curr = q_bnsd[i_B, i_N2 * G : (i_N2 + 1) * G, i_S1, :] |
| 196 | q_curr_fp32 = q_curr.to(dtype=torch.float32) | 179 | q_curr_fp32 = q_curr.to(dtype=torch.float32) |
| 197 | 180 | ||
| @@ -234,7 +217,7 @@ class GeneralizedSFA: | |||
| 234 | self.cmp_ratio, | 217 | self.cmp_ratio, |
| 235 | cur_cmp_act_kv, | 218 | cur_cmp_act_kv, |
| 236 | ) | 219 | ) |
| 237 | - if cur_cmp_k == []: | 220 | + if is_empty(cur_cmp_k): |
| 238 | cmp_s2_loop_time = 0 | 221 | cmp_s2_loop_time = 0 |
| 239 | cur_cmp_k_fp32 = [] | 222 | cur_cmp_k_fp32 = [] |
| 240 | else: | 223 | else: |
| @@ -265,7 +248,7 @@ class GeneralizedSFA: | |||
| 265 | 1, | 248 | 1, |
| 266 | cur_ori_act_kv, | 249 | cur_ori_act_kv, |
| 267 | ) | 250 | ) |
| 268 | - if cur_ori_k_bnsd == []: | 251 | + if is_empty(cur_ori_k_bnsd): |
| 269 | ori_s2_loop_time = 0 | 252 | ori_s2_loop_time = 0 |
| 270 | else: | 253 | else: |
| 271 | ori_s2_loop_time = math.ceil( | 254 | ori_s2_loop_time = math.ceil( |
| @@ -289,23 +272,42 @@ class GeneralizedSFA: | |||
| 289 | ) | 272 | ) |
| 290 | elif self.ori_mask_mode == 3: | 273 | elif self.ori_mask_mode == 3: |
| 291 | ori_threshold = cur_ori_act_kv - cur_act_q + i_S1 + 1 | 274 | ori_threshold = cur_ori_act_kv - cur_act_q + i_S1 + 1 |
| 292 | - ori_win_end = ori_threshold | 275 | + ori_win_end = max(ori_threshold, 0) |
| 293 | ori_win_start = 0 | 276 | ori_win_start = 0 |
| 294 | elif self.ori_mask_mode == 0: | 277 | elif self.ori_mask_mode == 0: |
| 295 | ori_win_start = 0 | 278 | ori_win_start = 0 |
| 296 | ori_win_end = cur_ori_act_kv | 279 | ori_win_end = cur_ori_act_kv |
| 297 | 280 | ||
| 298 | - cur_ori_k_bnsd = ori_k_bnsd[ | 281 | + if ori_win_start >= ori_win_end: |
| 299 | - i_B, i_N2, ori_win_start:ori_win_end, : | 282 | + cur_ori_k_bnsd = [] |
| 300 | - ] | 283 | + empty_flag_ori = True |
| 301 | - empty_flag_ori = False | 284 | + else: |
| 285 | + cur_ori_k_bnsd = ori_k_bnsd[ | ||
| 286 | + i_B, i_N2, ori_win_start:ori_win_end, : | ||
| 287 | + ] | ||
| 288 | + empty_flag_ori = False | ||
| 289 | + | ||
| 290 | + if is_empty(cur_ori_k_bnsd) and is_empty(cur_cmp_k): | ||
| 291 | + attn_out[i_B, i_N2 * G : (i_N2 + 1) * G, i_S1, :] = torch.zeros( | ||
| 292 | + [G, self.D], dtype=torch.float | ||
| 293 | + ) | ||
| 294 | + continue | ||
| 302 | 295 | ||
| 303 | cur_attn_out = attn_out[i_B, i_N2 * G : (i_N2 + 1) * G, i_S1, :] | 296 | cur_attn_out = attn_out[i_B, i_N2 * G : (i_N2 + 1) * G, i_S1, :] |
| 304 | if RUN_MODE == 0: | 297 | if RUN_MODE == 0: |
| 305 | if empty_flag: | 298 | if empty_flag: |
| 306 | - k_concat = cur_ori_k_bnsd | 299 | + k_concat = ( |
| 300 | + cur_ori_k_bnsd | ||
| 301 | + if not is_empty(cur_ori_k_bnsd) | ||
| 302 | + else cur_cmp_k | ||
| 303 | + ) | ||
| 307 | else: | 304 | else: |
| 308 | - k_concat = torch.concat([cur_ori_k_bnsd, cur_cmp_k], dim=0) | 305 | + if is_empty(cur_ori_k_bnsd): |
| 306 | + k_concat = cur_cmp_k | ||
| 307 | + else: | ||
| 308 | + k_concat = torch.concat( | ||
| 309 | + [cur_ori_k_bnsd, cur_cmp_k], dim=0 | ||
| 310 | + ) | ||
| 309 | 311 | ||
| 310 | k_concat_fp32 = k_concat.to(dtype=torch.float32) | 312 | k_concat_fp32 = k_concat.to(dtype=torch.float32) |
| 311 | v_concat_fp32 = k_concat_fp32.clone() | 313 | v_concat_fp32 = k_concat_fp32.clone() |
| @@ -323,15 +325,18 @@ class GeneralizedSFA: | |||
| 323 | dtype=q_bnsd.dtype | 325 | dtype=q_bnsd.dtype |
| 324 | ) | 326 | ) |
| 325 | if return_softmax_lse: | 327 | if return_softmax_lse: |
| 326 | - softmax_max[batch, n2Idx, s1Idx, :] = x_max[:, 0] | 328 | + softmax_lse[i_B, i_N2, i_S1, :] = x_max[:, 0] + torch.log( |
| 327 | - softmax_sum[batch, n2Idx, s1Idx, :] = x_sum[:, 0] | 329 | + x_sum[:, 0] + 1e-10 |
| 330 | + ) | ||
| 328 | elif RUN_MODE == 1: | 331 | elif RUN_MODE == 1: |
| 329 | if empty_flag_ori: | 332 | if empty_flag_ori: |
| 330 | ori_s2_loop_time = 0 | 333 | ori_s2_loop_time = 0 |
| 331 | - total_s2_loop_time = 0 | 334 | + total_s2_loop_time = cmp_s2_loop_time |
| 332 | row_sum = torch.empty((G), dtype=torch.float32).uniform_( | 335 | row_sum = torch.empty((G), dtype=torch.float32).uniform_( |
| 333 | 1.0, 1.0 | 336 | 1.0, 1.0 |
| 334 | ) | 337 | ) |
| 338 | + row_max = torch.empty((G, 1), dtype=torch.float32) | ||
| 339 | + row_max = cur_sinks | ||
| 335 | else: | 340 | else: |
| 336 | ori_s2_loop_time = math.ceil( | 341 | ori_s2_loop_time = math.ceil( |
| 337 | cur_ori_k_bnsd.size(0) / s2_base_size | 342 | cur_ori_k_bnsd.size(0) / s2_base_size |
| @@ -502,6 +507,10 @@ class GeneralizedSFA: | |||
| 502 | attn_out[i_B, i_N2 * G : (i_N2 + 1) * G, i_S1, :] = ( | 507 | attn_out[i_B, i_N2 * G : (i_N2 + 1) * G, i_S1, :] = ( |
| 503 | O_flash / row_sum_expand / tmp_scale_16_expand | 508 | O_flash / row_sum_expand / tmp_scale_16_expand |
| 504 | ).to(dtype=q_bnsd.dtype) | 509 | ).to(dtype=q_bnsd.dtype) |
| 510 | + if return_softmax_lse: | ||
| 511 | + softmax_lse[i_B, i_N2, i_S1, :] = row_max + torch.log( | ||
| 512 | + row_sum + 1e-10 | ||
| 513 | + ) | ||
| 505 | else: | 514 | else: |
| 506 | raise ValueError(f"unsupported RUN_MODE:{RUN_MODE}") | 515 | raise ValueError(f"unsupported RUN_MODE:{RUN_MODE}") |
| 507 | return attn_out, softmax_lse | 516 | return attn_out, softmax_lse |
| @@ -556,7 +565,7 @@ class GeneralizedSFA: | |||
| 556 | 565 | ||
| 557 | empty_flag = False | 566 | empty_flag = False |
| 558 | if len(s2_sparse) == 0: | 567 | if len(s2_sparse) == 0: |
| 559 | - cur_cmp_k = torch.tensor([]) | 568 | + cur_cmp_k = [] |
| 560 | empty_flag = True | 569 | empty_flag = True |
| 561 | else: | 570 | else: |
| 562 | cur_cmp_k = k_tensor[i_B, i_N2, s2_sparse, :] | 571 | cur_cmp_k = k_tensor[i_B, i_N2, s2_sparse, :] |
| @@ -11,7 +11,6 @@ | |||
| 11 | # ----------------------------------------------------------------------------------------------------------- | 11 | # ----------------------------------------------------------------------------------------------------------- |
| 12 | 12 | ||
| 13 | import torch | 13 | import torch |
| 14 | -import torch_npu | ||
| 15 | import result_compare_method | 14 | import result_compare_method |
| 16 | import utils | 15 | import utils |
| 17 | from batch import sparse_flash_mla_process | 16 | from batch import sparse_flash_mla_process |
| @@ -22,9 +21,11 @@ from pathlib import Path | |||
| 22 | import os | 21 | import os |
| 23 | 22 | ||
| 24 | pt_dir = os.getenv("SMLA_PT_LOAD_PATH", "./data") | 23 | pt_dir = os.getenv("SMLA_PT_LOAD_PATH", "./data") |
| 25 | -result_path = Path(os.getenv("SMLA_RESULT_SAVE_PATH", './result/smla_result.xlsx')) | 24 | +result_path = Path(os.getenv("SMLA_RESULT_SAVE_PATH", "./result/smla_result.xlsx")) |
| 26 | batch_test_mode = int(os.environ.get("SMLA_BATCH_TEST_MODE", 0)) | 25 | batch_test_mode = int(os.environ.get("SMLA_BATCH_TEST_MODE", 0)) |
| 27 | -excel_path = os.environ.get("SMLA_EXCEL_PATH", os.path.join(os.path.dirname(__file__), "excel", "example.xlsx")) | 26 | +excel_path = os.environ.get( |
| 27 | + "SMLA_EXCEL_PATH", os.path.join(os.path.dirname(__file__), "excel", "example.xlsx") | ||
| 28 | +) | ||
| 28 | excel_sheet = os.environ.get("SMLA_EXCEL_SHEET", "CSA") | 29 | excel_sheet = os.environ.get("SMLA_EXCEL_SHEET", "CSA") |
| 29 | device_id = int(os.environ.get("SMLA_DEVICE_ID", 0)) | 30 | device_id = int(os.environ.get("SMLA_DEVICE_ID", 0)) |
| 30 | 31 | ||
| @@ -33,19 +34,27 @@ _single_case_path = os.environ.get("QSAS_TESTCASE_PATH", "").strip() | |||
| 33 | locals()["testcase_files"] = [] | 34 | locals()["testcase_files"] = [] |
| 34 | if _single_case_path: | 35 | if _single_case_path: |
| 35 | if not os.path.isfile(_single_case_path): | 36 | if not os.path.isfile(_single_case_path): |
| 36 | - print(f"错误: 环境变量 QSAS_TESTCASE_PATH 指定的用例文件不存在: {_single_case_path}") | 37 | + print( |
| 38 | + f"错误: 环境变量 QSAS_TESTCASE_PATH 指定的用例文件不存在: {_single_case_path}" | ||
| 39 | + ) | ||
| 37 | else: | 40 | else: |
| 38 | print(f"单用例隔离模式, 仅执行: {_single_case_path}") | 41 | print(f"单用例隔离模式, 仅执行: {_single_case_path}") |
| 39 | locals()["testcase_files"].append(_single_case_path) | 42 | locals()["testcase_files"].append(_single_case_path) |
| 40 | elif os.path.isdir(pt_dir): | 43 | elif os.path.isdir(pt_dir): |
| 41 | - pt_files = [f for f in os.listdir(pt_dir) if f.endswith('.pt')] | 44 | + pt_files = [f for f in os.listdir(pt_dir) if f.endswith(".pt")] |
| 42 | if not pt_files: | 45 | if not pt_files: |
| 43 | print(f"错误: 目录中没有找到.pt文件: {pt_dir}") | 46 | print(f"错误: 目录中没有找到.pt文件: {pt_dir}") |
| 44 | elif batch_test_mode == 1: | 47 | elif batch_test_mode == 1: |
| 45 | df = pd.read_excel(excel_path, sheet_name=excel_sheet) | 48 | df = pd.read_excel(excel_path, sheet_name=excel_sheet) |
| 46 | - target_names = [str(name) for name in df['testcase_name'].dropna().tolist() if str(name) != 'None'] | 49 | + target_names = [ |
| 50 | + str(name) | ||
| 51 | + for name in df["testcase_name"].dropna().tolist() | ||
| 52 | + if str(name) != "None" | ||
| 53 | + ] | ||
| 47 | if not target_names: | 54 | if not target_names: |
| 48 | - print(f"错误: 表格中没有有效的testcase_name: {excel_path} sheet: {excel_sheet}") | 55 | + print( |
| 56 | + f"错误: 表格中没有有效的testcase_name: {excel_path} sheet: {excel_sheet}" | ||
| 57 | + ) | ||
| 49 | else: | 58 | else: |
| 50 | print(f"从表格[{excel_sheet}]中读取到 {len(target_names)} 个目标用例名") | 59 | print(f"从表格[{excel_sheet}]中读取到 {len(target_names)} 个目标用例名") |
| 51 | for target_name in target_names: | 60 | for target_name in target_names: |
| @@ -68,18 +77,25 @@ else: | |||
| 68 | 77 | ||
| 69 | print("files:", locals()["testcase_files"]) | 78 | print("files:", locals()["testcase_files"]) |
| 70 | 79 | ||
| 80 | + | ||
| 71 | def smla_graph(testcase_files): | 81 | def smla_graph(testcase_files): |
| 72 | test_data = torch.load(testcase_files, map_location="cpu") | 82 | test_data = torch.load(testcase_files, map_location="cpu") |
| 73 | npu_error_msg = None | 83 | npu_error_msg = None |
| 74 | try: | 84 | try: |
| 75 | - npu_result, softmax_lse = sparse_flash_mla_process.call_npu_graph(test_data, device_id=device_id) | 85 | + npu_result, softmax_lse = sparse_flash_mla_process.call_npu_graph( |
| 76 | - attn_result, attn_percent = result_compare_method.check_result(test_data['cpu_output'], npu_result) | 86 | + test_data, device_id=device_id |
| 87 | + ) | ||
| 88 | + attn_result, attn_percent = result_compare_method.check_result( | ||
| 89 | + test_data["cpu_output"], npu_result | ||
| 90 | + ) | ||
| 77 | lse_result, lse_percent = None, None | 91 | lse_result, lse_percent = None, None |
| 78 | fail_info = [] | 92 | fail_info = [] |
| 79 | min_fulfill = attn_percent | 93 | min_fulfill = attn_percent |
| 80 | - if test_data['params'].get('return_softmax_lse'): | 94 | + if test_data["params"].get("return_softmax_lse"): |
| 81 | print("return_softmax_lse is true!!!") | 95 | print("return_softmax_lse is true!!!") |
| 82 | - lse_result, lse_percent = result_compare_method.check_result(test_data['softmax_lse'], softmax_lse) | 96 | + lse_result, lse_percent = result_compare_method.check_result( |
| 97 | + test_data["softmax_lse"], softmax_lse | ||
| 98 | + ) | ||
| 83 | min_fulfill = min(min_fulfill, lse_percent) | 99 | min_fulfill = min(min_fulfill, lse_percent) |
| 84 | 100 | ||
| 85 | if attn_result != "Pass": | 101 | if attn_result != "Pass": |
| @@ -99,14 +115,22 @@ def smla_graph(testcase_files): | |||
| 99 | result = "NPU ERROR" | 115 | result = "NPU ERROR" |
| 100 | fulfill_percent = 0 | 116 | fulfill_percent = 0 |
| 101 | 117 | ||
| 102 | - utils.save_result(result, fulfill_percent, test_data['params'], result_path) | 118 | + utils.save_result(result, fulfill_percent, test_data["params"], result_path) |
| 103 | 119 | ||
| 104 | - if result == "Failed": | ||
| 105 | - pytest.fail(f"用例精度失败:{os.path.basename(testcase_files)} 精度:{fulfill_percent:.2f}%") | ||
| 106 | if result == "NPU ERROR": | 120 | if result == "NPU ERROR": |
| 107 | - pytest.fail(f"用例执行失败:{os.path.basename(testcase_files)} NPU ERROR: {npu_error_msg}") | 121 | + pytest.fail( |
| 122 | + f"用例执行失败:{os.path.basename(testcase_files)} NPU ERROR: {npu_error_msg}" | ||
| 123 | + ) | ||
| 124 | + elif result != "Pass": | ||
| 125 | + pytest.fail( | ||
| 126 | + f"用例精度失败:{os.path.basename(testcase_files)} 精度:{fulfill_percent:.2f}%" | ||
| 127 | + ) | ||
| 128 | + | ||
| 129 | + | ||
| 130 | +testcase_ids = [ | ||
| 131 | + os.path.splitext(os.path.basename(f))[0] for f in locals()["testcase_files"] | ||
| 132 | +] | ||
| 108 | 133 | ||
| 109 | -testcase_ids = [os.path.splitext(os.path.basename(f))[0] for f in locals()["testcase_files"]] | ||
| 110 | 134 | ||
| 111 | 135 | ||
| 112 | 136 | ||
| @@ -115,6 +139,6 @@ def test_sparse_flash_mla(testcase_files): | |||
| 115 | futures = executor.submit(smla_graph, testcase_files) | 139 | futures = executor.submit(smla_graph, testcase_files) |
| 116 | for future in concurrent.futures.as_completed([futures]): | 140 | for future in concurrent.futures.as_completed([futures]): |
| 117 | try: | 141 | try: |
| 118 | - result = future.result() | 142 | + future.result() |
| 119 | - except Exception as e: | 143 | + except Exception: |
| 120 | - pytest.fail(f"当前用例线程执行失败") | 144 | + pytest.fail("当前用例线程执行失败") |
这里引号为啥要改