Star1065
2570
代码
代码
Issues187
Pull Requests441
Actions
项目讨论
Wiki
分析
Star1065
2570
已合并
smla/qsmla/mqsmla pytest bugfix #9664
郑文惠创建于 17 天前
smla/qsmla/mqsmla pytest bugfix #9664
已合并
郑文惠创建于 17 天前
从zhengwenhui0817/ops-transformer_2pytest合入到cann/ops-transformermaster
共 10 个文件变更+69-61
Mattention/mixed_quant_sparse_flash_mla/tests/pytest/mixed_quant_sparse_flash_mla_golden.py+4-0
@@ -862,6 +862,8 @@ def gen_sparse_indices_bsnd(
862 kv_topk_mode,862 kv_topk_mode,
863 topk_length_override=None,863 topk_length_override=None,
864):864):
865+ if mask_mode != 0:
866+ kv_topk_mode = "no" # mask_mode != 0时，kv_topk_mode只能为no
865 if sparse_indices_mode is None:867 if sparse_indices_mode is None:
866 sparse_indices_mode = "full"868 sparse_indices_mode = "full"
867 if sparse_indices_mode not in ["full", "random"]:869 if sparse_indices_mode not in ["full", "random"]:
@@ -946,6 +948,8 @@ def gen_sparse_indices_tnd(
946 kv_topk_mode,948 kv_topk_mode,
947 topk_length_override=None,949 topk_length_override=None,
948):950):
951+ if mask_mode != 0:
952+ kv_topk_mode = "no" # mask_mode != 0时，kv_topk_mode只能为no
949 if sparse_indices_mode is None:953 if sparse_indices_mode is None:
950 sparse_indices_mode = "full"954 sparse_indices_mode = "full"
951 if sparse_indices_mode not in ["full", "random"]:955 if sparse_indices_mode not in ["full", "random"]:
Mattention/mixed_quant_sparse_flash_mla/tests/pytest/test_mixed_quant_sparse_flash_mla_batch.py+1-1
@@ -65,7 +65,7 @@ else:
65 print(f"错误: 输出目录不存在: {testcase_path}")65 print(f"错误: 输出目录不存在: {testcase_path}")
66 66 
67def mqsmla(testcase_files):67def mqsmla(testcase_files):
68- test_data = torch.load(testcase_files, map_location="cpu")68+ test_data = torch.load(testcase_files, map_location="cpu", weights_only=False)
69 npu_error_msg = None69 npu_error_msg = None
70 try:70 try:
71 npu_result, cpu_quant_result, cpu_lse, npu_lse = mixed_quant_sparse_flash_mla_process.test_mqsmla_quant_process_ci(71 npu_result, cpu_quant_result, cpu_lse, npu_lse = mixed_quant_sparse_flash_mla_process.test_mqsmla_quant_process_ci(
Mattention/mixed_quant_sparse_flash_mla/tests/pytest/test_mixed_quant_sparse_flash_mla_single.py+2-2
@@ -28,8 +28,8 @@ import utils
28 28 
29pt_save_path = "mqsmla_testcase"29pt_save_path = "mqsmla_testcase"
30device_id = 030device_id = 0
31-save_pt = False31+save_pt = os.environ.get("SAVE_PT", "0") == "1"
32-result_path = Path("result.xlsx")32+result_path = Path(os.environ.get("MQSMLA_RESULT_SAVE_PATH", "result.xlsx"))
33 33 
34param_combinations = []34param_combinations = []
35for params in ENABLED_PARAMS:35for params in ENABLED_PARAMS:
Mattention/quant_sparse_flash_mla/tests/pytest/quant_sparse_flash_mla_golden.py+4-0
@@ -778,6 +778,8 @@ def gen_sparse_indices_bsnd(
778 kv_topk_mode,778 kv_topk_mode,
779 topk_length_override=None,779 topk_length_override=None,
780):780):
781+ if mask_mode != 0:
782+ kv_topk_mode = "no" # mask_mode != 0时，kv_topk_mode只能为no
781 if sparse_indices_mode is None:783 if sparse_indices_mode is None:
782 sparse_indices_mode = "full"784 sparse_indices_mode = "full"
783 if sparse_indices_mode not in ["full", "random"]:785 if sparse_indices_mode not in ["full", "random"]:
@@ -863,6 +865,8 @@ def gen_sparse_indices_tnd(
863 kv_topk_mode,865 kv_topk_mode,
864 topk_length_override=None,866 topk_length_override=None,
865):867):
868+ if mask_mode != 0:
869+ kv_topk_mode = "no" # mask_mode != 0时，kv_topk_mode只能为no
866 if sparse_indices_mode is None:870 if sparse_indices_mode is None:
867 sparse_indices_mode = "full"871 sparse_indices_mode = "full"
868 if sparse_indices_mode not in ["full", "random"]:872 if sparse_indices_mode not in ["full", "random"]:
Mattention/quant_sparse_flash_mla/tests/pytest/test_quant_sparse_flash_mla_batch.py+1-1
@@ -47,7 +47,7 @@ else:
47 47 
48 48 
49def qsmla(testcase_files):49def qsmla(testcase_files):
50- test_data = torch.load(testcase_files, map_location="cpu")50+ test_data = torch.load(testcase_files, map_location="cpu", weights_only=False)
51 try:51 try:
52 if is_run_graph:52 if is_run_graph:
53 npu_result, cpu_quant_result, cpu_lse, npu_lse = (53 npu_result, cpu_quant_result, cpu_lse, npu_lse = (
Mattention/quant_sparse_flash_mla/tests/pytest/test_quant_sparse_flash_mla_single.py+4-4
@@ -25,8 +25,8 @@ import os
25 25 
26pt_save_path = "qsmla_testcase"26pt_save_path = "qsmla_testcase"
27device_id = 027device_id = 0
28-save_pt = False28+save_pt = os.environ.get("SAVE_PT", "0") == "1"
29-result_path = Path("result.xlsx")29+result_path = Path(os.environ.get("QSMLA_RESULT_SAVE_PATH", "result.xlsx"))
30is_run_graph = os.environ.get("RUN_GRAPH", "0") == "1"30is_run_graph = os.environ.get("RUN_GRAPH", "0") == "1"
31 31 
32# 日志配置32# 日志配置
@@ -152,10 +152,10 @@ def qsmla(param_combinations):
152 lse_res, lse_pct = result_compare_method.check_result(cpu_lse, npu_lse)152 lse_res, lse_pct = result_compare_method.check_result(cpu_lse, npu_lse)
153 153 
154 # 主输出失败记录154 # 主输出失败记录
155- if main_res != "PASS":155+ if main_res != "Pass":
156 fail_info.append(f"MAIN_FAILED:{main_res}")156 fail_info.append(f"MAIN_FAILED:{main_res}")
157 # LSE失败记录157 # LSE失败记录
158- if lse_res is not None and lse_res != "PASS":158+ if lse_res is not None and lse_res != "Pass":
159 fail_info.append(f"LSE_FAILED:{lse_res}")159 fail_info.append(f"LSE_FAILED:{lse_res}")
160 min_fulfill = min(min_fulfill, lse_pct)160 min_fulfill = min(min_fulfill, lse_pct)
161 161 
Mattention/sparse_flash_mla/tests/pytest/batch/sparse_flash_mla_process.py+2-2
@@ -104,7 +104,7 @@ def call_npu(input_data):
104 print("用例参数: ", params)104 print("用例参数: ", params)
105 105 
106 # metadata解析106 # metadata解析
107- K1 = metadata_input["K1"]107+ K1 = metadata_input.get("K1")
108 K = metadata_input["K"]108 K = metadata_input["K"]
109 cmp_ratio = metadata_input["cmp_ratio"]109 cmp_ratio = metadata_input["cmp_ratio"]
110 N1 = metadata_input["N1"]110 N1 = metadata_input["N1"]
@@ -290,7 +290,7 @@ def call_npu_graph(input_data, device_id=0):
290 )290 )
291 291 
292 # metadata解析292 # metadata解析
293- K1 = metadata_input["K1"]293+ K1 = metadata_input.get("K1")
294 K = metadata_input["K"]294 K = metadata_input["K"]
295 cmp_ratio = metadata_input["cmp_ratio"]295 cmp_ratio = metadata_input["cmp_ratio"]
296 N1 = metadata_input["N1"]296 N1 = metadata_input["N1"]
Mattention/sparse_flash_mla/tests/pytest/sparse_flash_mla_golden.py+48-48
@@ -724,50 +724,37 @@ class GeneralizedSFA:
724 else:724 else:
725 cmp_k_bnsd = None725 cmp_k_bnsd = None
726 cmp_k_bnsd_shape = None726 cmp_k_bnsd_shape = None
727+ 
728+ cmp_sparse_indices_bnsd = None
727 if template_idx == 2 or template_idx == 4:729 if template_idx == 2 or template_idx == 4:
728- cmp_sparse_indices_bnsd, cmp_sparse_indices_bnsd_shape = (730+ cmp_sparse_indices_bnsd, cmp_sparse_indices_bnsd_shape = self.trans_shape_to_bnsd(
729- self.trans_shape_to_bnsd(731+ cmp_sparse_indices,
730- cmp_sparse_indices,732+ cmp_sparse_indices.shape,
731- cmp_sparse_indices.shape,733+ self.layout_q,
732- self.layout_q,734+ cu_seqlens_q,
733- cu_seqlens_q,735+ seqused_q,
734- seqused_q,736+ )
735- )737+ 
736- )738+ ori_sparse_indices_bnsd = None
737- else:
738- cmp_sparse_indices_bnsd = None
739 if template_idx == 3 or template_idx == 4:739 if template_idx == 3 or template_idx == 4:
740- ori_sparse_indices_bnsd, ori_sparse_indices_bnsd_shape = (740+ ori_sparse_indices_bnsd, ori_sparse_indices_bnsd_shape = self.trans_shape_to_bnsd(
741- self.trans_shape_to_bnsd(741+ ori_sparse_indices,
742- ori_sparse_indices,742+ ori_sparse_indices.shape,
743- ori_sparse_indices.shape,743+ self.layout_q,
744- self.layout_q,744+ cu_seqlens_q,
745- cu_seqlens_q,745+ seqused_q,
746- seqused_q,746+ )
747- )747+ ori_topk_length_bnsd = None
748- )748+ if (template_idx == 3 or template_idx == 4) and ori_topk_length is not None:
749- else:749+ ori_topk_length_bnsd, _ = self.trans_topk_length_shape_to_bnsd(
750- ori_sparse_indices_bnsd = None750+ ori_topk_length, ori_topk_length.shape, self.layout_q, cu_seqlens_q)
751- if template_idx == 3 or template_idx == 4:751+ 
752- ori_topk_length_bnsd, _ = (752+ cmp_topk_length_bnsd = None
753- self.trans_topk_length_shape_to_bnsd(753+ if (template_idx == 2 or template_idx == 4) and cmp_topk_length is not None:
754- ori_topk_length, ori_topk_length.shape, self.layout_q, cu_seqlens_q754+ cmp_topk_length_bnsd, _ = self.trans_topk_length_shape_to_bnsd(
755- )755+ cmp_topk_length, cmp_topk_length.shape, self.layout_q, cu_seqlens_q)
756- if ori_topk_length is not None756+ 
757- else None757+ 
758- )
759- else:
760- ori_topk_length_bnsd = None
761- if template_idx == 2 or template_idx == 4:
762- cmp_topk_length_bnsd, _ = (
763- self.trans_topk_length_shape_to_bnsd(
764- cmp_topk_length, cmp_topk_length.shape, self.layout_q, cu_seqlens_q
765- )
766- if cmp_topk_length is not None
767- else None
768- )
769- else:
770- cmp_topk_length_bnsd = None
771 758 
772 attn_out, softmax_lse = self.calculate_by_bnsd(759 attn_out, softmax_lse = self.calculate_by_bnsd(
773 q_bnsd,760 q_bnsd,
@@ -878,6 +865,8 @@ def gen_sparse_indices_bsnd(
878 kv_topk_mode,865 kv_topk_mode,
879 topk_length_override=None,866 topk_length_override=None,
880):867):
868+ if mask_mode != 0:
869+ kv_topk_mode = "no" # mask_mode != 0时，kv_topk_mode只能为no
881 if sparse_indices_mode == None: # 不传默认取full, 兼容老用例870 if sparse_indices_mode == None: # 不传默认取full, 兼容老用例
882 sparse_indices_mode = "full"871 sparse_indices_mode = "full"
883 if sparse_indices_mode not in ["full", "random"]:872 if sparse_indices_mode not in ["full", "random"]:
@@ -968,6 +957,8 @@ def gen_sparse_indices_tnd(
968 kv_topk_mode,957 kv_topk_mode,
969 topk_length_override=None,958 topk_length_override=None,
970):959):
960+ if mask_mode != 0:
961+ kv_topk_mode = "no" # mask_mode != 0时，kv_topk_mode只能为no
971 if sparse_indices_mode == None:962 if sparse_indices_mode == None:
972 sparse_indices_mode = "full"963 sparse_indices_mode = "full"
973 if sparse_indices_mode not in ["full", "random"]:964 if sparse_indices_mode not in ["full", "random"]:
@@ -1782,6 +1773,19 @@ def gen_data(params, prepare_device_storage=True):
1782 cmp_topk_length,1773 cmp_topk_length,
1783 return_softmax_lse,1774 return_softmax_lse,
1784 )1775 )
1776+ 
1777+ max_seqlen_ori_kv = 0
1778+ if seqused_ori_kv is not None:
1779+ max_seqlen_ori_kv = seqused_ori_kv.max().item()
1780+ elif cu_seqlens_ori_kv is not None:
1781+ max_seqlen_ori_kv = get_max_adjacent_diff(cu_seqlens_ori_kv)
1782+ 
1783+ max_seqlen_cmp_kv = 0
1784+ if seqused_cmp_kv is not None:
1785+ max_seqlen_cmp_kv = seqused_cmp_kv.max().item()
1786+ elif cu_seqlens_cmp_kv is not None:
1787+ max_seqlen_cmp_kv = get_max_adjacent_diff(cu_seqlens_cmp_kv)
1788+ 
1785 # ORI_SPARSE/ORI_CMP_SPARSE: seqused (actualLength) not passed to op, determined by sparse_indices and topkLength1789 # ORI_SPARSE/ORI_CMP_SPARSE: seqused (actualLength) not passed to op, determined by sparse_indices and topkLength
1786 # cu_seqlens still passed for TND kv layout (needed for data addressing in kernel)1790 # cu_seqlens still passed for TND kv layout (needed for data addressing in kernel)
1787 no_actual_length = template_run_mode in ("ORI_SPARSE", "ORI_CMP_SPARSE")1791 no_actual_length = template_run_mode in ("ORI_SPARSE", "ORI_CMP_SPARSE")
@@ -1823,12 +1827,8 @@ def gen_data(params, prepare_device_storage=True):
1823 "cmp_topk_length": cmp_topk_length if cmp_topk_length is not None else None,1827 "cmp_topk_length": cmp_topk_length if cmp_topk_length is not None else None,
1824 "B": B,1828 "B": B,
1825 "max_seqlen_q": max_seqlen_q,1829 "max_seqlen_q": max_seqlen_q,
1826- "max_seqlen_ori_kv": max(seqused_ori_kv)1830+ "max_seqlen_ori_kv": max_seqlen_ori_kv,
1827- if seqused_ori_kv is not None1831+ "max_seqlen_cmp_kv": max_seqlen_cmp_kv,
1828- else (T2 if layout_kv == "TND" else S2),
1829- "max_seqlen_cmp_kv": max(seqused_cmp_kv)
1830- if seqused_cmp_kv is not None
1831- else 0,
1832 "K1": K1,1832 "K1": K1,
1833 "K": K,1833 "K": K,
1834 "cmp_ratio": cmp_ratio,1834 "cmp_ratio": cmp_ratio,
Mattention/sparse_flash_mla/tests/pytest/test_sparse_flash_mla_batch.py+1-1
@@ -73,7 +73,7 @@ print("files:", locals()["testcase_files"])
73def test_sparse_flash_mla(testcase_files):73def test_sparse_flash_mla(testcase_files):
74 print("执行文件: ", testcase_files)74 print("执行文件: ", testcase_files)
75 torch_npu.npu.set_device(0)75 torch_npu.npu.set_device(0)
76- test_data = torch.load(testcase_files, map_location="cpu")76+ test_data = torch.load(testcase_files, map_location="cpu", weights_only=False)
77 npu_result = None77 npu_result = None
78 try:78 try:
79 npu_result, softmax_lse = sparse_flash_mla_process.call_npu(test_data)79 npu_result, softmax_lse = sparse_flash_mla_process.call_npu(test_data)
Mattention/sparse_flash_mla/tests/pytest/test_sparse_flash_mla_single.py+2-2
@@ -19,8 +19,8 @@ import os
19import pytest19import pytest
20from batch import sparse_flash_mla_process20from batch import sparse_flash_mla_process
21 21 
22-save_pt = False22+save_pt = os.environ.get("SAVE_PT", "0") == "1"
23-pt_save_path = "data"23+pt_save_path = os.environ.get("SMLA_PT_SAVE_PATH", "data")
24 24 
25# 处理所有参数组合25# 处理所有参数组合
26result_path = os.getenv("SMLA_RESULT_SAVE_PATH", "./result/smla_result_all_sparse.xlsx")26result_path = os.getenv("SMLA_RESULT_SAVE_PATH", "./result/smla_result_all_sparse.xlsx")