Star1055
2549
代码
代码
Issues180
Pull Requests414
Actions
项目讨论
Wiki
分析
Star1055
2549
已合并
SALS V1 #3711
mayafei创建于 4月2日
SALS V1 #3711
已合并
mayafei创建于 4月2日
从mayafei/ops-transformer-SALS-experimentalmaster合入到cann/ops-transformermaster
共 58 个文件变更+18811-1
Aexperimental/attention/quant_sals_indexer/CMakeLists.txt+0-0
The file is empty
Aexperimental/attention/quant_sals_indexer/examples/test_npu_quant_sals_indexer_int4_metadata.py+447-0
@@ -0,0 +1,447 @@
1+# ---------------------------------------------------------------------------------------------------------
2+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+# CANN Open Software License Agreement Version 2.0 (the "License").
5+# Please refer to the License for details. You may not use this file except in compliance with the License.
6+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+# See LICENSE in the root of the software repository for the full text of the License.
9+# ---------------------------------------------------------------------------------------------------------
10+ 
11+import torch
12+import torch_npu
13+import torchair
14+import custom_ops
15+import numpy as np
16+import torch.nn as nn
17+import math
18+ 
19+from torch_npu.testing.testcase import TestCase, run_tests
20+ 
21+DEVICE_ID = 0
22+torch_npu.npu.set_device(int(DEVICE_ID))
23+ 
24+def convert_pa_nz_to_pa_bnsd(key_pa_nz, d=None):
25+ """
26+ 将PA_NZ格式转换为PA_BNSD格式
27+
28+ Args:
29+ key_pa_nz: PA_NZ格式的key，形状为 (B, N, d//64, S, 8)
30+ d: 可选参数，原始特征维度，用于验证
31+
32+ Returns:
33+ PA_BNSD格式的key，形状为 (B, N, S, d//8)
34+ """
35+ B, N, z_blocks, S, eight = key_pa_nz.shape
36+
37+ # 可选验证
38+ if d is not None:
39+ expected_z_blocks = d // 64
40+ assert z_blocks == expected_z_blocks, \
41+ f"z_blocks({z_blocks}) != d//64({expected_z_blocks})"
42+
43+ # 转换
44+ result = key_pa_nz.permute(0, 1, 3, 2, 4).reshape(B, N, S, -1)
45+
46+ return result
47+ 
48+ 
49+def _get_data_from_pa_cache(key, block_table, act_s2, layout_key):
50+ if layout_key == 'PA_BNSD' or layout_key == 'PA_NZ':
51+ block_num, n2, block_size, d = key.shape
52+ else:
53+ block_num, block_size, n2, d = key.shape
54+ need_blcok_num = (act_s2 + block_size - 1) // block_size
55+ act_s2_align = need_blcok_num * block_size
56+ if layout_key == 'PA_BNSD' or layout_key == 'PA_NZ':
57+ out = torch.zeros((n2, act_s2_align, d), dtype=key.dtype, device=key.device)
58+ else:
59+ out = torch.zeros((act_s2_align, n2, d), dtype=key.dtype, device=key.device)
60+ for i in range(need_blcok_num):
61+ if layout_key == 'PA_BNSD' or layout_key == 'PA_NZ':
62+ out[:, i * block_size:(i + 1) * block_size, :] = key[block_table[i], ...].reshape(n2, block_size, d)
63+ else:
64+ out[i * block_size:(i + 1) * block_size, :, :] = key[block_table[i], ...].reshape(block_size, n2, d)
65+ if layout_key == 'PA_BNSD' or layout_key == 'PA_NZ':
66+ return out[:, :act_s2, :]
67+ else:
68+ return out[:act_s2, :, :]
69+ 
70+ 
71+def _get_k_scale(key_dequant_scale, block_table, act_s2, layout_key):
72+ if layout_key == 'PA_BNSD' or layout_key == 'PA_NZ':
73+ block_num, n2, block_size = key_dequant_scale.shape
74+ else:
75+ block_num, block_size, n2 = key_dequant_scale.shape
76+ need_blcok_num = (act_s2 + block_size - 1) // block_size
77+ act_s2_align = need_blcok_num * block_size
78+ if layout_key == 'PA_BNSD' or layout_key == 'PA_NZ':
79+ out = torch.zeros((n2, act_s2_align), dtype=key_dequant_scale.dtype, device=key_dequant_scale.device)
80+ key_dequant_scale = key_dequant_scale.reshape(block_num, n2, block_size)
81+ else:
82+ out = torch.zeros((act_s2_align, n2), dtype=key_dequant_scale.dtype, device=key_dequant_scale.device)
83+ key_dequant_scale = key_dequant_scale.reshape(block_num, block_size, n2)
84+ for i in range(need_blcok_num):
85+ if layout_key == 'PA_BNSD' or layout_key == 'PA_NZ':
86+ out[:, i * block_size:(i + 1) * block_size] = key_dequant_scale[block_table[i], ...].reshape(n2, block_size)
87+ else:
88+ out[i * block_size:(i + 1) * block_size, :] = key_dequant_scale[block_table[i], ...].reshape(block_size, n2)
89+ if layout_key == 'PA_BNSD' or layout_key == 'PA_NZ':
90+ return out[:, :act_s2]
91+ else:
92+ return out[:act_s2, :]
93+ 
94+def trans_int32_2_int4_torch(key):
95+ """
96+ PyTorch版本的int32转int4转换
97+ """
98+ # 确保key是torch张量
99+ if not isinstance(key, torch.Tensor):
100+ key = torch.tensor(key)
101+
102+ # 将key转换为uint32类型进行位操作
103+ key_uint32 = key.to(torch.int32) # 使用int32而不是uint32，因为PyTorch对uint32支持有限
104+
105+ # 创建移位张量
106+ shifts = torch.arange(0, 32, 4, device=key.device, dtype=torch.int32)
107+
108+ # 使用广播进行向量化移位
109+ # 扩展维度以便广播 [..., 1] >> [8] -> [..., 8]
110+ expanded_key = key_uint32.unsqueeze(-1) # 在最后添加一个维度
111+ shifted = expanded_key >> shifts # 广播移位
112+
113+ # 提取低4位
114+ parts = shifted & 0xF
115+
116+ # 处理符号位和数值
117+ sign = (parts >> 3) & 1
118+ value = parts & 0x7
119+
120+ # 使用torch.where进行条件赋值
121+ result = torch.where(sign == 1, -((value ^ 0x7) + 1), value)
122+
123+ return result.to(torch.int8)
124+ 
125+def optimize_conversion_torch(key):
126+ """
127+ PyTorch优化的主函数
128+ """
129+ # 获取原始形状
130+ original_shape = key.shape
131+
132+ # 如果key是numpy数组，转换为torch张量
133+ if isinstance(key, np.ndarray):
134+ key = torch.from_numpy(key)
135+
136+ # 重塑为2D (batch, features)
137+ key_2d = key.reshape(-1, original_shape[-1])
138+
139+ # 向量化转换
140+ int4_data = trans_int32_2_int4_torch(key_2d)
141+
142+ # 重塑为最终形状
143+ final_shape = original_shape[:-1] + (original_shape[-1] * 8,)
144+ key_int4 = int4_data.reshape(final_shape)
145+
146+ return key_int4
147+ 
148+#将Int32类型的数据，拆成int4，存进int8
149+def trans_int32_2_int4(input_int32):
150+ # 将Int32类型的数据按bit位平均拆成8份，每份长度为4bit
151+ parts = [(input_int32 >> i) & 0xf for i in range(0, 32, 4)]
152+ output_int4 = []
153+ for part in parts:
154+ # 将每份数据构造成一个Int8类型的数据
155+ # 符号位为第一个bit的值
156+ sign = (part >> 3) & 0x1
157+ # 剩余的值作为int8最后3个bit的值
158+ value = part & 0x7
159+ if sign == 1:
160+ # 如果符号位为1，则需要将value转换成负数
161+ value = -((value ^ 0x7) + 1)
162+ output_int4.append(value)
163+ return output_int4
164+ 
165+def _quant_sals_indexer(query, key, query_dequant_scale, key_dequant_scale, actual_seq_lengths_key, block_table,
166+ sparse_block_size, sparse_ratio, fixed_tail_count, layout_key, max_seqlen_key):
167+ if layout_key == 'PA_NZ':
168+ key = convert_pa_nz_to_pa_bnsd(key)
169+ sparse_ratio = round(1.0 - sparse_ratio, 2)
170+ batch_size = query.shape[0]
171+ n2 = query.shape[1]
172+ # query转换
173+ q_shape_int4 = [query.shape[0], query.shape[1], query.shape[2] * 8]
174+ query_int4 = torch.zeros(q_shape_int4, dtype=torch.int8, device=query.device)
175+ query_int4 = optimize_conversion_torch(query)
176+
177+ # key转换
178+ k_shape_int4 = [key.shape[0], key.shape[1], key.shape[2], key.shape[3] * 8]
179+ key_int4 = torch.zeros(k_shape_int4, dtype=torch.int8, device=key.device)
180+ key_int4 = optimize_conversion_torch(key)
181+ d = query_int4.shape[-1]
182+ max_count = (max_seqlen_key + sparse_block_size - 1) // sparse_block_size
183+ if max_count - fixed_tail_count >= 0:
184+ sparse_count = min(int(math.ceil((max_count - fixed_tail_count) * sparse_ratio)), 2048) + fixed_tail_count
185+ else:
186+ sparse_count = max_count
187+ sparse_indices_shape = [batch_size, n2, sparse_count]
188+ # 初始化为全-1
189+ sparse_indices = torch.zeros(sparse_indices_shape, dtype=torch.int32, device=query.device) - 1
190+ 
191+ for batch_id in range(batch_size):
192+ act_s2 = actual_seq_lengths_key[batch_id]
193+ fixed_tail_count_tmp = fixed_tail_count
194+ 
195+ act_n_count = (act_s2 + sparse_block_size - 1) // sparse_block_size
196+ sort_n_count = act_n_count - fixed_tail_count_tmp if act_n_count - fixed_tail_count_tmp > 0 else 0
197+ topk_n_count = int(math.ceil(sort_n_count * sparse_ratio))
198+ fixed_tail_count_tmp = act_n_count if act_n_count - fixed_tail_count_tmp <= 0 else fixed_tail_count_tmp
199+ topk_n_count = min(topk_n_count, 2048)
200+ if topk_n_count > 0:
201+ # b, n2, d
202+ now_q = query_int4[batch_id, :, :].reshape(n2, 1, d).to(torch.int32)
203+ now_block_table = block_table[batch_id, :]
204+ # s2, n2, d -> n2, d, s2
205+ if layout_key == "PA_BNSD" or layout_key == "PA_NZ":
206+ now_k = _get_data_from_pa_cache(key_int4, now_block_table, act_s2, layout_key).permute(0, 2, 1).to(torch.int32)
207+ now_k_scale = _get_k_scale(key_dequant_scale, now_block_table, act_s2, layout_key).permute(1, 0)
208+ elif layout_key == "PA_BSND":
209+ now_k = _get_data_from_pa_cache(key_int4, now_block_table, act_s2, layout_key).permute(1, 2, 0).to(torch.int32)
210+ now_k_scale = _get_k_scale(key_dequant_scale, now_block_table, act_s2, layout_key)
211+ else:
212+ now_k = key_int4[batch_id, :act_s2, :, :].permute(1, 2, 0).to(torch.int32)
213+ now_k_scale = key_dequant_scale[batch_id, :act_s2, :].to(torch.float)
214+ now_q_scale = query_dequant_scale[batch_id, :].to(torch.float)
215+ now_qk_scale = now_q_scale * now_k_scale
216+ # n2,1,d @ d,s2 -> n2,1,s2
217+ s_out = torch.matmul(now_q, now_k)
218+ now_qk_scale = now_qk_scale.permute(1, 0).view(n2, 1, act_s2) # 广播到(n2, 1, s2)
219+ s_out = s_out * now_qk_scale
220+ s_max = torch.zeros(n2, 1, sort_n_count, dtype=torch.float32, device=key.device)
221+ s_lse = torch.zeros(n2, 1, sort_n_count, dtype=torch.float32, device=key.device)
222+ for sort_n_idx in range(sort_n_count):
223+ max_value, _ = torch.max(s_out[:,:, sort_n_idx * sparse_block_size : (sort_n_idx + 1)* sparse_block_size], axis = 2)
224+ s_max[:, :, sort_n_idx] = max_value
225+ s_lse[:, :, sort_n_idx] = s_max[:, :, sort_n_idx] + torch.log(torch.sum(torch.exp(s_out[:,:, sort_n_idx * sparse_block_size : (sort_n_idx + 1)* sparse_block_size] - s_max[:, :, sort_n_idx: sort_n_idx + 1]), axis=2))
226+ 
227+ 
228+ sorted_value, sorted_indices = torch.sort(s_lse, dim=2, descending=True, stable=True)
229+ sparse_indices[batch_id, :, :topk_n_count] = sorted_indices.to(torch.int32).permute(1,0,2)[:, :, :topk_n_count]
230+ for fixed_n_count_idx in range(fixed_tail_count_tmp):
231+ sparse_indices[batch_id, :, topk_n_count + fixed_n_count_idx] = sort_n_count + fixed_n_count_idx
232+ return sparse_indices
233+ 
234+ 
235+def _compare_res(golden_res, npu_res, ratio_thress_hold=0.999):
236+ npu_res = npu_res.reshape(-1)
237+ golden_res = golden_res.reshape(-1)
238+ total_res_num = golden_res.numel()
239+ diff_res = npu_res - golden_res
240+ match_ratio = (diff_res == 0).sum().float() / total_res_num
241+ if match_ratio >= ratio_thress_hold:
242+ print(f"Compare npu cpu res success! Match ratio is {match_ratio:.4%}")
243+ return True
244+ else:
245+ print(f"Match ratio {match_ratio:.4%} is under thress_hold {ratio_thress_hold} ",
246+ "Please Check!")
247+ non_zero_index = torch.nonzero(diff_res, as_tuple=False).squeeze(1)
248+ npu_index = npu_res[non_zero_index]
249+ golden_index = golden_res[non_zero_index]
250+ for i in non_zero_index:
251+ print(f"mismatch idx: {i}, golden and npu res is {golden_res[i]} || {npu_res[i]}")
252+ return False
253+ 
254+class QuantSINetwork(nn.Module):
255+ def __init__(self):
256+ super(QuantSINetwork, self).__init__()
257+ 
258+ def forward(self, b, s2, n2, d, query, key, query_dequant_scale, key_dequant_scale, actual_seq_lengths_key, block_table, max_seqlen_key, sparse_block_size, sparse_ratio,
259+ fixed_tail_count, layout_key):
260+ # super kernel test
261+ with torchair.scope.super_kernel("sp_QsQSI", ""):
262+ metadata = torch.ops.custom.npu_quant_sals_indexer_metadata(
263+ b,
264+ s2,
265+ n2,
266+ d,
267+ sparse_block_size=sparse_block_size,
268+ sparse_ratio=sparse_ratio,
269+ fixed_tail_count=fixed_tail_count,
270+ layout_key=layout_key,
271+ actual_seq_lengths_kv=actual_seq_lengths_key)
272+ 
273+ output0 = torch.ops.custom.npu_quant_sals_indexer(query, key, query_dequant_scale = query_dequant_scale,
274+ key_dequant_scale = key_dequant_scale, metadata=metadata,
275+ actual_seq_lengths_key=actual_seq_lengths_key,
276+ block_table=block_table, max_seqlen_key=max_seqlen_key,
277+ sparse_block_size=sparse_block_size,
278+ sparse_ratio=sparse_ratio, fixed_tail_count=fixed_tail_count,
279+ layout_key=layout_key)
280+ return output0
281+ 
282+ 
283+class TestCustomLightningIndexerQuant(TestCase):
284+ def cpu_op_exec(self, query, key, query_dequant_scale, key_dequant_scale, actual_seq_lengths_key, block_table,
285+ sparse_block_size, sparse_ratio, fixed_tail_count, layout_key, max_seqlen_key):
286+ output0 = _quant_sals_indexer(query, key, query_dequant_scale, key_dequant_scale, actual_seq_lengths_key, block_table,
287+ sparse_block_size, sparse_ratio, fixed_tail_count, layout_key, max_seqlen_key)
288+ output0 = output0.cpu()
289+ 
290+ return output0
291+ 
292+ def npu_op_exec_graph(self, b, s2, n2, d, query, key, query_dequant_scale, key_dequant_scale, actual_seq_lengths_key,
293+ block_table, max_seqlen_key, sparse_block_size, sparse_ratio, fixed_tail_count, layout_key):
294+ npu_mode = QuantSINetwork().to("npu:%s" % DEVICE_ID)
295+ from torchair.configs.compiler_config import CompilerConfig
296+ config = CompilerConfig()
297+ npu_backend = torchair.get_npu_backend(compiler_config=config)
298+ torch._dynamo.reset()
299+ npu_mode = torch.compile(npu_mode, fullgraph=True, backend=npu_backend, dynamic=False)
300+ npu_out0 = npu_mode(b, s2, n2, d, query, key,
301+ query_dequant_scale = query_dequant_scale,
302+ key_dequant_scale = key_dequant_scale,
303+ actual_seq_lengths_key=actual_seq_lengths_key,
304+ block_table=block_table,
305+ max_seqlen_key = max_seqlen_key,
306+ sparse_block_size=sparse_block_size,
307+ sparse_ratio=sparse_ratio,
308+ fixed_tail_count=fixed_tail_count,
309+ layout_key=layout_key)
310+ npu_out0 = npu_out0.cpu()
311+ return npu_out0
312+ 
313+ def npu_op_exec_eager(self, b, s2, n2, d, query, key, query_dequant_scale, key_dequant_scale, actual_seq_lengths_key,
314+ block_table, max_seqlen_key, sparse_block_size, sparse_ratio, fixed_tail_count, layout_key):
315+
316+ metadata = torch.ops.custom.npu_quant_sals_indexer_metadata(
317+ b,
318+ s2,
319+ n2,
320+ d,
321+ sparse_block_size=sparse_block_size,
322+ sparse_ratio=sparse_ratio,
323+ fixed_tail_count=fixed_tail_count,
324+ layout_key=layout_key,
325+ actual_seq_lengths_kv=actual_seq_lengths_key)
326+ npu_out0 = torch.ops.custom.npu_quant_sals_indexer(query, key,
327+ metadata=metadata,
328+ query_dequant_scale = query_dequant_scale,
329+ key_dequant_scale = key_dequant_scale,
330+ actual_seq_lengths_key=actual_seq_lengths_key,
331+ block_table=block_table,
332+ max_seqlen_key = max_seqlen_key,
333+ sparse_block_size=sparse_block_size,
334+ sparse_ratio=sparse_ratio,
335+ fixed_tail_count=fixed_tail_count,
336+ layout_key=layout_key)
337+ npu_out0 = npu_out0.cpu()
338+ return npu_out0
339+ 
340+ def quant_sals_indexer_result(self, b, s2, n2, d, act_seq_k, sparse_block_size, sparse_ratio, fixed_tail_count, max_seqlen_key):
341+ # -----固定参数--------
342+ block_size = 512
343+ layout_key = 'PA_BNSD'
344+ np.random.seed(0)
345+ # -------------
346+ max_block_table_num = (s2 + block_size - 1) // block_size
347+ block_table = torch.tensor([range(b * max_block_table_num)], dtype = torch.int32).reshape(b, -1)
348+ if layout_key == 'PA_BNSD':
349+ key = torch.tensor(np.random.uniform(-2147483648, 2147483647, (b * max_block_table_num, n2, block_size, d // 8))).to(torch.int32) # d轴8个数合并
350+ key_dequant_scale = torch.tensor(np.random.uniform(-1, 1, (b * max_block_table_num, n2, block_size))).to(torch.float)
351+ elif layout_key == 'PA_BSND':
352+ key = torch.tensor(np.random.uniform(-2147483648, 2147483647, (b * max_block_table_num, block_size, n2, d // 8))).to(torch.int32) # d轴8个数合并
353+ key_dequant_scale = torch.tensor(np.random.uniform(-1, 1, (b * max_block_table_num, block_size, n2))).to(torch.float)
354+ elif layout_key == 'PA_NZ':
355+ key = torch.tensor(np.random.uniform(-2147483648, 2147483647, (b * max_block_table_num, n2, ((d // 8) // 8), block_size, 8))).to(torch.int32)
356+ key_dequant_scale = torch.tensor(np.random.uniform(-1, 1, (b * max_block_table_num, n2, block_size))).to(torch.float)
357+ else:
358+ key = torch.tensor(np.random.uniform(-2147483648, 2147483647, (b, s2, n2, d // 8))).to(torch.int32) # d轴8个数合并
359+ key_dequant_scale = torch.tensor(np.random.uniform(-1, 1, (b, s2, n2))).to(torch.float)
360+ 
361+ query = torch.tensor(np.random.uniform(-2147483648, 2147483647, (b, n2, d // 8))).to(torch.int32) # d轴8个数合并
362+ actual_seq_lengths_key = torch.tensor(act_seq_k).to(torch.int32)
363+ query_dequant_scale = torch.tensor(np.random.uniform(-1, 1, (b, n2))).to(torch.float)
364+ print(f"------- test QuantSI BSND case b:{b} n2:{n2} s2:{s2} ----------")
365+ 
366+ cpu_out0 = self.cpu_op_exec(query, key, query_dequant_scale, key_dequant_scale, actual_seq_lengths_key,
367+ block_table, sparse_block_size, sparse_ratio, fixed_tail_count, layout_key, max_seqlen_key)
368+ 
369+ torch_npu.npu.set_device(int(DEVICE_ID))
370+ query = query.to("npu:%s" % DEVICE_ID)
371+ key = key.to("npu:%s" % DEVICE_ID)
372+ query_dequant_scale = query_dequant_scale.to("npu:%s" % DEVICE_ID)
373+ key_dequant_scale = key_dequant_scale.to("npu:%s" % DEVICE_ID)
374+ actual_seq_lengths_key = actual_seq_lengths_key.to("npu:%s" % DEVICE_ID)
375+ if layout_key == 'PA_BNSD' or layout_key == 'PA_BSND' or layout_key == 'PA_NZ':
376+ block_table = block_table.to("npu:%s" % DEVICE_ID)
377+ else:
378+ block_table = None
379+ print("run eager mode")
380+ sparse_indices = self.npu_op_exec_eager(b, s2, n2, d, query, key, query_dequant_scale, key_dequant_scale,
381+ actual_seq_lengths_key, block_table, max_seqlen_key,
382+ sparse_block_size, sparse_ratio, fixed_tail_count, layout_key)
383+ print("sparse_indices = ", sparse_indices)
384+ assert(_compare_res(cpu_out0, sparse_indices))
385+ 
386+ print("run graph mode")
387+ sparse_indices = self.npu_op_exec_graph(b, s2, n2, d, query, key, query_dequant_scale, key_dequant_scale,
388+ actual_seq_lengths_key, block_table, max_seqlen_key,
389+ sparse_block_size, sparse_ratio, fixed_tail_count, layout_key)
390+ 
391+ assert(_compare_res(cpu_out0, sparse_indices))
392+ 
393+ def test_quant_sals_indexer(self):
394+ # b, s2, n2, d, act_seq_k, sparse_block_size, sparsity, fixed_tail_count, max_seqlen_key
395+ import random
396+ b = 21 * 4
397+ n2 = 8 // 4
398+ base_s2 = 1024 * 10
399+ random_seed = 42
400+ unbalance_cache = 0
401+ random.seed(random_seed)
402+
403+ act_seq_k = []
404+ sumseqlength = 0
405+ for _ in range(b-1):
406+ offset_percent = random.randint(-unbalance_cache, unbalance_cache)
407+ offset = int(base_s2 * (offset_percent / 100))
408+ s2curent = base_s2 + offset
409+ sumseqlength = sumseqlength + s2curent
410+ act_seq_k.append(s2curent)
411+ totalseq = b * base_s2
412+ s2_last = totalseq - sumseqlength
413+ act_seq_k.append(s2_last)
414+ max_seqlen_key = 256*1024
415+ print('QSI test case, act_seq_k:', act_seq_k)
416+ print('QSI test case, max_seqlen_key:', max_seqlen_key)
417+ test_case_list = [
418+ #基础性能/功能用例
419+ (b, max_seqlen_key, n2, 64, act_seq_k, 16, 0.75, 32, max_seqlen_key),
420+ # (21, 10240, 8, 64, [10240] * 21, 16, 0.75, 16, 10240),
421+ # 泛化shape
422+ # (16, 512*64, 1, 64, [19,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0], 16, 0.75, 1, 19),
423+ # (16, 512*64, 1, 64, [33,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0], 16, 0.75, 1, 33),
424+ # (16, 512*64, 1, 64, [6892,1458,3356,2001,0,0,0,0,0,0,0,0,0,0,0,0], 16, 0.75, 1, 6892),
425+ # (1, 512*64, 1, 64, [100], 16, 0.75, 1, 100),
426+ # (1, 2289, 1, 64, [2289], 16, 0.75, 1, 2289),
427+ # (1, 4352, 1, 64, [4352], 16, 0.75, 1, 4352),
428+ # (35, 4727, 1, 64, [0, 0, 3013, 334, 3194, 4029, 116, 4419, 41, 1495, 4295, 4471, 666, 582, 2863, 847, 4147, 146, 591, 2206, 4318, 3770, 1412, 1782, 853, 4721, 561, 1198, 4043, 2166, 2605, 209, 7, 3864, 2289],
429+ # 16, 0.75, 16, 4727),
430+ # (4, 2358, 4, 64, [2358] * 4, 16, 0.75, 16, 2358),
431+ # (4, 1111, 4, 64, [1111] * 4, 16, 0.75, 16, 1111),
432+ # (4, 2003, 4, 64, [2003] * 4, 16, 0.75, 16, 2003),
433+ # (4, 7891, 4, 64, [7891] * 4, 16, 0.75, 16, 7891),
434+ # (4, 6100, 4, 64, [6100] * 4, 16, 0.75, 16, 6100),
435+ # (3, 10240, 4, 64, [10240] * 3, 16, 0.75, 16, 10240),
436+ # (11, 10240, 4, 64, [10240] * 11, 16, 0.75, 16, 10240),
437+ # (29, 10240, 4, 64, [10240] * 29, 16, 0.75, 16, 10240),
438+ # (29, 1234, 4, 64, [1234] * 29, 16, 0.78, 16, 1234),
439+ # (33, 887, 4, 64, [887] * 33, 16, 0.85, 16, 887),
440+ # (19, 9813, 4, 64, [9813] * 19, 16, 0.65, 16, 9813),
441+ ]
442+ for case in test_case_list:
443+ self.quant_sals_indexer_result(*case)
444+ 
445+ 
446+if __name__ == "__main__":
447+ run_tests()
Aexperimental/attention/quant_sals_indexer/op_host/quant_sals_indexer_def.cpp+77-0
@@ -0,0 +1,77 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file quant_sals_indexer_def.cpp
13+ * \brief
14+ */
15+#include <cstdint>
16+#include "register/op_def_registry.h"
17+ 
18+namespace ops {
19+class QuantSalsIndexer : public OpDef {
20+public:
21+ explicit QuantSalsIndexer(const char *name) : OpDef(name)
22+ {
23+ this->Input("query")
24+ .ParamType(REQUIRED)
25+ .DataType({ge::DT_INT8, ge::DT_INT32})
26+ .FormatList({ge::FORMAT_ND})
27+ .AutoContiguous();
28+ this->Input("key")
29+ .ParamType(REQUIRED)
30+ .DataType({ge::DT_INT8, ge::DT_INT32})
31+ .FormatList({ge::FORMAT_ND})
32+ .AutoContiguous();
33+ this->Input("query_dequant_scale")
34+ .ParamType(REQUIRED)
35+ .DataTypeList({ge::DT_FLOAT})
36+ .FormatList({ge::FORMAT_ND})
37+ .AutoContiguous();
38+ this->Input("key_dequant_scale")
39+ .ParamType(REQUIRED)
40+ .DataTypeList({ge::DT_FLOAT})
41+ .FormatList({ge::FORMAT_ND})
42+ .AutoContiguous();
43+ this->Input("actual_seq_lengths_key")
44+ .ParamType(OPTIONAL)
45+ .DataTypeList({ge::DT_INT32})
46+ .FormatList({ge::FORMAT_ND})
47+ .AutoContiguous();
48+ this->Input("block_table")
49+ .ParamType(OPTIONAL)
50+ .DataTypeList({ge::DT_INT32})
51+ .FormatList({ge::FORMAT_ND})
52+ .AutoContiguous();
53+ this->Input("metadata")
54+ .ParamType(OPTIONAL)
55+ .DataTypeList({ge::DT_INT32})
56+ .FormatList({ge::FORMAT_ND})
57+ .AutoContiguous();
58+ this->Output("sparse_indices").ParamType(REQUIRED).DataTypeList({ge::DT_INT32}).FormatList({ge::FORMAT_ND});
59+ this->Attr("max_seqlen_key").AttrType(REQUIRED).Int(0);
60+ this->Attr("sparse_block_size").AttrType(REQUIRED).Int(16);
61+ this->Attr("sparse_ratio").AttrType(REQUIRED).Float(0.25f);
62+ this->Attr("fixed_tail_count").AttrType(REQUIRED).Int(1);
63+ this->Attr("layout_key").AttrType(OPTIONAL).String("BSND");
64+ OpAICoreConfig aicore_config;
65+ aicore_config.DynamicCompileStaticFlag(true)
66+ .DynamicFormatFlag(true)
67+ .DynamicRankSupportFlag(true)
68+ .DynamicShapeSupportFlag(true)
69+ .NeedCheckSupportFlag(false)
70+ .PrecisionReduceFlag(true)
71+ .ExtendCfgInfo("aclnnSupport.value", "support_aclnn");
72+ this->AICore().AddConfig("ascend910b", aicore_config);
73+ this->AICore().AddConfig("ascend910_93", aicore_config);
74+ }
75+};
76+OP_ADD(QuantSalsIndexer);
77+} // namespace ops
Aexperimental/attention/quant_sals_indexer/op_host/quant_sals_indexer_proto.cpp+107-0
@@ -0,0 +1,107 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file quant_sals_indexer_proto.cpp
13+ * \brief
14+ */
15+#include <cmath>
16+#include <graph/utils/type_utils.h>
17+#include <register/op_impl_registry.h>
18+#include "error/ops_error.h"
19+ 
20+ 
21+using namespace ge;
22+ 
23+namespace ops {
24+constexpr uint32_t QUERY_INDEX = 0;
25+constexpr uint32_t KEY_INDEX = 1;
26+constexpr uint32_t ACTUAL_SEQ_K_INDEX = 4;
27+ 
28+constexpr uint32_t ATTR_MAX_SEQLEN_KEY_INDEX = 0;
29+constexpr uint32_t ATTR_SPARSE_BLOCK_SIZE_INDEX = 1;
30+constexpr uint32_t ATTR_SPARSE_RATIO_INDEX = 2;
31+constexpr uint32_t ATTR_FIXED_TAIL_COUNT_INDEX = 3;
32+constexpr uint32_t ATTR_KEY_LAYOUT_INDEX = 4;
33+ 
34+constexpr uint32_t OUTPUT_IDX_SPARSE_INDICES = 0;
35+constexpr uint32_t OUTPUT_IDX_SPARSE_SEQ_LENGTHS_KEY = 1;
36+ 
37+static ge::graphStatus InferShapeQuantSalsIndexer(gert::InferShapeContext *context)
38+{
39+ OPS_ERR_IF(context == nullptr, OPS_LOG_E("QuantSalsIndexer", "InferShapeContext is nullptr!"),
40+ return ge::GRAPH_FAILED);
41+ const gert::Shape *queryShape = context->GetInputShape(QUERY_INDEX);
42+ OPS_LOG_E_IF_NULL(context, queryShape, return ge::GRAPH_FAILED);
43+ const gert::Shape *keyShape = context->GetInputShape(KEY_INDEX);
44+ OPS_LOG_E_IF_NULL(context, keyShape, return ge::GRAPH_FAILED);
45+ 
46+ auto attrs = context->GetAttrs();
47+ OPS_LOG_E_IF_NULL(context, attrs, return ge::GRAPH_FAILED);
48+ const char *inputLayoutKeyPtr = attrs->GetAttrPointer<char>(ATTR_KEY_LAYOUT_INDEX);
49+ OPS_LOG_E_IF_NULL(context, inputLayoutKeyPtr, return ge::GRAPH_FAILED);
50+ 
51+ const int64_t *maxSeqlenKey = attrs->GetInt(ATTR_MAX_SEQLEN_KEY_INDEX);
52+ OPS_LOG_E_IF_NULL(context, maxSeqlenKey, return ge::GRAPH_FAILED);
53+ const int64_t *sparseBlockSize = attrs->GetInt(ATTR_SPARSE_BLOCK_SIZE_INDEX);
54+ OPS_LOG_E_IF_NULL(context, sparseBlockSize, return ge::GRAPH_FAILED);
55+ const int64_t *fixedTailCount = attrs->GetInt(ATTR_FIXED_TAIL_COUNT_INDEX);
56+ OPS_LOG_E_IF_NULL(context, fixedTailCount, return ge::GRAPH_FAILED);
57+ const float *sparse_ratio = attrs->GetFloat(ATTR_SPARSE_RATIO_INDEX);
58+ OPS_LOG_E_IF_NULL(context, sparse_ratio, return ge::GRAPH_FAILED);
59+ 
60+ int64_t totalNCount = ((*maxSeqlenKey) + ((*sparseBlockSize) - 1)) / (*sparseBlockSize);
61+ int64_t sparseCount;
62+ int64_t selectCount;
63+ if (totalNCount < (*fixedTailCount)) {
64+ sparseCount = totalNCount;
65+ selectCount = 0;
66+ } else {
67+ selectCount = static_cast<int64_t>(std::ceil((static_cast<double>(totalNCount) - static_cast<double>(*fixedTailCount)) * (std::round((1.0 - *sparse_ratio)*100.0) / 100.0)));
68+ if (selectCount > 2048) {
69+ selectCount = 2048;
70+ }
71+ sparseCount = selectCount + (*fixedTailCount);
72+ }
73+ 
74+ gert::Shape *sparseIndicesShape = context->GetOutputShape(OUTPUT_IDX_SPARSE_INDICES);
75+ OPS_LOG_E_IF_NULL(context, sparseIndicesShape, ge::GRAPH_FAILED);
76+ 
77+ sparseIndicesShape->SetDimNum(queryShape->GetDimNum());
78+ OPS_ERR_IF(
79+ queryShape->GetDimNum() != 3,
80+ OPS_LOG_E(context, "queryDims (%zu) must be 3!", queryShape->GetDimNum()),
81+ return ge::GRAPH_FAILED);
82+ sparseIndicesShape->SetDim(0, queryShape->GetDim(0)); // 0:Dim B
83+ sparseIndicesShape->SetDim(1, queryShape->GetDim(1)); // 1:Dim N2
84+ sparseIndicesShape->SetDim(2, sparseCount); // 2:Dim K
85+ 
86+ OPS_LOG_D(context->GetNodeName(), "QuantSalsIndexer InferShape end.");
87+ 
88+ return ge::GRAPH_SUCCESS;
89+}
90+ 
91+static ge::graphStatus InferDataTypeQuantSalsIndexer(gert::InferDataTypeContext *context)
92+{
93+ OPS_ERR_IF(context == nullptr, OPS_LOG_E("QuantSalsIndexer", "InferDataTypeContext is nullptr!"),
94+ return ge::GRAPH_FAILED);
95+ OPS_LOG_D(context->GetNodeName(), "Enter QuantSalsIndexer InferDataType impl.");
96+ // default set q's dtype as fia's output type
97+ ge::DataType outputType = ge::DT_INT32;
98+ // attention_out, outidx:0
99+ context->SetOutputDataType(0, outputType);
100+ OPS_LOG_D(context->GetNodeName(), "QuantSalsIndexer InferDataType end.");
101+ return GRAPH_SUCCESS;
102+}
103+ 
104+IMPL_OP_INFERSHAPE(QuantSalsIndexer)
105+ .InferShape(InferShapeQuantSalsIndexer)
106+ .InferDataType(InferDataTypeQuantSalsIndexer);
107+} // namespace ops
Aexperimental/attention/quant_sals_indexer/op_host/quant_sals_indexer_tiling.cpp+689-0
@@ -0,0 +1,689 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file quant_sals_indexer_tiling.cpp
13+ * \brief
14+ */
15+ 
16+#include "quant_sals_indexer_tiling.h"
17+#include "../op_kernel/quant_sals_indexer_template_tiling_key.h"
18+#include "split_core.h"
19+ 
20+using namespace ge;
21+using namespace AscendC;
22+using std::map;
23+using std::string;
24+using namespace optiling::qsi;
25+namespace optiling {
26+// --------------------------QSIInfoParser类成员函数定义-------------------------------------
27+ge::graphStatus QSIInfoParser::CheckRequiredInOutExistence() const
28+{
29+ OPS_ERR_IF(opParamInfo_.query.shape == nullptr, OPS_LOG_E(opName_, "Shape of tensor query is nullptr"),
30+ return ge::GRAPH_FAILED);
31+ OPS_ERR_IF(opParamInfo_.query.desc == nullptr, OPS_LOG_E(opName_, "Desc of tensor query is nullptr"),
32+ return ge::GRAPH_FAILED);
33+ OPS_ERR_IF(opParamInfo_.key.shape == nullptr, OPS_LOG_E(opName_, "Shape of tensor k is nullptr"),
34+ return ge::GRAPH_FAILED);
35+ OPS_ERR_IF(opParamInfo_.key.desc == nullptr, OPS_LOG_E(opName_, "Desc of tensor k is nullptr"),
36+ return ge::GRAPH_FAILED);
37+ 
38+ OPS_ERR_IF(opParamInfo_.query_dequant_scale.shape == nullptr,
39+ OPS_LOG_E(opName_, "Shape of tensor query_dequant_scale is nullptr"), return ge::GRAPH_FAILED);
40+ OPS_ERR_IF(opParamInfo_.query_dequant_scale.desc == nullptr,
41+ OPS_LOG_E(opName_, "Desc of tensor query_dequant_scale is nullptr"), return ge::GRAPH_FAILED);
42+ OPS_ERR_IF(opParamInfo_.key_dequant_scale.shape == nullptr,
43+ OPS_LOG_E(opName_, "Shape of tensor key_dequant_scale is nullptr"), return ge::GRAPH_FAILED);
44+ OPS_ERR_IF(opParamInfo_.key_dequant_scale.desc == nullptr,
45+ OPS_LOG_E(opName_, "Desc of tensor key_dequant_scale is nullptr"), return ge::GRAPH_FAILED);
46+ 
47+ OPS_ERR_IF(opParamInfo_.sparseIndices.shape == nullptr, OPS_LOG_E(opName_, "Shape of tensor output is nullptr"),
48+ return ge::GRAPH_FAILED);
49+ OPS_ERR_IF(opParamInfo_.sparseIndices.desc == nullptr, OPS_LOG_E(opName_, "Desc of tensor output is nullptr"),
50+ return ge::GRAPH_FAILED);
51+ return ge::GRAPH_SUCCESS;
52+}
53+ 
54+ge::graphStatus QSIInfoParser::CheckRequiredAttrExistence() const
55+{
56+ OPS_ERR_IF(opParamInfo_.layOutKey == nullptr, OPS_LOG_E(opName_, "attr layout_key is nullptr"),
57+ return ge::GRAPH_FAILED);
58+ 
59+ return ge::GRAPH_SUCCESS;
60+}
61+ 
62+ge::graphStatus QSIInfoParser::CheckRequiredParaExistence() const
63+{
64+ if (CheckRequiredInOutExistence() != ge::GRAPH_SUCCESS || CheckRequiredAttrExistence() != ge::GRAPH_SUCCESS) {
65+ return ge::GRAPH_FAILED;
66+ }
67+ 
68+ return ge::GRAPH_SUCCESS;
69+}
70+ 
71+ge::graphStatus QSIInfoParser::GetOpName()
72+{
73+ if (context_->GetNodeName() == nullptr) {
74+ OPS_LOG_E("SalsIndexer", "opName got from TilingContext is nullptr");
75+ return ge::GRAPH_FAILED;
76+ }
77+ opName_ = context_->GetNodeName();
78+ return ge::GRAPH_SUCCESS;
79+}
80+ 
81+ge::graphStatus QSIInfoParser::GetNpuInfo()
82+{
83+ platformInfo_ = context_->GetPlatformInfo();
84+ OPS_ERR_IF(platformInfo_ == nullptr, OPS_LOG_E(opName_, "GetPlatformInfo is nullptr."), return ge::GRAPH_FAILED);
85+ 
86+ auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfo_);
87+ uint32_t aivNum = ascendcPlatform.GetCoreNumAiv();
88+ uint32_t aicNum = ascendcPlatform.GetCoreNumAic();
89+ OPS_ERR_IF(aicNum == 0 || aivNum == 0, OPS_LOG_E(opName_, "num of core obtained is 0."), return GRAPH_FAILED);
90+ OPS_ERR_IF(aicNum != 24, OPS_LOG_E(opName_, "num of core only support 24."), return GRAPH_FAILED);
91+ 
92+ OPS_ERR_IF(context_->GetWorkspaceSizes(1) == nullptr, OPS_LOG_E(opName_, "workSpaceSize got from ge is nullptr"),
93+ return ge::GRAPH_FAILED);
94+ OPS_ERR_IF(context_->GetRawTilingData() == nullptr,
95+ OPS_LOG_E(context_->GetNodeName(), "RawTilingData got from GE context is nullptr."),
96+ return ge::GRAPH_FAILED);
97+ 
98+ return ge::GRAPH_SUCCESS;
99+}
100+ 
101+void QSIInfoParser::GetOptionalInputParaInfo()
102+{
103+ opParamInfo_.actualSeqLengths.tensor = context_->GetOptionalInputTensor(ACTUAL_SEQ_K_INDEX);
104+ opParamInfo_.actualSeqLengths.desc = context_->GetOptionalInputDesc(ACTUAL_SEQ_K_INDEX);
105+ opParamInfo_.blockTable.tensor = context_->GetOptionalInputTensor(BLOCK_TABLE_INDEX);
106+ opParamInfo_.blockTable.desc = context_->GetOptionalInputDesc(BLOCK_TABLE_INDEX);
107+}
108+ 
109+void QSIInfoParser::GetInputParaInfo()
110+{
111+ opParamInfo_.query.desc = context_->GetInputDesc(QUERY_INDEX);
112+ opParamInfo_.query.shape = context_->GetInputShape(QUERY_INDEX);
113+ opParamInfo_.key.desc = context_->GetInputDesc(KEY_INDEX);
114+ opParamInfo_.key.shape = context_->GetInputShape(KEY_INDEX);
115+ opParamInfo_.query_dequant_scale.desc = context_->GetInputDesc(QUERY_DEQUANT_SCALE_INDEX);
116+ opParamInfo_.query_dequant_scale.shape = context_->GetInputShape(QUERY_DEQUANT_SCALE_INDEX);
117+ opParamInfo_.key_dequant_scale.desc = context_->GetInputDesc(KEY_DEQUANT_SCALE_INDEX);
118+ opParamInfo_.key_dequant_scale.shape = context_->GetInputShape(KEY_DEQUANT_SCALE_INDEX);
119+ 
120+ GetOptionalInputParaInfo();
121+}
122+ 
123+void QSIInfoParser::GetOutputParaInfo()
124+{
125+ opParamInfo_.sparseIndices.desc = context_->GetOutputDesc(SPARSE_INDICES_INDEXER);
126+ opParamInfo_.sparseIndices.shape = context_->GetOutputShape(SPARSE_INDICES_INDEXER);
127+}
128+ 
129+ge::graphStatus QSIInfoParser::GetAndCheckAttrParaInfo()
130+{
131+ auto attrs = context_->GetAttrs();
132+ OPS_ERR_IF(attrs == nullptr, OPS_REPORT_VECTOR_INNER_ERR(context_->GetNodeName(), "attrs got from ge is nullptr"),
133+ return ge::GRAPH_FAILED);
134+ 
135+ OPS_LOG_I(context_->GetNodeName(), "GetAndCheckAttrParaInfo start");
136+ opParamInfo_.maxSeqlenKey = attrs->GetAttrPointer<int32_t>(ATTR_MAX_SEQLEN_KEY_INDEX);
137+ OPS_ERR_IF(opParamInfo_.maxSeqlenKey == nullptr, OPS_LOG_E(opName_, "input attr max_seqlen_key must not be null"),
138+ return ge::GRAPH_FAILED);
139+ 
140+ opParamInfo_.sparseBlockSize = attrs->GetAttrPointer<int32_t>(ATTR_SPARSE_BLOCK_SIZE_INDEX);
141+ OPS_ERR_IF(opParamInfo_.sparseBlockSize == nullptr,
142+ OPS_LOG_E(opName_, "input attr sparse_block_size must not be null"), return ge::GRAPH_FAILED);
143+ 
144+ opParamInfo_.sparseRatio = attrs->GetAttrPointer<float>(ATTR_SPARSE_RATIO_INDEX);
145+ OPS_ERR_IF(opParamInfo_.sparseRatio == nullptr, OPS_LOG_E(opName_, "input attr sparse_ratio must not be null"),
146+ return ge::GRAPH_FAILED);
147+ 
148+ opParamInfo_.fixedTailCount = attrs->GetAttrPointer<int32_t>(ATTR_FIXED_TAIL_COUNT_INDEX);
149+ OPS_ERR_IF(opParamInfo_.fixedTailCount == nullptr,
150+ OPS_LOG_E(opName_, "input attr fixed_tail_count must not be null"), return ge::GRAPH_FAILED);
151+ 
152+ OPS_ERR_IF(
153+ (*opParamInfo_.sparseBlockSize) <= 0 || (*opParamInfo_.sparseRatio) <= 0 ||
154+ (*opParamInfo_.maxSeqlenKey) <= 0 || (*opParamInfo_.fixedTailCount) <= 0,
155+ OPS_LOG_E(opName_,
156+ "input attr max_seqlen_key, sparse_block_size, sparse_ratio,fixed_tail_count must greater than "
157+ "0, but max_seqlen_key: %d, sparse_block_size: %d, sparse_ratio: %f, fixed_tail_count: %d.",
158+ (*opParamInfo_.maxSeqlenKey), (*opParamInfo_.sparseBlockSize), (*opParamInfo_.sparseRatio),
159+ (*opParamInfo_.fixedTailCount)),
160+ return ge::GRAPH_FAILED);
161+
162+ OPS_ERR_IF(
163+ (*opParamInfo_.sparseBlockSize) != 16,
164+ OPS_LOG_E(opName_,
165+ "input attr sparse_block_size, must equal %d"
166+ ", but sparse_block_size: %d.",
167+ 16, (*opParamInfo_.sparseBlockSize)),
168+ return ge::GRAPH_FAILED);
169+ 
170+ opParamInfo_.layOutKey = attrs->GetStr(ATTR_KEY_LAYOUT_INDEX);
171+ if (opParamInfo_.layOutKey != nullptr) {
172+ OPS_LOG_I(context_->GetNodeName(), "layout_key is:%s", opParamInfo_.layOutKey);
173+ }
174+ 
175+ OPS_ERR_IF(
Y
YYuZhengzhong4月2日

GetAndCheckAttrParaInfo() 限制 layout_key 只支持 PA_BNSD 和 PA_NZ，但 def.cpp 默认值为 "BSND"，layoutMap 和模板参数中也定义了 BSND/PA_BSND。tiling 却直接拒绝这些合法取值，接口定义和实际校验不一致。

likedislike
M
mayafei
4月3日 评论：
176+ ((std::string(opParamInfo_.layOutKey) != "PA_BNSD") && (std::string(opParamInfo_.layOutKey) != "PA_NZ")),
177+ OPS_LOG_E(opName_, "input attr layout_key only supported PA_BNSD, PA_NZ"), return ge::GRAPH_FAILED);
178+ 
179+ OPS_LOG_I(context_->GetNodeName(), "GetAndCheckAttrParaInfo end");
180+ return ge::GRAPH_SUCCESS;
181+}
182+ 
183+ge::graphStatus QSIInfoParser::GetOpParaInfo()
184+{
185+ GetInputParaInfo();
186+ GetOutputParaInfo();
187+ if (ge::GRAPH_SUCCESS != GetAndCheckAttrParaInfo()) {
188+ return ge::GRAPH_FAILED;
189+ }
190+ return ge::GRAPH_SUCCESS;
191+}
192+ 
193+ge::graphStatus QSIInfoParser::GetAndCheckInOutDataType()
194+{
195+ inputQType_ = opParamInfo_.query.desc->GetDataType();
196+ inputQType_ = (inputQType_== ge::DT_INT32) ? ge::DT_INT4 : ge::DT_INT8;
197+ inputKType_ = opParamInfo_.key.desc->GetDataType();
198+ inputKType_ = (inputKType_== ge::DT_INT32) ? ge::DT_INT4 : ge::DT_INT8;
199+ inputQueryScaleType_ = opParamInfo_.query_dequant_scale.desc->GetDataType();
200+ inputKeyScaleType_ = opParamInfo_.key_dequant_scale.desc->GetDataType();
201+ 
202+ sparseIndicesType_ = opParamInfo_.sparseIndices.desc->GetDataType();
203+ 
204+ bool inDTypeAllEqual = (inputQType_ == inputKType_);
205+ OPS_ERR_IF(!inDTypeAllEqual,
206+ OPS_LOG_E(opName_, "The data types of the input query and key must be the same."),
207+ return ge::GRAPH_FAILED);
208+ 
209+ OPS_ERR_IF((inputQType_ !=ge::DT_INT4),
210+ OPS_LOG_E(opName_, "The data types of the input query and key must be int4."),
211+ return ge::GRAPH_FAILED);
212+ 
213+ OPS_ERR_IF(!(inputQueryScaleType_ == inputKeyScaleType_),
214+ OPS_LOG_E(opName_, "The data types of the input query_dequant_scale and key_dequant_scale must be the same."),
215+ return ge::GRAPH_FAILED);
216+ 
217+ OPS_ERR_IF((inputQueryScaleType_ != ge::DT_FLOAT),
218+ OPS_LOG_E(opName_, "The data types of the input query_dequant_scale and key_dequant_scale must be float32."),
219+ return ge::GRAPH_FAILED);
220+ 
221+ OPS_ERR_IF(sparseIndicesType_ != ge::DT_INT32,
222+ OPS_LOG_E(opName_, "The data types of the output sparse_indices must be int32."),
223+ return ge::GRAPH_FAILED);
224+ 
225+ return ge::GRAPH_SUCCESS;
226+}
227+ 
228+ge::graphStatus QSIInfoParser::GetQueryKeyAndOutLayout()
229+{
230+ // 获取query,key的Layout基准值
231+ const map<string, DataLayout> layoutMap = {
232+ {"BSND", DataLayout::BSND},
233+ {"PA_BNSD", DataLayout::PA_BNSD},
234+ {"PA_BSND", DataLayout::PA_BSND},
235+ {"PA_NZ", DataLayout::PA_NZ},
236+ };
237+ 
238+ std::string layoutKey(opParamInfo_.layOutKey);
239+ auto itKey = layoutMap.find(layoutKey);
240+ if (itKey != layoutMap.end()) {
241+ kLayout_ = itKey->second;
242+ }
243+ 
244+ return ge::GRAPH_SUCCESS;
245+}
246+ 
247+ge::graphStatus QSIInfoParser::GetAndCheckOptionalInput()
248+{
249+ if (kLayout_ == DataLayout::PA_BNSD || kLayout_ == DataLayout::PA_NZ) {
250+ OPS_ERR_IF(opParamInfo_.blockTable.tensor == nullptr,
251+ OPS_LOG_E(opName_, "key layout supported PA_BNSD or PA_NZ, input block_table must not be null"),
252+ return ge::GRAPH_FAILED);
253+ OPS_ERR_IF(
254+ opParamInfo_.actualSeqLengths.tensor == nullptr,
255+ OPS_LOG_E(opName_, "key layout supported PA_BNSD or PA_NZ, input actual_seq_lengths_key must not be null"),
256+ return ge::GRAPH_FAILED);
257+ OPS_ERR_IF(opParamInfo_.blockTable.desc->GetDataType() != ge::DT_INT32,
258+ OPS_LOG_E(opName_, "input block_table data type only support int32"), return ge::GRAPH_FAILED);
259+ }
260+ OPS_ERR_IF(opParamInfo_.actualSeqLengths.tensor != nullptr &&
Y
YYuZhengzhong4月2日

GetAndCheckOptionalInput() 中对 actualSeqLengths 数据类型的检查完全重复了两次（第 259-262 行和第 263-266 行），应删除重复的那段。

likedislike
M
mayafei
4月3日 评论：
261+ opParamInfo_.actualSeqLengths.desc->GetDataType() != ge::DT_INT32,
262+ OPS_LOG_E(opName_, "input actual_seq_lengths_key data type only support int32"),
263+ return ge::GRAPH_FAILED);
264+ OPS_ERR_IF(kLayout_ != DataLayout::PA_NZ && kLayout_ != DataLayout::PA_BNSD && opParamInfo_.blockTable.tensor != nullptr,
265+ OPS_LOG_E(opName_, "when key layout is not PA_BNSD or PA_NZ, input block_table must be null"),
266+ return ge::GRAPH_FAILED);
267+ return ge::GRAPH_SUCCESS;
268+}
269+ 
270+ge::graphStatus QSIInfoParser::CheckShapeDim()
271+{
272+ OPS_ERR_IF((opParamInfo_.blockTable.tensor != nullptr) &&
273+ (opParamInfo_.blockTable.tensor->GetStorageShape().GetDimNum() != DIM_NUM_TWO),
274+ OPS_LOG_E(opName_, "the dim num of block_table's shape should be 2"), return ge::GRAPH_FAILED);
275+ 
276+ uint32_t kShapeDim = opParamInfo_.key.shape->GetStorageShape().GetDimNum();
277+ uint32_t qShapeDim = opParamInfo_.query.shape->GetStorageShape().GetDimNum();
278+ 
279+ uint32_t qExpectShapeDim = DIM_NUM_THREE;
280+ uint32_t kExpectShapeDim = DIM_NUM_FOUR;
281+ if (kLayout_ == DataLayout::PA_NZ) {
282+ kExpectShapeDim = DIM_NUM_FIVE;
283+ }
284+ 
285+ OPS_ERR_IF(kShapeDim != kExpectShapeDim,
286+ OPS_LOG_E(opName_, "the dim num of key's shape should be %u, but now is %u", kExpectShapeDim, kShapeDim),
287+ return ge::GRAPH_FAILED);
288+ OPS_ERR_IF(qShapeDim != qExpectShapeDim,
289+ OPS_LOG_E(opName_, "the dim num of query's shape should be %u, but now is %u",
290+ qExpectShapeDim, qShapeDim),
291+ return ge::GRAPH_FAILED);
292+
293+ uint32_t sparseIndicesShapeDim = opParamInfo_.sparseIndices.shape->GetStorageShape().GetDimNum();
294+ 
295+ OPS_ERR_IF(sparseIndicesShapeDim != qExpectShapeDim,
296+ OPS_LOG_E(opName_, "the dim num of sparse_indices's shape should be %u, but now is %u",
297+ qExpectShapeDim, sparseIndicesShapeDim),
298+ return ge::GRAPH_FAILED);
299+ 
300+ return ge::GRAPH_SUCCESS;
301+}
302+ 
303+ge::graphStatus QSIInfoParser::GetActualSeqLenSize(uint32_t &size, const gert::Tensor *tensor,
304+ const std::string &actualSeqLenName)
305+{
306+ size = static_cast<uint32_t>(tensor->GetShapeSize());
307+ if (size <= 0) {
308+ OPS_LOG_E(opName_, "%s's shape size is %u, it should be greater than 0.", actualSeqLenName.c_str(), size);
309+ return ge::GRAPH_FAILED;
310+ }
311+ return ge::GRAPH_SUCCESS;
312+}
313+ 
314+ge::graphStatus QSIInfoParser::GetAndCheckN2Size()
315+{
316+ uint32_t n2Index = DIM_IDX_TWO; // PA_BSND/BSND
317+ if (kLayout_ == DataLayout::PA_BNSD || kLayout_ == DataLayout::PA_NZ) { // PA_BNSD
318+ n2Index = DIM_IDX_ONE;
319+ }
320+ n2Size_ = static_cast<uint32_t>(opParamInfo_.key.shape->GetStorageShape().GetDim(n2Index));
321+ OPS_LOG_I(context_->GetNodeName(), "n2Size_ is %d", n2Size_);
322+ return ge::GRAPH_SUCCESS;
323+}
324+ 
325+ge::graphStatus QSIInfoParser::GetBatchSize()
326+{
327+ bSize_ = opParamInfo_.query.shape->GetStorageShape().GetDim(0);
328+ return ge::GRAPH_SUCCESS;
329+}
330+ 
331+ge::graphStatus QSIInfoParser::GetHeadDim()
332+{
333+ // 以query的D维度为基准
334+ headDim_ = opParamInfo_.query.shape->GetStorageShape().GetDim(DIM_IDX_TWO) * 8; // INT4伪装成IN32, D轴每8个数合成一个数，*8还原shape
335+ OPS_ERR_IF(headDim_ != 64,
336+ OPS_LOG_E(opName_, "the head dim of query and key should be %u, but now is %u",
337+ 64, headDim_),
338+ return ge::GRAPH_FAILED);
339+ return ge::GRAPH_SUCCESS;
340+}
341+ 
342+ge::graphStatus QSIInfoParser::GetSparseCount()
343+{
344+ sparseCount_ = opParamInfo_.sparseIndices.shape->GetStorageShape().GetDim(DIM_IDX_TWO);
345+ return ge::GRAPH_SUCCESS;
346+}
347+ 
348+ge::graphStatus QSIInfoParser::GetAndCheckBlockSize()
349+{
350+ if (kLayout_ == DataLayout::PA_BSND) {
351+ blockSize_ = static_cast<uint32_t>(opParamInfo_.key.shape->GetStorageShape().GetDim(1));
352+ } else if (kLayout_ == DataLayout::PA_BNSD) { // PA_BNSD
353+ blockSize_ = static_cast<uint32_t>(opParamInfo_.key.shape->GetStorageShape().GetDim(2));
354+ } else if (kLayout_ == DataLayout::PA_NZ) {
355+ blockSize_ = static_cast<uint32_t>(opParamInfo_.key.shape->GetStorageShape().GetDim(3));
356+ }
357+ OPS_LOG_I(context_->GetNodeName(), "blockSize_ is %d", blockSize_);
358+ 
359+ OPS_ERR_IF(((blockSize_ % 16 != 0) || (blockSize_ == 0) || (blockSize_ > 1024)),
360+ OPS_LOG_E(opName_, "input key's block_size must be a multiple of 16 and belong to (0, 1024]."),
361+ return ge::GRAPH_FAILED);
362+ 
363+ return ge::GRAPH_SUCCESS;
364+}
365+ 
366+ge::graphStatus QSIInfoParser::CheckBlockCount()
367+{
368+ int32_t blockCount_ = static_cast<uint32_t>(opParamInfo_.key.shape->GetStorageShape().GetDim(0));
369+ OPS_ERR_IF((blockCount_ == 0),
370+ OPS_LOG_E(opName_, "input key's block_count cannot be 0."),
371+ return ge::GRAPH_FAILED);
372+ return ge::GRAPH_SUCCESS;
373+}
374+ 
375+ge::graphStatus QSIInfoParser::GetS2SizeForPageAttention()
376+{
377+ if (GetAndCheckBlockSize() != ge::GRAPH_SUCCESS) {
378+ return ge::GRAPH_FAILED;
379+ }
380+ if (CheckBlockCount() != ge::GRAPH_SUCCESS) {
381+ return ge::GRAPH_FAILED;
382+ }
383+ maxBlockNumPerBatch_ = opParamInfo_.blockTable.tensor->GetStorageShape().GetDim(1);
384+ s2Size_ = maxBlockNumPerBatch_ * blockSize_;
385+ OPS_LOG_I(context_->GetNodeName(), "maxBlockNumPerBatch_ is %d, blockSize_ is %d, s2Size_ is %d",
386+ maxBlockNumPerBatch_, blockSize_, s2Size_);
387+ return ge::GRAPH_SUCCESS;
388+}
389+ 
390+ge::graphStatus QSIInfoParser::GetS2Size()
391+{
392+ // 获取S2基准值
393+ // 1、BATCH_CONTINUOUS时, 从key的S轴获取
394+ // 3、PAGE_ATTENTION时, S2 = block_table.dim1 * block_size
395+ if (kLayout_ == DataLayout::PA_BNSD || kLayout_ == DataLayout::PA_BSND || kLayout_ == DataLayout::PA_NZ) {
396+ return GetS2SizeForPageAttention();
397+ } else if (kLayout_ == DataLayout::BSND) {
398+ s2Size_ = opParamInfo_.key.shape->GetStorageShape().GetDim(1);
399+ }
400+ return ge::GRAPH_SUCCESS;
401+}
402+ 
403+ge::graphStatus QSIInfoParser::ValidateInputShapesMatch()
404+{
405+ /*
406+ BSND:
407+ query [BatchSize,S1,N1,D],
408+ key [BlockNum,BlockSize,N2,D],
409+ block_table [BatchSize, BatchMaxBlockNum],
410+ act_seq_k [BatchSize]
411+ act_seq_q [BatchSize] 可选
412+ out [BatchSize,S1,N2,topk]
413+ */
414+ uint32_t outN2Dim = DIM_IDX_ONE;
415+ // -----------------------check BatchSize-------------------
416+ // bSize_ 来源于query
417+ OPS_ERR_IF(
418+ ((opParamInfo_.blockTable.tensor != nullptr) &&
419+ (opParamInfo_.blockTable.tensor->GetStorageShape().GetDim(0) != bSize_)) ||
420+ (opParamInfo_.actualSeqLengths.tensor->GetShapeSize() != bSize_) ||
421+ (opParamInfo_.sparseIndices.shape->GetStorageShape().GetDim(0) != bSize_),
422+ OPS_LOG_E(opName_, "BSND case input query, actual_seq_lengths_key, block_table, sparse_indices dim 0 must be same."),
423+ return ge::GRAPH_FAILED);
424+ 
425+ // -----------------------check D-------------------
426+ uint32_t keyDDim = DIM_IDX_THREE;
427+ if (kLayout_ == DataLayout::PA_NZ) {
428+ keyDDim = DIM_IDX_TWO;
429+ OPS_ERR_IF((opParamInfo_.key.shape->GetStorageShape().GetDim(keyDDim) * 8 * 8 != headDim_),
430+ OPS_LOG_E(opName_, "input query, key shape last dim must be same."), return ge::GRAPH_FAILED); // INT4伪装成IN32, D轴每8个数合成一个数，*8还原shape
431+ } else {
432+ OPS_ERR_IF((opParamInfo_.key.shape->GetStorageShape().GetDim(keyDDim) * 8 != headDim_),
433+ OPS_LOG_E(opName_, "input query, key shape last dim must be same."), return ge::GRAPH_FAILED); // INT4伪装成IN32, D轴每8个数合成一个数，*8还原shape
434+ }
435+ // -----------------------check N2-------------------
436+ OPS_ERR_IF((opParamInfo_.sparseIndices.shape->GetStorageShape().GetDim(outN2Dim) != n2Size_),
437+ OPS_LOG_E(opName_, "input query and output sparse_indices shape n2 dim must be same."),
438+ return ge::GRAPH_FAILED);
439+ return ge::GRAPH_SUCCESS;
440+}
441+ 
442+ge::graphStatus QSIInfoParser::CheckScaleShape()
443+{
444+ uint32_t qShapeDim = opParamInfo_.query.shape->GetStorageShape().GetDimNum();
445+ uint32_t kShapeDim = opParamInfo_.key.shape->GetStorageShape().GetDimNum();
446+ uint32_t qDequantScaleShapeDim = opParamInfo_.query_dequant_scale.shape->GetStorageShape().GetDimNum();
447+ uint32_t kDequantScaleShapeDim = opParamInfo_.key_dequant_scale.shape->GetStorageShape().GetDimNum();
448+ OPS_ERR_IF(qDequantScaleShapeDim != (qShapeDim - 1),
449+ OPS_LOG_E(opName_, "the dim num of query_dequant_scale's shape should be %u, but now is %u",
450+ qShapeDim - 1, qDequantScaleShapeDim),
451+ return ge::GRAPH_FAILED);
452+ if (kLayout_ == DataLayout::PA_NZ) {
453+ OPS_ERR_IF(kDequantScaleShapeDim != (kShapeDim - 2),
454+ OPS_LOG_E(opName_, "the dim num of key_dequant_scale's shape should be %u, but now is %u", kShapeDim - 1,
455+ kDequantScaleShapeDim),
456+ return ge::GRAPH_FAILED);
457+ } else {
458+ OPS_ERR_IF(kDequantScaleShapeDim != (kShapeDim - 1),
459+ OPS_LOG_E(opName_, "the dim num of key_dequant_scale's shape should be %u, but now is %u", kShapeDim - 1,
460+ kDequantScaleShapeDim),
461+ return ge::GRAPH_FAILED);
462+ }
463+ // check q scale
464+ for (uint32_t i = 0; i < (qShapeDim - 2); i++) {
465+ uint32_t dimValueQueryScale = opParamInfo_.query_dequant_scale.shape->GetStorageShape().GetDim(i);
466+ uint32_t dimValueQuery = opParamInfo_.query.shape->GetStorageShape().GetDim(i);
467+ OPS_ERR_IF(dimValueQueryScale != dimValueQuery,
468+ OPS_LOG_E(opName_, "query_dequant_scale's shape[%u] %u and query's shape[%u] %u is not same", i,
469+ dimValueQueryScale, i, dimValueQuery),
470+ return ge::GRAPH_FAILED);
471+ }
472+ return ge::GRAPH_SUCCESS;
473+}
474+ 
475+void QSIInfoParser::GenerateInfo(QSITilingInfo &siInfo)
476+{
477+ siInfo.opName = opName_;
478+ siInfo.platformInfo = platformInfo_;
479+ siInfo.opParamInfo = opParamInfo_;
480+ siInfo.socVersion = socVersion_;
481+ 
482+ siInfo.bSize = bSize_;
483+ siInfo.dSize = headDim_;
484+ siInfo.n2Size = n2Size_;
485+ siInfo.s2Size = s2Size_;
486+ siInfo.sparseCount = sparseCount_;
487+ siInfo.maxSeqlenKey = *opParamInfo_.maxSeqlenKey;
488+ siInfo.sparseBlockSize = *opParamInfo_.sparseBlockSize;
489+ siInfo.fixedTailCount = *opParamInfo_.fixedTailCount;
490+ siInfo.sparseRatio = (int)((1.0 - *opParamInfo_.sparseRatio) * 100.0 + 0.5) / 100.0;
491+ 
492+ siInfo.inputQType = inputQType_;
493+ siInfo.inputKType = inputKType_;
494+ siInfo.outputType = sparseIndicesType_;
495+ 
496+ siInfo.blockSize = blockSize_;
497+ siInfo.maxBlockNumPerBatch = maxBlockNumPerBatch_;
498+ 
499+ std::string layOutKeyStr(opParamInfo_.layOutKey);
500+ siInfo.pageAttentionFlag = (layOutKeyStr == "PA_BNSD" || layOutKeyStr == "PA_BSND" || layOutKeyStr == "PA_NZ");
501+ 
502+ siInfo.inputKLayout = kLayout_;
503+}
504+ 
505+ge::graphStatus QSIInfoParser::ParseAndCheck(QSITilingInfo &siInfo)
506+{
507+ if (ge::GRAPH_SUCCESS != GetOpName() || ge::GRAPH_SUCCESS != GetNpuInfo() || ge::GRAPH_SUCCESS != GetOpParaInfo() ||
508+ ge::GRAPH_SUCCESS != CheckRequiredParaExistence()) {
509+ return ge::GRAPH_FAILED;
510+ }
511+ 
512+ if (ge::GRAPH_SUCCESS != GetAndCheckInOutDataType() || ge::GRAPH_SUCCESS != GetQueryKeyAndOutLayout() ||
513+ ge::GRAPH_SUCCESS != GetAndCheckOptionalInput()) {
514+ return ge::GRAPH_FAILED;
515+ }
516+ 
517+ if (ge::GRAPH_SUCCESS != CheckShapeDim() ||
518+ ge::GRAPH_SUCCESS != GetAndCheckN2Size()) {
519+ return ge::GRAPH_FAILED;
520+ }
521+ 
522+ if (ge::GRAPH_SUCCESS != GetBatchSize() || ge::GRAPH_SUCCESS != GetHeadDim() ||
523+ ge::GRAPH_SUCCESS != GetS2Size() || ge::GRAPH_SUCCESS != GetSparseCount()) {
524+ return ge::GRAPH_FAILED;
525+ }
526+ if (ge::GRAPH_SUCCESS != ValidateInputShapesMatch() || ge::GRAPH_SUCCESS != CheckScaleShape()) {
527+ return ge::GRAPH_FAILED;
528+ }
529+ 
530+ GenerateInfo(siInfo);
531+ 
532+ return ge::GRAPH_SUCCESS;
533+}
534+ 
535+void QuantSalsIndexerTiling::SplitCoreBN(uint32_t coreNum, QSITilingInfo *siInfo, SplitParams splitParams) {
536+ std::vector<uint32_t> s1GBaseNum(siInfo->bSize);
537+ std::vector<uint32_t> s2BaseNum(siInfo->bSize);
538+ std::vector<uint32_t> s1Size(siInfo->bSize);
539+ std::vector<uint32_t> s2Size(siInfo->bSize);
540+ uint32_t s2BaseSize = 2048;
541+ // 计算总基本块数
542+ uint32_t totalBaseNum = 0;
543+ for (uint32_t bIdx = 0; bIdx < siInfo->bSize; bIdx++) {
544+ s2Size[bIdx] = siInfo->s2Size;
545+ s1Size[bIdx] = 1;
546+ s1GBaseNum[bIdx] = 1;
547+ s2BaseNum[bIdx] = 1;
548+ totalBaseNum += s2BaseNum[bIdx] * siInfo->n2Size;
549+ }
550+ uint32_t avgBaseNum = 1;
551+ if (totalBaseNum > coreNum) {
552+ avgBaseNum = (totalBaseNum + coreNum - 1) / coreNum;
553+ }
554+ 
555+ uint32_t accumBaseNum = 0; // 当前累计的基本块数
556+ uint32_t targetBaseNum = 0;
557+ uint32_t currCoreIdx = 0;
558+ uint32_t lastValidBIdx = 0;
559+ // 分核流程，保存分核数据
560+ for (uint32_t bN2Idx = 0; bN2Idx < siInfo->bSize * siInfo->n2Size; bN2Idx++) {
561+ uint32_t bIdx = bN2Idx / siInfo->n2Size;
562+ for (uint32_t s1GIdx = 0; s1GIdx < s1GBaseNum[bIdx]; s1GIdx++) {
563+ uint32_t sInnerIndexStart = 0;
564+ uint32_t sInnerIndexEnd = s2BaseNum[bIdx];
565+ accumBaseNum = accumBaseNum + (sInnerIndexEnd - sInnerIndexStart);
566+ targetBaseNum = (currCoreIdx + 1) * avgBaseNum;
567+ if (accumBaseNum >= targetBaseNum) {
568+ // 更新当前核的End分核信息
569+ splitParams.bN2End[currCoreIdx] = bN2Idx;
570+ splitParams.gS1End[currCoreIdx] = s1GIdx;
571+ splitParams.s2End[currCoreIdx] = sInnerIndexEnd - 1;
572+ currCoreIdx += 1;
573+ }
574+ }
575+ if ((s1GBaseNum[bIdx] > 0) && (s2BaseNum[bIdx] > 0)) {
576+ lastValidBIdx = bIdx;
577+ }
578+ }
579+ if (accumBaseNum < targetBaseNum) {
580+ // 更新最后一个核的End分核信息
581+ splitParams.bN2End[currCoreIdx] = ((lastValidBIdx + 1) * (siInfo->n2Size)) - 1;
582+ splitParams.gS1End[currCoreIdx] = s1GBaseNum[lastValidBIdx] - 1;
583+ splitParams.s2End[currCoreIdx] = s2BaseNum[lastValidBIdx] - 1;
584+ currCoreIdx += 1;
585+ }
586+}
587+ 
588+// --------------------------TilingPrepare函数定义-------------------------------------
589+static ge::graphStatus TilingPrepareForQuantSalsIndexer(gert::TilingParseContext * /* context */)
590+{
591+ return ge::GRAPH_SUCCESS;
592+}
593+ 
594+// --------------------------QuantSalsIndexerTiling类成员函数定义-----------------------
595+ge::graphStatus QuantSalsIndexerTiling::DoTiling(QSITilingInfo *tilingInfo)
596+{
597+ // -------------set blockdim-----------------
598+ auto ascendcPlatform = platform_ascendc::PlatformAscendC(tilingInfo->platformInfo);
599+ uint32_t aivNum = ascendcPlatform.GetCoreNumAiv();
600+ uint32_t aicNum = ascendcPlatform.GetCoreNumAic();
601+ uint32_t blockDim = ascendcPlatform.CalcTschBlockDim(aivNum, aicNum, aivNum);
602+ context_->SetBlockDim(blockDim);
603+ 
604+ // ------------------------------------------------------------
605+ if (context_->GetOptionalInputDesc(METADATA_INDEX) == nullptr) {
606+ SplitParams splitParams;
607+ splitParams.bN2End = tilingData_.splitParams.get_bN2End();
608+ splitParams.gS1End = tilingData_.splitParams.get_gS1End();
609+ splitParams.s2End = tilingData_.splitParams.get_s2End();
610+ SplitCoreBN(aicNum, tilingInfo, splitParams);
611+ tilingData_.set_usedCoreNum(blockDim);
Y
YYuZhengzhong4月2日

这里把 blockDim（全部核数）设为 usedCoreNum，但 SplitCoreBN 可能只使用了部分核。当 metadata 不存在时，kernel 中 usedCoreNum 将大于实际有任务的核数，空闲核不会正确退出。

likedislike
M
mayafei
4月3日 评论：
612+ }
613+
614+ // ------------------------------------------------------------
615+ 
616+ 
617+ // -------------set workspacesize-----------------
618+ constexpr uint32_t MM1_RES_ELEM_SIZE = 4; // 4: fp32
619+ constexpr uint32_t DOUBLE_BUFFER = 2; // 双Buffer
620+ constexpr uint32_t M_BASE_SIZE = 512; // m轴基本块大小
621+ constexpr uint32_t S2_BASE_SIZE = 512; // S2轴基本块大小
622+ constexpr uint32_t V1_RES_ELEM_SIZE = 4; // 4: int32
623+ constexpr uint32_t V1_RES_ELEM_TYPE = 2; // 保留Index和Value 2种数据
624+ constexpr uint32_t V1_DECODE_PARAM_ELEM_SIZE = 8; // 8: int64
625+ constexpr uint32_t V1_DECODE_PARAM_NUM = 16; // Decode参数个数
626+ constexpr uint32_t V1_DECODE_DATA_NUM = 2; // Decode每个核需要存储头和尾部两块数据
627+ constexpr uint32_t S1_BASE_SIZE = 8; // S1轴基本块的大小
628+ constexpr uint32_t TOPK_MAX_SIZE = 2048; // TopK选取个数
629+ uint32_t workspaceSize = ascendcPlatform.GetLibApiWorkSpaceSize();
630+ // 主流程需Workspace大小
631+ uint32_t mm1ResSize = M_BASE_SIZE * S2_BASE_SIZE;
632+ workspaceSize += mm1ResSize * MM1_RES_ELEM_SIZE * DOUBLE_BUFFER * aicNum;
633+ // Decode流程(LD)需要Workspace大小
634+ // 临时存储Decode中间结果大小: 2(头/尾)*8(s1Base)*2(idx/value)*2048(K)*sizeof(int32)*24=6M
635+ workspaceSize += V1_DECODE_DATA_NUM * S1_BASE_SIZE * V1_RES_ELEM_TYPE * (TOPK_MAX_SIZE + tilingInfo->fixedTailCount) * V1_RES_ELEM_SIZE * aicNum;
636+ // 临时存储Decode中间参数信息大小: 2(头/尾)*8(s1Base)*16(paramNum)*sizeof(int64_t)*24=48k
637+ workspaceSize += V1_DECODE_DATA_NUM * S1_BASE_SIZE * V1_DECODE_PARAM_NUM * V1_DECODE_PARAM_ELEM_SIZE * aicNum;
638+ size_t *workSpaces = context_->GetWorkspaceSizes(1);
639+ workSpaces[0] = workspaceSize;
640+ 
641+ // -------------set tilingdata-----------------
642+ tilingData_.set_bSize(tilingInfo->bSize);
643+ tilingData_.set_s2Size(tilingInfo->s2Size);
644+ tilingData_.set_n2Size(tilingInfo->n2Size);
645+ tilingData_.set_fixedTailCount(tilingInfo->fixedTailCount);
646+ tilingData_.set_maxSeqlenKey(tilingInfo->maxSeqlenKey);
647+ tilingData_.set_sparseCount(tilingInfo->sparseCount);
648+ tilingData_.set_sparseBlockSize(tilingInfo->sparseBlockSize);
649+ tilingData_.set_sparseRatio(tilingInfo->sparseRatio);
650+ tilingData_.set_dSize(tilingInfo->dSize);
651+ tilingData_.set_blockSize(tilingInfo->blockSize);
652+ tilingData_.set_maxBlockNumPerBatch(tilingInfo->maxBlockNumPerBatch);
653+ tilingData_.SaveToBuffer(context_->GetRawTilingData()->GetData(), context_->GetRawTilingData()->GetCapacity());
654+ context_->GetRawTilingData()->SetDataSize(tilingData_.GetDataSize());
655+ 
656+ // -------------set tilingkey-----------------
657+ // DT_Q, DT_KV, DT_OUT, PAGE_ATTENTION, FLASH_DECODE, LAYOUT_T, KV_LAYOUT_T
658+ uint32_t inputQType = static_cast<uint32_t>(tilingInfo->inputQType);
659+ uint32_t inputKType = static_cast<uint32_t>(tilingInfo->inputKType);
660+ uint32_t outputType = static_cast<uint32_t>(tilingInfo->outputType);
661+ uint32_t pageAttentionFlag = static_cast<uint32_t>(tilingInfo->pageAttentionFlag);
662+ uint32_t inputKLayout = static_cast<uint32_t>(tilingInfo->inputKLayout);
663+ uint32_t tilingKey =
664+ GET_TPL_TILING_KEY(inputQType, inputKType, outputType, pageAttentionFlag, inputKLayout);
665+ context_->SetTilingKey(tilingKey);
666+ 
667+ return ge::GRAPH_SUCCESS;
668+}
669+ 
670+// --------------------------Tiling函数定义---------------------------
671+ge::graphStatus TilingForQuantSalsIndexer(gert::TilingContext *context)
672+{
673+ OPS_ERR_IF(context == nullptr, OPS_REPORT_VECTOR_INNER_ERR("SalsIndexer", "Tiling context is null."),
674+ return ge::GRAPH_FAILED);
675+ QSITilingInfo siInfo;
676+ QSIInfoParser QSIInfoParser(context);
677+ if (QSIInfoParser.ParseAndCheck(siInfo) != ge::GRAPH_SUCCESS) {
678+ return ge::GRAPH_FAILED;
679+ }
680+ QuantSalsIndexerTiling siTiling(context);
681+ return siTiling.DoTiling(&siInfo);
682+}
683+ 
684+// --------------------------Tiling函数及TilingPrepare函数注册--------
685+IMPL_OP_OPTILING(QuantSalsIndexer)
686+ .Tiling(TilingForQuantSalsIndexer)
687+ .TilingParse<QSICompileInfo>(TilingPrepareForQuantSalsIndexer);
688+ 
689+} // namespace optiling
Aexperimental/attention/quant_sals_indexer/op_host/quant_sals_indexer_tiling.h+239-0
@@ -0,0 +1,239 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file quant_sals_indexer_tiling.h
13+ * \brief
14+ */
15+ 
16+#ifndef QUANT_SALS_INDEXER_TILING_H_
17+#define QUANT_SALS_INDEXER_TILING_H_
18+ 
19+#include "exe_graph/runtime/tiling_context.h"
20+#include "tiling/platform/platform_ascendc.h"
21+#include "register/op_def_registry.h"
22+#include "register/tilingdata_base.h"
23+#include "tiling/tiling_api.h"
24+#include "error/ops_error.h"
25+#include "platform/platform_info.h"
26+ 
27+namespace optiling {
28+const uint32_t NCAI_MAX_AIC_CORE_NUM = 24; // 样例代码，当前仅支持核数24
29+// ------------------公共定义--------------------------
30+struct TilingRequiredParaInfo {
31+ const gert::CompileTimeTensorDesc *desc;
32+ const gert::StorageShape *shape;
33+};
34+ 
35+struct TilingOptionalParaInfo {
36+ const gert::CompileTimeTensorDesc *desc;
37+ const gert::Tensor *tensor;
38+};
39+ 
40+struct SplitParams {
41+ uint32_t *bN2End;
42+ uint32_t *gS1End;
43+ uint32_t *s2End;
44+};
45+ 
46+enum class DataLayout : uint32_t {
47+ BSND = 0,
48+ PA_BNSD = 2,
49+ PA_BSND = 3,
50+ PA_NZ = 4
51+};
52+ 
53+// ------------------算子原型索引常量定义----------------
54+// Inputs Index
55+constexpr uint32_t QUERY_INDEX = 0;
56+constexpr uint32_t KEY_INDEX = 1;
57+constexpr uint32_t QUERY_DEQUANT_SCALE_INDEX = 2;
58+constexpr uint32_t KEY_DEQUANT_SCALE_INDEX = 3;
59+constexpr uint32_t ACTUAL_SEQ_K_INDEX = 4;
60+constexpr uint32_t BLOCK_TABLE_INDEX = 5;
61+constexpr uint32_t METADATA_INDEX = 6;
62+ 
63+// Outputs Index
64+constexpr uint32_t SPARSE_INDICES_INDEXER = 0;
65+ 
66+// Attributes Index
67+constexpr uint32_t ATTR_MAX_SEQLEN_KEY_INDEX = 0;
68+constexpr uint32_t ATTR_SPARSE_BLOCK_SIZE_INDEX = 1;
69+constexpr uint32_t ATTR_SPARSE_RATIO_INDEX = 2;
70+constexpr uint32_t ATTR_FIXED_TAIL_COUNT_INDEX = 3;
71+constexpr uint32_t ATTR_KEY_LAYOUT_INDEX = 4;
72+ 
73+// Dim Index
74+constexpr uint32_t DIM_IDX_ONE = 1;
75+constexpr uint32_t DIM_IDX_TWO = 2;
76+constexpr uint32_t DIM_IDX_THREE = 3;
77+// Dim Num
78+constexpr uint32_t DIM_NUM_ONE = 1;
79+constexpr uint32_t DIM_NUM_TWO = 2;
80+constexpr uint32_t DIM_NUM_THREE = 3;
81+constexpr uint32_t DIM_NUM_FOUR = 4;
82+constexpr uint32_t DIM_NUM_FIVE = 5;
83+ 
84+// -----------算子TilingData定义---------------
85+// 外切分核参数
86+BEGIN_TILING_DATA_DEF(QuantSalsIndexerSplitParams)
87+TILING_DATA_FIELD_DEF_ARR(uint32_t, NCAI_MAX_AIC_CORE_NUM, bN2End)
88+TILING_DATA_FIELD_DEF_ARR(uint32_t, NCAI_MAX_AIC_CORE_NUM, gS1End)
89+TILING_DATA_FIELD_DEF_ARR(uint32_t, NCAI_MAX_AIC_CORE_NUM, s2End)
90+END_TILING_DATA_DEF
91+REGISTER_TILING_DATA_CLASS(QuantSalsIndexerSplitParamsOp, QuantSalsIndexerSplitParams)
92+ 
93+BEGIN_TILING_DATA_DEF(QSITilingData)
94+TILING_DATA_FIELD_DEF(uint32_t, bSize)
95+TILING_DATA_FIELD_DEF(uint32_t, n2Size)
96+TILING_DATA_FIELD_DEF(uint32_t, s2Size)
97+TILING_DATA_FIELD_DEF(uint32_t, dSize)
98+TILING_DATA_FIELD_DEF(int32_t, fixedTailCount)
99+TILING_DATA_FIELD_DEF(int32_t, maxSeqlenKey)
100+TILING_DATA_FIELD_DEF(uint32_t, sparseCount)
101+TILING_DATA_FIELD_DEF(uint32_t, usedCoreNum)
102+TILING_DATA_FIELD_DEF(uint32_t, blockSize)
103+TILING_DATA_FIELD_DEF(uint32_t, maxBlockNumPerBatch)
104+TILING_DATA_FIELD_DEF(uint32_t, sparseBlockSize)
105+TILING_DATA_FIELD_DEF(float, sparseRatio)
106+TILING_DATA_FIELD_DEF_STRUCT(QuantSalsIndexerSplitParams, splitParams)
107+END_TILING_DATA_DEF
108+REGISTER_TILING_DATA_CLASS(QuantSalsIndexer, QSITilingData)
109+ 
110+// -----------算子CompileInfo定义-------------------
111+struct QSICompileInfo {};
112+ 
113+// -----------算子Tiling入参结构体定义---------------
114+struct QSiParaInfo {
115+ TilingRequiredParaInfo query = {nullptr, nullptr};
116+ TilingRequiredParaInfo key = {nullptr, nullptr};
117+ TilingRequiredParaInfo query_dequant_scale = {nullptr, nullptr};
118+ TilingRequiredParaInfo key_dequant_scale = {nullptr, nullptr};
119+ TilingOptionalParaInfo actualSeqLengths = {nullptr, nullptr};
120+ TilingOptionalParaInfo blockTable = {nullptr, nullptr};
121+ TilingRequiredParaInfo sparseIndices = {nullptr, nullptr};
122+
123+ const char *layOutKey = nullptr;
124+ const int32_t *blockSize = nullptr;
125+ const int32_t *maxSeqlenKey = nullptr;
126+ const int32_t *sparseBlockSize = nullptr;
127+ const float *sparseRatio = nullptr;
128+ const int32_t *fixedTailCount = nullptr;
129+};
130+
131+// -----------算子Tiling入参信息类---------------
132+class QSITilingInfo {
133+public:
134+ const char *opName = nullptr;
135+ fe::PlatFormInfos *platformInfo = nullptr;
136+ QSiParaInfo opParamInfo;
137+ // Base Param
138+ platform_ascendc::SocVersion socVersion = platform_ascendc::SocVersion::ASCEND910B;
139+ uint32_t bSize = 0;
140+ uint32_t n2Size = 0;
141+ int64_t s2Size = 0;
142+ int64_t dSize = 0;
143+ int32_t sparseCount = 0;
144+ int32_t sparseBlockSize = 0;
145+ int32_t fixedTailCount = 0;
146+ int32_t maxSeqlenKey = 0;
147+ float sparseRatio = 0;
148+ // PageAttention
149+ bool pageAttentionFlag = false;
150+ int32_t blockSize = 0;
151+ uint32_t maxBlockNumPerBatch = 0;
152+ // DType
153+ ge::DataType inputQType = ge::DT_FLOAT16;
154+ ge::DataType inputKType = ge::DT_FLOAT16;
155+ ge::DataType outputType = ge::DT_INT32;
156+ // Layout
157+ DataLayout inputKLayout = DataLayout::PA_BNSD;
158+};
159+ 
160+// -----------算子Tiling入参信息解析及Check类---------------
161+class QSIInfoParser {
162+public:
163+ explicit QSIInfoParser(gert::TilingContext *context) : context_(context)
164+ {
165+ }
166+ ~QSIInfoParser() = default;
167+ 
168+ ge::graphStatus CheckRequiredInOutExistence() const;
169+ ge::graphStatus CheckRequiredAttrExistence() const;
170+ ge::graphStatus CheckRequiredParaExistence() const;
171+ ge::graphStatus GetActualSeqLenSize(uint32_t &size, const gert::Tensor *tensor,
172+ const std::string &actualSeqLenName);
173+ ge::graphStatus GetOpName();
174+ ge::graphStatus GetNpuInfo();
175+ void GetOptionalInputParaInfo();
176+ void GetInputParaInfo();
177+ void GetOutputParaInfo();
178+ ge::graphStatus GetAndCheckAttrParaInfo();
179+ ge::graphStatus GetOpParaInfo();
180+ ge::graphStatus ValidateInputShapesMatch();
181+ ge::graphStatus GetAndCheckInOutDataType();
182+ ge::graphStatus GetBatchSize();
183+ ge::graphStatus GetHeadDim();
184+ ge::graphStatus GetSparseCount();
185+ ge::graphStatus GetAndCheckOptionalInput();
186+ ge::graphStatus CheckShapeDim();
187+ ge::graphStatus GetAndCheckBlockSize();
188+ ge::graphStatus CheckBlockCount();
189+ ge::graphStatus GetS2SizeForPageAttention();
190+ ge::graphStatus GetS2Size();
191+ ge::graphStatus GetQueryKeyAndOutLayout();
192+ ge::graphStatus GetAndCheckN2Size();
193+ ge::graphStatus GetAttenMaskInfo();
194+ ge::graphStatus GetActualSeqInfo();
195+ void GenerateInfo(QSITilingInfo &siInfo);
196+ ge::graphStatus CheckScaleShape();
197+ ge::graphStatus ParseAndCheck(QSITilingInfo &siInfo);
198+ 
199+public:
200+ gert::TilingContext *context_ = nullptr;
201+ const char *opName_;
202+ fe::PlatFormInfos *platformInfo_;
203+ QSiParaInfo opParamInfo_;
204+ 
205+ // BaseParams
206+ uint32_t bSize_ = 0;
207+ uint32_t n2Size_ = 0;
208+ int64_t s2Size_ = 0;
209+ uint32_t headDim_ = 0;
210+ uint32_t sparseCount_ = 0;
211+ 
212+ // Layout
213+ DataLayout kLayout_ = DataLayout::PA_BNSD;
214+ // PageAttention
215+ uint32_t maxBlockNumPerBatch_ = 0;
216+ int32_t blockSize_ = 0;
217+ platform_ascendc::SocVersion socVersion_ = platform_ascendc::SocVersion::ASCEND910B;
218+ ge::DataType inputQType_ = ge::DT_FLOAT16;
219+ ge::DataType inputKType_ = ge::DT_FLOAT16;
220+ ge::DataType inputQueryScaleType_ = ge::DT_FLOAT;
221+ ge::DataType inputKeyScaleType_ = ge::DT_FLOAT;
222+ ge::DataType blockTableType_ = ge::DT_FLOAT16;
223+ ge::DataType inputKRopeType_ = ge::DT_FLOAT16;
224+ ge::DataType sparseIndicesType_ = ge::DT_FLOAT16;
225+};
226+ 
227+// ---------------算子Tiling类---------------
228+class QuantSalsIndexerTiling {
229+public:
230+ explicit QuantSalsIndexerTiling(gert::TilingContext *context) : context_(context){};
231+ ge::graphStatus DoTiling(QSITilingInfo *tilingInfo);
232+ void SplitCoreBN(uint32_t coreNum, QSITilingInfo *siInfo, SplitParams splitParams);
233+private:
234+ gert::TilingContext *context_ = nullptr;
235+ QSITilingData tilingData_;
236+};
237+ 
238+} // namespace optiling
239+#endif // QUANT_SALS_INDEXER_TILING_H_
Aexperimental/attention/quant_sals_indexer/op_host/split_core.cpp+516-0
@@ -0,0 +1,516 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file split_core.cpp
13+ * \brief
14+ */
15+ 
16+#include "split_core.h"
17+#include <cstdio>
18+#include <math.h>
19+ 
20+ 
21+namespace optiling {
22+namespace qsi{
23+ 
24+uint32_t GetS1SeqSize(uint32_t bIdx, const BaseInfo &baseInfo)
25+{
26+ return 1U;
27+}
28+ 
29+uint32_t GetS2SeqSize(uint32_t bIdx, const BaseInfo &baseInfo)
30+{
31+ uint32_t s2Size = 0;
32+ if (baseInfo.actualSeqS2Size.empty()) {
33+ s2Size = baseInfo.s2Size;
34+ } else if (baseInfo.actualLenKvDims == 1U) {
35+ s2Size = static_cast<uint32_t>(baseInfo.actualSeqS2Size[0]);
36+ }else if (!baseInfo.isAccumSeqS2) {
37+ s2Size = static_cast<uint32_t>(baseInfo.actualSeqS2Size[bIdx]);
38+ } else {
39+ s2Size = (bIdx == 0) ? static_cast<uint32_t>(baseInfo.actualSeqS2Size[bIdx]) :
40+ static_cast<uint32_t>(baseInfo.actualSeqS2Size[bIdx] - baseInfo.actualSeqS2Size[bIdx - 1U]);
41+ }
42+ 
43+ int32_t totalNCount = (s2Size + baseInfo.sparseBlockSize - 1) / baseInfo.sparseBlockSize;
44+ if (totalNCount <= baseInfo.fixedTailCount) {
45+ return 0U;
46+ }
47+ return (totalNCount - baseInfo.fixedTailCount) * baseInfo.sparseBlockSize;
48+}
49+ 
50+int64_t CalcCost(uint32_t basicM, uint32_t basicS2)
51+{
52+ uint32_t alignCoefM = 16U;
53+ uint32_t alignCoefS2 = 64U;
54+ uint32_t alignBasicM = (basicM + alignCoefM - 1U) >> 4U; // 按alignCoefM对齐，向上取整，4：移位操作实现除16
55+ uint32_t alignBasicS2 = (basicS2 + alignCoefS2 - 1U) >> 6U; // 按alignCoefS2对齐，向上取整，6：移位操作实现除64
56+ return static_cast<int64_t>(6U * alignBasicM + 10U * alignBasicS2); // 6：M轴系数，10：S2轴系数
57+}
58+ 
59+BlockCost<int64_t> CalcCostTable(uint32_t s1NormalSize, uint32_t s2NormalSize, uint32_t s1GTailSize,
60+ uint32_t s2TailSize)
61+{
62+ BlockCost<int64_t> typeCost {};
63+ typeCost[NORMAL_BLOCK][NORMAL_BLOCK] = CalcCost(s1NormalSize, s2NormalSize);
64+ typeCost[TAIL_BLOCK][NORMAL_BLOCK] = (s1GTailSize == 0U) ? 0U : CalcCost(s1GTailSize, s2NormalSize);
65+ typeCost[NORMAL_BLOCK][TAIL_BLOCK] = (s2TailSize == 0U) ? 0U : CalcCost(s1NormalSize, s2TailSize);
66+ typeCost[TAIL_BLOCK][TAIL_BLOCK] = (s1GTailSize == 0U || s2TailSize == 0U) ? 0U : CalcCost(s1GTailSize, s2TailSize);
67+ return typeCost;
68+}
69+ 
70+Range<uint32_t> CalcS2Range(const SplitParam &splitParam, const BatchCache &batchCache)
71+{
72+ uint32_t s2Start = 0U;
73+ uint32_t s2End = 0U;
74+ 
75+ if (batchCache.s1Size == 0U || batchCache.s2Size == 0U) {
76+ return std::make_pair(s2Start, s2End);
77+ }
78+ 
79+ s2End = (batchCache.s2Size + splitParam.s2BaseSize - 1U) / splitParam.s2BaseSize;
80+ return std::make_pair(s2Start, s2End);
81+}
82+ 
83+void CalcSplitInfo(SplitContext &splitContext)
84+{
85+ const BaseInfo &baseInfo = splitContext.baseInfo;
86+ const SplitParam &splitParam = splitContext.splitParam;
87+ 
88+ // 计算每个batch的切分，统计是否为空batch，记录最后有效batch（每个batch的每个N2切分是一样的）
89+ SplitInfo &splitInfo = splitContext.splitInfo;
90+ for (uint32_t bIdx = 0; bIdx < baseInfo.bSize; bIdx++) {
91+ uint32_t s1Size = GetS1SeqSize(bIdx, baseInfo);
92+ uint32_t s2Size = GetS2SeqSize(bIdx, baseInfo);
93+ 
94+ splitInfo.s1GBaseNum[bIdx] = (s1Size * baseInfo.gSize + (splitParam.mBaseSize - 1U)) / splitParam.mBaseSize;
95+ splitInfo.s1GTailSize[bIdx] = (s1Size * baseInfo.gSize) % splitParam.mBaseSize;
96+ splitInfo.s2BaseNum[bIdx] = (s2Size + splitParam.s2BaseSize - 1U) / splitParam.s2BaseSize;
97+ splitInfo.s2TailSize[bIdx] = s2Size % splitParam.s2BaseSize;
98+ if (splitInfo.s1GBaseNum[bIdx] != 0U && splitInfo.s2BaseNum[bIdx] != 0U) {
99+ splitInfo.isKvSeqAllZero = false;
100+ }
101+ }
102+}
103+ 
104+void CalcBatchCache(uint32_t bIdx, const SplitContext &splitContext, BatchCache &batchCache)
105+{
106+ const BaseInfo &baseInfo = splitContext.baseInfo;
107+ const SplitParam &splitParam = splitContext.splitParam;
108+ const SplitInfo &splitInfo = splitContext.splitInfo;
109+ 
110+ batchCache.bIdx = bIdx;
111+ batchCache.s1Size = GetS1SeqSize(bIdx, baseInfo);
112+ batchCache.s2Size = GetS2SeqSize(bIdx, baseInfo);
113+ batchCache.typeCost = CalcCostTable(splitParam.mBaseSize, splitParam.s2BaseSize, splitInfo.s1GTailSize[bIdx],
114+ splitInfo.s2TailSize[bIdx]);
115+}
116+ 
117+void CalcS1GCache(uint32_t s1GIdx, const SplitContext &splitContext, const BatchCache &batchCache, S1GCache &s1GCache)
118+{
119+ const BaseInfo &baseInfo = splitContext.baseInfo;
120+ const SplitParam &splitParam = splitContext.splitParam;
121+ const SplitInfo &splitInfo = splitContext.splitInfo;
122+ 
123+ s1GCache.bIdx = batchCache.bIdx;
124+ s1GCache.s1GIdx = s1GIdx;
125+ 
126+ auto s2Range = CalcS2Range(splitParam, batchCache);
127+ s1GCache.s2Start = s2Range.first;
128+ s1GCache.s2End = s2Range.second;
129+ 
130+ // 计算S2方向满块、尾块数量
131+ s1GCache.s1GBlock = s1GCache.s2End - s1GCache.s2Start;
132+ uint32_t curTailS2Num = (splitInfo.s2TailSize[batchCache.bIdx] != 0U &&
133+ s1GCache.s2End == splitInfo.s2BaseNum[batchCache.bIdx]) ? 1U : 0U;
134+ uint32_t curNormalS2Num = s1GCache.s1GBlock - curTailS2Num;
135+ if (splitInfo.s1GBaseNum[batchCache.bIdx] == 0) {
136+ s1GCache.s1GCost = 0;
137+ s1GCache.s1GLastBlockCost = 0;
138+ s1GCache.s1GNormalBlockCost = 0;
139+ } else if (s1GIdx == (splitInfo.s1GBaseNum[batchCache.bIdx] - 1U) && splitInfo.s1GTailSize[batchCache.bIdx] != 0U) {
140+ s1GCache.s1GCost = batchCache.typeCost[TAIL_BLOCK][NORMAL_BLOCK] * curNormalS2Num +
141+ batchCache.typeCost[TAIL_BLOCK][TAIL_BLOCK] * curTailS2Num;
142+ s1GCache.s1GLastBlockCost = curTailS2Num > 0U ? batchCache.typeCost[TAIL_BLOCK][TAIL_BLOCK] :
143+ batchCache.typeCost[TAIL_BLOCK][NORMAL_BLOCK];
144+ s1GCache.s1GNormalBlockCost = batchCache.typeCost[TAIL_BLOCK][NORMAL_BLOCK];
145+ } else {
146+ s1GCache.s1GCost = batchCache.typeCost[NORMAL_BLOCK][NORMAL_BLOCK] * curNormalS2Num +
147+ batchCache.typeCost[NORMAL_BLOCK][TAIL_BLOCK] * curTailS2Num;
148+ s1GCache.s1GLastBlockCost = curTailS2Num > 0U ? batchCache.typeCost[NORMAL_BLOCK][TAIL_BLOCK] :
149+ batchCache.typeCost[NORMAL_BLOCK][NORMAL_BLOCK];
150+ s1GCache.s1GNormalBlockCost = batchCache.typeCost[NORMAL_BLOCK][NORMAL_BLOCK];
151+ }
152+}
153+ 
154+void CopyTmpResult(SplitResult &tmpRes, SplitResult &splitRes)
155+{
156+ uint64_t len = tmpRes.bN2End.size();
157+ splitRes.usedCoreNum = tmpRes.usedCoreNum;
158+ splitRes.maxCost = tmpRes.maxCost;
159+ 
160+ for (size_t i = 0; i < len; ++i) {
161+ splitRes.s2End[i] = tmpRes.s2End[i];
162+ splitRes.gS1End[i] = tmpRes.gS1End[i];
163+ splitRes.bN2End[i] = tmpRes.bN2End[i];
164+ }
165+}
166+ 
167+void RollBackCursor(const BaseInfo &baseInfo, const SplitContext &splitContext, const CostInfo &costInfo, SplitResult &splitRes)
168+{
169+ for (size_t i = 0; i < splitRes.usedCoreNum; ++i) {
170+ // x, y, z
171+ if (splitRes.s2End[i] > 0U) {
172+ splitRes.s2End[i] = splitRes.s2End[i] - 1U;
173+ continue;
174+ }
175+ uint32_t bIdx = splitRes.bN2End[i] / baseInfo.n2Size;
176+ // x, y, 0
177+ if (splitRes.gS1End[i] > 0U) {
178+ splitRes.gS1End[i] = splitRes.gS1End[i] - 1U;
179+ splitRes.s2End[i] = splitContext.splitInfo.s2BaseNum[bIdx] - 1U;
180+ continue;
181+ }
182+ 
183+ // x, 0, 0
184+ uint32_t bN2Idx = splitRes.bN2End[i] > 0U ? splitRes.bN2End[i] - 1U : 0U;
185+ bIdx = bN2Idx / baseInfo.n2Size;
186+ while (bN2Idx > 0U && costInfo.bN2BlockOfEachBatch[bIdx] == 0U) {
187+ bN2Idx -= 1U;
188+ bIdx = bN2Idx / baseInfo.n2Size;
189+ }
190+ 
191+ if (costInfo.bN2BlockOfEachBatch[bIdx] != 0U) {
192+ splitRes.bN2End[i] = bN2Idx;
193+ splitRes.gS1End[i] = splitContext.splitInfo.s1GBaseNum[bIdx] - 1U;
194+ splitRes.s2End[i] = splitContext.splitInfo.s2BaseNum[bIdx] - 1U;
195+ } else {
196+ splitRes.bN2End[i] = 0U;
197+ splitRes.gS1End[i] = 0U;
198+ splitRes.s2End[i] = 0U;
199+ }
200+ }
201+}
202+ 
203+void ClearTmpResult(SplitResult &tmpResult)
204+{
205+ uint64_t len = tmpResult.bN2End.size();
206+ tmpResult.usedCoreNum = 0U;
207+ tmpResult.maxCost = 0;
208+ 
209+ for (size_t i = 0; i < len; ++i) {
210+ tmpResult.bN2End[i] = 0U;
211+ tmpResult.gS1End[i] = 0U;
212+ tmpResult.s2End[i] = 0U;
213+ }
214+}
215+ 
216+void CalcBatchCost(uint32_t bIdx, const SplitContext &splitContext, CostInfo &costInfo)
217+{
218+ const BaseInfo &baseInfo = splitContext.baseInfo;
219+ const SplitInfo &splitInfo = splitContext.splitInfo;
220+ 
221+ costInfo.bN2CostOfEachBatch[bIdx] = 0;
222+ costInfo.bN2BlockOfEachBatch[bIdx] = 0U;
223+ costInfo.bN2LastBlockCostOfEachBatch[bIdx] = 0U;
224+ 
225+ if (GetS1SeqSize(bIdx, baseInfo) == 0U || GetS2SeqSize(bIdx, baseInfo) == 0U) {
226+ return;
227+ }
228+ 
229+ BatchCache bCache;
230+ S1GCache s1GCache;
231+ CalcBatchCache(bIdx, splitContext, bCache);
232+ for (uint32_t s1GIdx = 0; s1GIdx < splitInfo.s1GBaseNum[bIdx]; s1GIdx++) {
233+ CalcS1GCache(s1GIdx, splitContext, bCache, s1GCache);
234+ costInfo.bN2CostOfEachBatch[bIdx] += s1GCache.s1GCost;
235+ costInfo.bN2BlockOfEachBatch[bIdx] += s1GCache.s1GBlock;
236+ }
237+ 
238+ costInfo.bN2LastBlockCostOfEachBatch[bIdx] = s1GCache.s1GLastBlockCost;
239+}
240+ 
241+void CalcCostInfo(SplitContext &splitContext)
242+{
243+ const BaseInfo &baseInfo = splitContext.baseInfo;
244+ const SplitInfo &splitInfo = splitContext.splitInfo;
245+ 
246+ CostInfo &costInfo = splitContext.costInfo;
247+ 
248+ if (splitInfo.isKvSeqAllZero) {
249+ costInfo.totalCost = 0;
250+ costInfo.totalBlockNum = 0U;
251+ return;
252+ }
253+ 
254+ // 计算batch的负载并记录，用于按batch分配，需要按行计算起止点，统计块数、负载
255+ for (uint32_t bIdx = 0; bIdx < baseInfo.bSize; bIdx++) {
256+ CalcBatchCost(bIdx, splitContext, costInfo);
257+ costInfo.totalCost += costInfo.bN2CostOfEachBatch[bIdx] * baseInfo.n2Size;
258+ costInfo.totalBlockNum += costInfo.bN2BlockOfEachBatch[bIdx] * baseInfo.n2Size;
259+ }
260+}
261+ 
262+void UpdateCursor(const SplitContext &splitContext, AssignContext &assignContext)
263+{
264+ const BaseInfo &baseInfo = splitContext.baseInfo;
265+ const SplitInfo &splitInfo = splitContext.splitInfo;
266+ const CostInfo &costInfo = splitContext.costInfo;
267+ 
268+ bool UpdateS1G = false;
269+ bool UpdateBatch = false;
270+ 
271+ // Update S2
272+ if (assignContext.curS2Idx >= assignContext.s1GCache.s2End) { // 边界assignInfo.s2End是取不到的开区间
273+ assignContext.curS2Idx = 0U;
274+ assignContext.curS1GIdx++;
275+ UpdateS1G = true;
276+ }
277+ 
278+ // Update S1G
279+ if (assignContext.curS1GIdx >= splitInfo.s1GBaseNum[assignContext.curBIdx]) {
280+ assignContext.curS1GIdx = 0U;
281+ assignContext.curBN2Idx++;
282+ }
283+ 
284+ // Update Batch
285+ if (assignContext.curBN2Idx == baseInfo.bSize * baseInfo.n2Size) { // 所有负载全部分配完，设置最后一个核的右开区间，返回
286+ assignContext.curS1GIdx = 0U;
287+ assignContext.curS2Idx = 0U;
288+ assignContext.isFinished = true;
289+ return;
290+ }
291+ 
292+ if (assignContext.curBN2Idx / baseInfo.n2Size != assignContext.curBIdx) {
293+ assignContext.curBIdx = assignContext.curBN2Idx / baseInfo.n2Size;
294+ assignContext.curS1GIdx = 0U;
295+ UpdateBatch = true;
296+ UpdateS1G = true;
297+ }
298+ 
299+ // Update Cache
300+ if (UpdateBatch) {
301+ CalcBatchCache(assignContext.curBIdx, splitContext, assignContext.batchCache);
302+ assignContext.bN2Cost = costInfo.bN2CostOfEachBatch[assignContext.curBIdx];
303+ assignContext.bN2Block = costInfo.bN2BlockOfEachBatch[assignContext.curBIdx];
304+ }
305+ if (UpdateS1G) {
306+ CalcS1GCache(assignContext.curS1GIdx, splitContext, assignContext.batchCache, assignContext.s1GCache);
307+ assignContext.curS2Idx = assignContext.s1GCache.s2Start;
308+ }
309+}
310+ 
311+void AssignByBatch(const SplitContext &splitContext, AssignContext &assignContext)
312+{
313+ if (assignContext.isFinished) {
314+ return;
315+ }
316+ 
317+ const BaseInfo &baseInfo = splitContext.baseInfo;
318+ const CostInfo &costInfo = splitContext.costInfo;
319+ 
320+ while (assignContext.bN2Cost == 0 || IsWithinTolerance(assignContext.coreCache.costLimit,
321+ costInfo.bN2LastBlockCostOfEachBatch[assignContext.curBIdx] / FA_TOLERANCE_RATIO,
322+ assignContext.coreCache.cost + assignContext.bN2Cost)) {
323+ assignContext.coreCache.cost += assignContext.bN2Cost;
324+ assignContext.coreCache.block += assignContext.bN2Block;
325+ assignContext.curBN2Idx++;
326+ 
327+ // to the end
328+ if (assignContext.curBN2Idx == baseInfo.bSize * baseInfo.n2Size) {
329+ assignContext.curS1GIdx = 0U;
330+ assignContext.curS2Idx = 0U;
331+ assignContext.isFinished = true;
332+ return;
333+ }
334+ 
335+ // next batch
336+ if (assignContext.curBN2Idx / baseInfo.n2Size != assignContext.curBIdx) {
337+ assignContext.curBIdx = assignContext.curBN2Idx / baseInfo.n2Size;
338+ CalcBatchCache(assignContext.curBIdx, splitContext, assignContext.batchCache);
339+ }
340+ 
341+ assignContext.bN2Cost = costInfo.bN2CostOfEachBatch[assignContext.curBIdx];
342+ assignContext.bN2Block = costInfo.bN2BlockOfEachBatch[assignContext.curBIdx];
343+ assignContext.curS1GIdx = 0U;
344+ CalcS1GCache(assignContext.curS1GIdx, splitContext, assignContext.batchCache, assignContext.s1GCache);
345+ assignContext.curS2Idx = assignContext.s1GCache.s2Start;
346+ }
347+}
348+ 
349+void AssignByRow(const SplitContext &splitContext, AssignContext &assignContext)
350+{
351+ if (assignContext.isFinished) {
352+ return;
353+ }
354+ 
355+ while (IsWithinTolerance(assignContext.coreCache.costLimit,
356+ assignContext.s1GCache.s1GLastBlockCost / FA_TOLERANCE_RATIO,
357+ assignContext.coreCache.cost + assignContext.s1GCache.s1GCost)) {
358+ assignContext.coreCache.cost += assignContext.s1GCache.s1GCost;
359+ assignContext.coreCache.block += assignContext.s1GCache.s1GBlock;
360+ 
361+ assignContext.curS1GIdx++;
362+ // 当前batch被分配一行出去，更新剩余负载
363+ assignContext.bN2Cost = assignContext.bN2Cost > assignContext.s1GCache.s1GCost ?
364+ assignContext.bN2Cost - assignContext.s1GCache.s1GCost : 0;
365+ assignContext.bN2Block = assignContext.bN2Block > assignContext.s1GCache.s1GBlock ?
366+ assignContext.bN2Block - assignContext.s1GCache.s1GBlock : 0U;
367+ // 计算新一行的信息
368+ CalcS1GCache(assignContext.curS1GIdx, splitContext, assignContext.batchCache, assignContext.s1GCache);
369+ assignContext.curS2Idx = assignContext.s1GCache.s2Start;
370+ }
371+}
372+ 
373+void AssignByBlock(const SplitContext &splitContext, AssignContext &assignContext)
374+{
375+ if (assignContext.isFinished) {
376+ return;
377+ }
378+ 
379+ int64_t curCost = assignContext.s1GCache.s1GNormalBlockCost;
380+ if (assignContext.curS2Idx == (assignContext.s1GCache.s2End - 1U)) {
381+ curCost = assignContext.s1GCache.s1GLastBlockCost;
382+ }
383+ 
384+ while (IsWithinTolerance(assignContext.coreCache.costLimit, curCost / FA_TOLERANCE_RATIO,
385+ assignContext.coreCache.cost + curCost)) { // (costLimit - curCostOnCore) * FA_TOLERANCE_RATIO > curCost；至少分配1块
386+ assignContext.coreCache.cost += curCost;
387+ assignContext.coreCache.block++;
388+ assignContext.curS2Idx++;
389+ // 当前batch被分配一块出去，更新剩余负载
390+ assignContext.bN2Cost = assignContext.bN2Cost - curCost;
391+ // 当前行被分配一块出去，更新剩余负载
392+ assignContext.s1GCache.s1GCost = assignContext.s1GCache.s1GCost - curCost;
393+ assignContext.bN2Block--;
394+ assignContext.s1GCache.s1GBlock--;
395+ }
396+}
397+ 
398+void ForceAssign(const SplitContext &splitContext, AssignContext &assignContext)
399+{
400+ if (assignContext.isFinished) {
401+ return;
402+ }
403+ 
404+ int64_t curCost = assignContext.s1GCache.s1GNormalBlockCost;
405+ if (assignContext.curS2Idx == (assignContext.s1GCache.s2End - 1U)) {
406+ curCost = assignContext.s1GCache.s1GLastBlockCost;
407+ }
408+ 
409+ assignContext.coreCache.cost += curCost;
410+ assignContext.coreCache.block++;
411+ assignContext.curS2Idx++;
412+ // 当前batch被分配一块出去，更新剩余负载
413+ assignContext.bN2Cost = assignContext.bN2Cost - curCost;
414+ assignContext.bN2Block--;
415+ // 当前行被分配一块出去，更新剩余负载
416+ assignContext.s1GCache.s1GCost = assignContext.s1GCache.s1GCost - curCost;
417+ assignContext.s1GCache.s1GBlock--;
418+ UpdateCursor(splitContext, assignContext);
419+}
420+ 
421+void CalcSplitPlan(uint32_t coreNum, int64_t costLimit, const SplitContext &splitContext, SplitResult &result)
422+{
423+ const CostInfo &costInfo = splitContext.costInfo;
424+ 
425+ if (coreNum == 0U) {
426+ return;
427+ }
428+ result.maxCost = 0U;
429+ result.usedCoreNum = 0U;
430+ 
431+ AssignContext assignContext {};
432+ assignContext.curBIdx = 0U;
433+ assignContext.curS1GIdx = 0U;
434+ assignContext.unassignedCost = costInfo.totalCost;
435+ assignContext.bN2Cost = costInfo.bN2CostOfEachBatch[assignContext.curBIdx];
436+ assignContext.bN2Block = costInfo.bN2BlockOfEachBatch[assignContext.curBIdx];
437+ CalcBatchCache(assignContext.curBIdx, splitContext, assignContext.batchCache);
438+ CalcS1GCache(assignContext.curS1GIdx, splitContext, assignContext.batchCache, assignContext.s1GCache);
439+ assignContext.curS2Idx = assignContext.s1GCache.s2Start;
440+ 
441+ for (uint32_t i = 0; i < coreNum; ++i) {
442+ if (result.maxCost >= costLimit) {
443+ return;
444+ }
445+ if (assignContext.isFinished || assignContext.unassignedCost <= 0) {
446+ break;
447+ }
448+ 
449+ assignContext.curCoreIdx = i;
450+ 
451+ assignContext.coreCache = {};
452+ assignContext.coreCache.costLimit = assignContext.unassignedCost / (coreNum - assignContext.curCoreIdx);
453+ 
454+ // 1、按整batch分配
455+ AssignByBatch(splitContext, assignContext);
456+ // 2、按行分配
457+ AssignByRow(splitContext, assignContext);
458+ // 3、按块分配
459+ AssignByBlock(splitContext, assignContext);
460+ // 4、强制分配
461+ if (assignContext.coreCache.block == 0) {
462+ ForceAssign(splitContext, assignContext);
463+ }
464+ 
465+ result.bN2End[i] = assignContext.curBN2Idx;
466+ result.gS1End[i] = assignContext.curS1GIdx;
467+ result.s2End[i] = assignContext.curS2Idx;
468+ result.maxCost = std::max(result.maxCost, assignContext.coreCache.cost);
469+ 
470+ assignContext.unassignedCost -= assignContext.coreCache.cost;
471+ }
472+ 
473+ result.usedCoreNum = assignContext.curCoreIdx + 1;
474+}
475+ 
476+void SplitCore(uint32_t coreNum, const BaseInfo &baseInfo, const SplitParam &param, SplitResult &result)
477+{
478+ SplitContext splitContext(baseInfo, param);
479+ 
480+ // 1、划分基本块，统计信息
481+ CalcSplitInfo(splitContext);
482+ // 全空case
483+ if (splitContext.splitInfo.isKvSeqAllZero) {
484+ result.usedCoreNum = 1U;
485+ result.bN2End[0] = baseInfo.bSize * baseInfo.n2Size;
486+ result.gS1End[0] = 0U;
487+ result.s2End[0] = 0U;
488+ return;
489+ }
490+ 
491+ CalcCostInfo(splitContext);
492+ 
493+ // 2、获取每个核的分配方案
494+ uint32_t maxCore = std::min(coreNum, splitContext.costInfo.totalBlockNum);
495+ uint32_t minCore = static_cast<uint32_t>(
496+ std::sqrt(static_cast<float>(splitContext.costInfo.totalBlockNum) + 0.25f) + 0.5f);
497+ minCore = std::min(minCore, maxCore);
498+ 
499+ result.maxCost = INT64_MAX;
500+ result.usedCoreNum = 1U;
501+ 
502+ SplitResult tmpResult {coreNum, result.vecCubeRatio};
503+ for (uint32_t i = minCore; i <= maxCore; ++i) {
504+ CalcSplitPlan(i, result.maxCost, splitContext, tmpResult);
505+ if (tmpResult.maxCost < result.maxCost) {
506+ CopyTmpResult(tmpResult, result);
507+ }
508+ ClearTmpResult(tmpResult);
509+ }
510+ RollBackCursor(baseInfo, splitContext, splitContext.costInfo, result);
511+ 
512+ result.usedCoreNum = std::max(result.usedCoreNum, 1U); // 至少使用1个core
513+ 
514+}
515+}
516+}
Aexperimental/attention/quant_sals_indexer/op_host/split_core.h+209-0
@@ -0,0 +1,209 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file split_core.h
13+ * \brief
14+ */
15+ 
16+#ifndef SPLIT_CORE_H
17+#define SPLIT_CORE_H
18+ 
19+#include <cstdint>
20+#include <vector>
21+#include <array>
22+#include <algorithm>
23+ 
24+namespace optiling {
25+namespace qsi{
26+constexpr int64_t FA_TOLERANCE_RATIO = 2;
27+ 
28+enum BlockType : uint32_t {
29+ NORMAL_BLOCK = 0,
30+ TAIL_BLOCK,
31+ BLOCK_MAX_TYPE
32+};
33+ 
34+template<class T>
35+using Range = std::pair<T, T>;
36+ 
37+template<class T>
38+using BlockCost = std::array<std::array<T, static_cast<size_t>(BLOCK_MAX_TYPE)>, static_cast<size_t>(BLOCK_MAX_TYPE)>;
39+ 
40+template<typename T>
41+inline bool IsWithinTolerance(T limit, T tolerance, T value)
42+{
43+ return limit + tolerance >= value;
44+}
45+ 
46+// 分核功能模块输入：输入case的基本信息
47+struct BaseInfo {
48+ uint32_t bSize { 0U };
49+ uint32_t n2Size { 0U };
50+ uint32_t gSize { 0U };
51+ uint32_t s1Size { 0U };
52+ uint32_t s2Size { 0U };
53+ bool isAccumSeqS2 { false };
54+ std::vector<int64_t> actualSeqS2Size {};
55+ uint32_t actualLenKvDims { 0U };
56+ uint32_t fixedTailCount { 0U };
57+ uint32_t sparseBlockSize { 0U };
58+};
59+ 
60+// 分核功能模块输入：切分属性，预留接口，可作为切分方案的参数入口
61+struct SplitParam {
62+ uint32_t mBaseSize { 1U };
63+ uint32_t s2BaseSize { 2048U };
64+};
65+ 
66+// 分核功能模块输出：FA阶段的核间分核信息
67+struct SplitResult {
68+ uint32_t usedCoreNum { 0U }; // 使用的核数量
69+ uint32_t vecCubeRatio { 0U }; // vec 与 cube 核数比例
70+ std::vector<uint32_t> bN2End {}; // 每个核处理数据的BN2结束点
71+ std::vector<uint32_t> gS1End {}; // 每个核处理数据的GS1结束点
72+ std::vector<uint32_t> s2End {}; // 每个核处理数据的S2结束点
73+ int64_t maxCost { 0 }; // 慢核开销
74+ 
75+ SplitResult(uint32_t coreNum, uint32_t ratio) :
76+ vecCubeRatio(ratio),
77+ bN2End(coreNum),
78+ gS1End(coreNum),
79+ s2End(coreNum) {}
80+};
81+ 
82+// 分核功能模块内部使用：记录切分信息
83+struct SplitInfo {
84+ std::vector<uint32_t> s1GBaseNum {}; // S1G方向，切了多少个基本块
85+ std::vector<uint32_t> s2BaseNum {}; // S2方向，切了多少个基本块
86+ std::vector<uint32_t> s1GTailSize {}; // S1G方向，尾块size
87+ std::vector<uint32_t> s2TailSize {}; // S2方向，尾块size
88+ bool isKvSeqAllZero { true };
89+ 
90+ explicit SplitInfo(uint32_t batchSize) :
91+ s1GBaseNum(batchSize),
92+ s2BaseNum(batchSize),
93+ s1GTailSize(batchSize),
94+ s2TailSize(batchSize) {}
95+};
96+ 
97+// 分核功能模块内部使用：记录batch的开销信息
98+struct CostInfo {
99+ std::vector<int64_t> bN2CostOfEachBatch {}; // 整个batch的开销
100+ std::vector<uint32_t> bN2BlockOfEachBatch {}; // 整个batch的开销
101+ std::vector<int64_t> bN2LastBlockCostOfEachBatch {}; // batch最后一块的开销
102+ uint32_t totalBlockNum { 0U };
103+ int64_t totalCost { 0 };
104+ 
105+ explicit CostInfo(uint32_t batchSize) :
106+ bN2CostOfEachBatch(batchSize),
107+ bN2BlockOfEachBatch(batchSize),
108+ bN2LastBlockCostOfEachBatch(batchSize) {}
109+};
110+ 
111+// 分核功能模块内部使用：分核过程中，case基本信息的上下文信息，组合以减少接口传参数量
112+struct SplitContext {
113+ const BaseInfo &baseInfo {};
114+ const SplitParam &splitParam {};
115+ SplitInfo splitInfo { 0U };
116+ CostInfo costInfo { 0U };
117+ 
118+ explicit SplitContext(const BaseInfo &info, const SplitParam &param) :
119+ baseInfo(info),
120+ splitParam(param),
121+ splitInfo(info.bSize),
122+ costInfo(info.bSize) {}
123+};
124+ 
125+// 分核功能模块内部使用：记录batch相关的临时信息
126+struct BatchCache {
127+ uint32_t bIdx { 0U };
128+ uint32_t s1Size { 0U };
129+ uint32_t s2Size { 0U };
130+ int64_t preTokenLeftUp { 0 };
131+ int64_t nextTokenLeftUp { 0 };
132+ BlockCost<int64_t> typeCost {};
133+};
134+ 
135+// 分核功能模块内部使用：记录当前行（S1G）的临时信息
136+struct S1GCache {
137+ uint32_t bIdx { 0U };
138+ uint32_t s1GIdx { 0U };
139+ uint32_t s2Start { 0U };
140+ uint32_t s2End { 0U };
141+ int64_t s1GCost { 0 };
142+ int64_t s1GLastBlockCost { 0 };
143+ uint32_t s1GBlock { 0U };
144+ int64_t s1GNormalBlockCost { 0 };
145+};
146+ 
147+// 分核功能模块内部使用：记录分配过程中，当前核的负载信息
148+struct CoreCache {
149+ int64_t costLimit { 0 }; // 负载上限
150+ int64_t cost { 0 }; // 已分配负载
151+ uint32_t block { 0U }; // 已分配块数
152+};
153+ 
154+// 分核功能模块内部使用：记录分配过程中的上下文信息
155+struct AssignContext {
156+ uint32_t curBIdx { 0U };
157+ uint32_t curBN2Idx { 0U };
158+ uint32_t curS1GIdx { 0U };
159+ uint32_t curS2Idx { 0U };
160+ uint32_t curCoreIdx { 0U };
161+ int64_t unassignedCost { 0 };
162+ uint32_t usedCoreNum { 0U };
163+ uint32_t curKvSplitPart { 1U };
164+ 
165+ int64_t bN2Cost { 0 };
166+ uint32_t bN2Block { 0U };
167+ bool isFinished { false };
168+ BatchCache batchCache {};
169+ S1GCache s1GCache {};
170+ CoreCache coreCache {};
171+};
172+ 
173+// util
174+uint32_t GetS1SeqSize(uint32_t bIdx, const BaseInfo &baseInfo);
175+uint32_t GetS2SeqSize(uint32_t bIdx, const BaseInfo &baseInfo);
176+int64_t CalcPreTokenLeftUp(uint32_t s1Size, uint32_t s2Size, const BaseInfo &baseInfo);
177+int64_t CalcNextTokenLeftUp(uint32_t s1Size, uint32_t s2Size, const BaseInfo &baseInfo);
178+Range<uint32_t> CalcS2Range(uint32_t s1GIdx, const BaseInfo &baseInfo, const SplitParam &splitParam,
179+ const BatchCache &batchCache);
180+int64_t CalcCost(uint32_t basicM, uint32_t basicS2);
181+BlockCost<int64_t> CalcCostTable(uint32_t s1NormalSize, uint32_t s2NormalSize, uint32_t s1GTailSize,
182+ uint32_t s2TailSize);
183+ 
184+// cache calculation
185+void CalcBatchCache(uint32_t bIdx, const SplitContext &splitContext, BatchCache &batchCache);
186+void CalcS1GCache(uint32_t s1GIdx, const SplitContext &splitContext, const BatchCache &batchCache, S1GCache &s1GCache);
187+void CopyTmpResult(const SplitContext &splitContext, SplitResult &tmpRes, SplitResult &splitRes);
188+void ClearTmpResult(SplitResult &tmpResult);
189+ 
190+// preprocess
191+void CalcSplitInfo(SplitContext &splitContext);
192+void CalcBatchCost(uint32_t bIdx, const SplitContext &splitContext, CostInfo &costInfo);
193+void CalcCostInfo(SplitContext &splitContext);
194+ 
195+// assign
196+void UpdateCursor(const SplitContext &splitContext, AssignContext &assignContext);
197+void AssignByBatch(const SplitContext &splitContext, AssignContext &assignContext);
198+void AssignByRow(const SplitContext &splitContext, AssignContext &assignContext);
199+void AssignByBlock(const SplitContext &splitContext, AssignContext &assignContext);
200+void ForceAssign(const SplitContext &splitContext, AssignContext &assignContext);
201+ 
202+// main
203+void SplitFD(SplitResult &result);
204+void CalcSplitPlan(uint32_t coreNum, int64_t costLimit, const SplitContext &splitContext, SplitResult &result);
205+void SplitCore(uint32_t coreNum, const BaseInfo &baseInfo, const SplitParam &splitParam, SplitResult &result);
206+void RollBackCursor(const BaseInfo &baseInfo, const SplitContext &splitContext, const CostInfo &costInfo, SplitResult &splitRes);
207+}
208+}
209+#endif
Aexperimental/attention/quant_sals_indexer/op_kernel/quant_sals_indexer.cpp+82-0
@@ -0,0 +1,82 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file quant_sals_indexer.cpp
13+ * \brief
14+ */
15+ 
16+#include "kernel_operator.h"
17+#include "lib/matmul_intf.h"
18+#include "quant_sals_indexer_kernel.h"
19+#include "quant_sals_indexer_template_tiling_key.h"
20+#include "quant_sals_indexer_metadata.h"
21+ 
22+using namespace QSIKernel;
23+using namespace optiling::detail;
24+ 
25+template <class T>
26+__inline__ __attribute__((always_inline)) __aicore__ void InitMetaData(const __gm__ uint8_t *p_metadata, T *metadata)
27+{
28+ constexpr uint64_t all_bytes = sizeof(T);
29+#if defined(ASCENDC_CPU_DEBUG) || defined(__DAV_C220_CUBE__) || defined(__DAV_C310_CUBE__) || defined(__DAV_310R6_CUBE__) || defined(__GET_CODE_CHANNEL__)
30+ copy_data_align64((uint8_t*)metadata, (__gm__ uint8_t *)p_metadata, all_bytes);
31+#else
32+ __ubuf__ uint8_t *metadata_in_ub = (__ubuf__ uint8_t *)get_imm(0);
33+ constexpr uint32_t len_burst = (all_bytes + 31) / 32;
34+ copy_gm_to_ubuf(((__ubuf__ uint8_t *)metadata_in_ub), p_metadata, 0, 1,len_burst, 0, 0);
35+ set_flag(PIPE_MTE2, PIPE_S, EVENT_ID0);
36+ wait_flag(PIPE_MTE2, PIPE_S, EVENT_ID0);
37+ copy_data_align64((uint8_t*)metadata, (__ubuf__ uint8_t *)metadata_in_ub, all_bytes);
38+#endif
39+ 
40+}
41+ 
42+#define INVOKE_QSI_NO_KFC_OP_IMPL(templateClass, ...) \
43+ do { \
44+ templateClass<QSIType<__VA_ARGS__>> op; \
45+ GET_TILING_DATA_WITH_STRUCT(QSITilingData, tiling_data_in, tiling); \
46+ const QSITilingData *__restrict tiling_data = &tiling_data_in; \
47+ QsiMetaData *__restrict meta_data = nullptr; \
48+ QsiMetaData metaDataTmp; \
49+ if (metaData != nullptr) { \
50+ InitMetaData<QsiMetaData>(metaData, &metaDataTmp); \
51+ meta_data = &metaDataTmp; \
52+ } \
53+ op.Init(query, key, query_dequant_scale, key_dequant_scale, actualSeqLengths, \
54+ blocktable, meta_data, sparseIndices, user, tiling_data, &tPipe); \
55+ op.Process(); \
56+ } while (0)
57+ 
58+ 
59+template <int DT_Q, int DT_K, int DT_OUT, int PAGE_ATTENTION, int K_LAYOUT_T>
60+__global__ __aicore__ void quant_sals_indexer(__gm__ uint8_t *query, __gm__ uint8_t *key,
61+ __gm__ uint8_t *query_dequant_scale, __gm__ uint8_t *key_dequant_scale,
62+ __gm__ uint8_t *actualSeqLengths, __gm__ uint8_t *blocktable,
63+ __gm__ uint8_t *metaData, __gm__ uint8_t *sparseIndices,
64+ __gm__ uint8_t *workspace, __gm__ uint8_t *tiling)
65+{
66+#if (__CCE_AICORE__ == 310) || (defined __DAV_310R6__) || (__CCE_AICORE__ == 200)
67+ 
68+#else
69+ TPipe tPipe;
70+ __gm__ uint8_t *user = GetUserWorkspace(workspace);
71+ KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_MIX_AIC_1_2);
72+ 
73+ if constexpr (DT_Q == QSI_TPL_INT4 && DT_K == QSI_TPL_INT4) {
74+ INVOKE_QSI_NO_KFC_OP_IMPL(QSIPreload, int4b_t, int4b_t, int32_t, PAGE_ATTENTION,
75+ QSI_LAYOUT(K_LAYOUT_T));
76+ } else {
77+ INVOKE_QSI_NO_KFC_OP_IMPL(QSIPreload, int8_t, int8_t, int32_t, PAGE_ATTENTION,
78+ QSI_LAYOUT(K_LAYOUT_T));
79+ }
80+ 
81+#endif
82+}
Aexperimental/attention/quant_sals_indexer/op_kernel/quant_sals_indexer_common.h+159-0
@@ -0,0 +1,159 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file quant_sals_indexer_common.h
13+ * \brief
14+ */
15+#ifndef QUANT_SALS_INDEXER_COMMON_H
16+#define QUANT_SALS_INDEXER_COMMON_H
17+ 
18+namespace QSICommon {
19+constexpr uint32_t MAX_TOPK = 2048;
20+constexpr uint32_t LD_PARAM_NUM = 16;
21+ 
22+// 与tiling的layout保持一致
23+enum class QSI_LAYOUT {
24+ BSND = 0,
25+ TND = 1,
26+ PA_BNSD = 2,
27+ PA_BSND = 3,
28+ PA_NZ = 4
29+};
30+ 
31+template <typename Q_T, typename K_T, typename OUT_T, const bool PAGE_ATTENTION = false,
32+ QSI_LAYOUT K_LAYOUT_T = QSI_LAYOUT::PA_BNSD, typename... Args>
33+struct QSIType {
34+ using queryType = Q_T;
35+ using keyType = K_T;
36+ using outputType = OUT_T;
37+ static constexpr bool pageAttention = PAGE_ATTENTION;
38+ static constexpr QSI_LAYOUT keyLayout = K_LAYOUT_T;
39+};
40+ 
41+struct RunInfo {
42+ uint32_t loop;
43+ uint32_t bN2Idx;
44+ uint32_t bIdx;
45+ uint32_t n2Idx = 0;
46+ uint32_t gS1Idx;
47+ uint32_t s2Idx;
48+ 
49+ uint32_t actS1Size = 1;
50+ uint32_t actS2Size = 1;
51+ int32_t needProcessS2Size = 0;
52+ int32_t targetTopKAlign = 0L;
53+ int32_t targetTopK = 0L;
54+ int32_t fixedTailCount = 0L;
55+ uint32_t actMBaseSize;
56+ uint32_t actualSingleProcessSInnerSize;
57+ uint32_t actualSingleProcessSInnerSizeAlign;
58+ 
59+ uint64_t tensorQueryOffset;
60+ uint64_t tensorKeyOffset;
61+ uint64_t tensorKeyScaleOffset;
62+ uint64_t indiceOutOffset;
63+ 
64+ float qScale = 0.0f;
65+ 
66+ bool isFirstS2InnerLoop;
67+ bool isLastS2InnerLoop;
68+ bool isAllLoopEnd = false;
69+};
70+ 
71+struct ConstInfo {
72+ // CUBE与VEC核间同步的模式
73+ static constexpr uint32_t FIA_SYNC_MODE2 = 2;
74+ // BUFFER的字节数
75+ static constexpr uint32_t BUFFER_SIZE_BYTE_32B = 32;
76+ static constexpr uint32_t BUFFER_SIZE_BYTE_64B = 64;
77+ static constexpr uint32_t BUFFER_SIZE_BYTE_256B = 256;
78+ static constexpr uint32_t BUFFER_SIZE_BYTE_512B = 512;
79+ static constexpr uint32_t BUFFER_SIZE_BYTE_1K = 1024;
80+ static constexpr uint32_t BUFFER_SIZE_BYTE_2K = 2048;
81+ static constexpr uint32_t BUFFER_SIZE_BYTE_4K = 4096;
82+ static constexpr uint32_t BUFFER_SIZE_BYTE_8K = 8192;
83+ static constexpr uint32_t BUFFER_SIZE_BYTE_16K = 16384;
84+ static constexpr uint32_t BUFFER_SIZE_BYTE_32K = 32768;
85+ // 无效索引
86+ static constexpr int INVALID_IDX = -1;
87+ 
88+ // CUBE和VEC的核间同步EventID
89+ uint32_t syncC1V1 = 0U;
90+ uint32_t syncV1C1 = 0U;
91+ 
92+ // 基本块大小
93+ uint32_t mBaseSize = 1ULL;
94+ uint32_t s1BaseSize = 1ULL;
95+ uint32_t s2BaseSize = 1ULL;
96+ 
97+ uint64_t batchSize = 0ULL;
98+ uint64_t qHeadNum = 0ULL;
99+ uint64_t kHeadNum;
100+ uint64_t headDim;
101+ uint64_t sparseCount; // topK选取大小
102+ uint64_t kSeqSize = 0ULL; // kv最大S长度
103+ uint64_t qSeqSize = 1ULL; // q最大S长度
104+ uint32_t kCacheBlockSize = 0; // PA场景的block size
105+ int32_t sparseBlockSize = 0;
106+ float sparseRatio = 0.0f;
107+ int32_t fixedTailCount = 1UL;
108+ int32_t maxSeqlenKey = 0ULL;
109+ uint32_t maxBlockNumPerBatch = 0; // PA场景的最大单batch block number
110+ QSI_LAYOUT outputLayout; // 输出的格式
111+ bool attenMaskFlag = false;
112+ 
113+ uint32_t actualLenQDims = 0U; // query的actualSeqLength 的维度
114+ uint32_t actualLenDims = 0U; // KV 的actualSeqLength 的维度
115+ bool isAccumSeqS1 = false; // 是否累加模式
116+ bool isAccumSeqS2 = false; // 是否累加模式
117+};
118+ 
119+struct SplitCoreInfo {
120+ uint32_t s2Start = 0U; // S2的起始位置
121+ uint32_t s2End = 0U; // S2循环index上限
122+ uint32_t bN2Start = 0U;
123+ uint32_t bN2End = 0U;
124+ uint32_t gS1Start = 0U;
125+ uint32_t gS1End = 0U;
126+ bool isLD = false; // 当前核是否需要进行Decode归约任务
127+};
128+ 
129+template <typename T>
130+__aicore__ inline T Align(T num, T rnd)
131+{
132+ return (((rnd) == 0) ? 0 : (((num) + (rnd)-1) / (rnd) * (rnd)));
133+}
134+ 
135+template <typename T1, typename T2>
136+__aicore__ inline T1 Min(T1 a, T2 b)
137+{
138+ return (a > b) ? (b) : (a);
139+}
140+ 
141+template <typename T1, typename T2>
142+__aicore__ inline T1 Max(T1 a, T2 b)
143+{
144+ return (a > b) ? (a) : (b);
145+}
146+ 
147+__aicore__ inline uint64_t QSiCeilDiv(uint64_t num, uint64_t rnd)
148+{
149+ return rnd == 0 ? 0 : (num + rnd-1) / rnd;
150+}
151+ 
152+template <typename T>
153+__aicore__ inline T QSiCeilAlign(T num, T rnd)
154+{
155+ return rnd == 0 ? 0 : (num + rnd-1) / rnd * rnd;
156+}
157+} // namespace QSICommon
158+ 
159+#endif // QUANT_SALS_INDEXER_COMMON_H
Aexperimental/attention/quant_sals_indexer/op_kernel/quant_sals_indexer_kernel.h+687-0
@@ -0,0 +1,687 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file quant_sals_indexer_kernel.h
13+ * \brief
14+ */
15+ 
16+#ifndef QUANT_SALS_INDEXER_KERNEL_H
17+#define QUANT_SALS_INDEXER_KERNEL_H
18+ 
19+#include "kernel_operator.h"
20+#include "kernel_operator_list_tensor_intf.h"
21+#include "kernel_tiling/kernel_tiling.h"
22+#include "lib/matmul_intf.h"
23+#include "lib/matrix/matmul/tiling.h"
24+#include "quant_sals_indexer_common.h"
25+#include "quant_sals_indexer_service_vector.h"
26+#include "quant_sals_indexer_service_cube_int4.h"
27+#include "quant_sals_indexer_metadata.h"
28+ 
29+#define ARRAY_SIZE(x) (sizeof(x) / sizeof(x[0]))
30+namespace QSIKernel {
31+using namespace QSICommon;
32+using namespace optiling::detail;
33+using namespace matmul;
34+using AscendC::CacheMode;
35+using AscendC::CrossCoreSetFlag;
36+using AscendC::CrossCoreWaitFlag;
37+ 
38+__aicore__ inline int32_t FloatCeil(float num) {
39+ int integer_part = static_cast<int>(num);
40+ if (num > 0 && num != integer_part) {
41+ integer_part += 1;
42+ }
43+ return integer_part;
44+}
45+ 
46+// 由于S2循环前，RunInfo还没有赋值，使用TempLoopInfo临时存放B、N、S1轴相关的信息；同时减少重复计算
47+struct TempLoopInfo {
48+ uint32_t bN2Idx = 0;
49+ uint32_t bIdx = 0U;
50+ uint32_t n2Idx = 0U;
51+ int32_t s2LoopEnd = 0; // S2方向循环的结束Idx
52+ uint32_t actS2Size = 0ULL;
53+ int32_t needProcessS2Size = 0L;
54+ int32_t targetTopKAlign = 0L;
55+ int32_t targetTopK = 0L;
56+ int32_t sparseBlockSize = 0UL;
57+ int32_t fixedTailCount = 0UL;
58+ float sparseRatio = 0.0f;
59+ float qScale = 0.0f;
60+ bool curActSeqLenIsZero = false;
61+ bool needCleanSparseCount = false; // S1的实际长度小于shape的S1长度时，是否需要清理输出
62+ bool OnlyFixTail = false;
63+ uint32_t actMBaseSize = 0U; // m轴(bN2)方向实际大小
64+ uint32_t s2BasicSizeTail = 0U; // S2方向循环的尾基本块大小
65+};
66+ 
67+template <typename QSIT>
68+class QSIPreload {
69+public:
70+ __aicore__ inline QSIPreload(){};
71+ __aicore__ inline void Init(__gm__ uint8_t *query, __gm__ uint8_t *key, __gm__ uint8_t *queryScale,
72+ __gm__ uint8_t *keyScale, __gm__ uint8_t *actualSeqLengths,
73+ __gm__ uint8_t *blockTable, QsiMetaData *metaData,
74+ __gm__ uint8_t *sparseIndices, __gm__ uint8_t *workspace,
75+ const QSITilingData *__restrict tiling, TPipe *tPipe);
76+ __aicore__ inline void Process();
77+ 
78+ // =================================类型定义区=================================
79+ using Q_T = typename QSIT::queryType;
80+ using K_T = typename QSIT::keyType;
81+ using OUT_T = typename QSIT::outputType;
82+ static constexpr bool PAGE_ATTENTION = QSIT::pageAttention;
83+ static constexpr QSI_LAYOUT K_LAYOUT_T = QSIT::keyLayout;
84+ 
85+ using MM1_OUT_T = int32_t;
86+ 
87+ QSIMatmulInt4<QSIT> matmulService;
88+ QSIVector<QSIT> vectorService;
89+ 
90+ // =================================常量区=================================
91+ static constexpr uint32_t SYNC_C1_V1_FLAG = 4;
92+ static constexpr uint32_t SYNC_V1_C1_FLAG = 5;
93+ 
94+ static constexpr uint32_t M_BASE_SIZE = 1;
95+ static constexpr uint32_t S2_BASE_SIZE = 2048;
96+ static constexpr uint32_t HEAD_DIM = 128;
97+ static constexpr uint32_t K_HEAD_NUM = 1;
98+ static constexpr uint32_t GM_ALIGN_BYTES = 512;
99+ static constexpr uint32_t SI_QUANT_PRELOAD_TASK_CACHE_SIZE = 2;
100+ 
101+ static constexpr int64_t LD_PREFETCH_LEN = 2;
102+ // for workspace double
103+ static constexpr uint32_t WS_DOBULE = 2;
104+ 
105+protected:
106+ TPipe *pipe = nullptr;
107+ QsiMetaData *metaDataPtr = nullptr;
108+ 
109+ // ================================Global Buffer区=================================
110+ GlobalTensor<Q_T> queryGm;
111+ GlobalTensor<K_T> keyGm;
112+ GlobalTensor<float> qScaleGm;
113+ GlobalTensor<float> kScaleGm;
114+ 
115+ GlobalTensor<int32_t> indiceOutGm;
116+ GlobalTensor<int32_t> blockTableGm;
117+ 
118+ GlobalTensor<uint32_t> actualSeqLengthsGm;
119+ // workspace
120+ GlobalTensor<MM1_OUT_T> mm1ResGm; // 存放S
121+ GlobalTensor<float> vec1ResGm; // 存放TopK计算中间结果
122+ GlobalTensor<int64_t> vec1ParamGm; // 存放LD参数信息
123+ 
124+ // ================================类成员变量====================================
125+ // aic、aiv核信息
126+ uint32_t tmpBlockIdx = 0U;
127+ uint32_t aiCoreIdx = 0U;
128+ uint32_t usedCoreNum = 0U;
129+ 
130+ uint64_t keyCoreOffset = 0ULL;
131+ uint64_t actualSeqKPrefixSum = 0ULL;
132+ QSICommon::ConstInfo constInfo{};
133+ TempLoopInfo tempLoopInfo{};
134+ QSICommon::SplitCoreInfo splitCoreInfo{};
135+ 
136+ // ================================Init functions==================================
137+ __aicore__ inline void InitTilingData(const QSITilingData *__restrict tilingData);
138+ __aicore__ inline void InitMetaData();
139+ __aicore__ inline void InitCalcParamsEach(const QSITilingData *__restrict tilingData,
140+ QSICommon::SplitCoreInfo &info);
141+ __aicore__ inline void GetAxisStartIdx(uint32_t bN2EndPrev, uint32_t s1GEndPrev,
142+ uint32_t s2EndPrev, QSICommon::SplitCoreInfo &info);
143+ __aicore__ inline void InitBuffers();
144+ __aicore__ inline void InitActualSeqLen(__gm__ uint8_t *actualSeqLengths);
145+ // ================================Split Core================================
146+ __aicore__ inline void SplitCore(uint32_t curCoreIdx, uint32_t &coreNum, QSICommon::SplitCoreInfo &info);
147+ __aicore__ inline int32_t GetS2NeedProcessSize(uint32_t actS2Size);
148+ __aicore__ inline uint32_t GetTotalBaseBlockNum();
149+ // ================================Process functions================================
150+ __aicore__ inline void ProcessMain();
151+ __aicore__ inline void ProcessBaseBlock(uint32_t loop, uint64_t s2LoopIdx, QSICommon::RunInfo &runInfo);
152+ __aicore__ inline void ProcessDecode();
153+ __aicore__ inline void ProcessInvalid();
154+ // ================================Params Calc=====================================
155+ __aicore__ inline void GetBN2Idx(uint32_t bN2Idx);
156+ __aicore__ inline uint32_t GetActualSeqLen(uint32_t bIdx);
157+ __aicore__ inline void CalcS2LoopParams(uint32_t bN2LoopIdx);
158+ __aicore__ inline void CalcRunInfo(uint32_t loop, uint32_t s2LoopIdx, QSICommon::RunInfo &runInfo);
159+ __aicore__ inline void CalcRunInfo(QSICommon::RunInfo &runInfo);
160+ __aicore__ inline void DealActSeqLenIsZero(uint32_t bIdx, uint32_t n2Idx);
161+};
162+ 
163+template <typename QSIT>
164+__aicore__ inline void QSIPreload<QSIT>::InitTilingData(const QSITilingData *__restrict tilingData)
165+{
166+ usedCoreNum = tilingData->usedCoreNum;
167+ constInfo.batchSize = tilingData->bSize;
168+ constInfo.kSeqSize = tilingData->s2Size;
169+ constInfo.kCacheBlockSize = tilingData->blockSize;
170+ constInfo.maxBlockNumPerBatch = tilingData->maxBlockNumPerBatch;
171+ constInfo.sparseCount = tilingData->sparseCount;
172+ constInfo.sparseBlockSize = tilingData->sparseBlockSize;
173+ constInfo.sparseRatio = tilingData->sparseRatio;
174+ constInfo.fixedTailCount = tilingData->fixedTailCount;
175+ constInfo.outputLayout = QSI_LAYOUT::BSND; // 输出和输入形状一致
176+ constInfo.maxSeqlenKey = tilingData->maxSeqlenKey;
177+ 
178+ if (K_LAYOUT_T == QSI_LAYOUT::TND) {
179+ constInfo.isAccumSeqS2 = true;
180+ }
181+ 
182+ constInfo.kHeadNum = tilingData->n2Size;
183+ constInfo.headDim = tilingData->dSize;
184+ 
185+ constInfo.mBaseSize = M_BASE_SIZE;
186+ constInfo.s2BaseSize = S2_BASE_SIZE;
187+ constInfo.s1BaseSize = M_BASE_SIZE;
188+}
189+ 
190+template <typename QSIT>
191+__aicore__ inline void QSIPreload<QSIT>::InitMetaData()
192+{
193+ usedCoreNum = metaDataPtr->usedCoreNum;
194+}
195+ 
196+template <typename QSIT>
197+__aicore__ inline void QSIPreload<QSIT>::InitCalcParamsEach(const QSITilingData *__restrict tilingData,
198+ QSICommon::SplitCoreInfo &info)
199+{
200+ // 计算总的基本块
201+ // 这里是编译器优化写法，定义一个局部数组变量coreSidxEnd(存在栈上)，使用copy_data_align64接口
202+ // 可以只从ub中拷贝tiling中coreSidxEnd的内容到栈上，而非将整个increFlashAttentionCoreParams
203+ // 内容拷贝到栈，减少拷贝时间
204+ if (metaDataPtr != nullptr) {
205+ const uint32_t *bN2End = metaDataPtr->bN2End;
206+ const uint32_t *gS1End = metaDataPtr->gS1End;
207+ const uint32_t *s2End = metaDataPtr->s2End;
208+ // TND分核信息
209+ info.bN2End = bN2End[aiCoreIdx];
210+ info.gS1End = gS1End[aiCoreIdx];
211+ info.s2End = s2End[aiCoreIdx];
212+ if (aiCoreIdx != 0) {
213+ GetAxisStartIdx(bN2End[aiCoreIdx - 1], gS1End[aiCoreIdx - 1], s2End[aiCoreIdx - 1], info);
214+ }
215+ // 当前核是否需要进行Decode规约任务
216+ uint32_t bEnd = info.bN2End / constInfo.kHeadNum;
217+ uint32_t s2BaseNum = (GetActualSeqLen(bEnd) + constInfo.s2BaseSize - 1) / constInfo.s2BaseSize;
218+ info.isLD = (info.s2Start > 0U || info.s2End < s2BaseNum - 1U);
219+ } else {
220+#ifdef ASCENDC_CPU_DEBUG
221+ const uint32_t *bN2End = tilingData->splitParams.bN2End;
222+ const uint32_t *gS1End = tilingData->splitParams.gS1End;
223+ const uint32_t *s2End = tilingData->splitParams.s2End;
224+#else
225+ uint32_t bN2End[ARRAY_SIZE(tilingData->splitParams.bN2End)];
226+ uint32_t gS1End[ARRAY_SIZE(tilingData->splitParams.gS1End)];
227+ uint32_t s2End[ARRAY_SIZE(tilingData->splitParams.s2End)];
228+ copy_data_align64((uint8_t*)bN2End, (uint8_t*)(tilingData->splitParams.bN2End), sizeof(bN2End));
229+ copy_data_align64((uint8_t*)gS1End, (uint8_t*)(tilingData->splitParams.gS1End), sizeof(gS1End));
230+ copy_data_align64((uint8_t*)s2End, (uint8_t*)(tilingData->splitParams.s2End), sizeof(s2End));
231+#endif
232+ // TND分核信息
233+ info.bN2End = bN2End[aiCoreIdx];
234+ info.gS1End = gS1End[aiCoreIdx];
235+ info.s2End = s2End[aiCoreIdx];
236+ if (aiCoreIdx != 0) {
237+ GetAxisStartIdx(bN2End[aiCoreIdx - 1], gS1End[aiCoreIdx - 1], s2End[aiCoreIdx - 1], info);
238+ }
239+ // 当前核是否需要进行Decode规约任务
240+ info.isLD = false; // 保持兼容性
241+ }
242+}
243+ 
244+template <typename QSIT>
245+__aicore__ inline void QSIPreload<QSIT>::GetAxisStartIdx(uint32_t bN2EndPrev, uint32_t s1GEndPrev,
246+ uint32_t s2EndPrev, QSICommon::SplitCoreInfo &info)
247+{
248+ if (metaDataPtr != nullptr) {
249+ uint32_t bEndPrev = bN2EndPrev / constInfo.kHeadNum;
250+ uint32_t actualSeqQPrev = 1;
251+ uint32_t s1GPrevBaseNum = (actualSeqQPrev + constInfo.mBaseSize - 1) / constInfo.mBaseSize;
252+ uint32_t s2PrevBaseNum = (GetActualSeqLen(bEndPrev) + constInfo.s2BaseSize - 1) / constInfo.s2BaseSize;
253+ info.bN2Start = bN2EndPrev;
254+ info.gS1Start = s1GEndPrev;
255+ info.s2Start = s2EndPrev + 1U;
256+ 
257+ if (info.s2Start >= s2PrevBaseNum) {
258+ info.gS1Start++;
259+ info.s2Start = 0;
260+ }
261+ if (info.gS1Start >= s1GPrevBaseNum) {
262+ info.bN2Start++;
263+ info.gS1Start = 0;
264+ }
265+ } else {
266+ uint32_t bEndPrev = bN2EndPrev / constInfo.kHeadNum;
267+ uint32_t actualSeqQPrev = 1;
268+ uint32_t s1GPrevBaseNum = (actualSeqQPrev + constInfo.mBaseSize - 1) / constInfo.mBaseSize;
269+ info.bN2Start = bN2EndPrev;
270+ info.gS1Start = s1GEndPrev;
271+ 
272+ info.s2Start = 0;
273+ if (s1GEndPrev >= s1GPrevBaseNum - 1) { // 上个核把S1G处理完了
274+ info.gS1Start = 0;
275+ info.bN2Start++;
276+ } else {
277+ info.gS1Start++;
278+ }
279+ }
280+}
281+ 
282+template <typename QSIT>
283+__aicore__ inline void QSIPreload<QSIT>::InitBuffers()
284+{
285+ if ASCEND_IS_AIV {
286+ vectorService.InitBuffers(pipe);
287+ } else {
288+ matmulService.InitBuffers(pipe);
289+ }
290+}
291+ 
292+template <typename QSIT>
293+__aicore__ inline void QSIPreload<QSIT>::InitActualSeqLen(__gm__ uint8_t *actualSeqLengths)
294+{
295+ // todo 删除
296+ constInfo.actualLenQDims = constInfo.batchSize;
297+ if (actualSeqLengths == nullptr) {
298+ constInfo.actualLenDims = 0;
299+ } else {
300+ constInfo.actualLenDims = constInfo.batchSize;
301+ actualSeqLengthsGm.SetGlobalBuffer((__gm__ uint32_t *)actualSeqLengths, constInfo.actualLenDims);
302+ }
303+}
304+ 
305+template <typename QSIT>
306+__aicore__ inline uint32_t QSIPreload<QSIT>::GetActualSeqLen(uint32_t bIdx)
307+{
308+ if (constInfo.actualLenDims == 0) {
309+ return constInfo.kSeqSize;
310+ } else if (constInfo.isAccumSeqS2 && bIdx > 0) {
311+ return actualSeqLengthsGm.GetValue(bIdx) - actualSeqLengthsGm.GetValue(bIdx - 1);
312+ } else {
313+ return actualSeqLengthsGm.GetValue(bIdx);
314+ }
315+}
316+ 
317+template <typename QSIT>
318+__aicore__ inline int32_t QSIPreload<QSIT>::GetS2NeedProcessSize(uint32_t actS2Size)
319+{
320+ int32_t totalNCount = (actS2Size + constInfo.sparseBlockSize - 1) / constInfo.sparseBlockSize;
321+ if (totalNCount <= constInfo.fixedTailCount) {
322+ return 0;
323+ }
324+ return (totalNCount - constInfo.fixedTailCount) * constInfo.sparseBlockSize;
325+}
326+ 
327+template <typename QSIT>
328+__aicore__ inline uint32_t QSIPreload<QSIT>::GetTotalBaseBlockNum()
329+{
330+ uint32_t totalBlockNum = 0;
331+ for (uint32_t bIdx = 0; bIdx < constInfo.batchSize; bIdx++) {
332+ uint32_t actS2Size = GetActualSeqLen(bIdx);
333+ 
334+ totalBlockNum += (GetS2NeedProcessSize(actS2Size)+ constInfo.s2BaseSize - 1) / constInfo.s2BaseSize * constInfo.kHeadNum;
335+ }
336+ return totalBlockNum;
337+}
338+ 
339+// 多核版本，双闭区间
340+template <typename QSIT>
341+__aicore__ void inline QSIPreload<QSIT>::SplitCore(uint32_t curCoreIdx, uint32_t &coreNum, QSICommon::SplitCoreInfo &info)
342+{
343+ // 计算每个核最少处理的块数, 剩余的部分前面的核每个核多处理一块
344+ uint32_t totalBlockNum = GetTotalBaseBlockNum();
345+ uint32_t minBlockPerCore = totalBlockNum / coreNum;
346+ uint32_t deal1MoreBlockCoreNum = totalBlockNum % coreNum;
347+ uint32_t coreIdx = 0;
348+ uint32_t lastBN2RemainBlockCnt = 0;
349+ uint32_t coreDealBlockCnt = coreIdx < deal1MoreBlockCoreNum ? minBlockPerCore + 1 : minBlockPerCore;
350+ coreNum = minBlockPerCore == 0 ? deal1MoreBlockCoreNum : coreNum;
351+ 
352+ bool findLastCoreEnd = true;
353+ uint32_t s2BaseNum;
354+ for (uint32_t bN2Idx = 0; bN2Idx < constInfo.batchSize * constInfo.kHeadNum; bN2Idx++) {
355+ uint32_t bIdx = bN2Idx / constInfo.kHeadNum;
356+ if (bN2Idx % constInfo.kHeadNum == 0) {
357+ uint32_t actS2Size = GetActualSeqLen(bIdx);
358+ s2BaseNum = (GetS2NeedProcessSize(actS2Size) + constInfo.s2BaseSize - 1) / constInfo.s2BaseSize;
359+ }
360+ if (findLastCoreEnd && s2BaseNum == 0U) {
361+ info.bN2Start = bN2Idx;
362+ info.s2Start = 0;
363+ findLastCoreEnd = false;
364+ }
365+ for (uint32_t s2Idx = 0; s2Idx < s2BaseNum;) {
366+ uint32_t s2RemainBaseNum = s2BaseNum - s2Idx;
367+ if (findLastCoreEnd) {
368+ info.bN2Start = bN2Idx;
369+ info.s2Start = s2Idx;
370+ findLastCoreEnd = false;
371+ }
372+ if (lastBN2RemainBlockCnt + s2RemainBaseNum >= coreDealBlockCnt) {
373+ info.bN2End = bN2Idx;
374+ info.s2End = s2Idx + coreDealBlockCnt - lastBN2RemainBlockCnt - 1;
375+ 
376+ if (coreIdx == curCoreIdx) {
377+ // S2被切N核，那么只有第一个核需要处理LD，其他核不用
378+ if (s2Idx == 0 && info.s2End + 1 < s2BaseNum) {
379+ info.isLD = true;
380+ }
381+ // 最后一个核处理的不是最后一个Batch，表明后面的Batch为空块(S2=0), 调整终点坐标以便清理输出
382+ if (coreIdx == coreNum - 1 && info.bN2End / constInfo.kHeadNum != constInfo.batchSize - 1) {
383+ info.bN2End = (constInfo.batchSize * constInfo.kHeadNum) - 1;
384+ info.s2End = 0;
385+ }
386+ return;
387+ }
388+ coreIdx++;
389+ findLastCoreEnd = true;
390+ s2Idx = info.s2End + 1;
391+ lastBN2RemainBlockCnt = 0;
392+ coreDealBlockCnt = coreIdx < deal1MoreBlockCoreNum ? minBlockPerCore + 1 : minBlockPerCore;
393+ } else {
394+ lastBN2RemainBlockCnt += s2RemainBaseNum;
395+ break;
396+ }
397+ }
398+ }
399+}
400+ 
401+template <typename QSIT>
402+__aicore__ inline void QSIPreload<QSIT>::DealActSeqLenIsZero(uint32_t bIdx, uint32_t n2Idx)
403+{
404+ if ASCEND_IS_AIV {
405+ // B,N2,K
406+ if (GetSubBlockIdx() == 0) {
407+ return;
408+ }
409+ uint64_t indiceOutOffset =
410+ bIdx * constInfo.kHeadNum * constInfo.sparseCount + n2Idx * constInfo.sparseCount; // N2轴偏移
411+ vectorService.CleanInvalidOutput(indiceOutOffset);
412+ }
413+}
414+ 
415+template <typename QSIT>
416+__aicore__ inline void QSIPreload<QSIT>::Init(__gm__ uint8_t *query, __gm__ uint8_t *key,
417+ __gm__ uint8_t *queryScale, __gm__ uint8_t *keyScale, __gm__ uint8_t *actualSeqLengths,
418+ __gm__ uint8_t *blockTable, QsiMetaData *metaData, __gm__ uint8_t *sparseIndices,
419+ __gm__ uint8_t *workspace, const QSITilingData *__restrict tiling,
420+ TPipe *tPipe)
421+{
422+ if ASCEND_IS_AIV {
423+ tmpBlockIdx = GetBlockIdx(); // vec:0-47
424+ aiCoreIdx = tmpBlockIdx / 2;
425+ } else {
426+ tmpBlockIdx = GetBlockIdx(); // cube:0-23
427+ aiCoreIdx = tmpBlockIdx;
428+ }
429+ 
430+ InitTilingData(tiling);
431+ InitActualSeqLen(actualSeqLengths);
432+ 
433+ // 计算分核
434+ // SplitCore(aiCoreIdx, usedCoreNum, splitCoreInfo);
435+ if (metaData != nullptr) {
436+ metaDataPtr = metaData;
437+ InitMetaData();
438+ }
439+ InitCalcParamsEach(tiling, splitCoreInfo);
440+ 
441+ pipe = tPipe;
442+ // workspace 内存排布
443+ // |mm1ResGm(存S)|vec1ResGm(存LD中间结果)|vec1ParamGm(存LD参数)
444+ // |Core0_mm1ResDB0-Core0_mm1ResDB1-Core1_mm1ResDB0....Core23_mm1ResDB0-Core23_mm1ResDB1|Core0_vec1Res...
445+ uint64_t offset = 0;
446+ 
447+ // mm1开DoubleBuffer
448+ uint64_t singleCoreMm1ResSize = WS_DOBULE * constInfo.mBaseSize * constInfo.s2BaseSize * sizeof(MM1_OUT_T);
449+ mm1ResGm.SetGlobalBuffer((__gm__ MM1_OUT_T *)(workspace + offset + aiCoreIdx * singleCoreMm1ResSize));
450+ offset += GetBlockNum() * singleCoreMm1ResSize;
451+ 
452+ // ld流程需要ws大小: [aicnum, 2, constInfo.mBaseSize, topkOut_*2]
453+ // (aic, 8, 2, 2, 2048)
454+ // (aic, s1_cube, 头尾, idx/value, K)
455+ vec1ResGm.SetGlobalBuffer((__gm__ float *)(workspace + offset));
456+ offset += GetBlockNum() * constInfo.s1BaseSize * WS_DOBULE * WS_DOBULE * MAX_TOPK * sizeof(float);
457+ 
458+ // (aic, 8, 2, 16)
459+ // (aic, s1_cube, 头尾，16ele)
460+ vec1ParamGm.SetGlobalBuffer((__gm__ int64_t *)(workspace + offset));
461+ offset += GetBlockNum() * constInfo.s1BaseSize * WS_DOBULE * LD_PARAM_NUM * sizeof(int64_t);
462+ 
463+ qScaleGm.SetGlobalBuffer((__gm__ float *)queryScale);
464+ if ASCEND_IS_AIV {
465+ vectorService.InitParams(constInfo, tiling);
466+ indiceOutGm.SetGlobalBuffer((__gm__ int32_t *)sparseIndices);
467+ kScaleGm.SetGlobalBuffer((__gm__ float *)keyScale);
468+ if constexpr (PAGE_ATTENTION) {
469+ blockTableGm.SetGlobalBuffer((__gm__ int32_t *)blockTable);
470+ }
471+ vectorService.InitVec1GlobalTensor(mm1ResGm, qScaleGm, kScaleGm, blockTableGm, vec1ResGm, vec1ParamGm,
472+ indiceOutGm);
473+ } else {
474+ matmulService.InitParams(constInfo);
475+ queryGm.SetGlobalBuffer((__gm__ Q_T *)query);
476+ if constexpr (PAGE_ATTENTION) {
477+ blockTableGm.SetGlobalBuffer((__gm__ int32_t *)blockTable);
478+ }
479+ keyGm.SetGlobalBuffer((__gm__ K_T *)key);
480+ keyGm.SetL2CacheHint(AscendC::CacheMode::CACHE_MODE_DISABLE);
481+ matmulService.InitMm1GlobalTensor(blockTableGm, keyGm, queryGm, mm1ResGm);
482+ }
483+ InitBuffers();
484+}
485+ 
486+template <typename QSIT>
487+__aicore__ inline void QSIPreload<QSIT>::GetBN2Idx(uint32_t bN2Idx)
488+{
489+ tempLoopInfo.bN2Idx = bN2Idx;
490+ tempLoopInfo.bIdx = bN2Idx / constInfo.kHeadNum;
491+ tempLoopInfo.n2Idx = bN2Idx % constInfo.kHeadNum;
492+}
493+ 
494+template <typename QSIT>
495+__aicore__ inline void QSIPreload<QSIT>::CalcS2LoopParams(uint32_t bN2LoopIdx)
496+{
497+ GetBN2Idx(bN2LoopIdx);
498+ tempLoopInfo.qScale = qScaleGm.GetValue(tempLoopInfo.bIdx * constInfo.kHeadNum + tempLoopInfo.n2Idx);
499+ tempLoopInfo.actMBaseSize = constInfo.mBaseSize;
500+ tempLoopInfo.actS2Size = GetActualSeqLen(bN2LoopIdx / constInfo.kHeadNum);
501+ 
502+ tempLoopInfo.needProcessS2Size = GetS2NeedProcessSize(tempLoopInfo.actS2Size);
503+ tempLoopInfo.s2LoopEnd =(tempLoopInfo.needProcessS2Size + constInfo.s2BaseSize - 1) / constInfo.s2BaseSize - 1;
504+ if (metaDataPtr != nullptr) {
505+ tempLoopInfo.s2LoopEnd = (bN2LoopIdx == splitCoreInfo.bN2End && !(tempLoopInfo.actS2Size == 0)) ? splitCoreInfo.s2End : tempLoopInfo.s2LoopEnd;
506+ }
507+ tempLoopInfo.needCleanSparseCount = (tempLoopInfo.actS2Size < constInfo.maxSeqlenKey);
508+ int32_t targetTopK = FloatCeil(((tempLoopInfo.needProcessS2Size + constInfo.sparseBlockSize - 1) / (constInfo.sparseBlockSize)) * constInfo.sparseRatio);
509+ targetTopK = targetTopK >= 2048 ? 2048 : targetTopK;
510+ tempLoopInfo.OnlyFixTail = (targetTopK == 0);
511+ if (targetTopK > constInfo.s2BaseSize / constInfo.sparseBlockSize) {
512+ tempLoopInfo.targetTopKAlign = QSiCeilAlign(static_cast<uint32_t>(targetTopK), 4 * constInfo.s2BaseSize / constInfo.sparseBlockSize);
513+ } else {
514+ tempLoopInfo.targetTopKAlign = QSiCeilAlign(targetTopK, 128);
515+ }
516+ tempLoopInfo.targetTopK = targetTopK;
517+}
518+ 
519+template <typename QSIT>
520+__aicore__ inline void QSIPreload<QSIT>::CalcRunInfo(QSICommon::RunInfo &runInfo)
521+{
522+ runInfo.needProcessS2Size = tempLoopInfo.needProcessS2Size;
523+ runInfo.targetTopKAlign = tempLoopInfo.targetTopKAlign;
524+ runInfo.targetTopK = tempLoopInfo.targetTopK;
525+ runInfo.indiceOutOffset = tempLoopInfo.bN2Idx * constInfo.sparseCount;
526+ int32_t totalNCount = (tempLoopInfo.actS2Size + constInfo.sparseBlockSize - 1) / constInfo.sparseBlockSize;
527+ runInfo.fixedTailCount = constInfo.fixedTailCount > totalNCount ? totalNCount : constInfo.fixedTailCount;
528+}
529+ 
530+template <typename QSIT>
531+__aicore__ inline void QSIPreload<QSIT>::CalcRunInfo(uint32_t loop, uint32_t s2LoopIdx, QSICommon::RunInfo &runInfo)
532+{
533+ runInfo.loop = loop;
534+ runInfo.bIdx = tempLoopInfo.bIdx;
535+ runInfo.s2Idx = s2LoopIdx;
536+ runInfo.bN2Idx = tempLoopInfo.bN2Idx;
537+ runInfo.n2Idx = tempLoopInfo.n2Idx;
538+ 
539+ runInfo.qScale = tempLoopInfo.qScale;
540+ runInfo.actS2Size = tempLoopInfo.actS2Size;
541+ runInfo.needProcessS2Size = tempLoopInfo.needProcessS2Size;
542+ runInfo.targetTopKAlign = tempLoopInfo.targetTopKAlign;
543+ runInfo.targetTopK = tempLoopInfo.targetTopK;
544+ // 计算实际基本块size
545+ runInfo.actMBaseSize = tempLoopInfo.actMBaseSize;
546+ runInfo.fixedTailCount = constInfo.fixedTailCount;
547+ runInfo.actualSingleProcessSInnerSize = constInfo.s2BaseSize;
548+ uint32_t s2SplitNum = (tempLoopInfo.needProcessS2Size + constInfo.s2BaseSize - 1) / constInfo.s2BaseSize;
549+ if (s2SplitNum == 0) {
550+ tempLoopInfo.s2BasicSizeTail = 0;
551+ runInfo.actualSingleProcessSInnerSize = 0;
552+ } else {
553+ tempLoopInfo.s2BasicSizeTail = tempLoopInfo.needProcessS2Size - ((s2SplitNum - 1) * constInfo.s2BaseSize);
554+ if (runInfo.s2Idx == s2SplitNum - 1) {
555+ runInfo.actualSingleProcessSInnerSize = tempLoopInfo.s2BasicSizeTail;
556+ }
557+ }
558+ runInfo.actualSingleProcessSInnerSizeAlign =
559+ QSICommon::Align((uint32_t)runInfo.actualSingleProcessSInnerSize, QSICommon::ConstInfo::BUFFER_SIZE_BYTE_32B);
560+ 
561+ runInfo.isFirstS2InnerLoop = s2LoopIdx == splitCoreInfo.s2Start;
562+ runInfo.isLastS2InnerLoop = s2LoopIdx == tempLoopInfo.s2LoopEnd;
563+ runInfo.isAllLoopEnd = (runInfo.bN2Idx == splitCoreInfo.bN2End) && (runInfo.s2Idx == tempLoopInfo.s2LoopEnd);
564+ 
565+ if (runInfo.isFirstS2InnerLoop) {
566+ actualSeqKPrefixSum = (runInfo.bIdx <= 0) ? 0 : runInfo.bIdx * constInfo.kSeqSize;
567+ uint64_t bsndKeyBIdxOffset = actualSeqKPrefixSum * constInfo.kHeadNum * constInfo.headDim;
568+ // B,N2,D
569+ runInfo.tensorQueryOffset = runInfo.bN2Idx * constInfo.headDim;
570+ 
571+ // bsnd or pa_bsnd
572+ keyCoreOffset = bsndKeyBIdxOffset + runInfo.n2Idx * constInfo.headDim;
573+ // B,N2,k
574+ runInfo.indiceOutOffset = runInfo.bN2Idx * constInfo.sparseCount;
575+ }
576+ runInfo.tensorKeyOffset = keyCoreOffset + runInfo.s2Idx * constInfo.s2BaseSize * constInfo.kHeadNum * constInfo.headDim;
577+ runInfo.tensorKeyScaleOffset = actualSeqKPrefixSum * constInfo.kHeadNum +
578+ runInfo.s2Idx * constInfo.s2BaseSize * constInfo.kHeadNum;
579+}
580+ 
581+template <typename QSIT>
582+__aicore__ inline void QSIPreload<QSIT>::Process()
583+{
584+ if (usedCoreNum == 0) {
585+ // 没有计算任务，直接清理输出
586+ ProcessInvalid();
587+ return;
588+ }
589+ ProcessMain();
590+ ProcessDecode();
591+}
592+ 
593+template <typename QSIT>
594+__aicore__ inline void QSIPreload<QSIT>::ProcessInvalid()
595+{
596+ if ASCEND_IS_AIV {
597+ uint32_t aivCoreNum = GetBlockNum() * 2; // 2 means c:v = 1:2
598+ uint64_t totalOutputSize =
599+ constInfo.batchSize * constInfo.kHeadNum * constInfo.sparseCount;
600+ uint64_t singleCoreSize =
601+ QSICommon::Align((totalOutputSize + aivCoreNum - 1) / aivCoreNum, GM_ALIGN_BYTES / sizeof(OUT_T));
602+ uint64_t baseSize = tmpBlockIdx * singleCoreSize;
603+ if (baseSize < totalOutputSize) {
604+ uint64_t dealSize =
605+ (baseSize + singleCoreSize < totalOutputSize) ? singleCoreSize : totalOutputSize - baseSize;
606+ GlobalTensor<OUT_T> output = indiceOutGm[baseSize];
607+ AscendC::InitGlobalMemory(output, dealSize, constInfo.INVALID_IDX);
608+ }
609+ }
610+}
611+ 
612+template <typename QSIT>
613+__aicore__ inline void QSIPreload<QSIT>::ProcessMain()
614+{
615+ if (aiCoreIdx >= usedCoreNum) {
616+ // 无任务核直接返回
617+ return;
618+ }
619+ 
620+ if ASCEND_IS_AIV {
621+ vectorService.AllocEventID();
622+ CrossCoreSetFlag<QSICommon::ConstInfo::FIA_SYNC_MODE2, PIPE_MTE2>(constInfo.syncV1C1);
623+ CrossCoreSetFlag<QSICommon::ConstInfo::FIA_SYNC_MODE2, PIPE_MTE2>(constInfo.syncV1C1);
624+ } else {
625+ matmulService.AllocEventID();
626+ }
627+ 
628+ QSICommon::RunInfo runInfo;
629+ uint32_t loopIdx = 0;
630+ for (uint32_t bN2LoopIdx = splitCoreInfo.bN2Start; bN2LoopIdx <= splitCoreInfo.bN2End; bN2LoopIdx++) {
631+ CalcS2LoopParams(bN2LoopIdx);
632+ if (tempLoopInfo.needCleanSparseCount && splitCoreInfo.s2Start == 0) {
633+ DealActSeqLenIsZero(tempLoopInfo.bIdx, tempLoopInfo.n2Idx);
634+ }
635+ if (tempLoopInfo.OnlyFixTail) {
636+ CalcRunInfo(runInfo);
637+ if ASCEND_IS_AIV {
638+ WaitFlag<HardEvent::MTE3_V>(0);
639+ vectorService.CopyOutResult(runInfo);
640+ SetFlag<HardEvent::MTE3_V>(0);
641+ }
642+ continue;
643+ }
644+ for (uint32_t s2LoopIdx = splitCoreInfo.s2Start; s2LoopIdx <= tempLoopInfo.s2LoopEnd; s2LoopIdx++) {
645+ ProcessBaseBlock(loopIdx, s2LoopIdx, runInfo);
646+ ++loopIdx;
647+ }
648+ splitCoreInfo.s2Start = 0;
649+ }
650+ 
651+ if ASCEND_IS_AIV {
652+ vectorService.FreeEventID();
653+ } else {
654+ matmulService.FreeEventID();
655+ CrossCoreWaitFlag(constInfo.syncV1C1);
656+ CrossCoreWaitFlag(constInfo.syncV1C1);
657+ }
658+}
659+ 
660+template <typename QSIT>
661+__aicore__ inline void QSIPreload<QSIT>::ProcessBaseBlock(uint32_t loop, uint64_t s2LoopIdx, QSICommon::RunInfo &runInfo)
662+{
663+ CalcRunInfo(loop, s2LoopIdx, runInfo);
664+ if ASCEND_IS_AIC {
665+ matmulService.ComputeMm1(runInfo);
666+ CrossCoreSetFlag<QSICommon::ConstInfo::FIA_SYNC_MODE2, PIPE_FIX>(constInfo.syncC1V1);
667+ } else {
668+ CrossCoreWaitFlag(constInfo.syncC1V1);
669+ vectorService.ProcessVec(runInfo);
670+ CrossCoreSetFlag<QSICommon::ConstInfo::FIA_SYNC_MODE2, PIPE_MTE2>(constInfo.syncV1C1);
671+ }
672+}
673+ 
674+template <typename QSIT>
675+__aicore__ inline void QSIPreload<QSIT>::ProcessDecode()
676+{
677+ if ASCEND_IS_AIV {
678+ SyncAll();
679+ if (splitCoreInfo.isLD) {
680+ vectorService.InitLDBuffers(pipe);
681+ ICachePreLoad(LD_PREFETCH_LEN);
682+ vectorService.ProcessLD();
683+ }
684+ }
685+}
686+} // namespace QSIKernel
687+#endif // QUANT_SALS_INDEXER_KERNEL_H
Aexperimental/attention/quant_sals_indexer/op_kernel/quant_sals_indexer_metadata.h+38-0
@@ -0,0 +1,38 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+
11+/*!
12+ * \file quant_sals_indexer_metadata.h
13+ * \brief
14+ */
15+ 
16+#ifndef QUANT_SALS_INDEXER_METADATA_H
17+#define QUANT_SALS_INDEXER_METADATA_H
18+ 
19+#include <cstdint>
20+ 
21+namespace optiling {
22+constexpr uint32_t CORE_NUM = 24; // 样例代码，当前仅支持核数24
23+constexpr uint32_t QSI_META_SIZE = 1024;
24+using QSI_METADATA_T = int32_t;
25+ 
26+namespace detail {
27+ struct QsiMetaData { // __attribute__((aligned(8)))
M
Mmiaofangzheng4月2日

QsiMetaData 中 bN2End、gS1End、s2End 三个数组（各 24 个 uint32_t）没有初始化，只有 usedCoreNum 有默认值。未使用的核的数组元素包含垃圾值，如果 kernel 侧读取了超过 usedCoreNum 的索引，会得到不可预期的数据。

likedislike
M
mayafei
4月3日 评论：
28+ uint32_t bN2End[CORE_NUM]; // 每个核处理数据的BN2结束点
29+ uint32_t gS1End[CORE_NUM]; // 每个核处理数据的GS1结束点
30+ uint32_t s2End[CORE_NUM]; // 每个核处理数据的S2结束点
31+ uint32_t usedCoreNum = 0U; // 使用的核数量
32+ };
33+};
34+ 
35+static_assert(QSI_META_SIZE * sizeof(QSI_METADATA_T) >= sizeof(detail::QsiMetaData));
36+};
37+ 
38+#endif
Aexperimental/attention/quant_sals_indexer/op_kernel/quant_sals_indexer_service_cube_int4.h+386-0
@@ -0,0 +1,386 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file quant_sals_indexer_service_cube.h
13+ * \brief use 5 buffer for matmul l1, better pipeline
14+ */
15+#ifndef QUANT_SALS_INDEXER_SERVICE_CUBE_H
16+#define QUANT_SALS_INDEXER_SERVICE_CUBE_H
17+ 
18+#include "kernel_operator.h"
19+#include "kernel_operator_list_tensor_intf.h"
20+#include "kernel_tiling/kernel_tiling.h"
21+#include "lib/matmul_intf.h"
22+#include "lib/matrix/matmul/tiling.h"
23+#include "quant_sals_indexer_common.h"
24+ 
25+namespace QSIKernel {
26+using namespace AscendC;
27+using namespace QSICommon;
28+ 
29+template <typename QSIT>
30+class QSIMatmulInt4 {
31+public:
32+ using Q_T = typename QSIT::queryType;
33+ using K_T = typename QSIT::keyType;
34+ 
35+ __aicore__ inline QSIMatmulInt4(){};
36+ __aicore__ inline void InitBuffers(TPipe *pipe);
37+ __aicore__ inline void InitMm1GlobalTensor(const GlobalTensor<int32_t> &blkTableGm, const GlobalTensor<K_T> &keyGm,
38+ const GlobalTensor<Q_T> &queryGm, const GlobalTensor<int32_t> &mm1ResGm);
39+ __aicore__ inline void InitParams(const ConstInfo &constInfo);
40+ __aicore__ inline void AllocEventID();
41+ __aicore__ inline void FreeEventID();
42+ __aicore__ inline void ComputeMm1(const QSICommon::RunInfo &runInfo);
43+ 
44+ static constexpr IsResetLoad3dConfig LOAD3DV2_CONFIG = {true, true}; // isSetFMatrix isSetPadding;
45+ static constexpr uint64_t KEY_BUF_NUM = 3;
46+ static constexpr uint64_t QUERY_BUF_NUM = 2;
47+ static constexpr uint64_t L0_BUF_NUM = 2;
48+ 
49+ static constexpr uint32_t KEY_MTE1_MTE2_EVENT = EVENT_ID2;
50+ static constexpr uint32_t QUERY_MTE1_MTE2_EVENT = EVENT_ID5; // KEY_MTE1_MTE2_EVENT + KEY_BUF_NUM;
51+ static constexpr uint32_t M_MTE1_EVENT = EVENT_ID3;
52+ 
53+ static constexpr uint32_t MTE2_MTE1_EVENT = EVENT_ID2;
54+ static constexpr uint32_t MTE1_M_EVENT = EVENT_ID2;
55+ 
56+ static constexpr uint64_t S8_BLOCK_CUBE = 32;
57+ static constexpr uint64_t S4_BLOCK_CUBE = 64;
58+ static constexpr uint64_t M_BASIC_BLOCK = 16;
59+ static constexpr uint64_t D_BASIC_BLOCK = 64;
60+ static constexpr uint64_t S2_BASIC_BLOCK = 1024;
61+ 
62+ static constexpr uint64_t M_BASIC_BLOCK_L0 = 16;
63+ static constexpr uint64_t D_BASIC_BLOCK_L0 = 64;
64+ static constexpr uint64_t S2_BASIC_BLOCK_L0 = 1024;
65+ 
66+ static constexpr uint64_t QUERY_BUFFER_OFFSET = M_BASIC_BLOCK * D_BASIC_BLOCK;
67+ static constexpr uint64_t KEY_BUFFER_OFFSET = 2048 * D_BASIC_BLOCK;
68+ static constexpr uint64_t L0B_BUFFER_OFFSET = D_BASIC_BLOCK_L0 * S2_BASIC_BLOCK_L0;
69+ static constexpr uint64_t L0A_BUFFER_OFFSET = M_BASIC_BLOCK_L0 * D_BASIC_BLOCK_L0;
70+ static constexpr uint64_t L0C_BUFFER_OFFSET = M_BASIC_BLOCK_L0 * S2_BASIC_BLOCK_L0;
71+ 
72+protected:
73+ __aicore__ inline void Fixp(uint64_t s1gL0RealSize,
74+ uint64_t s2L0RealSize, const QSICommon::RunInfo &runInfo);
75+ __aicore__ inline void ComputeL0c(uint64_t s2L0Offset, uint64_t s1gL0RealSize,
76+ uint64_t s2L0RealSize, const QSICommon::RunInfo &runInfo);
77+ __aicore__ inline void LoadKeyToL0b(uint64_t s2L0RealSize, uint64_t s2L1Offset,
78+ const QSICommon::RunInfo &runInfo);
79+ __aicore__ inline void LoadQueryToL0a(uint64_t s1gL1RealSize,
80+ uint64_t s1gL0RealSize, const QSICommon::RunInfo &runInfo);
81+ __aicore__ inline void QueryNd2Nz(uint64_t s1gL1RealSize, uint64_t s1gL1Offset, const QSICommon::RunInfo &runInfo);
82+ __aicore__ inline void KeyNd2Nz(uint64_t s2L1RealSize,
83+ uint64_t s2L1Offset, uint64_t s2GmOffset, const QSICommon::RunInfo &runInfo);
84+ __aicore__ inline void KeyNd2NzForPA(uint64_t s2L1RealSize,
85+ uint64_t s2L1Offset, uint64_t s2GmOffset, const QSICommon::RunInfo &runInfo);
86+ GlobalTensor<int32_t> blkTableGm_;
87+ GlobalTensor<K_T> keyGm_;
88+ GlobalTensor<Q_T> queryGm_;
89+ GlobalTensor<int32_t> mm1ResGm_;
90+ 
91+ TBuf<TPosition::A1> bufQL1_;
92+ LocalTensor<Q_T> queryL1_;
93+ TBuf<TPosition::B1> bufKeyL1_;
94+ LocalTensor<K_T> keyL1_;
95+ 
96+ TBuf<TPosition::A2> bufQL0_;
97+ LocalTensor<Q_T> queryL0_;
98+ TBuf<TPosition::B2> bufKeyL0_;
99+ LocalTensor<K_T> keyL0_;
100+ 
101+ TBuf<TPosition::CO1> bufL0C_;
102+ LocalTensor<int32_t> cL0_;
103+ 
104+ uint64_t keyL1BufIdx_ = 0;
105+ uint64_t queryL1Mte2BufIdx_ = 0;
106+ uint64_t l0BufIdx_ = 0;
107+ 
108+ ConstInfo constInfo_;
109+ 
110+private:
111+ static constexpr bool PAGE_ATTENTION = QSIT::pageAttention;
112+ static constexpr QSI_LAYOUT K_LAYOUT_T = QSIT::keyLayout;
113+};
114+ 
115+template <typename QSIT>
116+__aicore__ inline void QSIMatmulInt4<QSIT>::InitParams(const ConstInfo &constInfo)
117+{
118+ constInfo_ = constInfo;
119+}
120+ 
121+template <typename QSIT>
122+__aicore__ inline void QSIMatmulInt4<QSIT>::InitBuffers(TPipe *pipe)
123+{
124+ pipe->InitBuffer(bufQL1_, QUERY_BUF_NUM * M_BASIC_BLOCK * D_BASIC_BLOCK * sizeof(Q_T)); // (2) * 16 * 64 * 0.5 = 1K
125+ queryL1_ = bufQL1_.Get<Q_T>();
126+ pipe->InitBuffer(bufKeyL1_, KEY_BUF_NUM * 2048 * D_BASIC_BLOCK * sizeof(K_T)); // (3) * 2048 * 64 * 0.5 = 192K
127+ keyL1_ = bufKeyL1_.Get<K_T>();
128+ 
129+ pipe->InitBuffer(bufQL0_, L0_BUF_NUM * M_BASIC_BLOCK_L0 *
130+ D_BASIC_BLOCK_L0 * sizeof(Q_T)); // (2) * 16 * 64 * 0.5 = 1K
131+ queryL0_ = bufQL0_.Get<Q_T>();
132+ pipe->InitBuffer(bufKeyL0_, L0_BUF_NUM * D_BASIC_BLOCK_L0 *
133+ S2_BASIC_BLOCK_L0 * sizeof(K_T)); // (2) * 64 * 1024 * 0.5 = 64K
134+ keyL0_ = bufKeyL0_.Get<K_T>();
135+ 
136+ pipe->InitBuffer(bufL0C_, L0_BUF_NUM * M_BASIC_BLOCK_L0 *
137+ S2_BASIC_BLOCK_L0 * sizeof(int32_t)); // 16 * 2048 * 4 = 128K, 不开double buffer
138+ cL0_ = bufL0C_.Get<int32_t>();
139+}
140+ 
141+template <typename QSIT>
142+__aicore__ inline void
143+QSIMatmulInt4<QSIT>::InitMm1GlobalTensor(const GlobalTensor<int32_t> &blkTableGm, const GlobalTensor<K_T> &keyGm,
144+ const GlobalTensor<Q_T> &queryGm, const GlobalTensor<int32_t> &mm1ResGm)
145+{
146+ blkTableGm_ = blkTableGm;
147+ keyGm_ = keyGm;
148+ queryGm_ = queryGm;
149+ mm1ResGm_ = mm1ResGm;
150+}
151+ 
152+template <typename QSIT>
153+__aicore__ inline void QSIMatmulInt4<QSIT>::ComputeMm1(const QSICommon::RunInfo &runInfo)
154+{
155+ if (runInfo.isFirstS2InnerLoop) {
156+ WaitFlag<HardEvent::MTE1_MTE2>(QUERY_MTE1_MTE2_EVENT + queryL1Mte2BufIdx_ % QUERY_BUF_NUM);
157+ QueryNd2Nz(1, 0, runInfo);
158+ }
159+ 
160+ uint64_t s2GmBaseOffset = runInfo.s2Idx * constInfo_.s2BaseSize; // 当前基本块在S2方向上的偏移
161+ for (uint64_t s2GmOffset = 0; s2GmOffset < runInfo.actualSingleProcessSInnerSize; s2GmOffset += S2_BASIC_BLOCK) {
162+ // 切N轴: 基本块内S2方向上的偏移，每次处理1024个数，循环2次
163+ WaitFlag<HardEvent::MTE1_MTE2>(KEY_MTE1_MTE2_EVENT + keyL1BufIdx_ % KEY_BUF_NUM);
164+ uint64_t s2L1RealSize =
165+ s2GmOffset + S2_BASIC_BLOCK > runInfo.actualSingleProcessSInnerSize ?
166+ runInfo.actualSingleProcessSInnerSize - s2GmOffset : S2_BASIC_BLOCK; // key矩阵mte2的单次搬运量
167+ if constexpr (PAGE_ATTENTION) {
168+ KeyNd2NzForPA(s2L1RealSize, s2GmOffset, s2GmBaseOffset + s2GmOffset, runInfo);
169+ }else {
170+ KeyNd2Nz(s2L1RealSize, s2GmOffset, s2GmBaseOffset+s2GmOffset, runInfo);
171+ }
172+ 
173+ SetFlag<HardEvent::MTE2_MTE1>(MTE2_MTE1_EVENT);
174+ WaitFlag<HardEvent::MTE2_MTE1>(MTE2_MTE1_EVENT);
175+ 
176+ WaitFlag<HardEvent::M_MTE1>(M_MTE1_EVENT + l0BufIdx_ % L0_BUF_NUM);
177+ LoadQueryToL0a(1, 1, runInfo);
178+ LoadKeyToL0b(s2L1RealSize, s2GmOffset, runInfo);
179+ 
180+ ComputeL0c(s2GmOffset, 1, s2L1RealSize, runInfo);
181+ 
182+ SetFlag<HardEvent::M_MTE1>(M_MTE1_EVENT + l0BufIdx_ % L0_BUF_NUM);
183+ 
184+ l0BufIdx_++;
185+ 
186+ SetFlag<HardEvent::MTE1_MTE2>(KEY_MTE1_MTE2_EVENT + keyL1BufIdx_ % KEY_BUF_NUM);
187+ keyL1BufIdx_++;
188+ }
189+ CrossCoreWaitFlag(constInfo_.syncV1C1);
190+ Fixp(1, runInfo.actualSingleProcessSInnerSize, runInfo);
191+ if (runInfo.isLastS2InnerLoop) {
192+ SetFlag<HardEvent::MTE1_MTE2>(QUERY_MTE1_MTE2_EVENT + queryL1Mte2BufIdx_ % QUERY_BUF_NUM);
193+ queryL1Mte2BufIdx_++;
194+ }
195+}
196+ 
197+// bsnd
198+template <typename QSIT>
199+__aicore__ inline void QSIMatmulInt4<QSIT>::KeyNd2Nz(uint64_t s2L1RealSize, uint64_t s2L1Offset, uint64_t s2GmOffset,
200+ const QSICommon::RunInfo &runInfo)
201+{
202+ // DMA BSND
203+ DataCopyParams copyInParams;
204+ if (constInfo_.kHeadNum == 1) {
205+ copyInParams.blockCount = 1; // 待搬运的连续传输数据块个数
206+ copyInParams.blockLen = s2L1RealSize * constInfo_.headDim / S4_BLOCK_CUBE; // 待搬运的每个连续传输数据块长度，单位为DataBlock(32字节)
207+ copyInParams.srcStride = 0;
208+ } else {
209+ copyInParams.blockCount = s2L1RealSize; // 待搬运的连续传输数据块个数
210+ copyInParams.blockLen = constInfo_.headDim / S4_BLOCK_CUBE; // 待搬运的每个连续传输数据块长度，单位为DataBlock(32字节)
211+ copyInParams.srcStride = (constInfo_.kHeadNum - 1) * constInfo_.headDim / S4_BLOCK_CUBE;
212+ }
213+ copyInParams.dstStride = 0;
214+ DataCopy(keyL1_[(keyL1BufIdx_ % KEY_BUF_NUM) * KEY_BUFFER_OFFSET + s2L1Offset * constInfo_.headDim],
215+ keyGm_[runInfo.tensorKeyOffset + s2GmOffset * constInfo_.kHeadNum * constInfo_.headDim], copyInParams);
216+}
217+ 
218+// blkNum, blkSize, N2, D
219+template <typename QSIT>
220+__aicore__ inline void QSIMatmulInt4<QSIT>::KeyNd2NzForPA(uint64_t s2L1RealSize, uint64_t s2L1BaseOffset, uint64_t s2GmOffset,
221+ const QSICommon::RunInfo &runInfo)
222+{
223+ uint64_t s2L1Offset = 0; // 已经搬运了多少
224+ while (s2L1Offset < s2L1RealSize) {
225+ uint64_t s2BlkId = (s2L1Offset + s2GmOffset) / constInfo_.kCacheBlockSize; // 当前batch中blockTable中PA block ID
226+ uint64_t s2BlkOffset = (s2L1Offset + s2GmOffset) % constInfo_.kCacheBlockSize; // PA block中的偏移
227+ uint64_t keyGmOffset = 0;
228+ uint64_t s2Mte2Size = s2L1RealSize - s2L1Offset; // 本次PA搬运量
229+ s2Mte2Size = s2BlkOffset + s2Mte2Size >= constInfo_.kCacheBlockSize ? constInfo_.kCacheBlockSize - s2BlkOffset :
230+ s2Mte2Size; // 当前搬运量是否跨PA block，跨block时分多次搬运
231+ DataCopyParams copyInParams;
232+ if constexpr (K_LAYOUT_T == QSI_LAYOUT::PA_BNSD) { // PA_BNSD(blockNum, n2, blockSize, d)
233+ keyGmOffset = blkTableGm_.GetValue(runInfo.bIdx * constInfo_.maxBlockNumPerBatch + s2BlkId) *
234+ constInfo_.kHeadNum * constInfo_.kCacheBlockSize * constInfo_.headDim +
235+ runInfo.n2Idx * constInfo_.kCacheBlockSize * constInfo_.headDim +
236+ s2BlkOffset * constInfo_.headDim;
237+ // DMA
238+ copyInParams.blockCount = 1; // 待搬运的连续传输数据块个数
239+ copyInParams.blockLen = s2Mte2Size * constInfo_.headDim / S4_BLOCK_CUBE; // 待搬运的每个连续传输数据块长度，单位为DataBlock(32字节)
240+ copyInParams.srcStride = 0;
241+ copyInParams.dstStride = 0;
242+ } else if constexpr (K_LAYOUT_T == QSI_LAYOUT::PA_NZ) {
243+ uint32_t blockElementCntNZ = 32 / sizeof(K_T);
244+ if constexpr (IsSameType<K_T, int4b_t>::value) {
245+ blockElementCntNZ = 64;
246+ }
247+ keyGmOffset = blkTableGm_.GetValue(runInfo.bIdx * constInfo_.maxBlockNumPerBatch + s2BlkId) *
248+ constInfo_.kHeadNum * constInfo_.kCacheBlockSize * constInfo_.headDim +
249+ runInfo.n2Idx * constInfo_.kCacheBlockSize * constInfo_.headDim +
250+ s2BlkOffset * blockElementCntNZ;
251+ copyInParams.blockLen = s2Mte2Size; // 待搬运的每个连续传输数据块长度，单位为DataBlock(32字节)
252+ copyInParams.blockCount = constInfo_.headDim / blockElementCntNZ; // 待搬运的连续传输数据块个数
253+ copyInParams.dstStride = 0;
254+ copyInParams.srcStride = constInfo_.kCacheBlockSize - s2Mte2Size;
255+ } else { // PA_BSND(blockNum, blockSize, n2, d)
256+ keyGmOffset = blkTableGm_.GetValue(runInfo.bIdx * constInfo_.maxBlockNumPerBatch + s2BlkId) *
257+ constInfo_.kCacheBlockSize * constInfo_.kHeadNum * constInfo_.headDim +
258+ s2BlkOffset * constInfo_.kHeadNum * constInfo_.headDim + runInfo.n2Idx * constInfo_.headDim;
259+ // DMA
260+ copyInParams.blockCount = s2Mte2Size; // 待搬运的连续传输数据块个数
261+ copyInParams.blockLen = constInfo_.headDim / S4_BLOCK_CUBE; // 待搬运的每个连续传输数据块长度，单位为DataBlock(32字节)
262+ copyInParams.srcStride = (constInfo_.kHeadNum - 1) * constInfo_.headDim / S4_BLOCK_CUBE;
263+ copyInParams.dstStride = 0;
264+ }
265+ DataCopy(keyL1_[(keyL1BufIdx_ % KEY_BUF_NUM) * KEY_BUFFER_OFFSET + s2L1BaseOffset * constInfo_.headDim + s2L1Offset * constInfo_.headDim],
266+ keyGm_[keyGmOffset], copyInParams);
267+ s2L1Offset += s2Mte2Size;
268+ }
269+}
270+ 
271+// batch, s1, n2, g, d
272+template <typename QSIT>
273+__aicore__ inline void QSIMatmulInt4<QSIT>::QueryNd2Nz(uint64_t s1gL1RealSize, uint64_t s1gGmOffset,
274+ const QSICommon::RunInfo &runInfo)
275+{
276+ // m=1，按照gemv方式优化
277+ DataCopyParams copyInParams;
278+ copyInParams.blockCount = 1; // 待搬运的连续传输数据块个数
279+ copyInParams.blockLen = constInfo_.headDim / S4_BLOCK_CUBE; // 待搬运的每个连续传输数据块长度，单位为DataBlock(32字节)
280+ copyInParams.srcStride = 0;
281+ copyInParams.dstStride = 0;
282+ DataCopy(queryL1_[(queryL1Mte2BufIdx_ % QUERY_BUF_NUM) * QUERY_BUFFER_OFFSET],
283+ queryGm_[runInfo.tensorQueryOffset + s1gGmOffset * constInfo_.headDim], copyInParams);
284+}
285+ 
286+// 1, d
287+template <typename QSIT>
288+__aicore__ inline void QSIMatmulInt4<QSIT>::LoadQueryToL0a(uint64_t s1gL1RealSize,
289+ uint64_t s1gL0RealSize, const QSICommon::RunInfo &runInfo)
290+{
291+ // m=1，按照gemv方案优化
292+ LoadData2DParams loadData2DParams;
293+ loadData2DParams.startIndex = 0;
294+ loadData2DParams.repeatTimes = 1;
295+ loadData2DParams.srcStride = 1;
296+ loadData2DParams.dstGap = 0;
297+ loadData2DParams.ifTranspose = false;
298+ LoadData(queryL0_[(l0BufIdx_ % L0_BUF_NUM) * L0A_BUFFER_OFFSET],
299+ queryL1_[(queryL1Mte2BufIdx_ % QUERY_BUF_NUM) * QUERY_BUFFER_OFFSET], loadData2DParams);
300+}
301+ 
302+template <typename QSIT>
303+__aicore__ inline void QSIMatmulInt4<QSIT>::LoadKeyToL0b(uint64_t s2L0RealSize, uint64_t s2L1Offset,
304+ const QSICommon::RunInfo &runInfo)
305+{
306+ LoadData2DParams loadData2DParams;
307+ loadData2DParams.startIndex = 0;
308+ loadData2DParams.repeatTimes = QSiCeilDiv(s2L0RealSize, static_cast<uint64_t>(BLOCK_CUBE)) * QSiCeilDiv(constInfo_.headDim, static_cast<uint64_t>(S4_BLOCK_CUBE));
309+ loadData2DParams.srcStride = 1;
310+ loadData2DParams.dstGap = 0;
311+ loadData2DParams.ifTranspose = false;
312+ LoadData(keyL0_[(l0BufIdx_ % L0_BUF_NUM) * L0B_BUFFER_OFFSET],
313+ keyL1_[(keyL1BufIdx_ % KEY_BUF_NUM) * KEY_BUFFER_OFFSET + s2L1Offset * constInfo_.headDim], loadData2DParams);
314+}
315+ 
316+template <typename QSIT>
317+__aicore__ inline void QSIMatmulInt4<QSIT>::ComputeL0c(uint64_t s2L0Offset, uint64_t s1gL0RealSize, uint64_t s2L0RealSize,
318+ const QSICommon::RunInfo &runInfo)
319+{
320+ SetFlag<HardEvent::MTE1_M>(MTE1_M_EVENT);
321+ WaitFlag<HardEvent::MTE1_M>(MTE1_M_EVENT);
322+ MmadParams mmadParams;
323+ mmadParams.m = 1;
324+ mmadParams.n = s2L0RealSize;
325+ mmadParams.k = constInfo_.headDim;
326+ mmadParams.cmatrixInitVal = true;
327+ mmadParams.cmatrixSource = false;
328+ mmadParams.unitFlag = 0b11;
329+ Mmad(cL0_.template ReinterpretCast<int32_t>()[BLOCK_CUBE * s2L0Offset],
330+ queryL0_.template ReinterpretCast<int4b_t>()[(l0BufIdx_ % L0_BUF_NUM) * L0A_BUFFER_OFFSET],
331+ keyL0_.template ReinterpretCast<int4b_t>()[(l0BufIdx_ % L0_BUF_NUM) * L0B_BUFFER_OFFSET], mmadParams);
332+ if ((mmadParams.m / 16) * (mmadParams.n / 16) < 10) {
333+ PipeBarrier<PIPE_M>();
334+ }
335+}
336+ 
337+template <typename QSIT>
338+__aicore__ inline void QSIMatmulInt4<QSIT>::Fixp( uint64_t s1gL0RealSize,
339+ uint64_t s2L0RealSize, const QSICommon::RunInfo &runInfo)
340+{
341+ AscendC::DataCopyCO12DstParams intriParams;
342+ intriParams.mSize = 1;
343+ intriParams.nSize = s2L0RealSize;
344+ intriParams.dstStride = runInfo.actualSingleProcessSInnerSizeAlign; // 使能NZ2ND功能，dst同一ND矩阵的相邻行的偏移(头与头)，取值不为0，单位为元素
345+ intriParams.srcStride = QSiCeilAlign(s1gL0RealSize, static_cast<uint64_t>(BLOCK_CUBE)); // 使能NZ2ND功能，src同一NZ矩阵的相邻Z排布的偏移(头与头)，必须为16的倍数，取值范围：srcStride∈[0, 65535], 单位C0_size
346+ // set mode according to dtype
347+ intriParams.quantPre = QuantMode_t::NoQuant;
348+ intriParams.nz2ndEn = true;
349+ intriParams.unitFlag = 0b11; // 3 unitflag
350+ intriParams.reluPre = 0;
351+ AscendC::SetFixpipeNz2ndFlag(1, 1, 1);
352+ AscendC::DataCopy(mm1ResGm_[(runInfo.loop % 2) * constInfo_.mBaseSize * constInfo_.s2BaseSize],
353+ cL0_, intriParams);
354+}
355+ 
356+template <typename QSIT>
357+__aicore__ inline void QSIMatmulInt4<QSIT>::AllocEventID()
358+{
359+ SetMMLayoutTransform(true);
360+ SetFlag<HardEvent::MTE1_MTE2>(KEY_MTE1_MTE2_EVENT + 0);
361+ SetFlag<HardEvent::MTE1_MTE2>(KEY_MTE1_MTE2_EVENT + 1);
362+ SetFlag<HardEvent::MTE1_MTE2>(KEY_MTE1_MTE2_EVENT + 2);
363+ 
364+ SetFlag<HardEvent::MTE1_MTE2>(QUERY_MTE1_MTE2_EVENT + 0);
365+ SetFlag<HardEvent::MTE1_MTE2>(QUERY_MTE1_MTE2_EVENT + 1);
366+ 
367+ SetFlag<HardEvent::M_MTE1>(M_MTE1_EVENT + 0);
368+ SetFlag<HardEvent::M_MTE1>(M_MTE1_EVENT + 1);
369+}
370+ 
371+template <typename QSIT>
372+__aicore__ inline void QSIMatmulInt4<QSIT>::FreeEventID()
373+{
374+ SetMMLayoutTransform(false);
375+ WaitFlag<HardEvent::MTE1_MTE2>(KEY_MTE1_MTE2_EVENT + 0);
376+ WaitFlag<HardEvent::MTE1_MTE2>(KEY_MTE1_MTE2_EVENT + 1);
377+ WaitFlag<HardEvent::MTE1_MTE2>(KEY_MTE1_MTE2_EVENT + 2);
378+ 
379+ WaitFlag<HardEvent::MTE1_MTE2>(QUERY_MTE1_MTE2_EVENT + 0);
380+ WaitFlag<HardEvent::MTE1_MTE2>(QUERY_MTE1_MTE2_EVENT + 1);
381+ 
382+ WaitFlag<HardEvent::M_MTE1>(M_MTE1_EVENT + 0);
383+ WaitFlag<HardEvent::M_MTE1>(M_MTE1_EVENT + 1);
384+}
385+} // namespace QSIKernel
386+#endif
Aexperimental/attention/quant_sals_indexer/op_kernel/quant_sals_indexer_service_vector.h+880-0
@@ -0,0 +1,880 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file quant_sals_indexer_service_vector.h
13+ * \brief
14+ */
15+#ifndef QUANT_SALS_INDEXER_SERVICE_VECTOR_H
16+#define QUANT_SALS_INDEXER_SERVICE_VECTOR_H
17+ 
18+#include "kernel_operator.h"
19+#include "kernel_operator_list_tensor_intf.h"
20+#include "kernel_tiling/kernel_tiling.h"
21+#include "lib/matmul_intf.h"
22+#include "lib/matrix/matmul/tiling.h"
23+#include "quant_sals_indexer_common.h"
24+//#include "quant_sals_indexer_vector.h"
25+ 
26+namespace QSIKernel {
27+using namespace AscendC;
28+using namespace QSICommon;
29+ 
30+template <typename QSIT>
31+class QSIVector {
32+public:
33+ // =================================类型定义区=================================
34+ // 中间计算数据类型为float，高精度模式
35+ using K_T = typename QSIT::keyType;
36+ static constexpr QSI_LAYOUT LAYOUT_T = QSIT::layout;
37+ 
38+ // MM输出数据类型, 当前只支持float
39+ using MM1_OUT_T = int32_t;
40+ 
41+ __aicore__ inline QSIVector(){};
42+ __aicore__ inline void ProcessVec(const QSICommon::RunInfo &info);
43+ __aicore__ inline void ProcessLD();
44+ __aicore__ inline void InitBuffers(TPipe *pipe);
45+ __aicore__ inline void InitParams(const struct QSICommon::ConstInfo &constInfo,
46+ const QSITilingData *__restrict tilingData);
47+ __aicore__ inline void InitVec1GlobalTensor(GlobalTensor<MM1_OUT_T> mm1ResGm,
48+ GlobalTensor<float> qScaleGm,
49+ GlobalTensor<float> kScaleGm,
50+ GlobalTensor<int32_t> blkTableGm,
51+ GlobalTensor<float> vec1ResGm,
52+ GlobalTensor<int64_t> vec1ParamGm,
53+ GlobalTensor<int32_t> indiceOutGm);
54+ __aicore__ inline void CleanInvalidOutput(int64_t invalidS1offset);
55+ __aicore__ inline void AllocEventID();
56+ __aicore__ inline void FreeEventID();
57+ __aicore__ inline void InitLDBuffers(TPipe *pipe);
58+ __aicore__ inline void CopyMm1ResultIn(const QSICommon::RunInfo &runInfo);
59+ __aicore__ inline void GetKeyScale(const QSICommon::RunInfo &runInfo);
60+ __aicore__ inline void CopyOut(const GlobalTensor<int32_t> &dstGm,
61+ const LocalTensor<int32_t> &srcUb, int64_t copyCount);
62+ __aicore__ inline void ComputeLse(const QSICommon::RunInfo &runInfo);
63+ __aicore__ inline void SortSubTopK(const QSICommon::RunInfo &runInfo);
64+ __aicore__ inline void SortAll(const LocalTensor<float> &dst, const LocalTensor<float> &src, int64_t logitsNum, int64_t mrgElements);
65+ __aicore__ inline void SortBasicBlockTopKToSub(const QSICommon::RunInfo &runInfo);
66+ __aicore__ inline void CopyOutResult(const QSICommon::RunInfo &runInfo);
67+ __aicore__ inline void CopyOutFdResult(const QSICommon::RunInfo &runInfo);
68+ __aicore__ inline void ExtractIndex(const LocalTensor<int32_t> &idxULocal, const LocalTensor<int32_t> &sortLocal,
69+ int64_t extractNum);
70+protected:
71+ GlobalTensor<MM1_OUT_T> mm1ResGm_;
72+ GlobalTensor<float> qScaleGm_;
73+ GlobalTensor<float> kScaleGm_;
74+ GlobalTensor<int32_t> blkTableGm_;
75+ GlobalTensor<float> vec1ResGm;
76+ GlobalTensor<int64_t> vec1ParamGm_;
77+ GlobalTensor<int32_t> indiceOutGm_;
78+ // =================================常量区=================================
79+ 
80+private:
81+ // ================================Local Buffer区====================================
82+ // tmp buff for vector1
83+ TBuf<> outQueue_;
84+ TBuf<> sortOutBuf_;
85+ TBuf<> reduceOutBuf_;
86+ TBuf<> brcBuf_;
87+ TBuf<> paramBuf_;
88+ 
89+ // tmp buff for LD
90+ TBuf<> ldToBeMrgBuf_;
91+ TBuf<> ldTmpBuf_;
92+ TBuf<> ldOutValueBuf_;
93+ TBuf<> ldOutIdxBuf_;
94+ 
95+ LocalTensor<float> mmInUb_;
96+ LocalTensor<int32_t> globalTopkIndice_;
97+ LocalTensor<float> globalTopkUb_;
98+ LocalTensor<float> SortedBasicBlock_;
99+ LocalTensor<float> qScaleUb_;
100+ LocalTensor<float> kScaleUb_;
101+ LocalTensor<float> qkScaleUb_;
102+ 
103+ LocalTensor<float> nBlkMaxUb_;
104+ LocalTensor<float> nMaxUb_;
105+ LocalTensor<float> nMaxBrcbUb_;
106+ LocalTensor<float> nLseUb_;
107+ LocalTensor<float> nLseTobeSortUb_;
108+ LocalTensor<float> sortTensor_;
109+ LocalTensor<float> nSortOutUb_;
110+ LocalTensor<int32_t> nIdxUb_;
111+ 
112+ int32_t blockId_ = -1;
113+ // para for vector
114+ int32_t groupInner_ = 0;
115+ int32_t globalTopkNum_ = 0;
116+ int64_t blockS2StartIdx_ = 0;
117+ int32_t gSize_ = 0;
118+ int32_t kHeadNum_ = 0;
119+ int64_t mte2BufIdx_ = 0;
120+ int64_t mte3BufIdx_ = 0;
121+ int64_t nSubSortCountOffset_ = 0;
122+ int64_t nSortCountOffset_ = 0;
123+ int64_t sortedIdx_ = 0;
124+ 
125+ // para for LD
126+ uint32_t mrgListNum_ = 4;
127+ constexpr static uint32_t paramNum_ = 16;
128+ 
129+ constexpr static int64_t MTE2_BUF_SIZE = 2048;
130+ constexpr static int64_t SORT_BUF_SIZE = 64 * 256;
131+ constexpr static int64_t MAX_TOPK_WITH_ID = 2048 * 2;
132+ constexpr static int64_t DOUBLE_BUFFER_NUM = 2;
133+ constexpr static int64_t B32_VEC_BLK_NUM = 8;
134+ constexpr static int64_t B32_VEC_MAX_MASK = 64;
135+ constexpr static int64_t SORT_VALUE_ID_SIZE = 8;
136+ constexpr static int64_t SORT_MASK = 32;
137+ 
138+ constexpr static int32_t NEG_INF = 0xFF800000;
139+ constexpr static uint32_t REDUCE_BANK_CONFLICT_OFFSETS = 256;
140+ constexpr static uint32_t REDUCE_BANK_CONFLICT_NUM = REDUCE_BANK_CONFLICT_OFFSETS / sizeof(float);
141+ 
142+ constexpr static uint32_t V_MTE2_EVENT = 0;
143+ constexpr static uint32_t V_MTE2_EVENT_FD = 2;
144+ constexpr static uint32_t S_MTE2_EVENT = 0;
145+ 
146+ constexpr static uint32_t V_MTE3_EVENT = 0;
147+ constexpr static uint32_t S_MTE3_EVENT = 0;
148+ constexpr static uint32_t V_S_EVENT = 0;
149+ constexpr static uint32_t MTE3_S_EVENT = 0;
150+ 
151+ constexpr static uint32_t MTE2_V_EVENT = 0;
152+ constexpr static uint32_t MTE3_V_EVENT = 0;
153+ constexpr static uint32_t MTE3_V_EVENT_FD = 2;
154+ 
155+ constexpr static int64_t MRG_QUE_0 = 0;
156+ constexpr static int64_t MRG_QUE_1 = 1;
157+ constexpr static int64_t MRG_QUE_2 = 2;
158+ constexpr static int64_t MRG_QUE_3 = 3;
159+ constexpr static int64_t MRG_BLOCK_2 = 2;
160+ constexpr static int64_t MRG_BLOCK_3 = 3;
161+ constexpr static int64_t MRG_BLOCK_4 = 4;
162+ constexpr static int64_t VALUE_AND_INDEX_NUM = 2;
163+ static constexpr bool PAGE_ATTENTION = QSIT::pageAttention;
164+ static constexpr QSI_LAYOUT K_LAYOUT_T = QSIT::keyLayout;
165+ 
166+ struct QSICommon::ConstInfo constInfo_;
167+};
168+ 
169+template <typename QSIT>
170+__aicore__ inline void QSIVector<QSIT>::InitBuffers(TPipe *pipe)
171+{
172+ uint32_t outNeedBufSize = (MAX_TOPK * 2) * 2 * sizeof(float);
173+ uint32_t reduceCacheSize = REDUCE_BANK_CONFLICT_OFFSETS + groupInner_ * constInfo_.s2BaseSize * sizeof(float);
174+ outNeedBufSize = reduceCacheSize > outNeedBufSize ? reduceCacheSize : outNeedBufSize;
175+ 
176+ // 1, s2 -> 1 * 2048 输入开db 16k
177+ TBuf<> inputBuf;
178+ pipe->InitBuffer(inputBuf, DOUBLE_BUFFER_NUM * constInfo_.s2BaseSize * sizeof(float));
179+ mmInUb_ = inputBuf.Get<float>();
180+ 
181+ // topk的排序结果, mrgSort最多4条排序队列，空间预留128k
182+ TBuf<> sortBuf;
183+ pipe->InitBuffer(sortBuf, DOUBLE_BUFFER_NUM * MAX_TOPK_WITH_ID * 4 * sizeof(float));
184+ sortTensor_ = sortBuf.Get<float>();
185+ Duplicate(sortTensor_.template ReinterpretCast<int32_t>(), NEG_INF, 2 * SORT_BUF_SIZE);
186+ // 基本块的排序索引缓存, 最多支持256个待排序索引，空间预留1k
187+ TBuf<> indexBuf;
188+ pipe->InitBuffer(indexBuf, 256 * sizeof(int32_t));
189+ globalTopkIndice_ = indexBuf.Get<int32_t>();
190+ 
191+ TBuf<> tmpBuf;
192+ pipe->InitBuffer(tmpBuf, 46*1024);
193+ nBlkMaxUb_ = tmpBuf.Get<float>(); //8k
194+ qScaleUb_ = tmpBuf.Get<float>(); // 8k, 复用nBlkMaxUb_
195+ qkScaleUb_ = tmpBuf.Get<float>(); // 8k, 复用nBlkMaxUb_
196+ nMaxBrcbUb_ =tmpBuf.Get<float>()[8 * 256]; //8k;
197+ kScaleUb_ = tmpBuf.Get<float>()[8 * 256]; // 8k, 复用nMaxBrcbUb_
198+ nMaxUb_ =tmpBuf.Get<float>()[16 * 256]; //1k;
199+ nLseTobeSortUb_ =tmpBuf.Get<float>()[17 * 256]; //2k;
200+ 
201+ nSortOutUb_ =tmpBuf.Get<float>()[19 * 256]; //16k;
202+ nIdxUb_ =tmpBuf.Get<int32_t>()[36 * 256]; //8k;
203+ pipe->InitBuffer(paramBuf_, 1024);
204+ 
205+ ArithProgression<int32_t>(globalTopkIndice_, 0, 1,
206+ QSiCeilDiv(constInfo_.s2BaseSize, static_cast<uint32_t>(constInfo_.sparseBlockSize)));
207+ 
208+ LocalTensor<int32_t> tmpfBuff = sortBuf.Get<int32_t>();
209+ Duplicate(tmpfBuff, -1, 2 * (constInfo_.s1BaseSize / 2) * paramNum_ * 2);
210+ SetFlag<HardEvent::V_MTE3>(V_MTE3_EVENT);
211+ WaitFlag<HardEvent::V_MTE3>(V_MTE3_EVENT);
212+ int64_t wsInfoOffset = blockId_ * constInfo_.s1BaseSize * 2 * paramNum_;
213+ DataCopyPad(vec1ParamGm_[wsInfoOffset], tmpfBuff.template ReinterpretCast<int64_t>(),
214+ {1, static_cast<uint16_t>(constInfo_.s1BaseSize * 2 * paramNum_ * sizeof(int64_t)), 0, 0});
215+ SetFlag<HardEvent::MTE3_V>(MTE3_V_EVENT);
216+ WaitFlag<HardEvent::MTE3_V>(MTE3_V_EVENT);
217+}
218+ 
219+template <typename QSIT>
220+__aicore__ inline void QSIVector<QSIT>::InitLDBuffers(TPipe *pipe)
221+{
222+ pipe->Reset();
223+ pipe->InitBuffer(ldToBeMrgBuf_, 2 * MAX_TOPK * mrgListNum_ * sizeof(float)); // 2：value + index
224+ pipe->InitBuffer(ldTmpBuf_, 2 * MAX_TOPK * mrgListNum_ * sizeof(float)); // 2：value + index
225+ pipe->InitBuffer(ldOutValueBuf_, MAX_TOPK * sizeof(float));
226+ pipe->InitBuffer(ldOutIdxBuf_, MAX_TOPK * sizeof(int32_t) + 32);
227+}
228+ 
229+template <typename QSIT>
230+__aicore__ inline void QSIVector<QSIT>::InitParams(const struct QSICommon::ConstInfo &constInfo,
231+ const QSITilingData *__restrict tilingData)
232+{
233+ constInfo_ = constInfo;
234+ blockS2StartIdx_ = 0;
235+ // define N2 para
236+ kHeadNum_ = constInfo_.kHeadNum;
237+ 
238+ // group ub 切分因子当前按照UB空间强制为16
239+ groupInner_ = 16;
240+ 
241+ blockId_ = GetBlockIdx();
242+}
243+ 
244+template <typename QSIT>
245+__aicore__ inline void
246+QSIVector<QSIT>::InitVec1GlobalTensor(GlobalTensor<MM1_OUT_T> mm1ResGm,
247+ GlobalTensor<float> qScaleGm,
248+ GlobalTensor<float> kScaleGm,
249+ GlobalTensor<int32_t> blkTableGm,
250+ GlobalTensor<float> vec1ResGm,
251+ GlobalTensor<int64_t> vec1ParamGm,
252+ GlobalTensor<int32_t> indiceOutGm)
253+{
254+ mm1ResGm_ = mm1ResGm;
255+ qScaleGm_ = qScaleGm;
256+ blkTableGm_ = blkTableGm;
257+ kScaleGm_ = kScaleGm;
258+ this->vec1ResGm = vec1ResGm;
259+ vec1ParamGm_ = vec1ParamGm;
260+ indiceOutGm_ = indiceOutGm;
261+}
262+ 
263+template <typename QSIT>
264+__aicore__ inline void QSIVector<QSIT>::AllocEventID()
265+{
266+ SetFlag<HardEvent::V_MTE2>(V_MTE2_EVENT + 0);
267+ SetFlag<HardEvent::V_MTE2>(V_MTE2_EVENT + 1);
268+ 
269+ SetFlag<HardEvent::MTE3_V>(MTE3_V_EVENT + 0);
270+ SetFlag<HardEvent::MTE3_V>(MTE3_V_EVENT + 1);
271+}
272+ 
273+template <typename QSIT>
274+__aicore__ inline void QSIVector<QSIT>::FreeEventID()
275+{
276+ WaitFlag<HardEvent::V_MTE2>(V_MTE2_EVENT + 0);
277+ WaitFlag<HardEvent::V_MTE2>(V_MTE2_EVENT + 1);
278+ 
279+ WaitFlag<HardEvent::MTE3_V>(MTE3_V_EVENT + 0);
280+ WaitFlag<HardEvent::MTE3_V>(MTE3_V_EVENT + 1);
281+}
282+ 
283+template <typename QSIT>
284+__aicore__ inline void QSIVector<QSIT>::CopyOut(const GlobalTensor<int32_t> &dstGm,
285+ const LocalTensor<int32_t> &srcUb, int64_t copyCount)
286+{
287+ AscendC::DataCopyParams dataCopyOutyParams;
288+ dataCopyOutyParams.blockCount = 1;
289+ dataCopyOutyParams.blockLen = copyCount * sizeof(int32_t);
290+ dataCopyOutyParams.srcStride = 0;
291+ dataCopyOutyParams.dstStride = 0;
292+ AscendC::DataCopyPad(dstGm, srcUb, dataCopyOutyParams);
293+}
294+ 
295+template <typename QSIT>
296+__aicore__ inline void QSIVector<QSIT>::CopyMm1ResultIn(const QSICommon::RunInfo &runInfo)
297+{
298+ // 将MMout_gmoffset copy到UB上
299+ AscendC::DataCopyPadExtParams<int32_t> Mm1padParams{false, 0, 0, 0};
300+ AscendC::DataCopyExtParams Mm1dataCopymMoutParams;
301+ Mm1dataCopymMoutParams.blockCount = 1;
302+ Mm1dataCopymMoutParams.blockLen = runInfo.actualSingleProcessSInnerSizeAlign * sizeof(int32_t);
303+ Mm1dataCopymMoutParams.srcStride = 0;
304+ Mm1dataCopymMoutParams.dstStride = 0;
305+ Mm1dataCopymMoutParams.rsv = 0;
306+ int64_t mmGmOffset = (runInfo.loop % 2) * constInfo_.mBaseSize * constInfo_.s2BaseSize;
307+ AscendC::DataCopyPad(mmInUb_.template ReinterpretCast<int32_t>()[mte2BufIdx_ % DOUBLE_BUFFER_NUM * MTE2_BUF_SIZE],
308+ mm1ResGm_[mmGmOffset], Mm1dataCopymMoutParams, Mm1padParams);
309+}
310+ 
311+template <typename QSIT>
312+__aicore__ inline void QSIVector<QSIT>::GetKeyScale(const QSICommon::RunInfo &runInfo)
313+{
314+ // 读取kScale
315+ AscendC::DataCopyPadExtParams<float> padParams{false, 0, 0, 0};
316+ AscendC::DataCopyExtParams dataCopymMoutParams;
317+ if constexpr (PAGE_ATTENTION) {
318+ int32_t startS2 = runInfo.s2Idx * constInfo_.s2BaseSize;
319+ int32_t curS2Len = runInfo.actualSingleProcessSInnerSize;
320+ int32_t startBlockTableIdx = startS2 / constInfo_.kCacheBlockSize;
321+ int32_t startBlockTableOffset = startS2 % constInfo_.kCacheBlockSize;
322+ int32_t blockTableBatchOffset = runInfo.bIdx * constInfo_.maxBlockNumPerBatch;
323+ if constexpr (K_LAYOUT_T == QSI_LAYOUT::PA_BNSD || K_LAYOUT_T == QSI_LAYOUT::PA_NZ) {
324+ dataCopymMoutParams.blockCount = 1;
325+ dataCopymMoutParams.srcStride = 0;
326+ } else {
327+ dataCopymMoutParams.blockCount = 1;
328+ dataCopymMoutParams.srcStride = 0;
329+ }
330+ dataCopymMoutParams.dstStride = 0;
331+ dataCopymMoutParams.rsv = 0;
332+ int32_t resUbBaseOffset = 0;
333+ if (startBlockTableOffset > 0) {
334+ int32_t firstPartLen =
335+ constInfo_.kCacheBlockSize - startBlockTableOffset > curS2Len ? curS2Len :
336+ constInfo_.kCacheBlockSize - startBlockTableOffset;
337+ dataCopymMoutParams.blockLen = firstPartLen * sizeof(float);
338+ uint64_t keyScaleGmOffset = 0;
339+ if constexpr (K_LAYOUT_T == QSI_LAYOUT::PA_BNSD || K_LAYOUT_T == QSI_LAYOUT::PA_NZ) { // PA_BNSD(blockNum, n2, blockSize)
340+ keyScaleGmOffset = blkTableGm_.GetValue(blockTableBatchOffset + startBlockTableIdx) *
341+ constInfo_.kHeadNum * constInfo_.kCacheBlockSize +
342+ runInfo.n2Idx * constInfo_.kCacheBlockSize + startBlockTableOffset;
343+ } else { // PA_BSND(blockNum, blockSize, n2)
344+ keyScaleGmOffset = blkTableGm_.GetValue(blockTableBatchOffset + startBlockTableIdx) *
345+ constInfo_.kCacheBlockSize * constInfo_.kHeadNum +
346+ startBlockTableOffset * constInfo_.kHeadNum + runInfo.n2Idx;
347+ }
348+ SetFlag<HardEvent::S_MTE2>(S_MTE2_EVENT);
349+ WaitFlag<HardEvent::S_MTE2>(S_MTE2_EVENT);
350+ AscendC::DataCopyPad(kScaleUb_, kScaleGm_[keyScaleGmOffset], dataCopymMoutParams, padParams);
351+ startBlockTableIdx++;
352+ curS2Len -= firstPartLen;
353+ resUbBaseOffset = firstPartLen;
354+ }
355+ int32_t getLoopNum = QSiCeilDiv(curS2Len, constInfo_.kCacheBlockSize);
356+ dataCopymMoutParams.blockLen = constInfo_.kCacheBlockSize * sizeof(float);
357+ for (int32_t i = 0; i < getLoopNum; i++) {
358+ if (i == getLoopNum - 1) {
359+ dataCopymMoutParams.blockLen = (curS2Len - i * constInfo_.kCacheBlockSize) * sizeof(float);
360+ }
361+ uint64_t keyScaleGmOffset = 0;
362+ if constexpr (K_LAYOUT_T == QSI_LAYOUT::PA_BNSD || K_LAYOUT_T == QSI_LAYOUT::PA_NZ) { // PA_BNSD(blockNum, n2, blockSize)
363+ keyScaleGmOffset = blkTableGm_.GetValue(blockTableBatchOffset + startBlockTableIdx + i) *
364+ constInfo_.kHeadNum * constInfo_.kCacheBlockSize +
365+ runInfo.n2Idx * constInfo_.kCacheBlockSize;
366+ } else { // PA_BSND(blockNum, blockSize, n2)
367+ keyScaleGmOffset = blkTableGm_.GetValue(blockTableBatchOffset + startBlockTableIdx + i) *
368+ constInfo_.kCacheBlockSize * constInfo_.kHeadNum + runInfo.n2Idx;
369+ }
370+ SetFlag<HardEvent::S_MTE2>(S_MTE2_EVENT);
371+ WaitFlag<HardEvent::S_MTE2>(S_MTE2_EVENT);
372+ AscendC::DataCopyPad(kScaleUb_[resUbBaseOffset + i * constInfo_.kCacheBlockSize], kScaleGm_[keyScaleGmOffset],
373+ dataCopymMoutParams, padParams);
374+ }
375+ } else { // B S N
376+ if (constInfo_.kHeadNum == 1) {
377+ dataCopymMoutParams.blockCount = 1;
378+ dataCopymMoutParams.blockLen = runInfo.actualSingleProcessSInnerSize * sizeof(float);
379+ dataCopymMoutParams.srcStride = 0;
380+ } else {
381+ dataCopymMoutParams.blockCount = runInfo.actualSingleProcessSInnerSize;
382+ dataCopymMoutParams.blockLen = 1 * sizeof(float);
383+ dataCopymMoutParams.srcStride = (constInfo_.kHeadNum - 1) * sizeof(float);
384+ }
385+ dataCopymMoutParams.dstStride = 0;
386+ dataCopymMoutParams.rsv = 0;
387+ AscendC::DataCopyPad(kScaleUb_, kScaleGm_[runInfo.tensorKeyScaleOffset], dataCopymMoutParams, padParams);
388+ }
389+}
390+ 
391+template <typename QSIT>
392+__aicore__ inline void QSIVector<QSIT>::ComputeLse(const QSICommon::RunInfo &runInfo)
393+{
394+ LocalTensor<float> mmResUb = mmInUb_[mte2BufIdx_ % DOUBLE_BUFFER_NUM * MTE2_BUF_SIZE];
395+ int64_t nCount = QSiCeilDiv(runInfo.actualSingleProcessSInnerSize, static_cast<uint32_t>(constInfo_.sparseBlockSize));
396+ int64_t nRepeatTimes = QSiCeilDiv(nCount, B32_VEC_MAX_MASK / B32_VEC_BLK_NUM);
397+ // input n*sparseBlockSize
398+ CopyRepeatParams repeatParams;
399+ repeatParams.dstStride = 1;
400+ repeatParams.dstRepeatSize = B32_VEC_MAX_MASK / B32_VEC_BLK_NUM;
401+ repeatParams.srcStride = QSiCeilDiv(static_cast<int64_t>(constInfo_.sparseBlockSize), B32_VEC_MAX_MASK / B32_VEC_BLK_NUM);
402+ repeatParams.srcRepeatSize = repeatParams.srcStride * (B32_VEC_MAX_MASK / B32_VEC_BLK_NUM);
403+ Copy(nBlkMaxUb_, mmResUb, B32_VEC_MAX_MASK, nRepeatTimes, repeatParams);
404+ PipeBarrier<PIPE_V>();
405+ for (int64_t s2BlockOffset = B32_VEC_BLK_NUM; s2BlockOffset < constInfo_.sparseBlockSize;
406+ s2BlockOffset += B32_VEC_BLK_NUM) {
407+ BinaryRepeatParams binaryRepeatParams;
408+ binaryRepeatParams.dstBlkStride = 1;
409+ binaryRepeatParams.dstRepStride = B32_VEC_MAX_MASK / B32_VEC_BLK_NUM;
410+ binaryRepeatParams.src0BlkStride = 1;
411+ binaryRepeatParams.src0RepStride = B32_VEC_MAX_MASK / B32_VEC_BLK_NUM;
412+ binaryRepeatParams.src1BlkStride = QSiCeilDiv(static_cast<int64_t>(constInfo_.sparseBlockSize), B32_VEC_MAX_MASK / B32_VEC_BLK_NUM);
413+ binaryRepeatParams.src1RepStride = binaryRepeatParams.src1BlkStride * (B32_VEC_MAX_MASK / B32_VEC_BLK_NUM);
414+ AscendC::Max(nBlkMaxUb_, nBlkMaxUb_, mmResUb[s2BlockOffset], B32_VEC_MAX_MASK, nRepeatTimes,
415+ binaryRepeatParams);
416+ }
417+ PipeBarrier<PIPE_V>();
418+ BlockReduceMax(nMaxUb_, nBlkMaxUb_, nRepeatTimes, B32_VEC_MAX_MASK, 1, 1, B32_VEC_MAX_MASK / B32_VEC_BLK_NUM);
419+ 
420+ PipeBarrier<PIPE_V>();
421+ BrcbRepeatParams brcbRepeatParams;
422+ brcbRepeatParams.dstBlkStride = 1;
423+ brcbRepeatParams.dstRepStride = B32_VEC_MAX_MASK / B32_VEC_BLK_NUM;
424+ Brcb(nMaxBrcbUb_, nMaxUb_, nRepeatTimes, brcbRepeatParams); // n,1 -> n,8
425+ 
426+ PipeBarrier<PIPE_V>();
427+ for (int64_t s2BlockOffset = 0; s2BlockOffset < constInfo_.sparseBlockSize; s2BlockOffset += B32_VEC_BLK_NUM) {
428+ BinaryRepeatParams binaryRepeatParams;
429+ binaryRepeatParams.dstBlkStride = QSiCeilDiv(static_cast<int64_t>(constInfo_.sparseBlockSize), B32_VEC_MAX_MASK / B32_VEC_BLK_NUM);
430+ binaryRepeatParams.dstRepStride = binaryRepeatParams.dstBlkStride * (B32_VEC_MAX_MASK / B32_VEC_BLK_NUM);
431+ binaryRepeatParams.src0BlkStride = QSiCeilDiv(static_cast<int64_t>(constInfo_.sparseBlockSize), B32_VEC_MAX_MASK / B32_VEC_BLK_NUM);
432+ binaryRepeatParams.src0RepStride = binaryRepeatParams.dstBlkStride * (B32_VEC_MAX_MASK / B32_VEC_BLK_NUM);
433+ binaryRepeatParams.src1BlkStride = 1;
434+ binaryRepeatParams.src1RepStride = B32_VEC_MAX_MASK / B32_VEC_BLK_NUM;
435+ Sub(mmResUb[s2BlockOffset], mmResUb[s2BlockOffset], nMaxBrcbUb_, B32_VEC_MAX_MASK, nRepeatTimes,
436+ binaryRepeatParams); // n,sparseBlockSize - n,8
437+ }
438+ 
439+ PipeBarrier<PIPE_V>();
440+ Exp(mmResUb, mmResUb, runInfo.actualSingleProcessSInnerSize); // n, sparseBlockSize
441+ 
442+ PipeBarrier<PIPE_V>();
443+ // n,sparseBlockSize -> n, blk
444+ for (int64_t s2BlockOffset = B32_VEC_BLK_NUM; s2BlockOffset < constInfo_.sparseBlockSize;
445+ s2BlockOffset += B32_VEC_BLK_NUM) {
446+ BinaryRepeatParams binaryRepeatParams;
447+ binaryRepeatParams.dstBlkStride = 1;
448+ binaryRepeatParams.dstRepStride = B32_VEC_MAX_MASK / B32_VEC_BLK_NUM;
449+ binaryRepeatParams.src0BlkStride = QSiCeilDiv(static_cast<int64_t>(constInfo_.sparseBlockSize), B32_VEC_MAX_MASK / B32_VEC_BLK_NUM);
450+ binaryRepeatParams.src0RepStride = binaryRepeatParams.src0BlkStride * (B32_VEC_MAX_MASK / B32_VEC_BLK_NUM);
451+ binaryRepeatParams.src1BlkStride = binaryRepeatParams.src0BlkStride;
452+ binaryRepeatParams.src1RepStride = binaryRepeatParams.src0RepStride;
453+ Add(mmResUb, mmResUb[s2BlockOffset], mmResUb, B32_VEC_MAX_MASK, nRepeatTimes, binaryRepeatParams);
454+ }
455+ PipeBarrier<PIPE_V>();
456+ // n, blk -> n,1
457+ BlockReduceSum(mmResUb, mmResUb, nRepeatTimes, B32_VEC_MAX_MASK, 1, 1, B32_VEC_MAX_MASK / B32_VEC_BLK_NUM);
458+ PipeBarrier<PIPE_V>();
459+ Log(mmResUb, mmResUb, nCount); // n,1
460+ Duplicate(nLseTobeSortUb_.template ReinterpretCast<int32_t>(), NEG_INF, constInfo_.s2BaseSize / constInfo_.sparseBlockSize);
461+ PipeBarrier<PIPE_V>();
462+ Add(nLseTobeSortUb_, mmResUb, nMaxUb_, nCount); // n,1
463+}
464+ 
465+/**
466+ src: logits和索引，前logitsNum为logits，后logitsNum为索引
467+ dst: src最终存放的目的地址
468+ logitsNum: 排序的元素个数, 暂只支持[128,256,384,512,1024,1536,2048]
469+ */
470+template <typename QSIT>
471+__aicore__ inline void QSIVector<QSIT>::SortAll(const LocalTensor<float> &dst, const LocalTensor<float> &src, int64_t logitsNum, int64_t mrgElements)
472+{
473+ int64_t sort32Repeats = logitsNum / mrgElements;
474+ AscendC::PipeBarrier<PIPE_V>();
475+ 
476+ int64_t mrgGroups = sort32Repeats;
477+ int64_t i = 0;
478+ AscendC::LocalTensor<float> srcTensor;
479+ AscendC::LocalTensor<float> dstTensor;
480+ while (true) {
481+ if (i % 2 == 0) {
482+ srcTensor = src;
483+ dstTensor = dst;
484+ } else {
485+ srcTensor = dst;
486+ dstTensor = src;
487+ }
488+ AscendC::MrgSort4Info params;
489+ params.elementLengths[0] = mrgElements;
490+ params.elementLengths[MRG_QUE_1] = mrgElements;
491+ params.elementLengths[MRG_QUE_2] = mrgElements;
492+ params.elementLengths[MRG_QUE_3] = mrgElements;
493+ params.ifExhaustedSuspension = false;
494+ params.validBit = 0b1111;
495+ 
496+ AscendC::MrgSortSrcList<float> srcList;
497+ srcList.src1 = srcTensor[0];
498+ srcList.src2 = srcTensor[MRG_QUE_1 * VALUE_AND_INDEX_NUM * mrgElements];
499+ srcList.src3 = srcTensor[MRG_QUE_2 * VALUE_AND_INDEX_NUM * mrgElements];
500+ srcList.src4 = srcTensor[MRG_QUE_3 * VALUE_AND_INDEX_NUM * mrgElements];
501+ if (mrgGroups <= MRG_BLOCK_4) {
502+ params.repeatTimes = 1;
503+ if (mrgGroups == 1) {
504+ break;
505+ } else if (mrgGroups == MRG_BLOCK_2) {
506+ params.validBit = 0b0011;
507+ } else if (mrgGroups == MRG_BLOCK_3) {
508+ params.validBit = 0b0111;
509+ } else if (mrgGroups == MRG_BLOCK_4) {
510+ params.validBit = 0b1111;
511+ }
512+ AscendC::MrgSort<float>(dstTensor, srcList, params);
513+ i += 1;
514+ break;
515+ } else {
516+ params.repeatTimes = mrgGroups / MRG_BLOCK_4;
517+ AscendC::MrgSort<float>(dstTensor, srcList, params);
518+ i += 1;
519+ mrgElements = mrgElements * MRG_BLOCK_4;
520+ mrgGroups = mrgGroups / MRG_BLOCK_4;
521+ }
522+ AscendC::PipeBarrier<PIPE_V>();
523+ }
524+ if (i % 2 == 0) {
525+ AscendC::DataCopy(dst, src, logitsNum * VALUE_AND_INDEX_NUM);
526+ AscendC::PipeBarrier<PIPE_V>();
527+ }
528+}
529+ 
530+template <typename QSIT>
531+__aicore__ inline void QSIVector<QSIT>::SortBasicBlockTopKToSub(const QSICommon::RunInfo &runInfo)
532+{
533+ int64_t nCount = constInfo_.s2BaseSize / constInfo_.sparseBlockSize;
534+ int64_t nSortRepeatTimes = QSiCeilDiv(nCount, SORT_MASK);
535+ int64_t s2Offset = runInfo.s2Idx * constInfo_.s2BaseSize;
536+ PipeBarrier<PIPE_V>();
537+ // 生成索引, 按照最大规格，从256之后开始生成索引
538+ Adds(nLseTobeSortUb_.template ReinterpretCast<int32_t>()[nCount], globalTopkIndice_,
539+ static_cast<int32_t>(QSiCeilDiv(s2Offset, static_cast<int64_t>(constInfo_.sparseBlockSize))), nCount);
540+ 
541+ PipeBarrier<PIPE_V>();
542+ LocalTensor<float> subSortTensor = sortTensor_[((sortedIdx_ + 1) % DOUBLE_BUFFER_NUM) * SORT_BUF_SIZE + nSortCountOffset_ * VALUE_AND_INDEX_NUM];
543+ LocalTensor<float> subSortCacheTensor = sortTensor_[(sortedIdx_ % DOUBLE_BUFFER_NUM) * SORT_BUF_SIZE + nSortCountOffset_ * VALUE_AND_INDEX_NUM];
544+ // 8 * 32 or 4 * 32
545+ Sort32(subSortCacheTensor[nSubSortCountOffset_ * VALUE_AND_INDEX_NUM], nLseTobeSortUb_, nLseTobeSortUb_[nCount].ReinterpretCast<uint32_t>(), nSortRepeatTimes);
546+ 
547+ // src和dst不共地址
548+ SortAll(subSortTensor[nSubSortCountOffset_ * VALUE_AND_INDEX_NUM],
549+ subSortCacheTensor[nSubSortCountOffset_ * VALUE_AND_INDEX_NUM], nCount, SORT_MASK);
550+ nSubSortCountOffset_+= nCount;
551+}
552+ 
553+template <typename QSIT>
554+__aicore__ inline void QSIVector<QSIT>::SortSubTopK(const QSICommon::RunInfo &runInfo)
555+{
556+ if (nSubSortCountOffset_ >= runInfo.targetTopKAlign || runInfo.isLastS2InnerLoop) {
557+ LocalTensor<float> subSortTensor = sortTensor_[((sortedIdx_ + 1) % DOUBLE_BUFFER_NUM) * SORT_BUF_SIZE + nSortCountOffset_ * VALUE_AND_INDEX_NUM];
558+ LocalTensor<float> sortTensor = sortTensor_[(sortedIdx_ % DOUBLE_BUFFER_NUM) * SORT_BUF_SIZE + nSortCountOffset_ * VALUE_AND_INDEX_NUM];
559+ // 产生4的倍数个基本块 256->1024->2048/ 128->512->1024->1536->2048
560+ // src和dst非共地址
561+ SortAll(sortTensor,
562+ subSortTensor, runInfo.targetTopKAlign, constInfo_.s2BaseSize / constInfo_.sparseBlockSize);
563+ nSortCountOffset_ += nSubSortCountOffset_;
564+ nSubSortCountOffset_ = 0;
565+ }
566+ 
567+ if (nSortCountOffset_ >= 4 * runInfo.targetTopKAlign || runInfo.isLastS2InnerLoop) {
568+ LocalTensor<float> nextSortTensor = sortTensor_[((sortedIdx_ + 1) % DOUBLE_BUFFER_NUM) * SORT_BUF_SIZE];
569+ LocalTensor<float> sortTensor = sortTensor_[(sortedIdx_ % DOUBLE_BUFFER_NUM) * SORT_BUF_SIZE];
570+ if (nSortCountOffset_ <= runInfo.targetTopKAlign) {
571+ CopyRepeatParams repeatParams;
572+ repeatParams.dstStride = 1;
573+ repeatParams.dstRepeatSize = B32_VEC_MAX_MASK / B32_VEC_BLK_NUM;
574+ repeatParams.srcStride = 1;
575+ repeatParams.srcRepeatSize = repeatParams.srcStride * (B32_VEC_MAX_MASK / B32_VEC_BLK_NUM);
576+ Copy(nSortOutUb_, sortTensor,
577+ B32_VEC_MAX_MASK, (runInfo.targetTopKAlign * 2) / B32_VEC_MAX_MASK, repeatParams);
578+ sortedIdx_++;
579+ nSortCountOffset_ = runInfo.isLastS2InnerLoop ? 0 : runInfo.targetTopKAlign;
580+ PipeBarrier<PIPE_V>();
581+ Duplicate(sortTensor.template ReinterpretCast<int32_t>(), NEG_INF, SORT_BUF_SIZE);
582+ return;
583+ }
584+ // 只够缓存4条
585+ AscendC::MrgSort4Info params;
586+ params.elementLengths[0] = runInfo.targetTopKAlign;
587+ params.elementLengths[MRG_QUE_1] = runInfo.targetTopKAlign;
588+ params.elementLengths[MRG_QUE_2] = runInfo.targetTopKAlign;
589+ params.elementLengths[MRG_QUE_3] = runInfo.targetTopKAlign;
590+ params.ifExhaustedSuspension = true;
591+ 
592+ AscendC::MrgSortSrcList<float> srcList;
593+ srcList.src1 = sortTensor;
594+ srcList.src2 = sortTensor[MRG_QUE_1 * VALUE_AND_INDEX_NUM * runInfo.targetTopKAlign];
595+ srcList.src3 = sortTensor[MRG_QUE_2 * VALUE_AND_INDEX_NUM * runInfo.targetTopKAlign];
596+ srcList.src4 = sortTensor[MRG_QUE_3 * VALUE_AND_INDEX_NUM * runInfo.targetTopKAlign];
597+ params.repeatTimes = 1;
598+ params.validBit = (1 << QSiCeilDiv(nSortCountOffset_, static_cast<int64_t>(runInfo.targetTopKAlign))) - 1;
599+ AscendC::MrgSort<float>(nextSortTensor, srcList, params);
600+ if (runInfo.isLastS2InnerLoop) {
601+ PipeBarrier<PIPE_V>();
602+ DataCopy(nSortOutUb_, nextSortTensor, runInfo.targetTopKAlign * 2);
603+ PipeBarrier<PIPE_V>();
604+ Duplicate(nextSortTensor.template ReinterpretCast<int32_t>(), NEG_INF, SORT_BUF_SIZE);
605+ }
606+ sortedIdx_ += 1;
607+ nSortCountOffset_ = runInfo.isLastS2InnerLoop ? 0 : runInfo.targetTopKAlign;
608+ PipeBarrier<PIPE_V>();
609+ Duplicate(sortTensor.template ReinterpretCast<int32_t>(), NEG_INF, SORT_BUF_SIZE);
610+ PipeBarrier<PIPE_V>();
611+ }
612+}
613+ 
614+template <typename QSIT>
615+__aicore__ inline void QSIVector<QSIT>::ExtractIndex(const LocalTensor<int32_t> &idxULocal, const LocalTensor<int32_t> &sortLocal,
616+ int64_t extractNum)
617+{
618+ AscendC::GatherMaskParams gatherMaskParams;
619+ gatherMaskParams.repeatTimes = QSiCeilDiv(extractNum * VALUE_AND_INDEX_NUM, B32_VEC_MAX_MASK);
620+ gatherMaskParams.src0BlockStride = 1;
621+ gatherMaskParams.src0RepeatStride = B32_VEC_MAX_MASK / B32_VEC_BLK_NUM;
622+ gatherMaskParams.src1RepeatStride = 0;
623+ uint64_t rsvdCnt = 0; // 用于保存筛选后保留下来的元素个数
624+ uint8_t src1Pattern = 2; // 固定模式2,表示筛选出奇数索引的数
625+ AscendC::GatherMask(idxULocal, sortLocal, src1Pattern, false, static_cast<uint32_t>(0), gatherMaskParams, rsvdCnt);
626+ AscendC::PipeBarrier<PIPE_V>();
627+}
628+ 
629+template <typename QSIT>
630+__aicore__ inline void QSIVector<QSIT>::CopyOutResult(const QSICommon::RunInfo &runInfo)
631+{
632+ if (GetSubBlockIdx() == 0) {
633+ return;
634+ }
635+ // 1.写尾块
636+ PipeBarrier<PIPE_V>();
637+ ExtractIndex(nIdxUb_, nSortOutUb_.template ReinterpretCast<int32_t>(), runInfo.targetTopKAlign);
638+ 
639+ AscendC::DataCopyParams dataCopyOutyParams;
640+ dataCopyOutyParams.blockCount = 1;
641+ dataCopyOutyParams.blockLen = (runInfo.targetTopK + runInfo.fixedTailCount) * sizeof(int32_t);
642+ dataCopyOutyParams.srcStride = 0;
643+ dataCopyOutyParams.dstStride = 0;
644+ SetFlag<HardEvent::V_S>(V_S_EVENT);
645+ WaitFlag<HardEvent::V_S>(V_S_EVENT);
646+ for (int i = 0; i < runInfo.fixedTailCount; i++) {
647+ nIdxUb_.SetValue((runInfo.targetTopK + i),
648+ runInfo.needProcessS2Size / constInfo_.sparseBlockSize + i);
649+ }
650+ SetFlag<HardEvent::S_MTE3>(S_MTE3_EVENT);
651+ WaitFlag<HardEvent::S_MTE3>(S_MTE3_EVENT);
652+ AscendC::DataCopyPad(indiceOutGm_[runInfo.indiceOutOffset],nIdxUb_, dataCopyOutyParams);
653+}
654+ 
655+template <typename QSIT>
656+__aicore__ inline void QSIVector<QSIT>::CopyOutFdResult(const QSICommon::RunInfo &runInfo)
657+{
658+ // vec1Res Gm = [aic, constInfo_.s1BaseSize, 2, 2, topkOut_] float32
659+ // vec1Param Gm = [aic, constInfo_.s1BaseSize, 2, 16] int64
660+ // 16 = [needFd, s2AcSeq, s2Start, s2End, isS2End, bn2idx, TopkAlign, TopK, ......]
661+ int64_t wsOffset = blockId_ * constInfo_.s1BaseSize * 2 * MAX_TOPK_WITH_ID;
662+ int64_t wsInfoOffset = blockId_ * constInfo_.s1BaseSize * 2 * paramNum_;
663+ 
664+ LocalTensor<int64_t> tmpiBuff = paramBuf_.Get<int64_t>();
665+ SetFlag<HardEvent::MTE3_S>(MTE3_S_EVENT);
666+ WaitFlag<HardEvent::MTE3_S>(MTE3_S_EVENT);
667+ tmpiBuff.SetValue(0, static_cast<int64_t>(1));
668+ tmpiBuff.SetValue(1, static_cast<int64_t>(runInfo.needProcessS2Size));
669+ tmpiBuff.SetValue(2, static_cast<int64_t>(blockS2StartIdx_));
670+ tmpiBuff.SetValue(3, static_cast<int64_t>(runInfo.s2Idx * constInfo_.s2BaseSize + runInfo.actualSingleProcessSInnerSize));
671+ tmpiBuff.SetValue(4, static_cast<int64_t>((runInfo.s2Idx + 1) * constInfo_.s2BaseSize >= runInfo.needProcessS2Size));
672+ tmpiBuff.SetValue(5, static_cast<int64_t>(runInfo.bN2Idx));
673+ tmpiBuff.SetValue(6, static_cast<int64_t>(runInfo.targetTopKAlign));
674+ tmpiBuff.SetValue(7, static_cast<int64_t>(runInfo.targetTopK));
675+ tmpiBuff.SetValue(8, static_cast<int64_t>(runInfo.indiceOutOffset));
676+ // 写入头尾判断
677+ // [head, tail]
678+ // head: 与前面规约，与前后规约
679+ // tail: 与后面规约
680+ // WS偏移规则 blockS2StartIdx_ != 0
681+ // 跟前面块做规约 写到0偏移 不用做计算 blockS2StartIdx_ == 0 and !isS2End
682+ // 跟后面块做规约 写到1偏移 需要 + constInfo_.s1BaseSize, MAX_TOPK*2
683+ if (blockS2StartIdx_ == 0) { // S2不是最后结束的数据就需要往后做规约，放入第二块ws
684+ wsInfoOffset += paramNum_;
685+ wsOffset += MAX_TOPK_WITH_ID;
686+ }
687+ SetFlag<HardEvent::S_MTE3>(S_MTE3_EVENT);
688+ WaitFlag<HardEvent::S_MTE3>(S_MTE3_EVENT);
689+ AscendC::DataCopyParams dataCopyOutyParams;
690+ dataCopyOutyParams.blockCount = 1;
691+ dataCopyOutyParams.blockLen = 16 * sizeof(int64_t);
692+ dataCopyOutyParams.srcStride = 0;
693+ dataCopyOutyParams.dstStride = 0;
694+ AscendC::DataCopyPad(vec1ParamGm_[wsInfoOffset], tmpiBuff, dataCopyOutyParams);
695+ SetFlag<HardEvent::V_MTE3>(V_MTE3_EVENT);
696+ WaitFlag<HardEvent::V_MTE3>(V_MTE3_EVENT);
697+ 
698+ dataCopyOutyParams.blockLen = MAX_TOPK_WITH_ID * sizeof(float);
699+ AscendC::DataCopyPad(vec1ResGm[wsOffset], nSortOutUb_, dataCopyOutyParams);
700+ SetFlag<HardEvent::MTE3_V>(MTE3_V_EVENT_FD);
701+ WaitFlag<HardEvent::MTE3_V>(MTE3_V_EVENT_FD);
702+}
703+ 
704+template <typename QSIT>
705+__aicore__ inline void QSIVector<QSIT>::CleanInvalidOutput(int64_t invalidS1offset)
706+{
707+ // init -1 and copy to output
708+ WaitFlag<HardEvent::MTE3_V>(MTE3_V_EVENT);
709+ Duplicate(nIdxUb_, constInfo_.INVALID_IDX, constInfo_.sparseCount);
710+ SetFlag<HardEvent::V_MTE3>(V_MTE3_EVENT);
711+ WaitFlag<HardEvent::V_MTE3>(V_MTE3_EVENT);
712+ CopyOut(indiceOutGm_[invalidS1offset], nIdxUb_, constInfo_.sparseCount);
713+ SetFlag<HardEvent::MTE3_V>(MTE3_V_EVENT);
714+}
715+ 
716+template <typename QSIT>
717+__aicore__ inline void QSIVector<QSIT>::ProcessVec(const QSICommon::RunInfo &runInfo)
718+{
719+ if (GetSubBlockIdx() == 0) {
720+ return;
721+ }
722+ 
723+ if (runInfo.isFirstS2InnerLoop) {
724+ blockS2StartIdx_ = runInfo.s2Idx;
725+ }
726+ 
727+ WaitFlag<HardEvent::V_MTE2>(V_MTE2_EVENT + mte2BufIdx_ % DOUBLE_BUFFER_NUM);
728+ CopyMm1ResultIn(runInfo);
729+ GetKeyScale(runInfo);
730+ 
731+ AscendC::SetFlag<HardEvent::MTE2_V>(MTE2_V_EVENT);
732+ AscendC::WaitFlag<HardEvent::MTE2_V>(MTE2_V_EVENT);
733+ // 将qScale从1扩展到2048个数
734+ AscendC::Duplicate(qScaleUb_.template ReinterpretCast<float>(), runInfo.qScale, runInfo.actualSingleProcessSInnerSize);
735+ // 计算反量化参数Weight = qScale * kScale
736+ PipeBarrier<PIPE_V>();
737+ AscendC::Mul(qkScaleUb_, qScaleUb_, kScaleUb_, runInfo.actualSingleProcessSInnerSize);
738+ SetFlag<HardEvent::V_MTE2>(V_MTE2_EVENT);
739+ WaitFlag<HardEvent::V_MTE2>(V_MTE2_EVENT);
740+ 
741+ // bmm1结果从int32转换为fp32
742+ PipeBarrier<PIPE_V>();
743+ LocalTensor<int32_t> mmInUbInt = mmInUb_.template ReinterpretCast<int32_t>()[mte2BufIdx_ % DOUBLE_BUFFER_NUM * MTE2_BUF_SIZE];
744+ LocalTensor<float> mmInUb = mmInUb_[mte2BufIdx_ % DOUBLE_BUFFER_NUM * MTE2_BUF_SIZE];
745+ AscendC::Cast(mmInUb, mmInUbInt, RoundMode::CAST_NONE, MTE2_BUF_SIZE);
746+ SetFlag<HardEvent::V_S>(V_S_EVENT);
747+ WaitFlag<HardEvent::V_S>(V_S_EVENT);
748+ AscendC::Mul(mmInUb, mmInUb, qkScaleUb_, runInfo.actualSingleProcessSInnerSize);
749+ PipeBarrier<PIPE_V>();
750+ 
751+ ComputeLse(runInfo);
752+ SetFlag<HardEvent::V_MTE2>(V_MTE2_EVENT + mte2BufIdx_ % DOUBLE_BUFFER_NUM);
753+ mte2BufIdx_++;
754+ 
755+ if(runInfo.isLastS2InnerLoop) {
756+ WaitFlag<HardEvent::MTE3_V>(MTE3_V_EVENT);
757+ }
758+ SortBasicBlockTopKToSub(runInfo); // 基本块排满一个sub topk空间
759+ SortSubTopK(runInfo); // sub空间排满后，做一次整体排序，产生topk
760+ if (runInfo.isLastS2InnerLoop) {
761+ if (blockS2StartIdx_ == 0 && (runInfo.s2Idx + 1) * constInfo_.s2BaseSize >= runInfo.needProcessS2Size) {
762+ CopyOutResult(runInfo); // 最后一个循环，排序结果拷出
763+ } else {
764+ // 触发FD写出
765+ CopyOutFdResult(runInfo);
766+ }
767+ SetFlag<HardEvent::MTE3_V>(MTE3_V_EVENT);
768+ }
769+}
770+ 
771+template <typename QSIT>
772+__aicore__ inline void QSIVector<QSIT>::ProcessLD()
773+{
774+ if (GetSubBlockIdx() == 0) {
775+ return;
776+ }
777+ LocalTensor<float> curValueIdxUb = ldToBeMrgBuf_.Get<float>();
778+ LocalTensor<float> tmpUb = ldTmpBuf_.Get<float>();
779+ 
780+ // S2开头信息
781+ // 开始必然没有头规约，因此从尾规约开始处理，while循环读取下一个核的头规约
782+ // 存满4个list或者遇到S2结尾，则做merge，直到做完S2
783+ // 每个核都忽略自己的头规约，因为必然由前面的核做完
784+ // vec1Res Gm = [aiv, constInfo_.s1BaseSize, 2, 2, topkOut_] float32
785+ // vec1Param Gm = [aiv, constInfo_.s1BaseSize, 2, 16] int64
786+ int64_t needFd = vec1ParamGm_.GetValue(blockId_ * constInfo_.s1BaseSize * 2 * paramNum_ + paramNum_);
787+ 
788+ if (needFd == 0 || vec1ParamGm_.GetValue(blockId_ * constInfo_.s1BaseSize * 2 * paramNum_ + paramNum_ + 2) != 0) {
789+ return;
790+ }
791+ 
792+ // 搬入数据
793+ int64_t wsOffsetInit = blockId_ * constInfo_.s1BaseSize * 2 * MAX_TOPK_WITH_ID + MAX_TOPK_WITH_ID;
794+ SetFlag<HardEvent::V_MTE2>(V_MTE2_EVENT_FD);
795+ WaitFlag<HardEvent::V_MTE2>(V_MTE2_EVENT_FD);
796+ SetFlag<HardEvent::S_MTE2>(S_MTE2_EVENT);
797+ WaitFlag<HardEvent::S_MTE2>(S_MTE2_EVENT);
798+ DataCopyPad(curValueIdxUb, vec1ResGm[wsOffsetInit],
799+ {1, static_cast<uint16_t>(MAX_TOPK_WITH_ID * sizeof(int32_t)), 0, 0}, {true, 0, 0, 0});
800+ int64_t valueOffset = MAX_TOPK_WITH_ID;
801+ int64_t acc_list_num = vec1ParamGm_.GetValue(blockId_ * constInfo_.s1BaseSize * 2 * paramNum_ + paramNum_ + 6);
802+ // 获取下一个核规约信息
803+ int32_t tmpCubeId = blockId_ + 2;
804+ int64_t wsInfoOffsetInit = tmpCubeId * constInfo_.s1BaseSize * 2 * paramNum_;
805+ needFd = vec1ParamGm_.GetValue(wsInfoOffsetInit);
806+ int64_t isS2End = vec1ParamGm_.GetValue(wsInfoOffsetInit + 4);
807+ int64_t bN2Idx = vec1ParamGm_.GetValue(wsInfoOffsetInit + 5);
808+ int64_t targetTopKAlign = vec1ParamGm_.GetValue(wsInfoOffsetInit + 6);
809+ 
810+ while (needFd == 1) {
811+ int64_t wsOffset = tmpCubeId * constInfo_.s1BaseSize * 2 * MAX_TOPK_WITH_ID;
812+ // 搬入头规约数据
813+ SetFlag<HardEvent::V_MTE2>(V_MTE2_EVENT_FD);
814+ WaitFlag<HardEvent::V_MTE2>(V_MTE2_EVENT_FD);
815+ SetFlag<HardEvent::S_MTE2>(S_MTE2_EVENT);
816+ WaitFlag<HardEvent::S_MTE2>(S_MTE2_EVENT);
817+ DataCopyPad(curValueIdxUb[valueOffset], vec1ResGm[wsOffset],
818+ {1, static_cast<uint16_t>(MAX_TOPK_WITH_ID * sizeof(int32_t)), 0, 0}, {true, 0, 0, 0});
819+ valueOffset += MAX_TOPK_WITH_ID;
820+ acc_list_num+=targetTopKAlign;
821+ // 每满4个list，聚合 前2K为mrg结果
822+ if (acc_list_num >= 4*targetTopKAlign || isS2End == 1) {
823+ // MrgSort 四条2048的队列，Mrg成一条
824+ SetFlag<HardEvent::MTE2_V>(MTE2_V_EVENT);
825+ WaitFlag<HardEvent::MTE2_V>(MTE2_V_EVENT);
826+ AscendC::MrgSort4Info params;
827+ params.elementLengths[0] = targetTopKAlign;
828+ params.elementLengths[1] = targetTopKAlign;
829+ params.elementLengths[2] = targetTopKAlign;
830+ params.elementLengths[3] = targetTopKAlign;
831+ params.ifExhaustedSuspension = true;
832+ params.validBit = (1 << QSiCeilDiv(acc_list_num, static_cast<int64_t>(targetTopKAlign))) - 1;
833+ params.repeatTimes = 1;
834+ 
835+ AscendC::MrgSortSrcList<float> srcList;
836+ srcList.src1 = curValueIdxUb[0];
837+ srcList.src2 = curValueIdxUb[1 * MAX_TOPK_WITH_ID];
838+ srcList.src3 = curValueIdxUb[2 * MAX_TOPK_WITH_ID];
839+ srcList.src4 = curValueIdxUb[3 * MAX_TOPK_WITH_ID];
840+ MrgSort<float>(tmpUb, srcList, params);
841+ PipeBarrier<PIPE_V>();
842+ DataCopy(curValueIdxUb, tmpUb, MAX_TOPK_WITH_ID);
843+ PipeBarrier<PIPE_V>();
844+ acc_list_num = targetTopKAlign;
845+ valueOffset = MAX_TOPK_WITH_ID;
846+ }
847+ // reduce到S2末尾，则跳出
848+ if (isS2End == 1) {
849+ break;
850+ }
851+ tmpCubeId+=2;
852+ int64_t wsInfoOffset = tmpCubeId * constInfo_.s1BaseSize * 2 * paramNum_;
853+ needFd = vec1ParamGm_.GetValue(wsInfoOffset);
854+ isS2End = vec1ParamGm_.GetValue(wsInfoOffset + 4);
855+ }
856+ // 搬出
857+ // 1.写尾块
858+ int64_t needProcessS2Size = vec1ParamGm_.GetValue(wsInfoOffsetInit + 1);
859+ int64_t targetTopK = vec1ParamGm_.GetValue(wsInfoOffsetInit + 7);
860+ int64_t outOffset = vec1ParamGm_.GetValue(wsInfoOffsetInit + 8);
861+ LocalTensor<int32_t> outIdxUb = ldOutIdxBuf_.Get<int32_t>();
862+ ExtractIndex(outIdxUb, curValueIdxUb.template ReinterpretCast<int32_t>(), targetTopKAlign);
863+ 
864+ AscendC::DataCopyParams dataCopyOutyParams;
865+ dataCopyOutyParams.blockCount = 1;
866+ dataCopyOutyParams.blockLen = (targetTopK + constInfo_.fixedTailCount) * sizeof(int32_t);
867+ dataCopyOutyParams.srcStride = 0;
868+ dataCopyOutyParams.dstStride = 0;
869+ SetFlag<HardEvent::V_S>(V_S_EVENT);
870+ WaitFlag<HardEvent::V_S>(V_S_EVENT);
871+ for (int i = 0; i < constInfo_.fixedTailCount; i++) {
872+ outIdxUb.SetValue((targetTopK + i),
873+ needProcessS2Size / constInfo_.sparseBlockSize + i);
874+ }
875+ SetFlag<HardEvent::S_MTE3>(S_MTE3_EVENT);
876+ WaitFlag<HardEvent::S_MTE3>(S_MTE3_EVENT);
877+ AscendC::DataCopyPad(indiceOutGm_[outOffset], outIdxUb, dataCopyOutyParams);
878+}
879+} // namespace QSIKernel
880+#endif
Aexperimental/attention/quant_sals_indexer/op_kernel/quant_sals_indexer_template_tiling_key.h+85-0
@@ -0,0 +1,85 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file quant_sals_indexer_template_tiling_key.h
13+ * \brief
14+ */
15+ 
16+#ifndef TEMPLATE_TILING_KEY_QSI_H_
17+#define TEMPLATE_TILING_KEY_QSI_H_
18+ 
19+#include "ascendc/host_api/tiling/template_argument.h"
20+ 
21+#define QSI_TPL_FP16 1
22+#define QSI_TPL_INT8 2
23+#define QSI_TPL_INT32 3
24+#define QSI_TPL_INT4 29
25+#define QSI_TPL_BF16 27
26+ 
27+#define QSI_LAYOUT_BSND 0
28+#define QSI_LAYOUT_PA_BNSD 2
29+#define QSI_LAYOUT_PA_BSND 3
30+#define QSI_LAYOUT_PA_NZ 4
31+ 
32+#define ASCENDC_TPL_4_BW 4
33+ 
34+// 模板参数支持的范围定义
35+ASCENDC_TPL_ARGS_DECL(QuantSalsIndexer, // 算子OpType
36+ ASCENDC_TPL_DTYPE_DECL(DT_Q, QSI_TPL_INT8, QSI_TPL_INT4),
37+ ASCENDC_TPL_DTYPE_DECL(DT_K, QSI_TPL_INT8, QSI_TPL_INT4),
38+ ASCENDC_TPL_DTYPE_DECL(DT_OUT, QSI_TPL_INT32), ASCENDC_TPL_BOOL_DECL(PAGE_ATTENTION, 0, 1),
39+ ASCENDC_TPL_UINT_DECL(K_LAYOUT_T, ASCENDC_TPL_4_BW, ASCENDC_TPL_UI_LIST,
40+ QSI_LAYOUT_PA_BSND, QSI_LAYOUT_PA_BNSD, QSI_LAYOUT_BSND, QSI_LAYOUT_PA_NZ,), );
41+ 
42+// 支持的模板参数组合
43+// 用于调用GET_TPL_TILING_KEY获取TilingKey时，接口内部校验TilingKey是否合法
44+ASCENDC_TPL_SEL(
45+ ASCENDC_TPL_ARGS_SEL(ASCENDC_TPL_DTYPE_SEL(DT_Q, QSI_TPL_INT8), ASCENDC_TPL_DTYPE_SEL(DT_K, QSI_TPL_INT8),
46+ ASCENDC_TPL_DTYPE_SEL(DT_OUT, QSI_TPL_INT32),
47+ ASCENDC_TPL_BOOL_SEL(PAGE_ATTENTION, 1),
48+ ASCENDC_TPL_UINT_SEL(K_LAYOUT_T, ASCENDC_TPL_UI_LIST, QSI_LAYOUT_PA_BSND), ),
49+ 
50+ ASCENDC_TPL_ARGS_SEL(ASCENDC_TPL_DTYPE_SEL(DT_Q, QSI_TPL_INT8), ASCENDC_TPL_DTYPE_SEL(DT_K, QSI_TPL_INT8),
51+ ASCENDC_TPL_DTYPE_SEL(DT_OUT, QSI_TPL_INT32),
52+ ASCENDC_TPL_BOOL_SEL(PAGE_ATTENTION, 1),
53+ ASCENDC_TPL_UINT_SEL(K_LAYOUT_T, ASCENDC_TPL_UI_LIST, QSI_LAYOUT_PA_BNSD), ),
54+
55+ ASCENDC_TPL_ARGS_SEL(ASCENDC_TPL_DTYPE_SEL(DT_Q, QSI_TPL_INT8), ASCENDC_TPL_DTYPE_SEL(DT_K, QSI_TPL_INT8),
56+ ASCENDC_TPL_DTYPE_SEL(DT_OUT, QSI_TPL_INT32),
57+ ASCENDC_TPL_BOOL_SEL(PAGE_ATTENTION, 1),
58+ ASCENDC_TPL_UINT_SEL(K_LAYOUT_T, ASCENDC_TPL_UI_LIST, QSI_LAYOUT_PA_NZ), ),
59+ 
60+ ASCENDC_TPL_ARGS_SEL(ASCENDC_TPL_DTYPE_SEL(DT_Q, QSI_TPL_INT8), ASCENDC_TPL_DTYPE_SEL(DT_K, QSI_TPL_INT8),
61+ ASCENDC_TPL_DTYPE_SEL(DT_OUT, QSI_TPL_INT32),
62+ ASCENDC_TPL_BOOL_SEL(PAGE_ATTENTION, 0),
63+ ASCENDC_TPL_UINT_SEL(K_LAYOUT_T, ASCENDC_TPL_UI_LIST, QSI_LAYOUT_BSND), ),
64+ 
65+ ASCENDC_TPL_ARGS_SEL(ASCENDC_TPL_DTYPE_SEL(DT_Q, QSI_TPL_INT4), ASCENDC_TPL_DTYPE_SEL(DT_K, QSI_TPL_INT4),
66+ ASCENDC_TPL_DTYPE_SEL(DT_OUT, QSI_TPL_INT32),
67+ ASCENDC_TPL_BOOL_SEL(PAGE_ATTENTION, 1),
68+ ASCENDC_TPL_UINT_SEL(K_LAYOUT_T, ASCENDC_TPL_UI_LIST, QSI_LAYOUT_PA_BSND), ),
69+ 
70+ ASCENDC_TPL_ARGS_SEL(ASCENDC_TPL_DTYPE_SEL(DT_Q, QSI_TPL_INT4), ASCENDC_TPL_DTYPE_SEL(DT_K, QSI_TPL_INT4),
71+ ASCENDC_TPL_DTYPE_SEL(DT_OUT, QSI_TPL_INT32),
72+ ASCENDC_TPL_BOOL_SEL(PAGE_ATTENTION, 1),
73+ ASCENDC_TPL_UINT_SEL(K_LAYOUT_T, ASCENDC_TPL_UI_LIST, QSI_LAYOUT_PA_BNSD), ),
74+
75+ ASCENDC_TPL_ARGS_SEL(ASCENDC_TPL_DTYPE_SEL(DT_Q, QSI_TPL_INT4), ASCENDC_TPL_DTYPE_SEL(DT_K, QSI_TPL_INT4),
76+ ASCENDC_TPL_DTYPE_SEL(DT_OUT, QSI_TPL_INT32),
77+ ASCENDC_TPL_BOOL_SEL(PAGE_ATTENTION, 1),
78+ ASCENDC_TPL_UINT_SEL(K_LAYOUT_T, ASCENDC_TPL_UI_LIST, QSI_LAYOUT_PA_NZ), ),
79+ 
80+ ASCENDC_TPL_ARGS_SEL(ASCENDC_TPL_DTYPE_SEL(DT_Q, QSI_TPL_INT4), ASCENDC_TPL_DTYPE_SEL(DT_K, QSI_TPL_INT4),
81+ ASCENDC_TPL_DTYPE_SEL(DT_OUT, QSI_TPL_INT32),
82+ ASCENDC_TPL_BOOL_SEL(PAGE_ATTENTION, 0),
83+ ASCENDC_TPL_UINT_SEL(K_LAYOUT_T, ASCENDC_TPL_UI_LIST, QSI_LAYOUT_BSND), ), );
84+ 
85+#endif
Aexperimental/attention/quant_sals_indexer/op_kernel/quant_sals_indexer_vector.h+269-0
@@ -0,0 +1,269 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file quant_sals_indexer_vector.h
13+ * \brief
14+ */
15+#ifndef QUANT_SALS_INDEXER_VECTOR_H
16+#define QUANT_SALS_INDEXER_VECTOR_H
17+ 
18+#include "kernel_operator.h"
19+ 
20+namespace QSIServiceVec {
21+using namespace AscendC;
22+ 
23+ 
24+constexpr int32_t INVALID_INDEX = -1;
25+constexpr uint8_t VEC_REPEAT_MAX = 255;
26+ 
27+constexpr uint8_t B32_BLOCK_ALIGN_NUM = 8;
28+constexpr uint8_t B32_VEC_REPEAT_STRIDE = 8;
29+constexpr uint64_t VEC_REPEAT_BYTES = 256;
30+constexpr int32_t CONST_TWO = 2;
31+ 
32+ 
33+ 
34+ 
35+ 
36+__aicore__ inline void CopyIn(LocalTensor<float> &mmOutUb, GlobalTensor<float> &mMoutGm,
37+ int64_t MMout_gmoffset,
38+ int64_t groupInner, int64_t s2Inner, int64_t mmUbStride)
39+{
40+ // 将MMout_gmoffset copy到UB上
41+ AscendC::DataCopyPadExtParams<float> padParams{false, 0, 0, 0};
42+ AscendC::DataCopyExtParams dataCopymMoutParams;
43+ dataCopymMoutParams.blockCount = groupInner;
44+ dataCopymMoutParams.blockLen = s2Inner * sizeof(float);
45+ dataCopymMoutParams.srcStride = 0;
46+ dataCopymMoutParams.dstStride = mmUbStride;
47+ dataCopymMoutParams.rsv = 0;
48+ AscendC::DataCopyPad(mmOutUb, mMoutGm[MMout_gmoffset], dataCopymMoutParams, padParams);
49+}
50+ 
51+ 
52+template <typename T>
53+__aicore__ inline void CopyOut(const GlobalTensor<T> &dstGm, const LocalTensor<T> &srcUb, int64_t copyCount)
54+{
55+ AscendC::DataCopyParams dataCopyOutyParams;
56+ dataCopyOutyParams.blockCount = 1;
57+ dataCopyOutyParams.blockLen = copyCount * sizeof(T);
58+ dataCopyOutyParams.srcStride = 0;
59+ dataCopyOutyParams.dstStride = 0;
60+ AscendC::DataCopyPad(dstGm, srcUb, dataCopyOutyParams);
61+}
62+ 
63+__aicore__ inline void DoScale(const LocalTensor<float> &reduceCacheBuf, LocalTensor<float> &mmOutUb,
64+ LocalTensor<float> &tmpBuff,
65+ int64_t groupInner, int64_t s2Inner, int32_t outerGidx)
66+{
67+ // do scale: [groupInner, 8] * [groupInner, s2Inner]
68+ uint64_t countPerRepeat = VEC_REPEAT_BYTES / sizeof(float);
69+ uint64_t repeatTimes = s2Inner / countPerRepeat;
70+ for (int32_t i = 0; i < groupInner; i++) {
71+ if (outerGidx == 0) {
72+ AscendC::Mul(reduceCacheBuf[i * s2Inner], mmOutUb[i * s2Inner], tmpBuff[i * B32_BLOCK_ALIGN_NUM],
73+ countPerRepeat, repeatTimes, {1, 1, 0, B32_VEC_REPEAT_STRIDE, B32_VEC_REPEAT_STRIDE, 0});
74+ } else {
75+ AscendC::Mul(mmOutUb[i * s2Inner], mmOutUb[i * s2Inner], tmpBuff[i * B32_BLOCK_ALIGN_NUM], countPerRepeat,
76+ repeatTimes, {1, 1, 0, B32_VEC_REPEAT_STRIDE, B32_VEC_REPEAT_STRIDE, 0});
77+ }
78+ }
79+ 
80+ if (outerGidx != 0) {
81+ AscendC::PipeBarrier<PIPE_V>();
82+ AscendC::Add(reduceCacheBuf, mmOutUb, reduceCacheBuf, groupInner * s2Inner);
83+ }
84+ AscendC::PipeBarrier<PIPE_V>();
85+}
86+ 
87+ 
88+__aicore__ inline uint64_t FindNearestPower2(uint64_t value)
89+{
90+ if (value <= CONST_TWO) {
91+ return value;
92+ } else {
93+ const uint64_t pow = 63 - clz(value); // clz返回前导0的个数，对于64位整数，最大有效位位置 = 63 - 前导0个数
94+ return (1 << pow);
95+ }
96+}
97+ 
98+ 
99+// dstTensor 需要初始化0
100+__aicore__ inline void DoReduce(const LocalTensor<float> &srcTensor, LocalTensor<float> &dstTensor, int32_t rNum,
101+ int32_t aNum)
102+{
103+ if (rNum == 1) {
104+ AscendC::Adds<float>(dstTensor, srcTensor, 0, aNum);
105+ AscendC::PipeBarrier<PIPE_V>();
106+ return;
107+ }
108+ 
109+ uint32_t dichotomizeAddPow = FindNearestPower2(rNum);
110+ uint32_t dichotomizeAddDiffSize = rNum - dichotomizeAddPow;
111+ if (dichotomizeAddDiffSize != 0) {
112+ AscendC::Add(srcTensor, srcTensor, srcTensor[dichotomizeAddPow * aNum], dichotomizeAddDiffSize * aNum);
113+ AscendC::PipeBarrier<PIPE_V>();
114+ }
115+ int32_t nowRows = dichotomizeAddPow;
116+ while (nowRows > CONST_TWO) {
117+ nowRows = nowRows / CONST_TWO;
118+ AscendC::Add(srcTensor, srcTensor, srcTensor[nowRows * aNum], nowRows * aNum);
119+ AscendC::PipeBarrier<PIPE_V>();
120+ }
121+ AscendC::Add(dstTensor, srcTensor, srcTensor[aNum], aNum);
122+ AscendC::PipeBarrier<PIPE_V>();
123+}
124+ 
125+ 
126+/**
127+ src: 传入的初始化空间
128+ eleNum: 需要初始化的元素个数需为64整数倍，元素将被初始化为交错排布的-inf，-1
129+ */
130+__aicore__ inline void InitSortOutBuf(const LocalTensor<float> &src, int64_t eleNum)
131+{
132+ uint64_t mask1[2] = {0x5555555555555555, 0};
133+ uint64_t mask0[2] = {0xaaaaaaaaaaaaaaaa, 0};
134+ int64_t repeatNum = eleNum / B32_VEC_ELM_NUM;
135+ int64_t forLoop = repeatNum / VEC_REPEAT_MAX;
136+ int64_t forRemain = repeatNum % VEC_REPEAT_MAX;
137+ for (int i = 0; i < forLoop; i++) {
138+ AscendC::Duplicate(src.template ReinterpretCast<int32_t>(), NEG_INF, mask1, VEC_REPEAT_MAX, 1,
139+ B32_VEC_REPEAT_STRIDE);
140+ AscendC::Duplicate(src.template ReinterpretCast<int32_t>(), INVALID_INDEX, mask0, VEC_REPEAT_MAX, 1,
141+ B32_VEC_REPEAT_STRIDE);
142+ }
143+ if (forRemain > 0) {
144+ AscendC::Duplicate(src.template ReinterpretCast<int32_t>()[forLoop * VEC_REPEAT_MAX * B32_VEC_ELM_NUM], NEG_INF,
145+ mask1, forRemain, 1, B32_VEC_REPEAT_STRIDE);
146+ AscendC::Duplicate(src.template ReinterpretCast<int32_t>()[forLoop * VEC_REPEAT_MAX * B32_VEC_ELM_NUM],
147+ INVALID_INDEX, mask0, forRemain, 1, B32_VEC_REPEAT_STRIDE);
148+ }
149+ AscendC::PipeBarrier<PIPE_V>();
150+}
151+ 
152+ 
153+/**
154+ dst: 输出全排序的结果，排布方式为value，index
155+ srcValue：输入的待排序浮点数
156+ srcIndex：浮点数的索引
157+ tmp: 计算使用到的临时空间，大小为srcValue+srcIndex
158+ logitsNum: 排序的元素个数
159+ */
160+__aicore__ inline void SortAll(LocalTensor<float> &dst, LocalTensor<float> &srcValue, LocalTensor<uint32_t> &srcIndex,
161+ LocalTensor<float> &tmpTensor, int64_t logitsNum)
162+{
163+ int64_t sort32Repeats = logitsNum / BLOCK_BYTES;
164+ AscendC::Sort<float, true>(dst, srcValue, srcIndex, tmpTensor, sort32Repeats);
165+ AscendC::PipeBarrier<PIPE_V>();
166+}
167+ 
168+ 
169+/**
170+ mrgDst: 合并进的Tensor
171+ mrgSrc: 待合并的Tensor
172+ tmpTensor：空间为mrgDst+mrgSrc
173+ */
174+__aicore__ inline void MergeSort(const LocalTensor<float> &mrgDst, int32_t mrgDstNum, LocalTensor<float> &mrgSrc,
175+ int32_t mrgSrcNum, LocalTensor<float> &tmpTensor)
176+{
177+ AscendC::MrgSort4Info params;
178+ params.elementLengths[0] = mrgDstNum;
179+ params.elementLengths[1] = mrgSrcNum;
180+ params.ifExhaustedSuspension = false;
181+ params.validBit = 0b0011;
182+ params.repeatTimes = 1;
183+ 
184+ AscendC::MrgSortSrcList<float> srcList;
185+ srcList.src1 = mrgDst;
186+ srcList.src2 = mrgSrc;
187+ 
188+ AscendC::MrgSort<float>(tmpTensor, srcList, params);
189+ AscendC::PipeBarrier<PIPE_V>();
190+ AscendC::DataCopy(mrgDst, tmpTensor, mrgDstNum * VALUE_AND_INDEX_NUM);
191+ AscendC::PipeBarrier<PIPE_V>();
192+}
193+ 
194+ 
195+/**
196+ * @brief 合并基础块函数
197+ * @param dst 归并后的输出, 大小为blockNum * basicBlockSize * 2 * sizeof(float)
198+ * @param src 基本块输入
199+ * @param blockNum 基本块的数量
200+ * @param basicBlockSize 基础块的大小
201+ * @return 无
202+ */
203+__aicore__ inline void MrgBasicBlock(const LocalTensor<float> &dst, const LocalTensor<float> &src, int64_t blockNum,
204+ int64_t basicBlockSize)
205+{
206+ // 初始化合并排序参数
207+ AscendC::MrgSort4Info params;
208+ params.elementLengths[MRG_QUE_0] = basicBlockSize;
209+ params.elementLengths[MRG_QUE_1] = basicBlockSize;
210+ params.elementLengths[MRG_QUE_2] = basicBlockSize;
211+ params.elementLengths[MRG_QUE_3] = basicBlockSize;
212+ params.ifExhaustedSuspension = false;
213+ // 根据块的数量设置有效位
214+ if (blockNum == MRG_BLOCK_2) {
215+ params.validBit = 0b0011;
216+ } else if (blockNum == MRG_BLOCK_3) {
217+ params.validBit = 0b0111;
218+ } else if (blockNum == MRG_BLOCK_4) {
219+ params.validBit = 0b1111;
220+ } else {
221+ AscendC::DataCopy(dst, src, basicBlockSize * VALUE_AND_INDEX_NUM);
222+ return;
223+ }
224+ // 初始化源列表
225+ AscendC::MrgSortSrcList<float> srcList;
226+ srcList.src1 = src[0];
227+ srcList.src2 = src[basicBlockSize * VALUE_AND_INDEX_NUM * MRG_QUE_1];
228+ srcList.src3 = src[basicBlockSize * VALUE_AND_INDEX_NUM * MRG_QUE_2];
229+ srcList.src4 = src[basicBlockSize * VALUE_AND_INDEX_NUM * MRG_QUE_3];
230+ // 执行合并排序
231+ AscendC::MrgSort<float>(dst, srcList, params);
232+}
233+ 
234+ 
235+/**
236+ * @brief 从两个队列中选择topk
237+ * @param dst 已经归并好的topk数据
238+ * @param needsMerging 需要合并的有序数据
239+ * @param tmp 临时空间
240+ * @param topk topk的元素个数
241+ * @param mergSize 待合并的元素个数
242+ * @return 无
243+ */
244+template <bool needMrg = true>
245+__aicore__ inline void SparseTopK(const LocalTensor<float> &dst, const LocalTensor<float> &needsMerging,
246+ const LocalTensor<float> &tmp, int64_t topk, int64_t mergSize)
247+{
248+ // 如果不需要合并，则直接复制数据
249+ if (!needMrg) {
250+ AscendC::DataCopy(dst, needsMerging, mergSize * VALUE_AND_INDEX_NUM);
251+ return;
252+ }
253+ // 初始化合并排序参数
254+ AscendC::MrgSort4Info params;
255+ params.elementLengths[0] = topk;
256+ params.elementLengths[1] = mergSize;
257+ params.ifExhaustedSuspension = (topk == mergSize);
258+ params.validBit = 0b0011;
259+ // 初始化源列表
260+ AscendC::MrgSortSrcList<float> srcList;
261+ srcList.src1 = dst;
262+ srcList.src2 = needsMerging;
263+ // 执行合并排序
264+ AscendC::MrgSort<float>(tmp, srcList, params);
265+ // 将结果复制到目标张量
266+ AscendC::DataCopy(dst, tmp, topk * VALUE_AND_INDEX_NUM);
267+}
268+} // namespace QSIServiceVec
269+#endif // QUANT_SALS_INDEXER_VECTOR_H
Aexperimental/attention/quant_sals_indexer_metadata/CMakeLists.txt+0-0
The file is empty
Aexperimental/attention/quant_sals_indexer_metadata/examples/test_aclnn_quant_sals_indexer_metadata.cpp+168-0
@@ -0,0 +1,168 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+#include <stdio.h>
12+#include <iostream>
13+#include <tuple>
14+#include <vector>
15+#include "acl/acl.h"
16+//#include "aclnnop/aclnn_quant_sals_indexer_metadata.h"
17+#include "../../quant_sals_indexer/op_kernel/quant_sals_indexer_metadata.h"
18+#include "../op_host/aclnn_quant_sals_indexer_metadata.h"
19+ 
20+static const uint32_t batchSize = 3;
21+static const uint32_t kvSeqSize = 10240;
22+static const uint32_t kvHeadNum = 4;
23+static const uint32_t headDim = 128;
24+static const uint32_t sparseBlockSize = 16;
25+static const double sparseRatio = 0.25f;
26+static const uint32_t fixedTailCount = 16;
27+static std::string layoutKey = "BSND";
28+ 
29+static const std::vector<int32_t> actSeqLenKV = {10240, 10240, 10240};
30+static const std::vector<int64_t> actSeqLenKVShape = {batchSize};
31+static const std::vector<int64_t> actSeqLenKVStride = {1};
32+static const std::vector<int64_t> metadataShape = {optiling::QSI_META_SIZE};
33+static const std::vector<int64_t> metadataStride = {1};
34+ 
35+static const bool enableActLenKV = true;
36+ 
37+std::tuple<aclTensor*, void*> CreateTensor(size_t size, // in bytes
38+ std::vector<int64_t> shape,
39+ std::vector<int64_t> stride,
40+ aclDataType dType,
41+ const void* hostData = nullptr) {
42+ void* devicePtr = nullptr;
43+ auto ret = aclrtMalloc(&devicePtr, size, ACL_MEM_MALLOC_HUGE_FIRST);
44+ if (ret != ACL_SUCCESS) {
45+ printf("aclrtMalloc %d\n", ret);
46+ return {nullptr, nullptr};
47+ }
48+ 
49+ aclTensor* tensor = aclCreateTensor(&shape[0], shape.size(), dType,
50+ &stride[0], 0, aclFormat::ACL_FORMAT_ND,
51+ &shape[0], shape.size(), devicePtr);
52+ if (tensor == nullptr) {
53+ aclrtFree(devicePtr);
54+ return {nullptr, nullptr};
55+ }
56+ 
57+ if (hostData != nullptr) {
58+ aclrtMemcpy(devicePtr, size, hostData, size, ACL_MEMCPY_HOST_TO_DEVICE);
59+ }
60+ return {tensor, devicePtr};
61+}
62+ 
63+static void DumpMeta(void* data) {
64+ optiling::detail::QsiMetaData* metaDataPtr =
65+ (optiling::detail::QsiMetaData*)data;
66+ printf("usedCoreNum: %d \n", metaDataPtr->usedCoreNum);
67+ for (uint32_t i = 0; i < metaDataPtr->usedCoreNum; i++) {
68+ printf("bN2End[%d]: %d \n", i, metaDataPtr->bN2End[i]);
69+ printf("gS1End[%d]: %d \n", i, metaDataPtr->gS1End[i]);
70+ printf("s2End[%d]: %d \n", i, metaDataPtr->s2End[i]);
71+ }
72+}
73+ 
74+int main() {
75+ int32_t deviceId = 0;
76+ aclrtStream stream;
77+ aclError ret = 0;
78+ 
79+ aclTensor* kvSeqLenTensor = nullptr;
80+ void* kvSeqLenDevPtr = nullptr;
81+ 
82+ aclTensor* metadataTensor = nullptr;
83+ void* metadataDevPtr = nullptr;
84+ aclOpExecutor* executor = nullptr;
85+ uint64_t workspaceSize = 0;
86+ void* workspace = nullptr;
87+ 
88+ ret = aclInit(nullptr);
89+ if (ret != ACL_SUCCESS) {
90+ printf("aclInit %d\n", ret);
91+ return -1;
92+ }
93+ 
94+ ret = aclrtSetDevice(deviceId);
95+ if (ret != ACL_SUCCESS) {
96+ printf("aclrtSetDevice %d\n", ret);
97+ return -1;
98+ }
99+ 
100+ ret = aclrtCreateStream(&stream);
101+ if (ret != ACL_SUCCESS) {
102+ printf("aclrtCreateStream %d\n", ret);
103+ return -1;
104+ }
105+ 
106+ if (enableActLenKV) {
107+ std::tie(kvSeqLenTensor, kvSeqLenDevPtr) = CreateTensor(
108+ actSeqLenKV.size() * sizeof(actSeqLenKV[0]), actSeqLenKVShape,
109+ actSeqLenKVStride, aclDataType::ACL_INT32, &actSeqLenKV[0]);
110+ if (kvSeqLenTensor == nullptr) {
111+ return -1;
112+ }
113+ }
114+ 
115+ std::tie(metadataTensor, metadataDevPtr) =
116+ CreateTensor(sizeof(int32_t) * optiling::QSI_META_SIZE, metadataShape,
117+ metadataStride, aclDataType::ACL_INT32);
118+ if (metadataTensor == nullptr) {
119+ return -1;
120+ }
121+ 
122+ ret = aclnnQuantSalsIndexerMetadataGetWorkspaceSize(
123+ kvSeqLenTensor, batchSize, kvSeqSize, kvHeadNum, headDim,
124+ sparseBlockSize, sparseRatio, fixedTailCount, &layoutKey[0], metadataTensor, &workspaceSize,
125+ &executor);
126+ if (ret != ACL_SUCCESS) {
127+ printf("aclnnQuantSalsIndexerMetadataGetWorkspaceSize %d\n", ret);
128+ return -1;
129+ }
130+ 
131+ ret =
132+ aclnnQuantSalsIndexerMetadata(workspace, workspaceSize, executor, stream);
133+ if (ret != ACL_SUCCESS) {
134+ printf("aclnnQuantSalsIndexerMetadata %d\n", ret);
135+ return -1;
136+ }
137+ 
138+ ret = aclrtSynchronizeStream(stream);
139+ if (ret != ACL_SUCCESS) {
140+ printf("aclrtSynchronizeStream %d\n", ret);
141+ return -1;
142+ }
143+ 
144+ std::vector<int32_t> metdataHost(optiling::QSI_META_SIZE);
145+ ret = aclrtMemcpy(metdataHost.data(),
146+ metdataHost.size() * sizeof(metdataHost[0]), metadataDevPtr,
147+ optiling::QSI_META_SIZE * sizeof(int32_t),
148+ ACL_MEMCPY_DEVICE_TO_HOST);
149+ if (ret != ACL_SUCCESS) {
150+ printf("aclrtMemcpy %d\n", ret);
151+ return -1;
152+ }
153+ 
154+ DumpMeta(&metdataHost[0]);
155+ 
156+ aclDestroyTensor(kvSeqLenTensor);
157+ aclDestroyTensor(metadataTensor);
158+ 
159+ aclrtFree(kvSeqLenDevPtr);
160+ aclrtFree(metadataDevPtr);
161+ aclrtFree(workspace);
162+ 
163+ aclrtDestroyStream(stream);
164+ aclrtResetDevice(deviceId);
165+ aclFinalize();
166+ 
167+ return 0;
168+}
Aexperimental/attention/quant_sals_indexer_metadata/examples/test_geir_quant_sals_indexer_metadata.cpp+141-0
@@ -0,0 +1,141 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+#include <cstdint>
12+#include <vector>
13+#include <string>
14+#include <map>
15+ 
16+#include "graph.h"
17+#include "types.h"
18+#include "tensor.h"
19+#include "ge_error_codes.h"
20+#include "ge_api_types.h"
21+#include "ge_api.h"
22+#include "array_ops.h" // REG_OP(Data)
23+#include "ge_ir_build.h"
24+ 
25+// #include "experiment_ops.h"
26+#include "nn_other.h"
27+#include "../op_graph/quant_sals_indexer_metadata_proto.h"
28+#include "../../quant_sals_indexer/op_kernel/quant_sals_indexer_metadata.h"
29+ 
30+using namespace ge;
31+ 
32+static const uint32_t batchSize = 3;
33+static const uint32_t kvSeqSize = 10240;
34+static const uint32_t kvHeadNum = 4;
35+static const uint32_t sparseBlockSize = 16;
36+static const uint32_t fixedTailCount = 16;
37+static const uint32_t aicCoreNum = 24;
38+static const uint32_t aivCoreNum = 48;
39+ 
40+static const std::vector<int32_t> actSeqLenKV = {10240, 10240, 10240};
41+static const std::vector<int64_t> actSeqLenKVShape = {batchSize};
42+static const std::vector<int64_t> metadataShape = {optiling::QSI_META_SIZE};
43+static const std::string dumpFile = "./dump";
44+ 
45+static const bool enableActLenKV = true;
46+ 
47+using namespace ge;
48+ 
49+class GeEnv {
50+public:
51+ GeEnv() {
52+ std::map<AscendString, AscendString> opt = {{"ge.exec.deviceId", "0"},
53+ {"ge.graphRunMode", "1"}};
54+ inited_ = GEInitialize(opt) == SUCCESS;
55+ }
56+ ~GeEnv() {
57+ if (inited_)
58+ GEFinalize();
59+ }
60+ 
61+ bool Ok() { return inited_; }
62+ 
63+private:
64+ bool inited_;
65+};
66+ 
67+int main(int argc, char **argv) {
68+ GeEnv geEnv;
69+ if (!geEnv.Ok()) {
70+ return -1;
71+ }
72+ 
73+ Graph graph("GraphQuantSalsIndexerMetadata");
74+ 
75+ auto metaDataOp = op::QuantSalsIndexerMetadata(
76+ "QuantSalsIndexerMetadata-0");
77+ 
78+ // gen graph
79+ auto dataOp0 = op::Data("input0").set_attr_index(0); // Data 算子
80+ if (enableActLenKV) {
81+ TensorDesc desc(ge::Shape(actSeqLenKVShape), FORMAT_ND, DT_INT32);
82+ desc.SetPlacement(ge::kPlacementHost);
83+ desc.SetFormat(FORMAT_ND);
84+ desc.SetRealDimCnt(actSeqLenKVShape.size());
85+ dataOp0.update_input_desc_x(desc);
86+ graph.AddOp(dataOp0);
87+ metaDataOp.set_input_actual_seq_lengths_key(dataOp0);
88+ }
89+ 
90+ metaDataOp.update_output_desc_metadata(TensorDesc{ge::Shape(metadataShape), FORMAT_ND, DT_INT32});
91+ metaDataOp.set_attr_batch_size(batchSize);
92+ metaDataOp.set_attr_key_seq_size(kvSeqSize);
93+ metaDataOp.set_attr_key_head_num(kvHeadNum);
94+ metaDataOp.set_attr_sparse_block_size(sparseBlockSize);
95+ metaDataOp.set_attr_fixed_tail_count(fixedTailCount);
96+ metaDataOp.set_attr_aic_core_num(aicCoreNum);
97+ metaDataOp.set_attr_aiv_core_num(aivCoreNum);
98+ graph.AddOp(metaDataOp);
99+ 
100+ // run Graph
101+ std::vector<ge::Operator> inputOps = {dataOp0};
102+ std::vector<ge::Operator> outputOps = {metaDataOp};
103+ graph.SetInputs(inputOps).SetOutputs(outputOps);
104+ 
105+ aclgrphDumpGraph(graph, dumpFile.c_str(), dumpFile.length());
106+ 
107+ std::vector<ge::Tensor> inputTensors;
108+ std::vector<ge::Tensor> outputTensors;
109+ 
110+ if (enableActLenKV) {
111+ inputTensors.push_back(
112+ Tensor{dataOp0.get_input_desc_x(),
113+ reinterpret_cast<const uint8_t *>(&actSeqLenKV[0]),
114+ actSeqLenKV.size() * sizeof(actSeqLenKV[0])});
115+ }
116+ 
117+ 
118+ std::map<AscendString, AscendString> build_options;
119+ auto session = std::make_shared<Session>(build_options);
120+ std::map<AscendString, AscendString> graph_options;
121+ uint32_t graph_id = 0;
122+ session->AddGraph(graph_id, graph, graph_options);
123+ if (session->RunGraph(graph_id, inputTensors, outputTensors)) {
124+ printf("RunGraph Fail\n");
125+ return -1;
126+ }
127+ 
128+ auto tensor = outputTensors[0];
129+ auto data = tensor.GetData();
130+ auto dataSize = tensor.GetTensorDesc().GetShape().GetShapeSize();
131+
132+ optiling::detail::QsiMetaData *metaDataPtr = (optiling::detail::QsiMetaData*)data;
133+ printf("usedCoreNum: %d \n", metaDataPtr->usedCoreNum);
134+ for (uint32_t i = 0; i < metaDataPtr->usedCoreNum; i ++) {
135+ printf("bN2End[%d]: %d \n", i, metaDataPtr->bN2End[i]);
136+ printf("gS1End[%d]: %d \n", i, metaDataPtr->gS1End[i]);
137+ printf("s2End[%d]: %d \n", i, metaDataPtr->s2End[i]);
138+ }
139+ 
140+ return 0;
141+}
Aexperimental/attention/quant_sals_indexer_metadata/op_graph/quant_sals_indexer_metadata_proto.h+60-0
@@ -0,0 +1,60 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file quant_sals_indexer_metadata_proto.h
13+ * \brief
14+ */
15+#ifndef QUANT_SALS_INDEXER_METADATA_PROTO_H
16+#define QUANT_SALS_INDEXER_METADATA_PROTO_H
17+ 
18+#include "graph/operator_reg.h"
19+#include "graph/types.h"
20+ 
21+namespace ge {
22+ 
23+/**
24+* @brief Function QuantSalsIndexerMetadata.
25+ 
26+* @par Inputs:
27+* @li actual_seq_lengths_kv: A matrix tensor. The type support int32.
28+* Effective sequence length of key/value in different batches.
29+ 
30+* @par Attributes:
31+* @li batch_size: An int. batch size of key/value tensor.
32+* @li kv_seq_size: An int. sequence len of key/value Tensor.
33+* @li kv_head_num: An int. head-dim of key/value Tensor.
34+* @li fixed_tail_count: An int. fixed tail count.
35+* @li sparse_block_size: An int. The block size in the sparse phase. Default: 1.
36+* @li aic_core_num: An int. cube core num of device
37+* @li aiv_core_num: An int. vector core num of device
38+ 
39+* @par Outputs:
40+* @li metadata: A matrix tensor. The type support int32.
41+* The output of attention structure.
42+*/
43+REG_OP(QuantSalsIndexerMetadata)
44+ .OPTIONAL_INPUT(actual_seq_lengths_key, TensorType({DT_INT32}))
45+ .OUTPUT(metadata, TensorType({DT_INT32}))
46+ .REQUIRED_ATTR(batch_size, Int)
47+ .REQUIRED_ATTR(key_seq_size, Int)
48+ .REQUIRED_ATTR(key_head_num, Int)
49+ .REQUIRED_ATTR(head_dim, Int)
50+ .REQUIRED_ATTR(aic_core_num, Int)
51+ .REQUIRED_ATTR(aiv_core_num, Int)
52+ .REQUIRED_ATTR(sparse_block_size, Int)
53+ .REQUIRED_ATTR(sparse_ratio, Float)
54+ .REQUIRED_ATTR(fixed_tail_count, Int)
55+ .ATTR(layout_key, String, "BSND")
56+ .ATTR(soc_version, String, "ascend910B")
57+ .OP_END_FACTORY_REG(QuantSalsIndexerMetadata)
58+} // namespace ge
59+ 
60+#endif // QUANT_SALS_INDEXER_METADATA_PROTO_H
Aexperimental/attention/quant_sals_indexer_metadata/op_host/aclnn_quant_sals_indexer_metadata.cpp+113-0
@@ -0,0 +1,113 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file aclnn_quant_sals_indexer_metadata.cpp
13+ * \brief
14+ */
15+ 
16+#include "aclnn_quant_sals_indexer_metadata.h"
17+#include "aclnn/aclnn_base.h"
18+#include "aclnn_kernels/common/op_error_check.h"
19+#include "aclnn_kernels/contiguous.h"
20+#include "aclnn_kernels/reshape.h"
21+#include "l0_quant_sals_indexer_metadata.h"
22+#include "opdev/common_types.h"
23+#include "opdev/data_type_utils.h"
24+#include "opdev/format_utils.h"
25+#include "opdev/make_op_executor.h"
26+#include "opdev/op_dfx.h"
27+#include "opdev/op_executor.h"
28+#include "opdev/op_log.h"
29+#include "opdev/tensor_view_utils.h"
30+#include "opdev/platform.h"
31+ 
32+#ifdef __cplusplus
33+extern "C" {
34+#endif
35+ 
36+static aclnnStatus ParamsCheck(const aclTensor* actualSeqLengthsKvOptional,
37+ int64_t batchSize,
38+ int64_t keySeqSize,
39+ int64_t keyHeadNum,
40+ int64_t headDim,
41+ int64_t sparseBlockSize,
42+ double sparseRatio,
43+ int64_t fixedTailCount,
44+ char* layoutKeyOptional,
45+ const aclTensor* metaData) {
46+ if (headDim != 128) {
47+ return ACLNN_ERR_PARAM_INVALID;
48+ }
49+ 
50+ if (sparseBlockSize != 16) {
51+ return ACLNN_ERR_PARAM_INVALID;
52+ }
53+ if (batchSize < 0 || keySeqSize < 0 || keyHeadNum < 0 || sparseBlockSize < 0 || fixedTailCount < 0) {
54+ return ACLNN_ERR_PARAM_INVALID;
55+ }
56+ return ACLNN_SUCCESS;
M
Mmiaofangzheng4月2日

ParamsCheck 函数体为空（直接 return ACLNN_SUCCESS），没有做任何参数校验。这意味着用户传入任意非法参数（空指针、非法 shape、不支持的 dtype）都会透传到算子内部，可能导致不可预知的错误或 crash。即使是实验性算子，也应该至少做基本的空指针和 dtype 校验。

likedislike
M
mayafei
4月3日 评论：
57+}
58+ 
59+aclnnStatus aclnnQuantSalsIndexerMetadataGetWorkspaceSize(
60+ const aclTensor* actualSeqLengthsKvOptional,
61+ int64_t batchSize,
62+ int64_t keySeqSize,
63+ int64_t keyHeadNum,
64+ int64_t headDim,
65+ int64_t sparseBlockSize,
66+ double sparseRatio,
67+ int64_t fixedTailCount,
68+ char* layoutKeyOptional,
69+ const aclTensor* metaData,
70+ uint64_t* workspaceSize,
71+ aclOpExecutor** executor) {
72+ L2_DFX_PHASE_1(
73+ aclnnQuantSalsIndexerMetadata,
74+ DFX_IN(actualSeqLengthsKvOptional, batchSize, keySeqSize, keyHeadNum, headDim,
75+ sparseBlockSize, sparseRatio, fixedTailCount, layoutKeyOptional),
76+ DFX_OUT(metaData));
77+ 
78+ auto uniqueExecutor = CREATE_EXECUTOR();
79+ CHECK_RET(uniqueExecutor.get() != nullptr, ACLNN_ERR_INNER_CREATE_EXECUTOR);
80+ 
81+ auto ret = ParamsCheck(actualSeqLengthsKvOptional, batchSize, keySeqSize, keyHeadNum, headDim,
82+ sparseBlockSize, sparseRatio, fixedTailCount, layoutKeyOptional, metaData);
83+ CHECK_RET(ret == ACLNN_SUCCESS, ret);
84+ 
85+ const op::PlatformInfo &npuInfo = op::GetCurrentPlatformInfo();
86+ uint32_t aicCoreNum = npuInfo.GetCubeCoreNum();
87+ uint32_t aivCoreNum = npuInfo.GetVectorCoreNum();
88+ std::string socVersionStr = npuInfo.GetSocLongVersion();
89+ const char* socVersionOptional = socVersionStr.c_str();
90+ 
91+ auto output = l0op::QuantSalsIndexerMetadata(
92+ actualSeqLengthsKvOptional, batchSize, keySeqSize, keyHeadNum, headDim, aicCoreNum, aivCoreNum,
93+ sparseBlockSize, sparseRatio, fixedTailCount, layoutKeyOptional, socVersionOptional, metaData,
94+ uniqueExecutor.get());
95+ CHECK_RET(output != nullptr, ACLNN_ERR_INNER_NULLPTR);
96+ 
97+ *workspaceSize = 0;
98+ uniqueExecutor.ReleaseTo(executor);
99+ return ACLNN_SUCCESS;
100+}
101+ 
102+__attribute__((visibility("default"))) aclnnStatus
103+aclnnQuantSalsIndexerMetadata(void* workspace,
104+ uint64_t workspaceSize,
105+ aclOpExecutor* executor,
106+ aclrtStream stream) {
107+ L2_DFX_PHASE_2(aclnnQuantSalsIndexerMetadata);
108+ return CommonOpExecutorRun(workspace, workspaceSize, executor, stream);
109+}
110+ 
111+#ifdef __cplusplus
112+}
113+#endif
Aexperimental/attention/quant_sals_indexer_metadata/op_host/aclnn_quant_sals_indexer_metadata.h+64-0
@@ -0,0 +1,64 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+#ifndef ACLNN_QUANT_SALS_INDEXER_METADATA_AICPU_H
12+#define ACLNN_QUANT_SALS_INDEXER_METADATA_AICPU_H
13+ 
14+#include "aclnn/aclnn_base.h"
15+ 
16+#ifdef __cplusplus
17+extern "C" {
18+#endif
19+ 
20+/* funtion: aclnnQuantSalsIndexerMetadataGetWorkspaceSize
21+ * parameters :
22+ * actualSeqLengthsKvOptional : optional
23+ * batchSize : required
24+ * kvSeqSize : required
25+ * kvHeadNum : required
26+ * fixedTailCount : required
27+ * sparseBlockSize : required
28+ * out : required
29+ * workspaceSize : size of workspace(output).
30+ * executor : executor context(output).
31+ */
32+__attribute__((visibility("default"))) aclnnStatus
33+aclnnQuantSalsIndexerMetadataGetWorkspaceSize(
34+ const aclTensor* actualSeqLengthsKvOptional,
35+ int64_t batchSize,
36+ int64_t keySeqSize,
37+ int64_t keyHeadNum,
38+ int64_t headDim,
39+ int64_t sparseBlockSize,
40+ double sparseRatio,
41+ int64_t fixedTailCount,
42+ char* layoutKeyOptional,
43+ const aclTensor* metaData,
44+ uint64_t* workspaceSize,
45+ aclOpExecutor** executor);
46+ 
47+/* funtion: aclnnQuantSalsIndexerMetadata
48+ * parameters :
49+ * workspace : workspace memory addr(input).
50+ * workspaceSize : size of workspace(input).
51+ * executor : executor context(input).
52+ * stream : acl stream.
53+ */
54+__attribute__((visibility("default"))) aclnnStatus
55+aclnnQuantSalsIndexerMetadata(void* workspace,
56+ uint64_t workspaceSize,
57+ aclOpExecutor* executor,
58+ aclrtStream stream);
59+ 
60+#ifdef __cplusplus
61+}
62+#endif
63+ 
64+#endif // ACLNN_QUANT_SALS_INDEXER_METADATA_AICPU_H
Aexperimental/attention/quant_sals_indexer_metadata/op_host/l0_quant_sals_indexer_metadata.cpp+65-0
@@ -0,0 +1,65 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file l0_quant_sals_indexer_metadata.cpp
13+ * \brief
14+ */
15+ 
16+#include "l0_quant_sals_indexer_metadata.h"
17+#include "opdev/aicpu/aicpu_task.h"
18+#include "opdev/make_op_executor.h"
19+#include "opdev/op_def.h"
20+#include "opdev/op_dfx.h"
21+#include "opdev/op_executor.h"
22+#include "opdev/op_log.h"
23+#include "opdev/shape_utils.h"
24+ 
25+using namespace op;
26+namespace l0op {
27+OP_TYPE_REGISTER(QuantSalsIndexerMetadata);
28+ 
29+const aclTensor* QuantSalsIndexerMetadata(
30+ const aclTensor* actualSeqLengthsKvOptional,
31+ int64_t batchSize,
32+ int64_t keySeqSize,
33+ int64_t keyHeadNum,
34+ int64_t headDim,
35+ int64_t aicCoreNum,
36+ int64_t aivCoreNum,
37+ int64_t sparseBlockSize,
38+ double sparseRatio,
39+ int64_t fixedTailCount,
40+ char* layoutKeyOptional,
41+ const char* socVersionOptional,
42+ const aclTensor* metaData,
43+ aclOpExecutor* executor) {
44+ L0_DFX(QuantSalsIndexerMetadata, actualSeqLengthsKvOptional, batchSize,
45+ keySeqSize, keyHeadNum, headDim, aicCoreNum, aivCoreNum,
46+ sparseBlockSize, sparseRatio, fixedTailCount, layoutKeyOptional, socVersionOptional, metaData);
47+ 
48+ static internal::AicpuTaskSpace space("QuantSalsIndexerMetadata");
49+ 
50+ auto ret = ADD_TO_LAUNCHER_LIST_AICPU(
51+ QuantSalsIndexerMetadata,
52+ OP_ATTR_NAMES({"batch_size", "key_seq_size", "key_head_num", "head_dim", "aic_core_num",
53+ "aiv_core_num", "sparse_block_size", "sparse_ratio", "fixed_tail_count",
54+ "layout_key", "soc_version"}),
55+ OP_INPUT(actualSeqLengthsKvOptional), OP_OUTPUT(metaData),
56+ OP_ATTR(batchSize, keySeqSize, keyHeadNum, headDim, aicCoreNum, aivCoreNum,
57+ sparseBlockSize, sparseRatio, fixedTailCount, layoutKeyOptional, socVersionOptional));
58+ OP_CHECK(ret == ACL_SUCCESS,
59+ OP_LOGE(ACLNN_ERR_INNER_NULLPTR,
60+ "QuantSalsIndexerMetadata"
61+ " ADD_TO_LAUNCHER_LIST_AICPU failed."),
62+ return nullptr);
63+ return metaData;
64+}
65+} // namespace l0op
Aexperimental/attention/quant_sals_indexer_metadata/op_host/l0_quant_sals_indexer_metadata.h+34-0
@@ -0,0 +1,34 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+#ifndef L0_QUANT_SALS_INDEXER_METADATA_AICPU_H
12+#define L0_QUANT_SALS_INDEXER_METADATA_AICPU_H
13+ 
14+#include "opdev/op_executor.h"
15+ 
16+namespace l0op {
17+const aclTensor* QuantSalsIndexerMetadata(
18+ const aclTensor* actualSeqLengthsKvOptional,
19+ int64_t batchSize,
20+ int64_t keySeqSize,
21+ int64_t keyHeadNum,
22+ int64_t headDim,
23+ int64_t aicCoreNum,
24+ int64_t aivCoreNum,
25+ int64_t sparseBlockSize,
26+ double sparseRatio,
27+ int64_t fixedTailCount,
28+ char* layoutKeyOptional,
29+ const char* socVersionOptional,
30+ const aclTensor* metaData,
31+ aclOpExecutor* executor);
32+} // namespace l0op
33+ 
34+#endif
Aexperimental/attention/quant_sals_indexer_metadata/op_host/quant_sals_indexer_metadata_infershape.cpp+41-0
@@ -0,0 +1,41 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file quant_sals_indexer_metadata_infershape.cpp
13+ * \brief
14+ */
15+#include "register/op_impl_registry.h"
16+#include "error/ops_error.h"
17+#include "../../quant_sals_indexer/op_kernel/quant_sals_indexer_metadata.h"
18+#include "../op_graph/quant_sals_indexer_metadata_proto.h"
19+ 
20+using namespace ge;
21+ 
22+namespace ops {
23+static ge::graphStatus InferShapeQuantSalsIndexerMetadata(gert::InferShapeContext* context)
24+{
25+ gert::Shape* oShape = context->GetOutputShape(0);
26+ OPS_LOG_E_IF_NULL(context, oShape, ge::GRAPH_FAILED);
27+ oShape->SetDimNum(1);
28+ oShape->SetDim(0, optiling::QSI_META_SIZE);
29+ return GRAPH_SUCCESS;
30+}
31+ 
32+static ge::graphStatus InferDtypeQuantSalsIndexerMetadata(gert::InferDataTypeContext* context)
33+{
34+ context->SetOutputDataType(0, DT_INT32);
35+ return GRAPH_SUCCESS;
36+}
37+ 
38+IMPL_OP_INFERSHAPE(QuantSalsIndexerMetadata)
39+ .InferShape(InferShapeQuantSalsIndexerMetadata)
40+ .InferDataType(InferDtypeQuantSalsIndexerMetadata);
41+} // namespace ops
Aexperimental/attention/quant_sals_indexer_metadata/op_kernel_aicpu/quant_sals_indexer_metadata_aicpu.cpp+624-0
@@ -0,0 +1,624 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+
11+/*!
12+ * \file quant_sals_indexer_metadata_aicpu.cpp
13+ * \brief
14+ */
15+ 
16+#include "quant_sals_indexer_metadata_aicpu.h"
17+#include <cstdio>
18+#include <math.h>
19+#include "../../quant_sals_indexer/op_kernel/quant_sals_indexer_metadata.h"
20+#include "../../common/cpu_context_util.h"
21+ 
22+ 
23+namespace aicpu {
24+uint32_t QuantSalsIndexerMetaDataCpuKernel::Compute(CpuKernelContext &ctx) {
25+ context_ = &ctx;
26+ bool success = Prepare(ctx) && BalanceSchedule() && GenMetaData();
27+ return success ? KERNEL_STATUS_OK : KERNEL_STATUS_PARAM_INVALID;
28+}
29+ 
30+bool QuantSalsIndexerMetaDataCpuKernel::Prepare(CpuKernelContext &ctx) {
31+ // input
32+ actSeqLenKV_ = ctx.Input(static_cast<uint32_t>(ParamId::actSeqLenKV));
33+ // output
34+ metaData_ = ctx.Output(static_cast<uint32_t>(ParamId::metaData));
35+ 
36+ bool requiredAttrs =
37+ GetAttrValue(ctx, "batch_size", batchSize_) &&
38+ GetAttrValue(ctx, "key_seq_size", kvSeqSize_) &&
39+ GetAttrValue(ctx, "key_head_num", kvHeadNum_) &&
40+ GetAttrValue(ctx, "fixed_tail_count", fixedTailCount_) &&
41+ GetAttrValue(ctx, "sparse_block_size", sparseBlockSize_) &&
42+ GetAttrValue(ctx, "aic_core_num", aicCoreNum_) &&
43+ GetAttrValue(ctx, "aiv_core_num", aivCoreNum_);
44+ if (!requiredAttrs) {
45+ return false;
46+ }
47+ coreNum_ = aicCoreNum_;
48+ return (ParamsCheck() && ParamsInit());
49+}
50+ 
51+bool QuantSalsIndexerMetaDataCpuKernel::ParamsCheck() {
52+ if (actSeqLenKV_ != nullptr) {
53+ auto shape = actSeqLenKV_->GetTensorShape();
54+ auto data = actSeqLenKV_->GetData();
55+ auto dtype = actSeqLenKV_->GetDataType();
56+ 
57+ KERNEL_CHECK_NULLPTR(shape, false,
58+ "shape of actual_seq_lengths_kv is null");
59+ KERNEL_CHECK_NULLPTR(data, false, "data of actual_seq_lengths_kv is null");
60+ KERNEL_CHECK_FALSE((dtype == DataType::DT_INT32), false,
61+ "dtype of actual_seq_lengths_kv is not int32");
62+ 
63+ KERNEL_CHECK_FALSE(
64+ (shape->GetDims() == 1 && shape->GetDimSize(0) == batchSize_), false,
65+ "shape of actual_seq_lengths_query date is not {%u,}", batchSize_);
66+ }
67+ 
68+ KERNEL_CHECK_NULLPTR(metaData_, false, "metadata is null");
69+ auto shape = metaData_->GetTensorShape();
70+ auto data = metaData_->GetData();
71+ auto dtype = metaData_->GetDataType();
72+ 
73+ KERNEL_CHECK_NULLPTR(shape, false, "shape of metadata is null");
74+ KERNEL_CHECK_NULLPTR(data, false, "data of metadata is null");
75+ KERNEL_CHECK_FALSE((dtype == DataType::DT_INT32), false,
76+ "dtype of metadata is not int32");
77+ KERNEL_CHECK_FALSE((shape->GetDims() == 1 &&
78+ shape->GetDimSize(0) == optiling::QSI_META_SIZE),
79+ false, "shape of sparse_seq_lengths_kv date is not {%u,}",
80+ optiling::QSI_META_SIZE);
81+ KERNEL_CHECK_FALSE(
82+ (aicCoreNum_ != 0 && aivCoreNum_ != 0 && aivCoreNum_ % aicCoreNum_ == 0 &&
83+ aicCoreNum_ <= optiling::CORE_NUM &&
M
Mmiaofangzheng4月2日

校验用的 CORE_NUM（24）与 host 侧 tiling.h 中的 NCAI_MAX_AIC_CORE_NUM（28）不一致，可能导致 24 < aicCoreNum <= 28 的场景在 metadata 侧被拒绝但 tiling 侧能通过，两侧行为不一致。

likedislike
M
mayafei
4月3日 评论：
84+ aivCoreNum_ <= (2 * optiling::CORE_NUM)),
85+ false, "core num invalid aic:%u aiv:%u", aicCoreNum_,
86+ aivCoreNum_); // more limit check with platform-core
87+ 
88+ return true;
89+}
90+ 
91+bool QuantSalsIndexerMetaDataCpuKernel::ParamsInit() {
92+ groupSize_ = 1U;
93+ mBaseSize_ = 1U;
94+ s2BaseSize_ = 2048U;
95+ return true;
96+}
97+ 
98+uint32_t QuantSalsIndexerMetaDataCpuKernel::GetS1SeqSize(uint32_t bIdx)
99+{
100+ return 1U;
101+}
102+ 
103+uint32_t QuantSalsIndexerMetaDataCpuKernel::GetS2SeqSize(uint32_t bIdx)
104+{
105+ uint32_t s2Size = 0;
106+ if (actSeqLenKV_ == nullptr) {
107+ s2Size = kvSeqSize_;
108+ } else {
109+ const int32_t *s2Ptr = (int32_t*)actSeqLenKV_->GetData();
110+ s2Size = static_cast<uint32_t>(s2Ptr[bIdx]);
111+ }
112+ 
113+ // 即是S2为0，QSI依然存在刷-1和搬运动作。如果actual kv中存大量0，需要对刷新和搬运操作计算负载，避免集中在一个vector上操作
114+ // 而如果存在部分0，刷新和搬运操作理应会被正常计算流水掩盖，所以需要给一个合适小的值。
115+ // 此处返回512U是计算测试所得经验值。
116+ int32_t totalNCount = (s2Size + sparseBlockSize_ - 1) / sparseBlockSize_;
117+ if (totalNCount <= fixedTailCount_) {
118+ return 512U;
119+ }
120+ return (totalNCount - fixedTailCount_) * sparseBlockSize_;
121+}
122+ 
123+void QuantSalsIndexerMetaDataCpuKernel::CalcSplitInfo(SplitContext &splitContext)
124+{
125+ // 计算每个batch的切分，统计是否为空batch，记录最后有效batch（每个batch的每个N2切分是一样的）
126+ SplitInfo &splitInfo = splitContext.splitInfo;
127+ for (uint32_t bIdx = 0; bIdx < batchSize_; bIdx++) {
128+ uint32_t s1Size = GetS1SeqSize(bIdx);
129+ uint32_t s2Size = GetS2SeqSize(bIdx);
130+ splitInfo.s1GBaseNum[bIdx] = (s1Size * groupSize_ + (mBaseSize_ - 1U)) / mBaseSize_;
131+ splitInfo.s1GTailSize[bIdx] = (s1Size * groupSize_) % mBaseSize_;
132+ splitInfo.s2BaseNum[bIdx] = (s2Size + s2BaseSize_ - 1U) / s2BaseSize_;
133+ splitInfo.s2TailSize[bIdx] = s2Size % s2BaseSize_;
134+ if (splitInfo.s1GBaseNum[bIdx] != 0U && splitInfo.s2BaseNum[bIdx] != 0U) {
135+ splitInfo.isKvSeqAllZero = false;
136+ }
137+ }
138+ return;
139+}
140+ 
141+int64_t QuantSalsIndexerMetaDataCpuKernel::CalcCost(
142+ uint32_t basicM, uint32_t basicS2)
143+{
144+ uint32_t alignCoefM = 16U;
145+ uint32_t alignCoefS2 = 64U;
146+ uint32_t alignBasicM = (basicM + alignCoefM - 1U) >> 4U; // 按alignCoefM对齐，向上取整，4：移位操作实现除16
147+ uint32_t alignBasicS2 = (basicS2 + alignCoefS2 - 1U) >> 6U; // 按alignCoefS2对齐，向上取整，6：移位操作实现除64
148+ return static_cast<int64_t>(6U * alignBasicM + 10U * alignBasicS2); // 6：M轴系数，10：S2轴系数
149+}
150+ 
151+BlockCost<int64_t> QuantSalsIndexerMetaDataCpuKernel::CalcCostTable(uint32_t s1NormalSize,
152+ uint32_t s2NormalSize, uint32_t s1GTailSize, uint32_t s2TailSize)
153+{
154+ BlockCost<int64_t> typeCost {};
155+ typeCost[NORMAL_BLOCK][NORMAL_BLOCK] = CalcCost(s1NormalSize, s2NormalSize);
156+ typeCost[TAIL_BLOCK][NORMAL_BLOCK] = (s1GTailSize == 0U) ? 0U : CalcCost(s1GTailSize, s2NormalSize);
157+ typeCost[NORMAL_BLOCK][TAIL_BLOCK] = (s2TailSize == 0U) ? 0U : CalcCost(s1NormalSize, s2TailSize);
158+ typeCost[TAIL_BLOCK][TAIL_BLOCK] = (s1GTailSize == 0U || s2TailSize == 0U) ? 0U : CalcCost(s1GTailSize, s2TailSize);
159+ return typeCost;
160+}
161+ 
162+Range<uint32_t> QuantSalsIndexerMetaDataCpuKernel::CalcS2Range(
163+ uint32_t s1GIdx, const BatchCache &batchCache)
164+{
165+ uint32_t s2Start = 0U;
166+ uint32_t s2End = 0U;
167+ 
168+ if (batchCache.s1Size == 0U || batchCache.s2Size == 0U) {
169+ return std::make_pair(s2Start, s2End);
170+ }
171+ 
172+ s2End = (batchCache.s2Size + s2BaseSize_ - 1U) / s2BaseSize_;
173+ return std::make_pair(s2Start, s2End);
174+}
175+ 
176+void QuantSalsIndexerMetaDataCpuKernel::CalcBatchCache(
177+ uint32_t bIdx, const SplitContext &splitContext, BatchCache &batchCache)
178+{
179+ const SplitInfo &splitInfo = splitContext.splitInfo;
180+ 
181+ batchCache.bIdx = bIdx;
182+ batchCache.s1Size = GetS1SeqSize(bIdx);
183+ batchCache.s2Size = GetS2SeqSize(bIdx);
184+ batchCache.typeCost = CalcCostTable(mBaseSize_, s2BaseSize_, splitInfo.s1GTailSize[bIdx],
185+ splitInfo.s2TailSize[bIdx]);
186+}
187+ 
188+void QuantSalsIndexerMetaDataCpuKernel::CalcS1GCache(uint32_t s1GIdx,
189+ const SplitContext &splitContext, const BatchCache &batchCache, S1GCache &s1GCache)
190+{
191+ const SplitInfo &splitInfo = splitContext.splitInfo;
192+ 
193+ s1GCache.bIdx = batchCache.bIdx;
194+ s1GCache.s1GIdx = s1GIdx;
195+ 
196+ auto s2Range = CalcS2Range(s1GIdx, batchCache);
197+ s1GCache.s2Start = s2Range.first;
198+ s1GCache.s2End = s2Range.second;
199+ 
200+ if (s1GCache.s2Start >= s1GCache.s2End) {
201+ s1GCache.s1GBlock = 0;
202+ s1GCache.s1GCost = 0;
203+ s1GCache.s1GLastBlockCost = 0;
204+ s1GCache.s1GNormalBlockCost = 0;
205+ return;
206+ }
207+ 
208+ // 计算S2方向满块、尾块数量
209+ s1GCache.s1GBlock = s1GCache.s2End - s1GCache.s2Start;
210+ uint32_t curTailS2Num = (splitInfo.s2TailSize[batchCache.bIdx] != 0U &&
211+ s1GCache.s2End == splitInfo.s2BaseNum[batchCache.bIdx]) ? 1U : 0U;
212+ uint32_t curNormalS2Num = s1GCache.s1GBlock - curTailS2Num;
213+ if (splitInfo.s1GBaseNum[batchCache.bIdx] == 0) {
214+ s1GCache.s1GCost = 0;
215+ s1GCache.s1GLastBlockCost = 0;
216+ s1GCache.s1GNormalBlockCost = 0;
217+ } else if (s1GIdx == (splitInfo.s1GBaseNum[batchCache.bIdx] - 1U) && splitInfo.s1GTailSize[batchCache.bIdx] != 0U) {
218+ s1GCache.s1GCost = batchCache.typeCost[TAIL_BLOCK][NORMAL_BLOCK] * curNormalS2Num +
219+ batchCache.typeCost[TAIL_BLOCK][TAIL_BLOCK] * curTailS2Num;
220+ s1GCache.s1GLastBlockCost = curTailS2Num > 0U ? batchCache.typeCost[TAIL_BLOCK][TAIL_BLOCK] :
221+ batchCache.typeCost[TAIL_BLOCK][NORMAL_BLOCK];
222+ s1GCache.s1GNormalBlockCost = batchCache.typeCost[TAIL_BLOCK][NORMAL_BLOCK];
223+ } else {
224+ s1GCache.s1GCost = batchCache.typeCost[NORMAL_BLOCK][NORMAL_BLOCK] * curNormalS2Num +
225+ batchCache.typeCost[NORMAL_BLOCK][TAIL_BLOCK] * curTailS2Num;
226+ s1GCache.s1GLastBlockCost = curTailS2Num > 0U ? batchCache.typeCost[NORMAL_BLOCK][TAIL_BLOCK] :
227+ batchCache.typeCost[NORMAL_BLOCK][NORMAL_BLOCK];
228+ s1GCache.s1GNormalBlockCost = batchCache.typeCost[NORMAL_BLOCK][NORMAL_BLOCK];
229+ }
230+}
231+ 
232+void QuantSalsIndexerMetaDataCpuKernel::CalcBatchCost(
233+ uint32_t bIdx, const SplitContext &splitContext, CostInfo &costInfo)
234+{
235+ const SplitInfo &splitInfo = splitContext.splitInfo;
236+ 
237+ costInfo.bN2CostOfEachBatch[bIdx] = 0;
238+ costInfo.bN2BlockOfEachBatch[bIdx] = 0U;
239+ costInfo.bN2LastBlockCostOfEachBatch[bIdx] = 0U;
240+ 
241+ if (GetS1SeqSize(bIdx) == 0U || GetS2SeqSize(bIdx) == 0U) {
242+ return;
243+ }
244+ 
245+ BatchCache bCache;
246+ S1GCache s1GCache;
247+ CalcBatchCache(bIdx, splitContext, bCache);
248+ for (uint32_t s1GIdx = 0; s1GIdx < splitInfo.s1GBaseNum[bIdx]; s1GIdx++) {
249+ CalcS1GCache(s1GIdx, splitContext, bCache, s1GCache);
250+ costInfo.bN2CostOfEachBatch[bIdx] += s1GCache.s1GCost;
251+ costInfo.bN2BlockOfEachBatch[bIdx] += s1GCache.s1GBlock;
252+ if(s1GCache.s1GBlock > 0){
253+ costInfo.bN2LastBlockCostOfEachBatch[bIdx] = s1GCache.s1GLastBlockCost;
254+ }
255+ }
256+}
257+ 
258+void QuantSalsIndexerMetaDataCpuKernel::CalcCostInfo(SplitContext &splitContext)
259+{
260+ const SplitInfo &splitInfo = splitContext.splitInfo;
261+ CostInfo &costInfo = splitContext.costInfo;
262+ 
263+ if (splitInfo.isKvSeqAllZero) {
264+ costInfo.totalCost = 0;
265+ costInfo.totalBlockNum = 0U;
266+ return;
267+ }
268+ 
269+ // 计算batch的负载并记录，用于按batch分配，需要按行计算起止点，统计块数、负载
270+ for (uint32_t bIdx = 0; bIdx < batchSize_; bIdx++) {
271+ CalcBatchCost(bIdx, splitContext, costInfo);
272+ costInfo.totalCost += costInfo.bN2CostOfEachBatch[bIdx] * kvHeadNum_;
273+ costInfo.totalBlockNum += costInfo.bN2BlockOfEachBatch[bIdx] * kvHeadNum_;
274+ }
275+}
276+ 
277+void QuantSalsIndexerMetaDataCpuKernel::UpdateCursor(const SplitContext &splitContext, AssignContext &assignContext)
278+{
279+ const SplitInfo &splitInfo = splitContext.splitInfo;
280+ const CostInfo &costInfo = splitContext.costInfo;
281+ 
282+ bool UpdateS1G = false;
283+ bool UpdateBatch = false;
284+ 
285+ // Update S2
286+ if (assignContext.curS2Idx >= assignContext.s1GCache.s2End) { // 边界assignInfo.s2End是取不到的开区间
287+ assignContext.curS2Idx = 0U;
288+ assignContext.curS1GIdx++;
289+ UpdateS1G = true;
290+ }
291+ 
292+ // Update S1G
293+ if (assignContext.curS1GIdx >= splitInfo.s1GBaseNum[assignContext.curBIdx]) {
294+ assignContext.curS1GIdx = 0U;
295+ assignContext.curBN2Idx++;
296+ }
297+ 
298+ // Update Batch
299+ if (assignContext.curBN2Idx == batchSize_ * kvHeadNum_) { // 所有负载全部分配完，设置最后一个核的右开区间，返回
300+ assignContext.curS1GIdx = 0U;
301+ assignContext.curS2Idx = 0U;
302+ assignContext.isFinished = true;
303+ return;
304+ }
305+ 
306+ if (assignContext.curBN2Idx / kvHeadNum_ != assignContext.curBIdx) {
307+ assignContext.curBIdx = assignContext.curBN2Idx / kvHeadNum_;
308+ assignContext.curS1GIdx = 0U;
309+ UpdateBatch = true;
310+ UpdateS1G = true;
311+ }
312+ 
313+ // Update Cache
314+ if (UpdateBatch) {
315+ CalcBatchCache(assignContext.curBIdx, splitContext, assignContext.batchCache);
316+ assignContext.bN2Cost = costInfo.bN2CostOfEachBatch[assignContext.curBIdx];
317+ assignContext.bN2Block = costInfo.bN2BlockOfEachBatch[assignContext.curBIdx];
318+ }
319+ if (UpdateS1G) {
320+ CalcS1GCache(assignContext.curS1GIdx, splitContext, assignContext.batchCache, assignContext.s1GCache);
321+ assignContext.curS2Idx = assignContext.s1GCache.s2Start;
322+ }
323+}
324+ 
325+void QuantSalsIndexerMetaDataCpuKernel::AssignByBatch(const SplitContext &splitContext, AssignContext &assignContext)
326+{
327+ if (assignContext.isFinished) {
328+ return;
329+ }
330+ const CostInfo &costInfo = splitContext.costInfo;
331+ while (assignContext.bN2Cost == 0 || IsWithinTolerance(assignContext.coreCache.costLimit,
332+ costInfo.bN2LastBlockCostOfEachBatch[assignContext.curBIdx] / FA_TOLERANCE_RATIO,
333+ assignContext.coreCache.cost + assignContext.bN2Cost)) {
334+ assignContext.coreCache.cost += assignContext.bN2Cost;
335+ assignContext.coreCache.block += assignContext.bN2Block;
336+ assignContext.curBN2Idx++;
337+ 
338+ // to the end
339+ if (assignContext.curBN2Idx == batchSize_ * kvHeadNum_) {
340+ assignContext.curS1GIdx = 0U;
341+ assignContext.curS2Idx = 0U;
342+ assignContext.isFinished = true;
343+ return;
344+ }
345+ 
346+ // next batch
347+ if (assignContext.curBN2Idx / kvHeadNum_ != assignContext.curBIdx) {
348+ assignContext.curBIdx = assignContext.curBN2Idx / kvHeadNum_;
349+ CalcBatchCache(assignContext.curBIdx, splitContext, assignContext.batchCache);
350+ }
351+ 
352+ assignContext.bN2Cost = costInfo.bN2CostOfEachBatch[assignContext.curBIdx];
353+ assignContext.bN2Block = costInfo.bN2BlockOfEachBatch[assignContext.curBIdx];
354+ assignContext.curS1GIdx = 0U;
355+ CalcS1GCache(assignContext.curS1GIdx, splitContext, assignContext.batchCache, assignContext.s1GCache);
356+ assignContext.curS2Idx = assignContext.s1GCache.s2Start;
357+ }
358+}
359+ 
360+void QuantSalsIndexerMetaDataCpuKernel::AssignByRow(const SplitContext &splitContext, AssignContext &assignContext)
361+{
362+ if (assignContext.isFinished) {
363+ return;
364+ }
365+ 
366+ while (IsWithinTolerance(assignContext.coreCache.costLimit,
367+ assignContext.s1GCache.s1GLastBlockCost / FA_TOLERANCE_RATIO,
368+ assignContext.coreCache.cost + assignContext.s1GCache.s1GCost)) {
369+ assignContext.coreCache.cost += assignContext.s1GCache.s1GCost;
370+ assignContext.coreCache.block += assignContext.s1GCache.s1GBlock;
371+ 
372+ // 当前batch被分配一行出去，更新剩余负载
373+ assignContext.bN2Cost = assignContext.bN2Cost > assignContext.s1GCache.s1GCost ?
374+ assignContext.bN2Cost - assignContext.s1GCache.s1GCost : 0;
375+ assignContext.bN2Block = assignContext.bN2Block > assignContext.s1GCache.s1GBlock ?
376+ assignContext.bN2Block - assignContext.s1GCache.s1GBlock : 0U;
377+ // 计算新一行的信息
378+ do{
379+ assignContext.curS1GIdx++;
380+ CalcS1GCache(assignContext.curS1GIdx, splitContext, assignContext.batchCache, assignContext.s1GCache);
381+ }while(assignContext.s1GCache.s1GBlock == 0);
382+ assignContext.curS2Idx = assignContext.s1GCache.s2Start;
383+ }
384+}
385+ 
386+void QuantSalsIndexerMetaDataCpuKernel::AssignByBlock(const SplitContext &splitContext, AssignContext &assignContext)
387+{
388+ if (assignContext.isFinished) {
389+ return;
390+ }
391+ 
392+ int64_t curCost = assignContext.s1GCache.s1GNormalBlockCost;
393+ if (assignContext.curS2Idx == (assignContext.s1GCache.s2End - 1U)) {
394+ curCost = assignContext.s1GCache.s1GLastBlockCost;
395+ }
396+ 
397+ while (IsWithinTolerance(assignContext.coreCache.costLimit, curCost / FA_TOLERANCE_RATIO,
398+ assignContext.coreCache.cost + curCost)) { // (costLimit - curCostOnCore) * FA_TOLERANCE_RATIO > curCost；至少分配1块
399+ assignContext.coreCache.cost += curCost;
400+ assignContext.coreCache.block++;
401+ assignContext.curS2Idx++;
402+ // 当前batch被分配一块出去，更新剩余负载
403+ assignContext.bN2Cost = assignContext.bN2Cost - curCost;
404+ // 当前行被分配一块出去，更新剩余负载
405+ assignContext.s1GCache.s1GCost = assignContext.s1GCache.s1GCost - curCost;
406+ assignContext.bN2Block--;
407+ assignContext.s1GCache.s1GBlock--;
408+ }
409+}
410+ 
411+void QuantSalsIndexerMetaDataCpuKernel::ForceAssign(const SplitContext &splitContext, AssignContext &assignContext)
412+{
413+ if (assignContext.isFinished) {
414+ return;
415+ }
416+ 
417+ int64_t curCost = assignContext.s1GCache.s1GNormalBlockCost;
418+ if (assignContext.curS2Idx == (assignContext.s1GCache.s2End - 1U)) {
419+ curCost = assignContext.s1GCache.s1GLastBlockCost;
420+ }
421+ 
422+ assignContext.coreCache.cost += curCost;
423+ assignContext.coreCache.block++;
424+ assignContext.curS2Idx++;
425+ // 当前batch被分配一块出去，更新剩余负载
426+ assignContext.bN2Cost = assignContext.bN2Cost - curCost;
427+ assignContext.bN2Block--;
428+ // 当前行被分配一块出去，更新剩余负载
429+ assignContext.s1GCache.s1GCost = assignContext.s1GCache.s1GCost - curCost;
430+ assignContext.s1GCache.s1GBlock--;
431+ UpdateCursor(splitContext, assignContext);
432+}
433+ 
434+void QuantSalsIndexerMetaDataCpuKernel::CalcSplitPlan(uint32_t coreNum,
435+ int64_t costLimit, const SplitContext &splitContext, SplitResult &result)
436+{
437+ const CostInfo &costInfo = splitContext.costInfo;
438+ 
439+ if (coreNum == 0U) {
440+ return;
441+ }
442+ result.maxCost = 0U;
443+ result.usedCoreNum = 0U;
444+ 
445+ AssignContext assignContext {};
446+ assignContext.curBIdx = 0U;
447+ assignContext.curS1GIdx = 0U;
448+ assignContext.unassignedCost = costInfo.totalCost;
449+ assignContext.bN2Cost = costInfo.bN2CostOfEachBatch[assignContext.curBIdx];
450+ assignContext.bN2Block = costInfo.bN2BlockOfEachBatch[assignContext.curBIdx];
451+ CalcBatchCache(assignContext.curBIdx, splitContext, assignContext.batchCache);
452+ CalcS1GCache(assignContext.curS1GIdx, splitContext, assignContext.batchCache, assignContext.s1GCache);
453+ assignContext.curS2Idx = assignContext.s1GCache.s2Start;
454+ 
455+ for (uint32_t i = 0; i < coreNum; ++i) {
456+ if (result.maxCost > costLimit) {
457+ return;
458+ }
459+ if (assignContext.isFinished || assignContext.unassignedCost <= 0) {
460+ break;
461+ }
462+ 
463+ assignContext.curCoreIdx = i;
464+ 
465+ assignContext.coreCache = {};
466+ assignContext.coreCache.costLimit = assignContext.unassignedCost / (coreNum - assignContext.curCoreIdx);
467+ 
468+ // 1、按整batch分配
469+ AssignByBatch(splitContext, assignContext);
470+ // 2、按行分配
471+ AssignByRow(splitContext, assignContext);
472+ // 3、按块分配
473+ AssignByBlock(splitContext, assignContext);
474+ // 4、强制分配
475+ if (assignContext.coreCache.block == 0) {
476+ ForceAssign(splitContext, assignContext);
477+ }
478+ 
479+ result.bN2End[i] = assignContext.curBN2Idx;
480+ result.gS1End[i] = assignContext.curS1GIdx;
481+ result.s2End[i] = assignContext.curS2Idx;
482+ result.maxCost = std::max(result.maxCost, assignContext.coreCache.cost);
483+ 
484+ assignContext.unassignedCost -= assignContext.coreCache.cost;
485+ }
486+ 
487+ result.usedCoreNum = assignContext.curCoreIdx + 1;
488+}
489+ 
490+void QuantSalsIndexerMetaDataCpuKernel::CopyTmpResult(SplitResult &tmpRes, SplitResult &splitRes)
491+{
492+ uint64_t len = tmpRes.bN2End.size();
493+ splitRes.usedCoreNum = tmpRes.usedCoreNum;
494+ splitRes.maxCost = tmpRes.maxCost;
495+ 
496+ for (size_t i = 0; i < len; ++i) {
497+ splitRes.bN2End[i] = tmpRes.bN2End[i];
498+ splitRes.gS1End[i] = tmpRes.gS1End[i];
499+ splitRes.s2End[i] = tmpRes.s2End[i];
500+ }
501+}
502+ 
503+void QuantSalsIndexerMetaDataCpuKernel::ClearTmpResult(SplitResult &tmpResult)
504+{
505+ uint64_t len = tmpResult.bN2End.size();
506+ tmpResult.usedCoreNum = 0U;
507+ tmpResult.maxCost = 0;
508+ 
509+ for (size_t i = 0; i < len; ++i) {
510+ tmpResult.bN2End[i] = 0U;
511+ tmpResult.gS1End[i] = 0U;
512+ tmpResult.s2End[i] = 0U;
513+ }
514+}
515+ 
516+void QuantSalsIndexerMetaDataCpuKernel::RollBackCursor(const SplitContext &splitContext,
517+ const CostInfo &costInfo, SplitResult &splitRes)
518+{
519+ for (size_t i = 0; i < splitRes.usedCoreNum; ++i) {
520+ // x, y, z
521+ if (splitRes.s2End[i] > 0U) {
522+ splitRes.s2End[i] = splitRes.s2End[i] - 1U;
523+ continue;
524+ }
525+ uint32_t bIdx = splitRes.bN2End[i] / kvHeadNum_;
526+ // x, y, 0
527+ if (splitRes.gS1End[i] > 0U) {
528+ splitRes.gS1End[i] = splitRes.gS1End[i] - 1U;
529+ splitRes.s2End[i] = splitContext.splitInfo.s2BaseNum[bIdx] - 1U;
530+ continue;
531+ }
532+ 
533+ // x, 0, 0
534+ uint32_t bN2Idx = splitRes.bN2End[i] > 0U ? splitRes.bN2End[i] - 1U : 0U;
535+ bIdx = bN2Idx / kvHeadNum_;
536+ 
537+ // last end point
538+ if (i == splitRes.usedCoreNum - 1U && costInfo.bN2BlockOfEachBatch[bIdx] == 0U) {
539+ splitRes.bN2End[i] = batchSize_ * kvHeadNum_ - 1;
540+ splitRes.gS1End[i] = splitContext.splitInfo.s1GBaseNum[bIdx] > 0 ?
541+ splitContext.splitInfo.s1GBaseNum[bIdx] - 1U : 0U;
542+ splitRes.s2End[i] = splitContext.splitInfo.s2BaseNum[bIdx] > 0 ?
543+ splitContext.splitInfo.s2BaseNum[bIdx] - 1U : 0U;
544+ continue;
545+ }
546+ 
547+ while (bN2Idx > 0U && costInfo.bN2BlockOfEachBatch[bIdx] == 0U) {
548+ bN2Idx -= 1U;
549+ bIdx = bN2Idx / kvHeadNum_;
550+ }
551+ 
552+ if (costInfo.bN2BlockOfEachBatch[bIdx] != 0U) {
553+ splitRes.bN2End[i] = bN2Idx;
554+ splitRes.gS1End[i] = splitContext.splitInfo.s1GBaseNum[bIdx] - 1U;
555+ splitRes.s2End[i] = splitContext.splitInfo.s2BaseNum[bIdx] - 1U;
556+ } else {
557+ splitRes.bN2End[i] = 0U;
558+ splitRes.gS1End[i] = 0U;
559+ splitRes.s2End[i] = 0U;
560+ }
561+ }
562+}
563+ 
564+bool QuantSalsIndexerMetaDataCpuKernel::BalanceSchedule() {
565+ SplitContext splitContext(batchSize_);
566+ 
567+ // 1、划分基本块，统计信息
568+ CalcSplitInfo(splitContext);
569+ // 全空case
570+ if (splitContext.splitInfo.isKvSeqAllZero) {
571+ splitRes_.usedCoreNum = 1U;
572+ splitRes_.bN2End[0] = batchSize_ * kvHeadNum_ - 1;
573+ splitRes_.gS1End[0] = 0U;
574+ splitRes_.s2End[0] = 0U;
575+ return true;
576+ }
577+ CalcCostInfo(splitContext);
578+ 
579+ // 2、获取每个核的分配方案
580+ uint32_t maxCore = std::min(coreNum_, splitContext.costInfo.totalBlockNum);
581+ uint32_t minCore = static_cast<uint32_t>(
582+ std::sqrt(static_cast<float>(splitContext.costInfo.totalBlockNum) + 0.25f) + 0.5f);
583+ minCore = std::min(minCore, maxCore);
584+ 
585+ splitRes_.maxCost = INT64_MAX;
586+ splitRes_.usedCoreNum = 1U;
587+ SplitResult tmpResult {coreNum_, aivCoreNum_ / aicCoreNum_}; // C: V = 1: 2, TODO: C: V = 1: 1 ?
588+ for (uint32_t i = minCore; i <= maxCore; ++i) {
589+ CalcSplitPlan(i, splitRes_.maxCost, splitContext, tmpResult);
590+ if (tmpResult.maxCost < splitRes_.maxCost) {
591+ CopyTmpResult(tmpResult, splitRes_);
592+ }
593+ ClearTmpResult(tmpResult);
594+ }
595+ 
596+ splitRes_.usedCoreNum = std::max(splitRes_.usedCoreNum, 1U); // 至少使用1个core
597+ RollBackCursor(splitContext, splitContext.costInfo, splitRes_);
598+ if (GetS2SeqSize(batchSize_-1) == 0U) {
599+ splitRes_.bN2End[splitRes_.usedCoreNum-1] = batchSize_ * kvHeadNum_ - 1;
600+ splitRes_.gS1End[splitRes_.usedCoreNum-1] = 0U;
601+ splitRes_.s2End[splitRes_.usedCoreNum-1] = 0U;
602+ }
603+ return true;
604+}
605+ 
606+bool QuantSalsIndexerMetaDataCpuKernel::GenMetaData() {
607+ optiling::detail::QsiMetaData* metaDataPtr = (optiling::detail::QsiMetaData*)metaData_->GetData();
608+ metaDataPtr->usedCoreNum = splitRes_.usedCoreNum;
609+ 
610+ for (size_t i = 0; i < coreNum_; ++i) {
611+ metaDataPtr->bN2End[i] = splitRes_.bN2End[i];
612+ metaDataPtr->gS1End[i] = splitRes_.gS1End[i];
613+ metaDataPtr->s2End[i] = splitRes_.s2End[i];
614+ }
615+ 
616+ 
617+ return true;
618+}
619+ 
620+static const char *qsiKernelType = "QuantSalsIndexerMetadata";
621+REGISTER_CPU_KERNEL(qsiKernelType, QuantSalsIndexerMetaDataCpuKernel);
622+ 
623+}; // namespace aicpu
624+ 
Aexperimental/attention/quant_sals_indexer_metadata/op_kernel_aicpu/quant_sals_indexer_metadata_aicpu.h+234-0
@@ -0,0 +1,234 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file quant_sals_indexer_metadata_aicpu.h
13+ * \brief
14+ */
15+ 
16+#ifndef QUANT_SALS_INDEXER_METADATA_AICPU_H
17+#define QUANT_SALS_INDEXER_METADATA_AICPU_H
18+ 
19+#include "cpu_context.h"
20+#include "cpu_kernel.h"
21+#include "cpu_tensor.h"
22+#include <array>
23+#include <cstdint>
24+#include <string>
25+#include <vector>
26+ 
27+namespace aicpu {
28+constexpr int64_t FA_TOLERANCE_RATIO = 2;
29+constexpr uint32_t FD_TOLERANCE_RATIO = 2U;
30+ 
31+enum BlockType : uint32_t {
32+ NORMAL_BLOCK = 0,
33+ TAIL_BLOCK,
34+ BLOCK_MAX_TYPE
35+};
36+ 
37+template<class T>
38+using Range = std::pair<T, T>;
39+ 
40+template<class T>
41+using BlockCost = std::array<std::array<T, static_cast<size_t>(BLOCK_MAX_TYPE)>, static_cast<size_t>(BLOCK_MAX_TYPE)>;
42+ 
43+template<typename T>
44+inline bool IsWithinTolerance(T limit, T tolerance, T value)
45+{
46+ return limit + tolerance >= value;
47+}
48+ 
49+// 分核功能模块输出：FA阶段的核间分核信息
50+struct SplitResult {
51+ uint32_t usedCoreNum { 0U }; // 使用的核数量
52+ uint32_t vecCubeRatio { 0U }; // vec 与 cube 核数比例
53+ std::vector<uint32_t> bN2End {}; // 每个核处理数据的BN2结束点
54+ std::vector<uint32_t> gS1End {}; // 每个核处理数据的GS1结束点
55+ std::vector<uint32_t> s2End {}; // 每个核处理数据的S2结束点
56+ int64_t maxCost { 0 }; // 慢核开销
57+ 
58+ SplitResult(uint32_t coreNum, uint32_t ratio) :
59+ bN2End(coreNum),
60+ vecCubeRatio(ratio),
61+ gS1End(coreNum),
62+ s2End(coreNum) {};
63+};
64+ 
65+// 分核功能模块内部使用：记录切分信息
66+struct SplitInfo {
67+ std::vector<uint32_t> s1GBaseNum {}; // S1G方向，切了多少个基本块
68+ std::vector<uint32_t> s2BaseNum {}; // S2方向，切了多少个基本块
69+ std::vector<uint32_t> s1GTailSize {}; // S1G方向，尾块size
70+ std::vector<uint32_t> s2TailSize {}; // S2方向，尾块size
71+ bool isKvSeqAllZero { true };
72+ 
73+ explicit SplitInfo(uint32_t batchSize) :
74+ s1GBaseNum(batchSize),
75+ s2BaseNum(batchSize),
76+ s1GTailSize(batchSize),
77+ s2TailSize(batchSize) {}
78+};
79+ 
80+// 分核功能模块内部使用：记录batch的开销信息
81+struct CostInfo {
82+ std::vector<int64_t> bN2CostOfEachBatch {}; // 整个batch的开销
83+ std::vector<uint32_t> bN2BlockOfEachBatch {}; // 整个batch的开销
84+ std::vector<int64_t> bN2LastBlockCostOfEachBatch {}; // batch最后一块的开销
85+ uint32_t totalBlockNum { 0U };
86+ int64_t totalCost { 0 };
87+ 
88+ explicit CostInfo(uint32_t batchSize) :
89+ bN2CostOfEachBatch(batchSize),
90+ bN2BlockOfEachBatch(batchSize),
91+ bN2LastBlockCostOfEachBatch(batchSize) {}
92+};
93+ 
94+// 分核功能模块内部使用：分核过程中，case基本信息的上下文信息，组合以减少接口传参数量
95+struct SplitContext {
96+ SplitInfo splitInfo { 0U };
97+ CostInfo costInfo { 0U };
98+ 
99+ explicit SplitContext(uint32_t batchSize) :
100+ splitInfo(batchSize),
101+ costInfo(batchSize) {}
102+};
103+ 
104+// 分核功能模块内部使用：记录batch相关的临时信息
105+struct BatchCache {
106+ uint32_t bIdx { 0U };
107+ uint32_t s1Size { 0U };
108+ uint32_t s2Size { 0U };
109+ int64_t preTokenLeftUp { 0 };
110+ int64_t nextTokenLeftUp { 0 };
111+ BlockCost<int64_t> typeCost {};
112+};
113+ 
114+// 分核功能模块内部使用：记录当前行（S1G）的临时信息
115+struct S1GCache {
116+ uint32_t bIdx { 0U };
117+ uint32_t s1GIdx { 0U };
118+ uint32_t s2Start { 0U };
119+ uint32_t s2End { 0U };
120+ int64_t s1GCost { 0 };
121+ int64_t s1GLastBlockCost { 0 };
122+ uint32_t s1GBlock { 0U };
123+ int64_t s1GNormalBlockCost { 0 };
124+};
125+ 
126+// 分核功能模块内部使用：记录分配过程中，当前核的负载信息
127+struct CoreCache {
128+ int64_t costLimit { 0 }; // 负载上限
129+ int64_t cost { 0 }; // 已分配负载
130+ uint32_t block { 0U }; // 已分配块数
131+};
132+ 
133+// 分核功能模块内部使用：记录分配过程中的上下文信息
134+struct AssignContext {
135+ uint32_t curBIdx { 0U };
136+ uint32_t curBN2Idx { 0U };
137+ uint32_t curS1GIdx { 0U };
138+ uint32_t curS2Idx { 0U };
139+ uint32_t curCoreIdx { 0U };
140+ int64_t unassignedCost { 0 };
141+ uint32_t usedCoreNum { 0U };
142+ uint32_t curKvSplitPart { 1U };
143+ 
144+ int64_t bN2Cost { 0 };
145+ uint32_t bN2Block { 0U };
146+ bool isFinished { false };
147+ BatchCache batchCache {};
148+ S1GCache s1GCache {};
149+ CoreCache coreCache {};
150+};
151+ 
152+class QuantSalsIndexerMetaDataCpuKernel : public CpuKernel {
153+public:
154+ QuantSalsIndexerMetaDataCpuKernel() = default;
155+ ~QuantSalsIndexerMetaDataCpuKernel() = default;
156+ uint32_t Compute(CpuKernelContext &ctx) override;
157+ 
158+private:
159+ bool Prepare(CpuKernelContext &ctx);
160+ bool ParamsCheck();
161+ bool ParamsInit();
162+ bool BalanceSchedule();
163+ bool GenMetaData();
164+ 
165+ // util
166+ uint32_t GetS1SeqSize(uint32_t bIdx);
167+ uint32_t GetS2SeqSize(uint32_t bIdx);
168+ uint32_t GetSparseSeqSize(uint32_t bIdx);
169+ int64_t CalcPreTokenLeftUp(uint32_t s1Size, uint32_t s2Size);
170+ int64_t CalcNextTokenLeftUp(uint32_t s1Size, uint32_t s2Size);
171+ Range<uint32_t> CalcS2Range(uint32_t s1GIdx,const BatchCache &batchCache);
172+ int64_t CalcCost(uint32_t basicM, uint32_t basicS2);
173+ BlockCost<int64_t> CalcCostTable(uint32_t s1NormalSize, uint32_t s2NormalSize, uint32_t s1GTailSize,
174+ uint32_t s2TailSize);
175+ 
176+ // cache calculation
177+ void CalcBatchCache(uint32_t bIdx, const SplitContext &splitContext, BatchCache &batchCache);
178+ void CalcS1GCache(uint32_t s1GIdx, const SplitContext &splitContext, const BatchCache &batchCache, S1GCache &s1GCache);
179+ void CopyTmpResult(SplitResult &tmpRes, SplitResult &splitRes);
180+ void ClearTmpResult(SplitResult &tmpRes);
181+ 
182+ // preprocess
183+ void CalcSplitInfo(SplitContext &splitContext);
184+ void CalcBatchCost(uint32_t bIdx, const SplitContext &splitContext, CostInfo &costInfo);
185+ void CalcCostInfo(SplitContext &splitContext);
186+ 
187+ // assign
188+ void UpdateCursor(const SplitContext &splitContext, AssignContext &assignContext);
189+ void AssignByBatch(const SplitContext &splitContext, AssignContext &assignContext);
190+ void AssignByRow(const SplitContext &splitContext, AssignContext &assignContext);
191+ void AssignByBlock(const SplitContext &splitContext, AssignContext &assignContext);
192+ void ForceAssign(const SplitContext &splitContext, AssignContext &assignContext);
193+ 
194+ // FD
195+ bool IsNeedRecordFDInfo(const AssignContext &assignContext, const SplitResult &splitRes);
196+ void RecordFDInfo(const SplitContext &splitContext, const AssignContext &assignContext, SplitResult &result);
197+ 
198+ // main
199+ void SplitFD(SplitResult &result);
200+ void CalcSplitPlan(uint32_t coreNum, int64_t costLimit, const SplitContext &splitContext, SplitResult &result);
201+ void SplitCore();
202+ void RollBackCursor(const SplitContext &splitContext, const CostInfo &costInfo, SplitResult &splitRes);
203+ 
204+private:
205+ CpuKernelContext* context_ = nullptr;
206+ 
207+ Tensor *actSeqLenKV_ = nullptr;
208+ Tensor *metaData_ = nullptr;
209+ 
210+ uint32_t coreNum_ = 24U; // new
211+ uint32_t aicCoreNum_ = 24U;
212+ uint32_t aivCoreNum_ = 48U;
213+ uint32_t batchSize_ = 0;
214+ uint32_t kvSeqSize_ = 0;
215+ uint32_t kvHeadNum_ = 0;
216+ uint32_t fixedTailCount_ = 0;
217+ uint32_t sparseBlockSize_ = 0;
218+ 
219+ // SplitParam
220+ uint32_t groupSize_ = 0;
221+ uint32_t mBaseSize_ = 0;
222+ uint32_t s2BaseSize_ = 0;
223+ SplitResult splitRes_ {24, 2};
224+ 
225+private:
226+ enum class ParamId : uint32_t {
227+ // input
228+ actSeqLenKV = 0,
229+ // output
230+ metaData = 0,
231+ };
232+};
233+} // namespace aicpu
234+#endif
Aexperimental/attention/quant_sals_indexer_metadata/op_kernel_aicpu/quant_sals_indexer_metadata_aicpu.json+15-0
@@ -0,0 +1,15 @@
1+{
2+ "QuantSalsIndexerMetadata":{
3+ "opInfo":{
4+ "computeCost":"100",
5+ "engine":"DNN_VM_AICPU",
6+ "flagAsync":"False",
7+ "flagPartial":"False",
8+ "functionName":"RunCpuKernel",
9+ "kernelSo":"libtransformer_aicpu_kernels.so",
10+ "opKernelLib":"CUSTAICPUKernel",
11+ "userDefined":"True",
12+ "workspaceSize":"100"
13+ }
14+ }
15+}
Aexperimental/attention/sparse_flash_attention_antiquant/CMakeLists.txt+0-0
The file is empty
Aexperimental/attention/sparse_flash_attention_antiquant/examples/test_npu_sparse_flash_attention_antiquant_meta_gqa.py+654-0
@@ -0,0 +1,654 @@
1+# ---------------------------------------------------------------------------------------------------------
2+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+# CANN Open Software License Agreement Version 2.0 (the "License").
5+# Please refer to the License for details. You may not use this file except in compliance with the License.
6+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+# See LICENSE in the root of the software repository for the full text of the License.
9+# ---------------------------------------------------------------------------------------------------------
10+ 
11+import random
12+import torch
13+import torch_npu
14+import torchair
15+import math
16+import custom_ops
17+import numpy as np
18+import torch.nn as nn
19+from torch_npu.testing.testcase import TestCase, run_tests
20+ 
21+# import torch._dynamo
22+# torch._dynamo.reset()
23+ 
24+np.random.seed(234) # 固定随机种子
25+np.set_printoptions(suppress=True)
26+ 
27+DEVICE_ID = 0
28+# DEVICE_ID = 0
29+torch_npu.npu.set_device(int(DEVICE_ID))
30+ 
31+ 
32+def pa_to_bsnd(pa_in, block_table, actual_seq_lengths, layout_kv):
33+ if layout_kv == 'PA_BSND':
34+ block_num, block_size, n, d = pa_in.shape
35+ elif layout_kv == 'PA_BNSD':
36+ block_num, n, block_size, d = pa_in.shape
37+ elif layout_kv == 'PA_NZ':
38+ block_num, n, dn, block_size, ds = pa_in.shape
39+ d = dn * ds
40+ 
41+ b, max_block_num = block_table.shape
42+ out = torch.zeros((b, max(max_block_num, math.ceil(block_num / b)) * block_size, n, d)).to(pa_in.dtype)
43+ if layout_kv == 'PA_BSND':
44+ for i in range(b):
45+ loop = actual_seq_lengths[i] // block_size
46+ for j in range(loop):
47+ out[i, j * block_size: (j + 1) * block_size, :, :] = \
48+ pa_in[block_table[i][j], :, :, :].reshape(block_size, n, d)
49+ tail_len = actual_seq_lengths[i] % block_size
50+ if tail_len > 0:
51+ out[i, loop * block_size : actual_seq_lengths[i], :, :] = \
52+ pa_in[block_table[i][loop], : tail_len, :, :].reshape(tail_len, n, d)
53+ elif layout_kv == 'PA_BNSD':
54+ for i in range(b):
55+ loop = actual_seq_lengths[i] // block_size
56+ for j in range(loop):
57+ out[i, j * block_size: (j + 1) * block_size, :, :] = \
58+ pa_in[block_table[i][j], :, :, :].reshape(n, block_size, d).permute(1, 0, 2)
59+ tail_len = actual_seq_lengths[i] % block_size
60+ if tail_len > 0:
61+ out[i, loop * block_size : actual_seq_lengths[i], :, :] = \
62+ pa_in[block_table[i][loop], :, : tail_len, :].reshape(n, tail_len, d).permute(1, 0, 2)
63+ elif layout_kv == 'PA_NZ':
64+ for i in range(b):
65+ loop = actual_seq_lengths[i] // block_size
66+ for j in range(loop):
67+ out[i, j * block_size: (j + 1) * block_size, :, :] = \
68+ pa_in[block_table[i][j], :, :, :, :].reshape(n, dn, block_size, ds).permute(2, 0, 1, 3).reshape(block_size, n, d)
69+ tail_len = actual_seq_lengths[i] % block_size
70+ if tail_len > 0:
71+ out[i, loop * block_size : actual_seq_lengths[i], :, :] = \
72+ pa_in[block_table[i][loop], :, :, : tail_len, :].reshape(n, dn, tail_len, ds).permute(2, 0, 1, 3).reshape(tail_len, n, d)
73+ return out
74+ 
75+ 
76+def gather_kv(k_tensor, v_tensor, sparse_indices, sparse_block_size, sparse_count,
77+ batch, n2_idx, cur_actual_seq_lengths_q, cur_actual_seq_lengths_kv,
78+ sparse_mode, s1_shard_idx, cur_sparse_seq_lengths_kv):
79+ s2_sparse = list()
80+ if sparse_count > sparse_indices.numel():
81+ raise f"sparse_count({sparse_count}) should less than the length sparse_indices({sparse_indices.numel()})"
82+ if sparse_mode == 0:
83+ threshold = cur_actual_seq_lengths_kv
84+ elif sparse_mode == 3:
85+ threshold = cur_actual_seq_lengths_kv - cur_actual_seq_lengths_q + s1_shard_idx + 1
86+ sparse_kv_lengths_sum = 0
87+ for i in range(sparse_count):
88+ sparse_id = sparse_indices[i]
89+ if sparse_id == -1:
90+ break
91+ begin_idx = sparse_id * sparse_block_size
92+ end_idx = begin_idx + sparse_block_size \
93+ if begin_idx + sparse_block_size <= threshold else threshold
94+ sparse_kv_lengths_sum = sparse_kv_lengths_sum + (end_idx - begin_idx)
95+ if i == sparse_count - 1:
96+ end_idx = end_idx - (sparse_kv_lengths_sum - cur_sparse_seq_lengths_kv) \
97+ if sparse_kv_lengths_sum > cur_sparse_seq_lengths_kv else end_idx
98+ if begin_idx >= threshold:
99+ continue
100+ s2_sparse.extend(torch.arange(begin_idx, end_idx))
101+ 
102+ k_sparse, v_sparse = k_tensor[batch, n2_idx, s2_sparse, :], v_tensor[batch, n2_idx, s2_sparse, :]
103+ 
104+ return k_sparse, v_sparse, torch.tensor(s2_sparse)
105+ 
106+def softmax(x):
107+ x = x.astype(np.float32)
108+ x_max = x.max(axis=-1, keepdims=True)
109+ x_sub = x - x_max
110+ y = np.exp(x_sub)
111+ x_sum = y.sum(axis=-1, keepdims=True)
112+ ans = y / x_sum
113+ return ans
114+ 
115+class SFAANetwork(nn.Module):
116+ def __init__(self):
117+ super(SFAANetwork, self).__init__()
118+ self.dummy_version = 2 # 加个无意义的成员变量，改变代码指纹
119+ 
120+ def forward(self, b, s1, n1, s2, n2, dn,
121+ query, key, value, sparse_indices, key_dequant_scale, value_dequant_scale, scale_value, sparse_block_size,
122+ actual_seq_lengths_query, actual_seq_lengths_kv, sparse_seq_lengths_kv, layout_query, layout_kv, sparse_mode, block_table,
123+ attention_mode, quant_scale_repo_mode, tile_size, key_quant_mode, value_quant_mode, rope_head_dim, sparse_shard_size):
124+ # super kernel test
125+ # with torchair.scope.super_kernel("sp_QsSFAA", "stream-fusion=1:dcci-before-kernel-start=SparseFlashAttentionAntiquant"):
126+ metadata = torch.ops.custom.npu_sparse_flash_attention_antiquant_metadata(
127+ b, s1, n1, s2, n2, dn,
128+ 0,
129+ sparse_block_size,
130+ sparse_shared_size=sparse_shard_size,
131+ actual_seq_lengths_query=actual_seq_lengths_query,
132+ actual_seq_lengths_kv=actual_seq_lengths_kv,
133+ sparse_seq_lengths_kv=sparse_seq_lengths_kv)
134+ # y = torch.broadcast_to(torch.tensor(1, device='npu:14', dtype=torch.float32), (192*1024*1024,))
135+ # value_dq_scale_new = y.flatten()[:value_dequant_scale.numel()].reshape(value_dequant_scale.shape).to(torch.float32)
136+ output0 = torch_npu.npu_sparse_flash_attention_antiquant(query, key, value, sparse_indices,
137+ key_dequant_scale=key_dequant_scale, value_dequant_scale=value_dequant_scale,
138+ scale_value=scale_value, sparse_block_size=16,metadata=metadata,
139+ actual_seq_lengths_query=actual_seq_lengths_query, actual_seq_lengths_kv=actual_seq_lengths_kv,
140+ sparse_seq_lengths_kv = sparse_seq_lengths_kv, layout_query=layout_query, layout_kv=layout_kv,
141+ sparse_mode=sparse_mode, block_table=block_table, attention_mode=attention_mode,
142+ quant_scale_repo_mode=quant_scale_repo_mode, tile_size=tile_size, key_quant_mode=key_quant_mode,
143+ value_quant_mode=value_quant_mode, rope_head_dim=64, sparse_shard_size=sparse_shard_size)
144+ return output0
145+ 
146+ 
147+def cpu_sparse_flash_attention_antiquant_gqa(
148+ query, key, value, sparse_indices, key_dequant_scale, value_dequant_scale,
149+ scale_value, sparse_block_size,
150+ actual_seq_lengths_query, actual_seq_lengths_kv, sparse_seq_lengths_kv,
151+ layout_query='BSND', layout_kv='PA_BSND', sparse_mode=3, block_table=None,
152+ attention_mode=0, quant_scale_repo_mode=0, tile_size=0, key_quant_mode=0,
153+ value_quant_mode=0, rope_head_dim=0, sparse_shard_size=1):
154+ """
155+ CPU计算的sparse flash attention函数
156+ """
157+ query = query.to(torch.float32)
158+ query_type = query.dtype
159+ head_dim = query.shape[-1]
160+
161+ # 将PA格式转换为BSND格式
162+ key = pa_to_bsnd(key, block_table, actual_seq_lengths_kv, layout_kv)
163+ value = pa_to_bsnd(value, block_table, actual_seq_lengths_kv, layout_kv)
164+ 
165+ # 反量化
166+ key = (key.to(torch.float32) * key_dequant_scale).to(query_type)
167+ value = (value.to(torch.float32) * value_dequant_scale).to(query_type)
168+
169+ batch_size = actual_seq_lengths_query.shape[0]
170+ if layout_query == "TND":
171+ num_heads = query.shape[1]
172+ else:
173+ num_heads = query.shape[2]
174+ num_kv_heads = key.shape[2]
175+ g = num_heads // num_kv_heads
176+ 
177+ # 转置为B, N, S, D格式
178+ q_bnsd_tensor = torch.transpose(query, 1, 2)
179+ k_bnsd_tensor = torch.transpose(key, 1, 2)
180+ v_bnsd_tensor = torch.transpose(value, 1, 2)
181+ sparse_indices_tensor = torch.transpose(sparse_indices, 1, 2)
182+ out_shape_bnsd = list(q_bnsd_tensor.shape)
183+ y = torch.zeros(out_shape_bnsd, dtype=query_type)
184+ 
185+ # 遍历每个batch
186+ for batch in range(batch_size):
187+ cur_acutal_seq_lengths_q = actual_seq_lengths_query[batch]
188+ if layout_query == "TND" and batch > 0:
189+ cur_acutal_seq_lengths_q = actual_seq_lengths_query[batch] - actual_seq_lengths_query[batch - 1]
190+ cur_actual_seq_lengths_kv = actual_seq_lengths_kv[batch]
191+ cur_sparse_seq_lengths_kv = sparse_seq_lengths_kv[batch]
192+ sparse_count = math.ceil(sparse_seq_lengths_kv[batch] / sparse_block_size)
193+
194+ # 遍历每个KV头
195+ for n2_idx in range(num_kv_heads):
196+ # 遍历每个query token
197+ for s1_shard_idx in range(cur_acutal_seq_lengths_q):
198+ # 获取当前query
199+ q_curr = q_bnsd_tensor[batch, n2_idx * g: (n2_idx + 1) * g, s1_shard_idx, :].squeeze(1)
200+
201+ # 获取当前稀疏索引
202+ cur_sparse_indices = sparse_indices_tensor[batch, n2_idx, s1_shard_idx // sparse_shard_size, :]
203+
204+ # 根据稀疏索引收集key和value
205+ k_sparse, v_sparse, s2_index = gather_kv(k_bnsd_tensor, v_bnsd_tensor, cur_sparse_indices, sparse_block_size,
206+ sparse_count, batch, n2_idx, cur_acutal_seq_lengths_q, cur_actual_seq_lengths_kv,
207+ sparse_mode, s1_shard_idx, cur_sparse_seq_lengths_kv)
208+
209+ # 计算attention
210+ if k_sparse.shape[0] == 0: # 没有有效的key/value
211+ # 返回零向量
212+ mm2_res = torch.zeros((g, head_dim), dtype=query_type)
213+ else:
214+ mm1_res = torch.matmul(q_curr.to(torch.float32), k_sparse.to(torch.float32).T)
215+ scale_res = mm1_res * scale_value
216+ if scale_res.numel() != 0:
217+ softmax_res = softmax(scale_res.numpy())
218+ else:
219+ softmax_res = torch.zeros_like(scale_res)
220+ softmax_res = torch.tensor(softmax_res).to(query_type)
221+ mm2_res = torch.matmul(softmax_res.to(torch.float32), v_sparse.to(torch.float32))
222+
223+ # 处理输出维度
224+ if mm2_res.dim() == 1:
225+ mm2_res = mm2_res.unsqueeze(0)
226+ mm2_res = mm2_res.reshape(g, 1, head_dim) # 从(g, d)变为(g, 1, d)
227+ y[batch, n2_idx * g: (n2_idx + 1) * g, s1_shard_idx : s1_shard_idx + 1, :] = mm2_res
228+
229+ # 转置回B, S, N, D格式
230+ return torch.transpose(y, 1, 2)
231+ 
232+def cpu_sparse_flash_attention_antiquant_gqa_s2split(
233+ query, key, value, sparse_indices, key_dequant_scale, value_dequant_scale,
234+ scale_value, sparse_block_size,
235+ actual_seq_lengths_query, actual_seq_lengths_kv, sparse_seq_lengths_kv,
236+ layout_query='BSND', layout_kv='PA_BSND', sparse_mode=3, block_table=None,
237+ attention_mode=0, quant_scale_repo_mode=0, tile_size=0, key_quant_mode=0,
238+ value_quant_mode=0, rope_head_dim=0, sparse_shard_size=1):
239+ query_type = query.dtype
240+ head_dim = query.shape[-1]
241+ 
242+ key = pa_to_bsnd(key, block_table, actual_seq_lengths_kv).to(torch.float32)
243+ key = (key * key_dequant_scale).to(query_type)
244+
245+ value = pa_to_bsnd(value, block_table, actual_seq_lengths_kv).to(torch.float32)
246+ value = (value * value_dequant_scale).to(query_type)
247+
248+ batch_size = actual_seq_lengths_query.shape[0]
249+ if layout_query == "TND":
250+ num_heads = query.shape[1]
251+ else:
252+ num_heads = query.shape[2]
253+ num_kv_heads = key.shape[2]
254+ g = num_heads // num_kv_heads
255+ 
256+ q_bnsd_tensor = torch.transpose(query, 1, 2)
257+ k_bnsd_tensor = torch.transpose(key, 1, 2)
258+ v_bnsd_tensor = torch.transpose(value, 1, 2)
259+ sparse_indices_tensor = torch.transpose(sparse_indices, 1, 2)
260+ out_shape_bnsd = list(q_bnsd_tensor.shape)
261+ y = torch.zeros(out_shape_bnsd, dtype=query_type)
262+ 
263+ for batch in range(batch_size):
264+ cur_acutal_seq_lengths_q = actual_seq_lengths_query[batch]
265+ if layout_query == "TND" and batch > 0:
266+ cur_acutal_seq_lengths_q = actual_seq_lengths_query[batch] - actual_seq_lengths_query[batch - 1]
267+ cur_actual_seq_lengths_kv = actual_seq_lengths_kv[batch]
268+ sparse_count = sparse_seq_lengths_kv[batch]
269+ for n2_idx in range(num_kv_heads):
270+ s1_shard_loops = (cur_acutal_seq_lengths_q + sparse_shard_size - 1) // sparse_shard_size
271+ s1_shard_tail = cur_acutal_seq_lengths_q - (s1_shard_loops - 1) * sparse_shard_size
272+ for s1_shard_idx in range(s1_shard_loops):
273+ s1_shard_size = s1_shard_tail if (s1_shard_idx == s1_shard_loops -1) else sparse_shard_size
274+ q_curr = q_bnsd_tensor[batch, n2_idx * g: (n2_idx + 1) * g, s1_shard_idx * sparse_shard_size : s1_shard_idx * sparse_shard_size + s1_shard_size, :]
275+ cur_sparse_indices = sparse_indices_tensor[batch, n2_idx, s1_shard_idx, :]
276+ k_sparse, v_sparse, s2_index = gather_kv(k_bnsd_tensor, v_bnsd_tensor, cur_sparse_indices, sparse_block_size,
277+ sparse_count, batch, n2_idx, cur_acutal_seq_lengths_q, cur_actual_seq_lengths_kv,
278+ sparse_mode, s1_shard_idx)
279+ s2_base = 512
280+ s2 = k_sparse.shape[0]
281+ loop = (s2 + s2_base - 1) // s2_base
282+ tail = s2 - (loop - 1) * s2_base
283+ rowSum = torch.from_numpy(np.zeros(shape=(s1_shard_size * g, 1), dtype=np.float32))
284+ rowMax = torch.from_numpy(np.full(shape=(s1_shard_size * g, 1), fill_value=-np.inf, dtype=np.float32))
285+ mm2_res = torch.from_numpy(np.zeros(shape=(s1_shard_size * g, out_shape_bnsd[-1]), dtype=np.float32))
286+ threshold_base = cur_actual_seq_lengths_kv - cur_acutal_seq_lengths_q + s1_shard_idx * sparse_shard_size + 1
287+ for i in range(loop):
288+ if i < loop - 1:
289+ seq_start = i * s2_base
290+ seq_end = (i + 1) * s2_base
291+ else:
292+ seq_start = i * s2_base
293+ seq_end = i * s2_base + tail
294+ k_cur = k_sparse[seq_start: seq_end, :]
295+ v_cur = v_sparse[seq_start: seq_end, :]
296+ s2_index_cur = s2_index[seq_start: seq_end]
297+ mm1_res = torch.matmul(q_curr.reshape(g * s1_shard_size, head_dim).to(torch.float32), k_cur.to(torch.float32).T)
298+ scale_res = mm1_res * scale_value
299+ if sparse_mode == 3:
300+ for s1_idx in range(s1_shard_size):
301+ mask_index = s2_index_cur >= threshold_base + s1_idx
302+ scale_res[s1_idx * g: (s1_idx + 1) * g, mask_index] = -1e12
303+ max_local,_ = scale_res.max(axis=-1, keepdims=True)
304+ replace_idx = torch.where(rowMax < max_local)
305+ rowMax_old = rowMax.clone()
306+ rowMax[replace_idx] = max_local[replace_idx]
307+ update_mul = torch.exp(rowMax_old - rowMax)
308+ 
309+ softmax_res = torch.exp(scale_res - rowMax)
310+ sum_local = softmax_res.sum(axis=-1, keepdims=True)
311+ rowSum = update_mul * rowSum + sum_local
312+ 
313+ mm2_res_local = torch.matmul(softmax_res.to(torch.float32), v_cur.to(torch.float32))
314+ mm2_res = update_mul * mm2_res + mm2_res_local
315+ res = mm2_res / rowSum
316+ res = res.reshape(g, s1_shard_size, head_dim)
317+ y[batch, n2_idx * g: (n2_idx + 1) * g, s1_shard_idx * sparse_shard_size : s1_shard_idx * sparse_shard_size + s1_shard_size, :] = res.to(query_type)
318+ return torch.transpose(y, 1, 2)
319+ 
320+def calculate_new_sparse_seq_kv(act_seq_kv, sparse_ratio, sparse_block_size):
321+ """
322+ 计算新的sparse_seq_kv值
323+ 1. 先计算 act_seq_kv * sparse_ratio
324+ 2. 向上取整
325+ 3. 调整到满足条件：sparse_seq_kv % 16 == act_seq_kv % 16
326+ 4. 确保 sparse_seq_kv <= act_seq_kv
327+ """
328+ new_sparse_seq = []
329+
330+ for act_val in act_seq_kv:
331+ # 基本计算：act_seq_kv * sparse_ratio 向上取整
332+ base_val = math.ceil(act_val * sparse_ratio)
333+
334+ # 计算余数
335+ act_remainder = act_val % sparse_block_size
336+
337+ # 调整base_val，使其模16的余数与act_val模16的余数相等
338+ current_remainder = base_val % sparse_block_size
339+
340+ if current_remainder != act_remainder:
341+ # 计算需要调整的值
342+ diff = (act_remainder - current_remainder) % sparse_block_size
343+ base_val += diff
344+
345+ # 确保sparse_seq_kv <= act_seq_kv
346+ # 如果base_val大于act_val，尝试减小base_val
347+ while base_val > act_val and base_val >= sparse_block_size:
348+ base_val -= sparse_block_size
349+
350+ # 再次检查余数条件，确保调整后仍然满足
351+ if base_val % sparse_block_size != act_remainder:
352+ # 如果调整后不满足余数条件，尝试向下调整
353+ # 找到不大于act_val且满足余数条件的最大值
354+ candidate = act_val
355+ while candidate > 0 and candidate % sparse_block_size != act_remainder:
356+ candidate -= 1
357+
358+ # 确保candidate >= base_val的最小值（向上取整后的值）
359+ min_val = math.ceil(act_val * sparse_ratio)
360+ if candidate >= min_val:
361+ base_val = candidate
362+ else:
363+ # 如果找不到满足条件的值，使用原始计算值
364+ # 但确保不大于act_val
365+ base_val = min(min_val, act_val)
366+
367+ # 最终确保base_val不大于act_val
368+ base_val = min(base_val, act_val)
369+
370+ # 确保base_val非负
371+ base_val = max(base_val, 0)
372+
373+ new_sparse_seq.append(base_val)
374+
375+ return new_sparse_seq
376+ 
377+def compute_sparse_seq_len(qk_len_tensor, is_context, sparsity=4):
378+ sparse_block_size = 16
379+ max_select_count = 128 * 1024
380+ fixed_tail_count = 32
381+
382+ if not is_context:
383+ act_tail_seq = qk_len_tensor % sparse_block_size
384+ fixed_seq = torch.where(
385+ act_tail_seq == 0,
386+ torch.tensor(fixed_tail_count * sparse_block_size, device=qk_len_tensor.device),
387+ (fixed_tail_count - 1) * sparse_block_size + act_tail_seq,
388+ )
389+ # print("fixed_seq = ", fixed_seq)
390+ select_seq = torch.clamp(qk_len_tensor - fixed_seq, min=0)
391+ # print("select_seq = ", select_seq)
392+ select_N_count = torch.ceil((select_seq / sparse_block_size) * round(1 / sparsity, 2))
393+ select_N_count = torch.min(select_N_count, torch.tensor(max_select_count, device=qk_len_tensor.device))
394+ # print("select_N_count = ", select_N_count)
395+ sparse_seq_len = select_N_count * sparse_block_size + torch.min(qk_len_tensor, fixed_seq)
396+ # print("sparse_seq_len = ", sparse_seq_len)
397+ sparse_seq_len = sparse_seq_len.to(torch.int32)
398+ else:
399+ sparse_seq_len = None
400+ return sparse_seq_len
401+ 
402+class TestCustomSFA(TestCase):
403+ def test_sfa_eager(self):
404+ torch_npu.npu.set_device(int(DEVICE_ID))
405+ query_type = torch.bfloat16
406+ scale_value = 0.041666666666666664
407+ sparse_block_size = 16
408+ 
409+ # 典型shape性能/功能用例，其余泛化用例见文件末尾，可能需要修改一些参数
410+ sparse_ratio = 0.25
411+ b = 21*4
412+ s1 = 4
413+ n1 = 20
414+ n2 = 2
415+ dn = 128
416+ tile_size = 128
417+ block_size = 512
418+ import random
419+ random_seed = 42
420+ unbalance_cache = 100
421+ random.seed(random_seed)
422+ s2_list = []
423+ sumseqlength = 0
424+ for _ in range(b-1):
425+ offset_percent = random.randint(-unbalance_cache, unbalance_cache)
426+ offset = int(s2 * (offset_percent / 100))
427+ s2curent = s2 + offset
428+ sumseqlength = sumseqlength + s2curent
429+ s2_list.append(s2curent)
430+ totalseq = b * s2
431+ s2_last = max(totalseq - sumseqlength, 0)
432+ s2_list.append(s2_last)
433+ 
434+ act_seq_q_list = [s1] * b
435+ act_seq_q = torch.tensor(act_seq_q_list).to(torch.int32)
436+ act_seq_kv = torch.tensor(s2_list).to(torch.int32)
437+ print(act_seq_q)
438+ print('SFA test case, act_seq_kv:', act_seq_kv)
439+ # sparse_seq_kv = torch.tensor(calculate_new_sparse_seq_kv(act_seq_kv, sparse_ratio, sparse_block_size)).to(torch.int32)
440+ sparse_seq_kv = compute_sparse_seq_len(
441+ act_seq_kv,
442+ is_context=False,
443+ sparsity=4
444+ )
445+ print('SFA test case, sparse_seq_kv:', sparse_seq_kv)
446+ s2 = max(s2_list)
447+ print('SFA test case, s2:', s2)
448+ 
449+ layout_query = 'BSND'
450+ layout_kv = 'PA_NZ'
451+ sparse_shard_size = 4
452+ key_quant_mode = 0
453+ value_quant_mode = 0
454+ attention_mode = 0
455+ quant_scale_repo_mode = 0
456+ sparse_mode = 3
457+ maxsparse_block_count=2072
458+ sparse_block_count = torch.ceil(((sparse_seq_kv) / (sparse_block_size)))
459+ max_block_num = math.ceil(s2 / block_size)
460+ block_num = max_block_num * b
461+
462+ # max_block_num = 257
463+ # block_num = 2048*5
464+ print("max_block_num = ", max_block_num)
465+ print("block_num = ", block_num)
466+ 
467+ query = torch.tensor(np.random.uniform(-10, 10, (b, s1, n1, dn))).to(query_type)
468+ if layout_kv == 'PA_BSND':
469+ key = torch.tensor(np.random.uniform(-100, 100, math.ceil(b * (s2 / block_size), block_size, n2, dn))).to(torch.int8)
470+ value = torch.tensor(np.random.uniform(-100, 100, math.ceil(b * (s2 / block_size), block_size, n2, dn))).to(torch.int8)
471+ elif layout_kv == 'PA_BNSD':
472+ key = torch.tensor(np.random.uniform(-100, 100, (b * math.ceil(s2 / block_size), n2, block_size, dn))).to(torch.int8)
473+ value = torch.tensor(np.random.uniform(-100, 100, (b * math.ceil(s2 / block_size), n2, block_size, dn))).to(torch.int8)
474+ elif layout_kv == 'PA_NZ':
475+ key = torch.tensor(np.random.uniform(-100, 100, (block_num, n2, (dn//32), block_size, 32))).to(torch.int8)
476+ value = torch.tensor(np.random.uniform(-100, 100, (block_num, n2, (dn//32), block_size, 32))).to(torch.int8)
477+ elif layout_kv == 'BSND':
478+ key = torch.tensor(np.random.uniform(-100, 100, (b, s2, n2, dn))).to(torch.int8)
479+ value = torch.tensor(np.random.uniform(-100, 100, (b, s2, n2, dn))).to(torch.int8)
480+ key_antiquant_scale = torch.tensor(np.random.uniform(-100, 100, (n2, dn))).to(torch.float32)
481+ value_antiquant_scale = torch.tensor(np.random.uniform(-100, 100, (n2, dn))).to(torch.float32)
482+ sparse_indices = torch.full((b, n2, s1 // sparse_shard_size, maxsparse_block_count), -1).to(torch.int32)
483+ # generate sparse_indices
484+ for b_i in range(b):
485+ for n_i in range(n2):
486+ for s_i in range(s1 // sparse_shard_size):
487+ if sparse_mode == 0:
488+ threshold = act_seq_kv[b_i]
489+ elif sparse_mode == 3:
490+ threshold = act_seq_kv[b_i] - act_seq_q[b_i] + s_i * sparse_shard_size + sparse_shard_size
491+ if threshold <= 0:
492+ sparse_indices[b_i, n_i, s_i, :] = torch.tensor([-1] * maxsparse_block_count).to(torch.int32)
493+ continue
494+ valid_blocks_max = math.ceil(max(0, threshold) / sparse_block_size)
495+ # 处理边界情况：如果有效块数为0，跳过
496+ if valid_blocks_max == 0:
497+ continue
498+ valid_blocks_topk = min(valid_blocks_max, sparse_block_count[b_i])
499+ valid_blocks_topk_int = int(valid_blocks_topk)
500+ # 情况1: 只能选1个块
501+ if valid_blocks_topk_int <= 1:
502+ # 只选最后一个块
503+ sparse_indices[b_i, n_i, s_i, 0:1] = valid_blocks_max - 1
504+ # 情况2: 可以选2个或更多块
505+ else:
506+ # 计算要随机选择的块数
507+ random_blocks = max(0, valid_blocks_topk_int - 2)
508+ if random_blocks > 0 and valid_blocks_max - 2 > 0:
509+ # 从除了最后两个块之外的块中随机选择
510+ block_indices = torch.randperm(valid_blocks_max - 2).to(torch.int32)
511+ # 填充随机选择的块
512+ sparse_indices[b_i, n_i, s_i, :random_blocks] = block_indices[0:random_blocks]
513+ # 固定选择最后两个块
514+ if valid_blocks_max >= 2:
515+ # 先放倒数第二个块
516+ sparse_indices[b_i, n_i, s_i, random_blocks] = valid_blocks_max - 2
517+ # 再放最后一个块
518+ sparse_indices[b_i, n_i, s_i, random_blocks + 1] = valid_blocks_max - 1
519+ else:
520+ # 如果只有1个块，只选最后一个块
521+ sparse_indices[b_i, n_i, s_i, random_blocks] = valid_blocks_max - 1
522+ sparse_indices = torch.transpose(sparse_indices, 1, 2)
523+ print("sparse_indices = ", sparse_indices)
524+ # generate block_table
525+ if layout_kv == 'PA_BSND' or layout_kv == 'PA_BNSD' or layout_kv == 'PA_NZ':
526+ block_table = torch.full((b, max_block_num), -1).to(torch.int32)
527+ block_numPerBlock = []
528+ block_num_min = 0
529+ for actual_seq in act_seq_kv:
530+ block_numPerBlock.append(math.ceil(actual_seq / block_size))
531+ block_num_min += math.ceil(actual_seq / block_size)
532+ if block_num_min > block_num:
533+ raise RuntimeError(f"block_num{block_num} is too small, please increase block_num to at least {block_num_min}")
534+ block_idx_list = torch.randperm(block_num).to(torch.int32)
535+ block_idx = 0
536+ block_table_batch_idx = 0
537+ for idx in block_numPerBlock:
538+ for j in range(idx):
539+ block_table[block_table_batch_idx][j] = block_idx_list[block_idx]
540+ block_idx += 1
541+ block_table_batch_idx += 1
542+ block_table = block_table.to("npu:%s" % DEVICE_ID)
543+ else:
544+ block_table = None
545+ 
546+ # compare result
547+ cpu_out = cpu_sparse_flash_attention_antiquant_gqa(
548+ query, key, value, sparse_indices,
549+ key_dequant_scale=key_antiquant_scale, value_dequant_scale=value_antiquant_scale,
550+ scale_value=scale_value, sparse_block_size=sparse_block_size,
551+ actual_seq_lengths_query=act_seq_q, actual_seq_lengths_kv=act_seq_kv, sparse_seq_lengths_kv = sparse_seq_kv,
552+ layout_query=layout_query, layout_kv=layout_kv, sparse_mode=sparse_mode, block_table=block_table,
553+ attention_mode=attention_mode, quant_scale_repo_mode=quant_scale_repo_mode, tile_size=tile_size, key_quant_mode=key_quant_mode,
554+ value_quant_mode=value_quant_mode, rope_head_dim=64, sparse_shard_size=sparse_shard_size)
555+ 
556+ query = query.to("npu:%s" % DEVICE_ID)
557+ key = key.to("npu:%s" % DEVICE_ID)
558+ value = value.to("npu:%s" % DEVICE_ID)
559+ sparse_indices = sparse_indices.to("npu:%s" % DEVICE_ID)
560+ key_antiquant_scale = key_antiquant_scale.to("npu:%s" % DEVICE_ID)
561+ value_antiquant_scale = value_antiquant_scale.to("npu:%s" % DEVICE_ID)
562+ act_seq_q = act_seq_q.to("npu:%s" % DEVICE_ID)
563+ act_seq_kv = act_seq_kv.to("npu:%s" % DEVICE_ID)
564+ sparse_seq_kv = sparse_seq_kv.to("npu:%s" % DEVICE_ID)
565+ 
566+ 
567+ # print(f'======================== PTA eager BEGIN ========================')
568+ sals_sfa_meta = torch_npu.npu_sparse_flash_attention_antiquant_metadata(
569+ b, s1, n1, s2, n2, dn,
570+ 10,
571+ sparse_block_size,
572+ actual_seq_lengths_query=act_seq_q,
573+ actual_seq_lengths_kv=act_seq_kv,
574+ sparse_seq_lengths_kv=sparse_seq_kv,
575+ sparse_mode=sparse_mode,
576+ attention_mode=attention_mode,
577+ rope_head_dim=64,
578+ sparse_shared_size=sparse_shard_size)
579+ npu_out = torch_npu.npu_sparse_flash_attention_antiquant(
580+ query, key, value, sparse_indices,
581+ key_dequant_scale=key_antiquant_scale, value_dequant_scale=value_antiquant_scale,
582+ scale_value=scale_value, sparse_block_size=sparse_block_size,metadata=sals_sfa_meta,
583+ actual_seq_lengths_query=act_seq_q, actual_seq_lengths_kv=act_seq_kv, sparse_seq_lengths_kv = sparse_seq_kv,
584+ layout_query=layout_query, layout_kv=layout_kv, sparse_mode=sparse_mode, block_table=block_table,
585+ attention_mode=attention_mode, quant_scale_repo_mode=quant_scale_repo_mode, tile_size=tile_size, key_quant_mode=key_quant_mode,
586+ value_quant_mode=value_quant_mode, rope_head_dim=64, sparse_shard_size=sparse_shard_size)
587+ print("npu_out = ",npu_out)
588+ 
589+ npu_out = npu_out.cpu().to(torch.float32).numpy()
590+ cpu_out = cpu_out.cpu().to(torch.float32).numpy()
591+ 
592+ res = np.isclose(npu_out, cpu_out, rtol=0.005, atol=0.0001, equal_nan=False)
593+ diff_mask = ~res
594+ true_ratio = np.mean(res)
595+ print("npu output:\n", npu_out, npu_out.shape)
596+ print("cpu output:\n", cpu_out, cpu_out.shape)
597+ print("correct ratio of cpu vs npu is:", true_ratio * 100, "%")
598+ total_elements = npu_out.size
599+ diff_elements = np.sum(diff_mask)
600+ diff_perventage = (diff_elements / total_elements) * 100
601+ 
602+ print(f"总元素数: {total_elements}")
603+ print(f"不同元素数: {diff_elements}")
604+ print(f"不同比例: {diff_perventage:.2f}%")
605+
606+ #输出差别详情
607+ if diff_elements > 0:
608+ diff_indices = np.where(diff_mask)
609+ print("\n差别详情: ")
610+ for idx in range(len(diff_indices[0])):
611+ pos = tuple(dim[idx] for dim in diff_indices)
612+ a_val = npu_out[pos]
613+ b_val = cpu_out[pos]
614+ diff_val = abs(a_val - b_val)
615+ diff_val_relative = diff_val / abs(b_val)
616+ if diff_val_relative > 0.5:
617+ print(f"Warning! 位置 {pos}: npu_out={a_val:.6f}, cpu_out={b_val:.6f}, 差值={diff_val:.6f}, diff_value_relative={diff_val_relative:.6f}")
618+ else:
619+ print(f"位置 {pos}: npu_out={a_val:.6f}, cpu_out={b_val:.6f}, 差值={diff_val:.6f}, diff_value_relative={diff_val_relative:.6f}")
620+ self.assertTrue(true_ratio > 0.99, "precision compare fail")
621+ print(f'======================== PTA eager FINISH ========================')
622+ 
623+ print(f'======================== PTA eager Graph BEGIN ========================')
624+ print("sparse_indices = ", sparse_indices)
625+ print("block_table = ", block_table)
626+ npu_mode = SFAANetwork().to("npu:%s" % DEVICE_ID)
627+ from torchair.configs.compiler_config import CompilerConfig
628+ config = CompilerConfig()
629+ npu_backend = torchair.get_npu_backend(compiler_config=config)
630+ torch._dynamo.reset()
631+ npu_mode = torch.compile(npu_mode, fullgraph=False, backend=npu_backend, dynamic=False)
632+ npu_out0 = npu_mode(
633+ b, s1, n1, s2, n2, dn,
634+ query, key, value, sparse_indices,
635+ key_dequant_scale=key_antiquant_scale, value_dequant_scale=value_antiquant_scale,
636+ scale_value=scale_value, sparse_block_size=sparse_block_size,
637+ actual_seq_lengths_query=act_seq_q, actual_seq_lengths_kv=act_seq_kv, sparse_seq_lengths_kv = sparse_seq_kv,
638+ layout_query=layout_query, layout_kv=layout_kv, sparse_mode=sparse_mode, block_table=block_table,
639+ attention_mode=attention_mode, quant_scale_repo_mode=quant_scale_repo_mode, tile_size=tile_size, key_quant_mode=key_quant_mode,
640+ value_quant_mode=value_quant_mode, rope_head_dim=64, sparse_shard_size=sparse_shard_size)
641+ npu_out0 = npu_out0.cpu().to(torch.float32).numpy()
642+ cpu_out = cpu_out.to(torch.float32).numpy()
643+ res = np.isclose(npu_out0, cpu_out, rtol=0.005, atol=0.0001, equal_nan=False)
644+ true_ratio = np.mean(res)
645+ if true_ratio < 0.99:
646+ print("npu output:\n", npu_out0, npu_out0.shape)
647+ print("cpu output:\n", cpu_out, cpu_out.shape)
648+ print("correct ratio of cpu vs npu is:", true_ratio * 100, "%")
649+ self.assertTrue(true_ratio > 0.99, "precision compare fail")
650+ print(f'======================== PTA eager Graph FINISH ========================')
651+ 
652+ 
653+if __name__ == "__main__":
654+ run_tests()
Aexperimental/attention/sparse_flash_attention_antiquant/op_host/sparse_flash_attention_antiquant_def.cpp+107-0
@@ -0,0 +1,107 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+
11+/*!
12+ * \file sparse_flash_attention_antiquant_def.cpp
13+ * \brief
14+ */
15+ 
16+#include "register/op_def_registry.h"
17+ 
18+namespace ops {
19+class SparseFlashAttentionAntiquant : public OpDef {
20+public:
21+ explicit SparseFlashAttentionAntiquant(const char *name) : OpDef(name)
22+ {
23+ this->Input("query")
24+ .ParamType(REQUIRED)
25+ .DataType({ge::DT_FLOAT16, ge::DT_BF16})
26+ .Format({ge::FORMAT_ND, ge::FORMAT_ND})
27+ .AutoContiguous();
28+ this->Input("key")
29+ .ParamType(REQUIRED)
30+ .DataType({ge::DT_INT8, ge::DT_INT8})
31+ .Format({ge::FORMAT_ND, ge::FORMAT_ND})
32+ .AutoContiguous();
33+ this->Input("value")
34+ .ParamType(REQUIRED)
35+ .DataType({ge::DT_INT8, ge::DT_INT8})
36+ .Format({ge::FORMAT_ND, ge::FORMAT_ND})
37+ .AutoContiguous();
38+ this->Input("sparse_indices")
39+ .ParamType(REQUIRED)
40+ .DataType({ge::DT_INT32, ge::DT_INT32})
41+ .Format({ge::FORMAT_ND, ge::FORMAT_ND})
42+ .AutoContiguous();
43+ this->Input("key_dequant_scale")
44+ .ParamType(OPTIONAL)
45+ .DataType({ge::DT_FLOAT, ge::DT_FLOAT})
46+ .Format({ge::FORMAT_ND, ge::FORMAT_ND})
47+ .AutoContiguous();
48+ this->Input("value_dequant_scale")
49+ .ParamType(OPTIONAL)
50+ .DataType({ge::DT_FLOAT, ge::DT_FLOAT})
51+ .Format({ge::FORMAT_ND, ge::FORMAT_ND})
52+ .AutoContiguous();
53+ this->Input("block_table")
54+ .ParamType(OPTIONAL)
55+ .DataType({ge::DT_INT32, ge::DT_INT32})
56+ .Format({ge::FORMAT_ND, ge::FORMAT_ND})
57+ .AutoContiguous();
58+ this->Input("actual_seq_lengths_query")
59+ .ParamType(OPTIONAL)
60+ .DataType({ge::DT_INT32, ge::DT_INT32})
61+ .Format({ge::FORMAT_ND, ge::FORMAT_ND})
62+ .AutoContiguous();
63+ this->Input("actual_seq_lengths_kv")
64+ .ParamType(OPTIONAL)
65+ .DataType({ge::DT_INT32, ge::DT_INT32})
66+ .Format({ge::FORMAT_ND, ge::FORMAT_ND})
67+ .AutoContiguous();
68+ this->Input("sparse_seq_lengths_kv")
69+ .ParamType(OPTIONAL)
70+ .DataType({ge::DT_INT32, ge::DT_INT32})
71+ .Format({ge::FORMAT_ND, ge::FORMAT_ND})
72+ .AutoContiguous();
73+ this->Input("metadata")
74+ .ParamType(OPTIONAL)
75+ .DataType({ge::DT_INT32, ge::DT_INT32})
76+ .Format({ge::FORMAT_ND, ge::FORMAT_ND})
77+ .AutoContiguous();
78+ this->Output("attention_out")
79+ .ParamType(REQUIRED)
80+ .DataType({ge::DT_FLOAT16, ge::DT_BF16})
81+ .Format({ge::FORMAT_ND, ge::FORMAT_ND});
82+ this->Attr("scale_value").AttrType(REQUIRED).Float(1.0);
83+ this->Attr("sparse_block_size").AttrType(REQUIRED).Int(1);
84+ this->Attr("key_quant_mode").AttrType(REQUIRED).Int(1);
85+ this->Attr("value_quant_mode").AttrType(REQUIRED).Int(1);
86+ this->Attr("layout_query").AttrType(OPTIONAL).String("BSND");
87+ this->Attr("layout_kv").AttrType(OPTIONAL).String("PA_NZ");
88+ this->Attr("sparse_mode").AttrType(OPTIONAL).Int(3); // 3:默认值，只计算下三角
89+ this->Attr("attention_mode").AttrType(OPTIONAL).Int(0);
90+ this->Attr("quant_scale_repo_mode").AttrType(OPTIONAL).Int(0);
91+ this->Attr("tile_size").AttrType(OPTIONAL).Int(0);
92+ this->Attr("rope_head_dim").AttrType(OPTIONAL).Int(0);
93+ this->Attr("sparse_shard_size").AttrType(OPTIONAL).Int(1);
94+ OpAICoreConfig aicore_config;
95+ aicore_config.DynamicCompileStaticFlag(true)
96+ .DynamicFormatFlag(true)
97+ .DynamicRankSupportFlag(true)
98+ .DynamicShapeSupportFlag(true)
99+ .NeedCheckSupportFlag(false)
100+ .PrecisionReduceFlag(true)
101+ .ExtendCfgInfo("aclnnSupport.value", "support_aclnn");
102+ this->AICore().AddConfig("ascend910b", aicore_config);
103+ this->AICore().AddConfig("ascend910_93", aicore_config);
104+ }
105+};
106+OP_ADD(SparseFlashAttentionAntiquant);
107+} // namespace ops
Aexperimental/attention/sparse_flash_attention_antiquant/op_host/sparse_flash_attention_antiquant_proto.cpp+80-0
@@ -0,0 +1,80 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file sparse_flash_attention_antiquant_proto.cpp
13+ * \brief
14+ */
15+ 
16+#include <graph/utils/type_utils.h>
17+#include <register/op_impl_registry.h>
18+#include "error/ops_error.h"
19+ 
20+using namespace ge;
21+ 
22+namespace ops {
23+constexpr size_t QUERY_INPUT_INDEX = 0;
24+constexpr uint32_t LAYOUT_QUERY_ATTR_INDEX = 4;
25+constexpr uint32_t ROPE_HEAD_DIM_ATTR_INDEX = 10;
26+constexpr uint32_t ATTENTION_MODE_ATTR_INDEX = 7;
27+ 
28+ge::graphStatus InferShapeSparseFlashAttentionAntiquant(gert::InferShapeContext *context)
29+{
30+ OPS_ERR_IF(context == nullptr, OPS_LOG_E("SparseFlashAttentionAntiquant", "InferShapeContext is nullptr"),
31+ return ge::GRAPH_FAILED);
32+ const gert::Shape *queryShape = context->GetInputShape(QUERY_INPUT_INDEX);
33+ OPS_LOG_E_IF_NULL(context, queryShape, return ge::GRAPH_FAILED);
34+ gert::Shape *attentionOutShape = context->GetOutputShape(0);
35+ OPS_LOG_E_IF_NULL(context, attentionOutShape, return ge::GRAPH_FAILED);
36+ auto attrs = context->GetAttrs();
37+ const char *inputLayoutQueryPtr = attrs->GetAttrPointer<char>(LAYOUT_QUERY_ATTR_INDEX);
38+ OPS_LOG_E_IF_NULL(context, inputLayoutQueryPtr, return ge::GRAPH_FAILED);
39+ std::string inputLayoutQueryPtrStr = std::string(inputLayoutQueryPtr);
40+ 
41+ const int64_t *attentionModePtr = attrs->GetAttrPointer<int64_t>(ATTENTION_MODE_ATTR_INDEX);
42+ OPS_LOG_E_IF_NULL(context, attentionModePtr, return ge::GRAPH_FAILED);
43+ const int64_t attentionMode = *attentionModePtr;
44+ 
45+ const int64_t *ropeHeadDimPtr = attrs->GetAttrPointer<int64_t>(ROPE_HEAD_DIM_ATTR_INDEX);
46+ int64_t ropeHeadDim = 0;
47+ if (attentionMode != 0) {
48+ OPS_LOG_E_IF_NULL(context, ropeHeadDimPtr, return ge::GRAPH_FAILED);
49+ ropeHeadDim = *ropeHeadDimPtr;
50+ }
51+ attentionOutShape->SetDimNum(queryShape->GetDimNum());
52+ if (inputLayoutQueryPtrStr == "BSND") {
53+ attentionOutShape->SetDim(0, queryShape->GetDim(0));
54+ attentionOutShape->SetDim(1, queryShape->GetDim(1));
55+ attentionOutShape->SetDim(2, queryShape->GetDim(2)); // 2:dim2
56+ int64_t outHeadDim = (attentionMode == 0)? queryShape->GetDim(3) : queryShape->GetDim(3) - ropeHeadDim;
57+ attentionOutShape->SetDim(3, outHeadDim); // 3:dim3
58+ } else { // TND
59+ attentionOutShape->SetDim(0, queryShape->GetDim(0));
60+ attentionOutShape->SetDim(1, queryShape->GetDim(1));
61+ int64_t outHeadDim = (attentionMode == 0)? queryShape->GetDim(2) : queryShape->GetDim(2) - ropeHeadDim;
62+ attentionOutShape->SetDim(2, outHeadDim); // 2:dim2
63+ }
64+ return GRAPH_SUCCESS;
65+}
66+ 
67+ge::graphStatus InferDataTypeSparseFlashAttentionAntiquant(gert::InferDataTypeContext *context)
68+{
69+ OPS_ERR_IF(context == nullptr, OPS_LOG_E("SparseFlashAttentionAntiquant", "InferShapeContext is nullptr"),
70+ return ge::GRAPH_FAILED);
71+ const auto inputDataType = context->GetInputDataType(QUERY_INPUT_INDEX);
72+ context->SetOutputDataType(0, inputDataType);
73+ return ge::GRAPH_SUCCESS;
74+}
75+ 
76+IMPL_OP(SparseFlashAttentionAntiquant)
77+ .InferShape(InferShapeSparseFlashAttentionAntiquant)
78+ .InferDataType(InferDataTypeSparseFlashAttentionAntiquant);
79+} // namespace ops
80+
Aexperimental/attention/sparse_flash_attention_antiquant/op_host/sparse_flash_attention_antiquant_tiling.cpp+2112-0
@@ -0,0 +1,2112 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file sparse_flash_attention_antiquant_tiling.cpp
13+ * \brief
14+ */
15+ 
16+#include <map>
17+#include <vector>
18+#include <numeric>
19+#include <algorithm>
20+#include <graph/utils/type_utils.h>
21+#include "error/ops_error.h"
22+#include "register/op_def_registry.h"
23+#include "../op_kernel/sparse_flash_attention_antiquant_template_tiling_key.h"
24+#include "sparse_flash_attention_antiquant_tiling.h"
25+#include "split_core.h"
26+#include "split_core_BN.h"
27+ 
28+using std::map;
29+using std::string;
30+using std::pair;
31+ 
32+using namespace ge;
33+using namespace AscendC;
34+using namespace optiling::sfaa;
35+ 
36+namespace optiling {
37+ 
38+constexpr uint32_t PRE_LOAD_NUM = 2;
39+constexpr uint32_t BLOCK_TABLE_ELEM_BYTE = 4;
40+constexpr int32_t SPARSE_MODE_BAND = 4;
41+ 
42+static const std::string QUERY_NAME = "query";
43+static const std::string KEY_NAME = "key";
44+static const std::string VALUE_NAME = "value";
45+static const std::string SPARSE_INDICES_NAME = "sparse_indices";
46+static const std::string BLOCK_TABLE_NAME = "block_table";
47+static const std::string ATTEN_OUT_NAME = "attention_out";
48+ 
49+const std::map<std::string, std::vector<ge::DataType>> DTYPE_SUPPORT_MAP = {
50+ {QUERY_NAME, {ge::DT_FLOAT16, ge::DT_BF16}},
51+ {KEY_NAME, {ge::DT_INT8, ge::DT_INT8}},
52+ {VALUE_NAME, {ge::DT_INT8, ge::DT_INT8}},
53+ {ATTEN_OUT_NAME, {ge::DT_FLOAT16, ge::DT_BF16}},
54+ {SPARSE_INDICES_NAME, {ge::DT_INT32}}
55+};
56+ 
57+const std::map<std::string, std::vector<SFAALayout>> LAYOUT_SUPPORT_MAP = {
58+ {QUERY_NAME, {SFAALayout::BSND, SFAALayout::TND}},
59+ {KEY_NAME, {SFAALayout::BSND, SFAALayout::TND, SFAALayout::PA_BSND, SFAALayout::PA_BNSD, SFAALayout::PA_NZ}},
60+ {VALUE_NAME, {SFAALayout::BSND, SFAALayout::TND, SFAALayout::PA_BSND, SFAALayout::PA_BNSD, SFAALayout::PA_NZ}},
61+ {ATTEN_OUT_NAME, {SFAALayout::BSND, SFAALayout::TND}},
62+};
63+ 
64+const std::map<ge::DataType, std::string> DATATYPE_TO_STRING_MAP = {
65+ {ge::DT_UNDEFINED, "DT_UNDEFINED"}, // Used to indicate a DataType field has not been set.
66+ {ge::DT_FLOAT, "DT_FLOAT"}, // float type
67+ {ge::DT_FLOAT16, "DT_FLOAT16"}, // fp16 type
68+ {ge::DT_INT8, "DT_INT8"}, // int8 type
69+ {ge::DT_INT16, "DT_INT16"}, // int16 type
70+ {ge::DT_UINT16, "DT_UINT16"}, // uint16 type
71+ {ge::DT_UINT8, "DT_UINT8"}, // uint8 type
72+ {ge::DT_INT32, "DT_INT32"}, // uint32 type
73+ {ge::DT_INT64, "DT_INT64"}, // int64 type
74+ {ge::DT_UINT32, "DT_UINT32"}, // unsigned int32
75+ {ge::DT_UINT64, "DT_UINT64"}, // unsigned int64
76+ {ge::DT_BOOL, "DT_BOOL"}, // bool type
77+ {ge::DT_DOUBLE, "DT_DOUBLE"}, // double type
78+ {ge::DT_DUAL, "DT_DUAL"}, // dual output type
79+ {ge::DT_DUAL_SUB_INT8, "DT_DUAL_SUB_INT8"}, // dual output int8 type
80+ {ge::DT_DUAL_SUB_UINT8, "DT_DUAL_SUB_UINT8"}, // dual output uint8 type
81+ {ge::DT_COMPLEX32, "DT_COMPLEX32"}, // complex32 type
82+ {ge::DT_COMPLEX64, "DT_COMPLEX64"}, // complex64 type
83+ {ge::DT_COMPLEX128, "DT_COMPLEX128"}, // complex128 type
84+ {ge::DT_QINT8, "DT_QINT8"}, // qint8 type
85+ {ge::DT_QINT16, "DT_QINT16"}, // qint16 type
86+ {ge::DT_QINT32, "DT_QINT32"}, // qint32 type
87+ {ge::DT_QUINT8, "DT_QUINT8"}, // quint8 type
88+ {ge::DT_QUINT16, "DT_QUINT16"}, // quint16 type
89+ {ge::DT_RESOURCE, "DT_RESOURCE"}, // resource type
90+ {ge::DT_STRING_REF, "DT_STRING_REF"}, // string ref type
91+ {ge::DT_STRING, "DT_STRING"}, // string type
92+ {ge::DT_VARIANT, "DT_VARIANT"}, // dt_variant type
93+ {ge::DT_BF16, "DT_BFLOAT16"}, // dt_bfloat16 type
94+ {ge::DT_INT4, "DT_INT4"}, // dt_variant type
95+ {ge::DT_UINT1, "DT_UINT1"}, // dt_variant type
96+ {ge::DT_INT2, "DT_INT2"}, // dt_variant type
97+ {ge::DT_UINT2, "DT_UINT2"} // dt_variant type
98+};
99+ 
100+struct SparseFlashAttentionAntiquantCompileInfo {
101+ int64_t coreNum;
102+};
103+ 
104+static const std::map<SFAALayout, std::vector<SFAAAxis>> SFAA_LAYOUT_AXIS_MAP = {
105+ {SFAALayout::BSND, {SFAAAxis::B, SFAAAxis::S, SFAAAxis::N, SFAAAxis::D}},
106+ {SFAALayout::TND, {SFAAAxis::T, SFAAAxis::N, SFAAAxis::D}},
107+ {SFAALayout::PA_BSND, {SFAAAxis::Bn, SFAAAxis::Bs, SFAAAxis::N, SFAAAxis::D}},
108+ {SFAALayout::PA_BNSD, {SFAAAxis::Bn, SFAAAxis::N, SFAAAxis::Bs, SFAAAxis::D}},
109+ {SFAALayout::PA_NZ, {SFAAAxis::Bn, SFAAAxis::N, SFAAAxis::Dn, SFAAAxis::Bs, SFAAAxis::Ds}},
110+};
111+ 
112+static const std::map<SFAALayout, size_t> SFAA_LAYOUT_DIM_MAP = {
113+ {SFAALayout::BSND, DIM_NUM_FOUR},
114+ {SFAALayout::TND, DIM_NUM_THREE},
115+ {SFAALayout::PA_BSND, DIM_NUM_FOUR},
116+ {SFAALayout::PA_BNSD, DIM_NUM_FOUR},
117+ {SFAALayout::PA_NZ, DIM_NUM_FIVE},
118+};
119+ 
120+static std::string GetShapeStr(gert::Shape shape)
121+{
122+ std::ostringstream oss;
123+ oss << "[";
124+ if (shape.GetDimNum() > 0) {
125+ for (size_t i = 0; i < shape.GetDimNum() - 1; ++i) {
126+ oss << shape.GetDim(i) << ", ";
127+ }
128+ oss << shape.GetDim(shape.GetDimNum() - 1);
129+ }
130+ oss << "]";
131+ return oss.str();
132+}
133+ 
134+static std::string SFAADataTypeToSerialString(ge::DataType type)
135+{
136+ const auto it = DATATYPE_TO_STRING_MAP.find(type);
137+ if (it != DATATYPE_TO_STRING_MAP.end()) {
138+ return it->second;
139+ } else {
140+ OPS_LOG_E("SparseFlashAttention", "datatype %d not support", type);
141+ return "UNDEFINED";
142+ }
143+}
144+ 
145+string SFAATensorDesc2String(const gert::StorageShape *shape, const gert::CompileTimeTensorDesc *tensor)
146+{
147+ if (shape == nullptr || tensor == nullptr) {
148+ return "nil ";
149+ }
150+ 
151+ std::ostringstream oss;
152+ oss << "(dtype: " << ge::TypeUtils::DataTypeToAscendString(tensor->GetDataType()).GetString() << "),";
153+ oss << "(shape:" << SFAAShape2String(shape->GetStorageShape()) << "),";
154+ oss << "(ori_shape:" << SFAAShape2String(shape->GetOriginShape()) << "),";
155+ oss << "(format: "
156+ << ge::TypeUtils::FormatToAscendString(
157+ static_cast<ge::Format>(ge::GetPrimaryFormat(tensor->GetStorageFormat())))
158+ .GetString()
159+ << "),";
160+ oss << "(ori_format: " << ge::TypeUtils::FormatToAscendString(tensor->GetOriginFormat()).GetString() << ") ";
161+ 
162+ return oss.str();
163+}
164+ 
165+string SFAADebugTilingContext(const gert::TilingContext *context)
166+{
167+ std::ostringstream oss;
168+ for (size_t i = 0; i < context->GetComputeNodeInfo()->GetInputsNum(); ++i) {
169+ oss << "input" << i << ": ";
170+ oss << SFAATensorDesc2String(context->GetInputShape(i), context->GetInputDesc(i));
171+ }
172+ 
173+ for (size_t i = 0; i < context->GetComputeNodeInfo()->GetOutputsNum(); ++i) {
174+ oss << "output" << i << ": ";
175+ oss << SFAATensorDesc2String(context->GetOutputShape(i), context->GetOutputDesc(i));
176+ }
177+ return oss.str();
178+}
179+ 
180+std::string SFAALayoutToSerialString(SFAALayout layout)
181+{
182+ switch (layout) {
183+ case SFAALayout::BSND: return "BSND";
184+ case SFAALayout::TND: return "TND";
185+ case SFAALayout::PA_BSND: return "PA_BSND";
186+ case SFAALayout::PA_BNSD: return "PA_BNSD";
187+ default: return "UNKNOWN";
188+ }
189+}
190+ 
191+ge::graphStatus SFAAMlaTiling::SetBlockDim(uint32_t blockDim)
192+{
193+ context_->SetBlockDim(blockDim);
194+ return ge::GRAPH_SUCCESS;
195+}
196+ 
197+ge::graphStatus SFAAMlaTiling::SetTilingKey(uint64_t tilingKey)
198+{
199+ context_->SetTilingKey(tilingKey);
200+ return ge::GRAPH_SUCCESS;
201+}
202+ 
203+ge::graphStatus SFAAMlaTiling::SetWorkspaceSize(uint64_t workspaceSize)
204+{
205+ OPS_ERR_IF(context_->GetWorkspaceSizes(1) == nullptr,
206+ OPS_REPORT_VECTOR_INNER_ERR(context_->GetNodeName(), "workSpaceSize got from ge is nullptr"),
207+ return ge::GRAPH_FAILED);
208+ size_t *workSpaces = context_->GetWorkspaceSizes(1);
209+ workSpaces[0] = workspaceSize;
210+ return ge::GRAPH_SUCCESS;
211+}
212+ 
213+ge::graphStatus SFAAMlaTiling::SetTilingData(TilingDef &tilingData)
214+{
215+ OPS_ERR_IF(context_->GetRawTilingData() == nullptr,
216+ OPS_REPORT_VECTOR_INNER_ERR(context_->GetNodeName(), "RawTilingData got from GE context is nullptr."),
217+ return ge::GRAPH_FAILED);
218+ 
219+ tilingData.SaveToBuffer(context_->GetRawTilingData()->GetData(), context_->GetRawTilingData()->GetCapacity());
220+ context_->GetRawTilingData()->SetDataSize(tilingData.GetDataSize());
221+ 
222+ return ge::GRAPH_SUCCESS;
223+}
224+ 
225+ge::graphStatus SFAAMlaTiling::GetPlatformInfo()
226+{
227+ OPS_ERR_IF(sfaaInfo_->platformInfo == nullptr,
228+ OPS_REPORT_VECTOR_INNER_ERR(sfaaInfo_->opName, "GetPlatformInfo is nullptr."), return ge::GRAPH_FAILED);
229+ 
230+ auto ascendcPlatform = platform_ascendc::PlatformAscendC(sfaaInfo_->platformInfo);
231+ libapiSize_ = ascendcPlatform.GetLibApiWorkSpaceSize();
232+ aivNum_ = ascendcPlatform.GetCoreNumAiv();
233+ aicNum_ = ascendcPlatform.GetCoreNumAic();
234+ 
235+ OPS_ERR_IF(aicNum_ == 0 || aivNum_ == 0,
236+ OPS_REPORT_VECTOR_INNER_ERR(sfaaInfo_->opName, "num of core obtained is 0."), return GRAPH_FAILED);
237+ OPS_ERR_IF(aicNum_ != 24,
238+ OPS_REPORT_VECTOR_INNER_ERR(sfaaInfo_->opName, "num of core only supported 24."), return GRAPH_FAILED);
239+ 
240+ return ge::GRAPH_SUCCESS;
241+}
242+ 
243+void SFAAMlaTiling::GenTilingKey()
244+{
245+ uint32_t inputQType = static_cast<uint32_t>(sfaaInfo_->inputQType);
246+ uint32_t inputKvType = static_cast<uint32_t>(sfaaInfo_->inputKvType);
247+ uint32_t outputType = static_cast<uint32_t>(sfaaInfo_->outputType);
248+ uint32_t layoutQuery = static_cast<uint32_t>(sfaaInfo_->qLayout);
249+ uint32_t layoutKV = static_cast<uint32_t>(sfaaInfo_->kvLayout);
250+ uint32_t attentionMode = static_cast<uint32_t>(sfaaInfo_->attentionMode);
251+ // Tiling key含义修改：是否有meta传递
252+ if (context_->GetOptionalInputDesc(METADATA_INDEX) == nullptr) {
253+ tilingKey_ = GET_TPL_TILING_KEY(attentionMode, 0U, layoutQuery, layoutKV, perfMode_ == SFAAPerfMode::V_TEMPLATE_MODE);
254+ } else {
255+ tilingKey_ = GET_TPL_TILING_KEY(attentionMode, 1U, layoutQuery, layoutKV, perfMode_ == SFAAPerfMode::V_TEMPLATE_MODE);
256+ }
257+ 
258+ OPS_LOG_I(sfaaInfo_->opName, "SFAA tilingKey_: %lu.", tilingKey_);
259+}
260+ 
261+void SFAAMlaTiling::ZeroTensorProcess()
262+{
263+ if (sfaaInfo_->s2Size == 0) {
264+ /*
265+ * 1024，空tensor场景下，作为默认值完成后续计算
266+ * 避免matmal tiling softmax tiling异常
267+ * kernel计算使用真实的seqSize=0, 与actuseq_len流程归一
268+ */
269+ sfaaInfo_->s2Size = 1024;
270+ }
271+}
272+ 
273+void SFAAMlaTiling::InitParams()
274+{
275+ perfMode_ = SFAAPerfMode::V_TEMPLATE_MODE;
276+ coreNum_ = aicNum_;
277+ 
278+ headDimAlign_ = Align(sfaaInfo_->qkHeadDim, BYTE_BLOCK); // 元素个数按照基本块大小对齐
279+ ZeroTensorProcess();
280+}
281+ 
282+void SFAAMlaTiling::CalcUbBmm()
283+{
284+ uint32_t cubeMSize = sfaaInfo_->gSize * sfaaInfo_->s1Size;
285+ uint32_t maxMSize = mBaseSize_;
286+ if (cubeMSize > maxMSize) {
287+ cubeMSize = maxMSize;
288+ }
289+ mmResUbSize_ = sInnerSizeAlign_ * Align(cubeMSize, 16U); // kernel按照16对齐写出，tiling按照这个原则分配内存
290+ bmm2ResUbSize_ = headDimAlign_ * Align(cubeMSize, 16U); // kernel按照16对齐写出，tiling按照这个原则分配内存
291+ 
292+ qPreSizeMla_ = sfaaInfo_->gSize * (headDimAlign_) * sfaaInfo_->s1Size;
293+}
294+ 
295+void SFAAMlaTiling::CheckUbSpace()
296+{
297+ CalcUbBmm();
298+}
299+ 
300+void SFAAMlaTiling::CalcInnerSize(uint32_t s2Size)
301+{
302+ sInnerSize_ = 1024; // 1024:s2默认切分大小
303+ sInnerLoopTimes_ = (s2Size + sInnerSize_ - 1) / sInnerSize_;
304+ sInnerSizeTail_ = s2Size - (sInnerLoopTimes_ - 1) * sInnerSize_;
305+ if (sInnerSize_ > s2Size) {
306+ sInnerSize_ = s2Size;
307+ }
308+ sInnerSizeAlign_ = Align(sInnerSize_, BYTE_BLOCK); // 元素个数按照基本块大小对齐
309+ 
310+ CheckUbSpace();
311+}
312+ 
313+void SFAAMlaTiling::SplitBalancedBN()
314+{
315+ CalcInnerSize(sfaaInfo_->s2Size);
316+ 
317+ //构造分核输入参数
318+ BaseInfo baseInfo;
319+ baseInfo.bSize = sfaaInfo_->bSize;
320+ baseInfo.n2Size = sfaaInfo_->n2Size;
321+ baseInfo.gSize = sfaaInfo_->gSize;
322+ baseInfo.sliding = sfaaInfo_->sparseMode == 4;
323+ 
324+ InnerSplitParams innerSplitParams;
325+ innerSplitParams.s1GBaseSize = sfaaInfo_->gSize * sfaaInfo_->sparseShardSize;
326+ innerSplitParams.s2BaseSize = sInnerSize_;
327+ tilingData_.innerSplitParams.set_mBaseSize(innerSplitParams.s1GBaseSize);
328+ tilingData_.innerSplitParams.set_s2BaseSize(innerSplitParams.s2BaseSize);
329+ 
330+ //构造分核输出参数
331+ OuterSplitParams outerSplitParams;
332+ outerSplitParams.bN2End = tilingData_.outerSplitParams.get_bN2End();
333+ outerSplitParams.gS1End = tilingData_.outerSplitParams.get_gS1End();
334+ outerSplitParams.s2End = tilingData_.outerSplitParams.get_s2End();
335+ SfaSplitCore(baseInfo, innerSplitParams, aicNum_, outerSplitParams);
336+ 
337+ usedCoreNum_ = aicNum_;
338+}
339+ 
340+// void SFAAMlaTiling::CreateSplitInput(BaseInfo &baseInfo, SplitParam &splitParam)
341+// {
342+// //构造分核输入参数
343+// baseInfo.bSize = sfaaInfo_->bSize;
344+// baseInfo.n2Size = sfaaInfo_->n2Size;
345+// baseInfo.gSize = sfaaInfo_->gSize;
346+// baseInfo.s2Size = sfaaInfo_->s2Size;
347+// baseInfo.s1Size = sfaaInfo_->s1Size;
348+// baseInfo.actualLenQDims = sfaaInfo_->actualLenDimsQ;
349+// baseInfo.actualLenKvDims = sfaaInfo_->actualLenDimsKV;
350+// baseInfo.isS1G = sfaaInfo_->qLayout == SFAALayout::TND || sfaaInfo_->qLayout == SFAALayout::BSND || sfaaInfo_->qLayout == SFAALayout::PA_BSND; // 使用枚举映射
351+// baseInfo.sparseMode = sfaaInfo_->sparseMode;
352+// baseInfo.attenMaskFlag = sfaaInfo_->sparseMode != 0 ? true : false;
353+// baseInfo.sparseBlockSize = sfaaInfo_->sparseBlockSize;
354+// baseInfo.sparseBlockCount = sfaaInfo_->sparseBlockCount;
355+// baseInfo.sparseShardSize = sfaaInfo_->sparseShardSize;
356+ 
357+ 
358+// splitParam.mBaseSize = sfaaInfo_->gSize * sfaaInfo_->sparseShardSize;
359+// splitParam.s2BaseSize = sInnerSize_;
360+// splitParam.gS1BaseSizeOfFd = mFdBaseSize_;
361+// }
362+ 
363+// void SFAAMlaTiling::SetSplitOutput(const SplitResult &res)
364+// {
365+// uint32_t *bN2EndPtr = tilingData_.outerSplitParams.get_bN2End();
366+// uint32_t *gS1EndPtr = tilingData_.outerSplitParams.get_gS1End();
367+// uint32_t *s2EndPtr = tilingData_.outerSplitParams.get_s2End();
368+// uint32_t *bN2IdxOfFdHead = tilingData_.fdParams.get_bN2IdxOfFdHead();
369+// uint32_t *gS1IdxOfFdHead = tilingData_.fdParams.get_gS1IdxOfFdHead();
370+// uint32_t *s2SplitNumOfFdHead = tilingData_.fdParams.get_s2SplitNumOfFdHead();
371+// uint32_t *s2SplitStartIdxOfCore = tilingData_.fdParams.get_s2SplitStartIdxOfCore();
372+// uint32_t *gS1SplitNumOfFdHead = tilingData_.fdParams.get_gS1SplitNumOfFdHead();
373+// uint32_t *gS1LastPartSizeOfFdHead = tilingData_.fdParams.get_gS1LastPartSizeOfFdHead();
374+// uint32_t *gS1IdxEndOfFdHead = tilingData_.fdParams.get_gS1IdxEndOfFdHead();
375+// uint32_t *gS1IdxEndOfFdHeadSplit = tilingData_.fdParams.get_gS1IdxEndOfFdHeadSplit();
376+ 
377+// for (uint32_t i = 0; i < aicNum_; ++i) {
378+// bN2EndPtr[i] = res.bN2End[i];
379+// gS1EndPtr[i] = res.gS1End[i];
380+// s2EndPtr[i] = res.s2End[i];
381+// bN2IdxOfFdHead[i] = res.fdRes.bN2IdxOfFdHead[i];
382+// gS1IdxOfFdHead[i] = res.fdRes.gS1IdxOfFdHead[i];
383+// s2SplitNumOfFdHead[i] = res.fdRes.s2SplitNumOfFdHead[i];
384+// s2SplitStartIdxOfCore[i] = res.fdRes.s2SplitStartIdxOfCore[i];
385+// gS1SplitNumOfFdHead[i] = res.fdRes.gS1SplitNumOfFdHead[i];
386+// gS1LastPartSizeOfFdHead[i] = res.fdRes.gS1LastPartSizeOfFdHead[i];
387+// }
388+ 
389+// for (uint32_t i = 0; i < aicNum_ * 2U; ++i) { // 2: cube : vector = 1:2
390+// gS1IdxEndOfFdHead[i] = res.fdRes.gS1IdxEndOfFdHead[i];
391+// gS1IdxEndOfFdHeadSplit[i] = res.fdRes.gS1IdxEndOfFdHeadSplit[i];
392+// }
393+ 
394+// tilingData_.innerSplitParams.set_mBaseSize(mBaseSize_);
395+// tilingData_.innerSplitParams.set_s2BaseSize(sInnerSize_);
396+// tilingData_.fdParams.set_gS1BaseSizeOfFd(mFdBaseSize_);
397+// tilingData_.fdParams.set_numOfFdHead(res.numOfFdHead);
398+// usedCoreNum_ = res.usedCoreNum;
399+// }
400+ 
401+void SFAAMlaTiling::Split()
402+{
403+ if (context_->GetOptionalInputDesc(METADATA_INDEX) == nullptr) {
404+ SplitBalancedBN();
405+ } else {
406+ uint32_t s2SizeInput = static_cast<uint32_t>(sfaaInfo_->s2Size);
407+ CalcInnerSize(s2SizeInput);
408+ usedCoreNum_ = aicNum_;
409+ }
410+}
411+ 
412+void SFAAMlaTiling::FillTilingBaseParamsMla()
413+{
414+ tilingData_.baseParams.set_batchSize(sfaaInfo_->bSize);
415+ tilingData_.baseParams.set_seqSize(sfaaInfo_->s2Size);
416+ tilingData_.baseParams.set_qSeqSize(sfaaInfo_->s1Size);
417+ tilingData_.baseParams.set_kvHeadNum(sfaaInfo_->n2Size);
418+ tilingData_.baseParams.set_qkHeadDim(sfaaInfo_->qkHeadDim);
419+ tilingData_.baseParams.set_ropeHeadDim(sfaaInfo_->ropeHeadDim);
420+ tilingData_.baseParams.set_blockSize(sfaaInfo_->blockSize);
421+ tilingData_.baseParams.set_maxBlockNumPerBatch(sfaaInfo_->maxBlockNumPerBatch);
422+ tilingData_.baseParams.set_scaleValue(sfaaInfo_->scaleValue);
423+ tilingData_.baseParams.set_nNumOfQInOneGroup(sfaaInfo_->n1Size / sfaaInfo_->n2Size);
424+ tilingData_.baseParams.set_actualLenDimsQ(sfaaInfo_->actualLenDimsQ);
425+ tilingData_.baseParams.set_actualLenDimsKV(sfaaInfo_->actualLenDimsKV);
426+ tilingData_.baseParams.set_sparseLenDimsKV(sfaaInfo_->sparseLenDimsKV);
427+ tilingData_.baseParams.set_outputLayout(static_cast<uint32_t>(sfaaInfo_->outLayout));
428+ tilingData_.baseParams.set_sparseMode(sfaaInfo_->sparseMode);
429+ tilingData_.baseParams.set_sparseBlockSize(sfaaInfo_->sparseBlockSize);
430+ tilingData_.baseParams.set_sparseBlockCount(sfaaInfo_->sparseBlockCount);
431+ tilingData_.baseParams.set_sparseShardSize(sfaaInfo_->sparseShardSize);
432+ tilingData_.baseParams.set_attentionMode(sfaaInfo_->attentionMode);
433+ tilingData_.baseParams.set_keyQuantMode(sfaaInfo_->keyQuantMode);
434+ tilingData_.baseParams.set_valueQuantMode(sfaaInfo_->valueQuantMode);
435+ tilingData_.baseParams.set_quantScaleRepoMode(sfaaInfo_->quantScaleRepoMode);
436+}
437+ 
438+// for flash decode
439+void SFAAMlaTiling::FillTilingSplitKVMla()
440+{
441+ tilingData_.splitKVParams.set_s2(kvSplitPart_);
442+ if (context_->GetOptionalInputDesc(METADATA_INDEX) == nullptr) {
443+ // 2:每个核可能有头规约和尾规约，一共两份规约信息
444+ tilingData_.splitKVParams.set_accumOutSize((uint64_t)aicNum_ * 2 * sfaaInfo_->n2Size * mBaseSize_ * headDimAlign_);
445+ // 2:每个核可能有头规约和尾规约，一共两份规约信息;sum + max
446+ tilingData_.splitKVParams.set_logSumExpSize((uint64_t)2 * aicNum_ * 2 * sfaaInfo_->n2Size * mBaseSize_ *
447+ (BYTE_BLOCK / BLOCK_TABLE_ELEM_BYTE));
448+ } else {
449+ // 2:每个核可能有头规约和尾规约，一共两份规约信息
450+ tilingData_.splitKVParams.set_accumOutSize((uint64_t)aicNum_ * 2 * mBaseSize_ * headDimAlign_);
451+ // 2:每个核可能有头规约和尾规约，一共两份规约信息;sum + max
452+ tilingData_.splitKVParams.set_logSumExpSize((uint64_t)2 * aicNum_ * 2 * mBaseSize_ *
453+ (BYTE_BLOCK / BLOCK_TABLE_ELEM_BYTE));
454+ }
455+ 
456+ if (!splitKVFlag_) {
457+ tilingData_.splitKVParams.set_s2(0);
458+ }
459+}
460+ 
461+void SFAAMlaTiling::FillTilingSingleCoreParamsMla()
462+{
463+ tilingData_.singleCoreParams.set_usedCoreNum(usedCoreNum_);
464+}
465+ 
466+void SFAAMlaTiling::FillTilingSingleCoreTensorSizeMla()
467+{
468+ tilingData_.singleCoreTensorSize.set_mmResUbSize(mmResUbSize_);
469+ tilingData_.singleCoreTensorSize.set_bmm2ResUbSize(bmm2ResUbSize_);
470+}
471+ 
472+void SFAAMlaTiling::FillTiling()
473+{
474+ FillTilingBaseParamsMla();
475+ FillTilingSplitKVMla();
476+ FillTilingSingleCoreParamsMla();
477+ FillTilingSingleCoreTensorSizeMla();
478+}
479+ 
480+uint32_t SFAAMlaTiling::CalcBalanceFDParamNums(const uint32_t actCoreNum)
481+{
482+ if (context_->GetOptionalInputDesc(METADATA_INDEX) == nullptr) {
483+ return actCoreNum * 2 * sfaaInfo_->n2Size * mBaseSize_; // 2:每个核可能有头规约和尾规约，一共两份规约信息
484+ } else {
485+ return actCoreNum * 2 * mBaseSize_; // 2:每个核可能有头规约和尾规约，一共两份规约信息
486+ }
487+}
488+ 
489+void SFAAMlaTiling::NormalCalcFDWorkSpace(const uint32_t actCoreNum)
490+{
491+ if (splitKVFlag_) {
492+ uint32_t accumOutSize = 0;
493+ uint32_t logSumExpSize = 0;
494+ uint32_t FDParamNums = CalcBalanceFDParamNums(actCoreNum);
495+ accumOutSize = FDParamNums * headDimAlign_;
496+ logSumExpSize = 2 * FDParamNums * (BYTE_BLOCK / sfaaInfo_->blockTypeSize); // log和sum的存储空间一致，共需要2份内存
497+ workspaceSize_ += (accumOutSize + logSumExpSize) * sfaaInfo_->blockTypeSize;
498+ if (sfaaInfo_->socVersion == platform_ascendc::SocVersion::ASCEND310P) {
499+ workspaceSize_ += static_cast<size_t>(actCoreNum) * 32; // 每个核SyncAll软同步需要32Byte记录状态
500+ }
501+ }
502+}
503+ 
504+void SFAAMlaTiling::CalcFDWorkSpace(const uint32_t actCoreNum)
505+{
506+ NormalCalcFDWorkSpace(actCoreNum);
507+}
508+ 
509+void SFAAMlaTiling::GetWorkspaceSize()
510+{
511+ uint32_t mmResElemSize = 4; // 4:fp32
512+ uint32_t vec1ResElemSize = 2; // 2:fp16/bf16
513+ uint32_t bmm2ResElemSize = 4; // 4:fp32
514+ uint32_t qPreProcResElemSize = 2; // 普通场景不涉及Q预处理
515+ uint32_t softmaxSumElemSize = 4; // 4:int32
516+ 
517+ workspaceSize_ = libapiSize_;
518+ uint32_t preLoadNum = 1;
519+ uint32_t actCoreNum = coreNum_;
520+ preLoadNum = PRE_LOAD_NUM;
521+ 
522+ // query预处理结果queryPreProcessResGm所需空间
523+ workspaceSize_ += (uint64_t)preLoadNum * (bmm2ResUbSize_ * actCoreNum * qPreProcResElemSize);
524+ // mm1结果mm1ResGm所需空间
525+ workspaceSize_ += (uint64_t)preLoadNum * (mmResUbSize_ * actCoreNum * mmResElemSize);
526+ // vec1结果vec1ResGm所需空间
527+ workspaceSize_ += (uint64_t)preLoadNum * (mmResUbSize_ * actCoreNum * vec1ResElemSize);
528+ // mm2结果mm2ResGm所需空间
529+ workspaceSize_ += (uint64_t)preLoadNum * bmm2ResUbSize_ * actCoreNum * bmm2ResElemSize;
530+ workspaceSize_ += (uint64_t)preLoadNum * (qPreSizeMla_ * actCoreNum * qPreProcResElemSize);
531+ 
532+ // FD场景，softmaxSumGm所需空间
533+ workspaceSize_ += (uint64_t)preLoadNum * mBaseSize_ * actCoreNum * softmaxSumElemSize;
534+ // vec2临时结果vec2ResGm所需空间
535+ workspaceSize_ += (uint64_t)preLoadNum * bmm2ResUbSize_ * actCoreNum * bmm2ResElemSize; // vec2ResGm
536+ 
537+ // topk BlkSize == 1场景, 需要额外空间缓存离散聚合的值
538+ // bufNum s2Base D dRope sizeOf(half)
539+ workspaceSize_ += (uint64_t)4 * 512 * (128) * 2 * actCoreNum; // 4:bufNum 512:s2MergeSize 128:D 2:sizeOf(half)
540+ // 缓存有效mte2 size的长度 份数 512B对齐的长度 sizeof(int32_t) aiv核数
541+ workspaceSize_ += (uint64_t)4 * 128 * 4 * (2 * actCoreNum); // 4:缓存有效mte2 size的长度 128:份数 4:512B对齐的长度 2:aiv核数
542+ 
543+ // FD相关
544+ workspaceSize_ += (uint64_t)aicNum_ * 2 * sfaaInfo_->n2Size * mBaseSize_ * headDimAlign_;
545+ workspaceSize_ +=(uint64_t) 2 * aicNum_ * 2 * sfaaInfo_->n2Size * mBaseSize_ *
546+ (BYTE_BLOCK / BLOCK_TABLE_ELEM_BYTE);
547+ CalcFDWorkSpace(actCoreNum);
548+}
549+ 
550+void SFAAMlaTiling::CalcBlockDim()
551+{
552+ auto ascendcPlatform = platform_ascendc::PlatformAscendC(sfaaInfo_->platformInfo);
553+ auto aicNum = usedCoreNum_;
554+ auto aivNum = 2 * usedCoreNum_;
555+ 
556+ blockDim_ = ascendcPlatform.CalcTschBlockDim(aivNum, aicNum, aivNum);
557+ OPS_LOG_I(sfaaInfo_->opName, "SFAA block dim: %u aiv Num: %u aic Num: %u.", blockDim_, aivNum, aicNum);
558+}
559+ 
560+ge::graphStatus SFAAMlaTiling::DoOpTiling(SFAATilingInfo *sfaaInfo)
561+{
562+ sfaaInfo_ = sfaaInfo;
563+ if (GetPlatformInfo() != ge::GRAPH_SUCCESS) {
564+ return ge::GRAPH_FAILED;
565+ }
566+ InitParams();
567+ Split();
568+ FillTiling();
569+ CalcBlockDim();
570+ GetWorkspaceSize();
571+ GenTilingKey();
572+ 
573+ if ((SetBlockDim(blockDim_) != ge::GRAPH_SUCCESS) ||
574+ (SetTilingKey(tilingKey_) != ge::GRAPH_SUCCESS) ||
575+ (SetWorkspaceSize(workspaceSize_) != ge::GRAPH_SUCCESS) ||
576+ (SetTilingData(tilingData_) != ge::GRAPH_SUCCESS)) {
577+ return ge::GRAPH_FAILED;
578+ }
579+ 
580+ return ge::GRAPH_SUCCESS;
581+}
582+ 
583+ge::graphStatus TilingSparseFlashAttentionAntiquant(gert::TilingContext *context)
584+{
585+ SFAATilingInfo sfaaInfo;
586+ SFAAInfoParser sfaaInfoParser(context);
587+ if (sfaaInfoParser.Parse(sfaaInfo) != ge::GRAPH_SUCCESS) {
588+ return ge::GRAPH_FAILED;
589+ }
590+ 
591+ SFAATilingCheck tilingChecker(sfaaInfo);
592+ if (tilingChecker.Process() != ge::GRAPH_SUCCESS) {
593+ return ge::GRAPH_FAILED;
594+ }
595+ 
596+ SFAAMlaTiling tiling(context);
597+ return tiling.DoOpTiling(&sfaaInfo);
598+}
599+ 
600+ge::graphStatus TilingPrepareForSparseFlashAttentionAntiquant(gert::TilingParseContext *context)
601+{
602+ (void)context;
603+ return ge::GRAPH_SUCCESS;
604+}
605+ 
606+ge::graphStatus SFAATilingCheck::GetExpectedShape(gert::Shape &shapeExpected,
607+ const SFAATilingShapeCompareParam &param, const SFAALayout &layout) const
608+{
609+ if (layout == SFAALayout::BSND) {
610+ shapeExpected = gert::Shape({param.B, param.S, param.N, param.D});
611+ } else if (layout == SFAALayout::TND) {
612+ shapeExpected = gert::Shape({param.T, param.N, param.D});
613+ } else if (layout == SFAALayout::PA_BSND) {
614+ shapeExpected = gert::Shape({param.Bn, param.Bs, param.N, param.D});
615+ } else if (layout == SFAALayout::PA_BNSD) {
616+ shapeExpected = gert::Shape({param.Bn, param.N, param.Bs, param.D});
617+ } else {
618+ OPS_LOG_E(opName_, "layout %s is unsupported", SFAALayoutToSerialString(layout).c_str());
619+ return ge::GRAPH_FAILED;
620+ }
621+ return ge::GRAPH_SUCCESS;
622+}
623+ 
624+ge::graphStatus SFAATilingCheck::CompareShape(SFAATilingShapeCompareParam &param,
625+ const gert::Shape &shape, const SFAALayout &layout, const std::string &name) const
626+{
627+ gert::Shape shapeExpected;
628+ if (GetExpectedShape(shapeExpected, param, layout) != ge::GRAPH_SUCCESS) {
629+ return ge::GRAPH_FAILED;
630+ }
631+ 
632+ if (shape.GetDimNum() != shapeExpected.GetDimNum()) {
633+ OPS_LOG_E(opName_,
634+ "%s dimension is %zu, expected dimension is %zu.",
635+ name.c_str(), shape.GetDimNum(), shapeExpected.GetDimNum());
636+ return ge::GRAPH_FAILED;
637+ }
638+ 
639+ for (size_t i = 0; i < shape.GetDimNum(); i++) {
640+ if (shape.GetDim(i) != shapeExpected.GetDim(i)) {
641+ OPS_LOG_E(opName_, "%s layout is %s, shape is %s, expected shape is %s.",
642+ name.c_str(), SFAALayoutToSerialString(layout).c_str(),
643+ GetShapeStr(shape).c_str(), GetShapeStr(shapeExpected).c_str());
644+ return ge::GRAPH_FAILED;
645+ }
646+ }
647+ 
648+ return ge::GRAPH_SUCCESS;
649+}
650+ 
651+void SFAATilingCheck::LogErrorDtypeSupport(const std::vector<ge::DataType> &expectDtypeList,
652+ const ge::DataType &actualDtype, const std::string &name) const
653+{
654+ std::ostringstream oss;
655+ for (size_t i = 0; i < expectDtypeList.size(); ++i) {
656+ oss << SFAADataTypeToSerialString(expectDtypeList[i]);
657+ if (i < expectDtypeList.size() - 1) {
658+ oss << ", ";
659+ }
660+ }
661+ OPS_LOG_E(opName_, "Tensor %s only supports dtype %s, but got %s",
662+ name.c_str(), oss.str().c_str(), SFAADataTypeToSerialString(actualDtype).c_str());
663+}
664+ 
665+ge::graphStatus SFAATilingCheck::CheckDtypeSupport(const gert::CompileTimeTensorDesc *desc,
666+ const std::string &name) const
667+{
668+ if (desc != nullptr) {
669+ const auto& it = DTYPE_SUPPORT_MAP.find(name);
670+ OPS_ERR_IF(it == DTYPE_SUPPORT_MAP.end(),
671+ OPS_LOG_E(opName_, "%s datatype support list should be specify in DTYPE_SUPPORT_MAP", name.c_str()),
672+ return ge::GRAPH_FAILED);
673+ auto &expectDtypeList = it->second;
674+ OPS_ERR_IF(std::find(
675+ expectDtypeList.begin(), expectDtypeList.end(), desc->GetDataType()) == expectDtypeList.end(),
676+ LogErrorDtypeSupport(expectDtypeList, desc->GetDataType(), name),
677+ return ge::GRAPH_FAILED);
678+ }
679+ return ge::GRAPH_SUCCESS;
680+}
681+ 
682+template <typename T>
683+void SFAATilingCheck::LogErrorNumberSupport(const std::vector<T> &expectNumberList,
684+ const T &actualValue, const std::string &name, const std::string subName) const
685+{
686+ std::ostringstream oss;
687+ for (size_t i = 0; i < expectNumberList.size(); ++i) {
688+ oss << std::to_string(expectNumberList[i]);
689+ if (i < expectNumberList.size() - 1) {
690+ oss << ", ";
691+ }
692+ }
693+ 
694+ OPS_LOG_E(opName_, "%s %s only supports %s, but got %s",
695+ name.c_str(), subName.c_str(), oss.str().c_str(), std::to_string(actualValue).c_str());
696+}
697+ 
698+template <typename T>
699+void SFAATilingCheck::LogErrorDimNumSupport(const std::vector<T> &expectNumberList,
700+ const T &actualValue, const std::string &name) const
701+{
702+ LogErrorNumberSupport(expectNumberList, actualValue, name, "dimension");
703+}
704+ 
705+ge::graphStatus SFAATilingCheck::CheckDimNumInLayoutSupport(const SFAALayout &layout,
706+ const gert::StorageShape *shape, const std::string &name) const
707+{
708+ const auto& dimIt = SFAA_LAYOUT_DIM_MAP.find(layout);
709+ OPS_ERR_IF(shape->GetStorageShape().GetDimNum() != dimIt->second,
710+ OPS_LOG_E(opName_, "When layout is %s, %s dimension should be %zu, but it's %zu",
711+ SFAALayoutToSerialString(layout).c_str(), name.c_str(), dimIt->second,
712+ shape->GetStorageShape().GetDimNum()),
713+ return ge::GRAPH_FAILED);
714+ return ge::GRAPH_SUCCESS;
715+}
716+ 
717+ge::graphStatus SFAATilingCheck::CheckDimNumSupport(const gert::StorageShape *shape,
718+ const std::vector<size_t> &expectDimNumList, const std::string &name) const
719+{
720+ if (shape == nullptr) {
721+ return ge::GRAPH_SUCCESS;
722+ }
723+ 
724+ if (std::find(expectDimNumList.begin(), expectDimNumList.end(),
725+ shape->GetStorageShape().GetDimNum()) == expectDimNumList.end()) {
726+ LogErrorDimNumSupport(expectDimNumList, shape->GetStorageShape().GetDimNum(), name);
727+ return ge::GRAPH_FAILED;
728+ }
729+ 
730+ return ge::GRAPH_SUCCESS;
731+}
732+ 
733+ 
734+void SFAATilingCheck::LogErrorLayoutSupport(const std::vector<SFAALayout> &expectLayoutList,
735+ const SFAALayout &actualLayout, const std::string &name) const
736+{
737+ std::ostringstream oss;
738+ for (size_t i = 0; i < expectLayoutList.size(); ++i) {
739+ oss << SFAALayoutToSerialString(expectLayoutList[i]);
740+ if (i < expectLayoutList.size() - 1) {
741+ oss << ", ";
742+ }
743+ }
744+ OPS_LOG_E(opName_, "Tensor %s only supports layout %s, but got %s",
745+ name.c_str(), oss.str().c_str(), SFAALayoutToSerialString(actualLayout).c_str());
746+}
747+ 
748+ge::graphStatus SFAATilingCheck::CheckLayoutSupport(const SFAALayout &actualLayout, const std::string &name) const
749+{
750+ const auto& it = LAYOUT_SUPPORT_MAP.find(name);
751+ OPS_ERR_IF(it == LAYOUT_SUPPORT_MAP.end(),
752+ OPS_LOG_E(opName_, "%s layout support list should be specify in LAYOUT_SUPPORT_MAP", name.c_str()),
753+ return ge::GRAPH_FAILED);
754+ auto &expectLayoutList = it->second;
755+ OPS_ERR_IF(std::find(
756+ expectLayoutList.begin(), expectLayoutList.end(), actualLayout) == expectLayoutList.end(),
757+ LogErrorLayoutSupport(expectLayoutList, actualLayout, name),
758+ return ge::GRAPH_FAILED);
759+ 
760+ return ge::GRAPH_SUCCESS;
761+}
762+ 
763+ge::graphStatus SFAATilingCheck::CheckSingleParaQuery() const
764+{
765+ const std::vector<size_t> queryDimNumList = {DIM_NUM_THREE, DIM_NUM_FOUR};
766+ if (ge::GRAPH_SUCCESS != CheckDtypeSupport(opParamInfo_.query.desc, QUERY_NAME) ||
767+ ge::GRAPH_SUCCESS != CheckLayoutSupport(qLayout_, QUERY_NAME) ||
768+ ge::GRAPH_SUCCESS != CheckDimNumSupport(opParamInfo_.query.shape, queryDimNumList, QUERY_NAME) ||
769+ ge::GRAPH_SUCCESS != CheckDimNumInLayoutSupport(qLayout_, opParamInfo_.query.shape, QUERY_NAME)) {
770+ return ge::GRAPH_FAILED;
771+ }
772+ return ge::GRAPH_SUCCESS;
773+}
774+ 
775+ge::graphStatus SFAATilingCheck::CheckSingleParaKey() const
776+{
777+ const std::vector<size_t> keyDimNumList = {DIM_NUM_FOUR, DIM_NUM_FIVE};
778+ if (ge::GRAPH_SUCCESS != CheckDtypeSupport(opParamInfo_.key.desc, KEY_NAME) ||
779+ ge::GRAPH_SUCCESS != CheckLayoutSupport(kvLayout_, KEY_NAME) ||
780+ ge::GRAPH_SUCCESS != CheckDimNumSupport(opParamInfo_.key.shape, keyDimNumList, KEY_NAME) ||
781+ ge::GRAPH_SUCCESS != CheckDimNumInLayoutSupport(kvLayout_, opParamInfo_.key.shape, KEY_NAME)) {
782+ return ge::GRAPH_FAILED;
783+ }
784+ return ge::GRAPH_SUCCESS;
785+}
786+ 
787+ge::graphStatus SFAATilingCheck::CheckSingleParaNumHeads() const
788+{
789+ return ge::GRAPH_SUCCESS;
790+}
791+ 
792+ge::graphStatus SFAATilingCheck::CheckSingleParaKvHeadNums() const
793+{
794+ return ge::GRAPH_SUCCESS;
795+}
796+ 
797+ge::graphStatus SFAATilingCheck::CheckSingleParaSparseMode() const
798+{
799+ OPS_ERR_IF((*opParamInfo_.sparseMode != 3 && *opParamInfo_.sparseMode != 0),
800+ OPS_LOG_E(opName_, "sparseMode must == 0/3, but got: %ld.", *opParamInfo_.sparseMode),
801+ return ge::GRAPH_FAILED);
802+ return ge::GRAPH_SUCCESS;
803+}
804+ 
805+ge::graphStatus SFAATilingCheck::CheckSingleParaSparseBlockSize() const
806+{
807+ OPS_ERR_IF((*opParamInfo_.sparseBlockSize != 16),
808+ OPS_LOG_E(opName_, "sparseBlockSize should be equal 16, but got: %ld.", *opParamInfo_.sparseBlockSize),
809+ return ge::GRAPH_FAILED);
810+ return ge::GRAPH_SUCCESS;
811+}
812+ 
813+ge::graphStatus SFAATilingCheck::CheckSingleParaSparseIndices() const
814+{
815+ if (ge::GRAPH_SUCCESS != CheckDtypeSupport(opParamInfo_.sparseIndices.desc, SPARSE_INDICES_NAME)) {
816+ return ge::GRAPH_FAILED;
817+ }
818+ return ge::GRAPH_SUCCESS;
819+}
820+ 
821+ge::graphStatus SFAATilingCheck::CheckSinglePara() const
822+{
823+ if (ge::GRAPH_SUCCESS != CheckSingleParaQuery() ||
824+ ge::GRAPH_SUCCESS != CheckSingleParaKey() ||
825+ ge::GRAPH_SUCCESS != CheckSingleParaSparseIndices() ||
826+ ge::GRAPH_SUCCESS != CheckSingleParaNumHeads() ||
827+ ge::GRAPH_SUCCESS != CheckSingleParaKvHeadNums() ||
828+ ge::GRAPH_SUCCESS != CheckSingleParaSparseMode() ||
829+ ge::GRAPH_SUCCESS != CheckSingleParaSparseBlockSize()) {
830+ return ge::GRAPH_FAILED;
831+ }
832+ 
833+ return ge::GRAPH_SUCCESS;
834+}
835+ 
836+ge::graphStatus SFAATilingCheck::CheckDequantScaleNotExistence()
837+{
838+ if (quantScaleRepoMode_ == 1) {
839+ OPS_ERR_IF((opParamInfo_.keyDequantScale.tensor == nullptr || opParamInfo_.valueDequantScale.tensor == nullptr),
M
Mmiaofangzheng4月2日

函数名是 CheckDequantScaleNotExistence，语义是"检查 dequant scale 不存在"。但报错条件是"tensor 为 nullptr"，报错信息却说"should be both null"。逻辑矛盾：如果"不存在"意味着应该为空，那么 nullptr 时不应报错。请确认意图——如果是检查 combine 模式下 dequant scale 必须存在，应改为 CheckDequantScaleExistence 并检查非空；如果是检查必须不存在，则条件应为 != nullptr。

likedislike
M
mayafei
4月3日 评论：
840+ OPS_LOG_E(opName_,
841+ "When quant_scale_repo_mode is 1(combine), key_dequant_scale and value_dequant_scale should not be null."),
842+ return ge::GRAPH_FAILED);
843+ }
844+ return ge::GRAPH_SUCCESS;
845+}
846+ 
847+ge::graphStatus SFAATilingCheck::CheckExists(const void *pointer, const std::string &name) const
848+{
849+ OPS_ERR_IF(pointer == nullptr,
850+ OPS_LOG_E(opName_, "%s should not be null", name.c_str()),
851+ return ge::GRAPH_FAILED);
852+ return ge::GRAPH_SUCCESS;
853+}
854+ 
855+ge::graphStatus SFAATilingCheck::CheckNotExists(const void *pointer, const std::string &name) const
856+{
857+ OPS_ERR_IF(pointer != nullptr,
858+ OPS_LOG_E(opName_, "%s should be null", name.c_str()),
859+ return ge::GRAPH_FAILED);
860+ return ge::GRAPH_SUCCESS;
861+}
862+ 
863+ge::graphStatus SFAATilingCheck::CheckExistsByMap(const std::map<std::string, const void *> &paramMap) const
864+{
865+ for (const auto& kv : paramMap) {
866+ if (CheckExists(kv.second, kv.first) != ge::GRAPH_SUCCESS) {
867+ return ge::GRAPH_FAILED;
868+ }
869+ }
870+ return ge::GRAPH_SUCCESS;
871+}
872+ 
873+ge::graphStatus SFAATilingCheck::CheckNotExistsByMap(const std::map<std::string, const void *> &paramMap) const
874+{
875+ for (const auto& kv : paramMap) {
876+ if (CheckNotExists(kv.second, kv.first) != ge::GRAPH_SUCCESS) {
877+ return ge::GRAPH_FAILED;
878+ }
879+ }
880+ return ge::GRAPH_SUCCESS;
881+}
882+ 
883+ge::graphStatus SFAATilingCheck::CheckExistenceByMap(std::map<std::string, const void *> &existMap,
884+ std::map<std::string, const void *> &notExistMap) const
885+{
886+ if (CheckExistsByMap(existMap) != ge::GRAPH_SUCCESS) {
887+ return ge::GRAPH_FAILED;
888+ }
889+ if (CheckNotExistsByMap(notExistMap) != ge::GRAPH_SUCCESS) {
890+ return ge::GRAPH_FAILED;
891+ }
892+ return ge::GRAPH_SUCCESS;
893+}
894+ 
895+template <typename T>
896+ge::graphStatus SFAATilingCheck::CheckAttrValueByMap(std::map<std::string, std::pair<const T *, T>> &attrMap) const
897+{
898+ for (auto const &kv : attrMap) {
899+ const std::string &name = kv.first;
900+ const std::pair<const T *, T> &pointerValuePair = kv.second;
901+ if (pointerValuePair.first == nullptr) {
902+ OPS_LOG_E(opName_, "Attr %s should not be nullptr", name.c_str());
903+ return ge::GRAPH_FAILED;
904+ }
905+ 
906+ if (*(pointerValuePair.first) != pointerValuePair.second) {
907+ std::ostringstream ossExpect;
908+ ossExpect << std::to_string(pointerValuePair.second);
909+ std::ostringstream ossActual;
910+ ossActual << std::to_string(*(pointerValuePair.first));
911+ OPS_LOG_E(opName_,
912+ "%s value should be %s, but got %s",
913+ name.c_str(),
914+ ossExpect.str().c_str(),
915+ ossActual.str().c_str());
916+ return ge::GRAPH_FAILED;
917+ }
918+ }
919+ return ge::GRAPH_SUCCESS;
920+}
921+ 
922+ge::graphStatus SFAATilingCheck::CheckParaExistenceMlaAntiquant() const
923+{
924+ if (kvStorageMode_ != KvStorageMode::PAGE_ATTENTION) {
925+ return ge::GRAPH_SUCCESS;
926+ }
927+ std::map<std::string, const void *> mlaAntiquantParamExistMap = {
928+ {"actualSeqLengths", opParamInfo_.actualSeqLengths.tensor},
929+ {"blockTable", opParamInfo_.blockTable.tensor},
930+ };
931+ std::map<std::string, const void *> mlaAntiquantParamNotExistMap = {};
932+ if (CheckExistenceByMap(mlaAntiquantParamExistMap, mlaAntiquantParamNotExistMap) != ge::GRAPH_SUCCESS) {
933+ return ge::GRAPH_FAILED;
934+ }
935+ return ge::GRAPH_SUCCESS;
936+}
937+ 
938+ge::graphStatus SFAATilingCheck::CheckParaExistenceMla() const
939+{
940+ return CheckParaExistenceMlaAntiquant();
941+}
942+ 
943+ge::graphStatus SFAATilingCheck::CheckParaExistence()
944+{
945+ if (ge::GRAPH_SUCCESS != CheckDequantScaleNotExistence()) {
946+ return ge::GRAPH_FAILED;
947+ }
948+ 
949+ return CheckParaExistenceMla();
950+}
951+ 
952+ge::graphStatus SFAATilingCheck::GetActualSeqLenSize(uint32_t &size, const gert::Tensor *tensor,
953+ const SFAALayout &layoutQuery, const std::string &name)
954+{
955+ if (tensor == nullptr) {
956+ OPS_LOG_E(opName_, "when layout of query is %s, %s must be provided.",
957+ SFAALayoutToSerialString(layoutQuery).c_str(), name.c_str());
958+ return ge::GRAPH_FAILED;
959+ }
960+ int64_t shapeSize = tensor->GetShapeSize();
961+ if (shapeSize <= 0) {
962+ OPS_LOG_E(opName_, "the shape size of %s is %ld, it should be greater than 0.",
963+ name.c_str(), shapeSize);
964+ return ge::GRAPH_FAILED;
965+ }
966+ size = static_cast<uint32_t>(shapeSize);
967+ return ge::GRAPH_SUCCESS;
968+}
969+ 
970+void SFAATilingCheck::SetSFAAShapeCompare()
971+{
972+ queryShapeCmp_ = opParamInfo_.query.shape->GetStorageShape();
973+ topkShapeCmp_ = opParamInfo_.sparseIndices.shape->GetStorageShape();
974+ keyShapeCmp_ = opParamInfo_.key.shape->GetStorageShape();
975+ valueShapeCmp_ = opParamInfo_.value.shape->GetStorageShape();
976+ attenOutShapeCmp_ = opParamInfo_.attenOut.shape->GetStorageShape();
977+}
978+ 
979+ge::graphStatus SFAATilingCheck::CheckBlockTable() const
980+{
981+ if (kvStorageMode_ != KvStorageMode::PAGE_ATTENTION) {
982+ OPS_ERR_IF(opParamInfo_.blockTable.tensor != nullptr,
983+ OPS_LOG_E(opName_, "when the layout_kv is %s, %s should be null",
984+ SFAALayoutToSerialString(kvLayout_).c_str(), BLOCK_TABLE_NAME.c_str()),
985+ return ge::GRAPH_FAILED);
986+ return ge::GRAPH_SUCCESS;
987+ }
988+
989+ uint32_t blockTableBatch = opParamInfo_.blockTable.tensor->GetStorageShape().GetDim(0);
990+ OPS_ERR_IF(blockTableBatch != bSize_,
991+ OPS_LOG_E(opName_, "%s's first dimension(%u) should be equal to batch size(%u)",
992+ BLOCK_TABLE_NAME.c_str(), blockTableBatch, bSize_),
993+ return ge::GRAPH_FAILED);
994+ if (opParamInfo_.blockTable.desc->GetDataType() != ge::DT_INT32) {
995+ OPS_LOG_E(opName_, "blockTable's dtype is %s, it should be DT_INT32.",
996+ SFAADataTypeToSerialString(opParamInfo_.blockTable.desc->GetDataType()).c_str());
997+ return ge::GRAPH_FAILED;
998+ }
999+ return ge::GRAPH_SUCCESS;
1000+}
1001+ 
1002+ge::graphStatus SFAATilingCheck::CheckDTypeConsistency(const ge::DataType &actualDtype,
1003+ const ge::DataType &expectDtype, const std::string &name) const
1004+{
1005+ if (actualDtype != expectDtype) {
1006+ OPS_LOG_E(opName_, "%s dtype should be %s, but it's %s.", name.c_str(),
1007+ SFAADataTypeToSerialString(expectDtype).c_str(),
1008+ SFAADataTypeToSerialString(actualDtype).c_str());
1009+ return ge::GRAPH_FAILED;
1010+ }
1011+ return ge::GRAPH_SUCCESS;
1012+}
1013+ 
1014+ge::graphStatus SFAATilingCheck::CheckTopkShape()
1015+{
1016+ SFAATilingShapeCompareParam shapeParams;
1017+ shapeParams.B = bSize_;
1018+ shapeParams.N = n2Size_;
1019+ shapeParams.S = s1Size_ / sparseShardSize_;
1020+ shapeParams.D = sparseBlockCount_;
1021+ shapeParams.T = qTSize_;
1022+ return CompareShape(shapeParams, topkShapeCmp_, topkLayout_, SPARSE_INDICES_NAME);
1023+}
1024+ 
1025+ge::graphStatus SFAATilingCheck::CheckAttenOutShape()
1026+{
1027+ SFAATilingShapeCompareParam shapeParams;
1028+ shapeParams.B = bSize_;
1029+ shapeParams.N = n1Size_;
1030+ shapeParams.S = s1Size_;
1031+ shapeParams.D = (attentionMode_ == 0) ? 128 : 512;
1032+ shapeParams.T = qTSize_;
1033+ if (CompareShape(shapeParams, attenOutShapeCmp_, outLayout_, ATTEN_OUT_NAME) != ge::GRAPH_SUCCESS) {
1034+ return ge::GRAPH_FAILED;
1035+ }
1036+ return ge::GRAPH_SUCCESS;
1037+}
1038+ 
1039+ge::graphStatus SFAATilingCheck::CheckAttenOut()
1040+{
1041+ if (ge::GRAPH_SUCCESS != CheckDTypeConsistency(opParamInfo_.attenOut.desc->GetDataType(),
1042+ inputQType_, ATTEN_OUT_NAME) ||
1043+ ge::GRAPH_SUCCESS != CheckAttenOutShape()) {
1044+ return ge::GRAPH_FAILED;
1045+ }
1046+ return ge::GRAPH_SUCCESS;
1047+}
1048+ 
1049+ge::graphStatus SFAATilingCheck::CheckTopK()
1050+{
1051+ if (ge::GRAPH_SUCCESS != CheckTopkShape()) {
1052+ return ge::GRAPH_FAILED;
1053+ }
1054+ return ge::GRAPH_SUCCESS;
1055+}
1056+ 
1057+ge::graphStatus SFAATilingCheck::CheckKVShapeForBatchContinuous()
1058+{
1059+ SFAATilingShapeCompareParam shapeParams;
1060+ shapeParams.B = bSize_;
1061+ shapeParams.N = n2Size_;
1062+ shapeParams.S = s2Size_;
1063+ shapeParams.D = vHeadDim_;
1064+ shapeParams.T = kvTSize_;
1065+ if (CompareShape(shapeParams, valueShapeCmp_, kvLayout_, VALUE_NAME) != ge::GRAPH_SUCCESS) {
1066+ return ge::GRAPH_FAILED;
1067+ }
1068+ 
1069+ return ge::GRAPH_SUCCESS;
1070+}
1071+ 
1072+uint32_t SFAATilingCheck::GetTypeSize(ge::DataType dtype) const
1073+{
1074+ uint32_t typeSize = NUM_BYTES_FLOAT16;
1075+ switch (dtype) {
1076+ case ge::DT_FLOAT16:
1077+ typeSize = NUM_BYTES_FLOAT16;
1078+ break;
1079+ case ge::DT_BF16:
1080+ typeSize = NUM_BYTES_BF16;
1081+ break;
1082+ default:
1083+ typeSize = NUM_BYTES_FLOAT16;
1084+ }
1085+ return typeSize;
1086+}
1087+ 
1088+ge::graphStatus SFAATilingCheck::CheckKVShapeForPageAttention()
1089+{
1090+ uint32_t kvBlockElemNum = 32 / GetTypeSize(inputKvType_);
1091+ 
1092+ int64_t blockNum = keyShapeCmp_.GetDim(0);
1093+ SFAATilingShapeCompareParam shapeParams;
1094+ shapeParams.Bn = blockNum;
1095+ shapeParams.N = n2Size_;
1096+ shapeParams.Bs = blockSize_;
1097+ shapeParams.T = kvTSize_;
1098+ shapeParams.D = vHeadDim_;
1099+ if (CompareShape(shapeParams, valueShapeCmp_, kvLayout_, VALUE_NAME) != ge::GRAPH_SUCCESS) {
1100+ return ge::GRAPH_FAILED;
1101+ }
1102+ 
1103+ return ge::GRAPH_SUCCESS;
1104+}
1105+ 
1106+ge::graphStatus SFAATilingCheck::CheckKVShape()
1107+{
1108+ if (kvStorageMode_ == KvStorageMode::BATCH_CONTINUOUS) {
1109+ return CheckKVShapeForBatchContinuous();
1110+ }
1111+ 
1112+ if (kvStorageMode_ == KvStorageMode::PAGE_ATTENTION) {
1113+ return CheckKVShapeForPageAttention();
1114+ }
1115+ 
1116+ OPS_LOG_E(opName_, "storage mode of key and value is %u, it is incorrect.", static_cast<uint32_t>(kvStorageMode_));
1117+ return ge::GRAPH_FAILED;
1118+}
1119+ 
1120+ge::graphStatus SFAATilingCheck::CheckKV()
1121+{
1122+ if (ge::GRAPH_SUCCESS != CheckDTypeConsistency(opParamInfo_.value.desc->GetDataType(),
1123+ inputKvType_, VALUE_NAME)) {
1124+ return ge::GRAPH_FAILED;
1125+ }
1126+ return ge::GRAPH_SUCCESS;
1127+}
1128+ 
1129+ge::graphStatus SFAATilingCheck::CheckActualSeqLensQ()
1130+{
1131+ if (ge::GRAPH_SUCCESS != CheckActualSeqLensQDType() ||
1132+ ge::GRAPH_SUCCESS != CheckActualSeqLensQShape()) {
1133+ return ge::GRAPH_FAILED;
1134+ }
1135+ return ge::GRAPH_SUCCESS;
1136+}
1137+ 
1138+ge::graphStatus SFAATilingCheck::CheckActualSeqLensQDType()
1139+{
1140+ if (opParamInfo_.actualSeqLengthsQ.tensor == nullptr) {
1141+ return ge::GRAPH_SUCCESS;
1142+ }
1143+ if (opParamInfo_.actualSeqLengthsQ.desc == nullptr) {
1144+ OPS_LOG_E(opName_, "actualSeqLengthsQ is not empty,"
1145+ "but actualSeqLengthsQ's dtype is nullptr.");
1146+ return ge::GRAPH_FAILED;
1147+ }
1148+ if (opParamInfo_.actualSeqLengthsQ.desc->GetDataType() != ge::DT_INT32) {
1149+ OPS_LOG_E(opName_, "actualSeqLengthsQ's dtype is %s, it should be DT_INT32.",
1150+ SFAADataTypeToSerialString(opParamInfo_.actualSeqLengthsQ.desc->GetDataType()).c_str());
1151+ return ge::GRAPH_FAILED;
1152+ }
1153+ return ge::GRAPH_SUCCESS;
1154+}
1155+ 
1156+ge::graphStatus SFAATilingCheck::CheckActualSeqLensQShape()
1157+{
1158+ if (opParamInfo_.actualSeqLengthsQ.tensor == nullptr) {
1159+ return ge::GRAPH_SUCCESS;
1160+ }
1161+ uint32_t shapeSize = 0;
1162+ if (GetActualSeqLenSize(shapeSize, opParamInfo_.actualSeqLengthsQ.tensor, qLayout_, "actualSeqLengthsQ") !=
1163+ ge::GRAPH_SUCCESS) {
1164+ return ge::GRAPH_FAILED;
1165+ }
1166+ if (shapeSize != bSize_) {
1167+ OPS_LOG_E(opName_, "actualSeqLengthsQ shape size is %u, it should be equal to batch size[%u]",
1168+ shapeSize, bSize_);
1169+ return ge::GRAPH_FAILED;
1170+ }
1171+ return ge::GRAPH_SUCCESS;
1172+}
1173+ 
1174+ge::graphStatus SFAATilingCheck::CheckActualSeqLens()
1175+{
1176+ if (ge::GRAPH_SUCCESS != CheckActualSeqLensDType() ||
1177+ ge::GRAPH_SUCCESS != CheckActualSeqLensShape()) {
1178+ return ge::GRAPH_FAILED;
1179+ }
1180+ return ge::GRAPH_SUCCESS;
1181+}
1182+ 
1183+ge::graphStatus SFAATilingCheck::CheckActualSeqLensDType()
1184+{
1185+ if (opParamInfo_.actualSeqLengths.tensor == nullptr) {
1186+ return ge::GRAPH_SUCCESS;
1187+ }
1188+ if (opParamInfo_.actualSeqLengths.desc == nullptr) {
1189+ OPS_LOG_E(opName_, "actualSeqLengths is not empty,"
1190+ "but actualSeqLengths's dtype is nullptr.");
1191+ return ge::GRAPH_FAILED;
1192+ }
1193+ if (opParamInfo_.actualSeqLengths.desc->GetDataType() != ge::DT_INT32) {
1194+ OPS_LOG_E(opName_, "actualSeqLengths's dtype is %s, it should be DT_INT32.",
1195+ SFAADataTypeToSerialString(opParamInfo_.actualSeqLengths.desc->GetDataType()).c_str());
1196+ return ge::GRAPH_FAILED;
1197+ }
1198+ return ge::GRAPH_SUCCESS;
1199+}
1200+ 
1201+ge::graphStatus SFAATilingCheck::CheckActualSeqLensShape()
1202+{
1203+ if (opParamInfo_.actualSeqLengths.tensor == nullptr) {
1204+ return ge::GRAPH_SUCCESS;
1205+ }
1206+ uint32_t shapeSize = 0;
1207+ if (GetActualSeqLenSize(shapeSize, opParamInfo_.actualSeqLengths.tensor, kvLayout_, "actualSeqLengths") !=
1208+ ge::GRAPH_SUCCESS) {
1209+ return ge::GRAPH_FAILED;
1210+ }
1211+ if (shapeSize != bSize_) {
1212+ OPS_LOG_E(opName_, "actualSeqLengths shape size is %u, it should be equal to batch size[%u].",
1213+ shapeSize, bSize_);
1214+ return ge::GRAPH_FAILED;
1215+ }
1216+ return ge::GRAPH_SUCCESS;
1217+}
1218+ 
1219+ge::graphStatus SFAATilingCheck::CheckSparseSeqLens()
1220+{
1221+ if (ge::GRAPH_SUCCESS != CheckSparseSeqLensDType() ||
1222+ ge::GRAPH_SUCCESS != CheckSparseSeqLensShape()) {
1223+ return ge::GRAPH_FAILED;
1224+ }
1225+ return ge::GRAPH_SUCCESS;
1226+}
1227+ 
1228+ge::graphStatus SFAATilingCheck::CheckSparseSeqLensDType()
1229+{
1230+ if (opParamInfo_.sparseSeqLengths.tensor == nullptr) {
1231+ return ge::GRAPH_SUCCESS;
1232+ }
1233+ if (opParamInfo_.sparseSeqLengths.desc == nullptr) {
1234+ OPS_LOG_E(opName_, "sparseSeqLengths is not empty,"
1235+ "but sparseSeqLengths's dtype is nullptr.");
1236+ return ge::GRAPH_FAILED;
1237+ }
1238+ if (opParamInfo_.sparseSeqLengths.desc->GetDataType() != ge::DT_INT32) {
1239+ OPS_LOG_E(opName_, "sparseSeqLengths's dtype is %s, it should be DT_INT32.",
1240+ SFAADataTypeToSerialString(opParamInfo_.sparseSeqLengths.desc->GetDataType()).c_str());
1241+ return ge::GRAPH_FAILED;
1242+ }
1243+ return ge::GRAPH_SUCCESS;
1244+}
1245+ 
1246+ge::graphStatus SFAATilingCheck::CheckSparseSeqLensShape()
1247+{
1248+ if (opParamInfo_.sparseSeqLengths.tensor == nullptr) {
1249+ return ge::GRAPH_SUCCESS;
1250+ }
1251+ uint32_t shapeSize = 0;
1252+ if (GetActualSeqLenSize(shapeSize, opParamInfo_.sparseSeqLengths.tensor, kvLayout_, "sparseSeqLengths") !=
1253+ ge::GRAPH_SUCCESS) {
1254+ return ge::GRAPH_FAILED;
1255+ }
1256+ if (shapeSize != bSize_) {
1257+ OPS_LOG_E(opName_, "sparseSeqLengths shape size is %u, it should be equal to batch size[%u].",
1258+ shapeSize, bSize_);
1259+ return ge::GRAPH_FAILED;
1260+ }
1261+ return ge::GRAPH_SUCCESS;
1262+}
1263+ 
1264+ge::graphStatus SFAATilingCheck::CheckMultiParaConsistency()
1265+{
1266+ SetSFAAShapeCompare();
1267+ if (ge::GRAPH_SUCCESS != CheckKV() ||
1268+ ge::GRAPH_SUCCESS != CheckTopK() ||
1269+ ge::GRAPH_SUCCESS != CheckAttenOut() ||
1270+ ge::GRAPH_SUCCESS != CheckActualSeqLensQ() ||
1271+ ge::GRAPH_SUCCESS != CheckActualSeqLens() ||
1272+ ge::GRAPH_SUCCESS != CheckBlockTable()) {
1273+ return ge::GRAPH_FAILED;
1274+ }
1275+ 
1276+ // GQA模式需要传sparseSeqLens
1277+ if (attentionMode_ == 0 && ge::GRAPH_SUCCESS != CheckSparseSeqLens()){
1278+ return ge::GRAPH_FAILED;
1279+ }
1280+ 
1281+ return ge::GRAPH_SUCCESS;
1282+}
1283+ 
1284+ge::graphStatus SFAATilingCheck::CheckFeatureAntiquantShape() const
1285+{
1286+ OPS_ERR_IF(bSize_ <= 0,
1287+ OPS_LOG_E(opName_, "batch_size should be greater than 0, but got %u", bSize_),
1288+ return ge::GRAPH_FAILED);
1289+
1290+ OPS_ERR_IF(qTSize_ <= 0 && (qLayout_ == SFAALayout::TND),
1291+ OPS_LOG_E(opName_, "T_size of query should be greater than 0, but got %u", qTSize_),
1292+ return ge::GRAPH_FAILED);
1293+ 
1294+ OPS_ERR_IF(n1Size_ <= 0,
1295+ OPS_LOG_E(opName_, "q_head_num should be greater than 0, but got %u", n1Size_),
1296+ return ge::GRAPH_FAILED);
1297+
1298+ OPS_ERR_IF(n1Size_ % n2Size_ != 0,
1299+ OPS_LOG_E(opName_, "q_head_num(%u) must be divisible by kv_head_num(%u)", n1Size_, n2Size_),
1300+ return ge::GRAPH_FAILED);
1301+ 
1302+ return ge::GRAPH_SUCCESS;
1303+}
1304+ 
1305+ge::graphStatus SFAATilingCheck::CheckFeatureMlaAntiquantShape() const
1306+{
1307+ OPS_ERR_IF(n2Size_ != 1,
1308+ OPS_LOG_E(opName_, "kv_head_num should be 1, but got %u", n2Size_),
1309+ return ge::GRAPH_FAILED);
1310+ 
1311+ std::vector<uint32_t> gSizeSupportList = {1, 2, 4, 8, 16, 32, 64, 128};
1312+ OPS_ERR_IF(std::find(gSizeSupportList.begin(), gSizeSupportList.end(), gSize_) == gSizeSupportList.end(),
1313+ OPS_LOG_E(opName_, "group num should be in 1, 2, 4, 8, 16, 32, 64, 128, but got %u", gSize_),
1314+ return ge::GRAPH_FAILED);
1315+ 
1316+ OPS_ERR_IF(qkHeadDim_ != 576,
1317+ OPS_LOG_E(opName_, "qk_head_dim only support 576, but got %u", qkHeadDim_),
1318+ return ge::GRAPH_FAILED);
1319+ 
1320+ return CheckFeatureAntiquantShape();
1321+}
1322+ 
1323+ge::graphStatus SFAATilingCheck::CheckFeatureGqaAntiquantShape() const
1324+{
1325+ OPS_ERR_IF(n2Size_ <= 0,
1326+ OPS_LOG_E(opName_, "kv_head_num should be greater than 0, but got %u", n2Size_),
1327+ return ge::GRAPH_FAILED);
1328+ 
1329+ OPS_ERR_IF(gSize_ > 10,
1330+ OPS_LOG_E(opName_, "group num should not be greater than 10, but got %u", gSize_),
1331+ return ge::GRAPH_FAILED);
1332+ 
1333+ OPS_ERR_IF(qkHeadDim_ != 128,
1334+ OPS_LOG_E(opName_, "qk_head_dim only support 128, but got %u", qkHeadDim_),
1335+ return ge::GRAPH_FAILED);
1336+ 
1337+ return CheckFeatureAntiquantShape();
1338+}
1339+ 
1340+ 
1341+ge::graphStatus SFAATilingCheck::CheckFeatureAntiquantLayout() const
1342+{
1343+ const std::vector<std::string> layoutSupportList = {
1344+ "BSND",
1345+ // "TND"
1346+ };
1347+ std::string layoutQuery = opParamInfo_.layoutQuery;
1348+ OPS_ERR_IF(std::find(layoutSupportList.begin(), layoutSupportList.end(), layoutQuery) == layoutSupportList.end(),
1349+ OPS_LOG_E(opName_, "layoutQuery only supports BSND, but got %s", layoutQuery.c_str()),
1350+ return ge::GRAPH_FAILED);
1351+ return ge::GRAPH_SUCCESS;
1352+}
1353+ 
1354+ge::graphStatus SFAATilingCheck::CheckFeatureAntiquantDtype() const
1355+{
1356+ OPS_ERR_IF(inputQType_ != ge::DT_BF16 && inputQType_ != ge::DT_FLOAT16,
1357+ OPS_LOG_E(opName_, "query dtype only support %s and %s, but got %s",
1358+ SFAADataTypeToSerialString(ge::DT_BF16).c_str(), SFAADataTypeToSerialString(ge::DT_FLOAT16).c_str(),
1359+ SFAADataTypeToSerialString(inputQType_).c_str()),
1360+ return ge::GRAPH_FAILED);
1361+ OPS_ERR_IF(inputKvType_ != ge::DT_INT8,
1362+ OPS_LOG_E(opName_, "key and value dtype only support %s, but got %s",
1363+ SFAADataTypeToSerialString(ge::DT_INT8).c_str(), SFAADataTypeToSerialString(inputKvType_).c_str()),
1364+ return ge::GRAPH_FAILED);
1365+ return ge::GRAPH_SUCCESS;
1366+}
1367+ 
1368+ge::graphStatus SFAATilingCheck::CheckFeatureMlaAntiquantAttr() const
1369+{
1370+ OPS_ERR_IF(attentionMode_ != 2, // 2:MLA-absorb
1371+ OPS_LOG_E(opName_, "attention_mode should be 2(MLA-absorb), but got %u",
1372+ attentionMode_),
1373+ return ge::GRAPH_FAILED);
1374+ 
1375+ OPS_ERR_IF(keyQuantMode_ != 2, // 2:per-tile
1376+ OPS_LOG_E(opName_, "key_quant_mode should be 2(per-tile), but got %u",
1377+ keyQuantMode_),
1378+ return ge::GRAPH_FAILED);
1379+ 
1380+ OPS_ERR_IF(valueQuantMode_ != 2, // 2:per-tile
1381+ OPS_LOG_E(opName_, "value_quant_mode should be 2(per-tile), but got %u",
1382+ valueQuantMode_),
1383+ return ge::GRAPH_FAILED);
1384+ 
1385+ OPS_ERR_IF(quantScaleRepoMode_ != 1, // 1:combine
1386+ OPS_LOG_E(opName_, "quant_scale_repo_mode should be 1(combine), but got %u",
1387+ quantScaleRepoMode_),
1388+ return ge::GRAPH_FAILED);
1389+ 
1390+ OPS_ERR_IF(tileSize_ != 128, // 128:当前不泛化
1391+ OPS_LOG_E(opName_, "tile_size should be 128, but got %u",
1392+ tileSize_),
1393+ return ge::GRAPH_FAILED);
1394+ 
1395+ OPS_ERR_IF(ropeHeadDim_ != 64, // 64:当前不泛化
1396+ OPS_LOG_E(opName_, "rope_head_dim should be 64, but got %u",
1397+ ropeHeadDim_),
1398+ return ge::GRAPH_FAILED);
1399+ 
1400+ return ge::GRAPH_SUCCESS;
1401+}
1402+ 
1403+ge::graphStatus SFAATilingCheck::CheckFeatureGqaAntiquantAttr() const
1404+{
1405+ OPS_ERR_IF(attentionMode_ != 0, // 0:GQA/MHA
1406+ OPS_LOG_E(opName_, "attention_mode should be 0(GQA/MHA), but got %u",
1407+ attentionMode_),
1408+ return ge::GRAPH_FAILED);
1409+ 
1410+ OPS_ERR_IF(keyQuantMode_ != 0, // 0:per-channel
1411+ OPS_LOG_E(opName_, "key_quant_mode should be 0(per-channel), but got %u",
1412+ keyQuantMode_),
1413+ return ge::GRAPH_FAILED);
1414+ 
1415+ OPS_ERR_IF(valueQuantMode_ != 0, // 0:per-channel
1416+ OPS_LOG_E(opName_, "value_quant_mode should be 0(per-channel), but got %u",
1417+ valueQuantMode_),
1418+ return ge::GRAPH_FAILED);
1419+ 
1420+ OPS_ERR_IF(quantScaleRepoMode_ != 0, // 0:seprate
1421+ OPS_LOG_E(opName_, "quant_scale_repo_mode should be 0(seprate), but got %u",
1422+ quantScaleRepoMode_),
1423+ return ge::GRAPH_FAILED);
1424+ 
1425+ OPS_ERR_IF(sparseShardSize_ != s1Size_,
1426+ OPS_LOG_E(opName_, "sparse_shard_size should be equal %u, but got %u",
1427+ s1Size_, sparseShardSize_),
1428+ return ge::GRAPH_FAILED);
1429+
1430+ OPS_ERR_IF(s1Size_ > 6,
1431+ OPS_LOG_E(opName_, "s1Size_ should be smaller than 6, but got %u",
1432+ s1Size_),
1433+ return ge::GRAPH_FAILED);
1434+ 
1435+ return ge::GRAPH_SUCCESS;
1436+}
1437+ 
1438+ge::graphStatus SFAATilingCheck::CheckFeatureAntiquantPa() const
1439+{
1440+ if (kvStorageMode_ != KvStorageMode::PAGE_ATTENTION) {
1441+ return ge::GRAPH_SUCCESS;
1442+ }
1443+ 
1444+ OPS_ERR_IF(blockSize_ <= 0 || blockSize_ > static_cast<int32_t>(MAX_BLOCK_SIZE),
1445+ OPS_LOG_E(opName_, "when page attention is enabled, block_size(%d) should be in range (0, %u].",
1446+ blockSize_, MAX_BLOCK_SIZE), return ge::GRAPH_FAILED);
1447+
1448+ OPS_ERR_IF(blockSize_ % 16 > 0,
1449+ OPS_LOG_E(opName_, "when page attention is enabled, block_size(%d) should be 16-aligned.",
1450+ blockSize_), return ge::GRAPH_FAILED);
1451+
1452+ OPS_ERR_IF(blockSize_ % sparseBlockSize_ > 0,
1453+ OPS_LOG_E(opName_,
1454+ "when page attention is enabled, block_size(%d) must be divided by sparse_block_size(%d), but now the remainder is %d.",
1455+ blockSize_, sparseBlockSize_, blockSize_ % sparseBlockSize_), return ge::GRAPH_FAILED);
1456+ 
1457+ return ge::GRAPH_SUCCESS;
1458+}
1459+ 
1460+ge::graphStatus SFAATilingCheck::CheckFeatureMlaAntiquant() const
1461+{
1462+ if (ge::GRAPH_SUCCESS != CheckFeatureMlaAntiquantShape() ||
1463+ ge::GRAPH_SUCCESS != CheckFeatureAntiquantLayout() ||
1464+ ge::GRAPH_SUCCESS != CheckFeatureAntiquantDtype() ||
1465+ ge::GRAPH_SUCCESS != CheckFeatureAntiquantPa() ||
1466+ ge::GRAPH_SUCCESS != CheckFeatureMlaAntiquantAttr()) {
1467+ return ge::GRAPH_FAILED;
1468+ }
1469+ return ge::GRAPH_SUCCESS;
1470+}
1471+ 
1472+ge::graphStatus SFAATilingCheck::CheckFeatureGqaAntiquant() const
1473+{
1474+ if (ge::GRAPH_SUCCESS != CheckFeatureGqaAntiquantShape() ||
1475+ ge::GRAPH_SUCCESS != CheckFeatureAntiquantLayout() ||
1476+ ge::GRAPH_SUCCESS != CheckFeatureAntiquantDtype() ||
1477+ ge::GRAPH_SUCCESS != CheckFeatureAntiquantPa() ||
1478+ ge::GRAPH_SUCCESS != CheckFeatureGqaAntiquantAttr()) {
1479+ return ge::GRAPH_FAILED;
1480+ }
1481+ return ge::GRAPH_SUCCESS;
1482+}
1483+ 
1484+ge::graphStatus SFAATilingCheck::CheckFeature() const
1485+{
1486+ if (attentionMode_ == 0) {
1487+ return CheckFeatureGqaAntiquant();
1488+ } else {
1489+ return CheckFeatureMlaAntiquant();
1490+ }
1491+}
1492+ 
1493+void SFAATilingCheck::Init()
1494+{
1495+ opName_ = sfaaInfo_.opName;
1496+ platformInfo_ = sfaaInfo_.platformInfo;
1497+ opParamInfo_ = sfaaInfo_.opParamInfo;
1498+ socVersion_ = sfaaInfo_.socVersion;
1499+ 
1500+ bSize_ = sfaaInfo_.bSize;
1501+ n1Size_ = sfaaInfo_.n1Size;
1502+ n2Size_ = sfaaInfo_.n2Size;
1503+ s1Size_ = sfaaInfo_.s1Size;
1504+ s2Size_ = sfaaInfo_.s2Size;
1505+ gSize_ = sfaaInfo_.gSize;
1506+ qkHeadDim_ = sfaaInfo_.qkHeadDim;
1507+ vHeadDim_ = sfaaInfo_.vHeadDim;
1508+ ropeHeadDim_ = sfaaInfo_.ropeHeadDim;
1509+ maxBlockNumPerBatch_ = sfaaInfo_.maxBlockNumPerBatch;
1510+ qTSize_ = sfaaInfo_.qTSize;
1511+ kvTSize_ = sfaaInfo_.kvTSize;
1512+ blockSize_ = sfaaInfo_.blockSize;
1513+ sparseBlockCount_ = sfaaInfo_.sparseBlockCount;
1514+ sparseBlockSize_ = sfaaInfo_.sparseBlockSize;
1515+ sparseShardSize_ = sfaaInfo_.sparseShardSize;
1516+ 
1517+ attentionMode_ = sfaaInfo_.attentionMode;
1518+ keyQuantMode_ = sfaaInfo_.keyQuantMode;
1519+ valueQuantMode_ = sfaaInfo_.valueQuantMode;
1520+ quantScaleRepoMode_ = sfaaInfo_.quantScaleRepoMode;
1521+ tileSize_ = sfaaInfo_.tileSize;
1522+ 
1523+ inputQType_ = sfaaInfo_.inputQType;
1524+ inputKvType_ = sfaaInfo_.inputKvType;
1525+ outputType_ = sfaaInfo_.outputType;
1526+ 
1527+ qLayout_ = sfaaInfo_.qLayout;
1528+ topkLayout_ = sfaaInfo_.topkLayout;
1529+ kvLayout_ = sfaaInfo_.kvLayout;
1530+ outLayout_ = sfaaInfo_.outLayout;
1531+ 
1532+ kvStorageMode_ = sfaaInfo_.kvStorageMode;
1533+ l2CacheSize_ = sfaaInfo_.l2CacheSize;
1534+}
1535+ 
1536+ge::graphStatus SFAATilingCheck::Process()
1537+{
1538+ Init();
1539+ if (CheckSinglePara() != ge::GRAPH_SUCCESS ||
1540+ CheckParaExistence() != ge::GRAPH_SUCCESS ||
1541+ CheckFeature() != ge::GRAPH_SUCCESS ||
1542+ CheckMultiParaConsistency() != ge::GRAPH_SUCCESS) {
1543+ return ge::GRAPH_FAILED;
1544+ }
1545+ return ge::GRAPH_SUCCESS;
1546+}
1547+ 
1548+bool SFAAInfoParser::HasAxis(const SFAAAxis &axis, const SFAALayout &layout, const gert::Shape &shape) const
1549+{
1550+ const auto& layoutIt = SFAA_LAYOUT_AXIS_MAP.find(layout);
1551+ if (layoutIt == SFAA_LAYOUT_AXIS_MAP.end()) {
1552+ return false;
1553+ }
1554+ 
1555+ const std::vector<SFAAAxis>& axes = layoutIt->second;
1556+ const auto& axisIt = std::find(axes.begin(), axes.end(), axis);
1557+ if (axisIt == axes.end()) {
1558+ return false;
1559+ }
1560+ 
1561+ const auto& dimIt = SFAA_LAYOUT_DIM_MAP.find(layout);
1562+ if (dimIt == SFAA_LAYOUT_DIM_MAP.end() || dimIt->second != shape.GetDimNum()) {
1563+ return false;
1564+ }
1565+ return true;
1566+}
1567+ 
1568+size_t SFAAInfoParser::GetAxisIdx(const SFAAAxis &axis, const SFAALayout &layout) const
1569+{
1570+ const std::vector<SFAAAxis>& axes = SFAA_LAYOUT_AXIS_MAP.find(layout)->second;
1571+ const auto& axisIt = std::find(axes.begin(), axes.end(), axis);
1572+ return std::distance(axes.begin(), axisIt);
1573+}
1574+ 
1575+uint32_t SFAAInfoParser::GetAxisNum(const gert::Shape &shape, const SFAAAxis &axis, const SFAALayout &layout) const
1576+{
1577+ return HasAxis(axis, layout, shape) ? shape.GetDim(GetAxisIdx(axis, layout)) : invalidDimValue_;
1578+}
1579+ 
1580+ge::graphStatus SFAAInfoParser::CheckRequiredInOutExistence() const
1581+{
1582+ OPS_ERR_IF(opParamInfo_.query.shape == nullptr, OPS_LOG_E(opName_, "Shape of tensor query is nullptr"),
1583+ return ge::GRAPH_FAILED);
1584+ OPS_ERR_IF(opParamInfo_.query.desc == nullptr, OPS_LOG_E(opName_, "Desc of tensor query is nullptr"),
1585+ return ge::GRAPH_FAILED);
1586+ OPS_ERR_IF(opParamInfo_.key.shape == nullptr, OPS_LOG_E(opName_, "Shape of tensor k is nullptr"),
1587+ return ge::GRAPH_FAILED);
1588+ OPS_ERR_IF(opParamInfo_.key.desc == nullptr, OPS_LOG_E(opName_, "Desc of tensor k is nullptr"),
1589+ return ge::GRAPH_FAILED);
1590+ OPS_ERR_IF(opParamInfo_.value.shape == nullptr, OPS_LOG_E(opName_, "Shape of tensor value is nullptr"),
1591+ return ge::GRAPH_FAILED);
1592+ OPS_ERR_IF(opParamInfo_.value.desc == nullptr, OPS_LOG_E(opName_, "Desc of tensor value is nullptr"),
1593+ return ge::GRAPH_FAILED);
1594+ OPS_ERR_IF(opParamInfo_.sparseIndices.shape == nullptr,
1595+ OPS_LOG_E(opName_, "Shape of tensor sparseIndices is nullptr"),
1596+ return ge::GRAPH_FAILED);
1597+ OPS_ERR_IF(opParamInfo_.sparseIndices.desc == nullptr,
1598+ OPS_LOG_E(opName_, "Desc of tensor sparseIndices is nullptr"),
1599+ return ge::GRAPH_FAILED);
1600+ OPS_ERR_IF(opParamInfo_.attenOut.shape == nullptr, OPS_LOG_E(opName_, "Shape of tensor output is nullptr"),
1601+ return ge::GRAPH_FAILED);
1602+ OPS_ERR_IF(opParamInfo_.attenOut.desc == nullptr, OPS_LOG_E(opName_, "Desc of tensor output is nullptr"),
1603+ return ge::GRAPH_FAILED);
1604+ 
1605+ return ge::GRAPH_SUCCESS;
1606+}
1607+ 
1608+ge::graphStatus SFAAInfoParser::CheckRequiredAttrExistence() const
1609+{
1610+ OPS_ERR_IF(opParamInfo_.layoutQuery == nullptr, OPS_LOG_E(opName_, "attr layoutQuery is nullptr"),
1611+ return ge::GRAPH_FAILED);
1612+ OPS_ERR_IF(opParamInfo_.layoutKV == nullptr, OPS_LOG_E(opName_, "attr layoutKV is nullptr"),
1613+ return ge::GRAPH_FAILED);
1614+ OPS_ERR_IF(opParamInfo_.sparseBlockSize == nullptr, OPS_LOG_E(opName_, "attr sparseBlockSize is nullptr"),
1615+ return ge::GRAPH_FAILED);
1616+ OPS_ERR_IF(opParamInfo_.scaleValue == nullptr, OPS_LOG_E(opName_, "attr scaleValue is nullptr"),
1617+ return ge::GRAPH_FAILED);
1618+ OPS_ERR_IF(opParamInfo_.sparseMode == nullptr, OPS_LOG_E(opName_, "attr sparseMode is nullptr"),
1619+ return ge::GRAPH_FAILED);
1620+ return ge::GRAPH_SUCCESS;
1621+}
1622+ 
1623+ge::graphStatus SFAAInfoParser::CheckRequiredParaExistence() const
1624+{
1625+ if (CheckRequiredInOutExistence() != ge::GRAPH_SUCCESS ||
1626+ CheckRequiredAttrExistence() != ge::GRAPH_SUCCESS) {
1627+ return ge::GRAPH_FAILED;
1628+ }
1629+ 
1630+ return ge::GRAPH_SUCCESS;
1631+}
1632+ 
1633+ge::graphStatus SFAAInfoParser::GetActualSeqLenSize(uint32_t &size, const gert::Tensor *tensor,
1634+ SFAALayout &layout, const std::string &name)
1635+{
1636+ if ((tensor == nullptr)) {
1637+ OPS_LOG_E(opName_, "when layout of query is %s, %s must be provided.",
1638+ SFAALayoutToSerialString(layout).c_str(), name.c_str());
1639+ return ge::GRAPH_FAILED;
1640+ }
1641+ int64_t shapeSize = tensor->GetShapeSize();
1642+ if (shapeSize <= 0) {
1643+ OPS_LOG_E(opName_, "the shape size of %s is %ld, it should be greater than 0.",
1644+ name.c_str(), shapeSize);
1645+ return ge::GRAPH_FAILED;
1646+ }
1647+ size = static_cast<uint32_t>(shapeSize);
1648+ return ge::GRAPH_SUCCESS;
1649+}
1650+ 
1651+ge::graphStatus SFAAInfoParser::GetActualSeqLenQSize(uint32_t &size)
1652+{
1653+ return GetActualSeqLenSize(size, opParamInfo_.actualSeqLengthsQ.tensor, qLayout_, "actualSeqLengthsQ");
1654+}
1655+ 
1656+ge::graphStatus SFAAInfoParser::GetOpName()
1657+{
1658+ if (context_->GetNodeName() == nullptr) {
1659+ OPS_LOG_E("SparseFlashAttentionAntiquant", "opName got from TilingContext is nullptr");
1660+ return ge::GRAPH_FAILED;
1661+ }
1662+ opName_ = context_->GetNodeName();
1663+ return ge::GRAPH_SUCCESS;
1664+}
1665+ 
1666+ge::graphStatus SFAAInfoParser::GetNpuInfo()
1667+{
1668+ platformInfo_ = context_->GetPlatformInfo();
1669+ OPS_ERR_IF(platformInfo_ == nullptr,
1670+ OPS_REPORT_VECTOR_INNER_ERR(opName_, "GetPlatformInfo is nullptr."), return ge::GRAPH_FAILED);
1671+ 
1672+ auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfo_);
1673+ uint32_t aivNum = ascendcPlatform.GetCoreNumAiv();
1674+ uint32_t aicNum = ascendcPlatform.GetCoreNumAic();
1675+ OPS_ERR_IF(aicNum == 0 || aivNum == 0,
1676+ OPS_REPORT_VECTOR_INNER_ERR(opName_, "num of core obtained is 0."), return GRAPH_FAILED);
1677+ 
1678+ ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::L2, l2CacheSize_);
1679+ 
1680+ return ge::GRAPH_SUCCESS;
1681+}
1682+ 
1683+void SFAAInfoParser::GetOptionalInputParaInfo()
1684+{
1685+ opParamInfo_.blockTable.tensor = context_->GetOptionalInputTensor(BLOCK_TABLE_INPUT_INDEX);
1686+ opParamInfo_.blockTable.desc = context_->GetOptionalInputDesc(BLOCK_TABLE_INPUT_INDEX);
1687+ opParamInfo_.actualSeqLengthsQ.tensor = context_->GetOptionalInputTensor(ACT_SEQ_LEN_Q_INPUT_INDEX);
1688+ opParamInfo_.actualSeqLengthsQ.desc = context_->GetOptionalInputDesc(ACT_SEQ_LEN_Q_INPUT_INDEX);
1689+ opParamInfo_.actualSeqLengths.tensor = context_->GetOptionalInputTensor(ACT_SEQ_LEN_KV_INPUT_INDEX);
1690+ opParamInfo_.actualSeqLengths.desc = context_->GetOptionalInputDesc(ACT_SEQ_LEN_KV_INPUT_INDEX);
1691+ opParamInfo_.sparseSeqLengths.tensor = context_->GetOptionalInputTensor(SPARSE_SEQ_LEN_KV_INPUT_INDEX);
1692+ opParamInfo_.sparseSeqLengths.desc = context_->GetOptionalInputDesc(SPARSE_SEQ_LEN_KV_INPUT_INDEX);
1693+ opParamInfo_.keyDequantScale.tensor = context_->GetOptionalInputTensor(KEY_DEQUANT_SCALE_INPUT_INDEX);
1694+ opParamInfo_.valueDequantScale.tensor = context_->GetOptionalInputTensor(VALUE_DEQUANT_SCALE_INPUT_INDEX);
1695+}
1696+ 
1697+void SFAAInfoParser::GetInputParaInfo()
1698+{
1699+ opParamInfo_.query.desc = context_->GetInputDesc(QUERY_INPUT_INDEX);
1700+ opParamInfo_.query.shape = context_->GetInputShape(QUERY_INPUT_INDEX);
1701+ opParamInfo_.key.desc = context_->GetInputDesc(KEY_INPUT_INDEX);
1702+ opParamInfo_.key.shape = context_->GetInputShape(KEY_INPUT_INDEX);
1703+ opParamInfo_.value.desc = context_->GetInputDesc(VALUE_INPUT_INDEX);
1704+ opParamInfo_.value.shape = context_->GetInputShape(VALUE_INPUT_INDEX);
1705+ opParamInfo_.sparseIndices.desc = context_->GetInputDesc(SPARSE_INDICES_INPUT_INDEX);
1706+ opParamInfo_.sparseIndices.shape = context_->GetInputShape(SPARSE_INDICES_INPUT_INDEX);
1707+ GetOptionalInputParaInfo();
1708+}
1709+ 
1710+void SFAAInfoParser::GetOutputParaInfo()
1711+{
1712+ opParamInfo_.attenOut.desc = context_->GetOutputDesc(OUTPUT_INDEX);
1713+ opParamInfo_.attenOut.shape = context_->GetOutputShape(OUTPUT_INDEX);
1714+}
1715+ 
1716+ge::graphStatus SFAAInfoParser::GetAttrParaInfo()
1717+{
1718+ auto attrs = context_->GetAttrs();
1719+ OPS_ERR_IF(attrs == nullptr, OPS_REPORT_VECTOR_INNER_ERR(context_->GetNodeName(), "attrs got from ge is nullptr"),
1720+ return ge::GRAPH_FAILED);
1721+ 
1722+ opParamInfo_.layoutQuery = attrs->GetStr(LAYOUT_QUERY_ATTR_INDEX);
1723+ opParamInfo_.layoutKV = attrs->GetStr(LAYOUT_KV_ATTR_INDEX);
1724+ opParamInfo_.sparseBlockSize = attrs->GetAttrPointer<int64_t>(SPARSE_BLOCK_SIZE_ATTR_INDEX);
1725+ opParamInfo_.scaleValue = attrs->GetAttrPointer<float>(SCALE_VALUE_ATTR_INDEX);
1726+ opParamInfo_.sparseMode = attrs->GetAttrPointer<int64_t>(SPARSE_MODE_ATTR_INDEX);
1727+ opParamInfo_.keyQuantMode = attrs->GetAttrPointer<int64_t>(KEY_QUANT_MODE_ATTR_INDEX);
1728+ opParamInfo_.valueQuantMode = attrs->GetAttrPointer<int64_t>(VALUE_QUANT_MODE_ATTR_INDEX);
1729+ opParamInfo_.attentionMode = attrs->GetAttrPointer<int64_t>(ATTENTION_MODE_ATTR_INDEX);
1730+ opParamInfo_.quantScaleRepoMode = attrs->GetAttrPointer<int64_t>(QUANT_SCALE_REPO_MODE_ATTR_INDEX);
1731+ opParamInfo_.tileSize = attrs->GetAttrPointer<int64_t>(TILE_SIZE_ATTR_INDEX);
1732+ opParamInfo_.ropeHeadDim = attrs->GetAttrPointer<int64_t>(ROPE_HEAD_DIM_ATTR_INDEX);
1733+ opParamInfo_.sparseShardSize = attrs->GetAttrPointer<int64_t>(SPARSE_SHARD_SIZE_ATTR_INDEX);
1734+ 
1735+ return ge::GRAPH_SUCCESS;
1736+}
1737+ 
1738+ge::graphStatus SFAAInfoParser::GetOpParaInfo()
1739+{
1740+ GetInputParaInfo();
1741+ GetOutputParaInfo();
1742+ if (ge::GRAPH_SUCCESS != GetAttrParaInfo()) {
1743+ return ge::GRAPH_FAILED;
1744+ }
1745+ return ge::GRAPH_SUCCESS;
1746+}
1747+ 
1748+ge::graphStatus SFAAInfoParser::GetInOutDataType()
1749+{
1750+ inputQType_ = opParamInfo_.query.desc->GetDataType();
1751+ inputKvType_ = opParamInfo_.key.desc->GetDataType();
1752+ outputType_ = opParamInfo_.attenOut.desc->GetDataType();
1753+ return ge::GRAPH_SUCCESS;
1754+}
1755+ 
1756+ge::graphStatus SFAAInfoParser::GetBatchSize()
1757+{
1758+ // 获取B基准值
1759+ // 1、非TND时, 以query的batch_size维度为基准;
1760+ // 2、TND时, actual_seq_lens_q必须传入, 以actual_seq_lens_q数组的长度为B轴大小
1761+ if (qLayout_ == SFAALayout::TND) {
1762+ return GetActualSeqLenQSize(bSize_);
1763+ } else { // BSND
1764+ bSize_ = GetAxisNum(queryShape_, SFAAAxis::B, qLayout_);
1765+ return ge::GRAPH_SUCCESS;
1766+ }
1767+}
1768+ 
1769+ge::graphStatus SFAAInfoParser::GetQTSize()
1770+{
1771+ // 获取query的T基准值
1772+ // 1、非TND时, 以query的batch_size维度为基准;
1773+ // 2、TND时, actual_seq_lens_q必须传入, 以actual_seq_lens_q数组的长度为B轴大小
1774+ qTSize_ = (qLayout_ == SFAALayout::TND) ? GetAxisNum(queryShape_, SFAAAxis::T, qLayout_) : 0;
1775+ return ge::GRAPH_SUCCESS;
1776+}
1777+ 
1778+ge::graphStatus SFAAInfoParser::GetKVTSize()
1779+{
1780+ // 获取query的T基准值
1781+ // 1、非TND时, 以key的batch_size维度为基准;
1782+ // 2、TND时, actual_seq_lens_q必须传入, 以actual_seq_lens_q数组的长度为B轴大小
1783+ kvTSize_ = (kvLayout_ == SFAALayout::TND) ? GetAxisNum(keyShape_, SFAAAxis::T, kvLayout_) : 0;
1784+ return ge::GRAPH_SUCCESS;
1785+}
1786+ 
1787+ge::graphStatus SFAAInfoParser::GetQkHeadDim()
1788+{
1789+ // 获取qkHeadDim基准值
1790+ // 以query的D维度为基准
1791+ qkHeadDim_ = GetAxisNum(queryShape_, SFAAAxis::D, qLayout_);
1792+ return ge::GRAPH_SUCCESS;
1793+}
1794+ 
1795+ge::graphStatus SFAAInfoParser::GetS1Size()
1796+{
1797+ // 获取S1基准值
1798+ // 1、非TND时, 以query的S维度为基准;
1799+ // 2、TND时, actual_seq_lens_q必须传入, 以actual_seq_lens_q数组中的最大值为基准
1800+ if (qLayout_ == SFAALayout::TND) {
1801+ s1Size_ = GetAxisNum(queryShape_, SFAAAxis::T, qLayout_);
1802+ return ge::GRAPH_SUCCESS;
1803+ } else { // BSND
1804+ s1Size_ = GetAxisNum(queryShape_, SFAAAxis::S, qLayout_);
1805+ }
1806+ return ge::GRAPH_SUCCESS;
1807+}
1808+ 
1809+ge::graphStatus SFAAInfoParser::GetKvStorageMode()
1810+{
1811+ if (kvLayout_ == SFAALayout::PA_BSND || kvLayout_ == SFAALayout::PA_BNSD || kvLayout_ == SFAALayout::PA_NZ) {
1812+ kvStorageMode_ = KvStorageMode::PAGE_ATTENTION;
1813+ } else {
1814+ kvStorageMode_ = KvStorageMode::BATCH_CONTINUOUS;
1815+ }
1816+ // kv存储模式基准值
1817+ return ge::GRAPH_SUCCESS;
1818+}
1819+ 
1820+ge::graphStatus SFAAInfoParser::GetKvLayout()
1821+{
1822+ const map<string, SFAALayout> layoutKVMap = {
1823+ {"BSND", SFAALayout::BSND},
1824+ {"PA_BSND", SFAALayout::PA_BSND},
1825+ {"PA_BNSD", SFAALayout::PA_BNSD},
1826+ {"PA_NZ", SFAALayout::PA_NZ},
1827+ {"TND", SFAALayout::TND}
1828+ };
1829+ 
1830+ std::string layout(opParamInfo_.layoutKV);
1831+ auto it = layoutKVMap.find(layout);
1832+ if (it != layoutKVMap.end()) {
1833+ kvLayout_ = it->second;
1834+ } else {
1835+ OPS_LOG_E(opName_, "layoutKV is %s, it is unsupported.", layout.c_str());
1836+ return ge::GRAPH_FAILED;
1837+ }
1838+ if (qLayout_ != SFAALayout::BSND ) {
1839+ OPS_LOG_E(opName_, "layoutQ supports BSND, but now is not.");
1840+ return ge::GRAPH_FAILED;
1841+ }
1842+ if (kvLayout_ != SFAALayout::PA_NZ ) {
1843+ OPS_LOG_E(opName_, "layoutKV supports PA_NZ, but now is %s.", layout.c_str());
1844+ return ge::GRAPH_FAILED;
1845+ }
1846+ if (qLayout_ == SFAALayout::BSND && kvLayout_ != SFAALayout::PA_NZ) {
1847+ OPS_LOG_E(opName_, "When layoutQ is TND, layoutKV supports PA_BSND, PA_BNSD and PA_NZ, but now is %s.", layout.c_str());
1848+ return ge::GRAPH_FAILED;
1849+ }
1850+ uint32_t keyDimNum = opParamInfo_.key.shape->GetStorageShape().GetDimNum();
1851+ if ((kvLayout_ == SFAALayout::PA_BSND || kvLayout_ == SFAALayout::PA_BNSD) && keyDimNum != 4U) {
1852+ OPS_LOG_E(opName_, "When layoutKV is PA_BSND or PA_BSND, kvDimNum must be 4, but now is %d.", keyDimNum);
1853+ return ge::GRAPH_FAILED;
1854+ } else if (kvLayout_ == SFAALayout::PA_BSND && keyDimNum != 5U) {
1855+ OPS_LOG_E(opName_, "When layoutKV is PA_NZ, kvDimNum must be 5, but now is %d.", keyDimNum);
1856+ return ge::GRAPH_FAILED;
1857+ }
1858+ return ge::GRAPH_SUCCESS;
1859+}
1860+ 
1861+ge::graphStatus SFAAInfoParser::GetS2SizeForBatchContinuous()
1862+{
1863+ if (kvLayout_ != SFAALayout::BSND) {
1864+ OPS_LOG_E(opName_, "the layout of key is %s, it is unsupported.", SFAALayoutToSerialString(kvLayout_).c_str());
1865+ return ge::GRAPH_FAILED;
1866+ } else if (kvLayout_ == SFAALayout::BSND) { // BSND
1867+ s2Size_ = GetAxisNum(keyShape_, SFAAAxis::S, kvLayout_);
1868+ }
1869+ return ge::GRAPH_SUCCESS;
1870+}
1871+ 
1872+ge::graphStatus SFAAInfoParser::GetMaxBlockNumPerBatch()
1873+{
1874+ if (opParamInfo_.blockTable.tensor == nullptr) {
1875+ OPS_LOG_E(opName_, "the layout_kv is %s, blockTable must be provided.",
1876+ SFAALayoutToSerialString(kvLayout_).c_str());
1877+ return ge::GRAPH_FAILED;
1878+ }
1879+ uint32_t dimNum = opParamInfo_.blockTable.tensor->GetStorageShape().GetDimNum();
1880+ if (dimNum != DIM_NUM_TWO) {
1881+ OPS_LOG_E(opName_, "the dim num of block_table is %u, it should be %u.", dimNum, DIM_NUM_TWO);
1882+ return ge::GRAPH_FAILED;
1883+ }
1884+ if (opParamInfo_.blockTable.tensor->GetStorageShape().GetDim(1) <= 0) {
1885+ OPS_LOG_E(opName_, "%s's second dimension(%ld) should be greater than 0",
1886+ BLOCK_TABLE_NAME.c_str(), opParamInfo_.blockTable.tensor->GetStorageShape().GetDim(1));
1887+ return ge::GRAPH_FAILED;
1888+ }
1889+ maxBlockNumPerBatch_ = opParamInfo_.blockTable.tensor->GetStorageShape().GetDim(1);
1890+ return ge::GRAPH_SUCCESS;
1891+}
1892+ 
1893+ge::graphStatus SFAAInfoParser::GetBlockSize()
1894+{
1895+ blockSize_ = GetAxisNum(keyShape_, SFAAAxis::Bs, kvLayout_);
1896+ return ge::GRAPH_SUCCESS;
1897+}
1898+ 
1899+ge::graphStatus SFAAInfoParser::GetSparseBlockCount()
1900+{
1901+ sparseBlockCount_ = GetAxisNum(sparseIndicesShape_, SFAAAxis::K, qLayout_);
1902+ 
1903+ return ge::GRAPH_SUCCESS;
1904+}
1905+ 
1906+ge::graphStatus SFAAInfoParser::GetS2SizeForPageAttention()
1907+{
1908+ if (GetMaxBlockNumPerBatch() != ge::GRAPH_SUCCESS || GetBlockSize() != ge::GRAPH_SUCCESS) {
1909+ return ge::GRAPH_FAILED;
1910+ }
1911+ s2Size_ = maxBlockNumPerBatch_ * blockSize_;
1912+ return ge::GRAPH_SUCCESS;
1913+}
1914+ 
1915+ge::graphStatus SFAAInfoParser::GetS2Size()
1916+{
1917+ // 获取S2基准值
1918+ // 1、BATCH_CONTINUOUS时, 从key的S轴获取
1919+ // 2、PAGE_ATTENTION时, S2 = block_table.dim1 * block_size
1920+ if (kvStorageMode_ == KvStorageMode::BATCH_CONTINUOUS) {
1921+ return GetS2SizeForBatchContinuous();
1922+ }
1923+ return GetS2SizeForPageAttention();
1924+}
1925+ 
1926+ge::graphStatus SFAAInfoParser::GetValueHeadDim()
1927+{
1928+ // 获取vHeadDim基准值
1929+ // 以value的D维度为基准
1930+ vHeadDim_ = GetAxisNum(valueShape_, SFAAAxis::D, kvLayout_);
1931+ return ge::GRAPH_SUCCESS;
1932+}
1933+ 
1934+ge::graphStatus SFAAInfoParser::GetQueryAndOutLayout()
1935+{
1936+ // 获取query和attentionOut的Layout基准值
1937+ // layoutQuery: {qLayout, outLayout}
1938+ const map<string, pair<SFAALayout, SFAALayout>> layoutMap = {
1939+ {"BSND", {SFAALayout::BSND, SFAALayout::BSND}},
1940+ {"TND", {SFAALayout::TND, SFAALayout::TND }},
1941+ };
1942+ 
1943+ std::string layout(opParamInfo_.layoutQuery);
1944+ auto it = layoutMap.find(layout);
1945+ if (it != layoutMap.end()) {
1946+ qLayout_ = it->second.first;
1947+ outLayout_ = it->second.second;
1948+ } else {
1949+ OPS_LOG_E(opName_, "layoutQuery is %s, it is unsupported.", layout.c_str());
1950+ return ge::GRAPH_FAILED;
1951+ }
1952+ return ge::GRAPH_SUCCESS;
1953+}
1954+ 
1955+ge::graphStatus SFAAInfoParser::GetTopkLayout()
1956+{
1957+ topkLayout_ = qLayout_;
1958+ return ge::GRAPH_SUCCESS;
1959+}
1960+ 
1961+ge::graphStatus SFAAInfoParser::GetN1Size()
1962+{
1963+ n1Size_ = GetAxisNum(queryShape_, SFAAAxis::N, qLayout_);
1964+ return ge::GRAPH_SUCCESS;
1965+}
1966+ 
1967+ge::graphStatus SFAAInfoParser::GetN2Size()
1968+{
1969+ n2Size_ = GetAxisNum(keyShape_, SFAAAxis::N, kvLayout_);
1970+ return ge::GRAPH_SUCCESS;
1971+}
1972+ 
1973+void SFAAInfoParser::SetSFAAShape()
1974+{
1975+ queryShape_ = opParamInfo_.query.shape->GetStorageShape();
1976+ keyShape_ = opParamInfo_.key.shape->GetStorageShape();
1977+ valueShape_ = opParamInfo_.value.shape->GetStorageShape();
1978+ sparseIndicesShape_ = opParamInfo_.sparseIndices.shape->GetStorageShape();
1979+}
1980+ 
1981+ge::graphStatus SFAAInfoParser::GetGSize()
1982+{
1983+ if (n2Size_ != 0) {
1984+ gSize_ = n1Size_ / n2Size_;
1985+ }
1986+ return ge::GRAPH_SUCCESS;
1987+}
1988+ 
1989+ge::graphStatus SFAAInfoParser::GetActualseqInfo()
1990+{
1991+ maxActualseq_ = static_cast<uint32_t>(s2Size_);
1992+ if (opParamInfo_.actualSeqLengths.tensor != nullptr) {
1993+ actualLenDimsKV_ = opParamInfo_.actualSeqLengths.tensor->GetShapeSize();
1994+ }
1995+ if (opParamInfo_.actualSeqLengthsQ.tensor != nullptr) {
1996+ actualLenDimsQ_ = opParamInfo_.actualSeqLengthsQ.tensor->GetShapeSize();
1997+ }
1998+ if (opParamInfo_.sparseSeqLengths.tensor != nullptr) {
1999+ sparseLenDimsKV_ = opParamInfo_.sparseSeqLengths.tensor->GetShapeSize();
2000+ }
2001+ return ge::GRAPH_SUCCESS;
2002+}
2003+ 
2004+void SFAAInfoParser::GenerateInfo(SFAATilingInfo &sfaaInfo)
2005+{
2006+ sfaaInfo.opName = opName_;
2007+ sfaaInfo.platformInfo = platformInfo_;
2008+ sfaaInfo.opParamInfo = opParamInfo_;
2009+ sfaaInfo.socVersion = socVersion_;
2010+ 
2011+ sfaaInfo.bSize = bSize_;
2012+ sfaaInfo.n1Size = n1Size_;
2013+ sfaaInfo.n2Size = n2Size_;
2014+ sfaaInfo.s1Size = s1Size_;
2015+ sfaaInfo.s2Size = s2Size_;
2016+ sfaaInfo.gSize = gSize_;
2017+ sfaaInfo.qkHeadDim = qkHeadDim_;
2018+ sfaaInfo.vHeadDim = vHeadDim_;
2019+ sfaaInfo.qTSize = qTSize_;
2020+ sfaaInfo.kvTSize = kvTSize_;
2021+ sfaaInfo.sparseBlockSize = *opParamInfo_.sparseBlockSize;
2022+ sfaaInfo.sparseBlockCount = sparseBlockCount_;
2023+ 
2024+ sfaaInfo.inputQType = inputQType_;
2025+ sfaaInfo.inputKvType = inputKvType_;
2026+ sfaaInfo.outputType = outputType_;
2027+ 
2028+ sfaaInfo.kvStorageMode = kvStorageMode_;
2029+ sfaaInfo.l2CacheSize = l2CacheSize_;
2030+ 
2031+ sfaaInfo.totalBlockNum = opParamInfo_.key.shape->GetStorageShape().GetDim(0);
2032+ sfaaInfo.scaleValue = *opParamInfo_.scaleValue;
2033+ sfaaInfo.pageAttentionFlag = (kvStorageMode_ == KvStorageMode::PAGE_ATTENTION);
2034+ sfaaInfo.blockSize = blockSize_;
2035+ sfaaInfo.blockTypeSize = sizeof(float);
2036+ sfaaInfo.maxBlockNumPerBatch = maxBlockNumPerBatch_;
2037+ 
2038+ sfaaInfo.actualLenDimsQ = actualLenDimsQ_;
2039+ sfaaInfo.actualLenDimsKV = actualLenDimsKV_;
2040+ sfaaInfo.sparseLenDimsKV = sparseLenDimsKV_;
2041+ sfaaInfo.maxActualseq = maxActualseq_;
2042+ sfaaInfo.isSameSeqAllKVTensor = isSameSeqAllKVTensor_;
2043+ sfaaInfo.isSameActualseq = isSameActualseq_;
2044+ 
2045+ OPS_ERR_IF(opParamInfo_.sparseBlockSize == nullptr, OPS_LOG_E(opName_, "attr sparseBlockSize is nullptr"), return ge::GRAPH_FAILED);
2046+ OPS_ERR_IF(opParamInfo_.sparseShardSize == nullptr, OPS_LOG_E(opName_, "attr sparseShardSize is nullptr"), return ge::GRAPH_FAILED);
2047+ sfaaInfo.sparseMode = *opParamInfo_.sparseMode;
M
Mmiaofangzheng4月2日

连续 7 行对 opParamInfo_ 的属性指针解引用（sparseMode、attentionMode、keyQuantMode、valueQuantMode、quantScaleRepoMode、tileSize、ropeHeadDim、sparseShardSize），但没有任何空指针检查。如果任何一个属性未设置，会触发空指针解引用导致 host 侧 crash。

likedislike
M
mayafei
4月3日 评论：
2048+ sfaaInfo.attentionMode = *opParamInfo_.attentionMode;
2049+ sfaaInfo.keyQuantMode = *opParamInfo_.keyQuantMode;
2050+ sfaaInfo.valueQuantMode = *opParamInfo_.valueQuantMode;
2051+ sfaaInfo.quantScaleRepoMode = *opParamInfo_.quantScaleRepoMode;
2052+ sfaaInfo.tileSize = *opParamInfo_.tileSize;
2053+ sfaaInfo.ropeHeadDim = *opParamInfo_.ropeHeadDim;
2054+ sfaaInfo.sparseShardSize = *opParamInfo_.sparseShardSize;
2055+ 
2056+ sfaaInfo.qLayout = qLayout_;
2057+ sfaaInfo.topkLayout = topkLayout_;
2058+ sfaaInfo.kvLayout = kvLayout_;
2059+ sfaaInfo.outLayout = outLayout_;
2060+}
2061+ 
2062+ge::graphStatus SFAAInfoParser::Parse(SFAATilingInfo &sfaaInfo)
2063+{
2064+ if (context_ == nullptr) {
2065+ OPS_LOG_E("SparseFlashAttentionAntiquant", "tiling context is nullptr!");
2066+ return ge::GRAPH_FAILED;
2067+ }
2068+ OPS_LOG_FULL(DLOG_INFO, "SparseFlashAttentionAntiquant", "TilingContext: %s",
2069+ SFAADebugTilingContext(context_).c_str());
2070+ if (ge::GRAPH_SUCCESS != GetOpName() ||
2071+ ge::GRAPH_SUCCESS != GetNpuInfo() ||
2072+ ge::GRAPH_SUCCESS != GetOpParaInfo() ||
2073+ ge::GRAPH_SUCCESS != CheckRequiredParaExistence()) {
2074+ return ge::GRAPH_FAILED;
2075+ }
2076+ 
2077+ if (ge::GRAPH_SUCCESS != GetInOutDataType() ||
2078+ ge::GRAPH_SUCCESS != GetQueryAndOutLayout() ||
2079+ ge::GRAPH_SUCCESS != GetTopkLayout() ||
2080+ ge::GRAPH_SUCCESS != GetKvLayout() ||
2081+ ge::GRAPH_SUCCESS != GetKvStorageMode()) {
2082+ return ge::GRAPH_FAILED;
2083+ }
2084+ 
2085+ SetSFAAShape();
2086+ if (
2087+ ge::GRAPH_SUCCESS != GetN1Size() ||
2088+ ge::GRAPH_SUCCESS != GetN2Size() ||
2089+ ge::GRAPH_SUCCESS != GetGSize() ||
2090+ ge::GRAPH_SUCCESS != GetBatchSize() ||
2091+ ge::GRAPH_SUCCESS != GetQTSize() ||
2092+ ge::GRAPH_SUCCESS != GetKVTSize() ||
2093+ ge::GRAPH_SUCCESS != GetS1Size() ||
2094+ ge::GRAPH_SUCCESS != GetQkHeadDim() ||
2095+ ge::GRAPH_SUCCESS != GetS2Size() ||
2096+ ge::GRAPH_SUCCESS != GetValueHeadDim() ||
2097+ ge::GRAPH_SUCCESS != GetSparseBlockCount()) {
2098+ return ge::GRAPH_FAILED;
2099+ }
2100+ 
2101+ if (ge::GRAPH_SUCCESS != GetActualseqInfo()) {
2102+ return ge::GRAPH_FAILED;
2103+ }
2104+ 
2105+ GenerateInfo(sfaaInfo);
2106+ return ge::GRAPH_SUCCESS;
2107+}
2108+ 
2109+IMPL_OP_OPTILING(SparseFlashAttentionAntiquant)
2110+ .Tiling(TilingSparseFlashAttentionAntiquant)
2111+ .TilingParse<SparseFlashAttentionAntiquantCompileInfo>(TilingPrepareForSparseFlashAttentionAntiquant);
2112+} // namespace optiling
Aexperimental/attention/sparse_flash_attention_antiquant/op_host/sparse_flash_attention_antiquant_tiling.h+673-0
@@ -0,0 +1,673 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file sparse_flash_attention_antiquant_tiling.h
13+ * \brief
14+ */
15+#ifndef SPARSE_FLASH_ATTENTION_ANTIQUANT_TILING_H
16+#define SPARSE_FLASH_ATTENTION_ANTIQUANT_TILING_H
17+ 
18+#include <sstream>
19+#include <graph/utils/type_utils.h>
20+#include <exe_graph/runtime/tiling_context.h>
21+#include <tiling/platform/platform_ascendc.h>
22+#include "register/tilingdata_base.h"
23+#include "exe_graph/runtime/tiling_context.h"
24+#include "split_core.h"
25+ 
26+using std::map;
27+using std::string;
28+using std::pair;
29+using namespace optiling::sfaa;
30+ 
31+namespace optiling {
32+// ------------------算子原型索引常量定义----------------
33+// Inputs Index
34+constexpr uint32_t QUERY_INPUT_INDEX = 0;
35+constexpr uint32_t KEY_INPUT_INDEX = 1;
36+constexpr uint32_t VALUE_INPUT_INDEX = 2;
37+constexpr uint32_t SPARSE_INDICES_INPUT_INDEX = 3;
38+constexpr uint32_t KEY_DEQUANT_SCALE_INPUT_INDEX = 4;
39+constexpr uint32_t VALUE_DEQUANT_SCALE_INPUT_INDEX = 5;
40+constexpr uint32_t BLOCK_TABLE_INPUT_INDEX = 6;
41+constexpr uint32_t ACT_SEQ_LEN_Q_INPUT_INDEX = 7;
42+constexpr uint32_t ACT_SEQ_LEN_KV_INPUT_INDEX = 8;
43+constexpr uint32_t SPARSE_SEQ_LEN_KV_INPUT_INDEX = 9;
44+constexpr uint32_t METADATA_INDEX = 10;
45+// Outputs Index
46+constexpr uint32_t OUTPUT_INDEX = 0;
47+// Attributes Index
48+constexpr uint32_t SCALE_VALUE_ATTR_INDEX = 0;
49+constexpr uint32_t SPARSE_BLOCK_SIZE_ATTR_INDEX = 1;
50+constexpr uint32_t KEY_QUANT_MODE_ATTR_INDEX = 2;
51+constexpr uint32_t VALUE_QUANT_MODE_ATTR_INDEX = 3;
52+constexpr uint32_t LAYOUT_QUERY_ATTR_INDEX = 4;
53+constexpr uint32_t LAYOUT_KV_ATTR_INDEX = 5;
54+constexpr uint32_t SPARSE_MODE_ATTR_INDEX = 6;
55+constexpr uint32_t ATTENTION_MODE_ATTR_INDEX = 7;
56+constexpr uint32_t QUANT_SCALE_REPO_MODE_ATTR_INDEX = 8;
57+constexpr uint32_t TILE_SIZE_ATTR_INDEX = 9;
58+constexpr uint32_t ROPE_HEAD_DIM_ATTR_INDEX = 10;
59+constexpr uint32_t SPARSE_SHARD_SIZE_ATTR_INDEX = 11;
60+const uint32_t FIA_MAX_AIC_CORE_NUM = 26; // 25 + 1 保证数组8字节对齐
61+// Dim Num
62+constexpr size_t DIM_NUM_TWO = 2;
63+constexpr size_t DIM_NUM_THREE = 3;
64+constexpr size_t DIM_NUM_FOUR = 4;
65+constexpr size_t DIM_NUM_FIVE = 5;
66+// 常量
67+constexpr uint32_t MAX_BLOCK_SIZE = 1024;
68+constexpr uint32_t COPYND2NZ_SRC_STRIDE_LIMITATION = 65535;
69+constexpr uint32_t NUM_BYTES_FLOAT = 4;
70+constexpr uint32_t NUM_BYTES_FLOAT16 = 2;
71+constexpr uint32_t NUM_BYTES_BF16 = 2;
72+constexpr uint32_t BYTE_BLOCK = 32;
73+const uint32_t SFAA_MAX_AIC_CORE_NUM = 26; // 25 + 1 保证数组8字节对齐
74+ 
75+// ------------------公共定义--------------------------
76+enum class SFAALayout : uint32_t {
77+ BSND = 0,
78+ TND = 1,
79+ PA_BSND = 2,
80+ PA_BNSD = 3,
81+ PA_NZ = 4,
82+};
83+ 
84+struct SFAATilingShapeCompareParam {
85+ int64_t B = 1;
86+ int64_t S = 1;
87+ int64_t N = 1;
88+ int64_t D = 1;
89+ int64_t T = 1;
90+ // PA
91+ int64_t Bs = 1;
92+ int64_t Bn = 1;
93+};
94+ 
95+enum class KvStorageMode : uint32_t {
96+ BATCH_CONTINUOUS = 0,
97+ PAGE_ATTENTION = 1
98+};
99+ 
100+enum class SFAAPerfMode : uint32_t {
101+ C_TEMPLATE_MODE = 0,
102+ V_TEMPLATE_MODE
103+};
104+ 
105+enum class SFAAAxis : uint32_t {
106+ B = 0,
107+ S = 1,
108+ N = 2,
109+ D = 3,
110+ K = 3, // sparse_indices的K和key的D枚举值相同，表达相同位置, 最后一维
111+ T = 5,
112+ Bn = 6, // block number
113+ Bs = 7, // block size
114+ Dn = 8, // d block number
115+ Ds = 9, // d block size
116+};
117+ 
118+struct SFAARequiredParaInfo {
119+ const gert::CompileTimeTensorDesc *desc;
120+ const gert::StorageShape *shape;
121+};
122+ 
123+struct SFAAOptionalParaInfo {
124+ const gert::CompileTimeTensorDesc *desc;
125+ const gert::Tensor *tensor;
126+};
127+ 
128+// -----------算子Tiling入参结构体定义---------------
129+struct SFAAParaInfo {
130+ SFAARequiredParaInfo query = {nullptr, nullptr};
131+ SFAARequiredParaInfo key = {nullptr, nullptr};
132+ SFAARequiredParaInfo value = {nullptr, nullptr};
133+ SFAARequiredParaInfo sparseIndices = {nullptr, nullptr};
134+ SFAAOptionalParaInfo blockTable = {nullptr, nullptr};
135+ SFAAOptionalParaInfo actualSeqLengthsQ = {nullptr, nullptr};
136+ SFAAOptionalParaInfo actualSeqLengths = {nullptr, nullptr};
137+ SFAAOptionalParaInfo sparseSeqLengths = {nullptr, nullptr};
138+ SFAAOptionalParaInfo queryRope = {nullptr, nullptr};
139+ SFAAOptionalParaInfo keyRope = {nullptr, nullptr};
140+ SFAAOptionalParaInfo keyDequantScale = {nullptr, nullptr};
141+ SFAAOptionalParaInfo valueDequantScale = {nullptr, nullptr};
142+ SFAARequiredParaInfo attenOut = {nullptr, nullptr};
143+ 
144+ const char *layoutQuery = nullptr;
145+ const char *layoutKV = nullptr;
146+ const int64_t *sparseBlockSize = nullptr;
147+ const uint32_t *sparseBlockCount = nullptr;
148+ const uint32_t *blockSize = nullptr;
149+ const float *scaleValue = nullptr;
150+ const int64_t *sparseMode = nullptr;
151+ const int64_t *attentionMode = nullptr;
152+ const int64_t *keyQuantMode = nullptr;
153+ const int64_t *valueQuantMode = nullptr;
154+ const int64_t *quantScaleRepoMode = nullptr;
155+ const int64_t *tileSize = nullptr;
156+ const int64_t *ropeHeadDim = nullptr;
157+ const int64_t *sparseShardSize = nullptr;
158+};
159+ 
160+// -----------算子TilingData定义---------------
161+BEGIN_TILING_DATA_DEF(SparseFlashAttentionAntiquantBaseParamsMla)
162+TILING_DATA_FIELD_DEF(uint32_t, batchSize)
163+TILING_DATA_FIELD_DEF(uint32_t, seqSize)
164+TILING_DATA_FIELD_DEF(uint32_t, qSeqSize)
165+TILING_DATA_FIELD_DEF(uint32_t, kvHeadNum)
166+TILING_DATA_FIELD_DEF(uint32_t, qkHeadDim)
167+TILING_DATA_FIELD_DEF(uint32_t, ropeHeadDim)
168+TILING_DATA_FIELD_DEF(int64_t, blockSize)
169+TILING_DATA_FIELD_DEF(uint32_t, maxBlockNumPerBatch)
170+TILING_DATA_FIELD_DEF(float, scaleValue)
171+TILING_DATA_FIELD_DEF(uint32_t, nNumOfQInOneGroup)
172+TILING_DATA_FIELD_DEF(uint32_t, actualLenDimsQ)
173+TILING_DATA_FIELD_DEF(uint32_t, actualLenDimsKV)
174+TILING_DATA_FIELD_DEF(uint32_t, sparseLenDimsKV)
175+TILING_DATA_FIELD_DEF(uint32_t, outputLayout)
176+TILING_DATA_FIELD_DEF(uint32_t, sparseMode)
177+TILING_DATA_FIELD_DEF(int64_t, sparseBlockSize)
178+TILING_DATA_FIELD_DEF(uint32_t, sparseBlockCount)
179+TILING_DATA_FIELD_DEF(uint32_t, sparseShardSize)
180+TILING_DATA_FIELD_DEF(uint32_t, attentionMode)
181+TILING_DATA_FIELD_DEF(uint32_t, keyQuantMode)
182+TILING_DATA_FIELD_DEF(uint32_t, valueQuantMode)
183+TILING_DATA_FIELD_DEF(uint32_t, quantScaleRepoMode)
184+END_TILING_DATA_DEF
185+REGISTER_TILING_DATA_CLASS(SparseFlashAttentionAntiquantBaseParamsMlaOp, SparseFlashAttentionAntiquantBaseParamsMla)
186+ 
187+BEGIN_TILING_DATA_DEF(SparseFlashAttentionAntiquantSingleCoreParamsMla)
188+TILING_DATA_FIELD_DEF(uint32_t, usedCoreNum);
189+END_TILING_DATA_DEF
190+REGISTER_TILING_DATA_CLASS(SparseFlashAttentionAntiquantSingleCoreParamsMlaOp,
191+ SparseFlashAttentionAntiquantSingleCoreParamsMla)
192+ 
193+BEGIN_TILING_DATA_DEF(SparseFlashAttentionAntiquantSingleCoreTensorSizeMla)
194+TILING_DATA_FIELD_DEF(uint32_t, mmResUbSize);
195+TILING_DATA_FIELD_DEF(uint32_t, bmm2ResUbSize);
196+END_TILING_DATA_DEF
197+REGISTER_TILING_DATA_CLASS(SparseFlashAttentionAntiquantSingleCoreTensorSizeMlaOp,
198+ SparseFlashAttentionAntiquantSingleCoreTensorSizeMla)
199+ 
200+BEGIN_TILING_DATA_DEF(SparseFlashAttentionAntiquantSplitKVParamsMla)
201+TILING_DATA_FIELD_DEF(uint32_t, s2) // S2切分份数
202+TILING_DATA_FIELD_DEF(uint32_t, accumOutSize) // FD workspace
203+TILING_DATA_FIELD_DEF(uint32_t, logSumExpSize) // FD workspace
204+END_TILING_DATA_DEF
205+REGISTER_TILING_DATA_CLASS(SparseFlashAttentionAntiquantSplitKVParamsMlaOp,
206+ SparseFlashAttentionAntiquantSplitKVParamsMla)
207+ 
208+// 内切基本块参数
209+BEGIN_TILING_DATA_DEF(SparseFlashAttentionAntiquantInnerSplitParams)
210+TILING_DATA_FIELD_DEF(uint32_t, mBaseSize)
211+TILING_DATA_FIELD_DEF(uint32_t, s2BaseSize)
212+END_TILING_DATA_DEF
213+REGISTER_TILING_DATA_CLASS(SparseFlashAttentionAntiquantInnerSplitParamsOp,
214+ SparseFlashAttentionAntiquantInnerSplitParams)
215+ 
216+// 外切分核参数
217+BEGIN_TILING_DATA_DEF(SparseFlashAttentionAntiquantOuterSplitParams)
218+TILING_DATA_FIELD_DEF_ARR(uint32_t, FIA_MAX_AIC_CORE_NUM, bN2End)
219+TILING_DATA_FIELD_DEF_ARR(uint32_t, FIA_MAX_AIC_CORE_NUM, gS1End)
220+TILING_DATA_FIELD_DEF_ARR(uint32_t, FIA_MAX_AIC_CORE_NUM, s2End)
221+END_TILING_DATA_DEF
222+REGISTER_TILING_DATA_CLASS(SparseFlashAttentionAntiquantOuterSplitParamsOp,
223+ SparseFlashAttentionAntiquantOuterSplitParams)
224+ 
225+// FlashDecode规约参数
226+BEGIN_TILING_DATA_DEF(SparseFlashAttentionAntiquantFlashDecodeParams)
227+TILING_DATA_FIELD_DEF(uint32_t, numOfFdHead)
228+TILING_DATA_FIELD_DEF(uint32_t, reserved)
229+TILING_DATA_FIELD_DEF(uint32_t, gS1BaseSizeOfFd) // FD负载均衡中，每个FD任务按gS1切分的基本size
230+TILING_DATA_FIELD_DEF(uint32_t, usedVecNumOfFd) // FD负载均衡中，用到的vector数
231+TILING_DATA_FIELD_DEF_ARR(uint32_t, FIA_MAX_AIC_CORE_NUM, bN2IdxOfFdHead)
232+TILING_DATA_FIELD_DEF_ARR(uint32_t, FIA_MAX_AIC_CORE_NUM, gS1IdxOfFdHead)
233+TILING_DATA_FIELD_DEF_ARR(uint32_t, FIA_MAX_AIC_CORE_NUM, s2SplitNumOfFdHead)
234+TILING_DATA_FIELD_DEF_ARR(uint32_t, FIA_MAX_AIC_CORE_NUM, s2SplitStartIdxOfCore)
235+TILING_DATA_FIELD_DEF_ARR(uint32_t, FIA_MAX_AIC_CORE_NUM, gS1SplitNumOfFdHead) // FD负载均衡中，每个FD任务按gS1基本size切分后的份数
236+TILING_DATA_FIELD_DEF_ARR(uint32_t, FIA_MAX_AIC_CORE_NUM, gS1LastPartSizeOfFdHead) // FD负载均衡中，每个FD任务按gS1基本size切分后，最后一份的gS1大小，即尾块大小
237+TILING_DATA_FIELD_DEF_ARR(uint32_t, FIA_MAX_AIC_CORE_NUM * 2, gS1IdxEndOfFdHead) // FD负载均衡中，每个vector核处理的最后一个FD任务的序号
238+TILING_DATA_FIELD_DEF_ARR(uint32_t, FIA_MAX_AIC_CORE_NUM * 2, gS1IdxEndOfFdHeadSplit) // FD负载均衡中，每个vector核处理的最后一个FD任务的子划分的序号
239+END_TILING_DATA_DEF
240+REGISTER_TILING_DATA_CLASS(SparseFlashAttentionAntiquantFlashDecodeParamsOp, SparseFlashAttentionAntiquantFlashDecodeParams)
241+ 
242+BEGIN_TILING_DATA_DEF(SparseFlashAttentionAntiquantTilingDataMla)
243+TILING_DATA_FIELD_DEF_STRUCT(SparseFlashAttentionAntiquantBaseParamsMla, baseParams);
244+TILING_DATA_FIELD_DEF_STRUCT(SparseFlashAttentionAntiquantSplitKVParamsMla, splitKVParams);
245+TILING_DATA_FIELD_DEF_STRUCT(SparseFlashAttentionAntiquantSingleCoreParamsMla, singleCoreParams);
246+TILING_DATA_FIELD_DEF_STRUCT(SparseFlashAttentionAntiquantSingleCoreTensorSizeMla, singleCoreTensorSize);
247+TILING_DATA_FIELD_DEF_STRUCT(SparseFlashAttentionAntiquantInnerSplitParams, innerSplitParams);
248+TILING_DATA_FIELD_DEF_STRUCT(SparseFlashAttentionAntiquantOuterSplitParams, outerSplitParams);
249+TILING_DATA_FIELD_DEF_STRUCT(SparseFlashAttentionAntiquantFlashDecodeParams, fdParams);
250+END_TILING_DATA_DEF
251+REGISTER_TILING_DATA_CLASS(SparseFlashAttentionAntiquant, SparseFlashAttentionAntiquantTilingDataMla)
252+ 
253+template <typename T> inline T Align(T num, T rnd)
254+{
255+ return (((rnd) == 0) ? 0 : (((num) + (rnd) - 1) / (rnd) * (rnd)));
256+}
257+ 
258+template <typename T>
259+std::string SFAAShape2String(const T &shape)
260+{
261+ std::ostringstream oss;
262+ oss << "[";
263+ if (shape.GetDimNum() > 0) {
264+ for (size_t i = 0; i < shape.GetDimNum() - 1; ++i) {
265+ oss << shape.GetDim(i) << ", ";
266+ }
267+ oss << shape.GetDim(shape.GetDimNum() - 1);
268+ }
269+ oss << "]";
270+ return oss.str();
271+}
272+ 
273+static std::string GetShapeStr(gert::Shape shape);
274+static std::string SFAADataTypeToSerialString(ge::DataType type);
275+string SFAATensorDesc2String(const gert::StorageShape *shape, const gert::CompileTimeTensorDesc *tensor);
276+string SFAADebugTilingContext(const gert::TilingContext *context);
277+std::string SFAALayoutToSerialString(SFAALayout layout);
278+ 
279+// -----------算子Tiling入参信息类---------------
280+struct SFAATilingInfo {
281+ const char *opName = nullptr;
282+ fe::PlatFormInfos *platformInfo = nullptr;
283+ SFAAParaInfo opParamInfo;
284+ 
285+ // Base Param
286+ platform_ascendc::SocVersion socVersion = platform_ascendc::SocVersion::ASCEND910B;
287+ uint32_t bSize = 0;
288+ uint32_t n1Size = 0;
289+ uint32_t n2Size = 0;
290+ uint32_t s1Size = 0;
291+ int64_t s2Size = 0;
292+ uint32_t qkHeadDim = 0;
293+ uint32_t vHeadDim = 0;
294+ uint32_t gSize = 0;
295+ uint32_t ropeHeadDim = 0;
296+ uint32_t qTSize = 0; // 仅TND时生效
297+ uint32_t kvTSize = 0; // 仅TND时生效
298+ float scaleValue = 0;
299+ uint32_t sparseShardSize = 0;
300+ uint32_t innerPrecise = 0;
301+ uint32_t l2CacheOffFlag = 0;
302+ int64_t sparseBlockSize = 0;
303+ int64_t sparseBlockCount = 0;
304+ 
305+ bool pageAttentionFlag = false;
306+ int64_t blockSize = 0;
307+ uint32_t blockTypeSize = 0;
308+ uint32_t maxBlockNumPerBatch = 0;
309+ uint32_t totalBlockNum = 0;
310+ 
311+ uint32_t actualLenDimsQ = 0;
312+ uint32_t maxActualseq = 0;
313+ 
314+ bool isSameSeqAllKVTensor = true;
315+ bool isSameActualseq = true;
316+ uint32_t actualLenDimsKV = 0;
317+ uint32_t sparseLenDimsKV = 0;
318+ std::vector<int64_t> kvListSeqLens {};
319+ 
320+ uint32_t sparseMode = 0;
321+ 
322+ int64_t attentionMode = 0;
323+ int64_t keyQuantMode = 0;
324+ int64_t valueQuantMode = 0;
325+ int64_t quantScaleRepoMode = 0;
326+ int64_t tileSize = 0;
327+ 
328+ ge::DataType inputQType = ge::DT_FLOAT16;
329+ ge::DataType inputKvType = ge::DT_FLOAT16;
330+ ge::DataType outputType = ge::DT_FLOAT16;
331+ 
332+ KvStorageMode kvStorageMode = KvStorageMode::BATCH_CONTINUOUS;
333+ 
334+ SFAALayout qLayout = SFAALayout::BSND;
335+ SFAALayout topkLayout = SFAALayout::BSND;
336+ SFAALayout outLayout = SFAALayout::BSND;
337+ SFAALayout kvLayout = SFAALayout::BSND;
338+ 
339+ ge::DataType inputQRopeType = ge::DT_FLOAT16;
340+ ge::DataType inputKRopeType = ge::DT_FLOAT16;
341+ 
342+ uint64_t l2CacheSize = 0;
343+};
344+ 
345+// ---------------算子Tiling类---------------
346+class SFAAMlaTiling {
347+public:
348+ explicit SFAAMlaTiling(gert::TilingContext *context) : context_(context) {}
349+ ge::graphStatus DoOpTiling(SFAATilingInfo *sfaaInfo);
350+ 
351+private:
352+ ge::graphStatus SetBlockDim(uint32_t blockDim);
353+ ge::graphStatus SetTilingKey(uint64_t tilingKey);
354+ ge::graphStatus SetWorkspaceSize(uint64_t workspaceSize);
355+ ge::graphStatus SetTilingData(TilingDef &tilingData);
356+ gert::TilingContext *context_ = nullptr;
357+ ge::graphStatus GetPlatformInfo();
358+ void GenTilingKey();
359+ bool DealSameSeqEachBatch();
360+ 
361+ void ZeroTensorProcess();
362+ void InitParams();
363+ 
364+ void Split();
365+ bool IsBalanceSplitCore();
366+ 
367+ void SplitBalancedBN();
368+ void CalcInnerSize(uint32_t s2Size);
369+ void SetSplitOutput(const SplitResult &res);
370+ void CreateSplitInput(BaseInfo &baseInfo, SplitParam &splitParam);
371+ 
372+ bool IsFlashDecode(uint32_t coreNum);
373+ 
374+ void FillTilingBaseParamsMla();
375+ void FillTilingSplitKVMla();
376+ 
377+ void FillTilingSingleCoreParamsMla();
378+ void FillTilingSingleCoreTensorSizeMla();
379+ void FillTiling();
380+ 
381+ void CalcUbBmm();
382+ void CheckUbSpace();
383+ void NormalCalcFDWorkSpace(const uint32_t actCoreNum);
384+ void CalcFDWorkSpace(const uint32_t actCoreNum);
385+ void GetWorkspaceSize();
386+ 
387+ uint32_t CalcBalanceFDParamNums(const uint32_t actCoreNum);
388+ 
389+ void CalcBlockDim();
390+ 
391+ bool balanceModeFlag_ = false;
392+ bool splitKVFlag_ = false;
393+ 
394+ uint32_t coreNum_ = 0;
395+ SFAAPerfMode perfMode_ = SFAAPerfMode::V_TEMPLATE_MODE;
396+ uint32_t kvSplitPart_ = 1;
397+ size_t mmResUbSize_ = 0;
398+ size_t bmm2ResUbSize_ = 0;
399+ size_t qPreSizeMla_ = 0;
400+ uint32_t sInnerLoopTimes_ = 0;
401+ uint32_t sInnerSize_ = 0;
402+ uint32_t sInnerSizeTail_ = 0;
403+ uint32_t sInnerSizeAlign_ = 0;
404+ uint32_t kvSplit_ = 0;
405+ uint32_t usedCoreNum_ = 0;
406+ uint32_t formerCoreNum_ = 0;
407+ uint32_t blockSplitBn2Range_ = 0;
408+ uint32_t tailSplitedBatchRange_ = 0;
409+ 
410+ uint32_t aicNum_ = 0;
411+ uint32_t aivNum_ = 0;
412+ size_t libapiSize_ = 0;
413+ 
414+ SparseFlashAttentionAntiquantTilingDataMla tilingData_;
415+ uint32_t blockDim_{0};
416+ uint64_t workspaceSize_{0};
417+ uint64_t tilingKey_{0};
418+ 
419+ uint32_t headDimAlign_ = 0;
420+ uint32_t mBaseSize_ = 128;
421+ uint32_t mFdBaseSize_ = 8;
422+ 
423+ SFAATilingInfo *sfaaInfo_ = nullptr;
424+};
425+ 
426+// -----------算子Tiling入参信息解析及Check类---------------
427+class SFAATilingCheck {
428+public:
429+ explicit SFAATilingCheck(const SFAATilingInfo &sfaaInfo) : sfaaInfo_(sfaaInfo) {};
430+ ~SFAATilingCheck() = default;
431+ virtual ge::graphStatus Process();
432+private:
433+ void Init();
434+ void LogErrorDtypeSupport(const std::vector<ge::DataType> &expectDtypeList,
435+ const ge::DataType &actualDtype, const std::string &name) const;
436+ ge::graphStatus CheckDtypeSupport(const gert::CompileTimeTensorDesc *desc,
437+ const std::string &name) const;
438+ template <typename T> void LogErrorNumberSupport(const std::vector<T> &expectNumberList,
439+ const T &actualValue, const std::string &name, const std::string subName) const;
440+ template <typename T> void LogErrorDimNumSupport(const std::vector<T> &expectNumberList,
441+ const T &actualValue, const std::string &name) const;
442+ ge::graphStatus CheckDimNumSupport(const gert::StorageShape *shape,
443+ const std::vector<size_t> &expectDimNumList, const std::string &name) const;
444+ ge::graphStatus CheckDimNumInLayoutSupport(const SFAALayout &layout,
445+ const gert::StorageShape *shape, const std::string &name) const;
446+ void LogErrorLayoutSupport(const std::vector<SFAALayout> &expectLayoutList,
447+ const SFAALayout &actualLayout, const std::string &name) const;
448+ ge::graphStatus GetExpectedShape(gert::Shape &shapeExpected,
449+ const SFAATilingShapeCompareParam &param, const SFAALayout &layout) const;
450+ ge::graphStatus CompareShape(SFAATilingShapeCompareParam &param,
451+ const gert::Shape &shape, const SFAALayout &layout, const std::string &name) const;
452+ ge::graphStatus CheckLayoutSupport(const SFAALayout &actualLayout, const std::string &name) const;
453+ ge::graphStatus CheckSingleParaQuery() const;
454+ ge::graphStatus CheckSingleParaKey() const;
455+ ge::graphStatus CheckSingleParaValue() const;
456+ ge::graphStatus CheckSingleParaAttenOut() const;
457+ ge::graphStatus CheckSingleParaNumHeads() const;
458+ ge::graphStatus CheckSingleParaKvHeadNums() const;
459+ ge::graphStatus CheckSingleParaLayout() const;
460+ ge::graphStatus CheckSingleParaSparseMode() const;
461+ ge::graphStatus CheckSingleParaSparseBlockSize() const;
462+ ge::graphStatus CheckSingleParaSparseIndices() const;
463+ ge::graphStatus CheckSinglePara() const;
464+ ge::graphStatus CheckMultiParaConsistency() const;
465+ ge::graphStatus CheckDequantScaleNotExistence();
466+ ge::graphStatus CheckExists(const void *pointer, const std::string &name) const;
467+ ge::graphStatus CheckNotExists(const void *pointer, const std::string &name) const;
468+ ge::graphStatus CheckExistsByMap(const std::map<std::string, const void *> &paramMap) const;
469+ ge::graphStatus CheckNotExistsByMap(const std::map<std::string, const void *> &paramMap) const;
470+ ge::graphStatus CheckExistenceByMap(std::map<std::string, const void *> &existMap,
471+ std::map<std::string, const void *> &notExistMap) const;
472+ template <typename T> ge::graphStatus CheckAttrValueByMap(
473+ std::map<std::string, std::pair<const T *, T>> &attrMap) const;
474+ ge::graphStatus CheckParaExistenceMlaAntiquant() const;
475+ ge::graphStatus CheckParaExistenceGqaAntiquant() const;
476+ ge::graphStatus CheckParaExistenceMla() const;
477+ ge::graphStatus CheckParaExistence();
478+ ge::graphStatus GetActualSeqLenSize(uint32_t &size, const gert::Tensor *tensor,
479+ const SFAALayout &layout, const std::string &name);
480+ void SetSFAAShapeCompare();
481+ ge::graphStatus CheckKVDType();
482+ ge::graphStatus CheckKVShapeForBatchContinuous();
483+ uint32_t GetTypeSize(ge::DataType dtype) const;
484+ ge::graphStatus CheckKVShapeForPageAttention();
485+ ge::graphStatus CheckKVShape();
486+ ge::graphStatus CheckKV();
487+ ge::graphStatus CheckTopK();
488+ ge::graphStatus CheckTopkShape();
489+ ge::graphStatus CheckBlockTable() const;
490+ ge::graphStatus CheckDTypeConsistency(const ge::DataType &actualDtype,
491+ const ge::DataType &expectDtype, const std::string &name) const;
492+ 
493+ ge::graphStatus CheckAttenOut();
494+ ge::graphStatus CheckAttenOutShape();
495+ ge::graphStatus CheckActualSeqLensQ();
496+ ge::graphStatus CheckActualSeqLensQShape();
497+ ge::graphStatus CheckActualSeqLensQDType();
498+ ge::graphStatus CheckActualSeqLens();
499+ ge::graphStatus CheckActualSeqLensDType();
500+ ge::graphStatus CheckActualSeqLensShape();
501+ ge::graphStatus CheckSparseSeqLens();
502+ ge::graphStatus CheckSparseSeqLensDType();
503+ ge::graphStatus CheckSparseSeqLensShape();
504+ ge::graphStatus CheckMultiParaConsistency();
505+ 
506+ ge::graphStatus CheckFeature() const;
507+ ge::graphStatus CheckFeatureAntiquantShape() const;
508+ ge::graphStatus CheckFeatureAntiquantLayout() const;
509+ ge::graphStatus CheckFeatureAntiquantDtype() const;
510+ ge::graphStatus CheckFeatureAntiquantAttr() const;
511+ ge::graphStatus CheckFeatureAntiquantPa() const;
512+ ge::graphStatus CheckFeatureMlaAntiquantShape() const;
513+ ge::graphStatus CheckFeatureMlaAntiquantAttr() const;
514+ ge::graphStatus CheckFeatureMlaAntiquant() const;
515+ ge::graphStatus CheckFeatureGqaAntiquantShape() const;
516+ ge::graphStatus CheckFeatureGqaAntiquantAttr() const;
517+ ge::graphStatus CheckFeatureGqaAntiquant() const;
518+ 
519+private:
520+ const char *opName_;
521+ fe::PlatFormInfos *platformInfo_;
522+ SFAAParaInfo opParamInfo_;
523+ const SFAATilingInfo &sfaaInfo_;
524+ 
525+ uint32_t bSize_ = 0;
526+ uint32_t n1Size_ = 0;
527+ uint32_t n2Size_ = 0;
528+ uint32_t gSize_ = 0;
529+ uint32_t s1Size_ = 0;
530+ int64_t s2Size_ = 0;
531+ uint32_t qkHeadDim_ = 0;
532+ uint32_t vHeadDim_ = 0;
533+ uint32_t ropeHeadDim_ = 0;
534+ uint32_t qTSize_ = 0; // 仅TND时生效
535+ uint32_t kvTSize_ = 0; // 仅TND时生效
536+ KvStorageMode kvStorageMode_ = KvStorageMode::BATCH_CONTINUOUS;
537+ uint32_t sparseBlockCount_ = 0;
538+ int64_t sparseBlockSize_ = 0;
539+ int64_t attentionMode_ = 0;
540+ int64_t keyQuantMode_ = 0;
541+ int64_t valueQuantMode_ = 0;
542+ int64_t quantScaleRepoMode_ = 0;
543+ int64_t tileSize_ = 0;
544+ int64_t sparseShardSize_ = 0;
545+ 
546+ SFAALayout qLayout_ = SFAALayout::BSND;
547+ SFAALayout topkLayout_ = SFAALayout::BSND;
548+ SFAALayout outLayout_ = SFAALayout::BSND;
549+ SFAALayout kvLayout_ = SFAALayout::BSND;
550+ 
551+ uint32_t maxBlockNumPerBatch_ = 0;
552+ int64_t blockSize_ = 0;
553+ 
554+ uint32_t aicNum_ = 0;
555+ uint32_t aivNum_ = 0;
556+ platform_ascendc::SocVersion socVersion_ = platform_ascendc::SocVersion::ASCEND910B;
557+ uint64_t l2CacheSize_ = 0;
558+ 
559+ ge::DataType inputQType_ = ge::DT_FLOAT16;
560+ ge::DataType inputKvType_ = ge::DT_FLOAT16;
561+ ge::DataType outputType_ = ge::DT_FLOAT16;
562+ 
563+ gert::Shape queryShapeCmp_{};
564+ gert::Shape keyShapeCmp_{};
565+ gert::Shape valueShapeCmp_{};
566+ gert::Shape topkShapeCmp_{};
567+ gert::Shape attenOutShapeCmp_{};
568+};
569+ 
570+class SFAAInfoParser {
571+public:
572+ explicit SFAAInfoParser(const gert::TilingContext *context) : context_(context) {}
573+ ~SFAAInfoParser() = default;
574+ 
575+ ge::graphStatus CheckRequiredInOutExistence() const;
576+ ge::graphStatus CheckRequiredAttrExistence() const;
577+ ge::graphStatus CheckRequiredParaExistence() const;
578+ 
579+ ge::graphStatus GetActualSeqLenSize(uint32_t &size, const gert::Tensor *tensor,
580+ SFAALayout &layout, const std::string &name);
581+ ge::graphStatus GetActualSeqLenQSize(uint32_t &size);
582+ ge::graphStatus GetOpName();
583+ ge::graphStatus GetNpuInfo();
584+ void GetOptionalInputParaInfo();
585+ void GetInputParaInfo();
586+ void GetOutputParaInfo();
587+ ge::graphStatus GetAttrParaInfo();
588+ ge::graphStatus GetKvCache();
589+ ge::graphStatus GetOpParaInfo();
590+ 
591+ ge::graphStatus GetInOutDataType();
592+ ge::graphStatus GetBatchSize();
593+ ge::graphStatus GetQTSize();
594+ ge::graphStatus GetKVTSize();
595+ ge::graphStatus GetQkHeadDim();
596+ ge::graphStatus GetS1Size();
597+ ge::graphStatus GetKvStorageMode();
598+ ge::graphStatus GetKvLayout();
599+ void SetSFAAShape();
600+ ge::graphStatus GetS2SizeForBatchContinuous();
601+ ge::graphStatus GetMaxBlockNumPerBatch();
602+ ge::graphStatus GetBlockSize();
603+ ge::graphStatus GetS2SizeForPageAttention();
604+ ge::graphStatus GetS2Size();
605+ ge::graphStatus GetValueHeadDim();
606+ ge::graphStatus GetRopeHeadDim();
607+ ge::graphStatus GetQueryAndOutLayout();
608+ ge::graphStatus GetTopkLayout();
609+ ge::graphStatus GetN1Size();
610+ ge::graphStatus GetN2Size();
611+ ge::graphStatus GetGSize();
612+ ge::graphStatus GetSparseBlockCount();
613+ ge::graphStatus GetActualseqInfo();
614+ void GenerateInfo(SFAATilingInfo &sfaaInfo);
615+ ge::graphStatus Parse(SFAATilingInfo &sfaaInfo);
616+ 
617+public:
618+ bool HasAxis(const SFAAAxis &axis, const SFAALayout &layout, const gert::Shape &shape) const;
619+ size_t GetAxisIdx(const SFAAAxis &axis, const SFAALayout &layout) const;
620+ uint32_t GetAxisNum(const gert::Shape &shape, const SFAAAxis &axis, const SFAALayout &layout) const;
621+ 
622+ const gert::TilingContext *context_ = nullptr;
623+ 
624+ const char *opName_;
625+ fe::PlatFormInfos *platformInfo_;
626+ SFAAParaInfo opParamInfo_;
627+ static constexpr int64_t invalidDimValue_ = std::numeric_limits<int64_t>::min();
628+ 
629+ uint32_t bSize_ = 0;
630+ uint32_t n1Size_ = 0;
631+ uint32_t n2Size_ = 0;
632+ uint32_t gSize_ = 0;
633+ uint32_t s1Size_ = 0;
634+ int64_t s2Size_ = 0;
635+ uint32_t qkHeadDim_ = 0;
636+ uint32_t vHeadDim_ = 0;
637+ uint32_t ropeHeadDim_ = 0;
638+ uint32_t qTSize_ = 0; // 仅TND时生效
639+ uint32_t kvTSize_ = 0; // 仅TND时生效
640+ KvStorageMode kvStorageMode_ = KvStorageMode::BATCH_CONTINUOUS;
641+ uint32_t sparseBlockCount_ = 0;
642+ 
643+ SFAALayout qLayout_ = SFAALayout::BSND;
644+ SFAALayout topkLayout_ = SFAALayout::BSND;
645+ SFAALayout outLayout_ = SFAALayout::BSND;
646+ SFAALayout kvLayout_ = SFAALayout::BSND;
647+ 
648+ uint32_t maxBlockNumPerBatch_ = 0;
649+ uint32_t blockSize_ = 0;
650+ 
651+ platform_ascendc::SocVersion socVersion_ = platform_ascendc::SocVersion::ASCEND910B;
652+ 
653+ ge::DataType inputQType_ = ge::DT_FLOAT16;
654+ ge::DataType inputKvType_ = ge::DT_FLOAT16;
655+ ge::DataType outputType_ = ge::DT_FLOAT16;
656+ 
657+ uint64_t l2CacheSize_ = 0;
658+ 
659+ bool isSameSeqAllKVTensor_ = true;
660+ bool isSameActualseq_ = true;
661+ uint32_t maxActualseq_ = 0;
662+ 
663+ uint32_t actualLenDimsQ_ = 0;
664+ uint32_t actualLenDimsKV_ = 0;
665+ uint32_t sparseLenDimsKV_ = 0;
666+ 
667+ gert::Shape queryShape_{};
668+ gert::Shape keyShape_{};
669+ gert::Shape valueShape_{};
670+ gert::Shape sparseIndicesShape_{};
671+};
672+} // namespace optiling
673+#endif // SPARSE_FLASH_ATTENTION_ANTIQUANT_TILING_H
Aexperimental/attention/sparse_flash_attention_antiquant/op_host/split_core.cpp+742-0
@@ -0,0 +1,742 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file split_core.cpp
13+ * \brief
14+ */
15+ 
16+#include "split_core.h"
17+#include <cstdio>
18+#include <math.h>
19+ 
20+ 
21+namespace optiling {
22+namespace sfaa {
23+ 
24+uint32_t GetS1SeqSize(uint32_t bIdx, const BaseInfo &baseInfo)
25+{
26+ if (baseInfo.actualSeqS1Size.empty()) {
27+ return baseInfo.s1Size;
28+ }
29+ 
30+ if (baseInfo.actualLenQDims == 1U) {
31+ return static_cast<uint32_t>(baseInfo.actualSeqS1Size[0]);
32+ }
33+ 
34+ if (!baseInfo.isAccumSeqS1) {
35+ return static_cast<uint32_t>(baseInfo.actualSeqS1Size[bIdx]);
36+ }
37+ 
38+ return (bIdx == 0) ? static_cast<uint32_t>(baseInfo.actualSeqS1Size[bIdx]) :
39+ static_cast<uint32_t>(baseInfo.actualSeqS1Size[bIdx] - baseInfo.actualSeqS1Size[bIdx - 1U]);
40+}
41+ 
42+uint32_t GetS2SeqSize(uint32_t bIdx, const BaseInfo &baseInfo)
43+{
44+ int64_t s2Size = 0;
45+ if (baseInfo.actualSeqS2Size.empty()) {
46+ s2Size = baseInfo.s2Size;
47+ } else if (baseInfo.actualLenKvDims == 1U) {
48+ s2Size = static_cast<uint32_t>(baseInfo.actualSeqS2Size[0]);
49+ }else if (!baseInfo.isAccumSeqS2) {
50+ s2Size = static_cast<uint32_t>(baseInfo.actualSeqS2Size[bIdx]);
51+ } else {
52+ s2Size = (bIdx == 0) ? static_cast<uint32_t>(baseInfo.actualSeqS2Size[bIdx]) :
53+ static_cast<uint32_t>(baseInfo.actualSeqS2Size[bIdx] - baseInfo.actualSeqS2Size[bIdx - 1U]);
54+ }
55+ int64_t actualS2WithSparse = baseInfo.sparseBlockCount * baseInfo.sparseBlockSize;
56+ if (actualS2WithSparse < s2Size) {
57+ s2Size = actualS2WithSparse;
58+ }
59+ return s2Size;
60+}
61+ 
62+uint32_t GetSparseSeqSize(uint32_t bIdx, const BaseInfo &baseInfo)
63+{
64+ return baseInfo.sparseBlockCount;
65+}
66+ 
67+int64_t CalcPreTokenLeftUp(uint32_t s1Size, uint32_t s2Size, const BaseInfo &baseInfo)
68+{
69+ auto mode = static_cast<SparseMode>(baseInfo.sparseMode);
70+ if (mode == SparseMode::BAND) {
71+ return static_cast<int64_t>(s1Size) - static_cast<int64_t>(s2Size) + baseInfo.preToken;
72+ }
73+ return baseInfo.preToken;
74+}
75+ 
76+int64_t CalcNextTokenLeftUp(uint32_t s1Size, uint32_t s2Size, const BaseInfo &baseInfo)
77+{
78+ auto mode = static_cast<SparseMode>(baseInfo.sparseMode);
79+ switch (mode) {
80+ case SparseMode::DEFAULT_MASK:
81+ case SparseMode::ALL_MASK:
82+ case SparseMode::LEFT_UP_CAUSAL:
83+ return baseInfo.nextToken;
84+ case SparseMode::RIGHT_DOWN_CAUSAL:
85+ return static_cast<int64_t>(s2Size) - static_cast<int64_t>(s1Size);
86+ case SparseMode::BAND:
87+ return static_cast<int64_t>(s2Size) - static_cast<int64_t>(s1Size) + baseInfo.nextToken;
88+ default:
89+ return baseInfo.nextToken;
90+ }
91+}
92+ 
93+int64_t CalcCost(uint32_t basicM, uint32_t basicS2)
94+{
95+ uint32_t alignCoefM = 16U;
96+ uint32_t alignCoefS2 = 64U;
97+ uint32_t alignBasicM = (basicM + alignCoefM - 1U) >> 4U; // 按alignCoefM对齐，向上取整，4：移位操作实现除16
98+ uint32_t alignBasicS2 = (basicS2 + alignCoefS2 - 1U) >> 6U; // 按alignCoefS2对齐，向上取整，6：移位操作实现除64
99+ return static_cast<int64_t>(6U * alignBasicM + 10U * alignBasicS2); // 6：M轴系数，10：S2轴系数
100+}
101+ 
102+BlockCost<int64_t> CalcCostTable(uint32_t s1NormalSize, uint32_t s2NormalSize, uint32_t s1GTailSize,
103+ uint32_t s2TailSize)
104+{
105+ BlockCost<int64_t> typeCost {};
106+ typeCost[NORMAL_BLOCK][NORMAL_BLOCK] = CalcCost(s1NormalSize, s2NormalSize);
107+ typeCost[TAIL_BLOCK][NORMAL_BLOCK] = (s1GTailSize == 0U) ? 0U : CalcCost(s1GTailSize, s2NormalSize);
108+ typeCost[NORMAL_BLOCK][TAIL_BLOCK] = (s2TailSize == 0U) ? 0U : CalcCost(s1NormalSize, s2TailSize);
109+ typeCost[TAIL_BLOCK][TAIL_BLOCK] = (s1GTailSize == 0U || s2TailSize == 0U) ? 0U : CalcCost(s1GTailSize, s2TailSize);
110+ return typeCost;
111+}
112+ 
113+Range<uint32_t> CalcS2Range(uint32_t s1GIdx, const BaseInfo &baseInfo, const SplitParam &splitParam,
114+ const BatchCache &batchCache)
115+{
116+ uint32_t s2Start = 0U;
117+ uint32_t s2End = 0U;
118+ 
119+ // actual seq == 0
120+ if (batchCache.s1Size == 0U || batchCache.s2Size == 0U) {
121+ return std::make_pair(s2Start, s2End);
122+ }
123+ 
124+ // no mask
125+ if (!baseInfo.attenMaskFlag) {
126+ s2Start = 0U;
127+ s2End = (batchCache.s2Size + splitParam.s2BaseSize - 1U) / splitParam.s2BaseSize;
128+ return std::make_pair(s2Start, s2End);
129+ }
130+ 
131+ // 1. calc index of s2FirstToken, s2LastToken by index of s1GFirstToken, s1GLastToken
132+ int64_t s1GFirstToken = static_cast<int64_t>(s1GIdx) * static_cast<int64_t>(splitParam.mBaseSize);
133+ int64_t s1GLastToken = std::min(s1GFirstToken + static_cast<int64_t>(splitParam.mBaseSize),
134+ static_cast<int64_t>(batchCache.s1Size) * static_cast<int64_t>(baseInfo.gSize)) - 1;
135+
136+ int64_t s1FirstToken = 0;
137+ int64_t s1LastToken = 0;
138+ if (baseInfo.isS1G) {
139+ s1FirstToken = s1GFirstToken / static_cast<int64_t>(baseInfo.gSize);
140+ s1LastToken = s1GLastToken / static_cast<int64_t>(baseInfo.gSize);
141+ } else {
142+ if (s1GFirstToken / batchCache.s1Size == s1GLastToken / batchCache.s1Size) {
143+ // start and end locate in one G
144+ s1FirstToken = s1GFirstToken % static_cast<int64_t>(batchCache.s1Size);
145+ s1LastToken = s1GLastToken % static_cast<int64_t>(batchCache.s1Size);
146+ } else {
147+ // start and end locate in tow or more G, but working same as crossing a complete block
148+ s1FirstToken = 0;
149+ s1LastToken = batchCache.s1Size;
150+ }
151+ }
152+ 
153+ int64_t s2FirstToken = s1FirstToken - batchCache.preTokenLeftUp;
154+ int64_t s2LastToken = s1LastToken + batchCache.nextTokenLeftUp;
155+ 
156+ // 2. trans index of token to index of block
157+ // no valid token
158+ if (s2FirstToken >= static_cast<int64_t>(batchCache.s2Size) || s2LastToken < 0 || s2LastToken < s2FirstToken) {
159+ s2Start = 0U;
160+ s2End = 0U;
161+ return std::make_pair(s2Start, s2End);
162+ }
163+ // get valid range
164+ s2FirstToken = Clip(s2FirstToken, static_cast<int64_t>(0), static_cast<int64_t>(batchCache.s2Size - 1U));
165+ s2LastToken = Clip(s2LastToken, static_cast<int64_t>(0), static_cast<int64_t>(batchCache.s2Size - 1U));
166+ 
167+ s2Start = static_cast<uint32_t>(s2FirstToken) / splitParam.s2BaseSize;
168+ s2End = static_cast<uint32_t>(s2LastToken) / splitParam.s2BaseSize + 1U; // end of block index, Right-open interval
169+ 
170+ return std::make_pair(s2Start, s2End);
171+}
172+ 
173+void CalcSplitInfo(SplitContext &splitContext)
174+{
175+ const BaseInfo &baseInfo = splitContext.baseInfo;
176+ const SplitParam &splitParam = splitContext.splitParam;
177+ 
178+ // 计算每个batch的切分，统计是否为空batch，记录最后有效batch（每个batch的每个N2切分是一样的）
179+ SplitInfo &splitInfo = splitContext.splitInfo;
180+ for (uint32_t bIdx = 0; bIdx < baseInfo.bSize; bIdx++) {
181+ uint32_t s1Size = GetS1SeqSize(bIdx, baseInfo);
182+ uint32_t s2Size = GetS2SeqSize(bIdx, baseInfo);
183+ 
184+ splitInfo.s1GBaseNum[bIdx] = (s1Size * baseInfo.gSize + (splitParam.mBaseSize - 1U)) / splitParam.mBaseSize;
185+ splitInfo.s1GTailSize[bIdx] = (s1Size * baseInfo.gSize) % splitParam.mBaseSize;
186+ splitInfo.s2BaseNum[bIdx] = (s2Size + splitParam.s2BaseSize - 1U) / splitParam.s2BaseSize;
187+ splitInfo.s2TailSize[bIdx] = s2Size % splitParam.s2BaseSize;
188+ if (splitInfo.s1GBaseNum[bIdx] != 0U && splitInfo.s2BaseNum[bIdx] != 0U) {
189+ splitInfo.isKvSeqAllZero = false;
190+ }
191+ }
192+}
193+ 
194+void CalcBatchCache(uint32_t bIdx, const SplitContext &splitContext, BatchCache &batchCache)
195+{
196+ const BaseInfo &baseInfo = splitContext.baseInfo;
197+ const SplitParam &splitParam = splitContext.splitParam;
198+ const SplitInfo &splitInfo = splitContext.splitInfo;
199+ 
200+ batchCache.bIdx = bIdx;
201+ batchCache.s1Size = GetS1SeqSize(bIdx, baseInfo);
202+ batchCache.s2Size = GetS2SeqSize(bIdx, baseInfo);
203+ batchCache.preTokenLeftUp = CalcPreTokenLeftUp(batchCache.s1Size, batchCache.s2Size, baseInfo);
204+ batchCache.nextTokenLeftUp = CalcNextTokenLeftUp(batchCache.s1Size, batchCache.s2Size, baseInfo);
205+ batchCache.typeCost = CalcCostTable(splitParam.mBaseSize, splitParam.s2BaseSize, splitInfo.s1GTailSize[bIdx],
206+ splitInfo.s2TailSize[bIdx]);
207+}
208+ 
209+void CalcS1GCache(uint32_t s1GIdx, const SplitContext &splitContext, const BatchCache &batchCache, S1GCache &s1GCache)
210+{
211+ const BaseInfo &baseInfo = splitContext.baseInfo;
212+ const SplitParam &splitParam = splitContext.splitParam;
213+ const SplitInfo &splitInfo = splitContext.splitInfo;
214+ 
215+ s1GCache.bIdx = batchCache.bIdx;
216+ s1GCache.s1GIdx = s1GIdx;
217+ 
218+ auto s2Range = CalcS2Range(s1GIdx, baseInfo, splitParam, batchCache);
219+ s1GCache.s2Start = s2Range.first;
220+ s1GCache.s2End = s2Range.second;
221+ 
222+ if (s1GCache.s2Start >= s1GCache.s2End) {
223+ s1GCache.s1GBlock = 0;
224+ s1GCache.s1GCost = 0;
225+ s1GCache.s1GLastBlockCost = 0;
226+ s1GCache.s1GNormalBlockCost = 0;
227+ return;
228+ }
229+ 
230+ // 计算S2方向满块、尾块数量
231+ s1GCache.s1GBlock = s1GCache.s2End - s1GCache.s2Start;
232+ uint32_t curTailS2Num = (splitInfo.s2TailSize[batchCache.bIdx] != 0U &&
233+ s1GCache.s2End == splitInfo.s2BaseNum[batchCache.bIdx]) ? 1U : 0U;
234+ uint32_t curNormalS2Num = s1GCache.s1GBlock - curTailS2Num;
235+ if (splitInfo.s1GBaseNum[batchCache.bIdx] == 0) {
236+ s1GCache.s1GCost = 0;
237+ s1GCache.s1GLastBlockCost = 0;
238+ s1GCache.s1GNormalBlockCost = 0;
239+ } else if (s1GIdx == (splitInfo.s1GBaseNum[batchCache.bIdx] - 1U) && splitInfo.s1GTailSize[batchCache.bIdx] != 0U) {
240+ s1GCache.s1GCost = batchCache.typeCost[TAIL_BLOCK][NORMAL_BLOCK] * curNormalS2Num +
241+ batchCache.typeCost[TAIL_BLOCK][TAIL_BLOCK] * curTailS2Num;
242+ s1GCache.s1GLastBlockCost = curTailS2Num > 0U ? batchCache.typeCost[TAIL_BLOCK][TAIL_BLOCK] :
243+ batchCache.typeCost[TAIL_BLOCK][NORMAL_BLOCK];
244+ s1GCache.s1GNormalBlockCost = batchCache.typeCost[TAIL_BLOCK][NORMAL_BLOCK];
245+ } else {
246+ s1GCache.s1GCost = batchCache.typeCost[NORMAL_BLOCK][NORMAL_BLOCK] * curNormalS2Num +
247+ batchCache.typeCost[NORMAL_BLOCK][TAIL_BLOCK] * curTailS2Num;
248+ s1GCache.s1GLastBlockCost = curTailS2Num > 0U ? batchCache.typeCost[NORMAL_BLOCK][TAIL_BLOCK] :
249+ batchCache.typeCost[NORMAL_BLOCK][NORMAL_BLOCK];
250+ s1GCache.s1GNormalBlockCost = batchCache.typeCost[NORMAL_BLOCK][NORMAL_BLOCK];
251+ }
252+}
253+ 
254+void CopyTmpResult(SplitResult &tmpRes, SplitResult &splitRes)
255+{
256+ uint64_t len = tmpRes.bN2End.size();
257+ splitRes.usedCoreNum = tmpRes.usedCoreNum;
258+ splitRes.maxCost = tmpRes.maxCost;
259+ splitRes.numOfFdHead = tmpRes.numOfFdHead;
260+ splitRes.maxS2SplitNum = tmpRes.maxS2SplitNum;
261+ 
262+ for (size_t i = 0; i < len; ++i) {
263+ splitRes.bN2End[i] = tmpRes.bN2End[i];
264+ splitRes.gS1End[i] = tmpRes.gS1End[i];
265+ splitRes.s2End[i] = tmpRes.s2End[i];
266+ 
267+ splitRes.fdRes.bN2IdxOfFdHead[i] = tmpRes.fdRes.bN2IdxOfFdHead[i];
268+ splitRes.fdRes.gS1IdxOfFdHead[i] = tmpRes.fdRes.gS1IdxOfFdHead[i];
269+ splitRes.fdRes.s2SplitNumOfFdHead[i] = tmpRes.fdRes.s2SplitNumOfFdHead[i];
270+ splitRes.fdRes.s2SplitStartIdxOfCore[i] = tmpRes.fdRes.s2SplitStartIdxOfCore[i];
271+ splitRes.fdRes.gS1SplitNumOfFdHead[i] = tmpRes.fdRes.gS1SplitNumOfFdHead[i];
272+ splitRes.fdRes.gS1LastPartSizeOfFdHead[i] = tmpRes.fdRes.gS1LastPartSizeOfFdHead[i];
273+ }
274+}
275+ 
276+void ClearTmpResult(SplitResult &tmpResult)
277+{
278+ uint64_t len = tmpResult.bN2End.size();
279+ tmpResult.usedCoreNum = 0U;
280+ tmpResult.maxCost = 0;
281+ tmpResult.numOfFdHead = 0U;
282+ tmpResult.maxS2SplitNum = 0U;
283+ tmpResult.usedVecNumOfFd = 0U;
284+ 
285+ for (size_t i = 0; i < len; ++i) {
286+ tmpResult.bN2End[i] = 0U;
287+ tmpResult.gS1End[i] = 0U;
288+ tmpResult.s2End[i] = 0U;
289+ tmpResult.fdRes.bN2IdxOfFdHead[i] = 0U;
290+ tmpResult.fdRes.gS1IdxOfFdHead[i] = 0U;
291+ tmpResult.fdRes.s2SplitNumOfFdHead[i] = 0U;
292+ tmpResult.fdRes.s2SplitStartIdxOfCore[i] = 0U;
293+ tmpResult.fdRes.gS1SplitNumOfFdHead[i] = 0U;
294+ tmpResult.fdRes.gS1LastPartSizeOfFdHead[i] = 0U;
295+ }
296+}
297+ 
298+void RollBackCursor(const BaseInfo &baseInfo, const SplitContext &splitContext, const CostInfo &costInfo, SplitResult &splitRes)
299+{
300+ for (size_t i = 0; i < splitRes.usedCoreNum; ++i) {
301+ // x, y, z
302+ if (splitRes.s2End[i] > 0U) {
303+ splitRes.s2End[i] = splitRes.s2End[i] - 1U;
304+ continue;
305+ }
306+ uint32_t bIdx = splitRes.bN2End[i] / baseInfo.n2Size;
307+ // x, y, 0
308+ if (splitRes.gS1End[i] > 0U) {
309+ splitRes.gS1End[i] = splitRes.gS1End[i] - 1U;
310+ splitRes.s2End[i] = splitContext.splitInfo.s2BaseNum[bIdx] - 1U;
311+ continue;
312+ }
313+ 
314+ // x, 0, 0
315+ uint32_t bN2Idx = splitRes.bN2End[i] > 0U ? splitRes.bN2End[i] - 1U : 0U;
316+ bIdx = bN2Idx / baseInfo.n2Size;
317+ while (bN2Idx > 0U && costInfo.bN2BlockOfEachBatch[bIdx] == 0U) {
318+ bN2Idx -= 1U;
319+ bIdx = bN2Idx / baseInfo.n2Size;
320+ }
321+ 
322+ if (costInfo.bN2BlockOfEachBatch[bIdx] != 0U) {
323+ splitRes.bN2End[i] = bN2Idx;
324+ splitRes.gS1End[i] = splitContext.splitInfo.s1GBaseNum[bIdx] - 1U;
325+ splitRes.s2End[i] = splitContext.splitInfo.s2BaseNum[bIdx] - 1U;
326+ } else {
327+ splitRes.bN2End[i] = 0U;
328+ splitRes.gS1End[i] = 0U;
329+ splitRes.s2End[i] = 0U;
330+ }
331+ }
332+}
333+ 
334+void CalcBatchCost(uint32_t bIdx, const SplitContext &splitContext, CostInfo &costInfo)
335+{
336+ const BaseInfo &baseInfo = splitContext.baseInfo;
337+ const SplitInfo &splitInfo = splitContext.splitInfo;
338+ 
339+ costInfo.bN2CostOfEachBatch[bIdx] = 0;
340+ costInfo.bN2BlockOfEachBatch[bIdx] = 0U;
341+ costInfo.bN2LastBlockCostOfEachBatch[bIdx] = 0U;
342+ 
343+ if (GetS1SeqSize(bIdx, baseInfo) == 0U || GetS2SeqSize(bIdx, baseInfo) == 0U) {
344+ return;
345+ }
346+ 
347+ BatchCache bCache;
348+ S1GCache s1GCache;
349+ CalcBatchCache(bIdx, splitContext, bCache);
350+ for (uint32_t s1GIdx = 0; s1GIdx < splitInfo.s1GBaseNum[bIdx]; s1GIdx++) {
351+ CalcS1GCache(s1GIdx, splitContext, bCache, s1GCache);
352+ costInfo.bN2CostOfEachBatch[bIdx] += s1GCache.s1GCost;
353+ costInfo.bN2BlockOfEachBatch[bIdx] += s1GCache.s1GBlock;
354+ if(s1GCache.s1GBlock > 0){
355+ costInfo.bN2LastBlockCostOfEachBatch[bIdx] = s1GCache.s1GLastBlockCost;
356+ }
357+ }
358+}
359+ 
360+void CalcCostInfo(SplitContext &splitContext)
361+{
362+ const BaseInfo &baseInfo = splitContext.baseInfo;
363+ const SplitInfo &splitInfo = splitContext.splitInfo;
364+ 
365+ CostInfo &costInfo = splitContext.costInfo;
366+ 
367+ if (splitInfo.isKvSeqAllZero) {
368+ costInfo.totalCost = 0;
369+ costInfo.totalBlockNum = 0U;
370+ return;
371+ }
372+ 
373+ // 计算batch的负载并记录，用于按batch分配，需要按行计算起止点，统计块数、负载
374+ for (uint32_t bIdx = 0; bIdx < baseInfo.bSize; bIdx++) {
375+ CalcBatchCost(bIdx, splitContext, costInfo);
376+ costInfo.totalCost += costInfo.bN2CostOfEachBatch[bIdx] * baseInfo.n2Size;
377+ costInfo.totalBlockNum += costInfo.bN2BlockOfEachBatch[bIdx] * baseInfo.n2Size;
378+ }
379+}
380+ 
381+void UpdateCursor(const SplitContext &splitContext, AssignContext &assignContext)
382+{
383+ const BaseInfo &baseInfo = splitContext.baseInfo;
384+ const SplitInfo &splitInfo = splitContext.splitInfo;
385+ const CostInfo &costInfo = splitContext.costInfo;
386+ 
387+ bool UpdateS1G = false;
388+ bool UpdateBatch = false;
389+ 
390+ // Update S2
391+ if (assignContext.curS2Idx >= assignContext.s1GCache.s2End) { // 边界assignInfo.s2End是取不到的开区间
392+ assignContext.curS2Idx = 0U;
393+ assignContext.curS1GIdx++;
394+ UpdateS1G = true;
395+ }
396+ 
397+ // Update S1G
398+ if (assignContext.curS1GIdx >= splitInfo.s1GBaseNum[assignContext.curBIdx]) {
399+ assignContext.curS1GIdx = 0U;
400+ assignContext.curBN2Idx++;
401+ }
402+ 
403+ // Update Batch
404+ if (assignContext.curBN2Idx == baseInfo.bSize * baseInfo.n2Size) { // 所有负载全部分配完，设置最后一个核的右开区间，返回
405+ assignContext.curS1GIdx = 0U;
406+ assignContext.curS2Idx = 0U;
407+ assignContext.isFinished = true;
408+ return;
409+ }
410+ 
411+ if (assignContext.curBN2Idx / baseInfo.n2Size != assignContext.curBIdx) {
412+ assignContext.curBIdx = assignContext.curBN2Idx / baseInfo.n2Size;
413+ assignContext.curS1GIdx = 0U;
414+ UpdateBatch = true;
415+ UpdateS1G = true;
416+ }
417+ 
418+ // Update Cache
419+ if (UpdateBatch) {
420+ CalcBatchCache(assignContext.curBIdx, splitContext, assignContext.batchCache);
421+ assignContext.bN2Cost = costInfo.bN2CostOfEachBatch[assignContext.curBIdx];
422+ assignContext.bN2Block = costInfo.bN2BlockOfEachBatch[assignContext.curBIdx];
423+ }
424+ if (UpdateS1G) {
425+ CalcS1GCache(assignContext.curS1GIdx, splitContext, assignContext.batchCache, assignContext.s1GCache);
426+ assignContext.curS2Idx = assignContext.s1GCache.s2Start;
427+ }
428+}
429+ 
430+void AssignByBatch(const SplitContext &splitContext, AssignContext &assignContext)
431+{
432+ if (assignContext.isFinished) {
433+ return;
434+ }
435+ 
436+ const BaseInfo &baseInfo = splitContext.baseInfo;
437+ const CostInfo &costInfo = splitContext.costInfo;
438+ 
439+ while (assignContext.bN2Cost == 0 || IsWithinTolerance(assignContext.coreCache.costLimit,
440+ costInfo.bN2LastBlockCostOfEachBatch[assignContext.curBIdx] / FA_TOLERANCE_RATIO,
441+ assignContext.coreCache.cost + assignContext.bN2Cost)) {
442+ assignContext.coreCache.cost += assignContext.bN2Cost;
443+ assignContext.coreCache.block += assignContext.bN2Block;
444+ assignContext.curBN2Idx++;
445+ 
446+ // to the end
447+ if (assignContext.curBN2Idx == baseInfo.bSize * baseInfo.n2Size) {
448+ assignContext.curS1GIdx = 0U;
449+ assignContext.curS2Idx = 0U;
450+ assignContext.isFinished = true;
451+ return;
452+ }
453+ 
454+ // next batch
455+ if (assignContext.curBN2Idx / baseInfo.n2Size != assignContext.curBIdx) {
456+ assignContext.curBIdx = assignContext.curBN2Idx / baseInfo.n2Size;
457+ CalcBatchCache(assignContext.curBIdx, splitContext, assignContext.batchCache);
458+ }
459+ 
460+ assignContext.bN2Cost = costInfo.bN2CostOfEachBatch[assignContext.curBIdx];
461+ assignContext.bN2Block = costInfo.bN2BlockOfEachBatch[assignContext.curBIdx];
462+ assignContext.curS1GIdx = 0U;
463+ CalcS1GCache(assignContext.curS1GIdx, splitContext, assignContext.batchCache, assignContext.s1GCache);
464+ assignContext.curS2Idx = assignContext.s1GCache.s2Start;
465+ }
466+}
467+ 
468+void AssignByRow(const SplitContext &splitContext, AssignContext &assignContext)
469+{
470+ if (assignContext.isFinished) {
471+ return;
472+ }
473+ 
474+ while (IsWithinTolerance(assignContext.coreCache.costLimit,
475+ assignContext.s1GCache.s1GLastBlockCost / FA_TOLERANCE_RATIO,
476+ assignContext.coreCache.cost + assignContext.s1GCache.s1GCost)) {
477+ assignContext.coreCache.cost += assignContext.s1GCache.s1GCost;
478+ assignContext.coreCache.block += assignContext.s1GCache.s1GBlock;
479+ 
480+ // 当前batch被分配一行出去，更新剩余负载
481+ assignContext.bN2Cost = assignContext.bN2Cost > assignContext.s1GCache.s1GCost ?
482+ assignContext.bN2Cost - assignContext.s1GCache.s1GCost : 0;
483+ assignContext.bN2Block = assignContext.bN2Block > assignContext.s1GCache.s1GBlock ?
484+ assignContext.bN2Block - assignContext.s1GCache.s1GBlock : 0U;
485+ // 计算新一行的信息
486+ do{
487+ assignContext.curS1GIdx++;
488+ CalcS1GCache(assignContext.curS1GIdx, splitContext, assignContext.batchCache, assignContext.s1GCache);
489+ }while(assignContext.s1GCache.s1GBlock == 0);
490+ assignContext.curS2Idx = assignContext.s1GCache.s2Start;
491+ }
492+}
493+ 
494+void AssignByBlock(const SplitContext &splitContext, AssignContext &assignContext)
495+{
496+ if (assignContext.isFinished) {
497+ return;
498+ }
499+ 
500+ int64_t curCost = assignContext.s1GCache.s1GNormalBlockCost;
501+ if (assignContext.curS2Idx == (assignContext.s1GCache.s2End - 1U)) {
502+ curCost = assignContext.s1GCache.s1GLastBlockCost;
503+ }
504+ 
505+ while (IsWithinTolerance(assignContext.coreCache.costLimit, curCost / FA_TOLERANCE_RATIO,
506+ assignContext.coreCache.cost + curCost)) { // (costLimit - curCostOnCore) * FA_TOLERANCE_RATIO > curCost；至少分配1块
507+ assignContext.coreCache.cost += curCost;
508+ assignContext.coreCache.block++;
509+ assignContext.curS2Idx++;
510+ // 当前batch被分配一块出去，更新剩余负载
511+ assignContext.bN2Cost = assignContext.bN2Cost - curCost;
512+ // 当前行被分配一块出去，更新剩余负载
513+ assignContext.s1GCache.s1GCost = assignContext.s1GCache.s1GCost - curCost;
514+ assignContext.bN2Block--;
515+ assignContext.s1GCache.s1GBlock--;
516+ }
517+}
518+ 
519+void ForceAssign(const SplitContext &splitContext, AssignContext &assignContext)
520+{
521+ if (assignContext.isFinished) {
522+ return;
523+ }
524+ 
525+ int64_t curCost = assignContext.s1GCache.s1GNormalBlockCost;
526+ if (assignContext.curS2Idx == (assignContext.s1GCache.s2End - 1U)) {
527+ curCost = assignContext.s1GCache.s1GLastBlockCost;
528+ }
529+ 
530+ assignContext.coreCache.cost += curCost;
531+ assignContext.coreCache.block++;
532+ assignContext.curS2Idx++;
533+ // 当前batch被分配一块出去，更新剩余负载
534+ assignContext.bN2Cost = assignContext.bN2Cost - curCost;
535+ assignContext.bN2Block--;
536+ // 当前行被分配一块出去，更新剩余负载
537+ assignContext.s1GCache.s1GCost = assignContext.s1GCache.s1GCost - curCost;
538+ assignContext.s1GCache.s1GBlock--;
539+ UpdateCursor(splitContext, assignContext);
540+}
541+ 
542+bool IsNeedRecordFDInfo(const AssignContext &assignContext, const SplitResult &splitRes)
543+{
544+ // 切分点大概率不会刚好在行尾，因此滞后处理归约信息的统计，到下一个切分点再判断是否需要归约
545+ // 核0无需处理
546+ if (assignContext.curCoreIdx == 0U) {
547+ return false;
548+ }
549+ // 无跨核行，无需处理
550+ if (assignContext.curKvSplitPart <= 1U) {
551+ return false;
552+ }
553+ // 需要归约的行还未处理完
554+ if (assignContext.curBN2Idx == splitRes.bN2End[assignContext.curCoreIdx - 1U] &&
555+ assignContext.curS1GIdx == splitRes.gS1End[assignContext.curCoreIdx - 1U]) {
556+ return false;
557+ }
558+ return true;
559+}
560+ 
561+void RecordFDInfo(const SplitContext &splitContext, const AssignContext &assignContext, SplitResult &result)
562+{
563+ const BaseInfo &baseInfo = splitContext.baseInfo;
564+ const SplitParam &splitParam = splitContext.splitParam;
565+ const SplitInfo &splitInfo = splitContext.splitInfo;
566+ // 需要规约的行是上一个核的切分点所在位置
567+ uint32_t splitBIdx = result.bN2End[assignContext.curCoreIdx - 1U] / baseInfo.n2Size;
568+ uint32_t splitS1GIdx = result.gS1End[assignContext.curCoreIdx - 1U];
569+ uint32_t s1Size = GetS1SeqSize(splitBIdx, baseInfo);
570+ 
571+ // 计算归约数据的FD均衡划分信息
572+ uint32_t curFdS1gSize = (splitS1GIdx == splitInfo.s1GBaseNum[splitBIdx] - 1U) ?
573+ (s1Size * baseInfo.gSize - splitS1GIdx * splitParam.mBaseSize) : splitParam.mBaseSize;
574+ uint32_t curFdS1gSplitPart = (curFdS1gSize + splitParam.gS1BaseSizeOfFd - 1U) / splitParam.gS1BaseSizeOfFd;
575+ uint32_t curFdS1gLastPartSize = curFdS1gSize - (splitParam.gS1BaseSizeOfFd * (curFdS1gSplitPart - 1U));
576+ // 记录
577+ result.maxS2SplitNum = std::max(result.maxS2SplitNum, assignContext.curKvSplitPart);
578+ // 若存在头归约，则切分点一定为上一个核结束的位置
579+ result.fdRes.bN2IdxOfFdHead[result.numOfFdHead] = result.bN2End[assignContext.curCoreIdx - 1U];
580+ result.fdRes.gS1IdxOfFdHead[result.numOfFdHead] = result.gS1End[assignContext.curCoreIdx - 1U];
581+ result.fdRes.s2SplitNumOfFdHead[result.numOfFdHead] = assignContext.curKvSplitPart;
582+ result.fdRes.gS1SplitNumOfFdHead[result.numOfFdHead] = curFdS1gSplitPart;
583+ result.fdRes.gS1LastPartSizeOfFdHead[result.numOfFdHead] = curFdS1gLastPartSize;
584+ result.numOfFdHead++;
585+}
586+ 
587+void CalcSplitPlan(uint32_t coreNum, int64_t costLimit, const SplitContext &splitContext, SplitResult &result)
588+{
589+ const CostInfo &costInfo = splitContext.costInfo;
590+ 
591+ if (coreNum == 0U) {
592+ return;
593+ }
594+ result.maxCost = 0U;
595+ result.usedCoreNum = 0U;
596+ 
597+ AssignContext assignContext {};
598+ assignContext.curBIdx = 0U;
599+ assignContext.curS1GIdx = 0U;
600+ assignContext.unassignedCost = costInfo.totalCost;
601+ assignContext.bN2Cost = costInfo.bN2CostOfEachBatch[assignContext.curBIdx];
602+ assignContext.bN2Block = costInfo.bN2BlockOfEachBatch[assignContext.curBIdx];
603+ CalcBatchCache(assignContext.curBIdx, splitContext, assignContext.batchCache);
604+ CalcS1GCache(assignContext.curS1GIdx, splitContext, assignContext.batchCache, assignContext.s1GCache);
605+ assignContext.curS2Idx = assignContext.s1GCache.s2Start;
606+ 
607+ for (uint32_t i = 0; i < coreNum; ++i) {
608+ if (result.maxCost > costLimit) {
609+ return;
610+ }
611+ if (assignContext.isFinished || assignContext.unassignedCost <= 0) {
612+ break;
613+ }
614+ 
615+ assignContext.curCoreIdx = i;
616+ result.fdRes.s2SplitStartIdxOfCore[assignContext.curCoreIdx] = assignContext.curKvSplitPart - 1U;
617+ 
618+ assignContext.coreCache = {};
619+ assignContext.coreCache.costLimit = assignContext.unassignedCost / (coreNum - assignContext.curCoreIdx);
620+ 
621+ // 1、按整batch分配
622+ AssignByBatch(splitContext, assignContext);
623+ // 2、按行分配
624+ AssignByRow(splitContext, assignContext);
625+ // 3、按块分配
626+ AssignByBlock(splitContext, assignContext);
627+ // 4、强制分配
628+ if (assignContext.coreCache.block == 0) {
629+ ForceAssign(splitContext, assignContext);
630+ }
631+ 
632+ result.bN2End[i] = assignContext.curBN2Idx;
633+ result.gS1End[i] = assignContext.curS1GIdx;
634+ result.s2End[i] = assignContext.curS2Idx;
635+ result.maxCost = std::max(result.maxCost, assignContext.coreCache.cost);
636+ 
637+ assignContext.unassignedCost -= assignContext.coreCache.cost;
638+ 
639+ // 对之前的归约信息进行记录并清理
640+ if (IsNeedRecordFDInfo(assignContext, result)) {
641+ RecordFDInfo(splitContext, assignContext, result);
642+ assignContext.curKvSplitPart = 1U;
643+ }
644+ 
645+ // 更新S2切分信息
646+ if (assignContext.curS2Idx > assignContext.s1GCache.s2Start &&
647+ assignContext.curS2Idx <= assignContext.s1GCache.s2End) {
648+ assignContext.curKvSplitPart++;
649+ }
650+ }
651+ 
652+ result.usedCoreNum = assignContext.curCoreIdx + 1;
653+}
654+ 
655+void SplitFD(SplitResult &result)
656+{
657+ uint32_t totalFDLoad = 0;
658+ uint32_t totalFDHeadSplit = 0;
659+ // 计算FD的总数据量
660+ for (uint32_t i = 0; i < result.numOfFdHead; i++) {
661+ totalFDLoad += result.fdRes.s2SplitNumOfFdHead[i] * result.fdRes.gS1SplitNumOfFdHead[i];
662+ totalFDHeadSplit += result.fdRes.gS1SplitNumOfFdHead[i];
663+ }
664+ 
665+ // 基于FA开核数量，计算每个Vector需要计算的FD数据量
666+ // FD均衡的最小单位为一个归约任务的一个split，所以最多占用totalFDHeadSplit个vector
667+ uint32_t maxVectorNum = std::min(totalFDHeadSplit, result.usedCoreNum * result.vecCubeRatio);
668+ double loadThrOfVector = static_cast<double>(totalFDLoad) / static_cast<double>(maxVectorNum); // 初始化vector的负载上限
669+ int64_t loadOfCurVector = 0;
670+ uint32_t curCoreIndex = 0;
671+ uint32_t preTmpFDIndexEndOfFdHead = 0;
672+ uint32_t preTmpFDIndexEndOfFdHeadSplit = 0;
673+ for (uint32_t i = 0; i < result.numOfFdHead; i++) {
674+ uint32_t fDKVSplitNum = result.fdRes.s2SplitNumOfFdHead[i];
675+ for (uint32_t gS1SplitIdx = 0; gS1SplitIdx < result.fdRes.gS1SplitNumOfFdHead[i]; gS1SplitIdx++) {
676+ double remainSpace = loadThrOfVector - static_cast<double>(loadOfCurVector); // 计算当前vector剩余负载空间
677+ // 判断是否放在当前vector的标准是剩余空间是否能容纳一半当前归约块
678+ if (fDKVSplitNum > remainSpace * FD_TOLERANCE_RATIO) {
679+ result.fdRes.gS1IdxEndOfFdHead[curCoreIndex] = preTmpFDIndexEndOfFdHead;
680+ result.fdRes.gS1IdxEndOfFdHeadSplit[curCoreIndex] = preTmpFDIndexEndOfFdHeadSplit;
681+ curCoreIndex += 1U;
682+ totalFDLoad -= static_cast<uint32_t>(loadOfCurVector); // 当前未分配的总负载
683+ // 根据剩余负载和剩余可用vector更新负载上限，保证最后一个vector能分配所有负载
684+ loadThrOfVector = static_cast<double>(totalFDLoad) / static_cast<double>(maxVectorNum - curCoreIndex);
685+ loadOfCurVector = 0;
686+ }
687+ loadOfCurVector += fDKVSplitNum;
688+ preTmpFDIndexEndOfFdHead = i;
689+ preTmpFDIndexEndOfFdHeadSplit = gS1SplitIdx;
690+ }
691+ }
692+ result.fdRes.gS1IdxEndOfFdHead[curCoreIndex] = preTmpFDIndexEndOfFdHead;
693+ result.fdRes.gS1IdxEndOfFdHeadSplit[curCoreIndex] = preTmpFDIndexEndOfFdHeadSplit;
694+ result.usedVecNumOfFd = curCoreIndex + 1;
695+}
696+ 
697+void SplitCore(uint32_t coreNum, const BaseInfo &baseInfo, const SplitParam &param, SplitResult &result)
698+{
699+ SplitContext splitContext(baseInfo, param);
700+ 
701+ // 1、划分基本块，统计信息
702+ CalcSplitInfo(splitContext);
703+ // 全空case
704+ if (splitContext.splitInfo.isKvSeqAllZero) {
705+ result.usedCoreNum = 1U;
706+ result.bN2End[0] = baseInfo.bSize * baseInfo.n2Size;
707+ result.gS1End[0] = 0U;
708+ result.s2End[0] = 0U;
709+ return;
710+ }
711+ 
712+ CalcCostInfo(splitContext);
713+ 
714+ // 2、获取每个核的分配方案
715+ uint32_t maxCore = std::min(coreNum, splitContext.costInfo.totalBlockNum);
716+ uint32_t minCore = static_cast<uint32_t>(
717+ std::sqrt(static_cast<float>(splitContext.costInfo.totalBlockNum) + 0.25f) + 0.5f);
718+ minCore = std::min(minCore, maxCore);
719+ 
720+ result.maxCost = INT64_MAX;
721+ result.usedCoreNum = 1U;
722+ 
723+ SplitResult tmpResult {coreNum, result.vecCubeRatio};
724+ for (uint32_t i = minCore; i <= maxCore; ++i) {
725+ CalcSplitPlan(i, result.maxCost, splitContext, tmpResult);
726+ if (tmpResult.maxCost < result.maxCost) {
727+ CopyTmpResult(tmpResult, result);
728+ }
729+ ClearTmpResult(tmpResult);
730+ }
731+ 
732+ // 3、存在FD任务，对FD进行负载均衡分配
733+ if (result.numOfFdHead > 0U) {
734+ SplitFD(result);
735+ }
736+ result.usedCoreNum = std::max(result.usedCoreNum, 1U); // 至少使用1个core
737+ 
738+ RollBackCursor(baseInfo, splitContext, splitContext.costInfo, result);
739+}
740+ 
741+}
742+}
Aexperimental/attention/sparse_flash_attention_antiquant/op_host/split_core.h+282-0
@@ -0,0 +1,282 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file split_core.h
13+ * \brief
14+ */
15+ 
16+#ifndef SPLIT_CORE_H
17+#define SPLIT_CORE_H
18+ 
19+#include <cstdint>
20+#include <vector>
21+#include <array>
22+#include <algorithm>
23+ 
24+namespace optiling {
25+namespace sfaa {
26+constexpr int64_t FA_TOLERANCE_RATIO = 2;
27+constexpr uint32_t FD_TOLERANCE_RATIO = 2U;
28+ 
29+enum BlockType : uint32_t {
30+ NORMAL_BLOCK = 0,
31+ TAIL_BLOCK,
32+ BLOCK_MAX_TYPE
33+};
34+ 
35+enum class SparseMode : uint8_t {
36+ DEFAULT_MASK = 0,
37+ ALL_MASK,
38+ LEFT_UP_CAUSAL,
39+ RIGHT_DOWN_CAUSAL,
40+ BAND,
41+ SPARSE_BUTT,
42+};
43+ 
44+template<class T>
45+using Range = std::pair<T, T>;
46+ 
47+template<class T>
48+using BlockCost = std::array<std::array<T, static_cast<size_t>(BLOCK_MAX_TYPE)>, static_cast<size_t>(BLOCK_MAX_TYPE)>;
49+ 
50+template<typename T>
51+T Clip(T value, T minValue, T maxValue)
52+{
53+ if (value < minValue) {
54+ return minValue;
55+ }
56+ if (value > maxValue) {
57+ return maxValue;
58+ }
59+ return value;
60+}
61+ 
62+template<typename T>
63+inline bool IsWithinTolerance(T limit, T tolerance, T value)
64+{
65+ return limit + tolerance >= value;
66+}
67+ 
68+// 分核功能模块输入：输入case的基本信息
69+struct BaseInfo {
70+ uint32_t bSize { 0U };
71+ uint32_t n2Size { 0U };
72+ uint32_t gSize { 0U };
73+ uint32_t s1Size { 0U };
74+ uint32_t s2Size { 0U };
75+ bool isS1G { true };
76+ bool isAccumSeqS1 { false };
77+ bool isAccumSeqS2 { false };
78+ std::vector<int64_t> actualSeqS1Size {};
79+ std::vector<int64_t> actualSeqS2Size {};
80+ uint32_t actualLenQDims { 0U };
81+ uint32_t actualLenKvDims { 0U };
82+ bool attenMaskFlag { false };
83+ int32_t sparseMode { 0U };
84+ int64_t preToken { 0 };
85+ int64_t nextToken { 0 };
86+ int64_t actualSeqPrefixSize { 0 };
87+ 
88+ uint32_t actualLenSparseDims { 0U };
89+ std::vector<int64_t> actualSeqSparseSize {};
90+ int64_t sparseBlockSize = { 0 };
91+ int64_t sparseBlockCount = { 0 };
92+ uint32_t sparseShardSize = { 0 };
93+};
94+ 
95+// 分核功能模块输入：切分属性，预留接口，可作为切分方案的参数入口
96+struct SplitParam {
97+ uint32_t mBaseSize { 1U };
98+ uint32_t s2BaseSize { 1U };
99+ uint32_t gS1BaseSizeOfFd { 8U }; // FD阶段分核，m轴切分基本块大小
100+};
101+ 
102+ 
103+// 分核功能模块输出：FD信息，包含需要归约的数据索引及其分核信息
104+struct FlashDecodeResult {
105+ // 1、归约任务的索引信息
106+ std::vector<uint32_t> bN2IdxOfFdHead {}; // 每个归约任务的BN2索引，脚标为归约任务的序号，最大为核数-1
107+ std::vector<uint32_t> gS1IdxOfFdHead {}; // 每个归约任务的GS1索引，脚标为归约任务的序号
108+ std::vector<uint32_t> s2SplitNumOfFdHead {}; // 每个归约任务的S2核间切分份数，脚标为归约任务的序号
109+ // 2、FD负载均衡阶段，归约任务的分核（vec）信息
110+ std::vector<uint32_t> gS1SplitNumOfFdHead {}; // 每个归约任务m轴切分份数，脚标为归约任务的序号
111+ std::vector<uint32_t> gS1LastPartSizeOfFdHead {}; // 每个归约任务m轴切分的最后一份的大小，脚标为归约任务的序号
112+ std::vector<uint32_t> gS1IdxEndOfFdHead {}; // FD负载均衡阶段，每个vector的一级索引，脚标为vector ID，值为归约任务的ID
113+ std::vector<uint32_t> gS1IdxEndOfFdHeadSplit {}; // FD负载均衡阶段，每个vector的二级索引，脚标为vector ID，值为归约任务的m轴切分ID
114+ // 3、每个core处理的第1个归约任务的数据应存放的workspace位置
115+ std::vector<uint32_t> s2SplitStartIdxOfCore {};
116+ 
117+ FlashDecodeResult(uint32_t coreNum, uint32_t vecCubeRatio) :
118+ bN2IdxOfFdHead(coreNum),
119+ gS1IdxOfFdHead(coreNum),
120+ s2SplitNumOfFdHead(coreNum),
121+ gS1SplitNumOfFdHead(coreNum),
122+ gS1LastPartSizeOfFdHead(coreNum),
123+ gS1IdxEndOfFdHead(coreNum * vecCubeRatio),
124+ gS1IdxEndOfFdHeadSplit(coreNum * vecCubeRatio),
125+ s2SplitStartIdxOfCore(coreNum) {}
126+};
127+ 
128+// 分核功能模块输出：FA阶段的核间分核信息
129+struct SplitResult {
130+ uint32_t usedCoreNum { 0U }; // 使用的核数量
131+ uint32_t vecCubeRatio { 0U }; // vec 与 cube 核数比例
132+ std::vector<uint32_t> bN2End {}; // 每个核处理数据的BN2结束点
133+ std::vector<uint32_t> gS1End {}; // 每个核处理数据的GS1结束点
134+ std::vector<uint32_t> s2End {}; // 每个核处理数据的S2结束点
135+ int64_t maxCost { 0 }; // 慢核开销
136+ uint32_t numOfFdHead { 0U }; // 归约任务数量
137+ uint32_t maxS2SplitNum { 0U }; // 单个归约任务最大分核数量
138+ uint32_t usedVecNumOfFd { 0U }; // 归约过程使用的vector数量
139+ FlashDecodeResult fdRes { 0U, 0U }; // FD信息
140+ 
141+ SplitResult(uint32_t coreNum, uint32_t ratio) :
142+ bN2End(coreNum),
143+ vecCubeRatio(ratio),
144+ gS1End(coreNum),
145+ s2End(coreNum),
146+ fdRes(coreNum, ratio) {};
147+};
148+ 
149+ 
150+// 分核功能模块内部使用：记录切分信息
151+struct SplitInfo {
152+ std::vector<uint32_t> s1GBaseNum {}; // S1G方向，切了多少个基本块
153+ std::vector<uint32_t> s2BaseNum {}; // S2方向，切了多少个基本块
154+ std::vector<uint32_t> s1GTailSize {}; // S1G方向，尾块size
155+ std::vector<uint32_t> s2TailSize {}; // S2方向，尾块size
156+ bool isKvSeqAllZero { true };
157+ 
158+ explicit SplitInfo(uint32_t batchSize) :
159+ s1GBaseNum(batchSize),
160+ s2BaseNum(batchSize),
161+ s1GTailSize(batchSize),
162+ s2TailSize(batchSize) {}
163+};
164+ 
165+// 分核功能模块内部使用：记录batch的开销信息
166+struct CostInfo {
167+ std::vector<int64_t> bN2CostOfEachBatch {}; // 整个batch的开销
168+ std::vector<uint32_t> bN2BlockOfEachBatch {}; // 整个batch的开销
169+ std::vector<int64_t> bN2LastBlockCostOfEachBatch {}; // batch最后一块的开销
170+ uint32_t totalBlockNum { 0U };
171+ int64_t totalCost { 0 };
172+ 
173+ explicit CostInfo(uint32_t batchSize) :
174+ bN2CostOfEachBatch(batchSize),
175+ bN2BlockOfEachBatch(batchSize),
176+ bN2LastBlockCostOfEachBatch(batchSize) {}
177+};
178+ 
179+// 分核功能模块内部使用：分核过程中，case基本信息的上下文信息，组合以减少接口传参数量
180+struct SplitContext {
181+ const BaseInfo &baseInfo {};
182+ const SplitParam &splitParam {};
183+ SplitInfo splitInfo { 0U };
184+ CostInfo costInfo { 0U };
185+ 
186+ explicit SplitContext(const BaseInfo &info, const SplitParam &param) :
187+ baseInfo(info),
188+ splitParam(param),
189+ splitInfo(info.bSize),
190+ costInfo(info.bSize) {}
191+};
192+ 
193+// 分核功能模块内部使用：记录batch相关的临时信息
194+struct BatchCache {
195+ uint32_t bIdx { 0U };
196+ uint32_t s1Size { 0U };
197+ uint32_t s2Size { 0U };
198+ int64_t preTokenLeftUp { 0 };
199+ int64_t nextTokenLeftUp { 0 };
200+ BlockCost<int64_t> typeCost {};
201+};
202+ 
203+// 分核功能模块内部使用：记录当前行（S1G）的临时信息
204+struct S1GCache {
205+ uint32_t bIdx { 0U };
206+ uint32_t s1GIdx { 0U };
207+ uint32_t s2Start { 0U };
208+ uint32_t s2End { 0U };
209+ int64_t s1GCost { 0 };
210+ int64_t s1GLastBlockCost { 0 };
211+ uint32_t s1GBlock { 0U };
212+ int64_t s1GNormalBlockCost { 0 };
213+};
214+ 
215+// 分核功能模块内部使用：记录分配过程中，当前核的负载信息
216+struct CoreCache {
217+ int64_t costLimit { 0 }; // 负载上限
218+ int64_t cost { 0 }; // 已分配负载
219+ uint32_t block { 0U }; // 已分配块数
220+};
221+ 
222+// 分核功能模块内部使用：记录分配过程中的上下文信息
223+struct AssignContext {
224+ uint32_t curBIdx { 0U };
225+ uint32_t curBN2Idx { 0U };
226+ uint32_t curS1GIdx { 0U };
227+ uint32_t curS2Idx { 0U };
228+ uint32_t curCoreIdx { 0U };
229+ int64_t unassignedCost { 0 };
230+ uint32_t usedCoreNum { 0U };
231+ uint32_t curKvSplitPart { 1U };
232+ 
233+ int64_t bN2Cost { 0 };
234+ uint32_t bN2Block { 0U };
235+ bool isFinished { false };
236+ BatchCache batchCache {};
237+ S1GCache s1GCache {};
238+ CoreCache coreCache {};
239+};
240+ 
241+// util
242+uint32_t GetS1SeqSize(uint32_t bIdx, const BaseInfo &baseInfo);
243+uint32_t GetS2SeqSize(uint32_t bIdx, const BaseInfo &baseInfo);
244+uint32_t GetSparseSeqSize(uint32_t bIdx, const BaseInfo &baseInfo);
245+int64_t CalcPreTokenLeftUp(uint32_t s1Size, uint32_t s2Size, const BaseInfo &baseInfo);
246+int64_t CalcNextTokenLeftUp(uint32_t s1Size, uint32_t s2Size, const BaseInfo &baseInfo);
247+Range<uint32_t> CalcS2Range(uint32_t s1GIdx, const BaseInfo &baseInfo, const SplitParam &splitParam,
248+ const BatchCache &batchCache);
249+int64_t CalcCost(uint32_t basicM, uint32_t basicS2);
250+BlockCost<int64_t> CalcCostTable(uint32_t s1NormalSize, uint32_t s2NormalSize, uint32_t s1GTailSize,
251+ uint32_t s2TailSize);
252+ 
253+// cache calculation
254+void CalcBatchCache(uint32_t bIdx, const SplitContext &splitContext, BatchCache &batchCache);
255+void CalcS1GCache(uint32_t s1GIdx, const SplitContext &splitContext, const BatchCache &batchCache, S1GCache &s1GCache);
256+void CopyTmpResult(SplitResult &tmpRes, SplitResult &splitRes);
257+void ClearTmpResult(SplitResult &tmpResult);
258+ 
259+// preprocess
260+void CalcSplitInfo(SplitContext &splitContext);
261+void CalcBatchCost(uint32_t bIdx, const SplitContext &splitContext, CostInfo &costInfo);
262+void CalcCostInfo(SplitContext &splitContext);
263+ 
264+// assign
265+void UpdateCursor(const SplitContext &splitContext, AssignContext &assignContext);
266+void AssignByBatch(const SplitContext &splitContext, AssignContext &assignContext);
267+void AssignByRow(const SplitContext &splitContext, AssignContext &assignContext);
268+void AssignByBlock(const SplitContext &splitContext, AssignContext &assignContext);
269+void ForceAssign(const SplitContext &splitContext, AssignContext &assignContext);
270+ 
271+// FD
272+bool IsNeedRecordFDInfo(const AssignContext &assignContext, const SplitResult &splitRes);
273+void RecordFDInfo(const SplitContext &splitContext, const AssignContext &assignContext, SplitResult &result);
274+ 
275+// main
276+void SplitFD(SplitResult &result);
277+void CalcSplitPlan(uint32_t coreNum, int64_t costLimit, const SplitContext &splitContext, SplitResult &result);
278+void SplitCore(uint32_t coreNum, const BaseInfo &baseInfo, const SplitParam &splitParam, SplitResult &result);
279+void RollBackCursor(const BaseInfo &baseInfo, const SplitContext &splitContext, const CostInfo &costInfo, SplitResult &splitRes);
280+}
281+}
282+#endif
Aexperimental/attention/sparse_flash_attention_antiquant/op_host/split_core_BN.h+251-0
@@ -0,0 +1,251 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/* !
12+ * \file split_core_BN.h
13+ * \brief
14+ */
15+ 
16+#ifndef SPLIT_CORE_BN_H
17+#define SPLIT_CORE_BN_H
18+ 
19+ 
20+namespace optiling {
21+ 
22+struct BaseInfo {
23+ uint32_t bSize;
24+ uint32_t n2Size;
25+ uint32_t gSize;
26+ uint32_t s1Size = 0;
27+ uint32_t s2Size = 0;
28+ bool isAccumSeqS1 = false;
29+ bool isAccumSeqS2 = false;
30+ bool sliding = false;
31+ const int64_t *actualSeqS1Size = nullptr;
32+ const int64_t *actualSeqS2Size = nullptr;
33+ uint32_t actualLenDimsQ = 0;
34+ uint32_t actualLenDimsKV = 0;
35+ int64_t preToken = 0;
36+ int64_t nextToken = 0;
37+};
38+ 
39+struct InnerSplitParams {
40+ uint32_t s1GBaseSize = 1;
41+ uint32_t s2BaseSize = 1;
42+};
43+ 
44+struct OuterSplitParams {
45+ uint32_t *bN2End;
46+ uint32_t *gS1End;
47+ uint32_t *s2End;
48+};
49+ 
50+struct FlashDecodeParams {
51+ uint32_t *bN2IdxOfFdHead;
52+ uint32_t *gS1IdxOfFdHead;
53+ uint32_t *s2SplitNumOfFdHead;
54+ uint32_t *s2SplitStartIdxOfCore;
55+ uint32_t gS1BaseSizeOfFd;
56+ uint32_t *gS1SplitNumOfFdHead;
57+ uint32_t *gS1LastPartSizeOfFdHead;
58+ uint32_t *gS1IdxEndOfFdHead;
59+ uint32_t *gS1IdxEndOfFdHeadSplit;
60+};
61+ 
62+struct SplitCoreRes {
63+ uint32_t numOfFdHead;
64+ uint32_t maxS2SplitNum;
65+ uint32_t usedCoreNum;
66+ uint32_t usedVecNumOfFd;
67+};
68+ 
69+int64_t SfaClipSInnerTokenCube(int64_t sInnerToken, int64_t minValue, int64_t maxValue)
70+{
71+ sInnerToken = sInnerToken > minValue ? sInnerToken : minValue;
72+ sInnerToken = sInnerToken < maxValue ? sInnerToken : maxValue;
73+ return sInnerToken;
74+}
75+
76+
77+void SfaUpDateInnerLoop(const BaseInfo &baseInfo, const InnerSplitParams &innerSplitParams, uint32_t &s2Start, uint32_t &s2End,
78+ uint32_t sOuterLoopIdx, int64_t curActualSeqLen,int64_t preTokensLeftUp, int64_t nextTokensLeftUp, uint32_t s2Loop){
79+ if (!baseInfo.sliding) {
80+ return;
81+ }
82+ uint32_t sOuterSize = innerSplitParams.s1GBaseSize / baseInfo.gSize; // BNSD需要额外适配
83+ uint32_t sOuterOffset = sOuterLoopIdx * sOuterSize;
84+ int64_t sInnerFirstToken = SfaClipSInnerTokenCube(static_cast<int64_t>(sOuterOffset) - preTokensLeftUp,
85+ 0, curActualSeqLen);
86+ int64_t s2StartIdx = sInnerFirstToken / static_cast<int64_t>(innerSplitParams.s2BaseSize);
87+ if (s2StartIdx <= 0) {
88+ s2StartIdx = 0;
89+ }
90+ s2Start = s2StartIdx;
91+ int64_t sInnerLastToken = SfaClipSInnerTokenCube(static_cast<int64_t>(sOuterOffset) + nextTokensLeftUp +
92+ static_cast<int64_t>(sOuterSize), 0, curActualSeqLen);
93+ s2End= (sInnerLastToken + static_cast<int64_t>(innerSplitParams.s2BaseSize) - 1) /
94+ static_cast<int64_t>(innerSplitParams.s2BaseSize);
95+ if (s2End > s2Loop) {
96+ s2End = s2Loop;
97+ }
98+}
99+
100+void SfaGetPreNextTokensLeftUp(const BaseInfo &baseInfo,
101+ int64_t actualSeqLength, int64_t actualSeqLengthKV, int64_t& preTokensLeftUp, int64_t& nextTokensLeftUp) {
102+ if (baseInfo.sliding) {
103+ preTokensLeftUp = baseInfo.preToken - actualSeqLengthKV + actualSeqLength;
104+ nextTokensLeftUp = baseInfo.nextToken + actualSeqLengthKV - actualSeqLength;
105+ }
106+}
107+
108+uint32_t SfaGetCalcBlockNumsOneHead(const BaseInfo &baseInfo, const InnerSplitParams &innerSplitParams, int64_t outerBlockNums, int64_t innerBlockNums,
109+ int64_t actualSeqLength, int64_t actualSeqLengthKV, int64_t preTokensLeftUp, int64_t nextTokensLeftUp) {
110+ if (!baseInfo.sliding) { // 会不会影响原来
111+ return innerBlockNums * outerBlockNums;
112+ } else {
113+ uint32_t toCalcBlockNums = 0;
114+ for (uint32_t s1OuterIdx = 0; s1OuterIdx < outerBlockNums; s1OuterIdx++) {
115+ uint32_t sInnerIndexStart = 0;
116+ uint32_t sInnerIndexEnd = 0;
117+ SfaUpDateInnerLoop(baseInfo, innerSplitParams, sInnerIndexStart, sInnerIndexEnd,
118+ s1OuterIdx, actualSeqLengthKV,preTokensLeftUp, nextTokensLeftUp, innerBlockNums);
119+ toCalcBlockNums += (sInnerIndexEnd - sInnerIndexStart);
120+ }
121+ return toCalcBlockNums;
122+ }
123+}
124+
125+void SfaGetActualSeqLength(const BaseInfo &baseInfo, uint32_t &actualSeqLengths, uint32_t &actualSeqLengthsKV, uint32_t bIdx) {
126+ actualSeqLengths = baseInfo.s1Size;
127+ actualSeqLengthsKV = baseInfo.s2Size;
128+}
129+ 
130+void SfaSplitCore(const BaseInfo &baseInfo, const InnerSplitParams &innerSplitParams, uint32_t coreNum, OuterSplitParams outerSplitParams) {
131+ std::vector<uint32_t> s1GBaseNum(baseInfo.bSize); // S1G方向，切了多少个基本块
132+ std::vector<uint32_t> s2BaseNum(baseInfo.bSize); // S2方向，切了多少个基本块
133+ std::vector<uint32_t> s1Size(baseInfo.bSize);
134+ std::vector<uint32_t> s2Size(baseInfo.bSize);
135+ bool seqZeroFlag = true;
136+ for (uint32_t bIdx = 0; bIdx < baseInfo.bSize; bIdx++) {
137+ SfaGetActualSeqLength(baseInfo, s1Size[bIdx], s2Size[bIdx], bIdx);
138+ if (s1Size[bIdx] > 0) {
139+ seqZeroFlag = false;
140+ }
141+ }
142+ // 计算总基本块数
143+ uint32_t totalBaseNum = 0;
144+ for (uint32_t bIdx = 0; bIdx < baseInfo.bSize; bIdx++) {
145+ //SfaGetActualSeqLength(baseInfo,s1Size[bIdx], s2Size[bIdx], bIdx);
146+ if (seqZeroFlag == false) {
147+ s1GBaseNum[bIdx] = (s1Size[bIdx]* baseInfo.gSize + (innerSplitParams.s1GBaseSize - 1)) / innerSplitParams.s1GBaseSize;
148+ } else {
149+ s1GBaseNum[bIdx] = 1;
150+ }
151+
152+ s2BaseNum[bIdx] = 1;
153+ int64_t preTokensLeftUp = 0;
154+ int64_t nextTokensLeftUp = 0;
155+ SfaGetPreNextTokensLeftUp(baseInfo, s1Size[bIdx], s2Size[bIdx], preTokensLeftUp, nextTokensLeftUp);
156+ totalBaseNum += SfaGetCalcBlockNumsOneHead(baseInfo,innerSplitParams,s1GBaseNum[bIdx], s2BaseNum[bIdx],
157+ s1Size[bIdx], s2Size[bIdx], preTokensLeftUp, nextTokensLeftUp) * baseInfo.n2Size;
158+ }
159+ uint32_t avgBaseNum = 1;
160+ if (totalBaseNum > coreNum) {
161+ avgBaseNum = (totalBaseNum + coreNum - 1) / coreNum;
162+ }
163+ 
164+ uint32_t accumBaseNum = 0; // 当前累积的基本块数
165+ uint32_t targetBaseNum = 0;
166+ uint32_t currCoreIdx = 0;
167+ uint32_t lastValidBIdx = 0;
168+ // res.numOfFdHead = 0;
169+ // res.maxS2SplitNum = 1;
170+ // fDParams.s2SplitStartIdxOfCore[0] = 0; //每核头块所处当前线段被切的第几部分
171+ //分核流程，保存分核数据
172+ for (uint32_t bN2Idx = 0; bN2Idx < baseInfo.bSize * baseInfo.n2Size; bN2Idx++) {
173+ uint32_t bIdx = bN2Idx / baseInfo.n2Size;
174+ int64_t preTokensLeftUp = 0;
175+ int64_t nextTokensLeftUp = 0;
176+ SfaGetPreNextTokensLeftUp(baseInfo, s1Size[bIdx], s2Size[bIdx], preTokensLeftUp, nextTokensLeftUp);
177+ for (uint32_t s1GIdx = 0; s1GIdx < s1GBaseNum[bIdx]; s1GIdx++) {
178+ uint32_t sInnerIndexStart = 0;
179+ uint32_t sInnerIndexEnd = s2BaseNum[bIdx];
180+ SfaUpDateInnerLoop(baseInfo, innerSplitParams,sInnerIndexStart, sInnerIndexEnd,
181+ s1GIdx, s2Size[bIdx],preTokensLeftUp, nextTokensLeftUp, s2BaseNum[bIdx]);
182+ uint32_t currKvSplitPart = 1; // [B,N2,S1]确定后，S2被切了几份
183+ 
184+ for (uint32_t s2Idx = sInnerIndexStart; s2Idx < sInnerIndexEnd; s2Idx++) {
185+ accumBaseNum += 1;
186+ targetBaseNum = (currCoreIdx + 1) * avgBaseNum; // 计算当前的目标权重
187+ if (accumBaseNum >= targetBaseNum) {
188+ // 更新当前核的End分核信息
189+ outerSplitParams.bN2End[currCoreIdx] = bN2Idx;
190+ outerSplitParams.gS1End[currCoreIdx] = s1GIdx;
191+ outerSplitParams.s2End[currCoreIdx] = s2Idx;
192+ currCoreIdx += 1;
193+ }
194+ }
195+ }
196+ if ((s1GBaseNum[bIdx] > 0) && (s2BaseNum[bIdx] > 0)) {
197+ lastValidBIdx = bIdx;
198+ }
199+ }
200+ if (accumBaseNum < targetBaseNum) {
201+ // 更新最后一个核的End分核信息
202+ outerSplitParams.bN2End[currCoreIdx] = ((lastValidBIdx + 1) * (baseInfo.n2Size)) - 1;
203+ outerSplitParams.gS1End[currCoreIdx] = s1GBaseNum[lastValidBIdx] - 1;
204+ outerSplitParams.s2End[currCoreIdx] = s2BaseNum[lastValidBIdx] - 1;
205+ currCoreIdx += 1;
206+ }
207+ // res.usedCoreNum = currCoreIdx;
208+}
209+ 
210+void SfaSplitFD(SplitCoreRes &res, FlashDecodeParams fDParams, uint32_t coreNum)
211+{
212+ uint64_t totalFDLoad = 0;
213+ uint32_t totalFDHeadSplit = 0;
214+ // 计算FD的总数据量
215+ for (uint32_t i = 0; i < res.numOfFdHead; i++) {
216+ totalFDLoad += fDParams.s2SplitNumOfFdHead[i] * fDParams.gS1SplitNumOfFdHead[i];
217+ totalFDHeadSplit += fDParams.gS1SplitNumOfFdHead[i];
218+ }
219+ 
220+ // 基于FA开核数量，计算每个Vector需要计算的FD数据量
221+ uint32_t maxVectorNum = std::min(totalFDHeadSplit, coreNum * 2); // FD均衡的最小单位为一个归约任务的一个split，所以最多占用totalFDHeadSplit个vector
222+ double loadThrOfVector = static_cast<double>(totalFDLoad) / static_cast<double>(maxVectorNum); // 初始化vector的负载上限
223+ int64_t loadOfCurVector = 0;
224+ uint32_t curCoreIndex = 0;
225+ uint32_t preTmpFDIndexEndOfFdHead = 0;
226+ uint32_t preTmpFDIndexEndOfFdHeadSplit = 0;
227+ for (uint32_t i = 0; i < res.numOfFdHead; i++) {
228+ uint32_t fDKVSplitNum = fDParams.s2SplitNumOfFdHead[i]; // gs上多少次需要规约的次数
229+ for (uint32_t gS1SplitIdx = 0; gS1SplitIdx < fDParams.gS1SplitNumOfFdHead[i]; gS1SplitIdx++) {
230+ double remainSpace = loadThrOfVector - loadOfCurVector; // 计算当前vector剩余负载空间
231+ // 判断是否放在当前vector的标准是剩余空间是否能容纳一半当前归约块
232+ if (fDKVSplitNum > remainSpace * 2) {
233+ fDParams.gS1IdxEndOfFdHead[curCoreIndex] = preTmpFDIndexEndOfFdHead; // 记录寄几个上
234+ fDParams.gS1IdxEndOfFdHeadSplit[curCoreIndex] = preTmpFDIndexEndOfFdHeadSplit; // 记录有gs/8的end
235+ curCoreIndex += 1;
236+ totalFDLoad -= loadOfCurVector; // 当前未分配的总负载
237+ loadThrOfVector = static_cast<double>(totalFDLoad) / static_cast<double>(maxVectorNum - curCoreIndex); // 根据剩余负载和剩余可用vector更新负载上限，保证最后一个vector能分配所有负载
238+ loadOfCurVector = 0;
239+ }
240+ loadOfCurVector += fDKVSplitNum;
241+ preTmpFDIndexEndOfFdHead = i;
242+ preTmpFDIndexEndOfFdHeadSplit = gS1SplitIdx;
243+ }
244+ }
245+ fDParams.gS1IdxEndOfFdHead[curCoreIndex] = preTmpFDIndexEndOfFdHead;
246+ fDParams.gS1IdxEndOfFdHeadSplit[curCoreIndex] = preTmpFDIndexEndOfFdHeadSplit;
247+ res.usedVecNumOfFd = curCoreIndex + 1;
248+}
249+ 
250+}
251+#endif
Aexperimental/attention/sparse_flash_attention_antiquant/op_kernel/sparse_flash_attention_antiquant.cpp+83-0
@@ -0,0 +1,83 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+ /*!
12+ * \file sparse_flash_attention_antiquant.cpp
13+ * \brief
14+ */
15+ 
16+#include "kernel_operator.h"
17+#include "sparse_flash_attention_antiquant_template_tiling_key.h"
18+#include "sparse_flash_attention_antiquant_kernel_gqa.h"
19+#include "sparse_flash_attention_antiquant_metadata.h"
20+// #include "log.h"
21+ 
22+using namespace AscendC;
23+using namespace optiling::detail;
24+ 
25+ 
26+template <class T>
27+__inline__ __attribute__((always_inline)) __aicore__ void InitMetaData(const __gm__ uint8_t *p_metadata, T *metadata)
28+{
29+ constexpr uint64_t all_bytes = sizeof(T);
30+#if defined(ASCENDC_CPU_DEBUG) || defined(__DAV_C220_CUBE__) || defined(__DAV_C310_CUBE__) || defined(__DAV_310R6_CUBE__) || defined(__GET_CODE_CHANNEL__)
31+ copy_data_align64((uint8_t*)metadata, (__gm__ uint8_t *)p_metadata, all_bytes);
32+#else
33+ __ubuf__ uint8_t *metadata_in_ub = (__ubuf__ uint8_t *)get_imm(0);
34+ constexpr uint32_t len_burst = (all_bytes + 31) / 32;
35+ copy_gm_to_ubuf(((__ubuf__ uint8_t *)metadata_in_ub), p_metadata, 0, 1,len_burst, 0, 0);
36+ set_flag(PIPE_MTE2, PIPE_S, EVENT_ID0);
37+ wait_flag(PIPE_MTE2, PIPE_S, EVENT_ID0);
38+ copy_data_align64((uint8_t*)metadata, (__ubuf__ uint8_t *)metadata_in_ub, all_bytes);
39+#endif
40+}
41+ 
42+#define SFAA_OP_IMPL_GQA(templateClass, tilingdataClass, ...) \
43+ do { \
44+ templateClass<SFAAType<__VA_ARGS__>> op; \
45+ GET_TILING_DATA_WITH_STRUCT(tilingdataClass, tiling_data_in, tiling); \
46+ const tilingdataClass *__restrict tiling_data = &tiling_data_in; \
47+ SfaMetaData *__restrict meta_data = nullptr; \
48+ SfaMetaData metaDataTmp; \
49+ if (metaData != nullptr) { \
50+ InitMetaData<SfaMetaData>(metaData, &metaDataTmp); \
51+ meta_data = &metaDataTmp; \
52+ } \
53+ op.Init(query, key, value, sparseIndices, keyScale, valueScale, blocktable, \
54+ actualSeqLengthsQuery, actualSeqLengthsKV, sparseSeqLengthsKV, meta_data, \
55+ attentionOut, user, tiling_data, tiling, &tPipe); \
56+ op.Process(); \
57+ } while (0)
58+ 
59+template<int ATTENTION_MODE, int FLASH_DECODE, int LAYOUT_T, int KV_LAYOUT_T, int TEMPLATE_MODE>
60+ __global__ __aicore__ void
61+sparse_flash_attention_antiquant(__gm__ uint8_t *query, __gm__ uint8_t *key, __gm__ uint8_t *value,
62+ __gm__ uint8_t *sparseIndices, __gm__ uint8_t* keyScale, __gm__ uint8_t* valueScale,
63+ __gm__ uint8_t *blocktable, __gm__ uint8_t *actualSeqLengthsQuery,
64+ __gm__ uint8_t *actualSeqLengthsKV, __gm__ uint8_t *sparseSeqLengthsKV, __gm__ uint8_t *metaData,
65+ __gm__ uint8_t *attentionOut, __gm__ uint8_t *workspace, __gm__ uint8_t *tiling)
66+{
67+ KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_MIX_AIC_1_2);
68+ TPipe tPipe;
69+ __gm__ uint8_t *user = GetUserWorkspace(workspace);
70+ 
71+ if constexpr (ATTENTION_MODE == ATTENTION_GQA_MHA) {
72+ if constexpr (ORIG_DTYPE_QUERY == DT_FLOAT16 && ORIG_DTYPE_KEY == DT_INT8 &&
73+ ORIG_DTYPE_ATTENTION_OUT == DT_FLOAT16) {
74+ SFAA_OP_IMPL_GQA(SparseFlashAttentionAntiquantGqa, SparseFlashAttentionAntiquantTilingDataMla, half, int8_t,
75+ half, FLASH_DECODE, static_cast<SFAA_LAYOUT>(LAYOUT_T), static_cast<SFAA_LAYOUT>(KV_LAYOUT_T),
76+ TEMPLATE_MODE, false);
77+ } else { // bf16
78+ SFAA_OP_IMPL_GQA(SparseFlashAttentionAntiquantGqa, SparseFlashAttentionAntiquantTilingDataMla, bfloat16_t, int8_t,
79+ bfloat16_t, FLASH_DECODE, static_cast<SFAA_LAYOUT>(LAYOUT_T), static_cast<SFAA_LAYOUT>(KV_LAYOUT_T),
80+ TEMPLATE_MODE, true);
81+ }
82+ }
83+}
Aexperimental/attention/sparse_flash_attention_antiquant/op_kernel/sparse_flash_attention_antiquant_common.h+452-0
@@ -0,0 +1,452 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file sparse_flash_attention_antiquant_common.h
13+ * \brief
14+ */
15+ 
16+#ifndef SPARSE_FLASH_ATTENTION_ANTIQUANT_COMMON_H
17+#define SPARSE_FLASH_ATTENTION_ANTIQUANT_COMMON_H
18+ 
19+#include "kernel_operator.h"
20+#include "lib/matmul_intf.h"
21+#include "lib/matrix/matmul/tiling.h"
22+ 
23+using namespace AscendC;
24+// 将isCheckTiling设置为false, 输入输出的max&sum&exp的shape为(m, 1)
25+constexpr SoftmaxConfig SFAA_SOFTMAX_FLASHV2_CFG_WITHOUT_BRC = {false, 0, 0, SoftmaxMode::SOFTMAX_OUTPUT_WITHOUT_BRC};
26+ 
27+enum class SFAA_LAYOUT {
28+ BSND = 0,
29+ TND = 1,
30+ PA_BSND = 2,
31+ PA_BNSD = 3,
32+ PA_NZ = 4,
33+};
34+ 
35+enum class QUANT_MODE {
36+ PER_CHANNEL = 0, // GQA支持
37+ PER_TOKEN_HEAD = 1, // GQA支持
38+ PER_TILE = 2, // MLA支持
39+};
40+ 
41+enum class ATTENTION_MODE {
42+ GQA_MHA = 0, // QKV headDim相等
43+ MLA_NAIVE = 1, // Dn=128, Dr=64
44+ MLA_ABSORB = 2, // Dn=512, Dr=64
45+};
46+ 
47+enum class QUANT_SCALE_REPO_MODE {
48+ SEPARATE = 0, // 分开存储
49+ COMBINE = 1, // 合并存储，量化模式是PER_TOKEN_HEAD/PER_TILE时支持COMBINE模式，参数顺序为：Nope+Rope+DequantScale
50+};
51+ 
52+template <typename Q_T, typename KV_T, typename OUT_T, const bool FLASH_DECODE = false,
53+ SFAA_LAYOUT LAYOUT_T = SFAA_LAYOUT::BSND, SFAA_LAYOUT KV_LAYOUT_T = SFAA_LAYOUT::BSND,
54+ const int TEMPLATE_MODE = C_TEMPLATE, const bool MSD_DD = false, typename... Args>
55+struct SFAAType {
56+ using queryType = Q_T;
57+ using kvType = KV_T;
58+ using kRopeType = Q_T;
59+ using outputType = OUT_T;
60+ static constexpr bool isMsdDD = MSD_DD;
61+ static constexpr bool flashDecode = FLASH_DECODE;
62+ static constexpr SFAA_LAYOUT layout = LAYOUT_T;
63+ static constexpr SFAA_LAYOUT kvLayout = KV_LAYOUT_T;
64+ static constexpr int templateMode = TEMPLATE_MODE;
65+ static constexpr bool pageAttention = (KV_LAYOUT_T == SFAA_LAYOUT::PA_BSND ||
66+ KV_LAYOUT_T == SFAA_LAYOUT::PA_BNSD ||
67+ KV_LAYOUT_T == SFAA_LAYOUT::PA_NZ);
68+};
69+ 
70+// ================================Util functions==================================
71+__aicore__ inline uint64_t SFAAAlign(uint64_t num, uint64_t rnd)
72+{
73+ return (((rnd) == 0) ? 0 : (((num) + (rnd) - 1) / (rnd) * (rnd)));
74+}
75+ 
76+__aicore__ inline uint64_t SfaaCeilDiv(uint64_t num, uint64_t rnd)
77+{
78+ return rnd == 0 ? 0 : (num + rnd-1) / rnd;
79+}
80+ 
81+template <typename T1, typename T2> __aicore__ inline T1 Min(T1 a, T2 b)
82+{
83+ return (a > b) ? (b) : (a);
84+}
85+ 
86+template <typename T> __aicore__ inline size_t BlockAlign(size_t s)
87+{
88+ if constexpr (IsSameType<T, int4b_t>::value) {
89+ return (s + 63) / 64 * 64;
90+ }
91+ size_t n = (32 / sizeof(T));
92+ return (s + n - 1) / n * n;
93+}
94+ 
95+struct FDparams {
96+ uint32_t *bN2IdxOfFdHead;
97+ uint32_t *gS1IdxOfFdHead;
98+ uint32_t *s2SplitNumOfFdHead;
99+ uint32_t *gS1SplitNumOfFdHead;
100+ uint32_t *gS1LastPartSizeOfFdHead;
101+ uint32_t *gS1IdxEndOfFdHead;
102+ uint32_t *gS1IdxEndOfFdHeadSplit;
103+ uint32_t usedVecNumOfFd;
104+ uint32_t gS1BaseSizeOfFd;
105+};
106+ 
107+struct RunInfo {
108+ uint32_t loop;
109+ uint32_t bIdx;
110+ uint32_t n2Idx;
111+ uint32_t gIdx;
112+ uint32_t s1Idx;
113+ uint32_t s2Idx;
114+ uint32_t bN2Idx;
115+ uint32_t curSInnerLoopTimes;
116+ uint64_t tndBIdxOffsetQ;
117+ uint64_t tndBIdxOffsetKV;
118+ uint64_t tensorAOffset;
119+ uint64_t tensorBOffset;
120+ uint64_t tensorARopeOffset;
121+ uint64_t tensorBRopeOffset;
122+ uint64_t attenOutOffset;
123+ uint64_t attenMaskOffset;
124+ uint64_t topKBaseOffset;
125+ uint64_t curS2BaseOffset;
126+ uint64_t nextS2BaseOffset;
127+ uint32_t actualSingleProcessSInnerSize;
128+ uint32_t actualSingleProcessSInnerSizeAlign;
129+ bool isFirstSInnerLoop;
130+ bool isChangeBatch;
131+ uint32_t s2BatchOffset;
132+ uint32_t gSize;
133+ uint32_t s1Size;
134+ uint32_t s2Size;
135+ uint32_t mSize;
136+ uint32_t mSizeV;
137+ uint32_t aicS2AccessSize;
138+ uint32_t aivS2AccessSize;
139+ uint32_t mSizeVStart;
140+ uint32_t tndIsS2SplitCore;
141+ uint32_t tndCoreStartKVSplitPos;
142+ bool isBmm2Output;
143+ bool isValid = false;
144+ 
145+ uint64_t actS1Size = 1;
146+ uint64_t curActualSeqLenOri = 0ULL;
147+ 
148+ uint32_t gS1Idx;
149+ uint64_t actS2Size = 1;
150+ uint32_t actMBaseSize;
151+ bool isLastS2Loop;
152+ int32_t nextTokensPerBatch = 0;
153+ int64_t threshold;
154+ uint32_t sparseBlockCount;
155+ int64_t topkGmBaseOffset = 0;
156+ uint32_t maskStart = 0;
157+ uint32_t maskEnd = 0;
158+};
159+ 
160+struct ConstInfo {
161+ // CUBE与VEC核间同步的模式
162+ static constexpr uint32_t SFAA_SYNC_MODE2 = 2;
163+ // BUFFER的字节数
164+ static constexpr uint32_t BUFFER_SIZE_BYTE_32B = 32;
165+ static constexpr uint32_t BUFFER_SIZE_BYTE_64B = 64;
166+ static constexpr uint32_t BUFFER_SIZE_BYTE_256B = 256;
167+ static constexpr uint32_t BUFFER_SIZE_BYTE_512B = 512;
168+ static constexpr uint32_t BUFFER_SIZE_BYTE_1K = 1024;
169+ static constexpr uint32_t BUFFER_SIZE_BYTE_2K = 2048;
170+ static constexpr uint32_t BUFFER_SIZE_BYTE_4K = 4096;
171+ static constexpr uint32_t BUFFER_SIZE_BYTE_8K = 8192;
172+ static constexpr uint32_t BUFFER_SIZE_BYTE_16K = 16384;
173+ static constexpr uint32_t BUFFER_SIZE_BYTE_32K = 32768;
174+ // FP32的0值和极大值
175+ static constexpr float FLOAT_ZERO = 0;
176+ static constexpr float FLOAT_MAX = 3.402823466e+38F;
177+ 
178+ // preLoad的总次数
179+ uint32_t preLoadNum = 0U;
180+ uint32_t nBufferMBaseSize = 0U;
181+ // CUBE和VEC的核间同步EventID
182+ uint32_t syncV1NupdateC2 = 0U;
183+ uint32_t syncV0C1 = 0U;
184+ uint32_t syncC1V1 = 0U;
185+ uint32_t syncV1C2 = 0U;
186+ uint32_t syncC2V2 = 0U;
187+ uint32_t syncC2V1 = 0U;
188+ 
189+ uint32_t mmResUbSize = 0U; // Matmul1输出结果GM上的大小
190+ uint32_t vec1ResUbSize = 0U; // Vector1输出结果GM上的大小
191+ uint32_t bmm2ResUbSize = 0U; // Matmul2输出结果GM上的大小
192+ uint64_t batchSize = 0ULL;
193+ uint64_t gSize = 0ULL;
194+ uint64_t qHeadNum = 0ULL;
195+ uint64_t kvHeadNum;
196+ uint64_t headDim;
197+ uint64_t headDimRope;
198+ uint64_t headDimAlign;
199+ uint64_t combineHeadDim; // quantScaleRepoMode为Combine模式时=headDim+headDimRope, 否则=headDim
200+ uint64_t kvSeqSize = 0ULL; // kv最大S长度
201+ uint64_t qSeqSize = 1ULL; // q最大S长度
202+ int64_t kvCacheBlockSize = 0; // PA场景的block size
203+ uint32_t maxBlockNumPerBatch = 0; // PA场景的最大单batch block number
204+ uint32_t splitKVNum = 0U; // S2核间切分的切分份数
205+ SFAA_LAYOUT outputLayout; // 输出的Transpose格式
206+ uint32_t sparseMode = 0;
207+ bool needInit = false;
208+ 
209+ // FlashDecoding
210+ uint32_t actualCombineLoopSize = 0U; // FlashDecoding场景, S2在核间切分的最大份数
211+ uint64_t combineLseOffset = 0ULL;
212+ uint64_t combineAccumOutOffset = 0ULL;
213+ 
214+ uint32_t actualLenDimsQ = 0U; // query的actualSeqLength 的维度
215+ uint32_t actualLenDimsKV = 0U; // KV 的actualSeqLength 的维度
216+ uint32_t sparseLenDimsKV = 0U; // TopK的实际K大小的维度
217+ 
218+ // TND
219+ uint32_t s2Start = 0U; // TND场景下，S2的起始位置
220+ uint32_t s2End = 0U; // 单核TND场景下S2循环index上限
221+ 
222+ uint32_t bN2Start = 0U;
223+ uint32_t bN2End = 0U;
224+ uint32_t gS1Start = 0U;
225+ uint32_t gS1End = 0U;
226+ 
227+ uint32_t tndFDCoreArrLen = 0U; // TNDFlashDecoding相关分核信息array的长度
228+ uint32_t coreStartKVSplitPos = 0U; // TNDFlashDecoding kv起始位置
229+ 
230+ uint32_t mBaseSize = 1ULL;
231+ uint32_t s2BaseSize = 1ULL;
232+ 
233+ // sparse attr
234+ int64_t sparseBlockSize = 0;
235+ uint32_t sparseBlockCount = 0;
236+ uint32_t sparseShardSize = 0; // 多少个S1共用同一组sparse_indices
237+ 
238+ // attention模式与量化模式
239+ ATTENTION_MODE attentionMode = ATTENTION_MODE::MLA_ABSORB;
240+ QUANT_MODE keyQuantMode = QUANT_MODE::PER_TILE;
241+ QUANT_MODE valueQuantMode = QUANT_MODE::PER_TILE;
242+ QUANT_SCALE_REPO_MODE quantScaleRepoMode = QUANT_SCALE_REPO_MODE::COMBINE;
243+ uint64_t tileSize = 128ULL;
244+};
245+ 
246+struct MSplitInfo {
247+ uint32_t nBufferIdx = 0U;
248+ uint32_t nBufferStartM = 0U;
249+ uint32_t nBufferDealM = 0U;
250+ uint32_t vecStartM = 0U;
251+ uint32_t vecDealM = 0U;
252+};
253+ 
254+ 
255+// BLOCK和REPEAT的字节数
256+constexpr uint64_t BYTE_BLOCK = 32UL;
257+constexpr uint32_t REPEAT_BLOCK_BYTE = 256U;
258+// BLOCK和REPEAT的FP32元素数
259+constexpr uint32_t FP32_BLOCK_ELEMENT_NUM = BYTE_BLOCK / sizeof(float);
260+constexpr uint32_t FP32_REPEAT_ELEMENT_NUM = REPEAT_BLOCK_BYTE / sizeof(float);
261+// repeat stride不能超过256
262+constexpr uint32_t REPEATE_STRIDE_UP_BOUND = 256;
263+constexpr int64_t HALF_NUM = 2;
264+constexpr int64_t STRIDE_LENGTH = 8;
265+constexpr int64_t MAX_VALID_LENGTH = 1024;
266+ 
267+template <typename T>
268+__aicore__ inline void RowMuls(LocalTensor<T> dstUb, LocalTensor<T> src0Ub, LocalTensor<T> src1Ub,
269+ uint32_t dealRowCount, uint32_t columnCount, uint32_t actualColumnCount)
270+{
271+ // muls by row, 每行的元素乘以相同的元素
272+ // dstUb[i, (j * 8) : (j * 8 + 7)] = src0Ub[i, (j * 8) : (j * 8 + 7)] * src1Ub[i, 0 : 7]
273+ // src0Ub:[dealRowCount, columnCount] src1Ub:[dealRowCount, FP32_BLOCK_ELEMENT_NUM] dstUb:[dealRowCount,
274+ // columnCount]
275+ // dealRowCount is repeat times, must be less 256
276+ uint32_t repeatElementNum = FP32_REPEAT_ELEMENT_NUM;
277+ uint32_t blockElementNum = FP32_BLOCK_ELEMENT_NUM;
278+ 
279+ if constexpr (std::is_same<T, half>::value) {
280+ // 此限制由于每个repeat至多连续读取256B数据
281+ repeatElementNum = FP32_REPEAT_ELEMENT_NUM * 2; // 256/4 * 2=128
282+ blockElementNum = FP32_BLOCK_ELEMENT_NUM * 2; // 32/4 * 2 = 16
283+ }
284+ 
285+ // 每次只能连续读取256B的数据进行计算，故每次只能处理256B/sizeof(dType)=
286+ // 列方向分dLoop次，每次处理8列数据
287+ uint32_t dLoop = actualColumnCount / repeatElementNum;
288+ uint32_t dRemain = actualColumnCount % repeatElementNum;
289+ // REPEATE_STRIDE_UP_BOUND=256， 此限制由于src0RepStride数据类型为uint8之多256个datablock间距
290+ if (columnCount < REPEATE_STRIDE_UP_BOUND * blockElementNum) {
291+ BinaryRepeatParams repeatParams;
292+ repeatParams.src0BlkStride = 1;
293+ repeatParams.src1BlkStride = 0;
294+ repeatParams.dstBlkStride = 1;
295+ repeatParams.src0RepStride = columnCount / blockElementNum;
296+ repeatParams.src1RepStride = 1;
297+ repeatParams.dstRepStride = columnCount / blockElementNum;
298+ 
299+ // 如果以列为repeat所处理的次数小于行处理次数，则以列方式处理。反之则以行进行repeat处理
300+ if (dLoop <= dealRowCount) {
301+ uint32_t offset = 0;
302+ for (uint32_t i = 0; i < dLoop; i++) {
303+ Mul(dstUb[offset], src0Ub[offset], src1Ub, repeatElementNum, dealRowCount, repeatParams);
304+ offset += repeatElementNum;
305+ }
306+ } else {
307+ BinaryRepeatParams columnRepeatParams;
308+ columnRepeatParams.src0BlkStride = 1;
309+ columnRepeatParams.src1BlkStride = 0;
310+ columnRepeatParams.dstBlkStride = 1;
311+ columnRepeatParams.src0RepStride = 8; // 列方向上两次repeat起始地址间隔dtypeMask=64个元素，即8个block
312+ columnRepeatParams.src1RepStride = 0;
313+ columnRepeatParams.dstRepStride = 8; // 列方向上两次repeat起始地址间隔dtypeMask=64个元素，即8个block
314+ for (uint32_t i = 0; i < dealRowCount; i++) {
315+ Mul(dstUb[i * columnCount], src0Ub[i * columnCount], src1Ub[i * blockElementNum], repeatElementNum,
316+ dLoop, columnRepeatParams);
317+ }
318+ }
319+ 
320+ // 最后一次完成[dealRowCount, dRemain] * [dealRowCount, blockElementNum] 只计算有效部分
321+ if (dRemain > 0) {
322+ Mul(dstUb[dLoop * repeatElementNum], src0Ub[dLoop * repeatElementNum], src1Ub, dRemain, dealRowCount,
323+ repeatParams);
324+ }
325+ } else {
326+ BinaryRepeatParams repeatParams;
327+ repeatParams.src0RepStride = 8; // 每个repeat为256B数据，正好8个datablock
328+ repeatParams.src0BlkStride = 1;
329+ repeatParams.src1RepStride = 0;
330+ repeatParams.src1BlkStride = 0;
331+ repeatParams.dstRepStride = 8;
332+ repeatParams.dstBlkStride = 1;
333+ // 每次计算一行，共计算dealRowCount行
334+ for (uint32_t i = 0; i < dealRowCount; i++) {
335+ // 计算一行中的dLoop个repeat, 每个repeat计算256/block_size 个data_block
336+ Mul(dstUb[i * columnCount], src0Ub[i * columnCount], src1Ub[i * blockElementNum], repeatElementNum, dLoop,
337+ repeatParams);
338+ // 计算一行中的尾块
339+ if (dRemain > 0) {
340+ Mul(dstUb[i * columnCount + dLoop * repeatElementNum],
341+ src0Ub[i * columnCount + dLoop * repeatElementNum], src1Ub[i * blockElementNum], dRemain, 1,
342+ repeatParams);
343+ }
344+ }
345+ }
346+}
347+ 
348+__aicore__ inline void MatDivsVec(LocalTensor<float> dstUb, LocalTensor<float> src0Ub, LocalTensor<float> src1Ub,
349+ uint32_t dealRowCount, uint32_t columnCount, uint32_t actualColumnCount)
350+{
351+ uint32_t dtypeMask = FP32_REPEAT_ELEMENT_NUM;
352+ uint32_t dLoop = actualColumnCount / dtypeMask;
353+ uint32_t dRemain = actualColumnCount % dtypeMask;
354+ 
355+ BinaryRepeatParams repeatParamsDiv;
356+ repeatParamsDiv.src0BlkStride = 1;
357+ repeatParamsDiv.src1BlkStride = 1;
358+ repeatParamsDiv.dstBlkStride = 1;
359+ repeatParamsDiv.src0RepStride = columnCount / FP32_BLOCK_ELEMENT_NUM;
360+ repeatParamsDiv.src1RepStride = 0;
361+ repeatParamsDiv.dstRepStride = columnCount / FP32_BLOCK_ELEMENT_NUM;
362+ uint32_t columnRepeatCount = dLoop;
363+ uint32_t offset = 0;
364+ for (uint32_t i = 0; i < dLoop; i++) {
365+ Div(dstUb[offset], src0Ub[offset], src1Ub[offset], dtypeMask, dealRowCount, repeatParamsDiv);
366+ offset += dtypeMask;
367+ }
368+ 
369+ if (dRemain > 0) {
370+ Div(dstUb[dLoop * dtypeMask], src0Ub[dLoop * dtypeMask], src1Ub[dLoop * dtypeMask], dRemain, dealRowCount, repeatParamsDiv);
371+ }
372+}
373+ 
374+__aicore__ inline void RowSub(LocalTensor<float> dstUb, LocalTensor<float> src0Ub, LocalTensor<float> src1Ub,
375+ uint32_t dealRowCount, uint32_t columnCount, uint32_t actualColumnCount)
376+{
377+ uint32_t dtypeMask = FP32_REPEAT_ELEMENT_NUM;
378+ uint32_t dLoop = actualColumnCount / dtypeMask;
379+ uint32_t dRemain = actualColumnCount % dtypeMask;
380+ 
381+ BinaryRepeatParams repeatParamsSub;
382+ repeatParamsSub.src0BlkStride = 1;
383+ repeatParamsSub.src1BlkStride = 1;
384+ repeatParamsSub.dstBlkStride = 1;
385+ repeatParamsSub.src0RepStride = columnCount / FP32_BLOCK_ELEMENT_NUM;
386+ repeatParamsSub.src1RepStride = 0;
387+ repeatParamsSub.dstRepStride = columnCount / FP32_BLOCK_ELEMENT_NUM;
388+ uint32_t columnRepeatCount = dLoop;
389+ uint32_t offset = 0;
390+ for (uint32_t i = 0; i < dLoop; i++) {
391+ Sub(dstUb[offset], src0Ub[offset], src1Ub[offset], dtypeMask, dealRowCount, repeatParamsSub);
392+ offset += dtypeMask;
393+ }
394+ 
395+ if (dRemain > 0) {
396+ Sub(dstUb[dLoop * dtypeMask], src0Ub[dLoop * dtypeMask], src1Ub[dLoop * dtypeMask], dRemain, dealRowCount, repeatParamsSub);
397+ }
398+}
399+ 
400+__aicore__ inline void ColMax(LocalTensor<float> dstUb, LocalTensor<float> src0Ub, LocalTensor<float> src1Ub,
401+ uint32_t dealRowCount, uint32_t columnCount, uint32_t actualColumnCount)
402+{
403+ uint32_t dtypeMask = FP32_REPEAT_ELEMENT_NUM;
404+ uint32_t dLoop = actualColumnCount / dtypeMask;
405+ uint32_t dRemain = actualColumnCount % dtypeMask;
406+ 
407+ BinaryRepeatParams repeatParamsMax;
408+ repeatParamsMax.src0BlkStride = 1;
409+ repeatParamsMax.src1BlkStride = 1;
410+ repeatParamsMax.dstBlkStride = 1;
411+ repeatParamsMax.src0RepStride = columnCount / FP32_BLOCK_ELEMENT_NUM;
412+ repeatParamsMax.src1RepStride = 0;
413+ repeatParamsMax.dstRepStride = 0;
414+ uint32_t columnRepeatCount = dLoop;
415+ uint32_t offset = 0;
416+ for (uint32_t i = 0; i < dLoop; i++) {
417+ Max(dstUb[offset], src0Ub[offset], src1Ub[offset], dtypeMask, dealRowCount, repeatParamsMax);
418+ offset += dtypeMask;
419+ }
420+ 
421+ if (dRemain > 0) {
422+ Max(dstUb[dLoop * dtypeMask], src0Ub[dLoop * dtypeMask], src1Ub[dLoop * dtypeMask], dRemain, dealRowCount, repeatParamsMax);
423+ }
424+}
425+ 
426+__aicore__ inline void ColAdd(LocalTensor<float> dstUb, LocalTensor<float> src0Ub, LocalTensor<float> src1Ub,
427+ uint32_t dealRowCount, uint32_t columnCount, uint32_t actualColumnCount)
428+{
429+ uint32_t dtypeMask = FP32_REPEAT_ELEMENT_NUM;
430+ uint32_t dLoop = actualColumnCount / dtypeMask;
431+ uint32_t dRemain = actualColumnCount % dtypeMask;
432+ 
433+ BinaryRepeatParams repeatParamsAdd;
434+ repeatParamsAdd.src0BlkStride = 1;
435+ repeatParamsAdd.src1BlkStride = 1;
436+ repeatParamsAdd.dstBlkStride = 1;
437+ repeatParamsAdd.src0RepStride = columnCount / FP32_BLOCK_ELEMENT_NUM;
438+ repeatParamsAdd.src1RepStride = 0;
439+ repeatParamsAdd.dstRepStride = 0;
440+ uint32_t columnRepeatCount = dLoop;
441+ uint32_t offset = 0;
442+ for (uint32_t i = 0; i < dLoop; i++) {
443+ Add(dstUb[offset], src0Ub[offset], src1Ub[offset], dtypeMask, dealRowCount, repeatParamsAdd);
444+ offset += dtypeMask;
445+ }
446+ 
447+ if (dRemain > 0) {
448+ Add(dstUb[dLoop * dtypeMask], src0Ub[dLoop * dtypeMask], src1Ub[dLoop * dtypeMask], dRemain, dealRowCount, repeatParamsAdd);
449+ }
450+}
451+ 
452+#endif // SPARSE_FLASH_ATTENTION_ANTIQUANT_COMMON_H
Aexperimental/attention/sparse_flash_attention_antiquant/op_kernel/sparse_flash_attention_antiquant_kernel_gqa.h+1050-0
@@ -0,0 +1,1050 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file sparse_flash_attention_antiquant_kernel_gqa.h
13+ * \brief
14+ */
15+ 
16+#ifndef SPARSE_FLASH_ATTENTION_ANTIQUANT_KERNEL_GQA_H
17+#define SPARSE_FLASH_ATTENTION_ANTIQUANT_KERNEL_GQA_H
18+ 
19+#include "kernel_operator.h"
20+#include "kernel_operator_list_tensor_intf.h"
21+#include "kernel_tiling/kernel_tiling.h"
22+#include "lib/matmul_intf.h"
23+#include "lib/matrix/matmul/tiling.h"
24+#include "sparse_flash_attention_antiquant_common.h"
25+#include "sparse_flash_attention_antiquant_service_cube_gqa_msd.h"
26+#include "sparse_flash_attention_antiquant_service_vector_gqa_msd.h"
27+#include "sparse_flash_attention_antiquant_service_flashdecode.h"
28+#include "sparse_flash_attention_antiquant_metadata.h"
29+ 
30+using namespace matmul;
31+using AscendC::CacheMode;
32+using AscendC::CrossCoreSetFlag;
33+using AscendC::CrossCoreWaitFlag;
34+using namespace optiling::detail;
35+ 
36+#define ARRAY_SIZE(x) (sizeof(x) / sizeof((x)[0]))
37+#define SALSK_S2BASEIZE 1024
38+#define SALSK_S2BASEIZE_1 1023
39+ 
40+// 由于S2循环前，RunInfo还没有赋值，使用Bngs1Param临时存放B、N、S1轴相关的信息；同时减少重复计算
41+struct TempLoopInfoGqa {
42+ uint32_t bN2Idx = 0U;
43+ uint32_t bIdx = 0U;
44+ uint32_t n2Idx = 0U;
45+ uint64_t s2BasicSizeTail = 0U; // S2方向循环的尾基本块大小
46+ uint32_t s2LoopTimes = 0U; // S2方向循环的总次数，无论TND还是BXXD都是等于实际次数，不用减1
47+ uint64_t curActualSeqLen = 0ULL;
48+ uint64_t curActualSeqLenOri = 0ULL;
49+ bool curActSeqLenIsZero = false;
50+ int32_t nextTokensPerBatch = 0;
51+ 
52+ uint64_t actS1Size = 1ULL; // TND场景下当前Batch循环处理的S1轴的大小，非TND场景下不要用这个字段
53+ uint64_t sparseBlockCount = 1Ull; // TopK的K大小
54+ uint32_t tndCoreStartKVSplitPos;
55+ bool tndIsS2SplitCore;
56+ int32_t threshold = 0;
57+ 
58+ uint32_t gS1Idx = 0U;
59+ uint64_t mBasicSizeTail = 0U; // gS1方向循环的尾基本块大小
60+ uint32_t sparseBlockActualSeqSizeKv = 0U;
61+};
62+ 
63+template <typename SFAAT> class SparseFlashAttentionAntiquantGqa {
64+public:
65+ // 中间计算数据类型为float，高精度模式
66+ using T = float;
67+ using Q_T = typename SFAAT::queryType;
68+ using KV_T = typename SFAAT::kvType;
69+ using OUT_T = typename SFAAT::outputType;
70+ using Q_ROPE_T = Q_T;
71+ using K_ROPE_T = typename AscendC::Conditional<SFAAT::isMsdDD, KV_T, half>::type;
72+ using UPDATE_T = T;
73+ using MM1_OUT_T = typename AscendC::Conditional<SFAAT::isMsdDD, half, T>::type;
74+ using MM2_OUT_T = typename AscendC::Conditional<SFAAT::isMsdDD, half, T>::type;
75+ 
76+ __aicore__ inline SparseFlashAttentionAntiquantGqa(){};
77+ __aicore__ inline void Init(__gm__ uint8_t *query, __gm__ uint8_t *key, __gm__ uint8_t *value,
78+ __gm__ uint8_t *sparseIndices, __gm__ uint8_t* keyScale,
79+ __gm__ uint8_t* valueScale, __gm__ uint8_t *blockTable,
80+ __gm__ uint8_t *actualSeqLengthsQ, __gm__ uint8_t *actualSeqLengths,
81+ __gm__ uint8_t *sparseSeqLengthsKv, SfaMetaData *metaData,
82+ __gm__ uint8_t *attentionOut, __gm__ uint8_t *workspace,
83+ const SparseFlashAttentionAntiquantTilingDataMla *__restrict tiling,
84+ __gm__ uint8_t *gmTiling, TPipe *tPipe);
85+ 
86+ __aicore__ inline void Process();
87+ 
88+private:
89+ static constexpr bool PAGE_ATTENTION = SFAAT::pageAttention;
90+ static constexpr int TEMPLATE_MODE = SFAAT::templateMode;
91+ static constexpr SFAA_LAYOUT LAYOUT_T = SFAAT::layout;
92+ static constexpr SFAA_LAYOUT KV_LAYOUT_T = SFAAT::kvLayout;
93+ bool FLASH_DECODE = false;
94+ static constexpr bool IS_META = SFAAT::flashDecode;
95+ 
96+ static constexpr int64_t fdPrefetchLen = 2;
97+ static constexpr uint32_t PRELOAD_NUM = 2;
98+ static constexpr uint32_t N_BUFFER_M_BASIC_SIZE = 256;
99+ static constexpr uint32_t SFAA_PRELOAD_TASK_CACHE_SIZE = 3;
100+ static constexpr uint32_t SFAA_KVPREMERGE_CACHE_SIZE = 2;
101+ 
102+ static constexpr uint32_t SYNC_V0_C1_FLAG = 6;
103+ static constexpr uint32_t SYNC_C1_V1_FLAG = 7;
104+ static constexpr uint32_t SYNC_V1_C2_FLAG = 8;
105+ static constexpr uint32_t SYNC_C2_V2_FLAG = 9;
106+ static constexpr uint32_t SYNC_C2_V1_FLAG = 4;
107+ static constexpr uint32_t SYNC_V1_NUPDATE_C2_FLAG = 5;
108+ 
109+ static constexpr uint64_t SYNC_MM2RES_BUF1_FLAG = 10;
110+ static constexpr uint64_t SYNC_MM2RES_BUF2_FLAG = 11;
111+ static constexpr uint64_t SYNC_FDOUTPUT_BUF_FLAG = 12;
112+ 
113+ static constexpr uint32_t BLOCK_ELEMENT_NUM = SFAAVectorServiceGqa<SFAAT>::BYTE_BLOCK / sizeof(T);
114+ 
115+ static constexpr uint32_t dbWorkspaceRatio = PRELOAD_NUM;
116+ 
117+ const SparseFlashAttentionAntiquantTilingDataMla *__restrict tilingData = nullptr;
118+ 
119+ TPipe *pipe = nullptr;
120+ 
121+ uint64_t mSizeVStart = 0ULL;
122+ int64_t threshold = 0;
123+ uint64_t s2BatchBaseOffset = 0;
124+ uint64_t tensorACoreOffset = 0ULL;
125+ uint64_t tensorARopeCoreOffset = 0ULL;
126+ uint64_t tensorBCoreOffset = 0ULL;
127+ int64_t topkGmBaseOffset = 0;
128+ 
129+ uint32_t tmpBlockIdx = 0U;
130+ uint32_t aiCoreIdx = 0U;
131+ uint32_t usedCoreNum = 0U;
132+ 
133+ __gm__ uint8_t *keyPtr = nullptr;
134+ __gm__ uint8_t *valuePtr = nullptr;
135+ SfaMetaData *metaDataPtr = nullptr;
136+ 
137+ ConstInfo constInfo{};
138+ TempLoopInfoGqa tempLoopInfo{};
139+ 
140+ SFAAMatmulServiceGqaMsd<SFAAT> matmulServiceMsd;
141+ SFAAVectorServiceGqaMsd<SFAAT> vectorServiceMsd;
142+ SFAAFlashDecodeServiceGqa<SFAAT> fdService;
143+ 
144+ GlobalTensor<Q_T> queryGm;
145+ GlobalTensor<KV_T> keyGm;
146+ GlobalTensor<KV_T> valueGm;
147+ GlobalTensor<T> keyDequantScaleGm;
148+ GlobalTensor<T> valueDequantScaleGm;
149+ 
150+ GlobalTensor<OUT_T> attentionOutGm;
151+ GlobalTensor<int32_t> blockTableGm;
152+ GlobalTensor<int32_t> topKGm;
153+ 
154+ GlobalTensor<int32_t> actualSeqLengthsQGm;
155+ GlobalTensor<int32_t> actualSeqLengthsKVGm;
156+ GlobalTensor<int32_t> sparseSeqLengthsKVGm;
157+ 
158+ // workspace
159+ GlobalTensor<KV_T> queryPreProcessResGm;
160+ GlobalTensor<MM1_OUT_T> mm1ResGm;
161+ GlobalTensor<K_ROPE_T> vec1ResGm;
162+ GlobalTensor<MM2_OUT_T> mm2ResGm;
163+ GlobalTensor<K_ROPE_T> keyMergeGm_;
164+ GlobalTensor<K_ROPE_T> valueMergeGm_;
165+ 
166+ GlobalTensor<int32_t> mm2ResInt32Gm;
167+ GlobalTensor<T> vec2ResGm;
168+ 
169+ GlobalTensor<T> accumOutGm;
170+ GlobalTensor<T> lseSumFdGm;
171+ GlobalTensor<T> lseMaxFdGm;
172+ 
173+ template <typename T>
174+ __aicore__ inline T Align(T num, T rnd)
175+ {
176+ return (((rnd) == 0) ? 0 : (((num) + (rnd)-1) / (rnd) * (rnd)));
177+ }
178+ 
179+ // ================================Init functions==================================
180+ __aicore__ inline void InitTilingData();
181+ __aicore__ inline void InitMetaData();
182+ __aicore__ inline void InitCalcParamsEach();
183+ __aicore__ inline void InitBuffers();
184+ __aicore__ inline void InitActualSeqLen(__gm__ uint8_t *actualSeqLengthsQ, __gm__ uint8_t *actualSeqLengths,
185+ __gm__ uint8_t *sparseSeqLengthsKv);
186+ __aicore__ inline void InitOutputSingleCore();
187+ // ================================Process functions================================
188+ __aicore__ inline void ProcessBalance();
189+ __aicore__ inline void PreloadPipeline(uint32_t loop, uint64_t s2Start, uint64_t s2LoopIdx,
190+ RunInfo extraInfo[SFAA_PRELOAD_TASK_CACHE_SIZE]);
191+ __aicore__ inline void FlashDecode();
192+ // ================================Offset Calc=====================================
193+ __aicore__ inline void GetActualSeqLen(uint32_t bIdx, uint32_t s1Idx = 0);
194+ __aicore__ inline void GetSparseActualSeqLen(uint32_t bIdx, uint32_t s1Idx, uint32_t n2Idx);
195+ __aicore__ inline void UpdateInnerLoopCond();
196+ __aicore__ inline void DealActSeqLenIsZero(uint32_t bIdx, uint32_t s1StartIdx, uint32_t s1Size, uint32_t n2Idx);
197+ __aicore__ inline void CalcParams(uint32_t loop, uint64_t s2Start, uint32_t s2LoopIdx, RunInfo &info);
198+ __aicore__ inline void GetAxisStartIdx(uint32_t bN2EndPrev, uint32_t gS1EndPrev, uint32_t s2EndPrev);
199+ __aicore__ inline uint64_t GetBalanceActualSeqLengths(GlobalTensor<int32_t> &actualSeqLengths, uint32_t bIdx);
200+ __aicore__ inline uint32_t GetActualSeqLenKV(uint32_t bIdx);
201+ __aicore__ inline void GetSparseBlockCountAndActualSeqKv(uint32_t bIdx);
202+ __aicore__ inline void GetBN2Idx(uint32_t bN2Idx, uint32_t &bIdx, uint32_t &n2Idx);
203+ __aicore__ inline void UpdateInner(uint32_t &s2End, uint32_t &curS2End, uint32_t s1Idx, bool isEnd);
204+ __aicore__ inline void GetPreNextTokensLeftUp();
205+ // ================================Mm1==============================================
206+ __aicore__ inline void ComputeMm1(const RunInfo &info);
207+ // ================================Mm2==============================================
208+ __aicore__ inline void ComputeMm2(const RunInfo &info);
209+ __aicore__ inline void Bmm2DataCopyOut(uint64_t attenOutOffset, LocalTensor<OUT_T> &attenOutUb, uint32_t startRow,
210+ uint32_t dealRowCount, uint32_t columnCount, uint32_t actualColumnCount);
211+ __aicore__ inline void InitAllZeroOutput(uint32_t bIdx, uint32_t s1Idx, uint32_t n2Idx);
212+};
213+ 
214+template <typename SFAAT> __aicore__ inline void SparseFlashAttentionAntiquantGqa<SFAAT>::InitTilingData()
215+{
216+ usedCoreNum = tilingData->singleCoreParams.usedCoreNum;
217+ constInfo.splitKVNum = tilingData->splitKVParams.s2;
218+ constInfo.mmResUbSize = tilingData->singleCoreTensorSize.mmResUbSize;
219+ constInfo.bmm2ResUbSize = tilingData->singleCoreTensorSize.bmm2ResUbSize;
220+ constInfo.vec1ResUbSize = constInfo.mmResUbSize;
221+ 
222+ constInfo.batchSize = tilingData->baseParams.batchSize;
223+ constInfo.gSize = tilingData->baseParams.nNumOfQInOneGroup;
224+ constInfo.kvHeadNum = tilingData->baseParams.kvHeadNum;
225+ constInfo.qHeadNum = constInfo.gSize * tilingData->baseParams.kvHeadNum;
226+ constInfo.kvSeqSize = tilingData->baseParams.seqSize;
227+ constInfo.qSeqSize = tilingData->baseParams.qSeqSize;
228+ constInfo.maxBlockNumPerBatch = tilingData->baseParams.maxBlockNumPerBatch;
229+ constInfo.kvCacheBlockSize = tilingData->baseParams.blockSize;
230+ constInfo.outputLayout = static_cast<SFAA_LAYOUT>(tilingData->baseParams.outputLayout);
231+ 
232+ constInfo.sparseBlockSize = tilingData->baseParams.sparseBlockSize;
233+ constInfo.sparseBlockCount = tilingData->baseParams.sparseBlockCount;
234+ constInfo.sparseShardSize = tilingData->baseParams.sparseShardSize;
235+ constInfo.sparseMode = tilingData->baseParams.sparseMode;
236+ constInfo.attentionMode = static_cast<ATTENTION_MODE>(tilingData->baseParams.attentionMode);
237+ constInfo.keyQuantMode = static_cast<QUANT_MODE>(tilingData->baseParams.keyQuantMode);
238+ constInfo.valueQuantMode = static_cast<QUANT_MODE>(tilingData->baseParams.valueQuantMode);
239+ constInfo.quantScaleRepoMode = static_cast<QUANT_SCALE_REPO_MODE>(tilingData->baseParams.quantScaleRepoMode);
240+ constInfo.combineHeadDim = tilingData->baseParams.qkHeadDim;
241+ constInfo.headDimRope = tilingData->baseParams.ropeHeadDim;
242+ constInfo.headDim = (constInfo.quantScaleRepoMode == QUANT_SCALE_REPO_MODE::COMBINE) ?
243+ constInfo.combineHeadDim - constInfo.headDimRope : constInfo.combineHeadDim;
244+ constInfo.headDimAlign = Align(constInfo.headDim, (uint64_t)BYTE_BLOCK);
245+ 
246+ constInfo.s2BaseSize = SALSK_S2BASEIZE;
247+ constInfo.mBaseSize = constInfo.gSize * constInfo.sparseShardSize;
248+ 
249+ constInfo.preLoadNum = PRELOAD_NUM;
250+ constInfo.nBufferMBaseSize = N_BUFFER_M_BASIC_SIZE; // 256
251+ constInfo.syncV0C1 = SYNC_V0_C1_FLAG;
252+ constInfo.syncC1V1 = SYNC_C1_V1_FLAG;
253+ constInfo.syncV1C2 = SYNC_V1_C2_FLAG;
254+ constInfo.syncC2V2 = SYNC_C2_V2_FLAG;
255+ // constInfo.syncC2V1 = SYNC_C2_V1_FLAG;
256+ constInfo.syncV1NupdateC2 = SYNC_V1_NUPDATE_C2_FLAG;
257+}
258+ 
259+template <typename SFAAT> __aicore__ inline void SparseFlashAttentionAntiquantGqa<SFAAT>::InitMetaData()
260+{
261+ FLASH_DECODE = metaDataPtr->numOfFdHead > 0U;
262+ usedCoreNum = metaDataPtr->usedCoreNum;
263+ constInfo.s2BaseSize = SALSK_S2BASEIZE;
264+ constInfo.mBaseSize = metaDataPtr->mBaseSize;
265+}
266+ 
267+template <typename SFAAT> __aicore__ inline void SparseFlashAttentionAntiquantGqa<SFAAT>::InitBuffers()
268+{
269+ if constexpr (SFAAT::isMsdDD) {
270+ if ASCEND_IS_AIV {
271+ vectorServiceMsd.InitBuffers(pipe);
272+ } else {
273+ matmulServiceMsd.InitBuffers(pipe);
274+ }
275+ }
276+}
277+ 
278+template <typename SFAAT>
279+__aicore__ inline void
280+SparseFlashAttentionAntiquantGqa<SFAAT>::InitActualSeqLen(__gm__ uint8_t *actualSeqLengthsQ,
281+ __gm__ uint8_t *actualSeqLengths,
282+ __gm__ uint8_t *sparseSeqLengthsKv)
283+{
284+ constInfo.actualLenDimsQ = tilingData->baseParams.actualLenDimsQ;
285+ constInfo.actualLenDimsKV = tilingData->baseParams.actualLenDimsKV;
286+ constInfo.sparseLenDimsKV = tilingData->baseParams.sparseLenDimsKV;
287+ if (constInfo.actualLenDimsKV != 0) {
288+ actualSeqLengthsKVGm.SetGlobalBuffer((__gm__ int32_t *)actualSeqLengths, constInfo.actualLenDimsKV);
289+ }
290+ if (constInfo.actualLenDimsQ != 0) {
291+ actualSeqLengthsQGm.SetGlobalBuffer((__gm__ int32_t *)actualSeqLengthsQ, constInfo.actualLenDimsQ);
292+ }
293+ if (constInfo.sparseLenDimsKV != 0) {
294+ sparseSeqLengthsKVGm.SetGlobalBuffer((__gm__ int32_t *)sparseSeqLengthsKv, constInfo.sparseLenDimsKV);
295+ }
296+}
297+ 
298+template <typename SFAAT>
299+__aicore__ inline void SparseFlashAttentionAntiquantGqa<SFAAT>::InitAllZeroOutput(uint32_t bIdx, uint32_t s1Idx,
300+ uint32_t n2Idx)
301+{
302+ if (constInfo.outputLayout == SFAA_LAYOUT::TND) {
303+ uint32_t tBase = bIdx == 0 ? 0 : actualSeqLengthsQGm.GetValue(bIdx - 1);
304+ uint32_t s1Count = tempLoopInfo.actS1Size;
305+ 
306+ uint64_t attenOutOffset = (tBase + s1Idx) * constInfo.kvHeadNum * constInfo.gSize * constInfo.headDim + // T轴、s1轴偏移
307+ n2Idx * constInfo.gSize * constInfo.headDim; // N2轴偏移
308+ matmul::InitOutput<OUT_T>(attentionOutGm[attenOutOffset], constInfo.gSize * constInfo.headDim, 0);
309+ } else if (constInfo.outputLayout == SFAA_LAYOUT::BSND) {
310+ uint64_t attenOutOffset = bIdx * constInfo.qSeqSize * constInfo.kvHeadNum * constInfo.gSize * constInfo.headDim +
311+ s1Idx * constInfo.kvHeadNum * constInfo.gSize * constInfo.headDim + // B轴、S1轴偏移
312+ n2Idx * constInfo.gSize * constInfo.headDim; // N2轴偏移
313+ matmul::InitOutput<OUT_T>(attentionOutGm[attenOutOffset], constInfo.gSize * constInfo.headDim, 0);
314+ }
315+}
316+ 
317+template <typename SFAAT>
318+__aicore__ inline void SparseFlashAttentionAntiquantGqa<SFAAT>::InitOutputSingleCore()
319+{
320+ uint32_t coreNum = GetBlockNum();
321+ if (coreNum != 0) {
322+ uint64_t totalOutputSize = constInfo.batchSize * constInfo.qHeadNum * constInfo.qSeqSize * constInfo.headDim;
323+ uint64_t singleCoreSize = (totalOutputSize + (2 * coreNum) - 1) / (2 * coreNum); // 2 means c:v = 1:2
324+ uint64_t tailSize = totalOutputSize - tmpBlockIdx * singleCoreSize;
325+ uint64_t singleInitOutputSize = tailSize < singleCoreSize ? tailSize : singleCoreSize;
326+ if (singleInitOutputSize > 0) {
327+ matmul::InitOutput<OUT_T>(attentionOutGm[tmpBlockIdx * singleCoreSize], singleInitOutputSize, 0);
328+ }
329+ SyncAll();
330+ }
331+}
332+ 
333+template <typename SFAAT>
334+__aicore__ inline void SparseFlashAttentionAntiquantGqa<SFAAT>::GetActualSeqLen(uint32_t bIdx, uint32_t s1Idx)
335+{
336+ tempLoopInfo.curActualSeqLenOri = GetActualSeqLenKV(bIdx);
337+ tempLoopInfo.actS1Size = GetBalanceActualSeqLengths(actualSeqLengthsQGm, bIdx);
338+ GetSparseBlockCountAndActualSeqKv(bIdx);
339+}
340+ 
341+template <typename SFAAT>
342+__aicore__ inline void SparseFlashAttentionAntiquantGqa<SFAAT>::GetSparseActualSeqLen(uint32_t bIdx, uint32_t s1Idx,
343+ uint32_t n2Idx)
344+{
345+ if (tempLoopInfo.nextTokensPerBatch < 0 && s1Idx < (-tempLoopInfo.nextTokensPerBatch)) { // 存在行无效
346+ tempLoopInfo.curActualSeqLen = 0;
347+ return;
348+ }
349+ tempLoopInfo.threshold = tempLoopInfo.curActualSeqLenOri;
350+ if (constInfo.sparseMode == 3) {
351+ tempLoopInfo.threshold = static_cast<int64_t>(tempLoopInfo.nextTokensPerBatch) + s1Idx + 1;
352+ }
353+ tempLoopInfo.curActualSeqLen = (tempLoopInfo.sparseBlockActualSeqSizeKv > tempLoopInfo.threshold) ?
354+ tempLoopInfo.threshold : tempLoopInfo.sparseBlockActualSeqSizeKv;
355+}
356+ 
357+template <typename SFAAT>
358+__aicore__ inline void SparseFlashAttentionAntiquantGqa<SFAAT>::GetSparseBlockCountAndActualSeqKv(uint32_t bIdx)
359+{
360+ if (constInfo.sparseLenDimsKV == 0) {
361+ tempLoopInfo.sparseBlockCount = constInfo.sparseBlockCount;
362+ tempLoopInfo.sparseBlockActualSeqSizeKv = constInfo.sparseBlockCount * constInfo.sparseBlockSize;
363+ } else {
364+ tempLoopInfo.sparseBlockActualSeqSizeKv = sparseSeqLengthsKVGm.GetValue(bIdx);
365+ tempLoopInfo.sparseBlockCount = (tempLoopInfo.sparseBlockActualSeqSizeKv + constInfo.sparseBlockSize - 1) / constInfo.sparseBlockSize;
366+ }
367+}
368+ 
369+template <typename SFAAT>
370+__aicore__ inline uint32_t SparseFlashAttentionAntiquantGqa<SFAAT>::GetActualSeqLenKV(uint32_t bIdx)
371+{
372+ if (constInfo.actualLenDimsKV == 0) {
373+ return constInfo.kvSeqSize;
374+ } else if (constInfo.actualLenDimsKV == 1) {
375+ return actualSeqLengthsKVGm.GetValue(0);
376+ } else {
377+ return actualSeqLengthsKVGm.GetValue(bIdx);
378+ }
379+}
380+ 
381+template <typename SFAAT>
382+__aicore__ inline void SparseFlashAttentionAntiquantGqa<SFAAT>::DealActSeqLenIsZero(uint32_t bIdx, uint32_t s1StartIdx,
383+ uint32_t s1Size, uint32_t n2Idx)
384+{
385+ if ASCEND_IS_AIV {
386+ for (uint32_t i = 0; i < s1Size; i++) {
387+ InitAllZeroOutput(bIdx, s1StartIdx + i, n2Idx);
388+ }
389+ }
390+}
391+ 
392+template <typename SFAAT>
393+__aicore__ inline void SparseFlashAttentionAntiquantGqa<SFAAT>::GetPreNextTokensLeftUp()
394+{
395+ if (constInfo.sparseMode == 3) {
396+ tempLoopInfo.nextTokensPerBatch =
397+ static_cast<int32_t>(tempLoopInfo.curActualSeqLenOri) - static_cast<int32_t>(tempLoopInfo.actS1Size);
398+ }
399+}
400+ 
401+template <typename SFAAT> __aicore__ inline void SparseFlashAttentionAntiquantGqa<SFAAT>::UpdateInnerLoopCond()
402+{
403+ if ((tempLoopInfo.curActualSeqLen == 0) || (tempLoopInfo.actS1Size == 0)) {
404+ tempLoopInfo.curActSeqLenIsZero = true;
405+ return;
406+ }
407+ tempLoopInfo.curActSeqLenIsZero = false;
408+ tempLoopInfo.s2BasicSizeTail = tempLoopInfo.curActualSeqLen & (SALSK_S2BASEIZE_1);
409+ tempLoopInfo.s2BasicSizeTail =
410+ (tempLoopInfo.s2BasicSizeTail == 0) ? SALSK_S2BASEIZE : tempLoopInfo.s2BasicSizeTail;
411+ tempLoopInfo.mBasicSizeTail = (tempLoopInfo.actS1Size * constInfo.gSize) % constInfo.mBaseSize;
412+ tempLoopInfo.mBasicSizeTail =
413+ (tempLoopInfo.mBasicSizeTail == 0) ? constInfo.mBaseSize : tempLoopInfo.mBasicSizeTail;
414+ tempLoopInfo.s2LoopTimes = 0;
415+}
416+ 
417+template <typename SFAAT>
418+__aicore__ inline void SparseFlashAttentionAntiquantGqa<SFAAT>::UpdateInner(uint32_t &s2End, uint32_t &curS2End,
419+ uint32_t s1Idx, bool isEnd)
420+{
421+ uint32_t s1BaseSize = 1;
422+ int64_t s1Offset = s1BaseSize * s1Idx;
423+ int64_t s2LastToken = Min(s1Offset + tempLoopInfo.nextTokensPerBatch + s1BaseSize, tempLoopInfo.curActualSeqLenOri);
424+ s2LastToken = Min(constInfo.sparseBlockSize * tempLoopInfo.sparseBlockCount, s2LastToken);
425+ curS2End = (s2LastToken + SALSK_S2BASEIZE - 1) / SALSK_S2BASEIZE;
426+ tempLoopInfo.s2LoopTimes = isEnd ? constInfo.s2End + 1 : curS2End;
427+}
428+ 
429+template <typename SFAAT>
430+__aicore__ inline void SparseFlashAttentionAntiquantGqa<SFAAT>::Init(__gm__ uint8_t *query,
431+ __gm__ uint8_t *key, __gm__ uint8_t *value, __gm__ uint8_t *sparseIndices, __gm__ uint8_t* keyScale,
432+ __gm__ uint8_t* valueScale, __gm__ uint8_t *blockTable, __gm__ uint8_t *actualSeqLengthsQ,
433+ __gm__ uint8_t *actualSeqLengths, __gm__ uint8_t *sparseSeqLengthsKv, SfaMetaData *metaData,
434+ __gm__ uint8_t *attentionOut, __gm__ uint8_t *workspace,
435+ const SparseFlashAttentionAntiquantTilingDataMla *__restrict tiling,
436+ __gm__ uint8_t *gmTiling, TPipe *tPipe)
437+{
438+ if ASCEND_IS_AIV {
439+ tmpBlockIdx = GetBlockIdx(); // vec:0-47
440+ aiCoreIdx = tmpBlockIdx / 2;
441+ } else {
442+ tmpBlockIdx = GetBlockIdx(); // cube:0-23
443+ aiCoreIdx = tmpBlockIdx;
444+ }
445+ 
446+ // init tiling data
447+ tilingData = tiling;
448+ InitTilingData();
449+ InitActualSeqLen(actualSeqLengthsQ, actualSeqLengths, sparseSeqLengthsKv);
450+ 
451+ // 初始化计算参数
452+ if constexpr (IS_META) {
453+ if (metaData != nullptr) {
454+ metaDataPtr = metaData;
455+ InitMetaData();
456+ }
457+ }
458+ InitCalcParamsEach();
459+ pipe = tPipe;
460+ keyPtr = key;
461+ valuePtr = value;
462+ 
463+ // init global buffer
464+ queryGm.SetGlobalBuffer((__gm__ Q_T *)query);
465+ keyGm.SetGlobalBuffer((__gm__ KV_T *)keyPtr);
466+ valueGm.SetGlobalBuffer((__gm__ KV_T *)valuePtr);
467+ keyDequantScaleGm.SetGlobalBuffer((__gm__ T *)keyScale);
468+ valueDequantScaleGm.SetGlobalBuffer((__gm__ T *)valueScale);
469+ keyGm.SetL2CacheHint(AscendC::CacheMode::CACHE_MODE_DISABLE);
470+ valueGm.SetL2CacheHint(AscendC::CacheMode::CACHE_MODE_DISABLE);
471+ 
472+ attentionOutGm.SetGlobalBuffer((__gm__ OUT_T *)attentionOut);
473+ 
474+ if ASCEND_IS_AIV {
475+ if (constInfo.needInit && LAYOUT_T != SFAA_LAYOUT::TND) {
476+ InitOutputSingleCore();
477+ }
478+ }
479+ 
480+ if constexpr (PAGE_ATTENTION) {
481+ blockTableGm.SetGlobalBuffer((__gm__ int32_t *)blockTable);
482+ }
483+ topKGm.SetGlobalBuffer((__gm__ int32_t *)sparseIndices);
484+ 
485+ // workspace 内存排布
486+ // |Q--|mm1ResGm(存S)|vec1ResGm(存A1,A2)|mm2ResGm(存O)|vec2ResGm
487+ // |Core0_Q1-Core0_Q2-Core1_Q1-Core1_Q2....Core32_Q1-Core32_Q2|Core0_mmRes
488+ uint64_t offset = 0;
489+ queryPreProcessResGm.SetGlobalBuffer(
490+ (__gm__ KV_T *)(workspace + offset +
491+ aiCoreIdx * dbWorkspaceRatio * constInfo.bmm2ResUbSize * 2 * sizeof(KV_T)));
492+ offset += GetBlockNum() * dbWorkspaceRatio * constInfo.bmm2ResUbSize * 2 * sizeof(KV_T);
493+ 
494+ mm1ResGm.SetGlobalBuffer(
495+ (__gm__ MM1_OUT_T *)(workspace + offset +
496+ aiCoreIdx * dbWorkspaceRatio * constInfo.mmResUbSize * sizeof(MM1_OUT_T)));
497+ offset += GetBlockNum() * dbWorkspaceRatio * constInfo.mmResUbSize * sizeof(MM1_OUT_T);
498+ 
499+ vec1ResGm.SetGlobalBuffer(
500+ (__gm__ K_ROPE_T *)(workspace + offset + aiCoreIdx * dbWorkspaceRatio * constInfo.mmResUbSize * 2 *
501+ sizeof(K_ROPE_T)));
502+ offset += GetBlockNum() * dbWorkspaceRatio * constInfo.mmResUbSize * sizeof(K_ROPE_T) * 2;
503+ 
504+ mm2ResGm.SetGlobalBuffer(
505+ (__gm__ MM2_OUT_T *)(workspace + offset +
506+ aiCoreIdx * dbWorkspaceRatio * constInfo.bmm2ResUbSize * sizeof(MM2_OUT_T)));
507+ offset += GetBlockNum() * dbWorkspaceRatio * constInfo.bmm2ResUbSize * sizeof(MM2_OUT_T);
508+ mm2ResInt32Gm.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t *>(mm2ResGm.GetPhyAddr(0)));
509+ 
510+ vec2ResGm.SetGlobalBuffer((__gm__ T *)(workspace + offset +
511+ aiCoreIdx * dbWorkspaceRatio * constInfo.bmm2ResUbSize * sizeof(T)));
512+ offset += GetBlockNum() * dbWorkspaceRatio * constInfo.bmm2ResUbSize * sizeof(MM2_OUT_T);
513+ 
514+ // s2 d bufNum
515+ // 384 is mergesize
516+ // 128 is headdim
517+ // 4 is buffernum
518+ keyMergeGm_.SetGlobalBuffer((__gm__ K_ROPE_T *)(workspace + offset + aiCoreIdx * 384 * 128 * 4 *
M
Mmiaofangzheng4月2日

384 * 128 * 4 这些魔数在 workspace 分配中重复出现，但没有对应的常量定义或注释说明 384 和 128 代表什么。如果 tiling 侧修改了 workspace 计算但 kernel 侧忘了同步修改，就会导致 buffer 重叠或越界。建议通过 tiling data 传递这些参数，而不是在 kernel 侧硬编码。

likedislike
M
mayafei
4月3日 评论：
519+ sizeof(K_ROPE_T)));
520+ offset += GetBlockNum() * 384 * 128 * 4 * sizeof(K_ROPE_T);
521+ valueMergeGm_.SetGlobalBuffer((__gm__ K_ROPE_T *)(workspace + offset + aiCoreIdx * 384 * 128 * 4 *
522+ sizeof(K_ROPE_T)));
523+ offset += GetBlockNum() * 384 * 128 * 4 * sizeof(K_ROPE_T);
524+ 
525+ if constexpr (IS_META) {
526+ if (FLASH_DECODE) {
527+ accumOutGm.SetGlobalBuffer((__gm__ float *)(workspace + offset));
528+ offset = offset + tilingData->splitKVParams.accumOutSize * sizeof(float);
529+ lseSumFdGm.SetGlobalBuffer((__gm__ float *)(workspace + offset));
530+ lseMaxFdGm.SetGlobalBuffer((__gm__ float *)(workspace + offset) + tilingData->splitKVParams.logSumExpSize / 2);
531+ offset = offset + tilingData->splitKVParams.logSumExpSize * sizeof(float);
532+ }
533+ }
534+ 
535+ if ASCEND_IS_AIV {
536+ if constexpr (SFAAT::isMsdDD) {
537+ vectorServiceMsd.InitParams(constInfo, tilingData, metaDataPtr);
538+ vectorServiceMsd.InitVec0GlobalTensor(keyMergeGm_, valueMergeGm_, queryPreProcessResGm, queryGm,
539+ keyGm, valueGm, blockTableGm, keyDequantScaleGm);
540+ vectorServiceMsd.InitVec1GlobalTensor(mm1ResGm, vec1ResGm, actualSeqLengthsQGm,
541+ actualSeqLengthsKVGm, lseMaxFdGm, lseSumFdGm, topKGm);
542+ vectorServiceMsd.InitVec2GlobalTensor(valueDequantScaleGm, accumOutGm, mm2ResGm, attentionOutGm);
543+ }
544+ if constexpr (IS_META) {
545+ if (FLASH_DECODE) {
546+ fdService.InitParams(constInfo);
547+ fdService.InitGlobalTensor(lseMaxFdGm, lseSumFdGm, accumOutGm, attentionOutGm,
548+ actualSeqLengthsQGm, sparseSeqLengthsKVGm);
549+ }
550+ }
551+ }
552+ 
553+ if ASCEND_IS_AIC {
554+ if constexpr (SFAAT::isMsdDD) {
555+ matmulServiceMsd.InitParams(constInfo);
556+ matmulServiceMsd.InitMm1GlobalTensor(queryGm, keyGm, mm1ResGm);
557+ matmulServiceMsd.InitMm2GlobalTensor(vec1ResGm, valueGm, mm2ResGm, attentionOutGm);
558+ matmulServiceMsd.InitPageAttentionInfo(keyMergeGm_, valueMergeGm_, queryPreProcessResGm, blockTableGm, topKGm,
559+ constInfo.kvCacheBlockSize, constInfo.maxBlockNumPerBatch);
560+ }
561+ }
562+ // 要在InitParams之后执行
563+ if (pipe != nullptr) {
564+ InitBuffers();
565+ }
566+}
567+ 
568+template <typename SFAAT> __aicore__ inline void SparseFlashAttentionAntiquantGqa<SFAAT>::InitCalcParamsEach()
569+{
570+ if constexpr (IS_META) {
571+ const uint32_t *bN2End = metaDataPtr->bN2End;
572+ const uint32_t *gS1End = metaDataPtr->gS1End;
573+ const uint32_t *s2End = metaDataPtr->s2End;
574+ const uint32_t *s2SplitStartIdxOfCore = metaDataPtr->fdRes.s2SplitStartIdxOfCore;
575+ constInfo.bN2End = bN2End[aiCoreIdx];
576+ constInfo.gS1End = gS1End[aiCoreIdx];
577+ constInfo.s2End = s2End[aiCoreIdx];
578+ if (aiCoreIdx != 0) {
579+ GetAxisStartIdx(bN2End[aiCoreIdx - 1], gS1End[aiCoreIdx - 1], s2End[aiCoreIdx - 1]);
580+ }
581+ constInfo.coreStartKVSplitPos = s2SplitStartIdxOfCore[aiCoreIdx];
582+ } else {
583+#ifdef ASCENDC_CPU_DEBUG
584+ const uint32_t *bN2End = tilingData->outerSplitParams.bN2End;
585+ const uint32_t *gS1End = tilingData->outerSplitParams.gS1End;
586+ const uint32_t *s2End = tilingData->outerSplitParams.s2End;
587+#else
588+ uint32_t bN2End[ARRAY_SIZE(tilingData->outerSplitParams.bN2End)];
589+ uint32_t gS1End[ARRAY_SIZE(tilingData->outerSplitParams.gS1End)];
590+ uint32_t s2End[ARRAY_SIZE(tilingData->outerSplitParams.s2End)];
591+ copy_data_align64((uint8_t *)bN2End, (uint8_t *)(tilingData->outerSplitParams.bN2End), sizeof(bN2End));
592+ copy_data_align64((uint8_t *)gS1End, (uint8_t *)(tilingData->outerSplitParams.gS1End), sizeof(gS1End));
593+ copy_data_align64((uint8_t *)s2End, (uint8_t *)(tilingData->outerSplitParams.s2End), sizeof(s2End));
594+#endif
595+ // TND分核信息
596+ constInfo.bN2End = bN2End[aiCoreIdx];
597+ constInfo.gS1End = gS1End[aiCoreIdx];
598+ constInfo.s2End = s2End[aiCoreIdx];
599+ if (aiCoreIdx != 0) {
600+ GetAxisStartIdx(bN2End[aiCoreIdx - 1], gS1End[aiCoreIdx - 1], s2End[aiCoreIdx - 1]);
601+ }
602+ }
603+}
604+ 
605+template <typename SFAAT>
606+__aicore__ inline void
607+SparseFlashAttentionAntiquantGqa<SFAAT>::Bmm2DataCopyOut(uint64_t attenOutOffset, LocalTensor<OUT_T> &attenOutUb,
608+ uint32_t startRow, uint32_t dealRowCount,
609+ uint32_t columnCount, uint32_t actualColumnCount)
610+{
611+ DataCopyExtParams dataCopyParams;
612+ dataCopyParams.blockCount = dealRowCount;
613+ dataCopyParams.blockLen = actualColumnCount * sizeof(OUT_T);
614+ dataCopyParams.srcStride = (columnCount - actualColumnCount) / (SFAAVectorServiceGqa<SFAAT>::BYTE_BLOCK /
615+ sizeof(OUT_T));
616+ dataCopyParams.dstStride = 0;
617+ DataCopyPad(attentionOutGm[attenOutOffset + (mSizeVStart + startRow) * actualColumnCount], attenOutUb,
618+ dataCopyParams);
619+}
620+ 
621+ 
622+template <typename SFAAT>
623+__aicore__ inline void SparseFlashAttentionAntiquantGqa<SFAAT>::CalcParams(uint32_t loop, uint64_t s2Start,
624+ uint32_t s2LoopIdx, RunInfo &info)
625+{
626+ info.loop = loop;
627+ info.bN2Idx = tempLoopInfo.bN2Idx;
628+ info.bIdx = tempLoopInfo.bIdx;
629+ info.n2Idx = tempLoopInfo.n2Idx;
630+ info.gS1Idx = tempLoopInfo.gS1Idx;
631+ info.s2Idx = s2LoopIdx;
632+ info.curSInnerLoopTimes = tempLoopInfo.s2LoopTimes;
633+ info.threshold = tempLoopInfo.threshold;
634+ info.sparseBlockCount = tempLoopInfo.sparseBlockCount;
635+ 
636+ info.tndIsS2SplitCore = tempLoopInfo.tndIsS2SplitCore;
637+ info.tndCoreStartKVSplitPos = tempLoopInfo.tndCoreStartKVSplitPos;
638+ info.isBmm2Output = false;
639+ 
640+ info.actS1Size = tempLoopInfo.actS1Size;
641+ info.actS2Size = tempLoopInfo.curActualSeqLen;
642+ info.nextTokensPerBatch =
643+ static_cast<int32_t>(info.actS2Size) - static_cast<int32_t>(info.actS1Size);
644+
645+ info.actMBaseSize = constInfo.mBaseSize;
646+ uint32_t remainedGS1Size = tempLoopInfo.actS1Size * constInfo.gSize - tempLoopInfo.gS1Idx;
647+ if (remainedGS1Size <= constInfo.mBaseSize && remainedGS1Size > 0) {
648+ info.actMBaseSize = tempLoopInfo.mBasicSizeTail;
649+ }
650+ 
651+ info.isValid = s2LoopIdx < tempLoopInfo.s2LoopTimes;
652+ info.mSize = info.actMBaseSize;
653+ if ASCEND_IS_AIV {
654+ // info.mSizeV = (info.mSize <= 16) ? info.mSize : (((info.mSize + 15) / 16 + 1) / 2 * 16);
655+ info.mSizeV = (info.mSize + 1) / 2;
656+ info.mSizeVStart = 0;
657+ if (tmpBlockIdx % 2 == 1) {
658+ info.mSizeVStart = info.mSizeV;
659+ info.mSizeV = info.mSize - info.mSizeV;
660+ }
661+ }
662+ 
663+ info.isChangeBatch = false;
664+ 
665+ info.isFirstSInnerLoop = s2LoopIdx == s2Start;
666+ info.isLastS2Loop = s2LoopIdx == tempLoopInfo.s2LoopTimes - 1;
667+ uint64_t actualSeqQPrefixSum;
668+ if constexpr (LAYOUT_T == SFAA_LAYOUT::TND) {
669+ actualSeqQPrefixSum = (info.bIdx <= 0) ? 0 : actualSeqLengthsQGm.GetValue(info.bIdx - 1);
670+ } else {
671+ actualSeqQPrefixSum = (info.bIdx <= 0) ? 0 : info.bIdx * constInfo.qSeqSize;
672+ }
673+ info.tndBIdxOffsetQ = actualSeqQPrefixSum * constInfo.qHeadNum * constInfo.combineHeadDim;
674+ uint64_t actualSeqKvSum;
675+ if constexpr (KV_LAYOUT_T == SFAA_LAYOUT::TND) {
676+ actualSeqKvSum = (info.bIdx <= 0) ? 0 : actualSeqLengthsKVGm.GetValue(info.bIdx - 1);
677+ } else {
678+ actualSeqKvSum = (info.bIdx <= 0) ? 0 : info.bIdx * constInfo.kvSeqSize;
679+ }
680+ info.tndBIdxOffsetKV = actualSeqKvSum * constInfo.kvHeadNum * constInfo.combineHeadDim;
681+ if (info.isFirstSInnerLoop) {
682+ // 支持BSND/TND
683+ tensorACoreOffset = info.tndBIdxOffsetQ + info.gS1Idx * constInfo.kvHeadNum * constInfo.combineHeadDim
684+ + info.n2Idx * constInfo.gSize * constInfo.combineHeadDim; // info.gS1Idx：前提是不切G，否则会有向下取整的问题
685+ tensorBCoreOffset = info.tndBIdxOffsetKV + info.n2Idx * constInfo.gSize * constInfo.combineHeadDim;
686+ topkGmBaseOffset = actualSeqQPrefixSum / constInfo.sparseShardSize * constInfo.kvHeadNum *
687+ constInfo.sparseBlockCount + info.gS1Idx / constInfo.mBaseSize *
688+ constInfo.sparseBlockCount * constInfo.kvHeadNum +
689+ info.n2Idx * constInfo.sparseBlockCount;
690+ }
691+ info.topkGmBaseOffset = topkGmBaseOffset;
692+ info.tensorAOffset = tensorACoreOffset;
693+ info.tensorBOffset = tensorBCoreOffset;
694+ info.attenOutOffset = tensorACoreOffset;
695+ 
696+ info.curS2BaseOffset = info.s2Idx * SALSK_S2BASEIZE;
697+ info.nextS2BaseOffset = info.curS2BaseOffset + SALSK_S2BASEIZE;
698+ info.s2BatchOffset = s2BatchBaseOffset + info.curS2BaseOffset;
699+ 
700+ info.curActualSeqLenOri = tempLoopInfo.curActualSeqLenOri;
701+ // 计算实际基本块size
702+ if (tempLoopInfo.curActualSeqLen > info.curS2BaseOffset) {
703+ info.actualSingleProcessSInnerSize = tempLoopInfo.curActualSeqLen - info.curS2BaseOffset;
704+ info.actualSingleProcessSInnerSize = info.actualSingleProcessSInnerSize > SALSK_S2BASEIZE ?
705+ SALSK_S2BASEIZE : info.actualSingleProcessSInnerSize;
706+ } else {
707+ info.actualSingleProcessSInnerSize = 0;
708+ }
709+ info.aicS2AccessSize = (info.actualSingleProcessSInnerSize * 3 / 4) / 16 * 16;
710+ info.aivS2AccessSize = info.actualSingleProcessSInnerSize - info.aicS2AccessSize;
711+ info.actualSingleProcessSInnerSizeAlign =
712+ SFAAAlign((uint32_t)info.actualSingleProcessSInnerSize, (uint32_t)SFAAVectorServiceGqa<SFAAT>::BYTE_BLOCK);
713+ info.maskEnd = tempLoopInfo.sparseBlockActualSeqSizeKv - (info.actS1Size - info.gS1Idx / constInfo.mBaseSize * constInfo.sparseShardSize) +
714+ constInfo.sparseShardSize;
715+ info.maskStart = info.maskEnd - constInfo.sparseShardSize + 1;
716+}
717+ 
718+template <typename SFAAT>
719+__aicore__ inline void SparseFlashAttentionAntiquantGqa<SFAAT>::ComputeMm1(const RunInfo &info)
720+{
721+ uint32_t nBufferLoopTimes = (info.actMBaseSize + constInfo.nBufferMBaseSize - 1) / constInfo.nBufferMBaseSize;
722+ uint32_t nBufferTail = info.actMBaseSize - (nBufferLoopTimes - 1) * constInfo.nBufferMBaseSize;
723+ for (uint32_t i = 0; i < nBufferLoopTimes; i++) {
724+ MSplitInfo mSplitInfo;
725+ mSplitInfo.nBufferStartM = i * constInfo.nBufferMBaseSize;
726+ mSplitInfo.nBufferDealM = (i + 1 != nBufferLoopTimes) ? constInfo.nBufferMBaseSize : nBufferTail;
727+ if constexpr (SFAAT::isMsdDD) {
728+ matmulServiceMsd.ComputeMm1(info, mSplitInfo);
729+ }
730+ CrossCoreSetFlag<ConstInfo::SFAA_SYNC_MODE2, PIPE_FIX>(constInfo.syncC1V1);
731+ }
732+}
733+ 
734+template <typename SFAAT>
735+__aicore__ inline void SparseFlashAttentionAntiquantGqa<SFAAT>::ComputeMm2(const RunInfo &info)
736+{
737+ uint32_t nBufferLoopTimes = (info.actMBaseSize + constInfo.nBufferMBaseSize - 1) / constInfo.nBufferMBaseSize;
738+ uint32_t nBufferTail = info.actMBaseSize - (nBufferLoopTimes - 1) * constInfo.nBufferMBaseSize;
739+ for (uint32_t i = 0; i < nBufferLoopTimes; i++) {
740+ MSplitInfo mSplitInfo;
741+ mSplitInfo.nBufferStartM = i * constInfo.nBufferMBaseSize;
742+ mSplitInfo.nBufferDealM = (i + 1 != nBufferLoopTimes) ? constInfo.nBufferMBaseSize : nBufferTail;
743+ CrossCoreWaitFlag(constInfo.syncV1C2);
744+ if constexpr (SFAAT::isMsdDD) {
745+ matmulServiceMsd.ComputeMm2(info, mSplitInfo);
746+ }
747+ CrossCoreSetFlag<ConstInfo::SFAA_SYNC_MODE2, PIPE_FIX>(constInfo.syncC2V2);
748+ }
749+}
750+ 
751+template <typename SFAAT>
752+__aicore__ inline void SparseFlashAttentionAntiquantGqa<SFAAT>::FlashDecode()
753+{
754+ fdService.InitBuffers(pipe);
755+ AscendC::ICachePreLoad(fdPrefetchLen);
756+ SyncAll();
757+ if ASCEND_IS_AIV {
758+ if constexpr (IS_META) {
759+ uint32_t *bN2IdxOfFdHead = metaDataPtr->fdRes.bN2IdxOfFdHead;
760+ uint32_t *gS1IdxOfFdHead = metaDataPtr->fdRes.gS1IdxOfFdHead;
761+ uint32_t *s2SplitNumOfFdHead = metaDataPtr->fdRes.s2SplitNumOfFdHead;
762+ uint32_t *gS1IdxEndOfFdHead = metaDataPtr->fdRes.gS1IdxEndOfFdHead;
763+ uint32_t *gS1IdxEndOfFdHeadSplit = metaDataPtr->fdRes.gS1IdxEndOfFdHeadSplit;
764+ uint32_t *gS1SplitNumOfFdHead = metaDataPtr->fdRes.gS1SplitNumOfFdHead;
765+ uint32_t *gS1LastPartSizeOfFdHead = metaDataPtr->fdRes.gS1LastPartSizeOfFdHead;
766+ FDparams fdParams = {bN2IdxOfFdHead, gS1IdxOfFdHead, s2SplitNumOfFdHead, gS1SplitNumOfFdHead, gS1LastPartSizeOfFdHead,
767+ gS1IdxEndOfFdHead, gS1IdxEndOfFdHeadSplit, metaDataPtr->usedVecNumOfFd,
768+ metaDataPtr->gS1BaseSizeOfFd};
769+ fdService.AllocEventID();
770+ fdService.InitDecodeParams();
771+ fdService.FlashDecode(fdParams);
772+ fdService.FreeEventID();
773+ } else {
774+#ifdef ASCENDC_CPU_DEBUG
775+ const uint32_t *bN2IdxOfFdHead = tilingData->fdParams.bN2IdxOfFdHead;
776+ const uint32_t *gS1IdxOfFdHead = tilingData->fdParams.gS1IdxOfFdHead;
777+ const uint32_t *s2SplitNumOfFdHead = tilingData->fdParams.s2SplitNumOfFdHead;
778+ const uint32_t *gS1IdxEndOfFdHead = tilingData->fdParams.gS1IdxEndOfFdHead;
779+ const uint32_t *gS1IdxEndOfFdHeadSplit = tilingData->fdParams.gS1IdxEndOfFdHeadSplit;
780+ const uint32_t *gS1SplitNumOfFdHead = tilingData->fdParams.gS1SplitNumOfFdHead;
781+ const uint32_t *gS1LastPartSizeOfFdHead = tilingData->fdParams.gS1LastPartSizeOfFdHead;
782+#else
783+ uint32_t bN2IdxOfFdHead[ARRAY_SIZE(tilingData->fdParams.bN2IdxOfFdHead)];
784+ uint32_t gS1IdxOfFdHead[ARRAY_SIZE(tilingData->fdParams.gS1IdxOfFdHead)];
785+ uint32_t s2SplitNumOfFdHead[ARRAY_SIZE(tilingData->fdParams.s2SplitNumOfFdHead)];
786+ uint32_t gS1IdxEndOfFdHead[ARRAY_SIZE(tilingData->fdParams.gS1IdxEndOfFdHead)];
787+ uint32_t gS1IdxEndOfFdHeadSplit[ARRAY_SIZE(tilingData->fdParams.gS1IdxEndOfFdHeadSplit)];
788+ uint32_t gS1SplitNumOfFdHead[ARRAY_SIZE(tilingData->fdParams.gS1SplitNumOfFdHead)];
789+ uint32_t gS1LastPartSizeOfFdHead[ARRAY_SIZE(tilingData->fdParams.gS1LastPartSizeOfFdHead)];
790+ copy_data_align64((uint8_t *)bN2IdxOfFdHead, (uint8_t *)(tilingData->fdParams.bN2IdxOfFdHead),
791+ sizeof(bN2IdxOfFdHead));
792+ copy_data_align64((uint8_t *)gS1IdxOfFdHead, (uint8_t *)(tilingData->fdParams.gS1IdxOfFdHead),
793+ sizeof(gS1IdxOfFdHead));
794+ copy_data_align64((uint8_t *)s2SplitNumOfFdHead, (uint8_t *)(tilingData->fdParams.s2SplitNumOfFdHead),
795+ sizeof(s2SplitNumOfFdHead));
796+ copy_data_align64((uint8_t *)gS1IdxEndOfFdHead, (uint8_t *)(tilingData->fdParams.gS1IdxEndOfFdHead),
797+ sizeof(gS1IdxEndOfFdHead));
798+ copy_data_align64((uint8_t *)gS1IdxEndOfFdHeadSplit,
799+ (uint8_t *)(tilingData->fdParams.gS1IdxEndOfFdHeadSplit),
800+ sizeof(gS1IdxEndOfFdHeadSplit));
801+ copy_data_align64((uint8_t *)gS1SplitNumOfFdHead, (uint8_t *)(tilingData->fdParams.gS1SplitNumOfFdHead),
802+ sizeof(gS1SplitNumOfFdHead));
803+ copy_data_align64((uint8_t *)gS1LastPartSizeOfFdHead,
804+ (uint8_t *)(tilingData->fdParams.gS1LastPartSizeOfFdHead),
805+ sizeof(gS1LastPartSizeOfFdHead));
806+#endif
807+ FDparams fdParams = {bN2IdxOfFdHead, gS1IdxOfFdHead, s2SplitNumOfFdHead, gS1SplitNumOfFdHead, gS1LastPartSizeOfFdHead,
808+ gS1IdxEndOfFdHead, gS1IdxEndOfFdHeadSplit, tilingData->fdParams.usedVecNumOfFd,
809+ tilingData->fdParams.gS1BaseSizeOfFd};
810+ fdService.AllocEventID();
811+ fdService.InitDecodeParams();
812+ fdService.FlashDecode(fdParams);
813+ fdService.FreeEventID();
814+ }
815+ }
816+}
817+ 
818+template <typename SFAAT> __aicore__ inline void SparseFlashAttentionAntiquantGqa<SFAAT>::Process()
819+{
820+ if (aiCoreIdx < usedCoreNum) {
821+ if constexpr (SFAAT::isMsdDD) {
822+ if ASCEND_IS_AIV {
823+ vectorServiceMsd.AllocEventID();
824+ vectorServiceMsd.InitSoftmaxDefaultBuffer();
825+ } else {
826+ matmulServiceMsd.AllocEventID();
827+ }
828+ }
829+
830+ ProcessBalance();
831+ 
832+ if constexpr (SFAAT::isMsdDD) {
833+ if ASCEND_IS_AIV {
834+ vectorServiceMsd.FreeEventID();
835+ } else {
836+ matmulServiceMsd.FreeEventID();
837+ }
838+ }
839+ }
840+ if constexpr (IS_META) {
841+ if (FLASH_DECODE) {
842+ FlashDecode();
843+ }
844+ }
845+}
846+ 
847+template <typename SFAAT>
848+__aicore__ inline void SparseFlashAttentionAntiquantGqa<SFAAT>::GetBN2Idx(uint32_t bN2Idx, uint32_t &bIdx,
849+ uint32_t &n2Idx)
850+{
851+ bIdx = bN2Idx / constInfo.kvHeadNum ;
852+ n2Idx = bN2Idx % constInfo.kvHeadNum ;
853+}
854+ 
855+template <typename SFAAT> __aicore__ inline void SparseFlashAttentionAntiquantGqa<SFAAT>::ProcessBalance()
856+{
857+ RunInfo extraInfo[SFAA_PRELOAD_TASK_CACHE_SIZE];
858+ uint32_t gloop = 0;
859+ int gS1LoopEnd;
860+ bool globalLoopStart = true;
861+ if ASCEND_IS_AIC {
862+ CrossCoreSetFlag<ConstInfo::SFAA_SYNC_MODE2, PIPE_MTE2>(3);
863+ CrossCoreSetFlag<ConstInfo::SFAA_SYNC_MODE2, PIPE_MTE2>(3);
864+ CrossCoreSetFlag<ConstInfo::SFAA_SYNC_MODE2, PIPE_MTE2>(3);
865+ CrossCoreSetFlag<ConstInfo::SFAA_SYNC_MODE2, PIPE_MTE2>(3);
866+ }
867+ for (uint32_t bN2LoopIdx = constInfo.bN2Start; bN2LoopIdx <= constInfo.bN2End; bN2LoopIdx++) {
868+ GetBN2Idx(bN2LoopIdx, tempLoopInfo.bIdx, tempLoopInfo.n2Idx);
869+ tempLoopInfo.bN2Idx = bN2LoopIdx;
870+ GetActualSeqLen(tempLoopInfo.bIdx); // 获取actualSeqLength及ActualSeqLengthKV
871+ GetPreNextTokensLeftUp();
872+ if (tempLoopInfo.actS1Size == 0 && !(bN2LoopIdx == constInfo.bN2End)) {
873+ continue;
874+ }
875+ uint32_t gS1SplitNum;
876+ uint32_t s1SizeTail;
877+ if (tempLoopInfo.actS1Size == 0) {
878+ gS1SplitNum = 0;
879+ s1SizeTail = 0;
880+ } else {
881+ gS1SplitNum = (tempLoopInfo.actS1Size + constInfo.sparseShardSize - 1) / constInfo.sparseShardSize;
882+ s1SizeTail = tempLoopInfo.actS1Size - (gS1SplitNum - 1) * constInfo.sparseShardSize;
883+ }
884+ gS1LoopEnd = (bN2LoopIdx == constInfo.bN2End) ? constInfo.gS1End : gS1SplitNum - 1;
885+ for (uint32_t gS1LoopIdx = constInfo.gS1Start; gS1LoopIdx <= gS1LoopEnd; gS1LoopIdx++) {
886+ tempLoopInfo.gS1Idx = gS1LoopIdx * constInfo.mBaseSize;
887+ uint32_t s1Size = (gS1LoopIdx == gS1LoopEnd) ? s1SizeTail : constInfo.sparseShardSize;
888+ // TopK值sparse完后的ActualSeqLengthKV
889+ uint32_t s1StartIdx = gS1LoopIdx * constInfo.sparseShardSize;
890+ uint32_t s1EndIdx;
891+ if (s1Size == 0) {
892+ s1EndIdx = 0;
893+ } else {
894+ s1EndIdx = s1StartIdx + s1Size - 1;
895+ }
896+ GetSparseActualSeqLen(tempLoopInfo.bIdx, s1EndIdx, tempLoopInfo.n2Idx);
897+ UpdateInnerLoopCond();
898+ 
899+ if (tempLoopInfo.curActSeqLenIsZero) {
900+ DealActSeqLenIsZero(tempLoopInfo.bIdx, s1StartIdx, s1Size, tempLoopInfo.n2Idx);
901+ }
902+ int s2SplitNum =
903+ (tempLoopInfo.curActualSeqLen + SALSK_S2BASEIZE - 1) / SALSK_S2BASEIZE; // S2切分份数
904+ bool isEnd = (bN2LoopIdx == constInfo.bN2End) && (gS1LoopIdx == constInfo.gS1End);
905+ if constexpr (IS_META) {
906+ tempLoopInfo.s2LoopTimes = (bN2LoopIdx == constInfo.bN2End && gS1LoopIdx == constInfo.gS1End && !tempLoopInfo.curActSeqLenIsZero) ? constInfo.s2End + 1 : s2SplitNum;
907+ } else {
908+ tempLoopInfo.s2LoopTimes = s2SplitNum;
909+ }
910+ // 分核修改后需要打开
911+ // 当前s2是否被切，决定了输出是否要写到attenOut上
912+ tempLoopInfo.tndIsS2SplitCore =
913+ ((constInfo.s2Start == 0) && (tempLoopInfo.s2LoopTimes == s2SplitNum)) ? false : true;
914+ tempLoopInfo.tndCoreStartKVSplitPos = globalLoopStart ? constInfo.coreStartKVSplitPos : 0;
915+ uint32_t extraLoop = isEnd ? 2 : 0;
916+ for (int s2LoopIdx = constInfo.s2Start; s2LoopIdx < (tempLoopInfo.s2LoopTimes + extraLoop); s2LoopIdx++) {
917+ // PreloadPipeline loop初始值要求为 PRELOAD_NUM
918+ PreloadPipeline(gloop, constInfo.s2Start, s2LoopIdx, extraInfo);
919+ ++gloop;
920+ }
921+ globalLoopStart = false;
922+ constInfo.s2Start = 0;
923+ }
924+ constInfo.gS1Start = 0;
925+ }
926+ if ASCEND_IS_AIV {
927+ CrossCoreWaitFlag(3);
928+ CrossCoreWaitFlag(3);
929+ CrossCoreWaitFlag(3);
930+ CrossCoreWaitFlag(3);
931+ }
932+}
933+ 
934+template <typename SFAAT>
935+__aicore__ inline void
936+SparseFlashAttentionAntiquantGqa<SFAAT>::PreloadPipeline(uint32_t loop, uint64_t s2Start, uint64_t s2LoopIdx,
937+ RunInfo extraInfo[SFAA_PRELOAD_TASK_CACHE_SIZE])
938+{
939+ RunInfo &extraInfo0 = extraInfo[loop % SFAA_PRELOAD_TASK_CACHE_SIZE]; // 本轮任务
940+ RunInfo &extraInfo2 = extraInfo[(loop + 2) % SFAA_PRELOAD_TASK_CACHE_SIZE]; // 上一轮任务
941+ RunInfo &extraInfo1 = extraInfo[(loop + 1) % SFAA_PRELOAD_TASK_CACHE_SIZE]; // 上两轮任务
942+ 
943+ CalcParams(loop, s2Start, s2LoopIdx, extraInfo0);
944+ 
945+ if (extraInfo0.isValid) {
946+ if ASCEND_IS_AIC {
947+ CrossCoreWaitFlag(constInfo.syncV0C1);
948+ ComputeMm1(extraInfo0);
949+ } else {
950+ if constexpr (SFAAT::isMsdDD) {
951+ vectorServiceMsd.ProcessVec0Msd(extraInfo0);
952+ }
953+ CrossCoreSetFlag<ConstInfo::SFAA_SYNC_MODE2, PIPE_MTE3>(constInfo.syncV0C1);
954+ }
955+ }
956+ if (extraInfo2.isValid) {
957+ if ASCEND_IS_AIV {
958+ if constexpr (SFAAT::isMsdDD) {
959+ vectorServiceMsd.ProcessVec1Msd(extraInfo2);
960+ }
961+ }
962+ if ASCEND_IS_AIC {
963+ ComputeMm2(extraInfo2);
964+ }
965+ }
966+ if (extraInfo1.isValid) {
967+ if ASCEND_IS_AIV {
968+ if constexpr (SFAAT::isMsdDD) {
969+ vectorServiceMsd.ProcessVec2Msd(extraInfo1);
970+ }
971+ }
972+ extraInfo1.isValid = false;
973+ }
974+}
975+ 
976+template <typename SFAAT>
977+__aicore__ inline uint64_t
978+SparseFlashAttentionAntiquantGqa<SFAAT>::GetBalanceActualSeqLengths(GlobalTensor<int32_t> &actualSeqLengths,
979+ uint32_t bIdx)
980+{
981+ if constexpr (LAYOUT_T == SFAA_LAYOUT::TND) {
982+ if (bIdx > 0) {
983+ return actualSeqLengths.GetValue(bIdx) - actualSeqLengths.GetValue(bIdx - 1);
984+ } else if (bIdx == 0) {
985+ return actualSeqLengths.GetValue(0);
986+ } else {
987+ return 0;
988+ }
989+ } else {
990+ if (constInfo.actualLenDimsQ == 0) {
991+ return constInfo.qSeqSize;
992+ } else if (constInfo.actualLenDimsQ == 1) {
993+ return actualSeqLengths.GetValue(0);
994+ } else {
995+ return actualSeqLengths.GetValue(bIdx);
996+ }
997+ }
998+}
999+ 
1000+template <typename SFAAT>
1001+__aicore__ inline void SparseFlashAttentionAntiquantGqa<SFAAT>::GetAxisStartIdx(uint32_t bN2EndPrev,
1002+ uint32_t s1GEndPrev,
1003+ uint32_t s2EndPrev)
1004+{
1005+ if constexpr (IS_META) {
1006+ uint32_t bEndPrev = bN2EndPrev / constInfo.kvHeadNum ;
1007+ uint32_t actualSeqQPrev = GetBalanceActualSeqLengths(actualSeqLengthsQGm, bEndPrev);
1008+ uint32_t s1GPrevBaseNum = (actualSeqQPrev * constInfo.gSize + constInfo.mBaseSize - 1) / constInfo.mBaseSize;
1009+
1010+ // get s2PrevBaseNum
1011+ uint32_t actualSeqKVPrev = GetActualSeqLenKV(bEndPrev);
1012+ 
1013+ uint32_t sparseBlockActualSeqSizeKv = 0;
1014+ if (constInfo.sparseLenDimsKV == 0) {
1015+ sparseBlockActualSeqSizeKv = constInfo.sparseBlockCount * constInfo.sparseBlockSize;
1016+ } else {
1017+ sparseBlockActualSeqSizeKv = sparseSeqLengthsKVGm.GetValue(bEndPrev);
1018+ }
1019+ uint32_t curActualSeqLen = (sparseBlockActualSeqSizeKv > actualSeqKVPrev) ? actualSeqKVPrev : sparseBlockActualSeqSizeKv;
1020+ uint32_t s2PrevBaseNum = (curActualSeqLen + SALSK_S2BASEIZE - 1) / SALSK_S2BASEIZE;
1021+ 
1022+ constInfo.bN2Start = bN2EndPrev;
1023+ constInfo.gS1Start = s1GEndPrev;
1024+ constInfo.s2Start = s2EndPrev + 1U;
1025+ 
1026+ if (constInfo.s2Start >= s2PrevBaseNum) {
1027+ constInfo.gS1Start++;
1028+ constInfo.s2Start = 0;
1029+ }
1030+ if (constInfo.gS1Start >= s1GPrevBaseNum) {
1031+ constInfo.bN2Start++;
1032+ constInfo.gS1Start = 0;
1033+ }
1034+ } else {
1035+ uint32_t bEndPrev = bN2EndPrev / constInfo.kvHeadNum ;
1036+ uint32_t actualSeqQPrev = GetBalanceActualSeqLengths(actualSeqLengthsQGm, bEndPrev);
1037+ uint32_t s1GPrevBaseNum = (actualSeqQPrev * constInfo.gSize + constInfo.mBaseSize - 1) / constInfo.mBaseSize;
1038+ constInfo.bN2Start = bN2EndPrev;
1039+ constInfo.gS1Start = s1GEndPrev;
1040+
1041+ constInfo.s2Start = 0;
1042+ if (s1GEndPrev >= s1GPrevBaseNum - 1) { // 上个核把S1G处理完了
1043+ constInfo.gS1Start = 0;
1044+ constInfo.bN2Start++;
1045+ } else {
1046+ constInfo.gS1Start++;
1047+ }
1048+ }
1049+}
1050+#endif // SPARSE_FLASH_ATTENTION_ANTIQUANT_KERNEL_GQA_H
Aexperimental/attention/sparse_flash_attention_antiquant/op_kernel/sparse_flash_attention_antiquant_metadata.h+58-0
@@ -0,0 +1,58 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+
11+/*!
12+ * \file sparse_flash_attention_antiquant_metadata.h
13+ * \brief
14+ */
15+ 
16+#ifndef SPARSE_FLASH_ATTENTION_ANTIQUANT_METADATA_H
17+#define SPARSE_FLASH_ATTENTION_ANTIQUANT_METADATA_H
18+ 
19+#include <cstdint>
20+ 
21+namespace optiling {
22+const uint32_t CORE_NUM = 24;
23+constexpr uint32_t SFA_META_SIZE = 1024;
24+using SFA_METADATA_T = int32_t;
25+ 
26+namespace detail {
27+ // 分核功能模块输出：FD信息，包含需要归约的数据索引及其分核信息
28+ struct FlashDecodeResult {
29+ // 1、归约任务的索引信息
30+ uint32_t bN2IdxOfFdHead[CORE_NUM]; // 每个归约任务的BN2索引，脚标为归约任务的序号，最大为核数-1
31+ uint32_t gS1IdxOfFdHead[CORE_NUM]; // 每个归约任务的GS1索引，脚标为归约任务的序号
32+ uint32_t s2SplitNumOfFdHead[CORE_NUM]; // 每个归约任务的S2核间切分份数，脚标为归约任务的序号
33+ // 2、FD负载均衡阶段，归约任务的分核（vec）信息
34+ uint32_t gS1SplitNumOfFdHead[CORE_NUM]; // 每个归约任务m轴切分份数，脚标为归约任务的序号
35+ uint32_t gS1LastPartSizeOfFdHead[CORE_NUM]; // 每个归约任务m轴切分的最后一份的大小，脚标为归约任务的序号
36+ uint32_t gS1IdxEndOfFdHead[CORE_NUM * 2]; // FD负载均衡阶段，每个vector的一级索引，脚标为vector ID，值为归约任务的ID
37+ uint32_t gS1IdxEndOfFdHeadSplit[CORE_NUM * 2]; // FD负载均衡阶段，每个vector的二级索引，脚标为vector ID，值为归约任务的m轴切分ID
38+ // 3、每个core处理的第1个归约任务的数据应存放的workspace位置
39+ uint32_t s2SplitStartIdxOfCore[CORE_NUM];
40+ };
41+ 
42+ struct SfaMetaData { // __attribute__((aligned(8)))
43+ uint32_t bN2End[CORE_NUM]; // 每个核处理数据的BN2结束点
44+ uint32_t gS1End[CORE_NUM]; // 每个核处理数据的GS1结束点
45+ uint32_t s2End[CORE_NUM]; // 每个核处理数据的S2结束点
46+ uint32_t usedCoreNum = 0U; // 使用的核数量
47+ uint32_t numOfFdHead = 0U; // 归约任务数量
48+ uint32_t usedVecNumOfFd = 0U; // 归约过程使用的vector数量
49+ uint32_t mBaseSize = 0U;
50+ uint32_t s2BaseSize = 0U;
51+ uint32_t gS1BaseSizeOfFd = 0U;
52+ struct FlashDecodeResult fdRes; // FD信息
53+ };
54+};
55+static_assert(SFA_META_SIZE * sizeof(SFA_METADATA_T) >= sizeof(detail::SfaMetaData));
56+};
57+ 
58+#endif
Aexperimental/attention/sparse_flash_attention_antiquant/op_kernel/sparse_flash_attention_antiquant_service_cube_gqa_msd.h+1363-0
@@ -0,0 +1,1363 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file sparse_flash_attention_antiquant_service_cube_gqa.h
13+ * \brief use 7 buffer for matmul l1, better pipeline
14+ */
15+#ifndef SPARSE_FLASH_ATTENTION_ANTIQUANT_SERVICE_CUBE_GQA_MSD_H
16+#define SPARSE_FLASH_ATTENTION_ANTIQUANT_SERVICE_CUBE_GQA_MSD_H
17+ 
18+#include "kernel_operator.h"
19+#include "kernel_operator_list_tensor_intf.h"
20+#include "kernel_tiling/kernel_tiling.h"
21+#include "lib/matmul_intf.h"
22+#include "lib/matrix/matmul/tiling.h"
23+#include "sparse_flash_attention_antiquant_common.h"
24+ 
25+#define SALSC_S2BASEIZE 1024
26+#define SALSC_S2MERGESIZE 384
27+#define SALSC_S2BASEIZE_1 1023
28+ 
29+template <typename SFAAT> class SFAAMatmulServiceGqaMsd {
30+public:
31+ // 中间计算数据类型为float, 高精度模式
32+ using T = float;
33+ using Q_T = typename SFAAT::queryType;
34+ using KV_T = typename SFAAT::kvType;
35+ using KV_ORIGIN_T = typename SFAAT::queryType;
36+ using OUT_T = typename SFAAT::outputType;
37+ using MM_OUT_T = half;
38+ using L0C_T = int32_t;
39+ 
40+ __aicore__ inline SFAAMatmulServiceGqaMsd(){};
41+ __aicore__ inline void InitParams(const ConstInfo &constInfo);
42+ __aicore__ inline void InitMm1GlobalTensor(GlobalTensor<Q_T> queryGm, GlobalTensor<KV_T> keyGm,
43+ GlobalTensor<MM_OUT_T> mm1ResGm);
44+ __aicore__ inline void InitMm2GlobalTensor(GlobalTensor<KV_T> vec1ResGm, GlobalTensor<KV_T> valueGm,
45+ GlobalTensor<MM_OUT_T> mm2ResGm, GlobalTensor<OUT_T> attentionOutGm);
46+ __aicore__ inline void InitPageAttentionInfo(const GlobalTensor<KV_T>& keyMergeGm, const GlobalTensor<KV_T>& valueMergeGm,
47+ const GlobalTensor<KV_T>& queryPreProcessResGm, GlobalTensor<int32_t> blockTableGm,
48+ GlobalTensor<int32_t> topKGm, uint32_t blockSize, uint32_t maxBlockNumPerBatch);
49+ __aicore__ inline void InitBuffers(TPipe *pipe);
50+ __aicore__ inline void UpdateKey(GlobalTensor<KV_T> keyGm);
51+ __aicore__ inline void UpdateValue(GlobalTensor<KV_T> valueGm);
52+ 
53+ __aicore__ inline void AllocEventID();
54+ __aicore__ inline void FreeEventID();
55+ __aicore__ inline void CalcTopKBlockInfo(const RunInfo &info, uint32_t &curTopKIdx,
56+ uint64_t &curOffsetInSparseBlock, uint32_t curSeqIdx,
57+ uint32_t &copyRowCnt, uint64_t &idInTopK);
58+ __aicore__ inline void ComputeMm1(const RunInfo &info, const MSplitInfo mSplitInfo);
59+ __aicore__ inline void ComputeMm2(const RunInfo &info, const MSplitInfo mSplitInfo);
60+ 
61+private:
62+ static constexpr bool PAGE_ATTENTION = SFAAT::pageAttention;
63+ static constexpr int TEMPLATE_MODE = SFAAT::templateMode;
64+ static constexpr bool FLASH_DECODE = SFAAT::flashDecode;
65+ static constexpr SFAA_LAYOUT LAYOUT_T = SFAAT::layout;
66+ static constexpr SFAA_LAYOUT KV_LAYOUT_T = SFAAT::kvLayout;
67+ 
68+ static constexpr float quantScaleC1S1 = 1.0 / (1024);
69+ static constexpr float quantScaleC1S2 = 1.0 / (1024 * 254);
70+ static constexpr float quantScaleC2O1 = 1.0 / (1024 * 127);
71+ static constexpr float quantScaleC2O2 = 1.0 / (1024 * 254 * 127);
72+ static constexpr uint32_t msdIterNum = 2;
73+ static constexpr uint32_t P_LOAD_TO_L1_ROW_NUM = 128 / sizeof(KV_T);
74+ static constexpr uint32_t KV_LOAD_TO_L1_ROW_NUM = 512 / sizeof(KV_T);
75+ static constexpr uint64_t DATABLOCK_BYTES = 32UL;
76+ // L1轴切分大小
77+ static constexpr uint32_t S2_SPLIT_SIZE = 512; // S2方向切分，对于mm1是N轴，对于mm2是K轴
78+ 
79+ // L0轴切分大小
80+ static constexpr uint32_t M_BASE_SIZE = 128; // m方向基本块大小
81+ static constexpr uint32_t K_BASE_SIZE = 128; // k方向基本块大小
82+ static constexpr uint32_t N_BASE_SIZE = 256; // n方向基本块大小
83+ static constexpr uint32_t L1_NZ_BLOCK_SIZE_MM1 = 256;
84+ static constexpr uint32_t L1_NZ_BLOCK_SIZE_MM1_SHIFT = 8;
85+ static constexpr uint32_t L1_NZ_BLOCK_SIZE_MM1_MASK = L1_NZ_BLOCK_SIZE_MM1 - 1;
86+ static constexpr uint32_t L1_NZ_BLOCK_SIZE_MM2 = 32;
87+ static constexpr uint32_t L1_NZ_BLOCK_SIZE_MM2_SHIFT = 5;
88+ static constexpr uint32_t L1_NZ_BLOCK_SIZE_MM2_MASK = L1_NZ_BLOCK_SIZE_MM2 - 1;
89+ static constexpr uint32_t L1_NZ_KV_BLOCK_ELEMENT = 32;
90+ 
91+ static constexpr uint32_t L1Q_BLOCK_SIZE= 64 * 1024; // 64K
92+ static constexpr uint32_t L1KP_BLOCK_SIZE = 64 * 1024; // 64k
93+ static constexpr uint32_t L1V_BLOCK_SIZE = 64 * 1024; // 64k
94+ 
95+ static constexpr uint32_t L0A_PP_SIZE = (32 * 1024);
96+ static constexpr uint32_t L0B_PP_SIZE = (32 * 1024);
97+ static constexpr uint32_t L0C_PP_SIZE = (64 * 1024);
98+ static constexpr uint32_t L0AB_BLOCK_OFFSET = L0A_PP_SIZE / sizeof(KV_T);
99+ static constexpr uint32_t L0C_BLOCK_OFFSET = L0C_PP_SIZE / sizeof(MM_OUT_T);
100+ 
101+ // mte2 <> mte1 EventID
102+ // L1 3buf, 使用3个eventId
103+ static constexpr uint32_t L1Q_EVENT0 = EVENT_ID0;
104+ static constexpr uint32_t L1KP_EVENT0 = EVENT_ID1;
105+ static constexpr uint32_t L1KP_EVENT1 = EVENT_ID2;
106+ static constexpr uint32_t L1KP_EVENT2 = EVENT_ID3;
107+ static constexpr uint32_t L1V_EVENT0 = EVENT_ID4;
108+ static constexpr uint32_t L1V_EVENT1 = EVENT_ID5;
109+ static constexpr uint32_t L1V_EVENT2 = EVENT_ID6;
110+ static constexpr uint32_t L1V_EVENT3 = EVENT_ID7;
111+ static constexpr uint32_t L1V_BUFFER_NUM = 4;
112+ static constexpr uint32_t L1KP_BUFFER_NUM = 3;
113+ 
114+ // m <> mte1 EventID
115+ static constexpr uint32_t L0A_EVENT0 = EVENT_ID3;
116+ static constexpr uint32_t L0A_EVENT1 = EVENT_ID4;
117+ static constexpr uint32_t L0B_EVENT0 = EVENT_ID5;
118+ static constexpr uint32_t L0B_EVENT1 = EVENT_ID6;
119+ 
120+ // fix <> m
121+ static constexpr uint32_t L0C_EVENT0 = EVENT_ID3;
M
Mmiaofangzheng4月2日

MSD 变体同样存在 event ID 复用：L0A_EVENT0 和 L0C_EVENT0 共用 EVENT_ID3，L0A_EVENT1 和 L0C_EVENT1 共用 EVENT_ID4。如果 M↔MTE1 同步和 FIX↔M 同步存在时序重叠，可能导致数据竞争。

likedislike
M
mayafei
4月3日 评论：
122+ static constexpr uint32_t L0C_EVENT1 = EVENT_ID4;
123+ 
124+ static constexpr IsResetLoad3dConfig LOAD3DV2_CONFIG = {true, true}; // isSetFMatrix isSetPadding;
125+ 
126+ static constexpr uint32_t BLOCK_ELEMENT_NUM = ConstInfo::BUFFER_SIZE_BYTE_32B / sizeof(KV_T);
127+ 
128+ uint32_t kvCacheBlockSize = 0;
129+ uint32_t maxBlockNumPerBatch = 0;
130+ uint32_t headDimAlign = 0;
131+ ConstInfo constInfo{};
132+ 
133+ // L1分成3块buf, 用于记录
134+ uint32_t qL1BufIter = 0;
135+ uint32_t kpL1BufIter = -1;
136+ uint32_t vL1BufIter = -1;
137+ uint32_t aL0BufIter = 0;
138+ uint32_t bL0BufIter = 0;
139+ uint32_t cL0BufIter = 0;
140+ 
141+ // mm1
142+ GlobalTensor<Q_T> queryGm;
143+ GlobalTensor<KV_T> keyGm;
144+ GlobalTensor<MM_OUT_T> mm1ResGm;
145+ GlobalTensor<KV_T> keyMergeGm_;
146+ GlobalTensor<KV_T> queryPreProcessResGm_;
147+ 
148+ // mm2
149+ GlobalTensor<KV_T> vec1ResGm;
150+ GlobalTensor<KV_T> valueGm;
151+ GlobalTensor<MM_OUT_T> mm2ResGm;
152+ GlobalTensor<OUT_T> attentionOutGm;
153+ GlobalTensor<KV_T> valueMergeGm_;
154+ 
155+ // block_table
156+ GlobalTensor<int32_t> blockTableGm;
157+ GlobalTensor<int32_t> topKGm;
158+ 
159+ TBuf<TPosition::A1> tmpBufL1Q;
160+ LocalTensor<KV_T> qL1Buffers;
161+ TBuf<TPosition::A1> tmpBufL1KP;
162+ LocalTensor<KV_T> kpL1Buffers;
163+ TBuf<TPosition::A1> tmpBufL1V;
164+ LocalTensor<KV_T> vL1Buffers;
165+ 
166+ TBuf<TPosition::A2> tmpBufL0A;
167+ LocalTensor<KV_T> aL0TensorPingPong;
168+ // L0B
169+ TBuf<TPosition::B2> tmpBufL0B;
170+ LocalTensor<KV_T> bL0TensorPingPong;
171+ // L0C
172+ TBuf<TPosition::CO1> tmpBufL0C;
173+ LocalTensor<L0C_T> cL0TensorPingPong;
174+ 
175+ __aicore__ inline void CopyGmToL1(LocalTensor<KV_T> &l1Tensor, GlobalTensor<KV_T> &gmSrcTensor,
176+ uint32_t srcN, uint32_t srcD, uint32_t srcDstride);
177+ __aicore__ inline void CopyInMm1AToL1(LocalTensor<KV_T> &l1Tensor, const RunInfo &info,
178+ uint32_t mSizeAct, uint32_t headSize);
179+ __aicore__ inline void CopyInMm1AToL1(LocalTensor<KV_T> &aL1Tensor, const RunInfo &info);
180+ __aicore__ inline void CopyInMm2AToL1(LocalTensor<KV_T> &aL1Tensor, const RunInfo &info, uint32_t mSeqIdx,
181+ uint32_t subMSizeAct, uint32_t nSize, uint32_t nOffset);
182+ __aicore__ inline void CopyInMm2AToL1(LocalTensor<KV_T>& aL1Tensor, const RunInfo &info, uint32_t mCopyIdx,
183+ uint32_t mCopyRowCount, uint32_t mActCopyRowCount,
184+ uint32_t kCopyIdx, uint32_t kCopyRowCount, uint32_t kActCopyRowCount);
185+ __aicore__ inline void CopyInMm2BToL1(LocalTensor<KV_T>& bL1Tensor, const RunInfo &info,
186+ uint32_t kCopyIdx, uint32_t kActCopyRowCountAlign, uint32_t kActCopyRowCount);
187+ __aicore__ inline void LoadDataMm1A(LocalTensor<KV_T> &aL0Tensor, LocalTensor<KV_T> &aL1Tensor,
188+ uint32_t idx, uint32_t kSplitSize, uint32_t mSize, uint32_t kSize);
189+ __aicore__ inline void LoadDataMm1B(LocalTensor<KV_T> &bL0Tensor, LocalTensor<KV_T> &bL1Tensor,
190+ uint32_t idx, uint32_t kSplitSize, uint32_t kSize, uint32_t nSize);
191+ __aicore__ inline void CopyInMmBToL1(LocalTensor<KV_T> &bl1Tensor, GlobalTensor<KV_T> &gmSrcTensor,
192+ uint32_t subNSizeAct, uint32_t nOffset, const RunInfo &info);
193+ __aicore__ inline void CopyInMm1BToL1VecMerge(LocalTensor<KV_T>& bL1Tensor, const RunInfo &info,
194+ uint32_t aicNZBlockBum, uint32_t aicNzBlockTail, uint32_t VecAccessSize, uint32_t alreadyVecAccessSize,
195+ uint32_t nActCopyRowCountAlign, uint32_t CubeAccessSize);
196+ __aicore__ inline void CopyInMm2BToL1VecMerge(LocalTensor<KV_T>& bL1Tensor, const RunInfo &info,
197+ uint32_t aicNZBlockBum, uint32_t aicNzBlockTail, uint32_t VecAccessSize, uint32_t alreadyVecAccessSize,
198+ uint32_t kActCopyRowCountAlign, uint32_t CubeAccessSize);
199+ __aicore__ inline void CopyInMm1BToL1CubeDisepr(LocalTensor<KV_T>& bL1Tensor, const RunInfo &info, uint32_t nCopyIdx,
200+ uint32_t nCopyRowCount, uint32_t nActCopyRowCount, uint32_t nActCopyRowCountAlign);
201+ __aicore__ inline void CopyInMm2BToL1CubeDisepr(LocalTensor<KV_T>& bL1Tensor, const RunInfo &info, uint32_t nCopyIdx,
202+ uint32_t nCopyRowCount, uint32_t nActCopyRowCount, uint32_t nActCopyRowCountAlign);
203+ __aicore__ inline void CopyInMm1BToL1(LocalTensor<KV_T>& bL1Tensor, const RunInfo &info, uint32_t nCopyIdx,
204+ uint32_t nCopyRowCount, uint32_t nActCopyRowCount, uint32_t nActCopyRowCountAlign);
205+ __aicore__ inline void CopyInMm1BToL1ForPA(LocalTensor<KV_T>& bL1Tensor, uint64_t keyGmBaseOffset,
206+ uint32_t copyTotalRowCnt, uint32_t copyStartRowCnt, uint32_t nActCopyRowCount, const RunInfo &info);
207+ __aicore__ inline void CopyInMm2BToL1ForPA(LocalTensor<KV_T>& bL1Tensor, uint64_t valueGmBaseOffset,
208+ uint32_t copyStartRowCnt, uint32_t kActCopyRowCount);
209+ __aicore__ inline void LoadDataMm2A(LocalTensor<KV_T> aL0Tensor, LocalTensor<KV_T> aL1Tensor, uint32_t kSize);
210+};
211+ 
212+template <typename SFAAT> __aicore__ inline void SFAAMatmulServiceGqaMsd<SFAAT>::InitParams(const ConstInfo &constInfo)
213+{
214+ this->constInfo = constInfo;
215+ headDimAlign = SFAAAlign(constInfo.headDim, ConstInfo::BUFFER_SIZE_BYTE_32B);
216+}
217+ 
218+template <typename SFAAT>
219+__aicore__ inline void
220+SFAAMatmulServiceGqaMsd<SFAAT>::InitMm1GlobalTensor(GlobalTensor<Q_T> queryGm, GlobalTensor<KV_T> keyGm,
221+ GlobalTensor<MM_OUT_T> mm1ResGm)
222+{
223+ // mm1
224+ this->queryGm = queryGm;
225+ this->keyGm = keyGm;
226+ this->mm1ResGm = mm1ResGm;
227+}
228+ 
229+template <typename SFAAT>
230+__aicore__ inline void
231+SFAAMatmulServiceGqaMsd<SFAAT>::InitMm2GlobalTensor(GlobalTensor<KV_T> vec1ResGm, GlobalTensor<KV_T> valueGm,
232+ GlobalTensor<MM_OUT_T> mm2ResGm, GlobalTensor<OUT_T> attentionOutGm)
233+{
234+ // mm2
235+ this->vec1ResGm = vec1ResGm;
236+ this->valueGm = valueGm;
237+ this->mm2ResGm = mm2ResGm;
238+ this->attentionOutGm = attentionOutGm;
239+}
240+ 
241+template <typename SFAAT>
242+__aicore__ inline void
243+SFAAMatmulServiceGqaMsd<SFAAT>::InitPageAttentionInfo(const GlobalTensor<KV_T>& keyMergeGm, const GlobalTensor<KV_T>& valueMergeGm,
244+ const GlobalTensor<KV_T>& queryPreProcessResGm, GlobalTensor<int32_t> blockTableGm,
245+ GlobalTensor<int32_t> topKGm, uint32_t blockSize, uint32_t maxBlockNumPerBatch)
246+{
247+ this->blockTableGm = blockTableGm;
248+ this->topKGm = topKGm;
249+ this->kvCacheBlockSize = blockSize;
250+ this->maxBlockNumPerBatch = maxBlockNumPerBatch;
251+ this->keyMergeGm_ = keyMergeGm;
252+ this->valueMergeGm_ = valueMergeGm;
253+ this->queryPreProcessResGm_ = queryPreProcessResGm;
254+}
255+ 
256+template <typename SFAAT> __aicore__ inline void SFAAMatmulServiceGqaMsd<SFAAT>::InitBuffers(TPipe *pipe)
257+{
258+ pipe->InitBuffer(tmpBufL1Q, L1Q_BLOCK_SIZE); // (64K) * 1
259+ qL1Buffers = tmpBufL1Q.Get<KV_T>();
260+ pipe->InitBuffer(tmpBufL1KP, L1KP_BLOCK_SIZE * L1KP_BUFFER_NUM); // 64K * 3
261+ kpL1Buffers = tmpBufL1KP.Get<KV_T>();
262+ pipe->InitBuffer(tmpBufL1V, L1V_BLOCK_SIZE * L1V_BUFFER_NUM); // 64K * 4
263+ vL1Buffers = tmpBufL1V.Get<KV_T>();
264+ 
265+ // L0A
266+ pipe->InitBuffer(tmpBufL0A, L0A_PP_SIZE * 2); // 64K
267+ aL0TensorPingPong = tmpBufL0A.Get<KV_T>();
268+ // L0B
269+ pipe->InitBuffer(tmpBufL0B, L0B_PP_SIZE * 2); // 64K
270+ bL0TensorPingPong = tmpBufL0B.Get<KV_T>();
271+ // L0C
272+ pipe->InitBuffer(tmpBufL0C, L0C_PP_SIZE * 2); // 128K
273+ cL0TensorPingPong = tmpBufL0C.Get<L0C_T>();
274+}
275+ 
276+template <typename SFAAT> __aicore__ inline void SFAAMatmulServiceGqaMsd<SFAAT>::UpdateKey(GlobalTensor<KV_T> keyGm)
277+{
278+ this->keyGm = keyGm;
279+}
280+ 
281+template <typename SFAAT> __aicore__ inline void SFAAMatmulServiceGqaMsd<SFAAT>::UpdateValue(GlobalTensor<KV_T> valueGm)
282+{
283+ this->valueGm = valueGm;
284+}
285+ 
286+template <typename SFAAT> __aicore__ inline void SFAAMatmulServiceGqaMsd<SFAAT>::AllocEventID()
287+{
288+ SetFlag<HardEvent::MTE1_MTE2>(L1Q_EVENT0);
289+ SetFlag<HardEvent::MTE1_MTE2>(L1KP_EVENT0);
290+ SetFlag<HardEvent::MTE1_MTE2>(L1KP_EVENT1);
291+ SetFlag<HardEvent::MTE1_MTE2>(L1KP_EVENT2);
292+ SetFlag<HardEvent::MTE1_MTE2>(L1V_EVENT0);
293+ SetFlag<HardEvent::MTE1_MTE2>(L1V_EVENT1);
294+ SetFlag<HardEvent::MTE1_MTE2>(L1V_EVENT2);
295+ SetFlag<HardEvent::MTE1_MTE2>(L1V_EVENT3);
296+ SetFlag<HardEvent::M_MTE1>(L0A_EVENT0);
297+ SetFlag<HardEvent::M_MTE1>(L0A_EVENT1);
298+ SetFlag<HardEvent::M_MTE1>(L0B_EVENT0);
299+ SetFlag<HardEvent::M_MTE1>(L0B_EVENT1);
300+ SetFlag<HardEvent::FIX_M>(L0C_EVENT0);
301+ SetFlag<HardEvent::FIX_M>(L0C_EVENT1);
302+}
303+ 
304+template <typename SFAAT> __aicore__ inline void SFAAMatmulServiceGqaMsd<SFAAT>::FreeEventID()
305+{
306+ WaitFlag<HardEvent::MTE1_MTE2>(L1Q_EVENT0);
307+ WaitFlag<HardEvent::MTE1_MTE2>(L1KP_EVENT0);
308+ WaitFlag<HardEvent::MTE1_MTE2>(L1KP_EVENT1);
309+ WaitFlag<HardEvent::MTE1_MTE2>(L1KP_EVENT2);
310+ WaitFlag<HardEvent::MTE1_MTE2>(L1V_EVENT0);
311+ WaitFlag<HardEvent::MTE1_MTE2>(L1V_EVENT1);
312+ WaitFlag<HardEvent::MTE1_MTE2>(L1V_EVENT2);
313+ WaitFlag<HardEvent::MTE1_MTE2>(L1V_EVENT3);
314+ WaitFlag<HardEvent::M_MTE1>(L0A_EVENT0);
315+ WaitFlag<HardEvent::M_MTE1>(L0A_EVENT1);
316+ WaitFlag<HardEvent::M_MTE1>(L0B_EVENT0);
317+ WaitFlag<HardEvent::M_MTE1>(L0B_EVENT1);
318+ WaitFlag<HardEvent::FIX_M>(L0C_EVENT0);
319+ WaitFlag<HardEvent::FIX_M>(L0C_EVENT1);
320+}
321+ 
322+template <typename SFAAT>
323+__aicore__ inline void SFAAMatmulServiceGqaMsd<SFAAT>::CopyGmToL1(LocalTensor<KV_T> &l1Tensor,
324+ GlobalTensor<KV_T> &gmSrcTensor, uint32_t srcN,
325+ uint32_t srcD, uint32_t srcDstride)
326+{
327+ Nd2NzParams nd2nzPara;
328+ nd2nzPara.ndNum = 1;
329+ nd2nzPara.nValue = srcN; // 行数
330+ nd2nzPara.dValue = srcD;
331+ nd2nzPara.srcDValue = srcDstride;
332+ nd2nzPara.dstNzC0Stride = SFAAAlign(srcN, 16U); // 对齐到16 单位block
333+ nd2nzPara.dstNzNStride = 1;
334+ nd2nzPara.srcNdMatrixStride = 0;
335+ nd2nzPara.dstNzMatrixStride = 0;
336+ DataCopy(l1Tensor, gmSrcTensor, nd2nzPara);
337+}
338+ 
339+template <typename SFAAT>
340+__aicore__ inline void SFAAMatmulServiceGqaMsd<SFAAT>::CopyInMm1AToL1(LocalTensor<KV_T> &l1Tensor, const RunInfo &info,
341+ uint32_t mSizeAct, uint32_t headSize)
342+{
343+ auto srcGm = queryPreProcessResGm_[(info.bN2Idx % constInfo.preLoadNum) * constInfo.bmm2ResUbSize * 2];
344+ // 不切G, mSizeAct必然是constInfo.gSize的整数倍
345+ uint32_t s1Size = mSizeAct / constInfo.gSize;
346+ Nd2NzParams nd2nzPara;
347+ nd2nzPara.ndNum = s1Size;
348+ nd2nzPara.nValue = constInfo.gSize; // 行数
349+ nd2nzPara.dValue = headSize;
350+ nd2nzPara.srcDValue = headSize;
351+ nd2nzPara.dstNzC0Stride = (mSizeAct + 15) / 16 * 16; // 对齐到16 单位block
352+ nd2nzPara.dstNzNStride = 1;
353+ nd2nzPara.srcNdMatrixStride = constInfo.qHeadNum * constInfo.headDim; // 这里小于65536
354+ nd2nzPara.dstNzMatrixStride = constInfo.gSize * BLOCK_ELEMENT_NUM;
355+ DataCopy(l1Tensor, srcGm, nd2nzPara);
356+}
357+ 
358+template <typename SFAAT>
359+__aicore__ inline void SFAAMatmulServiceGqaMsd<SFAAT>::LoadDataMm1A(LocalTensor<KV_T> &aL0Tensor,
360+ LocalTensor<KV_T> &aL1Tensor, uint32_t idx,
361+ uint32_t kSplitSize, uint32_t mSize, uint32_t kSize)
362+{
363+ LocalTensor<KV_T> srcTensor = aL1Tensor[mSize * kSplitSize * idx];
364+ LoadData3DParamsV2<KV_T> loadData3DParams;
365+ // SetFmatrixParams
366+ loadData3DParams.l1H = mSize / 16; // Hin=M1=8
367+ loadData3DParams.l1W = 16; // Win=M0
368+ loadData3DParams.padList[0] = 0;
369+ loadData3DParams.padList[1] = 0;
370+ loadData3DParams.padList[2] = 0;
371+ loadData3DParams.padList[3] = 255; // 尾部数据不影响滑窗的结果
372+ 
373+ // SetLoadToA0Params
374+ loadData3DParams.mExtension = mSize; // M
375+ loadData3DParams.kExtension = kSize; // K
376+ loadData3DParams.mStartPt = 0;
377+ loadData3DParams.kStartPt = 0;
378+ loadData3DParams.strideW = 1;
379+ loadData3DParams.strideH = 1;
380+ loadData3DParams.filterW = 1;
381+ loadData3DParams.filterSizeW = (1 >> 8) & 255;
382+ loadData3DParams.filterH = 1;
383+ loadData3DParams.filterSizeH = (1 >> 8) & 255;
384+ loadData3DParams.dilationFilterW = 1;
385+ loadData3DParams.dilationFilterH = 1;
386+ loadData3DParams.enTranspose = 0;
387+ loadData3DParams.fMatrixCtrl = 0;
388+ loadData3DParams.channelSize = kSize; // Cin=K
389+ LoadData<KV_T, LOAD3DV2_CONFIG>(aL0Tensor, srcTensor, loadData3DParams);
390+}
391+ 
392+template <typename SFAAT>
393+__aicore__ inline void SFAAMatmulServiceGqaMsd<SFAAT>::LoadDataMm1B(LocalTensor<KV_T> &l0Tensor,
394+ LocalTensor<KV_T> &l1Tensor, uint32_t idx,
395+ uint32_t kSplitSize, uint32_t kSize, uint32_t nSize)
396+{
397+ // N 方向全载
398+ LocalTensor<KV_T> srcTensor = l1Tensor[nSize * kSplitSize * idx];
399+ 
400+ LoadData2DParams loadData2DParams;
401+ loadData2DParams.startIndex = 0;
402+ loadData2DParams.repeatTimes = (nSize + 15) / 16 * kSize / (32 / sizeof(KV_T));
403+ loadData2DParams.srcStride = 1;
404+ loadData2DParams.dstGap = 0;
405+ loadData2DParams.ifTranspose = false;
406+ LoadData(l0Tensor, srcTensor, loadData2DParams);
407+}
408+ 
409+template <typename SFAAT>
410+__aicore__ inline void SFAAMatmulServiceGqaMsd<SFAAT>::CopyInMm2AToL1(LocalTensor<KV_T> &aL1Tensor, const RunInfo &info,
411+ uint32_t mSeqIdx, uint32_t subMSizeAct,
412+ uint32_t nSize, uint32_t nOffset)
413+{
414+ auto srcGm = vec1ResGm[(info.loop % constInfo.preLoadNum) * constInfo.mmResUbSize +
415+ mSeqIdx * info.actualSingleProcessSInnerSizeAlign + nOffset];
416+ CopyGmToL1(aL1Tensor, srcGm, subMSizeAct, nSize, info.actualSingleProcessSInnerSizeAlign);
417+}
418+ 
419+template <typename SFAAT>
420+__aicore__ inline void SFAAMatmulServiceGqaMsd<SFAAT>::CopyInMmBToL1(LocalTensor<KV_T> &bl1Tensor,
421+ GlobalTensor<KV_T> &gmSrcTensor,
422+ uint32_t subNSizeAct, uint32_t nOffset,
423+ const RunInfo &info)
424+{
425+ auto srcGm = gmSrcTensor[(info.loop % 4) * SALSC_S2BASEIZE * constInfo.headDim + nOffset];
426+ 
427+ CopyGmToL1(bl1Tensor, srcGm, subNSizeAct, constInfo.headDim, constInfo.headDim);
428+}
429+ 
430+template <typename SFAAT>
431+__aicore__ inline void SFAAMatmulServiceGqaMsd<SFAAT>::CopyInMm2AToL1(LocalTensor<KV_T>& aL1Tensor,
432+ const RunInfo &info, uint32_t mCopyIdx, uint32_t mCopyRowCount, uint32_t mActCopyRowCount,
433+ uint32_t kCopyIdx, uint32_t kCopyRowCount, uint32_t kActCopyRowCount)
434+{
435+ uint32_t mmRowCount = mActCopyRowCount;
436+ uint32_t copyStrideL1 = 16 * kActCopyRowCount;
437+ uint32_t copyStrideGm = 16 * info.actualSingleProcessSInnerSizeAlign;
438+ uint32_t copyIterNum = (mmRowCount + 15) / 16;
439+ for(int i = 0; i < copyIterNum; i++){
440+ Nd2NzParams mm1Nd2NzParamsForA;
441+ mm1Nd2NzParamsForA.ndNum = 1; // ND矩阵的个数
442+ if(i == copyIterNum - 1) {
443+ mm1Nd2NzParamsForA.nValue = mmRowCount - i * 16;
444+ }
445+ else {
446+ mm1Nd2NzParamsForA.nValue = 16; // 单个ND矩阵的行数, 单位为元素个数 16
447+ }
448+ mm1Nd2NzParamsForA.dValue = kActCopyRowCount; // 单个ND矩阵的列数, 单位为元素个数
449+ mm1Nd2NzParamsForA.srcNdMatrixStride = 0; // 相邻ND矩阵起始地址之间的偏移, 单位为元素个数
450+ mm1Nd2NzParamsForA.srcDValue = info.actualSingleProcessSInnerSizeAlign; // 同一个ND矩阵中相邻行起始地址之间的偏移, 单位为元素个数
451+ mm1Nd2NzParamsForA.dstNzC0Stride = 16; // 转换为NZ矩阵后，相邻Block起始地址之间的偏移, 单位为Block个数
452+ mm1Nd2NzParamsForA.dstNzNStride = 1; // 转换为NZ矩阵后，ND中之前相邻两行在NZ矩阵中起始地址之间的偏移
453+ mm1Nd2NzParamsForA.dstNzMatrixStride = 0; // 两个NZ矩阵，起始地址之间的偏移
454+ DataCopy(aL1Tensor[i * copyStrideL1],
455+ vec1ResGm[(info.loop % constInfo.preLoadNum) * constInfo.mmResUbSize * 2 +
456+ (mCopyIdx * mCopyRowCount) * info.actualSingleProcessSInnerSizeAlign +
457+ kCopyIdx * kCopyRowCount + i * copyStrideGm],
458+ mm1Nd2NzParamsForA);
459+ }
460+}
461+ 
462+// nCopyRowCount需要32元素对齐
463+// 先不化简吧，最后我们统一化简降低scaler
464+template <typename SFAAT>
465+__aicore__ inline void SFAAMatmulServiceGqaMsd<SFAAT>::CopyInMm1BToL1CubeDisepr(
466+ LocalTensor<KV_T>& bL1Tensor, const RunInfo &info, uint32_t nCopyIdx, uint32_t nCopyRowCount,
467+ uint32_t nActCopyRowCount, uint32_t nActCopyRowCountAlign)
468+{
469+ if (nActCopyRowCount == 0) {
470+ return;
471+ }
472+ uint32_t copySparseTotalBlockCount = (nActCopyRowCount + constInfo.sparseBlockSize - 1) / constInfo.sparseBlockSize;
473+ uint32_t copySparseBlockSizeTail = nActCopyRowCount - ((copySparseTotalBlockCount - 1) * constInfo.sparseBlockSize);
474+ uint64_t blockTableBaseOffset = info.bIdx * maxBlockNumPerBatch;
475+ uint32_t copyFinishRowCnt = 0;
476+ for (uint32_t sparseBlockCountIdx = 0; sparseBlockCountIdx < copySparseTotalBlockCount; sparseBlockCountIdx++) {
477+ // 计算当前block count需要搬入的block size
478+ int32_t copySparseBlockSize = constInfo.sparseBlockSize;
479+ if (sparseBlockCountIdx + 1 == copySparseTotalBlockCount) { //尾块
480+ copySparseBlockSize = copySparseBlockSizeTail;
481+ }
482+ uint32_t curLogicSeqIdx = info.s2BatchOffset + nCopyIdx * nCopyRowCount + sparseBlockCountIdx * constInfo.sparseBlockSize;
483+ uint32_t topkGmIdx = curLogicSeqIdx / constInfo.sparseBlockSize;
484+ uint32_t realS2Idx = topKGm.GetValue(info.topkGmBaseOffset + topkGmIdx) * static_cast<int64_t>(constInfo.sparseBlockSize) +
485+ (curLogicSeqIdx % constInfo.sparseBlockSize);
486+ PipeBarrier<PIPE_V>();
487+ uint64_t blockIdOffset = realS2Idx / kvCacheBlockSize;
488+ uint64_t reaminRowCnt = realS2Idx % kvCacheBlockSize;
489+ uint64_t idInBlockTable = blockTableGm.GetValue(blockTableBaseOffset + blockIdOffset); // 从block table上的获取编号
490+ uint32_t copyRowCnt = ((kvCacheBlockSize - reaminRowCnt) < (copySparseBlockSize)) ?
491+ (kvCacheBlockSize - reaminRowCnt) : copySparseBlockSize;
492+ if (copyFinishRowCnt + copyRowCnt > nActCopyRowCount) {
493+ copyRowCnt = nActCopyRowCount - copyFinishRowCnt;
494+ }
495+ uint64_t keyOffset = idInBlockTable * kvCacheBlockSize * constInfo.headDim * constInfo.kvHeadNum;
496+ keyOffset += (uint64_t)(info.n2Idx * constInfo.headDim * kvCacheBlockSize) + reaminRowCnt * L1_NZ_KV_BLOCK_ELEMENT;
497+ CopyInMm1BToL1ForPA(bL1Tensor, keyOffset, nActCopyRowCountAlign, copyFinishRowCnt, copyRowCnt, info);
498+ copyFinishRowCnt += copyRowCnt;
499+ curLogicSeqIdx += copyRowCnt;
500+ }
501+}
502+ 
503+template <typename SFAAT>
504+__aicore__ inline void SFAAMatmulServiceGqaMsd<SFAAT>::CopyInMm2BToL1CubeDisepr(
505+ LocalTensor<KV_T>& bL1Tensor, const RunInfo &info, uint32_t kCopyIdx, uint32_t kCopyRowCount,
506+ uint32_t kActCopyRowCount, uint32_t kActCopyRowCountAlign)
507+{
508+ if (info.aicS2AccessSize == 0) {
509+ return;
510+ }
511+ uint32_t copySparseTotalBlockCount = (kActCopyRowCount + constInfo.sparseBlockSize - 1) / constInfo.sparseBlockSize;
512+ uint32_t copySparseBlockSizeTail = kActCopyRowCount - ((copySparseTotalBlockCount - 1) * constInfo.sparseBlockSize);
513+ uint64_t blockTableBaseOffset = info.bIdx * maxBlockNumPerBatch;
514+ uint32_t copyFinishRowCnt = 0;
515+ for (uint32_t sparseBlockCountIdx = 0; sparseBlockCountIdx < copySparseTotalBlockCount; sparseBlockCountIdx++) {
516+ int32_t copySparseBlockSize = constInfo.sparseBlockSize;
517+ if (sparseBlockCountIdx + 1 == copySparseTotalBlockCount) { //尾块
518+ copySparseBlockSize = copySparseBlockSizeTail;
519+ }
520+ uint32_t curLogicSeqIdx = info.s2BatchOffset + kCopyIdx * kCopyRowCount + sparseBlockCountIdx * constInfo.sparseBlockSize;
521+ uint32_t topkGmIdx = curLogicSeqIdx / constInfo.sparseBlockSize;
522+ uint32_t realS2Idx = topKGm.GetValue(info.topkGmBaseOffset + topkGmIdx) * static_cast<int64_t>(constInfo.sparseBlockSize) +
523+ (curLogicSeqIdx % constInfo.sparseBlockSize);
524+ PipeBarrier<PIPE_V>();
525+ uint64_t blockIdOffset = realS2Idx / kvCacheBlockSize;
526+ uint64_t reaminRowCnt = realS2Idx % kvCacheBlockSize;
527+ uint64_t idInBlockTable = blockTableGm.GetValue(blockTableBaseOffset + blockIdOffset); // 从block table上的获取编号
528+ uint32_t copyRowCnt = ((kvCacheBlockSize - reaminRowCnt) < (copySparseBlockSize)) ?
529+ (kvCacheBlockSize - reaminRowCnt) : copySparseBlockSize;
530+ if (copyFinishRowCnt + copyRowCnt > kActCopyRowCount) {
531+ copyRowCnt = kActCopyRowCount - copyFinishRowCnt;
532+ }
533+ uint64_t valueOffset = idInBlockTable * kvCacheBlockSize * constInfo.headDim * constInfo.kvHeadNum;
534+ valueOffset += (uint64_t)(info.n2Idx * constInfo.headDim * kvCacheBlockSize) + reaminRowCnt * L1_NZ_KV_BLOCK_ELEMENT;
535+ CopyInMm2BToL1ForPA(bL1Tensor, valueOffset, copyFinishRowCnt, copyRowCnt);
536+ copyFinishRowCnt += copyRowCnt;
537+ curLogicSeqIdx += copyRowCnt;
538+ }
539+}
540+ 
541+ 
542+template <typename SFAAT>
543+__aicore__ inline void SFAAMatmulServiceGqaMsd<SFAAT>::CopyInMm2BToL1ForPA(
544+ LocalTensor<KV_T>& bL1Tensor, uint64_t valueGmBaseOffset, uint32_t copyStartRowCnt, uint32_t kActCopyRowCount)
545+{
546+ uint32_t blockK = 32 / sizeof(KV_T); // 单个ND矩阵的行数
547+ uint32_t blockElementCnt = 32 / sizeof(KV_T);
548+ 
549+ uint64_t step = blockElementCnt;
550+ 
551+ uint32_t nHead = 0;
552+ uint32_t nTail = 0;
553+ uint32_t midNdNum = 0;
554+ uint32_t copyStartRowCntInABlockK = copyStartRowCnt % blockK;
555+ uint32_t copyStartRowCntInBaseNTail = 0;
556+ if (copyStartRowCntInABlockK + kActCopyRowCount <= blockK) {
557+ nHead = 0;
558+ midNdNum = 0;
559+ nTail = kActCopyRowCount;
560+ copyStartRowCntInBaseNTail = copyStartRowCntInABlockK;
561+ } else {
562+ if (copyStartRowCntInABlockK == 0) {
563+ nHead = 0;
564+ } else {
565+ nHead = blockK - copyStartRowCntInABlockK;
566+ }
567+ midNdNum = (kActCopyRowCount - nHead) / blockK;
568+ nTail = (kActCopyRowCount - nHead) % blockK;
569+ copyStartRowCntInBaseNTail = 0;
570+ }
571+ 
572+ if (nHead != 0) {
573+ DataCopyParams intriParams;
574+ intriParams.blockLen = nHead;
575+ intriParams.blockCount = constInfo.headDim / blockElementCnt;
576+ intriParams.dstStride = blockK - nHead;
577+ intriParams.srcStride = kvCacheBlockSize - nHead;
578+ 
579+ uint32_t ndNumFinish = copyStartRowCnt / blockK;
580+ DataCopy(bL1Tensor[ndNumFinish * blockK * headDimAlign + copyStartRowCntInABlockK * (32 / sizeof(KV_T))],
581+ valueGm[valueGmBaseOffset], intriParams);
582+ }
583+ 
584+ if (midNdNum != 0) {
585+ DataCopyParams intriParams;
586+ intriParams.blockLen = blockK;
587+ intriParams.blockCount = constInfo.headDim / blockElementCnt;
588+ intriParams.dstStride = 0;
589+ intriParams.srcStride = kvCacheBlockSize - blockK;
590+ 
591+ int32_t ndNumFinish = (copyStartRowCnt + nHead) / blockK;
592+ for (uint32_t i = 0; i < midNdNum; i++) {
593+ DataCopy(bL1Tensor[(ndNumFinish + i) * blockK * headDimAlign],
594+ valueGm[valueGmBaseOffset + nHead * step + i * blockK * step], intriParams);
595+ }
596+ }
597+ 
598+ if (nTail != 0) {
599+ DataCopyParams intriParams;
600+ intriParams.blockLen = nTail;
601+ intriParams.blockCount = constInfo.headDim / blockElementCnt;
602+ intriParams.dstStride = blockK - nTail;
603+ intriParams.srcStride = kvCacheBlockSize - nTail;
604+ 
605+ int32_t ndNumFinish = (copyStartRowCnt + nHead) / blockK + midNdNum;
606+ DataCopy(bL1Tensor[ndNumFinish * blockK * headDimAlign + copyStartRowCntInBaseNTail * (32 / sizeof(KV_T))],
607+ valueGm[valueGmBaseOffset + (nHead + midNdNum * blockK) * step], intriParams);
608+ }
609+}
610+ 
611+template <typename SFAAT>
612+__aicore__ inline void SFAAMatmulServiceGqaMsd<SFAAT>::CopyInMm1BToL1ForPA(
613+ LocalTensor<KV_T>& bL1Tensor, uint64_t keyGmBaseOffset, uint32_t copyTotalRowCnt,
614+ uint32_t copyStartRowCnt, uint32_t nActCopyRowCount, const RunInfo &info)
615+{
616+ uint32_t baseN = 256 / sizeof(KV_T); // 256
617+ 
618+ uint32_t nHead = 0;
619+ uint32_t nTail = 0;
620+ uint32_t midNdNum = 0;
621+ uint32_t copyStartRowCntInBaseN = copyStartRowCnt % L1_NZ_BLOCK_SIZE_MM1;
622+ uint32_t copyStartRowCntInBaseNTail = 0;
623+ if (copyStartRowCntInBaseN + nActCopyRowCount <= L1_NZ_BLOCK_SIZE_MM1) {
624+ nHead = 0;
625+ midNdNum = 0;
626+ nTail = nActCopyRowCount;
627+ copyStartRowCntInBaseNTail = copyStartRowCntInBaseN;
628+ } else {
629+ if (copyStartRowCntInBaseN == 0) {
630+ nHead = 0;
631+ } else {
632+ nHead = L1_NZ_BLOCK_SIZE_MM1 - copyStartRowCntInBaseN;
633+ }
634+ midNdNum = (nActCopyRowCount - nHead) / L1_NZ_BLOCK_SIZE_MM1;
635+ nTail = (nActCopyRowCount - nHead) % L1_NZ_BLOCK_SIZE_MM1;
636+ copyStartRowCntInBaseNTail = 0;
637+ }
638+ 
639+ uint32_t dstNzC0StrideTail = L1_NZ_BLOCK_SIZE_MM1;
640+ if (copyTotalRowCnt % L1_NZ_BLOCK_SIZE_MM1 != 0) {
641+ if ((copyStartRowCnt + nActCopyRowCount) > (copyTotalRowCnt / L1_NZ_BLOCK_SIZE_MM1 * L1_NZ_BLOCK_SIZE_MM1)) {
642+ dstNzC0StrideTail = copyTotalRowCnt % L1_NZ_BLOCK_SIZE_MM1;
643+ }
644+ }
645+ Nd2NzParams mm1Nd2NzParamsForB;
646+ if (nHead != 0) {
647+ DataCopyParams intriParams;
648+ intriParams.blockLen = nHead;
649+ intriParams.blockCount = constInfo.headDim / L1_NZ_KV_BLOCK_ELEMENT;
650+ intriParams.dstStride = L1_NZ_BLOCK_SIZE_MM1 - nHead;
651+ intriParams.srcStride = kvCacheBlockSize - nHead;
652+ 
653+ uint32_t ndNumFinish = copyStartRowCnt / L1_NZ_BLOCK_SIZE_MM1;
654+ DataCopy(bL1Tensor[ndNumFinish * L1_NZ_BLOCK_SIZE_MM1 * headDimAlign + copyStartRowCntInBaseN * (32 / sizeof(KV_T))],
655+ keyGm[keyGmBaseOffset], intriParams);
656+ copyTotalRowCnt -= nHead;
657+ }
658+ 
659+ if (midNdNum != 0) {
660+ DataCopyParams intriParams;
661+ intriParams.blockLen = L1_NZ_BLOCK_SIZE_MM1;
662+ intriParams.blockCount = constInfo.headDim / L1_NZ_KV_BLOCK_ELEMENT;
663+ intriParams.dstStride = 0;
664+ intriParams.srcStride = kvCacheBlockSize - L1_NZ_BLOCK_SIZE_MM1;
665+ 
666+ int32_t ndNumFinish = (copyStartRowCnt + nHead) / L1_NZ_BLOCK_SIZE_MM1;
667+ for (uint32_t i = 0; i < midNdNum; i++) {
668+ DataCopy(bL1Tensor[(ndNumFinish + i) * L1_NZ_BLOCK_SIZE_MM1 * headDimAlign],
669+ keyGm[keyGmBaseOffset + nHead * L1_NZ_KV_BLOCK_ELEMENT + i * L1_NZ_BLOCK_SIZE_MM1 * L1_NZ_KV_BLOCK_ELEMENT], intriParams);
670+ }
671+ copyTotalRowCnt -= midNdNum * L1_NZ_BLOCK_SIZE_MM1;
672+ }
673+ 
674+ if (nTail != 0) {
675+ DataCopyParams intriParams;
676+ intriParams.blockLen = nTail; // 16
677+ intriParams.blockCount = constInfo.headDim / L1_NZ_KV_BLOCK_ELEMENT; // 4
678+ intriParams.dstStride = dstNzC0StrideTail - nTail;
679+ intriParams.srcStride = kvCacheBlockSize - nTail;
680+ 
681+ uint32_t ndNumFinish = (copyStartRowCnt + nHead) / L1_NZ_BLOCK_SIZE_MM1 + midNdNum;
682+ DataCopy(bL1Tensor[ndNumFinish * L1_NZ_BLOCK_SIZE_MM1 * headDimAlign + copyStartRowCntInBaseNTail * (L1_NZ_KV_BLOCK_ELEMENT / sizeof(KV_T))],
683+ keyGm[keyGmBaseOffset + (nHead + midNdNum * L1_NZ_BLOCK_SIZE_MM1) * L1_NZ_KV_BLOCK_ELEMENT], intriParams);
684+ }
685+}
686+ 
687+ 
688+// nCopyRowCount需要32元素对齐
689+template <typename SFAAT>
690+__aicore__ inline void SFAAMatmulServiceGqaMsd<SFAAT>::CopyInMm1BToL1VecMerge(
691+ LocalTensor<KV_T>& bL1Tensor, const RunInfo &info, uint32_t aicNZBlockBum, uint32_t aicNzBlockTail,
692+ uint32_t VecAccessSize, uint32_t alreadyVecAccesssSize, uint32_t nActCopyRowCountAlign, uint32_t CubeAccessSize)
693+{
694+ if (VecAccessSize == 0) {
695+ return;
696+ }
697+ uint64_t keyMergeGmBaseOffset = (info.loop % 4) * SALSC_S2MERGESIZE * constInfo.headDim + alreadyVecAccesssSize * L1_NZ_KV_BLOCK_ELEMENT;
698+ int32_t currentAivSize = 0;
699+ if constexpr (KV_LAYOUT_T == SFAA_LAYOUT::PA_NZ) {
700+ // head mid tail
701+ int32_t aivNzBlockHeadSize = aicNzBlockTail == 0 ? 0 : L1_NZ_BLOCK_SIZE_MM1 - aicNzBlockTail;
702+ int32_t remaining;
703+ int32_t aivNzBlockNum;
704+ int32_t aivNzBlockTail;
705+ int32_t aivNzBlockTailAlign;
706+ int32_t HeadBlockSize = L1_NZ_BLOCK_SIZE_MM1;
707+ if (aivNzBlockHeadSize > VecAccessSize) {
708+ aivNzBlockHeadSize = nActCopyRowCountAlign - CubeAccessSize;
709+ aivNzBlockNum = 0;
710+ aivNzBlockTail = 0;
711+ aivNzBlockTailAlign = 0;
712+ HeadBlockSize = nActCopyRowCountAlign - aicNZBlockBum * L1_NZ_BLOCK_SIZE_MM1;
713+ } else {
714+ remaining = VecAccessSize - aivNzBlockHeadSize;
715+ aivNzBlockNum = remaining >> L1_NZ_BLOCK_SIZE_MM1_SHIFT;
716+ aivNzBlockTail = remaining & L1_NZ_BLOCK_SIZE_MM1_MASK;
717+ aivNzBlockTailAlign = nActCopyRowCountAlign - CubeAccessSize - aivNzBlockHeadSize - (aivNzBlockNum) * L1_NZ_BLOCK_SIZE_MM1;
718+ }
719+ DataCopyParams mm1CopyParamsForB;
720+ mm1CopyParamsForB.blockCount = constInfo.headDim / L1_NZ_KV_BLOCK_ELEMENT;
721+ if (aivNzBlockHeadSize != 0) {
722+ mm1CopyParamsForB.blockLen = aivNzBlockHeadSize;
723+ mm1CopyParamsForB.dstStride = HeadBlockSize - aivNzBlockHeadSize;
724+ mm1CopyParamsForB.srcStride = SALSC_S2MERGESIZE - aivNzBlockHeadSize;
725+ DataCopy(bL1Tensor[aicNZBlockBum * L1_NZ_BLOCK_SIZE_MM1 * 128 + aicNzBlockTail * L1_NZ_KV_BLOCK_ELEMENT],
726+ keyMergeGm_[keyMergeGmBaseOffset], mm1CopyParamsForB);
727+ aicNZBlockBum = aicNZBlockBum + 1;
728+ currentAivSize = currentAivSize + aivNzBlockHeadSize;
729+ }
730+ if (aivNzBlockNum != 0) {
731+ for (uint32_t i = 0; i < aivNzBlockNum; ++i) {
732+ mm1CopyParamsForB.blockLen = L1_NZ_BLOCK_SIZE_MM1;
733+ mm1CopyParamsForB.dstStride = 0;
734+ mm1CopyParamsForB.srcStride = SALSC_S2MERGESIZE - L1_NZ_BLOCK_SIZE_MM1;
735+ DataCopy(bL1Tensor[aicNZBlockBum * L1_NZ_BLOCK_SIZE_MM1 * 128],
736+ keyMergeGm_[keyMergeGmBaseOffset + currentAivSize * L1_NZ_KV_BLOCK_ELEMENT], mm1CopyParamsForB);
737+ aicNZBlockBum = aicNZBlockBum + 1;
738+ currentAivSize = currentAivSize + L1_NZ_BLOCK_SIZE_MM1;
739+ }
740+ }
741+ if(aivNzBlockTail != 0) {
742+ mm1CopyParamsForB.blockLen = SfaaCeilDiv(aivNzBlockTailAlign * L1_NZ_KV_BLOCK_ELEMENT, DATABLOCK_BYTES);
743+ mm1CopyParamsForB.srcStride = SALSC_S2MERGESIZE - aivNzBlockTailAlign;
744+ mm1CopyParamsForB.dstStride = 0;
745+ DataCopy(bL1Tensor[aicNZBlockBum * L1_NZ_BLOCK_SIZE_MM1 * constInfo.headDim],
746+ keyMergeGm_[keyMergeGmBaseOffset + currentAivSize * L1_NZ_KV_BLOCK_ELEMENT], mm1CopyParamsForB);
747+ }
748+ }
749+}
750+ 
751+template <typename SFAAT>
752+__aicore__ inline void SFAAMatmulServiceGqaMsd<SFAAT>::CopyInMm2BToL1VecMerge(
753+ LocalTensor<KV_T>& bL1Tensor, const RunInfo &info, uint32_t aicNZBlockBum,
754+ uint32_t aicNzBlockTail, uint32_t VecAccessSize, uint32_t alreadyVecAccessSize,
755+ uint32_t kActCopyRowCountAlign, uint32_t CubeAccessSize)
756+{
757+ if (VecAccessSize == 0) {
758+ return;
759+ }
760+ uint64_t valueMergeGmBaseOffset = (info.loop % 4) * SALSC_S2MERGESIZE * constInfo.headDim + alreadyVecAccessSize * L1_NZ_KV_BLOCK_ELEMENT;
761+ if constexpr (KV_LAYOUT_T == SFAA_LAYOUT::PA_NZ) {
762+ // head mid tail
763+ int32_t aivNzBlockHeadSize = aicNzBlockTail == 0 ? 0 : L1_NZ_BLOCK_SIZE_MM2 - aicNzBlockTail;
764+ int32_t currentAivSize = 0;
765+ int32_t remaining;
766+ int32_t aivNzBlockNum;
767+ int32_t aivNzBlockTail;
768+ int32_t aivNzBlockTailAlign;
769+ int32_t HeadBlockSize = L1_NZ_BLOCK_SIZE_MM2;
770+ if (aivNzBlockHeadSize > VecAccessSize) {
771+ aivNzBlockHeadSize = kActCopyRowCountAlign - CubeAccessSize;
772+ aivNzBlockNum = 0;
773+ aivNzBlockTail = 0;
774+ aivNzBlockTailAlign = 0;
775+ HeadBlockSize = kActCopyRowCountAlign - aicNZBlockBum * L1_NZ_BLOCK_SIZE_MM2;
776+ } else {
777+ remaining = VecAccessSize - aivNzBlockHeadSize;
778+ aivNzBlockNum = remaining >> L1_NZ_BLOCK_SIZE_MM2_SHIFT;
779+ aivNzBlockTail = remaining & L1_NZ_BLOCK_SIZE_MM2_MASK;
780+ aivNzBlockTailAlign = kActCopyRowCountAlign - CubeAccessSize - aivNzBlockHeadSize - (aivNzBlockNum) * L1_NZ_BLOCK_SIZE_MM2;
781+ }
782+ DataCopyParams mm1CopyParamsForB;
783+ mm1CopyParamsForB.blockCount = constInfo.headDim / L1_NZ_KV_BLOCK_ELEMENT;
784+ if (aivNzBlockHeadSize != 0) {
785+ mm1CopyParamsForB.blockLen = aivNzBlockHeadSize;
786+ mm1CopyParamsForB.dstStride = HeadBlockSize - aivNzBlockHeadSize;
787+ mm1CopyParamsForB.srcStride = SALSC_S2MERGESIZE - aivNzBlockHeadSize;
788+ DataCopy(bL1Tensor[aicNZBlockBum * L1_NZ_BLOCK_SIZE_MM2 * 128 + aicNzBlockTail * L1_NZ_KV_BLOCK_ELEMENT],
789+ valueMergeGm_[valueMergeGmBaseOffset], mm1CopyParamsForB);
790+ aicNZBlockBum = aicNZBlockBum + 1;
791+ currentAivSize = currentAivSize + aivNzBlockHeadSize;
792+ }
793+ if (aivNzBlockNum != 0) {
794+ for (uint32_t i = 0; i < aivNzBlockNum; ++i) {
795+ mm1CopyParamsForB.blockLen = L1_NZ_BLOCK_SIZE_MM2;
796+ mm1CopyParamsForB.dstStride = 0;
797+ mm1CopyParamsForB.srcStride = SALSC_S2MERGESIZE - L1_NZ_BLOCK_SIZE_MM2;
798+ DataCopy(bL1Tensor[aicNZBlockBum * L1_NZ_BLOCK_SIZE_MM2 * 128],
799+ valueMergeGm_[valueMergeGmBaseOffset + currentAivSize * L1_NZ_KV_BLOCK_ELEMENT], mm1CopyParamsForB);
800+ aicNZBlockBum = aicNZBlockBum + 1;
801+ currentAivSize = currentAivSize + L1_NZ_BLOCK_SIZE_MM2;
802+ }
803+ }
804+ if(aivNzBlockTail != 0) {
805+ mm1CopyParamsForB.blockLen = SfaaCeilDiv(aivNzBlockTailAlign * L1_NZ_KV_BLOCK_ELEMENT, DATABLOCK_BYTES);
806+ mm1CopyParamsForB.srcStride = SALSC_S2MERGESIZE - aivNzBlockTailAlign;
807+ mm1CopyParamsForB.dstStride = 0;
808+ DataCopy(bL1Tensor[aicNZBlockBum * L1_NZ_BLOCK_SIZE_MM2 * constInfo.headDim],
809+ valueMergeGm_[valueMergeGmBaseOffset + currentAivSize * 32], mm1CopyParamsForB);
810+ }
811+ }
812+}
813+ 
814+template <typename SFAAT>
815+__aicore__ inline void SFAAMatmulServiceGqaMsd<SFAAT>::CopyInMm1BToL1(
816+ LocalTensor<KV_T>& bL1Tensor, const RunInfo &info, uint32_t nCopyIdx, uint32_t nCopyRowCount,
817+ uint32_t nActCopyRowCount, uint32_t nActCopyRowCountAlign)
818+{
819+ uint32_t baseN = 256 / sizeof(KV_T);
820+ uint64_t keyMergeGmBaseOffset = (info.loop % 4) * SALSC_S2BASEIZE * constInfo.headDim;
821+ uint32_t nTail = nActCopyRowCount % baseN;
822+ if constexpr (KV_LAYOUT_T == SFAA_LAYOUT::PA_NZ) {
823+ uint32_t blockElementCnt = 32 / sizeof(KV_T);
824+ uint32_t L1nLoopTimes = SfaaCeilDiv(nActCopyRowCount, baseN);
825+ uint32_t nTailAlign = nActCopyRowCountAlign - (L1nLoopTimes - 1) * baseN;
826+ 
827+ DataCopyParams mm1CopyParamsForB;
828+ mm1CopyParamsForB.blockCount = constInfo.headDim / blockElementCnt;
829+ mm1CopyParamsForB.dstStride = 0;
830+ for (uint32_t i = 0; i < L1nLoopTimes; ++i) {
831+ uint32_t n0RealSizeAlign = (i == L1nLoopTimes - 1) ? nTailAlign : baseN;
832+ mm1CopyParamsForB.blockLen = SfaaCeilDiv(n0RealSizeAlign * blockElementCnt * sizeof(KV_T), DATABLOCK_BYTES);
833+ mm1CopyParamsForB.srcStride = nActCopyRowCountAlign - n0RealSizeAlign;
834+ DataCopy(bL1Tensor[i * baseN * constInfo.headDim], keyMergeGm_[keyMergeGmBaseOffset + i * baseN * blockElementCnt],
835+ mm1CopyParamsForB);
836+ }
837+ } else {
838+ uint64_t step = constInfo.headDim;
839+ uint32_t ndNum_tmp = nActCopyRowCount / baseN;
840+ uint32_t nTailAlign = nActCopyRowCountAlign - ndNum_tmp * baseN;
841+ 
842+ Nd2NzParams mm1Nd2NzParamsForB;
843+ if (ndNum_tmp != 0) {
844+ mm1Nd2NzParamsForB.nValue = baseN; // 单个ND矩阵的行数, 单位为元素个数
845+ mm1Nd2NzParamsForB.dValue = constInfo.headDim; // 单个ND矩阵的列数, 单位为元素个数
846+ mm1Nd2NzParamsForB.srcDValue = step; // 同一个ND矩阵中相邻行起始地址之间的偏移, 单位为元素个数
847+ mm1Nd2NzParamsForB.dstNzC0Stride = baseN; // 转换为NZ矩阵后，相邻Block起始地址之间的偏移, 单位为Block个数
848+ mm1Nd2NzParamsForB.dstNzNStride = 1; // 转换为NZ矩阵后，ND中之前相邻两行在NZ矩阵中起始地址之间的偏移, 单位为Block个数
849+ mm1Nd2NzParamsForB.ndNum = ndNum_tmp;
850+ mm1Nd2NzParamsForB.srcNdMatrixStride = baseN * step; // 相邻ND矩阵起始地址之间的偏移, 单位为元素
851+ mm1Nd2NzParamsForB.dstNzMatrixStride = baseN * headDimAlign; // 两个NZ矩阵，起始地址之间的偏移, 单位为元素个数
852+ DataCopy(bL1Tensor, keyMergeGm_[keyMergeGmBaseOffset], mm1Nd2NzParamsForB);
853+ }
854+ 
855+ if (nTail != 0){
856+ mm1Nd2NzParamsForB.ndNum = 1;
857+ mm1Nd2NzParamsForB.nValue = nTail; // 单个ND矩阵的行数, 单位为元素个数
858+ mm1Nd2NzParamsForB.dValue = constInfo.headDim; // 单个ND矩阵的列数, 单位为元素个数
859+ mm1Nd2NzParamsForB.srcDValue = step; // 同一个ND矩阵中相邻行起始地址之间的偏移, 单位为元素个数
860+ mm1Nd2NzParamsForB.srcNdMatrixStride = 0; // 相邻ND矩阵起始地址之间的偏移, 单位为元素个数
861+ mm1Nd2NzParamsForB.dstNzC0Stride = nTailAlign; // 转换为NZ矩阵后，相邻Block起始地址之间的偏移, 单位为Block个数
862+ mm1Nd2NzParamsForB.dstNzNStride = 1; // 转换为NZ矩阵后，ND中之前相邻两行在NZ矩阵中起始地址之间的偏移, 单位为Block个数
863+ mm1Nd2NzParamsForB.dstNzMatrixStride = 0; // 两个NZ矩阵，起始地址之间的偏移, 单位为元素个数
864+ DataCopy(bL1Tensor[ndNum_tmp * baseN * headDimAlign], keyMergeGm_[keyMergeGmBaseOffset + ndNum_tmp * baseN * step],
865+ mm1Nd2NzParamsForB); //需要调整偏移地址，bL1Tensor的偏移， keyGm的偏移
866+ }
867+ }
868+}
869+ 
870+template <typename SFAAT>
871+__aicore__ inline void SFAAMatmulServiceGqaMsd<SFAAT>::CopyInMm1AToL1(LocalTensor<KV_T>& aL1Tensor, const RunInfo &info)
872+{
873+ uint32_t mmRowCount = msdIterNum * info.mSize;
874+ uint32_t copyStride = 16 * headDimAlign;
875+ uint32_t copyIterNum = (mmRowCount + 15) / 16;
876+ for(int i = 0; i < copyIterNum; i++) {
877+ Nd2NzParams mm1Nd2NzParamsForA;
878+ mm1Nd2NzParamsForA.ndNum = 1; // ND矩阵的个数
879+ if(i == copyIterNum - 1) {
880+ mm1Nd2NzParamsForA.nValue = msdIterNum * info.mSize - i * 16;
881+ }
882+ else {
883+ mm1Nd2NzParamsForA.nValue = 16; // 单个ND矩阵的行数, 单位为元素个数
884+ }
885+ mm1Nd2NzParamsForA.dValue = headDimAlign; // 单个ND矩阵的列数, 单位为元素个数
886+ mm1Nd2NzParamsForA.srcDValue = headDimAlign; // 同一个ND矩阵中相邻行起始地址之间的偏移, 单位为元素个数
887+ mm1Nd2NzParamsForA.srcNdMatrixStride = 0; // 相邻ND矩阵起始地址之间的偏移, 单位为元素个数
888+ mm1Nd2NzParamsForA.dstNzC0Stride = 16; // 转换为NZ矩阵后，同一行相邻Block起始地址之间的偏移, 单位为Block个数
889+ mm1Nd2NzParamsForA.dstNzNStride = 1; // 转换为NZ矩阵后，ND中之前相邻两行在NZ矩阵中起始地址之间的偏移
890+ mm1Nd2NzParamsForA.dstNzMatrixStride = 0; // 两个NZ矩阵，起始地址之间的偏移
891+ DataCopy(aL1Tensor[i * copyStride], queryPreProcessResGm_[(info.bN2Idx % constInfo.preLoadNum) * constInfo.bmm2ResUbSize * 2 + i * copyStride], mm1Nd2NzParamsForA);
892+ }
893+}
894+ 
895+template <typename SFAAT>
896+__aicore__ inline void SFAAMatmulServiceGqaMsd<SFAAT>::ComputeMm1(const RunInfo &info, const MSplitInfo mSplitInfo)
897+{
898+ LocalTensor<KV_T> aL1Tensor = qL1Buffers[qL1BufIter * L1Q_BLOCK_SIZE / sizeof(KV_T)];
899+ if (info.isFirstSInnerLoop) {
900+ WaitFlag<HardEvent::MTE1_MTE2>(L1Q_EVENT0 + qL1BufIter);
901+ CopyInMm1AToL1(aL1Tensor, info);
902+ SetFlag<HardEvent::MTE2_MTE1>(L1Q_EVENT0 + qL1BufIter);
903+ WaitFlag<HardEvent::MTE2_MTE1>(L1Q_EVENT0 + qL1BufIter);
904+ }
905+ constexpr uint32_t nCopyRowCount = KV_LOAD_TO_L1_ROW_NUM;
906+ uint32_t nCopyTimes = (info.actualSingleProcessSInnerSize + nCopyRowCount - 1) / nCopyRowCount;
907+ uint32_t nTailCopyRowCount = info.actualSingleProcessSInnerSize - (nCopyTimes - 1) * nCopyRowCount;
908+ uint32_t nTailCopyRowCountAlign = info.actualSingleProcessSInnerSizeAlign - (nCopyTimes - 1) * nCopyRowCount;
909+ uint32_t alreadyVecAccesssSize = 0;
910+ uint32_t intraCoreStart;
911+ uint32_t intraCoreEnd;
912+ uint32_t CubeAccessSize;
913+ uint32_t VecAccessSize;
914+ for (uint32_t nL1 = 0, nActCopyRowCount = nCopyRowCount, nActCopyRowCountAlign = nCopyRowCount; nL1 < nCopyTimes; nL1++) {
915+ if (nL1 + 1 == nCopyTimes) { // 尾块处理
916+ nActCopyRowCount = nTailCopyRowCount;
917+ nActCopyRowCountAlign = nTailCopyRowCountAlign;
918+ }
919+ intraCoreStart = nL1 * KV_LOAD_TO_L1_ROW_NUM; // 本轮循环的起始行号
920+ intraCoreEnd = intraCoreStart + nActCopyRowCount; // 本轮循环的结束行号
921+ if (info.aicS2AccessSize <= intraCoreStart) {
922+ CubeAccessSize = 0;
923+ } else if (info.aicS2AccessSize >= intraCoreEnd) {
924+ CubeAccessSize = nActCopyRowCount;
925+ } else {
926+ CubeAccessSize = info.aicS2AccessSize - intraCoreStart;
927+ }
928+ VecAccessSize = nActCopyRowCount - CubeAccessSize;
929+ kpL1BufIter++;
930+ LocalTensor<KV_T> bL1Tensor = kpL1Buffers[(kpL1BufIter % 3) * L1KP_BLOCK_SIZE / sizeof(KV_T)];
931+ WaitFlag<HardEvent::MTE1_MTE2>(L1KP_EVENT0 + (kpL1BufIter % 3));
932+ 
933+ CopyInMm1BToL1CubeDisepr(bL1Tensor, info, nL1, nCopyRowCount, CubeAccessSize, nActCopyRowCountAlign);
934+ uint32_t aicNZBlockBum = CubeAccessSize / L1_NZ_BLOCK_SIZE_MM1;
935+ uint32_t aicNzBlockTail = CubeAccessSize % L1_NZ_BLOCK_SIZE_MM1;
936+ CopyInMm1BToL1VecMerge(bL1Tensor, info, aicNZBlockBum, aicNzBlockTail, VecAccessSize, alreadyVecAccesssSize,
937+ nActCopyRowCountAlign, CubeAccessSize);
938+ alreadyVecAccesssSize = alreadyVecAccesssSize + VecAccessSize;
939+ SetFlag<HardEvent::MTE2_MTE1>(L1KP_EVENT0 + (kpL1BufIter % 3));
940+ WaitFlag<HardEvent::MTE2_MTE1>(L1KP_EVENT0 + (kpL1BufIter % 3));
941+ 
942+ constexpr uint32_t baseM = 64 / sizeof(KV_T); // 64
943+ uint32_t mActCopyRowCount = msdIterNum * info.mSize; // 8
944+ uint32_t mLoopTimes = (mActCopyRowCount + baseM - 1) / baseM; // 1
945+ uint32_t mTail = mActCopyRowCount - (mLoopTimes - 1) * baseM; // 8
946+ bool isHead = false;
947+ bool isMid = false;
948+ bool isTail = false;
949+ bool isOdd = (mLoopTimes % 2) != 0 ? true : false;
950+ bool hasTail = (mTail == baseM) ? false : true;
951+ // A:256*128，M方向按照64循环
952+ for (uint32_t i = 0, actualBaseM = baseM; i < mLoopTimes; i++) {
953+ if (i + 1 == mLoopTimes) {
954+ actualBaseM = mTail;
955+ }
956+ isHead = (!isOdd && ((!hasTail && (i < (mLoopTimes / 2))) || (i < (mLoopTimes / 2 - 1)))) || (isOdd && (i < (mLoopTimes / 2)));
957+ isMid = (!isOdd && hasTail && (i == (mLoopTimes / 2 - 1))) || (isOdd && (i == (mLoopTimes / 2)));
958+ isTail = (!isOdd && (!hasTail || (i > (mLoopTimes / 2 - 1)))) || (isOdd && (i > (mLoopTimes / 2)));
959+ LocalTensor<KV_T> aL0Tensor = aL0TensorPingPong[(aL0BufIter % 2) * L0A_PP_SIZE / sizeof(KV_T)];
960+ WaitFlag<HardEvent::M_MTE1>(L0A_EVENT0 + (aL0BufIter % 2));
961+ LoadData2DParams loadData2DParamsForA;
962+ loadData2DParamsForA.startIndex = 0;
963+ loadData2DParamsForA.repeatTimes = (actualBaseM + 15) / 16 * headDimAlign / (32 / sizeof(KV_T));
964+ loadData2DParamsForA.srcStride = 1;
965+ loadData2DParamsForA.dstGap = 0;
966+ loadData2DParamsForA.ifTranspose = false;
967+ LoadData(aL0Tensor, aL1Tensor[i * baseM * headDimAlign], loadData2DParamsForA);
968+ SetFlag<HardEvent::MTE1_M>(L0A_EVENT0 + (aL0BufIter % 2));
969+ WaitFlag<HardEvent::MTE1_M>(L0A_EVENT0 + (aL0BufIter % 2));
970+ 
971+ constexpr uint32_t baseN = 256 / sizeof(KV_T); // 256
972+ uint32_t nLoopTimes = (nActCopyRowCountAlign + baseN - 1) / baseN;
973+ uint32_t nTail = nActCopyRowCountAlign - (nLoopTimes - 1) * baseN;
974+ for (uint32_t j = 0, actualBaseN = baseN; j < nLoopTimes; j++) {
975+ if (j + 1 == nLoopTimes) {
976+ actualBaseN = nTail;
977+ }
978+ LocalTensor<KV_T> bL0Tensor = bL0TensorPingPong[(bL0BufIter % 2) * L0B_PP_SIZE / sizeof(KV_T)];
979+ WaitFlag<HardEvent::M_MTE1>(L0B_EVENT0 + (bL0BufIter % 2));
980+ uint32_t blockElementCnt = 32 / sizeof(KV_T);
981+ LoadData2DParams loadData2DParamsForB;
982+ loadData2DParamsForB.startIndex = 0;
983+ loadData2DParamsForB.srcStride = 1;
984+ loadData2DParamsForB.dstGap = 0;
985+ loadData2DParamsForB.ifTranspose = false;
986+ loadData2DParamsForB.repeatTimes = (actualBaseN / 16) * (headDimAlign / blockElementCnt);
987+ LoadData(bL0Tensor, bL1Tensor[baseN * headDimAlign * j], loadData2DParamsForB);
988+ SetFlag<HardEvent::MTE1_M>(L0B_EVENT0 + (bL0BufIter % 2));
989+ WaitFlag<HardEvent::MTE1_M>(L0B_EVENT0 + (bL0BufIter % 2));
990+ 
991+ MmadParams mmadParams;
992+ mmadParams.m = actualBaseM;
993+ if (mmadParams.m == 1) { // m等于1会默认开GEMV模式，且不可关闭GEMV，所以规避当作矩阵计算
994+ mmadParams.m = 16;
995+ }
996+ mmadParams.n = actualBaseN; // 无效数据不参与计算
997+ mmadParams.k = 128;
998+ mmadParams.cmatrixInitVal = true;
999+ mmadParams.cmatrixSource = false;
1000+ 
1001+ LocalTensor<L0C_T> cL0Tensor = cL0TensorPingPong[(cL0BufIter % 2) * L0C_PP_SIZE / sizeof(L0C_T)];
1002+ WaitFlag<HardEvent::FIX_M>(L0C_EVENT0 + (cL0BufIter % 2));
1003+ 
1004+ Mmad(cL0Tensor, aL0Tensor, bL0Tensor, mmadParams);
1005+ PipeBarrier<PIPE_M>();
1006+ SetFlag<HardEvent::M_FIX>(L0C_EVENT0 + (cL0BufIter % 2));
1007+ WaitFlag<HardEvent::M_FIX>(L0C_EVENT0 + (cL0BufIter % 2));
1008+ if (mLoopTimes == 1) {
1009+ for (uint32_t mIter = 0; mIter < msdIterNum; mIter++) {
1010+ float tmp = quantScaleC1S1;
1011+ if (mIter == 1) {
1012+ tmp = quantScaleC1S2;
1013+ }
1014+ FixpipeParamsV220 fixParams;
1015+ fixParams.nSize = actualBaseN;
1016+ fixParams.mSize = actualBaseM / msdIterNum; // 有效数据不足16行，只需要输出部分行即可
1017+ fixParams.srcStride = ((actualBaseM + 15) / 16) * 16;
1018+ fixParams.dstStride = info.actualSingleProcessSInnerSizeAlign; // mm1ResGm两行之间的间隔
1019+ fixParams.ndNum = 1;
1020+ fixParams.quantPre = QuantMode_t::DEQF16;
1021+ fixParams.deqScalar = static_cast<uint64_t>(*reinterpret_cast<int32_t *>(&tmp));
1022+ if (mIter == 1) {
1023+ SetAtomicAdd<half>();
1024+ }
1025+ Fixpipe(mm1ResGm[(info.loop % constInfo.preLoadNum) * constInfo.mmResUbSize + i * baseM * info.actualSingleProcessSInnerSizeAlign + nL1 * nCopyRowCount + j * baseN],
1026+ cL0Tensor[info.mSize * mIter * 16], fixParams);
1027+ if (mIter == 1) {
1028+ SetAtomicNone();
1029+ }
1030+ PipeBarrier<PIPE_FIX>();
1031+ }
1032+ } else if (isHead) {
1033+ float tmp = quantScaleC1S1;
1034+ FixpipeParamsV220 fixParams;
1035+ fixParams.nSize = actualBaseN;
1036+ fixParams.mSize = actualBaseM; // 有效数据不足16行，只需要输出部分行即可
1037+ fixParams.srcStride = ((actualBaseM + 15) / 16) * 16;
1038+ fixParams.dstStride = info.actualSingleProcessSInnerSizeAlign; // mm1ResGm两行之间的间隔
1039+ fixParams.ndNum = 1;
1040+ fixParams.quantPre = QuantMode_t::DEQF16;
1041+ fixParams.deqScalar = static_cast<uint64_t>(*reinterpret_cast<int32_t *>(&tmp));
1042+ Fixpipe(mm1ResGm[(info.loop % constInfo.preLoadNum) * constInfo.mmResUbSize + i * baseM * info.actualSingleProcessSInnerSizeAlign + nL1 * nCopyRowCount + j * baseN],
1043+ cL0Tensor, fixParams);
1044+ PipeBarrier<PIPE_FIX>();
1045+ } else if (isTail) {
1046+ float tmp = quantScaleC1S2;
1047+ FixpipeParamsV220 fixParams;
1048+ fixParams.nSize = actualBaseN;
1049+ fixParams.mSize = actualBaseM; // 有效数据不足16行，只需要输出部分行即可
1050+ fixParams.srcStride = ((actualBaseM + 15) / 16) * 16;
1051+ fixParams.dstStride = info.actualSingleProcessSInnerSizeAlign; // mm1ResGm两行之间的间隔
1052+ fixParams.ndNum = 1;
1053+ fixParams.quantPre = QuantMode_t::DEQF16;
1054+ fixParams.deqScalar = static_cast<uint64_t>(*reinterpret_cast<int32_t *>(&tmp));
1055+ SetAtomicAdd<half>();
1056+ Fixpipe(mm1ResGm[(info.loop % constInfo.preLoadNum) * constInfo.mmResUbSize + (i * baseM - info.mSize) * info.actualSingleProcessSInnerSizeAlign + nL1 * nCopyRowCount + j * baseN],
1057+ cL0Tensor, fixParams);
1058+ SetAtomicNone();
1059+ PipeBarrier<PIPE_FIX>();
1060+ } else {
1061+ float tmp = quantScaleC1S1;
1062+ FixpipeParamsV220 fixParams;
1063+ fixParams.nSize = actualBaseN;
1064+ fixParams.mSize = info.mSize - i * baseM; // 有效数据不足16行，只需要输出部分行即可
1065+ fixParams.srcStride = ((actualBaseM + 15) / 16) * 16;
1066+ fixParams.dstStride = info.actualSingleProcessSInnerSizeAlign; // mm1ResGm两行之间的间隔
1067+ fixParams.ndNum = 1;
1068+ fixParams.quantPre = QuantMode_t::DEQF16;
1069+ fixParams.deqScalar = static_cast<uint64_t>(*reinterpret_cast<int32_t *>(&tmp));
1070+ Fixpipe(mm1ResGm[(info.loop % constInfo.preLoadNum) * constInfo.mmResUbSize + i * baseM * info.actualSingleProcessSInnerSizeAlign + nL1 * nCopyRowCount + j * baseN],
1071+ cL0Tensor, fixParams);
1072+ PipeBarrier<PIPE_FIX>();
1073+ tmp = quantScaleC1S2;
1074+ fixParams.mSize = (i + 1) * baseM - info.mSize;
1075+ fixParams.deqScalar = static_cast<uint64_t>(*reinterpret_cast<int32_t *>(&tmp));
1076+ SetAtomicAdd<half>();
1077+ Fixpipe(mm1ResGm[(info.loop % constInfo.preLoadNum) * constInfo.mmResUbSize + nL1 * nCopyRowCount + j * baseN], cL0Tensor[(info.mSize - i * baseM) * 16], fixParams);
1078+ SetAtomicNone();
1079+ PipeBarrier<PIPE_FIX>();
1080+ }
1081+ SetFlag<HardEvent::FIX_M>(L0C_EVENT0 + (cL0BufIter % 2));
1082+ cL0BufIter++;
1083+ SetFlag<HardEvent::M_MTE1>(L0B_EVENT0 + (bL0BufIter % 2));
1084+ bL0BufIter++;
1085+ }
1086+ SetFlag<HardEvent::M_MTE1>(L0A_EVENT0 + (aL0BufIter % 2));
1087+ aL0BufIter++;
1088+ }
1089+ SetFlag<HardEvent::MTE1_MTE2>(L1KP_EVENT0 + (kpL1BufIter % 3));
1090+ }
1091+ // will change batch
1092+ if ((info.s2Idx + 1) == info.curSInnerLoopTimes) {
1093+ SetFlag<HardEvent::MTE1_MTE2>(L1Q_EVENT0 + qL1BufIter);
1094+ }
1095+}
1096+ 
1097+template <typename SFAAT>
1098+__aicore__ inline void SFAAMatmulServiceGqaMsd<SFAAT>::CopyInMm2BToL1(LocalTensor<KV_T>& bL1Tensor, const RunInfo &info,
1099+ uint32_t kCopyIdx, uint32_t kActCopyRowCountAlign, uint32_t kActCopyRowCount)
1100+{
1101+ uint32_t blockK = 32 / sizeof(KV_T); // 单个ND矩阵的行数
1102+ uint64_t valueMergeGmBaseOffset = (info.loop % 4) * SALSC_S2BASEIZE * constInfo.headDim;
1103+ 
1104+ if constexpr (KV_LAYOUT_T == SFAA_LAYOUT::PA_NZ) {
1105+ uint32_t blockElementCnt = 32 / sizeof(KV_T);
1106+ DataCopyParams mm2CopyParamsForB;
1107+ uint32_t L1kLoopTimes = SfaaCeilDiv(kActCopyRowCount, blockK);
1108+ uint32_t kTailAlign = kActCopyRowCountAlign - (L1kLoopTimes - 1) * blockK;
1109+ 
1110+ mm2CopyParamsForB.blockCount = constInfo.headDim / blockElementCnt;
1111+ mm2CopyParamsForB.dstStride = 0;
1112+ for (uint32_t i = 0; i < L1kLoopTimes; ++i) {
1113+ uint32_t k0RealSizeAlign = (i == L1kLoopTimes - 1) ? kTailAlign : blockK;
1114+ mm2CopyParamsForB.blockLen = SfaaCeilDiv(k0RealSizeAlign * blockElementCnt * sizeof(KV_T), DATABLOCK_BYTES);
1115+ mm2CopyParamsForB.srcStride = kActCopyRowCountAlign - k0RealSizeAlign;
1116+ DataCopy(bL1Tensor[i * blockK * constInfo.headDim], valueMergeGm_[valueMergeGmBaseOffset + i * blockK * blockElementCnt],
1117+ mm2CopyParamsForB);
1118+ }
1119+ } else {
1120+ uint64_t step = constInfo.headDim;
1121+ uint32_t kTail = kActCopyRowCount % blockK;
1122+ uint32_t ndNum_tmp = kActCopyRowCount / blockK;
1123+ 
1124+ Nd2NzParams mm1Nd2NzParamsForB;
1125+ mm1Nd2NzParamsForB.dValue = constInfo.headDim;
1126+ mm1Nd2NzParamsForB.srcDValue = step; // 同一个ND矩阵中相邻行起始地址之间的偏移, 单位为元素个数
1127+ mm1Nd2NzParamsForB.dstNzC0Stride = blockK;
1128+ mm1Nd2NzParamsForB.dstNzNStride = 1;
1129+ if (ndNum_tmp != 0){
1130+ mm1Nd2NzParamsForB.ndNum = ndNum_tmp;
1131+ mm1Nd2NzParamsForB.nValue = blockK;
1132+ mm1Nd2NzParamsForB.srcNdMatrixStride = blockK * step;
1133+ mm1Nd2NzParamsForB.dstNzMatrixStride = blockK * headDimAlign;
1134+ DataCopy(bL1Tensor, valueMergeGm_[valueMergeGmBaseOffset], mm1Nd2NzParamsForB);
1135+ }
1136+ if (kTail != 0){
1137+ mm1Nd2NzParamsForB.ndNum = 1;
1138+ mm1Nd2NzParamsForB.nValue = kTail;
1139+ mm1Nd2NzParamsForB.srcNdMatrixStride = 0;
1140+ mm1Nd2NzParamsForB.dstNzMatrixStride = 0;
1141+ DataCopy(bL1Tensor[ndNum_tmp * blockK * headDimAlign], valueMergeGm_[valueMergeGmBaseOffset + ndNum_tmp * blockK * step], mm1Nd2NzParamsForB);
1142+ }
1143+ }
1144+}
1145+ 
1146+template <typename SFAAT>
1147+__aicore__ inline void SFAAMatmulServiceGqaMsd<SFAAT>::LoadDataMm2A(LocalTensor<KV_T> aL0Tensor, LocalTensor<KV_T> aL1Tensor, uint32_t kSize)
1148+{
1149+ LoadData2DParams loadData2DParams;
1150+ loadData2DParams.startIndex = 0;
1151+ loadData2DParams.repeatTimes = kSize / (32 / sizeof(KV_T));
1152+ loadData2DParams.srcStride = 1;
1153+ loadData2DParams.dstGap = 0;
1154+ loadData2DParams.ifTranspose = false;
1155+ LoadData(aL0Tensor, aL1Tensor, loadData2DParams);
1156+}
1157+ 
1158+template <typename SFAAT>
1159+__aicore__ inline void SFAAMatmulServiceGqaMsd<SFAAT>::ComputeMm2(const RunInfo &info, const MSplitInfo mSplitInfo)
1160+{
1161+ constexpr uint32_t mCopyRowCount = P_LOAD_TO_L1_ROW_NUM; // 128
1162+ uint32_t mActRowCount = msdIterNum * info.mSize; // 2
1163+ uint32_t mCopyTimes = (mActRowCount + mCopyRowCount - 1) / mCopyRowCount; // 1
1164+ uint32_t mTailCopyRowCount = mActRowCount - (mCopyTimes - 1) * mCopyRowCount; // 2
1165+ 
1166+ constexpr uint32_t kCopyRowCount = KV_LOAD_TO_L1_ROW_NUM; // 512
1167+ uint32_t kCopyTimes = (info.actualSingleProcessSInnerSize + kCopyRowCount - 1) / kCopyRowCount;
1168+ uint32_t kTailCopyRowCount = info.actualSingleProcessSInnerSize - (kCopyTimes - 1) * kCopyRowCount;
1169+ uint32_t kTailCopyRowCountAlign = info.actualSingleProcessSInnerSizeAlign - (kCopyTimes - 1) * kCopyRowCount;
1170+ 
1171+ for (uint32_t mCopyIdx = 0, mActCopyRowCount = mCopyRowCount; mCopyIdx < mCopyTimes; mCopyIdx++) {
1172+ if (mCopyIdx + 1 == mCopyTimes) {
1173+ mActCopyRowCount = mTailCopyRowCount;
1174+ }
1175+ LocalTensor<L0C_T> cL0Tensor = cL0TensorPingPong[(cL0BufIter % 2) * L0C_PP_SIZE / sizeof(L0C_T)];
1176+ WaitFlag<HardEvent::FIX_M>(L0C_EVENT0 + (cL0BufIter % 2));
1177+ uint32_t CubeAccessSize;
1178+ uint32_t VecAccessSize;
1179+ uint32_t alreadyVecAccesssSize = 0;
1180+ uint32_t intraCoreStart;
1181+ uint32_t intraCoreEnd;
1182+ for (uint32_t kCopyIdx = 0, kActCopyRowCount = kCopyRowCount, kActCopyRowCountAlign = kCopyRowCount; kCopyIdx < kCopyTimes; kCopyIdx++) {
1183+ if (kCopyIdx + 1 == kCopyTimes) {
1184+ kActCopyRowCount = kTailCopyRowCount;
1185+ kActCopyRowCountAlign = kTailCopyRowCountAlign;
1186+ }
1187+ kpL1BufIter++;
1188+ LocalTensor<KV_T> aL1Tensor = kpL1Buffers[(kpL1BufIter % 3) * L1KP_BLOCK_SIZE / sizeof(KV_T)];
1189+ WaitFlag<HardEvent::MTE1_MTE2>(L1KP_EVENT0 + (kpL1BufIter % 3));
1190+ CopyInMm2AToL1(aL1Tensor, info, mCopyIdx, mCopyRowCount, mActCopyRowCount, kCopyIdx, kCopyRowCount, kActCopyRowCountAlign);
1191+ SetFlag<HardEvent::MTE2_MTE1>(L1KP_EVENT0 + (kpL1BufIter % 3));
1192+ WaitFlag<HardEvent::MTE2_MTE1>(L1KP_EVENT0 + (kpL1BufIter % 3));
1193+ LocalTensor<KV_T> bL1Tensor = vL1Buffers[(vL1BufIter % 4) * L1V_BLOCK_SIZE / sizeof(KV_T)];
1194+ uint32_t kb = 0;
1195+ if (mCopyIdx == 0) {
1196+ intraCoreStart = kCopyIdx * KV_LOAD_TO_L1_ROW_NUM;
1197+ intraCoreEnd = intraCoreStart + kActCopyRowCount;
1198+ if (info.aicS2AccessSize <= intraCoreStart) {
1199+ CubeAccessSize = 0;
1200+ } else if (info.aicS2AccessSize >= intraCoreEnd) {
1201+ CubeAccessSize = kActCopyRowCount;
1202+ } else {
1203+ CubeAccessSize = info.aicS2AccessSize - intraCoreStart;
1204+ }
1205+ VecAccessSize = kActCopyRowCount - CubeAccessSize;
1206+ vL1BufIter++;
1207+ bL1Tensor = vL1Buffers[(vL1BufIter % 4) * L1V_BLOCK_SIZE / sizeof(KV_T)];
1208+ WaitFlag<HardEvent::MTE1_MTE2>(L1V_EVENT0 + (vL1BufIter % 4));
1209+ CopyInMm2BToL1CubeDisepr(bL1Tensor, info, kCopyIdx, kCopyRowCount, CubeAccessSize, kActCopyRowCountAlign);
1210+ uint32_t aicNZBlockBum = CubeAccessSize / L1_NZ_BLOCK_SIZE_MM2;
1211+ uint32_t aicNzBlockTail = CubeAccessSize % L1_NZ_BLOCK_SIZE_MM2;
1212+ CopyInMm2BToL1VecMerge(bL1Tensor, info, aicNZBlockBum, aicNzBlockTail, VecAccessSize, alreadyVecAccesssSize,
1213+ kActCopyRowCountAlign, CubeAccessSize);
1214+ alreadyVecAccesssSize = alreadyVecAccesssSize + VecAccessSize;
1215+ SetFlag<HardEvent::MTE2_MTE1>(L1V_EVENT0 + (vL1BufIter % 4));
1216+ WaitFlag<HardEvent::MTE2_MTE1>(L1V_EVENT0 + (vL1BufIter % 4));
1217+ kb = vL1BufIter;
1218+ } else {
1219+ kb = vL1BufIter - (kCopyTimes-kCopyIdx-1);
1220+ bL1Tensor = vL1Buffers[(kb % 4) * L1V_BLOCK_SIZE / sizeof(KV_T)];
1221+ }
1222+ constexpr uint32_t baseK = 128 / sizeof(KV_T);
1223+ uint32_t kLoopTimes = (kActCopyRowCountAlign + baseK - 1) / baseK;
1224+ uint32_t kTailAlign = kActCopyRowCountAlign - (kLoopTimes - 1) * baseK;
1225+ uint32_t kTail = kActCopyRowCount - (kLoopTimes - 1) * baseK;
1226+ for (uint32_t i = 0, actualBaseKAlign = baseK, actualBaseK = baseK; i < kLoopTimes; i++) {
1227+ if (i + 1 == kLoopTimes) {
1228+ actualBaseKAlign = kTailAlign;
1229+ actualBaseK = kTail;
1230+ }
1231+ 
1232+ LocalTensor<KV_T> aL0Tensor = aL0TensorPingPong[(aL0BufIter % 2) * L0A_PP_SIZE / sizeof(KV_T)];
1233+ WaitFlag<HardEvent::M_MTE1>(L0A_EVENT0 + (aL0BufIter % 2));
1234+ LocalTensor<KV_T> bL0Tensor = bL0TensorPingPong[(bL0BufIter % 2) * L0B_PP_SIZE / sizeof(KV_T)];
1235+ WaitFlag<HardEvent::M_MTE1>(L0B_EVENT0 + (bL0BufIter % 2));
1236+ 
1237+ LocalTensor<KV_T> curAL1Tensor = aL1Tensor[16 * baseK * i];
1238+ 
1239+ uint32_t mmRowCount = mActCopyRowCount;
1240+ uint32_t copyStrideL0 = 16 * actualBaseKAlign;
1241+ uint32_t copyStrideL1 = 16 * kActCopyRowCountAlign;
1242+ uint32_t copyIterNum = (mmRowCount + 15) / 16;
1243+ for(int i = 0; i < copyIterNum; i++){
1244+ LoadDataMm2A(aL0Tensor[i * copyStrideL0], curAL1Tensor[i * copyStrideL1], actualBaseKAlign);
1245+ }
1246+
1247+ SetFlag<HardEvent::MTE1_M>(L0A_EVENT0 + (aL0BufIter % 2));
1248+ WaitFlag<HardEvent::MTE1_M>(L0A_EVENT0 + (aL0BufIter % 2));
1249+ 
1250+ uint32_t blockElementCnt = 32 / sizeof(KV_T);
1251+ LoadData2dTransposeParams loadData2DTransposeParamsForB;
1252+ loadData2DTransposeParamsForB.startIndex = 0;
1253+ loadData2DTransposeParamsForB.srcStride = 1;
1254+ loadData2DTransposeParamsForB.dstFracGap = 0;
1255+ loadData2DTransposeParamsForB.repeatTimes = (actualBaseKAlign / blockElementCnt) * (headDimAlign / blockElementCnt); // 16
1256+ loadData2DTransposeParamsForB.dstGap = blockElementCnt / 16 - 1;
1257+ uint32_t l1BaseOffset = baseK * headDimAlign * i;
1258+ LoadDataWithTranspose(bL0Tensor, bL1Tensor[l1BaseOffset], loadData2DTransposeParamsForB);
1259+ SetFlag<HardEvent::MTE1_M>(L0B_EVENT0 + (bL0BufIter % 2));
1260+ WaitFlag<HardEvent::MTE1_M>(L0B_EVENT0 + (bL0BufIter % 2));
1261+ 
1262+ MmadParams mmadParams;
1263+ mmadParams.m = mActCopyRowCount;
1264+ if (mmadParams.m == 1) { // m等于1会默认开GEMV模式，且不可关闭GEMV，所以规避当作矩阵计算
1265+ mmadParams.m = 16;
1266+ }
1267+ mmadParams.n = 128;
1268+ mmadParams.k = actualBaseK; // 无效数据不参与计算
1269+ mmadParams.cmatrixInitVal = (kCopyIdx == 0) && (i == 0);
1270+ mmadParams.cmatrixSource = false;
1271+ Mmad(cL0Tensor, aL0Tensor, bL0Tensor, mmadParams);
1272+ PipeBarrier<PIPE_M>();
1273+ 
1274+ SetFlag<HardEvent::M_MTE1>(L0A_EVENT0 + (aL0BufIter % 2));
1275+ aL0BufIter++;
1276+ SetFlag<HardEvent::M_MTE1>(L0B_EVENT0 + (bL0BufIter % 2));
1277+ bL0BufIter++;
1278+ }
1279+ 
1280+ SetFlag<HardEvent::MTE1_MTE2>(L1KP_EVENT0 + (kpL1BufIter % 3));
1281+ if ((mCopyIdx + 1) == mCopyTimes) {
1282+ SetFlag<HardEvent::MTE1_MTE2>(L1V_EVENT0 + (kb % 4));
1283+ }
1284+ }
1285+ SetFlag<HardEvent::M_FIX>(L0C_EVENT0 + (cL0BufIter % 2));
1286+ WaitFlag<HardEvent::M_FIX>(L0C_EVENT0 + (cL0BufIter % 2));
1287+ if (mCopyTimes == 1) {
1288+ for (uint32_t mIter = 0; mIter < msdIterNum; mIter++) {
1289+ float tmp = quantScaleC2O1;
1290+ if (mIter == 1) {
1291+ tmp = quantScaleC2O2;
1292+ }
1293+ FixpipeParamsV220 fixParams;
1294+ fixParams.nSize = 128;
1295+ fixParams.mSize = mActCopyRowCount / msdIterNum; // 有效数据不足16行，只需要输出部分行即可
1296+ fixParams.srcStride = ((msdIterNum * fixParams.mSize + 15) / 16) * 16;
1297+ fixParams.dstStride = 128;
1298+ fixParams.ndNum = 1;
1299+ fixParams.quantPre = QuantMode_t::DEQF16;
1300+ fixParams.deqScalar = static_cast<uint64_t>(*reinterpret_cast<int32_t *>(&tmp));
1301+ if (mIter == 1) {
1302+ SetAtomicAdd<half>();
1303+ }
1304+ Fixpipe(mm2ResGm[(info.loop % constInfo.preLoadNum) * constInfo.bmm2ResUbSize], cL0Tensor[mIter * fixParams.mSize * 16], fixParams);
1305+ if (mIter == 1) {
1306+ SetAtomicNone();
1307+ }
1308+ PipeBarrier<PIPE_FIX>();
1309+ }
1310+ } else {
1311+ if (mTailCopyRowCount != mCopyRowCount && mCopyIdx == 0) {
1312+ float tmp = quantScaleC2O1;
1313+ FixpipeParamsV220 fixParams;
1314+ fixParams.nSize = 128;
1315+ fixParams.mSize = info.mSize; // 有效数据不足16行，只需要输出部分行即可
1316+ fixParams.srcStride = ((mActCopyRowCount + 15) / 16) * 16;
1317+ fixParams.dstStride = 128;
1318+ fixParams.ndNum = 1;
1319+ fixParams.quantPre = QuantMode_t::DEQF16;
1320+ fixParams.deqScalar = static_cast<uint64_t>(*reinterpret_cast<int32_t *>(&tmp));
1321+ Fixpipe(mm2ResGm[(info.loop % constInfo.preLoadNum) * constInfo.bmm2ResUbSize], cL0Tensor, fixParams);
1322+ PipeBarrier<PIPE_FIX>();
1323+ tmp = quantScaleC2O2;
1324+ fixParams.mSize = mActCopyRowCount - info.mSize;
1325+ fixParams.deqScalar = static_cast<uint64_t>(*reinterpret_cast<int32_t *>(&tmp));
1326+ SetAtomicAdd<half>();
1327+ Fixpipe(mm2ResGm[(info.loop % constInfo.preLoadNum) * constInfo.bmm2ResUbSize], cL0Tensor[info.mSize * 16], fixParams);
1328+ SetAtomicNone();
1329+ PipeBarrier<PIPE_FIX>();
1330+ } else if (mCopyIdx == 0) {
1331+ float tmp = quantScaleC2O1;
1332+ FixpipeParamsV220 fixParams;
1333+ fixParams.nSize = 128;
1334+ fixParams.mSize = mActCopyRowCount; // 有效数据不足16行，只需要输出部分行即可
1335+ fixParams.srcStride = ((mActCopyRowCount + 15) / 16) * 16;
1336+ fixParams.dstStride = 128;
1337+ fixParams.ndNum = 1;
1338+ fixParams.quantPre = QuantMode_t::DEQF16;
1339+ fixParams.deqScalar = static_cast<uint64_t>(*reinterpret_cast<int32_t *>(&tmp));
1340+ Fixpipe(mm2ResGm[(info.loop % constInfo.preLoadNum) * constInfo.bmm2ResUbSize], cL0Tensor, fixParams);
1341+ PipeBarrier<PIPE_FIX>();
1342+ } else {
1343+ float tmp = quantScaleC2O2;
1344+ FixpipeParamsV220 fixParams;
1345+ fixParams.nSize = 128;
1346+ fixParams.mSize = mActCopyRowCount; // 有效数据不足16行，只需要输出部分行即可
1347+ fixParams.srcStride = ((mActCopyRowCount + 15) / 16) * 16;
1348+ fixParams.dstStride = 128;
1349+ fixParams.ndNum = 1;
1350+ fixParams.quantPre = QuantMode_t::DEQF16;
1351+ fixParams.deqScalar = static_cast<uint64_t>(*reinterpret_cast<int32_t *>(&tmp));
1352+ SetAtomicAdd<half>();
1353+ Fixpipe(mm2ResGm[(info.loop % constInfo.preLoadNum) * constInfo.bmm2ResUbSize + (mCopyIdx * mCopyRowCount - info.mSize) * headDimAlign], cL0Tensor, fixParams);
1354+ SetAtomicNone();
1355+ PipeBarrier<PIPE_FIX>();
1356+ }
1357+ }
1358+ SetFlag<HardEvent::FIX_M>(L0C_EVENT0 + (cL0BufIter % 2));
1359+ cL0BufIter++;
1360+ }
1361+}
1362+ 
1363+#endif // SPARSE_FLASH_ATTENTION_ANTIQUANT_SERVICE_CUBE_GQA_MSD_H
Aexperimental/attention/sparse_flash_attention_antiquant/op_kernel/sparse_flash_attention_antiquant_service_flashdecode.h+471-0
@@ -0,0 +1,471 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file sparse_flash_attention_antiquant_service_flashdecode.h
13+ * \brief
14+ */
15+#ifndef FIA_BLOCK_VEC_FLASHDECODE_H
16+#define FIA_BLOCK_VEC_FLASHDECODE_H
17+ 
18+#include "kernel_operator.h"
19+#include "kernel_operator_list_tensor_intf.h"
20+#include "kernel_tiling/kernel_tiling.h"
21+#include "lib/matmul_intf.h"
22+#include "lib/matrix/matmul/tiling.h"
23+#include "sparse_flash_attention_antiquant_common.h"
24+ 
25+struct TaskInfo {
26+ uint32_t bIdx;
27+ uint32_t n2Idx;
28+ uint32_t gS1Idx;
29+ uint32_t actualCombineLoopSize;
30+ int64_t attenOutOffset;
31+};
32+ 
33+template <typename SFAAT>
34+class SFAAFlashDecodeServiceGqa {
35+public:
36+ // =================================类型定义区=================================
37+ // 中间计算数据类型为float，高精度模式
38+ using T = float;
39+ using OUT_T = typename SFAAT::outputType;
40+ static constexpr SFAA_LAYOUT LAYOUT_T = SFAAT::layout;
41+ 
42+ __aicore__ inline void InitGlobalTensor(GlobalTensor<T> lseMaxFdGm, GlobalTensor<T> lseSumFdGm, GlobalTensor<T> accumOutGm,
43+ GlobalTensor<OUT_T> attentionOutGm, GlobalTensor<int32_t> actualSeqLengthsGmQ, GlobalTensor<int32_t> actualSeqLengthsGm);
44+ __aicore__ inline void InitSoftmaxLseGm(GlobalTensor<float> softmaxLseGm);
45+ __aicore__ inline void InitParams(const ConstInfo &constInfo);
46+ __aicore__ inline void InitDecodeParams();
47+ __aicore__ inline void InitBuffers(TPipe *pipe);
48+ __aicore__ inline void AllocEventID();
49+ __aicore__ inline void FreeEventID();
50+ __aicore__ inline void FlashDecode(FDparams &fd);
51+ 
52+private:
53+// =================================常量区=================================
54+ static constexpr uint64_t SYNC_LSE_SUM_BUF1_FLAG = 6;
55+ static constexpr uint64_t SYNC_LSE_SUM_BUF2_FLAG = 7;
56+ static constexpr uint64_t SYNC_LSE_MAX_BUF1_FLAG = 8;
57+ static constexpr uint64_t SYNC_LSE_MAX_BUF2_FLAG = 9;
58+ static constexpr uint64_t SYNC_MM2RES_BUF1_FLAG = 10;
59+ static constexpr uint64_t SYNC_MM2RES_BUF2_FLAG = 11;
60+ static constexpr uint64_t SYNC_FDOUTPUT_BUF_FLAG = 6;
61+ static constexpr uint64_t SYNC_LSEOUTPUT_BUF_FLAG = 7;
62+ static constexpr uint64_t SYNC_SINK_BUF1_FLAG = 12;
63+ static constexpr uint64_t SYNC_SINK_BUF2_FLAG = 13;
64+ 
65+ static constexpr uint32_t BLOCK_ELEMENT_NUM = BYTE_BLOCK / sizeof(T); // 32/4=8
66+ 
67+protected:
68+ GlobalTensor<T> lseSumFdGm;
69+ GlobalTensor<T> lseMaxFdGm;
70+ GlobalTensor<T> accumOutGm;
71+ GlobalTensor<OUT_T> attentionOutGm;
72+ GlobalTensor<float> softmaxLseGm;
73+ GlobalTensor<int32_t> actualSeqLengthsGmQ;
74+ GlobalTensor<int32_t> actualSeqLengthsGm;
75+ // =======================获取实际Act_S，用于行无效处理===========================
76+ static constexpr bool PAGE_ATTENTION = SFAAT::pageAttention;
77+ uint64_t actSeqLensKv = 0;
78+ uint64_t actSeqLensQ = 0;
79+
80+ int64_t preTokensPerBatch = 0;
81+ int64_t nextTokensPerBatch = 0;
82+
83+ static constexpr T BOOL_ATTEN_MASK_SCALAR_VALUE = -1000000000000.0; // 用于mask为bool类型
84+ uint32_t negativeIntScalar = *((uint32_t *)&BOOL_ATTEN_MASK_SCALAR_VALUE);
85+ bool learnableSinkFlag = false;
86+ // ================================类成员变量====================================
87+ // aic、aiv核信息
88+ uint32_t blockIdx = 0U;
89+ ConstInfo constInfo{};
90+ TaskInfo taskInfo{};
91+ __aicore__ inline void CopyAccumOutIn(LocalTensor<T> &accumOutLocal, uint32_t splitKVIndex, uint32_t startRow,
92+ uint32_t dealRowCount);
93+ __aicore__ inline void CopyLseIn(uint32_t startRow, uint32_t dealRowCount, uint64_t baseOffset, uint32_t cntM);
94+ __aicore__ inline void ComputeScaleValue(LocalTensor<T> &lseExp, uint32_t startRow, uint32_t dealRowCount,
95+ uint32_t cntM);
96+ __aicore__ inline void Bmm2DataCopyOutTrans(LocalTensor<OUT_T> &attenOutUb, uint32_t startRow,
97+ uint32_t dealRowCount, uint32_t columnCount);
98+ __aicore__ inline void Bmm2DataCopyOut(uint64_t attenOutOffset, LocalTensor<OUT_T> &attenOutUb, uint32_t startRow,
99+ uint32_t dealRowCount, uint32_t columnCount, uint32_t actualColumnCount);
100+ __aicore__ inline void ReduceFinalRes(LocalTensor<T> &reduceOut, LocalTensor<T> &mm2Res, LocalTensor<T> &lseLocal,
101+ uint32_t cntKV, uint32_t dealRowCount);
102+ __aicore__ inline void CopyFinalResOut(LocalTensor<T> &accumOutLocal, uint32_t startRow, uint32_t dealRowCount,
103+ uint32_t cntM);
104+private:
105+ // ================================FD Local Buffer区====================================
106+ TBuf<> fdSumBuf1; // 1.5k: 16*24*4
107+ TBuf<> fdSumBuf2; // 1.5k: 16*24*4
108+ TBuf<> fdMaxBuf1; // 1.5k: 16*24*4
109+ TBuf<> fdMaxBuf2; // 1.5k: 16*24*4
110+ TBuf<> fdLseExpBuf; // 1.5k: 16*24*4
111+ TBuf<> fdMm2ResBuf1; // 32k: 16*512*4
112+ TBuf<> fdMm2ResBuf2; // 32k: 16*512*4
113+ TBuf<> fdReduceBuf; // 32k: 16*512*4
114+ TBuf<> fdOutputBuf; // 32k: 16*512*4
115+ TBuf<> fdSinkCopyInBuf; // 2*1k: 2*128*8
116+ TBuf<> fdSinkValueBuf; // 2k
117+ TBuf<> fdSinkExpBuf; // 256B
118+ TBuf<> fdSinkTmpBuf; // 2k
119+ 
120+ TBuf<> fdLseMaxUbBuf1; // 64B: 16*4
121+ TBuf<> fdLseMaxUbBuf2; // 64B: 16*4
122+ TBuf<> fdLseSumUbBuf1; // 64B: 16*4
123+ TBuf<> fdLseSumUbBuf2; // 64B: 16*4
124+ TBuf<> fdLseUbBuf; // 64B: 16*4
125+};
126+ 
127+template <typename SFAAT> __aicore__ inline
128+void SFAAFlashDecodeServiceGqa<SFAAT>::InitGlobalTensor(GlobalTensor<T> lseMaxFdGm,
129+ GlobalTensor<T> lseSumFdGm,
130+ GlobalTensor<T> accumOutGm,
131+ GlobalTensor<OUT_T> attentionOutGm,
132+ GlobalTensor<int32_t> actualSeqLengthsGmQ,
133+ GlobalTensor<int32_t> actualSeqLengthsGm)
134+{
135+ this->lseMaxFdGm = lseMaxFdGm;
136+ this->lseSumFdGm = lseSumFdGm;
137+ this->accumOutGm = accumOutGm;
138+ this->attentionOutGm = attentionOutGm;
139+ this->actualSeqLengthsGmQ = actualSeqLengthsGmQ;
140+ this->actualSeqLengthsGm = actualSeqLengthsGm;
141+}
142+ 
143+template <typename SFAAT> __aicore__ inline
144+void SFAAFlashDecodeServiceGqa<SFAAT>::InitSoftmaxLseGm(GlobalTensor<float> softmaxLseGm)
145+{
146+ this->softmaxLseGm = softmaxLseGm;
147+}
148+ 
149+template <typename SFAAT> __aicore__ inline
150+void SFAAFlashDecodeServiceGqa<SFAAT>::InitParams(const ConstInfo &constInfo)
151+{
152+ this->constInfo = constInfo;
153+}
154+ 
155+ 
156+template <typename SFAAT>__aicore__ inline
157+void SFAAFlashDecodeServiceGqa<SFAAT>::InitDecodeParams()
158+{
159+ this->blockIdx = GetBlockIdx();
160+}
161+ 
162+template <typename SFAAT> __aicore__ inline
163+void SFAAFlashDecodeServiceGqa<SFAAT>::InitBuffers(TPipe *pipe)
164+{
165+ if ASCEND_IS_AIV {
166+ pipe->Reset();
167+ pipe->InitBuffer(fdSumBuf1, ConstInfo::BUFFER_SIZE_BYTE_4K + ConstInfo::BUFFER_SIZE_BYTE_2K);
168+ pipe->InitBuffer(fdSumBuf2, ConstInfo::BUFFER_SIZE_BYTE_4K + ConstInfo::BUFFER_SIZE_BYTE_2K);
169+ pipe->InitBuffer(fdMaxBuf1, ConstInfo::BUFFER_SIZE_BYTE_4K + ConstInfo::BUFFER_SIZE_BYTE_2K);
170+ pipe->InitBuffer(fdMaxBuf2, ConstInfo::BUFFER_SIZE_BYTE_4K + ConstInfo::BUFFER_SIZE_BYTE_2K);
171+ pipe->InitBuffer(fdLseExpBuf, ConstInfo::BUFFER_SIZE_BYTE_4K + ConstInfo::BUFFER_SIZE_BYTE_2K);
172+ pipe->InitBuffer(fdMm2ResBuf1, ConstInfo::BUFFER_SIZE_BYTE_16K);
173+ pipe->InitBuffer(fdMm2ResBuf2, ConstInfo::BUFFER_SIZE_BYTE_16K);
174+ pipe->InitBuffer(fdReduceBuf, ConstInfo::BUFFER_SIZE_BYTE_16K);
175+ pipe->InitBuffer(fdOutputBuf, ConstInfo::BUFFER_SIZE_BYTE_16K);
176+ pipe->InitBuffer(fdLseMaxUbBuf1, ConstInfo::BUFFER_SIZE_BYTE_256B);
177+ pipe->InitBuffer(fdLseSumUbBuf1, ConstInfo::BUFFER_SIZE_BYTE_256B);
178+ pipe->InitBuffer(fdLseMaxUbBuf2, ConstInfo::BUFFER_SIZE_BYTE_256B);
179+ pipe->InitBuffer(fdLseSumUbBuf2, ConstInfo::BUFFER_SIZE_BYTE_256B);
180+ pipe->InitBuffer(fdLseUbBuf, ConstInfo::BUFFER_SIZE_BYTE_256B);
181+ }
182+}
183+ 
184+template <typename SFAAT> __aicore__ inline
185+void SFAAFlashDecodeServiceGqa<SFAAT>::AllocEventID()
186+{
187+ SetFlag<AscendC::HardEvent::V_MTE2>(SYNC_LSE_SUM_BUF1_FLAG);
188+ SetFlag<AscendC::HardEvent::V_MTE2>(SYNC_LSE_SUM_BUF2_FLAG);
189+ SetFlag<AscendC::HardEvent::V_MTE2>(SYNC_LSE_MAX_BUF1_FLAG);
190+ SetFlag<AscendC::HardEvent::V_MTE2>(SYNC_LSE_MAX_BUF2_FLAG);
191+ SetFlag<AscendC::HardEvent::V_MTE2>(SYNC_MM2RES_BUF1_FLAG);
192+ SetFlag<AscendC::HardEvent::V_MTE2>(SYNC_MM2RES_BUF2_FLAG);
193+ SetFlag<AscendC::HardEvent::MTE3_V>(SYNC_FDOUTPUT_BUF_FLAG);
194+ SetFlag<AscendC::HardEvent::MTE3_V>(SYNC_LSEOUTPUT_BUF_FLAG);
195+ SetFlag<AscendC::HardEvent::V_MTE2>(SYNC_SINK_BUF1_FLAG);
196+ SetFlag<AscendC::HardEvent::V_MTE2>(SYNC_SINK_BUF2_FLAG);
197+}
198+ 
199+template <typename SFAAT> __aicore__ inline
200+void SFAAFlashDecodeServiceGqa<SFAAT>::FreeEventID()
201+{
202+ WaitFlag<AscendC::HardEvent::V_MTE2>(SYNC_LSE_SUM_BUF1_FLAG);
203+ WaitFlag<AscendC::HardEvent::V_MTE2>(SYNC_LSE_SUM_BUF2_FLAG);
204+ WaitFlag<AscendC::HardEvent::V_MTE2>(SYNC_LSE_MAX_BUF1_FLAG);
205+ WaitFlag<AscendC::HardEvent::V_MTE2>(SYNC_LSE_MAX_BUF2_FLAG);
206+ WaitFlag<AscendC::HardEvent::V_MTE2>(SYNC_MM2RES_BUF1_FLAG);
207+ WaitFlag<AscendC::HardEvent::V_MTE2>(SYNC_MM2RES_BUF2_FLAG);
208+ WaitFlag<AscendC::HardEvent::MTE3_V>(SYNC_FDOUTPUT_BUF_FLAG);
209+ WaitFlag<AscendC::HardEvent::MTE3_V>(SYNC_LSEOUTPUT_BUF_FLAG);
210+ WaitFlag<AscendC::HardEvent::V_MTE2>(SYNC_SINK_BUF1_FLAG);
211+ WaitFlag<AscendC::HardEvent::V_MTE2>(SYNC_SINK_BUF2_FLAG);
212+}
213+ 
214+template <typename SFAAT> __aicore__ inline
215+void SFAAFlashDecodeServiceGqa<SFAAT>::CopyAccumOutIn(LocalTensor<T> &accumOutLocal, uint32_t splitKVIndex,
216+ uint32_t startRow, uint32_t dealRowCount)
217+{
218+ DataCopyExtParams copyInParams;
219+ DataCopyPadExtParams<T> copyInPadParams;
220+ copyInParams.blockCount = dealRowCount;
221+ copyInParams.blockLen = constInfo.headDim * sizeof(T);
222+ copyInParams.srcStride = 0;
223+ copyInParams.dstStride = (constInfo.headDimAlign - constInfo.headDim) / BLOCK_ELEMENT_NUM;
224+ 
225+ copyInPadParams.isPad = true;
226+ copyInPadParams.leftPadding = 0;
227+ copyInPadParams.rightPadding = (constInfo.headDimAlign - constInfo.headDim) % BLOCK_ELEMENT_NUM;
228+ copyInPadParams.paddingValue = 0;
229+ uint64_t combineAccumOutOffset = startRow * constInfo.headDim + // taskoffset + g轴offset
230+ splitKVIndex * constInfo.mBaseSize * constInfo.headDim; // 份数offset
231+ DataCopyPad(accumOutLocal, accumOutGm[combineAccumOutOffset], copyInParams, copyInPadParams);
232+}
233+ 
234+template <typename SFAAT> __aicore__ inline
235+void SFAAFlashDecodeServiceGqa<SFAAT>::CopyLseIn(uint32_t startRow,
236+ uint32_t dealRowCount, uint64_t baseOffset, uint32_t cntM)
237+{
238+ LocalTensor<T> lseSum = cntM % 2 == 0 ? fdSumBuf1.Get<T>() : fdSumBuf2.Get<T>();
239+ LocalTensor<T> lseMax = cntM % 2 == 0 ? fdMaxBuf1.Get<T>() : fdMaxBuf2.Get<T>();
240+ 
241+ WaitFlag<AscendC::HardEvent::V_MTE2>(SYNC_LSE_SUM_BUF1_FLAG + cntM % 2);
242+ WaitFlag<AscendC::HardEvent::V_MTE2>(SYNC_LSE_MAX_BUF1_FLAG + cntM % 2);
243+ 
244+ uint64_t combineLseOffset = (baseOffset + startRow) * FP32_BLOCK_ELEMENT_NUM;
245+ uint64_t combineLoopOffset = constInfo.mBaseSize * FP32_BLOCK_ELEMENT_NUM;
246+ uint64_t dealRowCountAlign = dealRowCount * FP32_BLOCK_ELEMENT_NUM;
247+ for (uint32_t i = 0; i < taskInfo.actualCombineLoopSize; i++) {
248+ DataCopy(lseSum[i * dealRowCountAlign], lseSumFdGm[combineLseOffset + i * combineLoopOffset],
249+ dealRowCountAlign); // 份数offset
250+ DataCopy(lseMax[i * dealRowCountAlign], lseMaxFdGm[combineLseOffset + i * combineLoopOffset],
251+ dealRowCountAlign);
252+ }
253+ 
254+ SetFlag<AscendC::HardEvent::MTE2_V>(SYNC_LSE_SUM_BUF1_FLAG + cntM % 2);
255+ SetFlag<AscendC::HardEvent::MTE2_V>(SYNC_LSE_MAX_BUF1_FLAG + cntM % 2);
256+ WaitFlag<AscendC::HardEvent::MTE2_V>(SYNC_LSE_SUM_BUF1_FLAG + cntM % 2);
257+ WaitFlag<AscendC::HardEvent::MTE2_V>(SYNC_LSE_MAX_BUF1_FLAG + cntM % 2);
258+}
259+ 
260+template <typename SFAAT> __aicore__ inline void
261+SFAAFlashDecodeServiceGqa<SFAAT>::ComputeScaleValue(LocalTensor<T> &lseExp,
262+ uint32_t startRow,
263+ uint32_t dealRowCount,
264+ uint32_t cntM)
265+{
266+ LocalTensor<T> lseSum = cntM % 2 == 0 ? fdSumBuf1.Get<T>() : fdSumBuf2.Get<T>();
267+ LocalTensor<T> lseMax = cntM % 2 == 0 ? fdMaxBuf1.Get<T>() : fdMaxBuf2.Get<T>();
268+ 
269+ // 开双buff
270+ LocalTensor<T> lseMaxUb = cntM % 2 == 0 ? fdLseMaxUbBuf1.Get<T>() : fdLseMaxUbBuf2.Get<T>();
271+ LocalTensor<T> lseSumUb = cntM % 2 == 0 ? fdLseSumUbBuf1.Get<T>() : fdLseSumUbBuf2.Get<T>();
272+ uint64_t dealRowCountAlign = dealRowCount * FP32_BLOCK_ELEMENT_NUM;
273+ 
274+ Duplicate(lseMaxUb, -ConstInfo::FLOAT_MAX, dealRowCountAlign);
275+ Duplicate(lseSumUb, ConstInfo::FLOAT_ZERO, dealRowCountAlign);
276+ AscendC::PipeBarrier<PIPE_V>();
277+ 
278+ ColMax(lseMaxUb, lseMax, lseMaxUb, taskInfo.actualCombineLoopSize, dealRowCountAlign, dealRowCountAlign);
279+ AscendC::PipeBarrier<PIPE_V>();
280+ 
281+ RowSub(lseExp, lseMax, lseMaxUb, taskInfo.actualCombineLoopSize, dealRowCountAlign, dealRowCountAlign);
282+ AscendC::PipeBarrier<PIPE_V>();
283+ 
284+ Exp(lseExp, lseExp, taskInfo.actualCombineLoopSize * dealRowCountAlign);
285+ AscendC::PipeBarrier<PIPE_V>();
286+ 
287+ Mul(lseExp, lseSum, lseExp, taskInfo.actualCombineLoopSize * dealRowCountAlign);
288+ AscendC::PipeBarrier<PIPE_V>();
289+ 
290+ ColAdd(lseSumUb, lseExp, lseSumUb, taskInfo.actualCombineLoopSize, dealRowCountAlign, dealRowCountAlign);
291+ AscendC::PipeBarrier<PIPE_V>();
292+ 
293+ MatDivsVec(lseExp, lseExp, lseSumUb, taskInfo.actualCombineLoopSize, dealRowCountAlign, dealRowCountAlign);
294+ AscendC::PipeBarrier<PIPE_V>();
295+}
296+ 
297+template <typename SFAAT>
298+__aicore__ inline void SFAAFlashDecodeServiceGqa<SFAAT>::Bmm2DataCopyOutTrans(LocalTensor<OUT_T> &attenOutUb, uint32_t startRow,
299+ uint32_t dealRowCount, uint32_t columnCount)
300+{
301+ uint32_t s1StartIdx = startRow / constInfo.gSize;
302+ uint32_t startGOffset = startRow % constInfo.gSize;
303+ uint32_t s1EndIdx = CeilDiv(startRow + dealRowCount, static_cast<uint32_t>(constInfo.gSize)) - 1;
304+ uint32_t curStartRow = startRow;
305+ uint32_t curDealRowCount = 0;
306+ uint32_t ubOffset = 0;
307+ 
308+ uint64_t actualSeqQPrefixSum;
309+ if constexpr (LAYOUT_T == SFAA_LAYOUT::TND) {
310+ actualSeqQPrefixSum = (taskInfo.bIdx <= 0) ? 0 : static_cast<uint32_t>(actualSeqLengthsGmQ.GetValue(taskInfo.bIdx - 1));
311+ } else {
312+ actualSeqQPrefixSum = (taskInfo.bIdx <= 0) ? 0 : taskInfo.bIdx * constInfo.qSeqSize;
313+ }
314+ uint64_t attenOutOffset = actualSeqQPrefixSum * constInfo.qHeadNum * constInfo.headDim
315+ + taskInfo.gS1Idx * constInfo.kvHeadNum * constInfo.headDim
316+ + taskInfo.n2Idx * constInfo.gSize * constInfo.headDim; // gS1Idx：与V2操作保持一致，前提是不切G，否则会有向下取整的问题
317+
318+ for (uint32_t curS1idx = s1StartIdx; curS1idx <= s1EndIdx; curS1idx++) {
319+ uint32_t outOffset = attenOutOffset + curS1idx * constInfo.qHeadNum * constInfo.headDim + startGOffset * constInfo.headDim;
320+ if (curS1idx != s1EndIdx) {
321+ curDealRowCount = (curS1idx + 1) * constInfo.gSize - curStartRow;
322+ } else {
323+ curDealRowCount = startRow + dealRowCount - curStartRow;
324+ }
325+ ubOffset = (curStartRow - startRow) * columnCount;
326+ LocalTensor<OUT_T> curAttenOutUb = attenOutUb[ubOffset];
327+ DataCopyExtParams dataCopyParams;
328+ dataCopyParams.blockCount = curDealRowCount;
329+ dataCopyParams.blockLen = columnCount * sizeof(OUT_T);
330+ dataCopyParams.srcStride = (columnCount - columnCount) / (BYTE_BLOCK / sizeof(OUT_T));
331+ dataCopyParams.dstStride = 0;
332+ DataCopyPad(attentionOutGm[outOffset], curAttenOutUb, dataCopyParams);
333+ curStartRow += curDealRowCount;
334+ startGOffset = 0;
335+ }
336+}
337+ 
338+template <typename SFAAT>__aicore__ inline
339+void SFAAFlashDecodeServiceGqa<SFAAT>::Bmm2DataCopyOut(uint64_t attenOutOffset, LocalTensor<OUT_T> &attenOutUb,
340+ uint32_t startRow, uint32_t dealRowCount,
341+ uint32_t columnCount, uint32_t actualColumnCount)
342+{
343+ DataCopyExtParams dataCopyParams;
344+ dataCopyParams.blockCount = dealRowCount;
345+ dataCopyParams.blockLen = actualColumnCount * sizeof(OUT_T);
346+ dataCopyParams.srcStride = (columnCount - actualColumnCount) / (BYTE_BLOCK / sizeof(OUT_T));
347+ dataCopyParams.dstStride = 0;
348+ DataCopyPad(attentionOutGm[attenOutOffset + startRow * actualColumnCount], attenOutUb,
349+ dataCopyParams);
350+}
351+ 
352+template <typename SFAAT>__aicore__ inline
353+void SFAAFlashDecodeServiceGqa<SFAAT>::ReduceFinalRes(LocalTensor<T> &reduceOut,
354+ LocalTensor<T> &mm2Res,
355+ LocalTensor<T> &lseLocal,
356+ uint32_t cntKV,
357+ uint32_t dealRowCount)
358+{
359+ uint32_t dealRowCountAlign = dealRowCount * FP32_BLOCK_ELEMENT_NUM;
360+ LocalTensor<T> tmpRst =
361+ cntKV == 0 ? reduceOut : mm2Res; // 第一次mul结果直接写入reduceOut，否则在mm2Res原地进行mul，再加到reduceOut
362+ 
363+ RowMuls(tmpRst, mm2Res, lseLocal[cntKV * dealRowCountAlign], dealRowCount, constInfo.headDimAlign, constInfo.headDim);
364+ 
365+ if (cntKV != 0) {
366+ AscendC::PipeBarrier<PIPE_V>();
367+ Add(reduceOut, reduceOut, tmpRst, dealRowCount * constInfo.headDimAlign);
368+ AscendC::PipeBarrier<PIPE_V>();
369+ }
370+}
371+ 
372+template <typename SFAAT> __aicore__ inline
373+void SFAAFlashDecodeServiceGqa<SFAAT>::CopyFinalResOut(LocalTensor<T> &accumOutLocal,
374+ uint32_t startRow,
375+ uint32_t dealRowCount,
376+ uint32_t cntM)
377+{
378+ LocalTensor<OUT_T> tmpBmm2ResCastTensor = fdOutputBuf.Get<OUT_T>();
379+ WaitFlag<AscendC::HardEvent::MTE3_V>(SYNC_FDOUTPUT_BUF_FLAG);
380+ uint32_t shapeArray[] = {dealRowCount, (uint32_t)constInfo.headDim};
381+ tmpBmm2ResCastTensor.SetShapeInfo(ShapeInfo(2, shapeArray, DataFormat::ND));
382+ if constexpr (IsSameType<OUT_T, bfloat16_t>::value) { // bf16 采取四舍六入五成双模式
383+ Cast(tmpBmm2ResCastTensor, accumOutLocal, AscendC::RoundMode::CAST_RINT, dealRowCount * constInfo.headDimAlign);
384+ } else {
385+ Cast(tmpBmm2ResCastTensor, accumOutLocal, AscendC::RoundMode::CAST_ROUND, dealRowCount * constInfo.headDimAlign);
386+ }
387+ 
388+ SetFlag<AscendC::HardEvent::V_MTE3>(SYNC_FDOUTPUT_BUF_FLAG);
389+ WaitFlag<AscendC::HardEvent::V_MTE3>(SYNC_FDOUTPUT_BUF_FLAG);
390+ Bmm2DataCopyOutTrans(tmpBmm2ResCastTensor, startRow, dealRowCount, constInfo.headDimAlign);
391+ SetFlag<AscendC::HardEvent::MTE3_V>(SYNC_FDOUTPUT_BUF_FLAG);
392+}
393+ 
394+template <typename SFAAT> __aicore__ inline void
395+SFAAFlashDecodeServiceGqa<SFAAT>::FlashDecode(FDparams &fd)
396+{
397+ if (blockIdx >= fd.usedVecNumOfFd) {
398+ return;
399+ }
400+ uint32_t fdTaskPrevEnd = (blockIdx > 0) ? fd.gS1IdxEndOfFdHead[blockIdx - 1] : 0; // 上一个核末尾是第几个规约
401+ uint32_t fdS1gOuterMPrevEnd =
402+ (blockIdx > 0) ? fd.gS1IdxEndOfFdHeadSplit[blockIdx - 1] : 0; //上一个核末尾是该规约的第几个base行
403+ uint32_t fdTaskEnd = fd.gS1IdxEndOfFdHead[blockIdx]; // 当前核的末尾是第几个规约任务
404+ uint32_t fdS1gOuterMEnd = fd.gS1IdxEndOfFdHeadSplit[blockIdx]; // 当前核的末尾是该规约的第几个base行
405+ uint32_t tmpFdS1gOuterMStart = (blockIdx > 0) ? fdS1gOuterMPrevEnd + 1 : 0; // 当前核从第几个base行开始
406+ uint32_t tmpFdS1gOuterMEnd = 0;
407+ uint32_t reduceGlobaLoop = 0;
408+ uint32_t reduceMLoop = 0;
409+ 
410+ for (uint32_t fdTaskId = fdTaskPrevEnd; fdTaskId <= fdTaskEnd; fdTaskId++) {
411+ tmpFdS1gOuterMEnd = (fdTaskId == fdTaskEnd) ? fdS1gOuterMEnd : (fd.gS1SplitNumOfFdHead[fdTaskId] - 1);
412+ taskInfo.bIdx = fd.bN2IdxOfFdHead[fdTaskId] / constInfo.kvHeadNum;
413+ taskInfo.n2Idx = fd.bN2IdxOfFdHead[fdTaskId] % constInfo.kvHeadNum;
414+ taskInfo.gS1Idx = fd.gS1IdxOfFdHead[fdTaskId] * constInfo.mBaseSize;
415+ taskInfo.actualCombineLoopSize = fd.s2SplitNumOfFdHead[fdTaskId]; // 当前规约任务kv方向有几份
416+ // CalcPreNextTokens();
417+ 
418+ uint64_t combineTaskPrefixSum = 0;
419+ for (int i = 0; i < fdTaskId; i++) {
420+ // 计算此前规约数据的累计份数，每一份的数据大小为 kvHeadNum * constInfo.tndSgBasicSize
421+ // |Task0-0|Task0-1|Task0-3|Task1-0|Task1-2|...|
422+ combineTaskPrefixSum += fd.s2SplitNumOfFdHead[i];
423+ }
424+ 
425+ uint64_t taskOffset = combineTaskPrefixSum * constInfo.mBaseSize;
426+ 
427+ for (uint32_t fdS1gOuterMIdx = tmpFdS1gOuterMStart; fdS1gOuterMIdx <= tmpFdS1gOuterMEnd;
428+ fdS1gOuterMIdx++) { // 左闭右闭
429+ 
430+ uint32_t actualGSplitSize = fd.gS1BaseSizeOfFd;
431+ if (fdS1gOuterMIdx == fd.gS1SplitNumOfFdHead[fdTaskId] - 1) {
432+ actualGSplitSize = fd.gS1LastPartSizeOfFdHead[fdTaskId];
433+ }
434+ uint32_t startRow = fdS1gOuterMIdx * fd.gS1BaseSizeOfFd;
435+ 
436+ LocalTensor<T> lseExp = fdLseExpBuf.Get<T>();
437+ LocalTensor<T> reduceOut = fdReduceBuf.Get<T>();
438+ CopyLseIn(startRow, actualGSplitSize, taskOffset, reduceMLoop);
439+ 
440+ LocalTensor<T> mm2Res;
441+ for (uint32_t preLoadIdx = 0; preLoadIdx < constInfo.preLoadNum; preLoadIdx++) {
442+ mm2Res = (reduceGlobaLoop + preLoadIdx) % 2 == 0 ? fdMm2ResBuf1.Get<T>() : fdMm2ResBuf2.Get<T>();
443+ WaitFlag<AscendC::HardEvent::V_MTE2>(SYNC_MM2RES_BUF1_FLAG + (reduceGlobaLoop + preLoadIdx) % 2);
444+ CopyAccumOutIn(mm2Res, preLoadIdx, taskOffset + startRow, actualGSplitSize);
445+ SetFlag<AscendC::HardEvent::MTE2_V>(SYNC_MM2RES_BUF1_FLAG + (reduceGlobaLoop + preLoadIdx) % 2);
446+ }
447+ 
448+ ComputeScaleValue(lseExp, startRow, actualGSplitSize, reduceMLoop);
449+ SetFlag<AscendC::HardEvent::V_MTE2>(SYNC_LSE_SUM_BUF1_FLAG + reduceMLoop % 2);
450+ SetFlag<AscendC::HardEvent::V_MTE2>(SYNC_LSE_MAX_BUF1_FLAG + reduceMLoop % 2);
451+ 
452+ for (uint32_t i = 0; i < taskInfo.actualCombineLoopSize; i++) {
453+ mm2Res = reduceGlobaLoop % 2 == 0 ? fdMm2ResBuf1.Get<T>() : fdMm2ResBuf2.Get<T>();
454+ if (i >= constInfo.preLoadNum) {
455+ WaitFlag<AscendC::HardEvent::V_MTE2>(SYNC_MM2RES_BUF1_FLAG + reduceGlobaLoop % 2);
456+ CopyAccumOutIn(mm2Res, i, taskOffset + startRow, actualGSplitSize);
457+ SetFlag<AscendC::HardEvent::MTE2_V>(SYNC_MM2RES_BUF1_FLAG + reduceGlobaLoop % 2);
458+ }
459+ 
460+ WaitFlag<AscendC::HardEvent::MTE2_V>(SYNC_MM2RES_BUF1_FLAG + reduceGlobaLoop % 2);
461+ ReduceFinalRes(reduceOut, mm2Res, lseExp, i, actualGSplitSize);
462+ SetFlag<AscendC::HardEvent::V_MTE2>(SYNC_MM2RES_BUF1_FLAG + reduceGlobaLoop % 2);
463+ reduceGlobaLoop += 1;
464+ }
465+ CopyFinalResOut(reduceOut, startRow, actualGSplitSize, reduceMLoop);
466+ reduceMLoop += 1;
467+ }
468+ tmpFdS1gOuterMStart = 0;
469+ }
470+}
471+#endif
Aexperimental/attention/sparse_flash_attention_antiquant/op_kernel/sparse_flash_attention_antiquant_service_vector_gqa_msd.h+1763-0
@@ -0,0 +1,1763 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file sparse_flash_attention_antiquant_service_vector_mla.h
13+ * \brief
14+ */
15+#ifndef SPARSE_FLASH_ATTENTION_ANTIQUANT_SERVICE_VECTOR_GQA_MSD_H
16+#define SPARSE_FLASH_ATTENTION_ANTIQUANT_SERVICE_VECTOR_GQA_MSD_H
17+ 
18+#include "kernel_operator.h"
19+#include "kernel_operator_list_tensor_intf.h"
20+#include "kernel_tiling/kernel_tiling.h"
21+#include "lib/matmul_intf.h"
22+#include "lib/matrix/matmul/tiling.h"
23+#include "sparse_flash_attention_antiquant_common.h"
24+#include "sparse_flash_attention_antiquant_metadata.h"
25+ 
26+using AscendC::CrossCoreSetFlag;
27+using AscendC::CrossCoreWaitFlag;
28+ 
29+#define ARRAY_SIZE(x) (sizeof(x) / sizeof((x)[0]))
30+#define Q_AMAX_BUF_BYTES (ConstInfo::BUFFER_SIZE_BYTE_2K + ConstInfo::BUFFER_SIZE_BYTE_256B)
31+#define P_AMAX_BUF_BYTES (ConstInfo::BUFFER_SIZE_BYTE_2K + ConstInfo::BUFFER_SIZE_BYTE_256B)
32+#define Q_AMAX_BUF_SIZE (Q_AMAX_BUF_BYTES / sizeof(T))
33+#define SOFTMAX_MAX_BUF_BYTES ConstInfo::BUFFER_SIZE_BYTE_512B
34+#define SOFTMAX_EXP_BUF_BYTES ConstInfo::BUFFER_SIZE_BYTE_512B
35+#define SOFTMAX_SUM_BUF_BYTES ConstInfo::BUFFER_SIZE_BYTE_512B
36+#define SOFTMAX_MAX_BUF_SIZE (SOFTMAX_MAX_BUF_BYTES / sizeof(T))
37+#define SOFTMAX_EXP_BUF_SIZE (SOFTMAX_EXP_BUF_BYTES / sizeof(T))
38+#define SOFTMAX_SUM_BUF_SIZE (SOFTMAX_SUM_BUF_BYTES / sizeof(T))
39+#define SALSV_S2BASEIZE 1024
40+#define SALSV_MERGESIZE 384
41+#define SALSV_S2BASEIZE_1 1023
42+#define SFACV_SOFTMAX_WITHOUT_BRC_PRELOAD
43+constexpr SoftmaxConfig IFA_SOFTMAX_FLASHV2_CFG = {false}; // 将isCheckTiling设置为false
44+template <typename SFAAT> class SFAAVectorServiceGqaMsd {
45+public:
46+ // 中间计算数据类型为float，高精度模式
47+ using Q_T = typename SFAAT::queryType;
48+ using T = float;
49+ using KV_T = typename SFAAT::kvType;
50+ using KV_ORIGIN_T = typename SFAAT::queryType;
51+ using OUT_T = typename SFAAT::outputType;
52+ using UPDATE_T = half;
53+ using MM1_OUT_T = half;
54+ using MM2_OUT_T = half;
55+ 
56+ __aicore__ inline SFAAVectorServiceGqaMsd(){};
57+ __aicore__ inline void ProcessVec0Msd(const RunInfo &info);
58+ __aicore__ inline void ProcessVec1Msd(const RunInfo &info);
59+ __aicore__ inline void ProcessVec2Msd(const RunInfo &info);
60+ __aicore__ inline void InitBuffers(TPipe *pipe);
61+ __aicore__ inline void InitParams(const struct ConstInfo &constInfo,
62+ const SparseFlashAttentionAntiquantTilingDataMla *__restrict tilingData,
63+ SfaMetaData *metaDataPtr);
64+ __aicore__ inline void InitVec0GlobalTensor(const GlobalTensor<KV_T> &keyMergeGm, const GlobalTensor<KV_T> &valueMergeGm,
65+ const GlobalTensor<KV_T> &queryPreProcessResGm, const GlobalTensor<Q_T> &queryGm,
66+ const GlobalTensor<KV_T> &keyGm, const GlobalTensor<KV_T> &valueGm,
67+ const GlobalTensor<int32_t> &blkTableGm, const GlobalTensor<T> &keyDequantScaleGm);
68+ __aicore__ inline void InitVec1GlobalTensor(GlobalTensor<MM1_OUT_T> mm1ResGm, GlobalTensor<KV_T> vec1ResGm,
69+ GlobalTensor<int32_t> actualSeqLengthsQGm,
70+ GlobalTensor<int32_t> actualSeqLengthsKVGm, GlobalTensor<T> lseMaxFdGm,
71+ GlobalTensor<T> lseSumFdGm, GlobalTensor<int32_t> topKGm);
72+ __aicore__ inline void InitVec2GlobalTensor(const GlobalTensor<T> &valueDequantScaleGm, GlobalTensor<T> accumOutGm,
73+ GlobalTensor<MM2_OUT_T> mm2ResGm, GlobalTensor<OUT_T> attentionOutGm);
74+ __aicore__ inline void AllocEventID();
75+ __aicore__ inline void FreeEventID();
76+ __aicore__ inline void InitSoftmaxDefaultBuffer();
77+ 
78+private:
79+ // ================================Base Vector==========================================
80+ __aicore__ inline void RowDivs(LocalTensor<float> dstUb, LocalTensor<float> src0Ub, LocalTensor<float> src1Ub,
81+ uint32_t dealRowCount, uint32_t columnCount, uint32_t actualColumnCount);
82+ __aicore__ inline void VecMulMat(LocalTensor<float> dstUb, LocalTensor<float> src0Ub, LocalTensor<float> src1Ub,
83+ uint32_t dealRowCount, uint32_t columnCount, uint32_t actualColumnCount);
84+ __aicore__ inline void RowMaxForLongColumnCount(LocalTensor<float> &dstUb, LocalTensor<float> srcUb, uint32_t dealRowCount,
85+ uint32_t columnCount, uint32_t actualColumnCount);
86+ __aicore__ inline void RowMax(LocalTensor<float> &dstUb, LocalTensor<float> &srcUb, uint32_t dealRowCount, uint32_t columnCount,
87+ uint32_t actualColumnCount);
88+ // ================================Vector0==========================================
89+ __aicore__ inline void QueryPreProcess(const RunInfo &info);
90+ __aicore__ inline void DealQueryPreProcessBaseBlock(const RunInfo &info, uint32_t startRow, uint32_t dealRowCount,
91+ uint32_t columnCount, uint32_t actualColumnCount);
92+ __aicore__ inline void CopyAntiqQuery(LocalTensor<T> &queryCastUb, uint64_t qOffset, uint32_t dealRowCount,
93+ uint32_t columnCount, uint32_t actualColumnCount);
94+ __aicore__ inline void AntiquantMatmulPreProcess(const RunInfo &info, GlobalTensor<KV_T> dstGm, LocalTensor<T> aMaxResUb,
95+ LocalTensor<T> srcUb, LocalTensor<T> tmpAFloorUb, uint32_t startRow,
96+ uint32_t dealRowCount, uint32_t columnCount, uint32_t actualColumnCount);
97+ __aicore__ inline void AbsRowMax(LocalTensor<T> &tmpAMaxRes, LocalTensor<T> &srcUb, LocalTensor<T> tmpAUb,
98+ uint32_t dealRowCount, uint32_t columnCount, uint32_t actualColumnCount);
99+ __aicore__ inline void AntiquantAIterExpand(GlobalTensor<KV_T> &dstGm, LocalTensor<half> &tmpA1, LocalTensor<half> &tmpA2,
100+ uint32_t calcSize, bool isFirst, uint64_t outOffset);
101+ __aicore__ inline int64_t MergeKv(const RunInfo &runInfo, int64_t s2GmStartOffset, int64_t s2GmLimit,
102+ int64_t topkGmBaseOffset, bool isValue, int64_t mergeMte3Idx, int64_t s2GmStartOffset4Merge);
103+ __aicore__ inline int64_t GetKeyBNBOffset(int64_t realS2Idx, const RunInfo &runInfo, int64_t s2IdLimit);
104+ __aicore__ inline void GetRealS2Idx(int64_t s2GmOffset, int64_t &realS2Idx, int64_t topkGmBaseOffset,
105+ const RunInfo &runInfo);
106+ __aicore__ inline void CopyInKv(int64_t &mte2Size, int64_t mte3Size, int64_t mergeMte3Idx, int64_t realS2Idx1,
107+ int64_t realS2Idx2, const RunInfo &runInfo, bool &mask,
108+ int64_t s2GmOffsetArray, int64_t s2GmLimit);
109+ __aicore__ inline void CopyOutMrgeResult(int64_t mte2Size, int64_t mte3Size, int64_t s2StartGmOffset,
110+ int64_t mergeMte3Idx, const RunInfo &runInfo, bool isValue, bool &needWaitMte3ToMte2);
111+ __aicore__ inline void SetInfInBlk(const LocalTensor<T> &mmResUb, uint32_t dealRowCount, uint32_t columnCount,
112+ uint64_t startId, uint64_t endId);
113+ __aicore__ inline void SetInfInBlkHasTail(const LocalTensor<T> &mmResUb, uint32_t dealRowCount, uint32_t columnCount,
114+ uint64_t startId, uint64_t endId);
115+ __aicore__ inline void SetMidInf(const LocalTensor<T> &mmResUb, uint32_t dealRowCount, uint32_t columnCount,
116+ uint64_t startId, uint64_t endId);
117+ __aicore__ inline void CopyInSingleKv(int64_t &mte2Size, int64_t mte3Size, int64_t mergeMte3Idx, int64_t realS2Idx,
118+ int64_t keyBNBOffset, int64_t s2IdLimit, const RunInfo &runInfo, bool &mask,
119+ int64_t &s2GmOffsetArray, int64_t s2GmLimit);
120+ // ================================Vector1==========================================
121+ __aicore__ inline void ProcessVec1SingleBuf(const RunInfo &info, const MSplitInfo &mSplitInfo);
122+ __aicore__ inline void DealBmm1ResBaseBlock(const RunInfo &info, const MSplitInfo &mSplitInfo, uint32_t startRow,
123+ uint32_t dealRowCount, uint32_t columnCount, uint32_t actualColumnCount,
124+ uint32_t loopId, bool &needMask, uint32_t &maskStart, uint32_t &maskEnd);
125+ __aicore__ inline void AntiquantMatmulResCombineDD(const RunInfo &info, LocalTensor<T> bmmResUb, GlobalTensor<MM1_OUT_T> srcGm,
126+ uint32_t startRow, uint32_t dealRowCount, uint32_t columnCount, uint32_t actualColumnCount,
127+ float scaleC);
128+ __aicore__ inline void AntiquantSoftmaxResPreProcess(const RunInfo &info, GlobalTensor<KV_T> dstGm, LocalTensor<T> srcUb,
129+ LocalTensor<T> tmpAFloorUb, uint32_t startRow,
130+ uint32_t dealRowCount, uint32_t columnCount, uint32_t actualColumnCount);
131+ __aicore__ inline void SoftmaxFlashV2Compute(const RunInfo &info, const MSplitInfo &mSplitInfo,
132+ LocalTensor<T> &mmResUb, LocalTensor<uint8_t> &softmaxTmpUb,
133+ uint32_t startRow, uint32_t dealRowCount, uint32_t columnCount,
134+ uint32_t actualColumnCount);
135+ __aicore__ inline void ElewiseCompute(const RunInfo &info, const LocalTensor<T> &mmResUb, uint32_t dealRowCount,
136+ uint32_t columnCount);
137+ __aicore__ inline void AttentionMaskCompute(const RunInfo &info, const MSplitInfo &mSplitInfo,
138+ const LocalTensor<T> &mmResUb, uint32_t dealRowCount,
139+ uint32_t columnCount, uint32_t startRow, bool &needMask,
140+ uint32_t &maskStart, uint32_t &maskEnd);
141+ __aicore__ inline void ComputeLogSumExpAndCopyToGm(const RunInfo &info, const MSplitInfo &mSplitInfo,
142+ LocalTensor<T> &softmaxSumUb, LocalTensor<T> &softmaxMaxUb);
143+ // ================================Vecotr2==========================================
144+ __aicore__ inline void ProcessVec2SingleBuf(const RunInfo &info, const MSplitInfo &mSplitInfo);
145+ __aicore__ inline void DealBmm2ResBaseBlock(const RunInfo &info, const MSplitInfo &mSplitInfo, uint32_t startRow,
146+ uint32_t dealRowCount, uint32_t columnCount,
147+ uint32_t actualColumnCount);
148+ __aicore__ inline void AntiquantMM2ResCombine(const RunInfo &info, LocalTensor<MM2_OUT_T> bmmResUb,
149+ GlobalTensor<MM2_OUT_T> srcGm, uint32_t startRow, uint32_t dealRowCount,
150+ uint32_t columnCount, uint32_t actualColumnCount);
151+ __aicore__ inline void CopyAntiquantScale(LocalTensor<T> &castUb, GlobalTensor<T> srcGm, uint64_t offset);
152+ __aicore__ inline void ProcessVec2Inner(const RunInfo &info, const MSplitInfo &mSplitInfo, uint32_t mStartRow,
153+ uint32_t mDealSize);
154+ __aicore__ inline void Bmm2DataCopyOutTrans(const RunInfo &info, LocalTensor<OUT_T> &attenOutUb, uint32_t wsMStart,
155+ uint32_t dealRowCount, uint32_t columnCount,
156+ uint32_t actualColumnCount);
157+ __aicore__ inline void Bmm2ResCopyOut(const RunInfo &info, LocalTensor<T> &bmm2ResUb, uint32_t wsMStart,
158+ uint32_t dealRowCount, uint32_t columnCount, uint32_t actualColumnCount);
159+ __aicore__ inline void Bmm2CastAndCopyOut(const RunInfo &info, LocalTensor<T> &bmm2ResUb, uint32_t wsMStart,
160+ uint32_t dealRowCount, uint32_t columnCount, uint32_t actualColumnCount);
161+ __aicore__ inline void Bmm2FDDataCopyOut(const RunInfo &info, LocalTensor<T> &bmm2ResUb, uint32_t wsMStart,
162+ uint32_t dealRowCount, uint32_t columnCount, uint32_t actualColumnCount);
163+ __aicore__ inline uint64_t CalcAccumOffset(uint32_t bN2Idx, uint32_t gS1Idx);
164+ __aicore__ inline void GetConfusionTransposeTiling(int64_t numR, int64_t numC, const uint32_t stackBufferSize,
165+ const uint32_t typeSize, ConfusionTransposeTiling &tiling);
166+ 
167+ static constexpr bool PAGE_ATTENTION = SFAAT::pageAttention;
168+ static constexpr int TEMPLATE_MODE = SFAAT::templateMode;
169+ static constexpr bool IS_META = SFAAT::flashDecode;
170+ bool FLASH_DECODE = SFAAT::flashDecode;
171+ static constexpr SFAA_LAYOUT LAYOUT_T = SFAAT::layout;
172+ static constexpr SFAA_LAYOUT KV_LAYOUT_T = SFAAT::kvLayout;
173+ 
174+ static constexpr uint64_t MERGE_CACHE_GM_BUF_NUM = 4;
175+ static constexpr uint64_t SYNC_INPUT_BUF1_FLAG = 2;
176+ static constexpr uint64_t SYNC_INPUT_BUF2_FLAG = 7;
177+ static constexpr uint64_t SYNC_INPUT_BUF2_PONG_FLAG = 8;
178+ static constexpr uint64_t SYNC_INPUT_DEQUANT_SCALE_FLAG = 6;
179+ static constexpr uint64_t SYNC_OUTPUT_BUF1_FLAG = 4;
180+ static constexpr uint64_t SYNC_OUTPUT_BUF2_FLAG = 5;
181+ static constexpr uint32_t INPUT1_BUFFER_OFFSET = ConstInfo::BUFFER_SIZE_BYTE_8K;
182+ static constexpr uint32_t SOFTMAX_TMP_BUFFER_OFFSET = ConstInfo::BUFFER_SIZE_BYTE_512B / sizeof(T);
183+ static constexpr uint32_t BASE_BLOCK_MAX_ELEMENT_NUM = ConstInfo::BUFFER_SIZE_BYTE_32K / sizeof(T); // 32768/4=8192
184+ static constexpr uint32_t BLOCK_ELEMENT_NUM = BYTE_BLOCK / sizeof(T); // 32/4=8
185+ static constexpr T FLOAT_E_SCALAR = 8388608;
186+ static constexpr T LN2 = 0.6931471805599453094172;
187+ static constexpr T RECIP_OF_LN2 = 1 / LN2;
188+ static constexpr T SOFTMAX_MIN_NUM = -2e38;
189+ static constexpr T antiqCoeff1 = 127;
190+ static constexpr T antiqCoeff2 = 1 / antiqCoeff1;
191+ static constexpr float scaleC1 = 1024.0 / 127;
192+ static constexpr float scaleC2 = 1024.0;
193+ MM1_OUT_T antiquantExpandCoeff = 254;
194+ static constexpr uint32_t msdIterNum = 2;
195+ 
196+ const SparseFlashAttentionAntiquantTilingDataMla *__restrict tilingData;
197+ SfaMetaData *metaDataPtr;
198+ 
199+ uint32_t pingpongFlag = 0U;
200+ ConstInfo constInfo = {};
201+ uint32_t lastBN2Idx = 1 << 31;
202+ 
203+ GlobalTensor<int32_t> mm2ResInt32Gm;
204+ GlobalTensor<MM1_OUT_T> mm1ResGm;
205+ GlobalTensor<KV_T> vec1ResGm;
206+ GlobalTensor<T> lseSumFdGm;
207+ GlobalTensor<T> lseMaxFdGm;
208+ 
209+ GlobalTensor<int32_t> actualSeqLengthsQGm;
210+ GlobalTensor<int32_t> actualSeqLengthsKVGm;
211+ GlobalTensor<T> vec2ResGm;
212+ GlobalTensor<MM2_OUT_T> mm2ResGm;
213+ GlobalTensor<T> accumOutGm;
214+ GlobalTensor<OUT_T> attentionOutGm;
215+ GlobalTensor<int32_t> blkTableGm_;
216+ 
217+ GlobalTensor<KV_T> keyMergeGm_;
218+ GlobalTensor<KV_T> valueMergeGm_;
219+ GlobalTensor<KV_T> queryPreProcessResGm_;
220+ GlobalTensor<Q_T> queryGm_;
221+ GlobalTensor<KV_T> keyGm_;
222+ GlobalTensor<KV_T> valueGm_;
223+ GlobalTensor<T> keyDequantScaleGm_;
224+ GlobalTensor<T> valueDequantScaleGm_;
225+ GlobalTensor<int32_t> topkGm_;
226+ 
227+ // ================================Local Buffer区====================================
228+ 
229+ // queue
230+ TBuf<> inputBuf1; // 32K, inque
231+ TBuf<> inputBuf2; // 16K, inque
232+ TBuf<> outputBuf1; // 32K, outque
233+ TBuf<> outputBuf2; // 8K, outque
234+ 
235+ // 临时tbuf
236+ TBuf<> tmpBuff1; // 32K
237+ TBuf<> tmpBuff2; // 32K
238+ TBuf<> tmpBuff3; // 2K
239+ 
240+ // 常驻tbuf
241+ TBuf<> vec2ResBuff; // 16k 伪量化场景
242+ TBuf<> antiqKeyScaleBuff; // 2K
243+ TBuf<> antiqValueScaleBuff; // 2K
244+ TBuf<> qAmaxBuff; // constInfo.preLoadNum * (2K + 256B)
245+ TBuf<> softmaxResAmaxBuff; // 2K + 256B
246+ TBuf<> qRowSumBuff; // 2K + 256B
247+ TBuf<> softmaxResRowSumBuff; // 2K + 256B
248+ TBuf<> softmaxMaxBuff; // constInfo.preLoadNum * 2K
249+ TBuf<> softmaxExpBuff; // constInfo.preLoadNum * 2K
250+ TBuf<> softmaxSumBuff; // constInfo.preLoadNum * 2K
251+ TBuf<> softmaxMaxDefaultBuff; // 2K
252+ TBuf<> softmaxSumDefaultBuff; // 2K
253+ 
254+ LocalTensor<T> softmaxMaxUb;
255+ LocalTensor<T> softmaxSumUb;
256+ LocalTensor<T> softmaxExpUb;
257+ LocalTensor<T> softmaxMaxDefaultUb;
258+ LocalTensor<T> softmaxSumDefaultUb;
259+ 
260+ LocalTensor<KV_T> kvMergUb_;
261+ 
262+ // antiquant msd
263+ LocalTensor<T> aMaxBmm1Ub;
264+ LocalTensor<T> aMaxBmm2Ub;
265+ LocalTensor<T> softmaxScaleResRowSumUb;
266+ LocalTensor<half> vec2ResUb;
267+ LocalTensor<T> antiquantKScaleUb_;
268+ LocalTensor<T> antiquantVScaleUb_;
269+ LocalTensor<T> qRowSumUb;
270+};
271+ 
272+__aicore__ inline uint32_t GetMinPowerTwo(uint32_t cap)
273+{
274+ uint32_t i = 1;
275+ while (i < cap) {
276+ i = i << 1;
277+ }
278+ return i;
279+}
280+ 
281+template <typename SFAAT> __aicore__ inline void SFAAVectorServiceGqaMsd<SFAAT>::InitBuffers(TPipe *pipe)
282+{
283+ // 182.5k
284+ pipe->InitBuffer(inputBuf1, ConstInfo::BUFFER_SIZE_BYTE_16K * 2); // 32K // mm1 mm2 res
285+ pipe->InitBuffer(inputBuf2, ConstInfo::BUFFER_SIZE_BYTE_8K * 2); // 16K // kvmerge
286+ pipe->InitBuffer(outputBuf1, ConstInfo::BUFFER_SIZE_BYTE_16K); // 16K // attentionOut
287+ pipe->InitBuffer(vec2ResBuff, ConstInfo::BUFFER_SIZE_BYTE_16K); // 16K // flashupdate常驻ub
288+ pipe->InitBuffer(outputBuf2, ConstInfo::BUFFER_SIZE_BYTE_8K); // 8K
289+ 
290+ // tmpBuff
291+ pipe->InitBuffer(tmpBuff1, ConstInfo::BUFFER_SIZE_BYTE_32K); // 32K
292+ pipe->InitBuffer(tmpBuff2, ConstInfo::BUFFER_SIZE_BYTE_32K); // 32K
293+ pipe->InitBuffer(tmpBuff3, ConstInfo::BUFFER_SIZE_BYTE_2K); // 2K
294+ 
295+ // 常驻buffer
296+ pipe->InitBuffer(antiqKeyScaleBuff, ConstInfo::BUFFER_SIZE_BYTE_2K); // 2K // 实际0.5K
297+ pipe->InitBuffer(antiqValueScaleBuff, ConstInfo::BUFFER_SIZE_BYTE_2K); // 2K
298+ // 预留空间2K = 64 * 32，支持 gSize = 64
299+ // brcb 操作每次操作8*32字节输出，startRow接近64时，
300+ // 输出最多可能超出2k空间7*32字节， 这里预留256B防止越界
301+ pipe->InitBuffer(qAmaxBuff, Q_AMAX_BUF_BYTES * constInfo.preLoadNum); // 2.25*2 = 4.5K
302+ pipe->InitBuffer(softmaxMaxBuff, SOFTMAX_MAX_BUF_BYTES * constInfo.preLoadNum); // 0.5*2 = 1K
303+ pipe->InitBuffer(softmaxExpBuff, SOFTMAX_EXP_BUF_BYTES * constInfo.preLoadNum); // 0.5*2 = 1K
304+ pipe->InitBuffer(softmaxSumBuff, SOFTMAX_SUM_BUF_BYTES * constInfo.preLoadNum); // 0.5*2 = 1K
305+ 
306+ pipe->InitBuffer(softmaxMaxDefaultBuff, SOFTMAX_MAX_BUF_BYTES); // 0.5K
307+ pipe->InitBuffer(softmaxSumDefaultBuff, SOFTMAX_SUM_BUF_BYTES); // 0.5K
308+ 
309+ softmaxMaxUb = softmaxMaxBuff.Get<T>();
310+ softmaxSumUb = softmaxSumBuff.Get<T>();
311+ softmaxExpUb = softmaxExpBuff.Get<T>();
312+ 
313+ softmaxMaxDefaultUb = softmaxMaxDefaultBuff.Get<T>();
314+ softmaxSumDefaultUb = softmaxSumDefaultBuff.Get<T>();
315+ 
316+ kvMergUb_ = inputBuf2.Get<KV_T>();
317+ 
318+ antiquantKScaleUb_ = antiqKeyScaleBuff.Get<T>();
319+ antiquantVScaleUb_ = antiqValueScaleBuff.Get<T>();
320+ vec2ResUb = vec2ResBuff.Get<half>();
321+ aMaxBmm1Ub = qAmaxBuff.Get<T>();
322+}
323+ 
324+template <typename SFAAT>
325+__aicore__ inline void
326+SFAAVectorServiceGqaMsd<SFAAT>::InitParams(const struct ConstInfo &constInfo,
327+ const SparseFlashAttentionAntiquantTilingDataMla *__restrict tilingData,
328+ SfaMetaData *metaDataPtr)
329+{
330+ this->constInfo = constInfo;
331+ this->tilingData = tilingData;
332+ this->metaDataPtr = metaDataPtr;
333+ if constexpr (IS_META) {
334+ FLASH_DECODE = metaDataPtr->numOfFdHead > 0U;
335+ }
336+}
337+ 
338+template <typename SFAAT>
339+__aicore__ inline void SFAAVectorServiceGqaMsd<SFAAT>::InitVec0GlobalTensor(
340+ const GlobalTensor<KV_T> &keyMergeGm, const GlobalTensor<KV_T> &valueMergeGm,
341+ const GlobalTensor<KV_T> &queryPreProcessResGm, const GlobalTensor<Q_T> &queryGm, const GlobalTensor<KV_T> &keyGm, const GlobalTensor<KV_T> &valueGm,
342+ const GlobalTensor<int32_t> &blkTableGm, const GlobalTensor<T> &keyDequantScaleGm)
343+{
344+ this->keyMergeGm_ = keyMergeGm;
345+ this->valueMergeGm_ = valueMergeGm;
346+ this->queryPreProcessResGm_ = queryPreProcessResGm;
347+ this->queryGm_ = queryGm;
348+ this->keyGm_ = keyGm;
349+ this->valueGm_ = valueGm;
350+ this->blkTableGm_ = blkTableGm;
351+ this->keyDequantScaleGm_ = keyDequantScaleGm;
352+}
353+ 
354+template <typename SFAAT>
355+__aicore__ inline void SFAAVectorServiceGqaMsd<SFAAT>::InitVec1GlobalTensor(
356+ GlobalTensor<MM1_OUT_T> mm1ResGm, GlobalTensor<KV_T> vec1ResGm,
357+ GlobalTensor<int32_t> actualSeqLengthsQGm, GlobalTensor<int32_t> actualSeqLengthsKVGm, GlobalTensor<T> lseMaxFdGm,
358+ GlobalTensor<T> lseSumFdGm, GlobalTensor<int32_t> topKGm)
359+{
360+ this->mm1ResGm = mm1ResGm;
361+ this->vec1ResGm = vec1ResGm;
362+ this->actualSeqLengthsQGm = actualSeqLengthsQGm;
363+ this->actualSeqLengthsKVGm = actualSeqLengthsKVGm;
364+ this->lseMaxFdGm = lseMaxFdGm;
365+ this->lseSumFdGm = lseSumFdGm;
366+ this->topkGm_ = topKGm;
367+}
368+ 
369+template <typename SFAAT>
370+__aicore__ inline void SFAAVectorServiceGqaMsd<SFAAT>::InitVec2GlobalTensor(const GlobalTensor<T> &valueDequantScaleGm,
371+ GlobalTensor<T> accumOutGm,
372+ GlobalTensor<MM2_OUT_T> mm2ResGm,
373+ GlobalTensor<OUT_T> attentionOutGm)
374+{
375+ this->accumOutGm = accumOutGm;
376+ this->mm2ResGm = mm2ResGm;
377+ this->attentionOutGm = attentionOutGm;
378+ this->valueDequantScaleGm_ = valueDequantScaleGm;
379+}
380+ 
381+template <typename SFAAT> __aicore__ inline void SFAAVectorServiceGqaMsd<SFAAT>::AllocEventID()
382+{
383+ SetFlag<AscendC::HardEvent::V_MTE2>(SYNC_INPUT_BUF1_FLAG);
384+ SetFlag<AscendC::HardEvent::MTE3_MTE2>(SYNC_INPUT_BUF2_FLAG);
385+ SetFlag<AscendC::HardEvent::MTE3_MTE2>(SYNC_INPUT_BUF2_PONG_FLAG);
386+ SetFlag<AscendC::HardEvent::V_MTE2>(SYNC_INPUT_DEQUANT_SCALE_FLAG);
387+ SetFlag<AscendC::HardEvent::MTE3_V>(SYNC_OUTPUT_BUF1_FLAG);
388+ SetFlag<AscendC::HardEvent::MTE3_V>(SYNC_OUTPUT_BUF2_FLAG);
389+}
390+ 
391+template <typename SFAAT> __aicore__ inline void SFAAVectorServiceGqaMsd<SFAAT>::FreeEventID()
392+{
393+ WaitFlag<AscendC::HardEvent::V_MTE2>(SYNC_INPUT_BUF1_FLAG);
394+ WaitFlag<AscendC::HardEvent::MTE3_MTE2>(SYNC_INPUT_BUF2_FLAG);
395+ WaitFlag<AscendC::HardEvent::MTE3_MTE2>(SYNC_INPUT_BUF2_PONG_FLAG);
396+ WaitFlag<AscendC::HardEvent::V_MTE2>(SYNC_INPUT_DEQUANT_SCALE_FLAG);
397+ WaitFlag<AscendC::HardEvent::MTE3_V>(SYNC_OUTPUT_BUF1_FLAG);
398+ WaitFlag<AscendC::HardEvent::MTE3_V>(SYNC_OUTPUT_BUF2_FLAG);
399+}
400+ 
401+template <typename SFAAT> __aicore__ inline void SFAAVectorServiceGqaMsd<SFAAT>::InitSoftmaxDefaultBuffer()
402+{
403+ Duplicate(softmaxMaxDefaultUb, SOFTMAX_MIN_NUM, SOFTMAX_TMP_BUFFER_OFFSET);
404+ Duplicate(softmaxSumDefaultUb, ConstInfo::FLOAT_ZERO, SOFTMAX_TMP_BUFFER_OFFSET);
405+}
406+ 
407+template <typename SFAAT>
408+__aicore__ inline void SFAAVectorServiceGqaMsd<SFAAT>::ComputeLogSumExpAndCopyToGm(const RunInfo &info,
409+ const MSplitInfo &mSplitInfo,
410+ LocalTensor<T> &softmaxSumUb,
411+ LocalTensor<T> &softmaxMaxUb)
412+{
413+ if (mSplitInfo.vecDealM == 0) {
414+ return;
415+ }
416+ uint64_t baseOffset = mSplitInfo.nBufferStartM / 2;
417+ size_t size = mSplitInfo.vecDealM * FP32_BLOCK_ELEMENT_NUM;
418+ uint64_t accumTmpOutNum = CalcAccumOffset(info.bN2Idx, info.gS1Idx);
419+ uint64_t offset = (accumTmpOutNum * constInfo.mBaseSize + // taskoffset
420+ info.tndCoreStartKVSplitPos * constInfo.mBaseSize + // 份数offset
421+ mSplitInfo.nBufferStartM + mSplitInfo.vecStartM) *
422+ FP32_BLOCK_ELEMENT_NUM; // m轴offset
423+ if (info.actualSingleProcessSInnerSize != 0) {
424+ LocalTensor<T> tmp = outputBuf2.Get<T>();
425+ WaitFlag<AscendC::HardEvent::MTE3_V>(SYNC_OUTPUT_BUF2_FLAG);
426+ Brcb(tmp, softmaxSumUb[baseOffset], (mSplitInfo.vecDealM + 7) / 8, {1, 8});
427+ SetFlag<AscendC::HardEvent::V_MTE3>(SYNC_OUTPUT_BUF2_FLAG);
428+ WaitFlag<AscendC::HardEvent::V_MTE3>(SYNC_OUTPUT_BUF2_FLAG);
429+ DataCopy(lseSumFdGm[offset], tmp, size);
430+ SetFlag<AscendC::HardEvent::MTE3_V>(SYNC_OUTPUT_BUF2_FLAG);
431+ 
432+ tmp = outputBuf2.Get<T>();
433+ WaitFlag<AscendC::HardEvent::MTE3_V>(SYNC_OUTPUT_BUF2_FLAG);
434+ Brcb(tmp, softmaxMaxUb[baseOffset], (mSplitInfo.vecDealM + 7) / 8, {1, 8});
435+ SetFlag<AscendC::HardEvent::V_MTE3>(SYNC_OUTPUT_BUF2_FLAG);
436+ WaitFlag<AscendC::HardEvent::V_MTE3>(SYNC_OUTPUT_BUF2_FLAG);
437+ DataCopy(lseMaxFdGm[offset], tmp, size);
438+ SetFlag<AscendC::HardEvent::MTE3_V>(SYNC_OUTPUT_BUF2_FLAG);
439+ } else {
440+ matmul::InitOutput<T>(lseSumFdGm[offset], size, ConstInfo::FLOAT_ZERO);
441+ matmul::InitOutput<T>(lseMaxFdGm[offset], size, SOFTMAX_MIN_NUM);
442+ }
443+}
444+ 
445+template <typename SFAAT>
446+__aicore__ inline void SFAAVectorServiceGqaMsd<SFAAT>::ElewiseCompute(const RunInfo &info,
447+ const LocalTensor<T> &mmResUb,
448+ uint32_t dealRowCount, uint32_t columnCount)
449+{
450+}
451+template <typename SFAAT>
452+__aicore__ inline void SFAAVectorServiceGqaMsd<SFAAT>::AttentionMaskCompute(const RunInfo &info, const MSplitInfo &mSplitInfo,
453+ const LocalTensor<T> &mmResUb, uint32_t dealRowCount,
454+ uint32_t columnCount, uint32_t startRow, bool &needMask,
455+ uint32_t &maskStart, uint32_t &maskEnd)
456+{
457+ uint32_t maskEndCurrentPart = (maskEnd > info.nextS2BaseOffset) ?
458+ info.nextS2BaseOffset : maskEnd;
459+ uint32_t maskStartCurrentPart = (maskStart < info.curS2BaseOffset) ?
460+ info.curS2BaseOffset : maskStart;
461+ if (!(maskStartCurrentPart < info.actS2Size && needMask) || maskStartCurrentPart >= maskEndCurrentPart) {
462+ return;
463+ }
464+ if (info.nextS2BaseOffset <= maskStartCurrentPart) {
465+ return;
466+ }
467+ uint32_t maskStartSinglePro = (maskStartCurrentPart & SALSV_S2BASEIZE_1);
468+ uint32_t maskEndSinglePro = (maskEndCurrentPart & SALSV_S2BASEIZE_1) == 0 ? SALSV_S2BASEIZE : (maskEndCurrentPart & SALSV_S2BASEIZE_1);
469+ uint32_t gS1StartIdx = info.gS1Idx + mSplitInfo.nBufferStartM + mSplitInfo.vecStartM + startRow;
470+ uint32_t s1StartIdx = gS1StartIdx / constInfo.gSize;
471+ uint32_t gStartIdx = gS1StartIdx % constInfo.gSize;
472+ uint32_t gS1EndIdx = gS1StartIdx + dealRowCount - 1;
473+ uint32_t s1EndIdx = gS1EndIdx / constInfo.gSize;
474+ uint32_t s1Count = s1EndIdx - s1StartIdx + 1;
475+ uint32_t headGCount = s1Count > 1 ? (constInfo.gSize - gStartIdx) : dealRowCount;
476+ uint32_t dstMaskOffset = 0;
477+ int64_t s2StartCeilAlign = SFAAAlign(maskStartSinglePro, 8);
478+ int64_t s2MidFloorAlign = maskEndSinglePro / 8 * 8;
479+ uint64_t maskEndTmp = s2StartCeilAlign >= maskEndSinglePro ? maskEndSinglePro : s2StartCeilAlign;
480+ uint64_t maskStartTmp = s2StartCeilAlign <= s2MidFloorAlign ? s2MidFloorAlign : s2StartCeilAlign;
481+ SetInfInBlkHasTail(mmResUb, headGCount, columnCount, maskStartSinglePro, maskEndTmp);
482+ SetMidInf(mmResUb, headGCount, columnCount, s2StartCeilAlign, s2MidFloorAlign);
483+ SetInfInBlkHasTail(mmResUb, headGCount, columnCount, maskStartTmp, maskEndSinglePro);
484+ if ((headGCount < dealRowCount) || ((headGCount == dealRowCount) && (((gS1StartIdx + headGCount) % (constInfo.gSize)) == 0))) {
485+ maskStart++;
486+ maskStartCurrentPart = (maskStart < info.curS2BaseOffset) ?
487+ info.curS2BaseOffset : maskStart;
488+ if (!(maskStartCurrentPart < info.actS2Size) || maskStartCurrentPart >= maskEndCurrentPart) {
489+ return;
490+ }
491+ if (info.nextS2BaseOffset <= maskStartCurrentPart) {
492+ return;
493+ }
494+ maskStartSinglePro = (maskStartCurrentPart - info.curS2BaseOffset);
495+ }
496+ dstMaskOffset += headGCount * columnCount;
497+ uint32_t remainRowCount = dealRowCount - headGCount;
498+ uint32_t midS1Count = remainRowCount / constInfo.gSize;
499+ uint32_t tailGSize = remainRowCount % constInfo.gSize;
500+ for (uint32_t midIdx = 0; midIdx < midS1Count; ++midIdx) {
501+ s2StartCeilAlign = SFAAAlign(maskStartSinglePro, 8);
502+ s2MidFloorAlign = maskEndSinglePro / 8 * 8;
503+ maskEndTmp = s2StartCeilAlign >= maskEndSinglePro ? maskEndSinglePro : s2StartCeilAlign;
504+ maskStartTmp = s2StartCeilAlign <= s2MidFloorAlign ? s2MidFloorAlign : s2StartCeilAlign;
505+ SetInfInBlkHasTail(mmResUb[dstMaskOffset], constInfo.gSize, columnCount, maskStartSinglePro, maskEndTmp);
506+ SetMidInf(mmResUb[dstMaskOffset], constInfo.gSize, columnCount, s2StartCeilAlign, s2MidFloorAlign);
507+ SetInfInBlkHasTail(mmResUb[dstMaskOffset], constInfo.gSize, columnCount, maskStartTmp, maskEndSinglePro);
508+ dstMaskOffset += constInfo.gSize * columnCount;
509+ maskStart++;
510+ maskStartCurrentPart = (maskStart < info.curS2BaseOffset) ?
511+ info.curS2BaseOffset : maskStart;
512+ if (!(maskStartCurrentPart < info.actS2Size) || maskStartCurrentPart >= maskEndCurrentPart) {
513+ return;
514+ }
515+ if (info.nextS2BaseOffset <= maskStartCurrentPart) {
516+ return;
517+ }
518+ maskStartSinglePro = (maskStartCurrentPart & SALSV_S2BASEIZE_1);
519+ }
520+ if (tailGSize > 0) {
521+ s2StartCeilAlign = SFAAAlign(maskStartSinglePro, 8);
522+ s2MidFloorAlign = maskEndSinglePro / 8 * 8;
523+ maskEndTmp = s2StartCeilAlign >= maskEndSinglePro ? maskEndSinglePro : s2StartCeilAlign;
524+ maskStartTmp = s2StartCeilAlign <= s2MidFloorAlign ? s2MidFloorAlign : s2StartCeilAlign;
525+ 
526+ SetInfInBlkHasTail(mmResUb[dstMaskOffset], tailGSize, columnCount, maskStartSinglePro, maskEndTmp);
527+ SetMidInf(mmResUb[dstMaskOffset], tailGSize, columnCount, s2StartCeilAlign, s2MidFloorAlign);
528+ SetInfInBlkHasTail(mmResUb[dstMaskOffset], tailGSize, columnCount, maskStartTmp, maskEndSinglePro);
529+ }
530+}
531+ 
532+template <typename SFAAT>
533+__aicore__ inline void SFAAVectorServiceGqaMsd<SFAAT>::SetInfInBlk(const LocalTensor<T> &mmResUb,
534+ uint32_t dealRowCount, uint32_t columnCount,
535+ uint64_t startId, uint64_t endId)
536+{
537+ if (startId >= endId) {
538+ return;
539+ }
540+ 
541+ uint64_t startFloorAlignSize = startId / BLOCK_ELEMENT_NUM * BLOCK_ELEMENT_NUM;
542+ uint64_t notComputePreMaskOneBlk = (1ULL << (startId - startFloorAlignSize)) - 1;
543+ uint64_t notComputePostMaskOneBlk = ~((1ULL << (endId - startFloorAlignSize)) - 1);
544+ uint64_t notComputeMaskOneBlk = notComputePreMaskOneBlk ^ notComputePostMaskOneBlk;
545+ 
546+ uint64_t maskOneBlk = (~notComputeMaskOneBlk) & 0xFFFFFFFFULL;
547+ uint64_t mask[1] = {maskOneBlk};
548+ for (int i = 1; i < 8; i++) {
549+ mask[0] = mask[0] | (maskOneBlk << (i * 8));
550+ }
551+ for (uint64_t rowId = 0; rowId < dealRowCount; rowId += 8) {
552+ Duplicate(mmResUb[rowId * columnCount + startFloorAlignSize], SOFTMAX_MIN_NUM, mask,
553+ 1, SfaaCeilDiv(columnCount, 8), 0);
554+ }
555+}
556+ 
557+template <typename SFAAT>
558+__aicore__ inline void SFAAVectorServiceGqaMsd<SFAAT>::SetInfInBlkHasTail(const LocalTensor<T> &mmResUb, uint32_t dealRowCount, uint32_t columnCount,
559+ uint64_t startId, uint64_t endId)
560+{
561+ if (startId >= endId) {
562+ return;
563+ }
564+ uint64_t startFloorAlignSize = startId / BLOCK_ELEMENT_NUM * BLOCK_ELEMENT_NUM;
565+ uint64_t notComputePreMaskOneBlk = ((1ULL << (startId - startFloorAlignSize)) - 1);
566+ uint64_t notComputePostMaskOneBlk = ~((1ULL << (endId - startFloorAlignSize)) - 1);
567+ uint64_t notComputeMaskOneBlk = notComputePreMaskOneBlk ^ notComputePostMaskOneBlk;
568+ uint64_t maskOneBlk = (~notComputeMaskOneBlk) & 0xFFFFFFFFULL;
569+ uint64_t mask[1] = {maskOneBlk};
570+ for (int i = 1; i < 8; ++i) {
571+ mask[0] = mask[0] | (maskOneBlk << (i * 8));
572+ }
573+ uint32_t rowLoop = dealRowCount / 8;
574+ uint32_t rowTail = dealRowCount % 8;
575+ for (uint64_t rowId = 0; rowId < rowLoop; rowId ++) {
576+ Duplicate(mmResUb[rowId * columnCount * 8 + startFloorAlignSize], SOFTMAX_MIN_NUM, mask,
577+ 1, SfaaCeilDiv(columnCount, 8), 0);
578+ }
579+ 
580+ if (rowTail > 0) {
581+ mask[0] = maskOneBlk;
582+ for (int i = 1; i < rowTail; ++i) {
583+ mask[0] = mask[0] | (maskOneBlk << (i * 8));
584+ }
585+ Duplicate(mmResUb[rowLoop * columnCount * 8 + startFloorAlignSize], SOFTMAX_MIN_NUM, mask,
586+ 1, SfaaCeilDiv(columnCount, 8), 0);
587+ }
588+}
589+ 
590+template <typename SFAAT>
591+__aicore__ inline void SFAAVectorServiceGqaMsd<SFAAT>::SetMidInf(const LocalTensor<T> &mmResUb,
592+ uint32_t dealRowCount, uint32_t columnCount,
593+ uint64_t startId, uint64_t endId)
594+{
595+ if (startId >= endId) {
596+ return;
597+ }
598+ // startId endId
599+ // 0 ... 0
600+ // 从startId到endId部分置-inf, startId、endId为32B对齐的下标
601+ for (uint64_t rowId = 0; rowId < dealRowCount; rowId++) {
602+ Duplicate(mmResUb[rowId * columnCount + startId], SOFTMAX_MIN_NUM, endId - startId);
603+ }
604+}
605+ 
606+template <typename SFAAT>
607+__aicore__ inline void SFAAVectorServiceGqaMsd<SFAAT>::SoftmaxFlashV2Compute(
608+ const RunInfo &info, const MSplitInfo &mSplitInfo, LocalTensor<T> &mmResUb, LocalTensor<uint8_t> &softmaxTmpUb,
609+ uint32_t startRow, uint32_t dealRowCount, uint32_t columnCount, uint32_t actualColumnCount)
610+{
611+ LocalTensor<T> inSumTensor;
612+ LocalTensor<T> inMaxTensor;
613+#ifdef SFACV_SOFTMAX_WITHOUT_BRC_PRELOAD
614+ uint32_t baseOffset = mSplitInfo.nBufferStartM / 2 + startRow;
615+#else
616+ uint32_t baseOffset = mSplitInfo.nBufferStartM / 2 + startRow * BLOCK_ELEMENT_NUM;
617+#endif
618+ uint32_t outIdx = info.loop % (constInfo.preLoadNum);
619+ uint32_t softmaxOutOffset = outIdx * SOFTMAX_TMP_BUFFER_OFFSET + baseOffset;
620+ if (info.isFirstSInnerLoop) {
621+ inMaxTensor = softmaxMaxDefaultUb;
622+ inSumTensor = softmaxSumDefaultUb;
623+ } else {
624+ uint32_t inIdx = (info.loop - 1) % (constInfo.preLoadNum);
625+ inMaxTensor = softmaxMaxUb[inIdx * SOFTMAX_TMP_BUFFER_OFFSET + baseOffset];
626+ inSumTensor = softmaxSumUb[inIdx * SOFTMAX_TMP_BUFFER_OFFSET + baseOffset];
627+ }
628+ if (actualColumnCount !=0) {
629+ SoftMaxShapeInfo srcShape{dealRowCount, columnCount, dealRowCount, actualColumnCount};
630+ SoftMaxTiling newTiling =
631+ SoftMaxFlashV2TilingFunc(srcShape, sizeof(T), sizeof(T), softmaxTmpUb.GetSize(), true, false);
632+#ifdef SFACV_SOFTMAX_WITHOUT_BRC_PRELOAD
633+ SoftmaxFlashV2<T, true, true, false, false, SFAA_SOFTMAX_FLASHV2_CFG_WITHOUT_BRC>(
634+ mmResUb, softmaxSumUb[softmaxOutOffset], softmaxMaxUb[softmaxOutOffset], mmResUb,
635+ softmaxExpUb[softmaxOutOffset], inSumTensor, inMaxTensor, softmaxTmpUb, newTiling, srcShape);
636+#else
637+ SoftmaxFlashV2<T, true, true, false, false, IFA_SOFTMAX_FLASHV2_CFG>(
638+ mmResUb, softmaxSumUb[softmaxOutOffset], softmaxMaxUb[softmaxOutOffset], mmResUb,
639+ softmaxExpUb[softmaxOutOffset], inSumTensor, inMaxTensor, softmaxTmpUb, newTiling, srcShape);
640+#endif
641+ } else {
642+ DataCopy(softmaxSumUb[softmaxOutOffset], inSumTensor, dealRowCount);
643+ pipe_barrier(PIPE_V);
644+ DataCopy(softmaxMaxUb[softmaxOutOffset], inMaxTensor, dealRowCount);
645+ }
646+}
647+ 
648+// 由于S在GM上已完成atomatic，此处只需要把数据拷贝出来有fp16转成fp32
649+template <typename SFAAT>
650+__aicore__ inline void SFAAVectorServiceGqaMsd<SFAAT>::AntiquantMatmulResCombineDD(const RunInfo &info,
651+ LocalTensor<T> bmmResUb, GlobalTensor<MM1_OUT_T> srcGm, uint32_t startRow, uint32_t dealRowCount,
652+ uint32_t columnCount, uint32_t actualColumnCount, float scaleC)
653+{
654+ uint32_t baseOffset = startRow * columnCount;
655+ uint32_t copySize = dealRowCount * columnCount;
656+ 
657+ LocalTensor<MM1_OUT_T> tmpMMRes = inputBuf1.Get<MM1_OUT_T>();
658+ WaitFlag<AscendC::HardEvent::V_MTE2>(SYNC_INPUT_BUF1_FLAG);
659+ DataCopy(tmpMMRes, srcGm[baseOffset], copySize);
660+ SetFlag<AscendC::HardEvent::MTE2_V>(SYNC_INPUT_BUF1_FLAG);
661+ WaitFlag<AscendC::HardEvent::MTE2_V>(SYNC_INPUT_BUF1_FLAG);
662+ Cast(bmmResUb, tmpMMRes, AscendC::RoundMode::CAST_NONE, copySize);
663+ PipeBarrier<PIPE_V>();
664+ SetFlag<AscendC::HardEvent::V_MTE2>(SYNC_INPUT_BUF1_FLAG);
665+ Muls(bmmResUb, bmmResUb, scaleC, copySize);
666+ PipeBarrier<PIPE_V>();
667+}
668+ 
669+template <typename SFAAT>
670+__aicore__ inline void SFAAVectorServiceGqaMsd<SFAAT>::DealBmm1ResBaseBlock(
671+ const RunInfo &info, const MSplitInfo &mSplitInfo, uint32_t startRow, uint32_t dealRowCount,
672+ uint32_t columnCount, uint32_t actualColumnCount, uint32_t loopId, bool &needMask, uint32_t &maskStart, uint32_t &maskEnd)
673+{
674+ uint32_t computeSize = dealRowCount * columnCount;
675+ uint64_t inGmOffset = (info.loop % constInfo.preLoadNum) * constInfo.mmResUbSize +
676+ (mSplitInfo.nBufferStartM + mSplitInfo.vecStartM) * columnCount;
677+ uint64_t outGmOffset = (info.loop % constInfo.preLoadNum) * constInfo.mmResUbSize * 2 +
678+ (mSplitInfo.nBufferStartM + mSplitInfo.vecStartM) * columnCount;
679+ LocalTensor<T> mmResUb = tmpBuff1.Get<T>();
680+ AntiquantMatmulResCombineDD(info, mmResUb, mm1ResGm[inGmOffset], startRow, dealRowCount, columnCount,
681+ actualColumnCount, scaleC1 * static_cast<T>(tilingData->baseParams.scaleValue));
682+ LocalTensor<T> aMax = aMaxBmm1Ub[(info.bN2Idx % constInfo.preLoadNum) * Q_AMAX_BUF_SIZE + startRow * BLOCK_ELEMENT_NUM];
683+ RowMuls<T>(mmResUb, mmResUb, aMax, dealRowCount, columnCount, actualColumnCount);
684+ PipeBarrier<PIPE_V>();
685+ if (loopId == 0) {
686+ maskEnd = maskStart + constInfo.sparseShardSize - 1;
687+ maskStart += mSplitInfo.vecStartM / constInfo.gSize;
688+ }
689+ ElewiseCompute(info, mmResUb, dealRowCount, columnCount);
690+ pipe_barrier(PIPE_V);
691+ AttentionMaskCompute(info, mSplitInfo, mmResUb, dealRowCount, columnCount,
692+ startRow, needMask, maskStart, maskEnd);
693+ pipe_barrier(PIPE_V);
694+ LocalTensor<T> tmpAFloorUb = tmpBuff2.Get<T>();
695+ LocalTensor<uint8_t> softmaxTmpUb = tmpAFloorUb.template ReinterpretCast<uint8_t>();
696+ SoftmaxFlashV2Compute(info, mSplitInfo, mmResUb, softmaxTmpUb, startRow, dealRowCount,
697+ columnCount, actualColumnCount);
698+ pipe_barrier(PIPE_V);
699+ AntiquantSoftmaxResPreProcess(info, vec1ResGm[outGmOffset], mmResUb, tmpAFloorUb, startRow, dealRowCount, columnCount,
700+ actualColumnCount);
701+}
702+ 
703+template <typename SFAAT>
704+__aicore__ inline void SFAAVectorServiceGqaMsd<SFAAT>::AntiquantSoftmaxResPreProcess(const RunInfo &info,
705+ GlobalTensor<KV_T> dstGm, LocalTensor<T> srcUb, LocalTensor<T> tmpAFloorUb, uint32_t startRow,
706+ uint32_t dealRowCount, uint32_t columnCount, uint32_t actualColumnCount)
707+{
708+ uint32_t step = info.mSize * columnCount;
709+ uint32_t baseOffset = startRow * columnCount;
710+ uint32_t calcSize = dealRowCount * columnCount;
711+ 
712+ Muls(srcUb, srcUb, antiqCoeff1, calcSize); // 127 * A
713+ PipeBarrier<PIPE_V>();
714+ 
715+ Cast(tmpAFloorUb, srcUb, RoundMode::CAST_ROUND, calcSize); // fp32
716+ PipeBarrier<PIPE_V>();
717+ LocalTensor<half> tmpAFloorUbFp16 = tmpAFloorUb.template ReinterpretCast<half>();
718+ tmpAFloorUbFp16.SetSize(tmpAFloorUb.GetSize());
719+ Cast(tmpAFloorUbFp16, tmpAFloorUb, RoundMode::CAST_ROUND, calcSize); // A1:fp16
720+ PipeBarrier<PIPE_V>();
721+ // step5: 将Qtmp转成fp16
722+ LocalTensor<half> srcUbFp16 = srcUb.template ReinterpretCast<half>();
723+ srcUbFp16.SetSize(srcUb.GetSize());
724+ Cast(srcUbFp16, srcUb, RoundMode::CAST_ROUND, calcSize);
725+ PipeBarrier<PIPE_V>();
726+ 
727+ for (uint32_t i = 0; i < msdIterNum; i++) {
728+ AntiquantAIterExpand(dstGm, srcUbFp16, tmpAFloorUbFp16, calcSize, (i == 0 ? true : false), step * i + baseOffset);
729+ }
730+}
731+ 
732+template <typename SFAAT>
733+__aicore__ inline void SFAAVectorServiceGqaMsd<SFAAT>::ProcessVec1SingleBuf(const RunInfo &info,
734+ const MSplitInfo &mSplitInfo)
735+{
736+ if (mSplitInfo.vecDealM == 0) {
737+ return;
738+ }
739+ uint32_t mSplitSize = info.actualSingleProcessSInnerSize == 0 ?
740+ 16 : BASE_BLOCK_MAX_ELEMENT_NUM / info.actualSingleProcessSInnerSizeAlign;
741+#ifdef SFACV_SOFTMAX_WITHOUT_BRC_PRELOAD
742+ // 1. 向下8对齐是因为UB操作至少32B
743+ // 2. info.actualSingleProcessSInnerSizeAlign最大512, mSplitSize可以确保最小为16
744+ mSplitSize = mSplitSize / 8 * 8;
745+#endif
746+ 
747+ if (mSplitSize > mSplitInfo.vecDealM) {
748+ mSplitSize = mSplitInfo.vecDealM;
749+ }
750+ uint32_t loopCount = (mSplitInfo.vecDealM + mSplitSize - 1) / mSplitSize;
751+ uint32_t tailSplitSize = mSplitInfo.vecDealM - (loopCount - 1) * mSplitSize;
752+ uint32_t maskStart = info.maskStart;
753+ uint32_t maskEnd = info.maskEnd;
754+ bool needMask = (constInfo.sparseMode == 3);
755+ for (uint32_t i = 0, dealSize = mSplitSize; i < loopCount; i++) {
756+ if (i == (loopCount - 1)) {
757+ dealSize = tailSplitSize;
758+ }
759+ DealBmm1ResBaseBlock(info, mSplitInfo, i * mSplitSize, dealSize,
760+ info.actualSingleProcessSInnerSizeAlign, info.actualSingleProcessSInnerSize, i, needMask, maskStart, maskEnd);
761+ pingpongFlag ^= 1; // pingpong 0 1切换
762+ }
763+}
764+ 
765+template <typename SFAAT>
766+__aicore__ inline void SFAAVectorServiceGqaMsd<SFAAT>::GetRealS2Idx(int64_t s2GmOffset, int64_t &realS2Idx,
767+ int64_t topkGmBaseOffset, const RunInfo &runInfo)
768+{
769+ int64_t topkGmIdx = (s2GmOffset + runInfo.s2Idx * SALSV_S2BASEIZE) / constInfo.sparseBlockSize;
770+ if (unlikely(topkGmIdx >= runInfo.sparseBlockCount)) {
771+ realS2Idx = -1;
772+ return;
773+ }
774+ realS2Idx = topkGm_.GetValue(topkGmBaseOffset + topkGmIdx) * static_cast<int64_t>(constInfo.sparseBlockSize) +
775+ static_cast<int64_t>((s2GmOffset + runInfo.s2Idx * SALSV_S2BASEIZE) % constInfo.sparseBlockSize);
776+}
777+ 
778+template <typename SFAAT>
779+__aicore__ inline int64_t SFAAVectorServiceGqaMsd<SFAAT>::GetKeyBNBOffset(int64_t realS2Idx,
780+ const RunInfo &runInfo, int64_t s2IdLimit)
781+{
782+ if (realS2Idx < 0 || realS2Idx >= s2IdLimit) {
783+ return -1;
784+ }
785+ int64_t realKeyBNBOffset = 0;
786+ if constexpr (PAGE_ATTENTION) {
787+ int64_t blkTableIdx = realS2Idx / constInfo.kvCacheBlockSize;
788+ int64_t blkTableOffset = realS2Idx % constInfo.kvCacheBlockSize;
789+ if constexpr (KV_LAYOUT_T == SFAA_LAYOUT::PA_BSND) {
790+ realKeyBNBOffset = (blkTableGm_.GetValue(runInfo.bIdx * constInfo.maxBlockNumPerBatch + blkTableIdx) *
791+ static_cast<int64_t>(constInfo.kvCacheBlockSize) + blkTableOffset) *
792+ static_cast<int64_t>(constInfo.kvHeadNum) + runInfo.n2Idx;
793+ } else if constexpr (KV_LAYOUT_T == SFAA_LAYOUT::PA_BNSD) { // PA_BNSD (blockNum n2 blockSize)
794+ realKeyBNBOffset = blkTableGm_.GetValue(runInfo.bIdx * constInfo.maxBlockNumPerBatch + blkTableIdx) *
795+ static_cast<int64_t>(constInfo.kvHeadNum) * static_cast<int64_t>(constInfo.kvCacheBlockSize) +
796+ + runInfo.n2Idx * static_cast<int64_t>(constInfo.kvCacheBlockSize) + blkTableOffset;
797+ }
798+ } else {
799+ realKeyBNBOffset = runInfo.tensorBOffset / constInfo.headDim + realS2Idx * static_cast<int64_t>(constInfo.kvHeadNum);
800+ }
801+ 
802+ return realKeyBNBOffset;
803+}
804+ 
805+template <typename SFAAT>
806+__aicore__ inline void
807+SFAAVectorServiceGqaMsd<SFAAT>::CopyInSingleKv(int64_t &mte2Size, int64_t mte3Size, int64_t mergeMte3Idx, int64_t realS2Idx,
808+ int64_t keyBNBOffset, int64_t s2IdLimit, const RunInfo &runInfo, bool &mask,
809+ int64_t &s2GmOffsetArray, int64_t s2GmLimit)
810+{
811+ if (keyBNBOffset < 0) {
812+ return;
813+ }
814+ int64_t sparseBlockRealSize = (s2GmLimit - s2GmOffsetArray < constInfo.sparseBlockSize) ?
815+ (s2GmLimit - s2GmOffsetArray) : constInfo.sparseBlockSize;
816+ int64_t validS2Count =
817+ (realS2Idx + sparseBlockRealSize > s2IdLimit ? s2IdLimit - realS2Idx : sparseBlockRealSize);
818+ // 当前仅支持SEPARATE模式
819+ if (constInfo.quantScaleRepoMode == QUANT_SCALE_REPO_MODE::SEPARATE) {
820+ DataCopyExtParams intriParams;
821+ DataCopyPadExtParams<KV_T> padParams;
822+ if constexpr (KV_LAYOUT_T == SFAA_LAYOUT::PA_NZ) {
823+ uint32_t copyFinishRowCnt = 0;
824+ uint32_t blockElementCnt = 32 / sizeof(KV_T);
825+ uint32_t ubOffset = mergeMte3Idx % 2 * INPUT1_BUFFER_OFFSET / sizeof(KV_T) + (mte2Size - mte3Size) * blockElementCnt;
826+ uint32_t kvMergUbRowSize = ConstInfo::BUFFER_SIZE_BYTE_4K / constInfo.headDim;
827+ while (copyFinishRowCnt < validS2Count) {
828+ uint64_t blockIdOffset = realS2Idx / constInfo.kvCacheBlockSize; // 获取block table上的索引
829+ uint64_t reaminRowCnt = realS2Idx % constInfo.kvCacheBlockSize; // 获取在单个块上超出的行数
830+ uint64_t idInBlockTable = blkTableGm_.GetValue(runInfo.bIdx * constInfo.maxBlockNumPerBatch + blockIdOffset); // 从block table上获取编号
831+ // 计算可以拷贝的行数
832+ uint32_t copyRowCnt = constInfo.kvCacheBlockSize - reaminRowCnt;
833+ if (copyFinishRowCnt + copyRowCnt > validS2Count) {
834+ copyRowCnt = validS2Count - copyFinishRowCnt;
835+ }
836+ uint64_t keyOffset = idInBlockTable * constInfo.kvCacheBlockSize * constInfo.headDim * constInfo.kvHeadNum;
837+ keyOffset += (uint64_t)(runInfo.n2Idx * constInfo.headDim * constInfo.kvCacheBlockSize) + reaminRowCnt * blockElementCnt;
838+ intriParams.blockLen = copyRowCnt * blockElementCnt * sizeof(KV_T);
839+ intriParams.blockCount = constInfo.headDim / blockElementCnt;
840+ intriParams.dstStride = kvMergUbRowSize - copyRowCnt;
841+ intriParams.srcStride = (constInfo.kvCacheBlockSize - copyRowCnt) * blockElementCnt * sizeof(KV_T);
842+ padParams.isPad = false;
843+ DataCopyPad(kvMergUb_[ubOffset + copyFinishRowCnt * blockElementCnt], keyGm_[keyOffset], intriParams, padParams);
844+ DataCopyPad(kvMergUb_[ubOffset + ConstInfo::BUFFER_SIZE_BYTE_4K + copyFinishRowCnt * blockElementCnt],
845+ valueGm_[keyOffset], intriParams, padParams);
846+ // 更新循环变量
847+ copyFinishRowCnt += copyRowCnt;
848+ realS2Idx += copyRowCnt;
849+ }
850+ } else {
851+ uint32_t ubOffset = mergeMte3Idx % 2 * INPUT1_BUFFER_OFFSET / sizeof(KV_T) + (mte2Size - mte3Size) * constInfo.headDim;
852+ uint16_t srcStride = 0;
853+ if constexpr (KV_LAYOUT_T == SFAA_LAYOUT::PA_BNSD) {
854+ intriParams.blockCount = 1;
855+ intriParams.blockLen = validS2Count * constInfo.headDim * sizeof(KV_T);
856+ } else if constexpr (KV_LAYOUT_T == SFAA_LAYOUT::PA_BSND) {
857+ srcStride = (constInfo.kvHeadNum - 1) * constInfo.headDim * sizeof(KV_T);
858+ if (unlikely(srcStride == 0)) {
859+ intriParams.blockCount = 1;
860+ intriParams.blockLen = validS2Count * constInfo.headDim * sizeof(KV_T);
861+ } else {
862+ intriParams.blockCount = validS2Count;
863+ intriParams.blockLen = constInfo.headDim * sizeof(KV_T);
864+ }
865+ }
866+ intriParams.dstStride = 0;
867+ intriParams.srcStride = srcStride;
868+ padParams.isPad = false;
869+ DataCopyPad(kvMergUb_[ubOffset], keyGm_[keyBNBOffset * constInfo.headDim], intriParams, padParams);
870+ DataCopyPad(kvMergUb_[ubOffset + ConstInfo::BUFFER_SIZE_BYTE_4K], valueGm_[keyBNBOffset * constInfo.headDim], intriParams, padParams);
871+ }
872+ }
873+ mte2Size += validS2Count;
874+ s2GmOffsetArray += validS2Count;
875+}
876+ 
877+template <typename SFAAT>
878+__aicore__ inline void SFAAVectorServiceGqaMsd<SFAAT>::CopyInKv(int64_t &mte2Size, int64_t mte3Size, int64_t mergeMte3Idx,
879+ int64_t realS2Idx1, int64_t realS2Idx2, const RunInfo &runInfo,
880+ bool &mask, int64_t s2GmOffsetArray, int64_t s2GmLimit)
881+{
882+ int64_t s2IdLimit = runInfo.threshold;
883+ 
884+ int64_t keyBNBOffset1 = GetKeyBNBOffset(realS2Idx1, runInfo, s2IdLimit);
885+ int64_t keyBNBOffset2 = GetKeyBNBOffset(realS2Idx2, runInfo, s2IdLimit);
886+ if (unlikely(keyBNBOffset1 < 0 && keyBNBOffset2 < 0)) {
887+ return;
888+ }
889+ 
890+ int64_t blkTableSrcStride = (keyBNBOffset1 > keyBNBOffset2 ? (keyBNBOffset1 - keyBNBOffset2) :
891+ (keyBNBOffset2 - keyBNBOffset1)) - constInfo.sparseBlockSize;
892+ int64_t keySrcStride = blkTableSrcStride * constInfo.headDim;
893+ if (likely(KV_LAYOUT_T == SFAA_LAYOUT::PA_NZ || constInfo.kvHeadNum > 1 || keySrcStride >= INT32_MAX || keySrcStride < 0 ||
894+ realS2Idx1 + constInfo.sparseBlockSize >= s2IdLimit || realS2Idx2 + constInfo.sparseBlockSize >= s2IdLimit)) {
895+ // stride溢出、stride为负数、s2超长等异常场景，还原成2条搬运指令
896+ CopyInSingleKv(mte2Size, mte3Size, mergeMte3Idx, realS2Idx1, keyBNBOffset1, s2IdLimit, runInfo, mask, s2GmOffsetArray, s2GmLimit);
897+ CopyInSingleKv(mte2Size, mte3Size, mergeMte3Idx, realS2Idx2, keyBNBOffset2, s2IdLimit, runInfo, mask, s2GmOffsetArray, s2GmLimit);
898+ } else {
899+ DataCopyExtParams intriParams;
900+ intriParams.blockCount = (keyBNBOffset1 >= 0) + (keyBNBOffset2 >= 0);
901+ intriParams.dstStride = 0;
902+ intriParams.srcStride = keySrcStride;
903+ DataCopyPadExtParams<KV_T> padParams;
904+ 
905+ int64_t startGmOffset = keyBNBOffset1 > -1 ? keyBNBOffset1 : keyBNBOffset2;
906+ if (keyBNBOffset2 > -1 && keyBNBOffset2 < keyBNBOffset1) {
907+ startGmOffset = keyBNBOffset2;
908+ }
909+ 
910+ // 当前仅支持SEPARATE模式
911+ if (constInfo.quantScaleRepoMode == QUANT_SCALE_REPO_MODE::SEPARATE) {
912+ intriParams.blockLen = constInfo.sparseBlockSize * constInfo.headDim;
913+ uint32_t headDimAlign = SFAAAlign(constInfo.headDim, ConstInfo::BUFFER_SIZE_BYTE_32B) / sizeof(KV_T);
914+ padParams.isPad = true;
915+ padParams.leftPadding = 0;
916+ padParams.rightPadding = headDimAlign - constInfo.headDim;
917+ padParams.paddingValue = 0;
918+ uint32_t ubOffset = mergeMte3Idx % 2 * INPUT1_BUFFER_OFFSET / sizeof(KV_T) + (mte2Size - mte3Size) * headDimAlign;
919+ DataCopyPad(kvMergUb_[ubOffset], keyGm_[startGmOffset * constInfo.headDim], intriParams, padParams);
920+ DataCopyPad(kvMergUb_[ubOffset + ConstInfo::BUFFER_SIZE_BYTE_4K], valueGm_[startGmOffset * constInfo.headDim], intriParams, padParams);
921+ }
922+ mte2Size += ((keyBNBOffset1 > -1) + (keyBNBOffset2 > -1)) * constInfo.sparseBlockSize;
923+ }
924+}
925+ 
926+template <typename SFAAT>
927+__aicore__ inline void SFAAVectorServiceGqaMsd<SFAAT>::CopyOutMrgeResult(int64_t mte2Size, int64_t mte3Size,
928+ int64_t s2GmStartOffset, int64_t mergeMte3Idx,
929+ const RunInfo &runInfo, bool isValue, bool &needWaitMte3ToMte2)
930+{
931+ if (mte2Size <= mte3Size) {
932+ SetFlag<AscendC::HardEvent::MTE3_MTE2>(SYNC_INPUT_BUF2_FLAG + mergeMte3Idx % 2);
933+ return;
934+ }
935+ int32_t dealRow = mte2Size - mte3Size;
936+ LocalTensor<KV_T> srcTensor = kvMergUb_[mergeMte3Idx % 2 * INPUT1_BUFFER_OFFSET / sizeof(KV_T)];
937+
938+ SetFlag<AscendC::HardEvent::MTE2_MTE3>(SYNC_INPUT_BUF2_FLAG);
939+ WaitFlag<AscendC::HardEvent::MTE2_MTE3>(SYNC_INPUT_BUF2_FLAG);
940+ DataCopyExtParams dataCopyParams;
941+ if constexpr (KV_LAYOUT_T == SFAA_LAYOUT::PA_NZ) {
942+ uint32_t blockElementCnt = 32 / sizeof(KV_T);
943+ dataCopyParams.blockCount = constInfo.headDim / blockElementCnt;
944+ dataCopyParams.blockLen = 32 * blockElementCnt * sizeof(KV_T);
945+ dataCopyParams.srcStride = 0;
946+ dataCopyParams.dstStride = (SALSV_MERGESIZE - 32) * blockElementCnt * sizeof(KV_T);
947+ DataCopyPad(valueMergeGm_[runInfo.loop % MERGE_CACHE_GM_BUF_NUM * SALSV_MERGESIZE * constInfo.headDim +
948+ (s2GmStartOffset + mte3Size) * blockElementCnt], srcTensor[ConstInfo::BUFFER_SIZE_BYTE_4K], dataCopyParams);
949+ DataCopyPad(keyMergeGm_[runInfo.loop % MERGE_CACHE_GM_BUF_NUM * SALSV_MERGESIZE * constInfo.headDim +
950+ (s2GmStartOffset + mte3Size) * blockElementCnt], srcTensor, dataCopyParams);
951+
952+ } else {
953+ dataCopyParams.blockCount = dealRow;
954+ dataCopyParams.blockLen = constInfo.headDim * sizeof(KV_T);
955+ dataCopyParams.srcStride = 0;
956+ dataCopyParams.dstStride = 0;
957+ DataCopyPad(keyMergeGm_[runInfo.loop % MERGE_CACHE_GM_BUF_NUM * SALSV_MERGESIZE * constInfo.headDim + (s2GmStartOffset + mte3Size) *
958+ constInfo.headDim], srcTensor, dataCopyParams);
959+ DataCopyPad(valueMergeGm_[runInfo.loop % MERGE_CACHE_GM_BUF_NUM * SALSV_MERGESIZE * constInfo.headDim + (s2GmStartOffset + mte3Size) *
960+ constInfo.headDim], srcTensor[ConstInfo::BUFFER_SIZE_BYTE_4K], dataCopyParams);
961+ }
962+ SetFlag<AscendC::HardEvent::MTE3_MTE2>(SYNC_INPUT_BUF2_FLAG + mergeMte3Idx % 2);
963+ needWaitMte3ToMte2 = true;
964+}
965+ 
966+// b s1 k
967+template <typename SFAAT>
968+__aicore__ inline void SFAAVectorServiceGqaMsd<SFAAT>::ProcessVec0Msd(const RunInfo &info)
969+{
970+ QueryPreProcess(info);
971+ if (info.aivS2AccessSize == 0) {
972+ return;
973+ }
974+ int64_t s2ProcessSize = info.aivS2AccessSize;
975+ int64_t s2Pair = SfaaCeilDiv(s2ProcessSize, 2L * constInfo.sparseBlockSize);
976+ int64_t topkGmBaseOffset = info.topkGmBaseOffset;
977+ 
978+ int64_t s2GmStartOffset = GetSubBlockIdx() == 0 ?
979+ info.aicS2AccessSize : (s2Pair / 2L) * 2 * constInfo.sparseBlockSize + info.aicS2AccessSize;
980+ int64_t s2GmStartOffset4Merge = GetSubBlockIdx() == 0 ?
981+ 0 : (s2Pair / 2L) * 2 * constInfo.sparseBlockSize;
982+ int64_t s2GmLimit = GetSubBlockIdx() == 0 ? (s2Pair / 2L) * 2 * constInfo.sparseBlockSize + info.aicS2AccessSize :
983+ s2ProcessSize + info.aicS2AccessSize;
984+ if (s2GmLimit > s2ProcessSize + info.aicS2AccessSize) {
985+ s2GmLimit = s2ProcessSize + info.aicS2AccessSize;
986+ }
987+ int64_t mergeMte3Idx = 0;
988+ int64_t mte2Size = MergeKv(info, s2GmStartOffset, s2GmLimit, topkGmBaseOffset, false, mergeMte3Idx, s2GmStartOffset4Merge);
989+ return;
990+}
991+ 
992+template <typename SFAAT>
993+__aicore__ inline void SFAAVectorServiceGqaMsd<SFAAT>::QueryPreProcess(const RunInfo &info)
994+{
995+ if (!info.isFirstSInnerLoop) {
996+ return;
997+ }
998+ if (info.bN2Idx != lastBN2Idx) { // 搬入scaleK
999+ DataCopyExtParams dataCopyParams;
1000+ dataCopyParams.blockCount = 1;
1001+ dataCopyParams.blockLen = constInfo.headDim * sizeof(T);
1002+ dataCopyParams.srcStride = 0;
1003+ dataCopyParams.dstStride = 0;
1004+ DataCopyPadExtParams<T> padParams;
1005+ padParams.isPad = false;
1006+ WaitFlag<AscendC::HardEvent::V_MTE2>(SYNC_INPUT_DEQUANT_SCALE_FLAG);
1007+ uint32_t srcOffset = info.n2Idx * constInfo.headDim;
1008+ DataCopyPad(antiquantKScaleUb_, keyDequantScaleGm_[srcOffset], dataCopyParams, padParams);
1009+ SetFlag<AscendC::HardEvent::MTE2_V>(SYNC_INPUT_DEQUANT_SCALE_FLAG);
1010+ WaitFlag<AscendC::HardEvent::MTE2_V>(SYNC_INPUT_DEQUANT_SCALE_FLAG);
1011+ }
1012+
1013+ uint32_t headDimAlign = SFAAAlign(constInfo.headDim, ConstInfo::BUFFER_SIZE_BYTE_32B); // headDim对齐32B
1014+ uint32_t mSplitSize = BASE_BLOCK_MAX_ELEMENT_NUM / headDimAlign; // 单次处理query矩阵的行数（8192/128=64）
1015+ if (mSplitSize > info.mSizeV) {
1016+ mSplitSize = info.mSizeV; // 20
1017+ }
1018+ uint32_t loopCount = (info.mSizeV + mSplitSize - 1) / mSplitSize; // 单个vec核Q预处理的循环次数 // 1
1019+ uint32_t tailSplitSize = info.mSizeV - (loopCount - 1) * mSplitSize;
1020+ for (uint32_t i = 0, dealSize = mSplitSize; i < loopCount; i++) {
1021+ if (i == (loopCount - 1)) {
1022+ dealSize = tailSplitSize;
1023+ }
1024+ DealQueryPreProcessBaseBlock(info, i * mSplitSize, dealSize, headDimAlign, constInfo.headDim);
1025+ }
1026+ 
1027+ if (info.bN2Idx != lastBN2Idx) {
1028+ SetFlag<AscendC::HardEvent::V_MTE2>(SYNC_INPUT_DEQUANT_SCALE_FLAG);
1029+ lastBN2Idx = info.bN2Idx;
1030+ }
1031+}
1032+ 
1033+template <typename SFAAT>
1034+__aicore__ inline void SFAAVectorServiceGqaMsd<SFAAT>::DealQueryPreProcessBaseBlock(const RunInfo &info,
1035+ uint32_t startRow, uint32_t dealRowCount, uint32_t columnCount, uint32_t actualColumnCount)
1036+{
1037+ uint32_t bufferId = info.bN2Idx % constInfo.preLoadNum;
1038+ LocalTensor<T> queryUb = tmpBuff1.Get<T>();
1039+ LocalTensor<T> aFloorUb = tmpBuff2.Get<T>();
1040+ uint64_t qOffset = info.tensorAOffset + (info.mSizeVStart + startRow) * actualColumnCount;
1041+ 
1042+ if (constInfo.kvHeadNum != 1) {
1043+ uint32_t startS1Idx = (info.mSizeVStart + startRow) / constInfo.gSize; // V0: 0 V1: 1
1044+ uint32_t startGOfffset = (info.mSizeVStart + startRow) % constInfo.gSize; // V0: 0 V1: 0
1045+ uint32_t endS1Idx = (info.mSizeVStart + startRow + dealRowCount - 1) / constInfo.gSize; // V0: 0 V1: 1
1046+ uint32_t curStartRow = (info.mSizeVStart + startRow); // V0: 0 V1: 2
1047+ uint32_t curDealRowCount = 0;
1048+ uint32_t ubOffset = 0;
1049+ for (uint32_t curS1idx = startS1Idx; curS1idx <= endS1Idx; curS1idx++) {
1050+ qOffset = info.tensorAOffset + curS1idx * constInfo.qHeadNum * constInfo.headDim +
1051+ startGOfffset * constInfo.headDim;
1052+ if (curS1idx != endS1Idx) {
1053+ curDealRowCount = (curS1idx + 1) * constInfo.gSize - curStartRow;
1054+ } else {
1055+ curDealRowCount = info.mSizeVStart + startRow + dealRowCount - curStartRow;
1056+ }
1057+ ubOffset = (curStartRow - info.mSizeVStart) * columnCount;
1058+ LocalTensor<T> curQueryUb = queryUb[ubOffset];
1059+ CopyAntiqQuery(curQueryUb, qOffset, curDealRowCount, columnCount, actualColumnCount);
1060+ PipeBarrier<PIPE_V>();
1061+ curStartRow += curDealRowCount;
1062+ startGOfffset = 0;
1063+ }
1064+ } else {
1065+ CopyAntiqQuery(queryUb, qOffset, dealRowCount, columnCount, actualColumnCount);
1066+ PipeBarrier<PIPE_V>();
1067+ }
1068+ // mul scale
1069+ VecMulMat(queryUb, antiquantKScaleUb_, queryUb, dealRowCount, columnCount, actualColumnCount);
1070+ PipeBarrier<PIPE_V>();
1071+ 
1072+ // A pre process
1073+ size_t dstOffset = bufferId * constInfo.bmm2ResUbSize * 2 + info.mSizeVStart * columnCount;
1074+ 
1075+ AntiquantMatmulPreProcess(info, queryPreProcessResGm_[dstOffset], aMaxBmm1Ub[bufferId * Q_AMAX_BUF_SIZE],
1076+ queryUb, aFloorUb, startRow, dealRowCount, columnCount, actualColumnCount);
1077+}
1078+ 
1079+template <typename SFAAT>
1080+__aicore__ inline void SFAAVectorServiceGqaMsd<SFAAT>::VecMulMat(LocalTensor<float> dstUb, LocalTensor<float> src0Ub, LocalTensor<float> src1Ub,
1081+ uint32_t dealRowCount, uint32_t columnCount, uint32_t actualColumnCount)
1082+{
1083+ // vec mul by row
1084+ // dstUb[i, j] = src0Ub[j] * src1Ub[i, j],
1085+ // src0Ub:[1, columnCount] src1Ub:[dealRowCount, actualColumnCount] dstUb:[dealRowCount, columnCount]
1086+ if (columnCount < REPEATE_STRIDE_UP_BOUND * FP32_BLOCK_ELEMENT_NUM) { // dstRepStride为0~255,columnCount需要小于2048
1087+ BinaryRepeatParams repeatParams;
1088+ repeatParams.dstBlkStride = 1;
1089+ repeatParams.src0BlkStride = 1;
1090+ repeatParams.src1BlkStride = 1;
1091+ repeatParams.dstRepStride = columnCount / FP32_BLOCK_ELEMENT_NUM;
1092+ repeatParams.src0RepStride = 0;
1093+ repeatParams.src1RepStride = columnCount / FP32_BLOCK_ELEMENT_NUM;
1094+ uint32_t mask = FP32_REPEAT_ELEMENT_NUM;
1095+ uint32_t loopCount = actualColumnCount / mask;
1096+ uint32_t remainCount = actualColumnCount % mask;
1097+ uint32_t offset = 0;
1098+ for (int i = 0; i < loopCount; i++) {
1099+ // offset = i * mask
1100+ Mul(dstUb[offset], src0Ub[offset], src1Ub[offset], mask, dealRowCount, repeatParams);
1101+ offset += mask;
1102+ }
1103+ if (remainCount > 0) {
1104+ // offset = loopCount * mask
1105+ Mul(dstUb[offset], src0Ub[offset], src1Ub[offset], remainCount, dealRowCount, repeatParams);
1106+ }
1107+ } else {
1108+ uint32_t offset = 0;
1109+ for (int i = 0; i < dealRowCount; i++) {
1110+ Mul(dstUb[offset], src0Ub, src1Ub[offset], actualColumnCount);
1111+ offset += columnCount;
1112+ }
1113+ }
1114+}
1115+ 
1116+template <typename SFAAT>
1117+__aicore__ inline void SFAAVectorServiceGqaMsd<SFAAT>::CopyAntiqQuery(LocalTensor<T> &queryCastUb,
1118+ uint64_t qOffset, uint32_t dealRowCount, uint32_t columnCount, uint32_t actualColumnCount)
1119+{
1120+ uint32_t qTypeElementSize = BYTE_BLOCK / sizeof(Q_T); // 16 / 2 = 8
1121+ DataCopyExtParams copyInParams;
1122+ DataCopyPadExtParams<Q_T> copyInPadParams;
1123+ // antiq scale copy in
1124+ copyInParams.blockCount = dealRowCount;
1125+ copyInParams.blockLen = actualColumnCount * sizeof(Q_T);
1126+ copyInParams.srcStride = 0;
1127+ copyInParams.dstStride = (columnCount - actualColumnCount) / qTypeElementSize;
1128+ 
1129+ copyInPadParams.isPad = true;
1130+ copyInPadParams.leftPadding = 0;
1131+ copyInPadParams.rightPadding = (columnCount - actualColumnCount) % qTypeElementSize;
1132+ copyInPadParams.paddingValue = 0;
1133+ 
1134+ LocalTensor<Q_T> inputUb = inputBuf1.Get<Q_T>();
1135+ WaitFlag<AscendC::HardEvent::V_MTE2>(SYNC_INPUT_BUF1_FLAG);
1136+ DataCopyPad(inputUb, queryGm_[qOffset], copyInParams, copyInPadParams);
1137+ SetFlag<AscendC::HardEvent::MTE2_V>(SYNC_INPUT_BUF1_FLAG);
1138+ WaitFlag<AscendC::HardEvent::MTE2_V>(SYNC_INPUT_BUF1_FLAG);
1139+ Cast(queryCastUb, inputUb, RoundMode::CAST_NONE, dealRowCount * columnCount); // 将Query cast到fp32
1140+ SetFlag<AscendC::HardEvent::V_MTE2>(SYNC_INPUT_BUF1_FLAG);
1141+}
1142+ 
1143+template <typename SFAAT>
1144+__aicore__ inline void SFAAVectorServiceGqaMsd<SFAAT>::AntiquantMatmulPreProcess(const RunInfo &info,
1145+ GlobalTensor<KV_T> dstGm, LocalTensor<T> aMaxResUb, LocalTensor<T> srcUb, LocalTensor<T> tmpAFloorUb,
1146+ uint32_t startRow, uint32_t dealRowCount, uint32_t columnCount, uint32_t actualColumnCount)
1147+{
1148+ uint32_t step = info.mSize * columnCount;
1149+ uint32_t baseOffset = startRow * columnCount;
1150+ uint32_t calcSize = dealRowCount * columnCount;
1151+ 
1152+ LocalTensor<T> tmpAMaxRes = aMaxResUb[startRow * BLOCK_ELEMENT_NUM];
1153+ AbsRowMax(tmpAMaxRes, srcUb, tmpAFloorUb, dealRowCount, columnCount, actualColumnCount);
1154+ PipeBarrier<PIPE_V>();
1155+ 
1156+ // 128/(1.001*Amax)*A
1157+ Duplicate(tmpAFloorUb, antiqCoeff1, dealRowCount * BLOCK_ELEMENT_NUM);
1158+ PipeBarrier<PIPE_V>();
1159+ Div(tmpAFloorUb, tmpAFloorUb, tmpAMaxRes, dealRowCount * BLOCK_ELEMENT_NUM);
1160+ PipeBarrier<PIPE_V>();
1161+ RowMuls<T>(srcUb, srcUb, tmpAFloorUb, dealRowCount, columnCount, actualColumnCount);
1162+ PipeBarrier<PIPE_V>();
1163+ 
1164+ // step4: 取出Qtmp整数部分
1165+ Cast(tmpAFloorUb, srcUb, RoundMode::CAST_ROUND, calcSize);
1166+ PipeBarrier<PIPE_V>();
1167+ LocalTensor<half> tmpAFloorUbFp16 = tmpAFloorUb.template ReinterpretCast<half>();
1168+ tmpAFloorUbFp16.SetSize(tmpAFloorUb.GetSize());
1169+ Cast(tmpAFloorUbFp16, tmpAFloorUb, RoundMode::CAST_ROUND, calcSize);
1170+ PipeBarrier<PIPE_V>();
1171+ // step5: 将Qtmp转成fp16
1172+ LocalTensor<half> srcUbFp16 = srcUb.template ReinterpretCast<half>();
1173+ srcUbFp16.SetSize(srcUb.GetSize());
1174+ Cast(srcUbFp16, srcUb, RoundMode::CAST_ROUND, calcSize);
1175+ PipeBarrier<PIPE_V>();
1176+ 
1177+ for (uint32_t i = 0; i < msdIterNum; i++) {
1178+ AntiquantAIterExpand(dstGm, srcUbFp16, tmpAFloorUbFp16, calcSize, (i == 0 ? true : false), step * i + baseOffset);
1179+ }
1180+}
1181+ 
1182+template <typename SFAAT>
1183+__aicore__ inline void SFAAVectorServiceGqaMsd<SFAAT>::AbsRowMax(LocalTensor<T> &tmpAMaxRes, LocalTensor<T> &srcUb,
1184+ LocalTensor<T> tmpAUb, uint32_t dealRowCount,
1185+ uint32_t columnCount, uint32_t actualColumnCount)
1186+{
1187+ Abs(tmpAUb, srcUb, dealRowCount * columnCount);
1188+ PipeBarrier<PIPE_V>();
1189+ LocalTensor<T> tmpRowMaxUb = tmpBuff3.Get<T>();
1190+ RowMaxForLongColumnCount(tmpRowMaxUb, tmpAUb, dealRowCount, columnCount, actualColumnCount);
1191+ PipeBarrier<PIPE_V>();
1192+ Brcb(tmpAMaxRes, tmpRowMaxUb, (dealRowCount + 7) / 8, {1, 8});
1193+}
1194+ 
1195+template <typename SFAAT>
1196+__aicore__ inline void SFAAVectorServiceGqaMsd<SFAAT>::RowMaxForLongColumnCount(LocalTensor<float> &dstUb, LocalTensor<float> srcUb,
1197+ uint32_t dealRowCount, uint32_t columnCount,
1198+ uint32_t actualColumnCount)
1199+{
1200+ // max by row, 按行求最大值
1201+ // dstUb[i] = max(srcUb[i, :])
1202+ // src0Ub:[dealRowCount, columnCount] dstUb:[1, dealRowCount]
1203+ uint32_t newColumnCount = columnCount;
1204+ uint32_t newActualColumnCount = actualColumnCount;
1205+ if (columnCount >= REPEATE_STRIDE_UP_BOUND * FP32_BLOCK_ELEMENT_NUM) {
1206+ uint32_t split = GetMinPowerTwo(actualColumnCount);
1207+ split = split >> 1;
1208+ 
1209+ // deal tail
1210+ uint32_t offset = 0;
1211+ for (uint32_t i = 0; i < dealRowCount; i++) {
1212+ Max(srcUb[offset], srcUb[offset], srcUb[offset + split], actualColumnCount - split);
1213+ offset += columnCount;
1214+ }
1215+ PipeBarrier<PIPE_V>();
1216+ 
1217+ uint32_t validLen = split;
1218+ while (validLen > ConstInfo::BUFFER_SIZE_BYTE_1K) {
1219+ uint32_t copyLen = validLen / 2;
1220+ 
1221+ offset = 0;
1222+ for (uint32_t i = 0; i < dealRowCount; i++) {
1223+ Max(srcUb[offset], srcUb[offset], srcUb[offset + copyLen], copyLen);
1224+ offset += columnCount;
1225+ }
1226+ PipeBarrier<PIPE_V>();
1227+ 
1228+ validLen = copyLen;
1229+ }
1230+ 
1231+ for (uint32_t i = 0; i < dealRowCount; i++) {
1232+ DataCopy(srcUb[i * validLen], srcUb[i * columnCount], validLen);
1233+ PipeBarrier<PIPE_V>();
1234+ }
1235+ 
1236+ newColumnCount = validLen;
1237+ newActualColumnCount = validLen;
1238+ }
1239+ 
1240+ RowMax(dstUb, srcUb, dealRowCount, newColumnCount, newActualColumnCount);
1241+}
1242+ 
1243+template <typename SFAAT>
1244+__aicore__ inline void SFAAVectorServiceGqaMsd<SFAAT>::RowMax(LocalTensor<float> &dstUb, LocalTensor<float> &srcUb,
1245+ uint32_t dealRowCount, uint32_t columnCount, uint32_t actualColumnCount)
1246+{
1247+ // max by row, 按行求最大值
1248+ // dstUb[i] = max(srcUb[i, :])
1249+ // src0Ub:[dealRowCount, columnCount] dstUb:[1, dealRowCount]
1250+ uint32_t dtypeMask = FP32_REPEAT_ELEMENT_NUM;
1251+ uint32_t blockCount = actualColumnCount / dtypeMask;
1252+ uint32_t remain = actualColumnCount % dtypeMask;
1253+ 
1254+ BinaryRepeatParams repeatParamsMax;
1255+ repeatParamsMax.src0BlkStride = 1;
1256+ repeatParamsMax.src1BlkStride = 1;
1257+ repeatParamsMax.dstBlkStride = 1;
1258+ repeatParamsMax.src0RepStride = columnCount / FP32_BLOCK_ELEMENT_NUM;
1259+ repeatParamsMax.src1RepStride = columnCount / FP32_BLOCK_ELEMENT_NUM;
1260+ repeatParamsMax.dstRepStride = columnCount / FP32_BLOCK_ELEMENT_NUM;
1261+ if (blockCount > 0 && remain > 0) {
1262+ Max(srcUb, srcUb, srcUb[blockCount * dtypeMask], remain, dealRowCount, repeatParamsMax);
1263+ PipeBarrier<PIPE_V>();
1264+ }
1265+ 
1266+ for (uint32_t loopCount = blockCount / 2; loopCount > 0; loopCount = blockCount / 2) {
1267+ blockCount = (blockCount + 1) / 2;
1268+ for (uint32_t j = 0; j < loopCount; j++) {
1269+ Max(srcUb[j * dtypeMask], srcUb[j * dtypeMask], srcUb[(j + blockCount) * dtypeMask], dtypeMask,
1270+ dealRowCount, repeatParamsMax);
1271+ }
1272+ PipeBarrier<PIPE_V>();
1273+ }
1274+ 
1275+ WholeReduceMax(dstUb, srcUb, (actualColumnCount < dtypeMask) ? actualColumnCount : dtypeMask, dealRowCount, 1, 1,
1276+ columnCount / FP32_BLOCK_ELEMENT_NUM, ReduceOrder::ORDER_ONLY_VALUE);
1277+}
1278+ 
1279+template <typename SFAAT>
1280+__aicore__ inline void SFAAVectorServiceGqaMsd<SFAAT>::AntiquantAIterExpand(GlobalTensor<KV_T> &dstGm, LocalTensor<half> &tmpA1,
1281+ LocalTensor<half> &tmpA2, uint32_t calcSize,
1282+ bool isFirst, uint64_t outOffset)
1283+{
1284+ if (!isFirst) {
1285+ Sub(tmpA2, tmpA1, tmpA2, calcSize);
1286+ PipeBarrier<PIPE_V>();
1287+ Muls(tmpA2, tmpA2, antiquantExpandCoeff, calcSize);
1288+ PipeBarrier<PIPE_V>();
1289+ }
1290+ LocalTensor<KV_T> aResOutUbI8 = outputBuf1.Get<KV_T>();
1291+ WaitFlag<AscendC::HardEvent::MTE3_V>(SYNC_OUTPUT_BUF1_FLAG);
1292+ Cast(aResOutUbI8, tmpA2, RoundMode::CAST_ROUND, calcSize);
1293+ PipeBarrier<PIPE_V>();
1294+ SetFlag<AscendC::HardEvent::V_MTE3>(SYNC_OUTPUT_BUF1_FLAG);
1295+ WaitFlag<AscendC::HardEvent::V_MTE3>(SYNC_OUTPUT_BUF1_FLAG);
1296+ DataCopy(dstGm[outOffset], aResOutUbI8, calcSize);
1297+ SetFlag<AscendC::HardEvent::MTE3_V>(SYNC_OUTPUT_BUF1_FLAG);
1298+}
1299+ 
1300+template <typename SFAAT>
1301+__aicore__ inline int64_t SFAAVectorServiceGqaMsd<SFAAT>::MergeKv(const RunInfo &runInfo, int64_t s2GmStartOffset,
1302+ int64_t s2GmLimit, int64_t topkGmBaseOffset, bool isValue, int64_t mergeMte3Idx, int64_t s2GmStartOffset4Merge)
1303+{
1304+ int64_t mte2Size = 0;
1305+ int64_t mte3Size = 0;
1306+ int64_t s2IdxArray0 = -1;
1307+ int64_t s2IdxArray1 = -1;
1308+ bool needWaitMte3ToMte2 = true;
1309+ bool mask = false;
1310+ 
1311+ for (int64_t s2GmOffsetArray = s2GmStartOffset; s2GmOffsetArray < s2GmLimit; s2GmOffsetArray += 2 * constInfo.sparseBlockSize) {
1312+ if (needWaitMte3ToMte2) {
1313+ WaitFlag<AscendC::HardEvent::MTE3_MTE2>(SYNC_INPUT_BUF2_FLAG + mergeMte3Idx % 2);
1314+ needWaitMte3ToMte2 = false;
1315+ }
1316+ GetRealS2Idx(s2GmOffsetArray, s2IdxArray0, topkGmBaseOffset, runInfo);
1317+ if (unlikely(s2IdxArray0 < 0)) { // 当前sparseblock已经超过了sparseCount的范围
1318+ CopyOutMrgeResult(mte2Size, mte3Size, s2GmStartOffset4Merge, mergeMte3Idx, runInfo, isValue, needWaitMte3ToMte2);
1319+ mergeMte3Idx++;
1320+ break;
1321+ }
1322+ GetRealS2Idx(s2GmOffsetArray + constInfo.sparseBlockSize, s2IdxArray1, topkGmBaseOffset, runInfo);
1323+ CopyInKv(mte2Size, mte3Size, mergeMte3Idx, s2IdxArray0, s2IdxArray1, runInfo, mask, s2GmOffsetArray, s2GmLimit);
1324+ if ((mte2Size - mte3Size + 2 * constInfo.sparseBlockSize > 32) || // TODO，这里可以支持到128，但需要循环拷出
1325+ s2GmOffsetArray + 2 * constInfo.sparseBlockSize >= s2GmLimit) {
1326+ CopyOutMrgeResult(mte2Size, mte3Size, s2GmStartOffset4Merge, mergeMte3Idx, runInfo, isValue, needWaitMte3ToMte2);
1327+ mte3Size = mte2Size;
1328+ mergeMte3Idx++;
1329+ }
1330+ }
1331+ return mte2Size;
1332+}
1333+ 
1334+template <typename SFAAT>
1335+__aicore__ inline void SFAAVectorServiceGqaMsd<SFAAT>::ProcessVec1Msd(const RunInfo &info)
1336+{
1337+ uint32_t nBufferLoopTimes = (info.actMBaseSize + constInfo.nBufferMBaseSize - 1) / constInfo.nBufferMBaseSize;
1338+ uint32_t nBufferTail = info.actMBaseSize - (nBufferLoopTimes - 1) * constInfo.nBufferMBaseSize;
1339+ for (uint32_t i = 0; i < nBufferLoopTimes; i++) {
1340+ MSplitInfo mSplitInfo;
1341+ mSplitInfo.nBufferIdx = i;
1342+ mSplitInfo.nBufferStartM = i * constInfo.nBufferMBaseSize;
1343+ mSplitInfo.nBufferDealM = (i + 1 != nBufferLoopTimes) ? constInfo.nBufferMBaseSize : nBufferTail;
1344+ 
1345+ // mSplitInfo.vecDealM = (mSplitInfo.nBufferDealM <= 16) ? mSplitInfo.nBufferDealM :
1346+ // (((mSplitInfo.nBufferDealM + 15) / 16 + 1) / 2 * 16);
1347+ mSplitInfo.vecDealM = info.mSizeV;
1348+ mSplitInfo.vecStartM = 0;
1349+ if (GetBlockIdx() % 2 == 1) {
1350+ mSplitInfo.vecStartM = info.mSize - mSplitInfo.vecDealM;
1351+ }
1352+ 
1353+ CrossCoreWaitFlag(constInfo.syncC1V1);
1354+ // vec1 compute
1355+ ProcessVec1SingleBuf(info, mSplitInfo);
1356+ CrossCoreSetFlag<ConstInfo::SFAA_SYNC_MODE2, PIPE_MTE3>(constInfo.syncV1C2);
1357+ // move lse for flash decode
1358+ if constexpr (IS_META) {
1359+ if (info.s2Idx == info.curSInnerLoopTimes - 1) {
1360+ if (info.tndIsS2SplitCore) {
1361+ if (FLASH_DECODE) {
1362+ uint32_t outIdx = info.loop % (constInfo.preLoadNum);
1363+ auto sumTensor = softmaxSumUb[outIdx * SOFTMAX_TMP_BUFFER_OFFSET];
1364+ auto maxTensor = softmaxMaxUb[outIdx * SOFTMAX_TMP_BUFFER_OFFSET];
1365+ ComputeLogSumExpAndCopyToGm(info, mSplitInfo, sumTensor, maxTensor);
1366+ }
1367+ }
1368+ }
1369+ }
1370+ }
1371+}
1372+ 
1373+template <typename SFAAT>
1374+__aicore__ inline uint64_t SFAAVectorServiceGqaMsd<SFAAT>::CalcAccumOffset(uint32_t bN2Idx, uint32_t gS1Idx)
1375+{
1376+ uint64_t accumTmpOutNum = 0;
1377+ uint32_t taskId = 0;
1378+ 
1379+ if constexpr (IS_META) {
1380+ const uint32_t *bN2IdxOfFdHead = metaDataPtr->fdRes.bN2IdxOfFdHead;
1381+ const uint32_t *gS1IdxOfFdHead = metaDataPtr->fdRes.gS1IdxOfFdHead;
1382+ const uint32_t *s2SplitNumOfFdHead = metaDataPtr->fdRes.s2SplitNumOfFdHead;
1383+ uint32_t usedCoreNum = metaDataPtr->usedCoreNum;
1384+ while (taskId < usedCoreNum && (bN2IdxOfFdHead[taskId] != bN2Idx || gS1IdxOfFdHead[taskId] * constInfo.mBaseSize != gS1Idx)) { // 考虑在tiling阶段直接算出accumOut的偏置，则可以省略CalcAccumOffset()
1385+ accumTmpOutNum += s2SplitNumOfFdHead[taskId]; // 计算前面的workspace数
1386+ taskId++;
1387+ }
1388+ } else {
1389+ return 0;
1390+ }
1391+ return accumTmpOutNum;
1392+}
1393+ 
1394+template <typename SFAAT>
1395+__aicore__ inline void SFAAVectorServiceGqaMsd<SFAAT>::ProcessVec2SingleBuf(const RunInfo &info,
1396+ const MSplitInfo &mSplitInfo)
1397+{
1398+ if (mSplitInfo.vecDealM == 0) {
1399+ return;
1400+ }
1401+ 
1402+ uint32_t gPreSplitSize = BASE_BLOCK_MAX_ELEMENT_NUM / constInfo.headDim; // 64
1403+ if (gPreSplitSize > mSplitInfo.vecDealM) {
1404+ gPreSplitSize = mSplitInfo.vecDealM;
1405+ }
1406+ uint32_t loopCount = (mSplitInfo.vecDealM + gPreSplitSize - 1) / gPreSplitSize; // 1
1407+ uint32_t tailSplitSize = mSplitInfo.vecDealM - (loopCount - 1) * gPreSplitSize; // 20
1408+ 
1409+ for (uint32_t i = 0, dealSize = gPreSplitSize; i < loopCount; i++) {
1410+ if (i == (loopCount - 1)) {
1411+ dealSize = tailSplitSize;
1412+ }
1413+ DealBmm2ResBaseBlock(info, mSplitInfo, i * gPreSplitSize, dealSize, constInfo.headDim, constInfo.headDim);
1414+ pingpongFlag ^= 1; // pingpong 0 1切换
1415+ }
1416+}
1417+ 
1418+template <typename SFAAT>
1419+__aicore__ inline void SFAAVectorServiceGqaMsd<SFAAT>::AntiquantMM2ResCombine(const RunInfo &info,
1420+ LocalTensor<MM2_OUT_T> bmmResUb, GlobalTensor<MM2_OUT_T> srcGm, uint32_t startRow, uint32_t dealRowCount,
1421+ uint32_t columnCount, uint32_t actualColumnCount)
1422+{
1423+ uint32_t baseOffset = startRow * columnCount;
1424+ uint32_t copySize = dealRowCount * columnCount;
1425+ LocalTensor<MM2_OUT_T> tmpCInt = inputBuf1.Get<MM2_OUT_T>();
1426+ WaitFlag<AscendC::HardEvent::V_MTE2>(SYNC_INPUT_BUF1_FLAG);
1427+ DataCopy(tmpCInt, srcGm[baseOffset], copySize);
1428+ SetFlag<AscendC::HardEvent::MTE2_V>(SYNC_INPUT_BUF1_FLAG);
1429+ WaitFlag<AscendC::HardEvent::MTE2_V>(SYNC_INPUT_BUF1_FLAG);
1430+ DataCopy(bmmResUb, tmpCInt, copySize);
1431+ SetFlag<AscendC::HardEvent::V_MTE2>(SYNC_INPUT_BUF1_FLAG);
1432+}
1433+ 
1434+template <typename SFAAT>
1435+__aicore__ inline void SFAAVectorServiceGqaMsd<SFAAT>::DealBmm2ResBaseBlock(const RunInfo &info, const MSplitInfo &mSplitInfo,
1436+ uint32_t startRow, uint32_t dealRowCount,
1437+ uint32_t columnCount, uint32_t actualColumnCount)
1438+{
1439+ uint32_t vec2ComputeSize = dealRowCount * columnCount;
1440+#ifdef SFACV_SOFTMAX_WITHOUT_BRC_PRELOAD
1441+ uint32_t baseOffset = startRow;
1442+#else
1443+ uint32_t baseOffset = startRow * BLOCK_ELEMENT_NUM;
1444+#endif
1445+ size_t batchBase = 0;
1446+ uint64_t inOutBaseOffset = (mSplitInfo.vecStartM + startRow) * columnCount;
1447+ uint64_t srcGmOffset = (info.loop % constInfo.preLoadNum) * constInfo.bmm2ResUbSize + inOutBaseOffset;
1448+ LocalTensor<half> bmm2ResUb = tmpBuff1.Get<half>();
1449+ bmm2ResUb.SetSize(vec2ComputeSize);
1450+ AntiquantMM2ResCombine(info, bmm2ResUb, mm2ResGm[srcGmOffset], startRow, dealRowCount, columnCount, actualColumnCount);
1451+ // 除第一个循环外，均需要更新中间计算结果
1452+ if (!info.isFirstSInnerLoop) {
1453+ //step2: 将softmaxExpUb转换为FP16
1454+#ifdef SFACV_SOFTMAX_WITHOUT_BRC_PRELOAD
1455+ LocalTensor<half> tmpSoftmaxFp16 = tmpBuff3.Get<half>();
1456+ Cast(tmpSoftmaxFp16, softmaxExpUb[(info.loop % constInfo.preLoadNum) * SOFTMAX_TMP_BUFFER_OFFSET + baseOffset],
1457+ AscendC::RoundMode::CAST_ROUND, dealRowCount);
1458+ PipeBarrier<PIPE_V>();
1459+ LocalTensor<half> tmpExpBrcbResUb = tmpBuff2.Get<half>();
1460+ // repeatTime:指令迭代次数，每次迭代完成8个datablock的数据收集
1461+ // repeatParams:{单次迭代内，矢量目的操作数不同datablock间地址步长, 相邻迭代间，矢量目的操作数相同datablock地址步长}
1462+ Brcb(tmpExpBrcbResUb, tmpSoftmaxFp16, (mSplitInfo.vecDealM + 7) / 8, {1, 8});
1463+ // step3: Oi = (Oi - 1)*tmpSoftmaxFp16 + Otmp
1464+ PipeBarrier<PIPE_V>();
1465+ RowMuls<half>(vec2ResUb, vec2ResUb, tmpExpBrcbResUb, dealRowCount, columnCount, actualColumnCount);
1466+#else
1467+ LocalTensor<T> tmpSoftmaxFp32 = tmpBuff3.Get<T>();
1468+ DataCopyParams dataCopyParams;
1469+ dataCopyParams.blockCount = dealRowCount;
1470+ dataCopyParams.blockLen = 1;
1471+ dataCopyParams.srcStride = 0;
1472+ dataCopyParams.dstStride = 1;
1473+ DataCopy(tmpSoftmaxFp32, softmaxExpUb[(info.loop % constInfo.preLoadNum) * SOFTMAX_TMP_BUFFER_OFFSET + baseOffset], dataCopyParams);
1474+ DataCopy(tmpSoftmaxFp32[BLOCK_ELEMENT_NUM], softmaxExpUb[(info.loop % constInfo.preLoadNum) * SOFTMAX_TMP_BUFFER_OFFSET + baseOffset],
1475+ dataCopyParams);
1476+ LocalTensor<half> tmpSoftmaxFp16 = tmpBuff3.Get<half>();
1477+ PipeBarrier<PIPE_V>();
1478+ Cast(tmpSoftmaxFp16, tmpSoftmaxFp32, AscendC::RoundMode::CAST_ROUND, dealRowCount * BLOCK_ELEMENT_NUM * 2);
1479+ PipeBarrier<PIPE_V>();
1480+ RowMuls(vec2ResUb, vec2ResUb, tmpSoftmaxFp16,
1481+ dealRowCount, columnCount, actualColumnCount);
1482+#endif
1483+ PipeBarrier<PIPE_V>();
1484+ Add(bmm2ResUb, bmm2ResUb, vec2ResUb, vec2ComputeSize);
1485+ }
1486+ // 最后一次输出计算结果，否则将中间结果暂存至workspace
1487+ if (info.s2Idx + 1 == info.curSInnerLoopTimes) {
1488+ PipeBarrier<PIPE_V>();
1489+ LocalTensor<T> bmm2ResUbFp32 = tmpBuff2.Get<T>();
1490+ Cast(bmm2ResUbFp32, bmm2ResUb, RoundMode::CAST_NONE, vec2ComputeSize);
1491+ PipeBarrier<PIPE_V>();
1492+ uint32_t idx = info.loop % constInfo.preLoadNum;
1493+#ifdef SFACV_SOFTMAX_WITHOUT_BRC_PRELOAD
1494+ LocalTensor<T> tmpSumBrcbResUb = tmpBuff1.Get<T>();
1495+ Brcb(tmpSumBrcbResUb, softmaxSumUb[(info.loop % constInfo.preLoadNum) * SOFTMAX_TMP_BUFFER_OFFSET + baseOffset],
1496+ (mSplitInfo.vecDealM + 7) / 8, {1, 8});
1497+ PipeBarrier<PIPE_V>();
1498+ RowDivs(bmm2ResUbFp32, bmm2ResUbFp32, tmpSumBrcbResUb, dealRowCount, columnCount, actualColumnCount);
1499+#else
1500+ RowDivs(bmm2ResUbFp32, bmm2ResUbFp32, softmaxSumUb[(info.loop % constInfo.preLoadNum) * SOFTMAX_TMP_BUFFER_OFFSET + baseOffset],
1501+ dealRowCount, columnCount, actualColumnCount);
1502+#endif
1503+ pipe_barrier(PIPE_V);
1504+ Muls(bmm2ResUbFp32, bmm2ResUbFp32, scaleC2, vec2ComputeSize);
1505+ PipeBarrier<PIPE_V>();
1506+ CopyAntiquantScale(antiquantVScaleUb_, valueDequantScaleGm_, info.n2Idx * constInfo.headDim);
1507+ PipeBarrier<PIPE_V>();
1508+ // ScaleV * bmm2res
1509+ VecMulMat(bmm2ResUbFp32, antiquantVScaleUb_, bmm2ResUbFp32, dealRowCount, columnCount, actualColumnCount);
1510+ PipeBarrier<PIPE_V>();
1511+ Bmm2ResCopyOut(info, bmm2ResUbFp32, mSplitInfo.vecStartM + startRow, dealRowCount, columnCount, actualColumnCount);
1512+ } else {
1513+ PipeBarrier<PIPE_V>();
1514+ DataCopy(vec2ResUb, bmm2ResUb, dealRowCount * columnCount);
1515+ }
1516+}
1517+ 
1518+template <typename SFAAT>
1519+__aicore__ inline void SFAAVectorServiceGqaMsd<SFAAT>::CopyAntiquantScale(LocalTensor<T> &castUb, GlobalTensor<T> srcGm,
1520+ uint64_t offset)
1521+{
1522+ uint32_t qTypeElementSize = BYTE_BLOCK / sizeof(T);
1523+ uint32_t headDimAlign = SFAAAlign(constInfo.headDim, ConstInfo::BUFFER_SIZE_BYTE_32B) / sizeof(KV_T);
1524+ DataCopyExtParams copyInParams;
1525+ DataCopyPadExtParams<T> copyInPadParams;
1526+ // antiq scale copy in
1527+ copyInParams.blockCount = 1;
1528+ copyInParams.blockLen = constInfo.headDim * sizeof(T);
1529+ copyInParams.srcStride = 0;
1530+ copyInParams.dstStride = (headDimAlign - constInfo.headDim) / qTypeElementSize;
1531+ 
1532+ copyInPadParams.isPad = true;
1533+ copyInPadParams.leftPadding = 0;
1534+ copyInPadParams.rightPadding = (headDimAlign - constInfo.headDim) % qTypeElementSize;
1535+ copyInPadParams.paddingValue = 0;
1536+ 
1537+ WaitFlag<AscendC::HardEvent::V_MTE2>(SYNC_INPUT_DEQUANT_SCALE_FLAG);
1538+ DataCopyPad(castUb, srcGm[offset], copyInParams, copyInPadParams);
1539+ SetFlag<AscendC::HardEvent::MTE2_V>(SYNC_INPUT_DEQUANT_SCALE_FLAG);
1540+ WaitFlag<AscendC::HardEvent::MTE2_V>(SYNC_INPUT_DEQUANT_SCALE_FLAG);
1541+ SetFlag<AscendC::HardEvent::V_MTE2>(SYNC_INPUT_DEQUANT_SCALE_FLAG);
1542+}
1543+ 
1544+template <typename SFAAT> __aicore__ inline void SFAAVectorServiceGqaMsd<SFAAT>::ProcessVec2Msd(const RunInfo &info)
1545+{
1546+ uint32_t nBufferLoopTimes = (info.actMBaseSize + constInfo.nBufferMBaseSize - 1) / constInfo.nBufferMBaseSize;
1547+ uint32_t nBufferTail = info.actMBaseSize - (nBufferLoopTimes - 1) * constInfo.nBufferMBaseSize;
1548+ for (uint32_t i = 0; i < nBufferLoopTimes; i++) {
1549+ MSplitInfo mSplitInfo;
1550+ mSplitInfo.nBufferIdx = i;
1551+ mSplitInfo.nBufferStartM = i * constInfo.nBufferMBaseSize;
1552+ mSplitInfo.nBufferDealM = (i + 1 != nBufferLoopTimes) ? constInfo.nBufferMBaseSize : nBufferTail;
1553+ 
1554+ mSplitInfo.vecDealM = info.mSizeV;
1555+ mSplitInfo.vecStartM = 0;
1556+ if (GetBlockIdx() % 2 == 1) {
1557+ mSplitInfo.vecStartM = info.mSize - mSplitInfo.vecDealM;
1558+ }
1559+ CrossCoreWaitFlag(constInfo.syncC2V2);
1560+ ProcessVec2SingleBuf(info, mSplitInfo);
1561+ }
1562+}
1563+ 
1564+template <typename SFAAT>
1565+__aicore__ inline void SFAAVectorServiceGqaMsd<SFAAT>::ProcessVec2Inner(const RunInfo &info,
1566+ const MSplitInfo &mSplitInfo,
1567+ uint32_t mStartRow, uint32_t mDealSize)
1568+{
1569+ uint32_t mSplitSize = BASE_BLOCK_MAX_ELEMENT_NUM / constInfo.headDim;
1570+ if (mSplitSize > mDealSize) {
1571+ mSplitSize = mDealSize;
1572+ }
1573+ 
1574+ uint32_t loopCount = (mDealSize + mSplitSize - 1) / mSplitSize;
1575+ uint32_t tailSplitSize = mDealSize - (loopCount - 1) * mSplitSize;
1576+ for (uint32_t i = 0, dealSize = mSplitSize; i < loopCount; i++) {
1577+ if (i == (loopCount - 1)) {
1578+ dealSize = tailSplitSize;
1579+ }
1580+ DealBmm2ResBaseBlock(info, mSplitInfo, i * mSplitSize + mStartRow, dealSize,
1581+ constInfo.headDim, constInfo.headDim);
1582+ pingpongFlag ^= 1; // pingpong 0 1切换
1583+ }
1584+}
1585+ 
1586+ 
1587+template <typename SFAAT>
1588+__aicore__ inline void SFAAVectorServiceGqaMsd<SFAAT>::GetConfusionTransposeTiling(
1589+ int64_t numR, int64_t numC, const uint32_t stackBufferSize, const uint32_t typeSize,
1590+ ConfusionTransposeTiling &tiling)
1591+{
1592+ (void)stackBufferSize;
1593+ uint32_t blockSize = ONE_BLK_SIZE / typeSize;
1594+ uint32_t height = numC;
1595+ uint32_t width = numR;
1596+ uint32_t highBlock = height / BLOCK_CUBE;
1597+ uint32_t stride = height * blockSize * typeSize / ONE_BLK_SIZE;
1598+ uint32_t repeat = width / blockSize;
1599+ 
1600+ tiling.param0 = blockSize;
1601+ tiling.param1 = height;
1602+ tiling.param2 = width;
1603+ tiling.param3 = highBlock;
1604+ tiling.param4 = stride;
1605+ tiling.param5 = repeat;
1606+}
1607+ 
1608+template <typename SFAAT>
1609+__aicore__ inline void
1610+SFAAVectorServiceGqaMsd<SFAAT>::Bmm2FDDataCopyOut(const RunInfo &info, LocalTensor<T> &bmm2ResUb,
1611+ uint32_t wsMStart, uint32_t dealRowCount, uint32_t columnCount,
1612+ uint32_t actualColumnCount)
1613+{
1614+ LocalTensor<T> tmp = outputBuf1.Get<T>();
1615+ WaitFlag<AscendC::HardEvent::MTE3_V>(SYNC_OUTPUT_BUF1_FLAG);
1616+ DataCopy(tmp, bmm2ResUb, columnCount * dealRowCount);
1617+ SetFlag<AscendC::HardEvent::V_MTE3>(SYNC_OUTPUT_BUF1_FLAG);
1618+ WaitFlag<AscendC::HardEvent::V_MTE3>(SYNC_OUTPUT_BUF1_FLAG);
1619+ uint64_t accumTmpOutNum = CalcAccumOffset(info.bN2Idx, info.gS1Idx);
1620+ uint64_t offset = accumTmpOutNum * constInfo.mBaseSize * constInfo.headDim + // taskoffset
1621+ // 份数offset
1622+ info.tndCoreStartKVSplitPos * constInfo.mBaseSize * constInfo.headDim +
1623+ wsMStart * actualColumnCount; // m轴offset
1624+ GlobalTensor<T> dst = accumOutGm[offset];
1625+ if (info.actualSingleProcessSInnerSize != 0) {
1626+ DataCopyExtParams dataCopyParams;
1627+ dataCopyParams.blockCount = dealRowCount;
1628+ dataCopyParams.blockLen = actualColumnCount * sizeof(T);
1629+ dataCopyParams.srcStride = (columnCount - actualColumnCount) / (BYTE_BLOCK / sizeof(T));
1630+ dataCopyParams.dstStride = 0;
1631+ DataCopyPad(dst, tmp, dataCopyParams);
1632+ } else {
1633+ matmul::InitOutput<T>(dst, dealRowCount * actualColumnCount, ConstInfo::FLOAT_ZERO);
1634+ }
1635+ SetFlag<AscendC::HardEvent::MTE3_V>(SYNC_OUTPUT_BUF1_FLAG);
1636+}
1637+ 
1638+template <typename SFAAT>
1639+__aicore__ inline void
1640+SFAAVectorServiceGqaMsd<SFAAT>::Bmm2DataCopyOutTrans(const RunInfo &info, LocalTensor<OUT_T> &attenOutUb,
1641+ uint32_t wsMStart, uint32_t dealRowCount,
1642+ uint32_t columnCount, uint32_t actualColumnCount)
1643+{
1644+ uint32_t s1StartIdx = wsMStart / constInfo.gSize;
1645+ uint32_t startGOffset = wsMStart % constInfo.gSize;
1646+ uint32_t s1EndIdx = SfaaCeilDiv(wsMStart + dealRowCount, static_cast<uint32_t>(constInfo.gSize)) - 1;
1647+ uint32_t curStartRow = wsMStart;
1648+ uint32_t curDealRowCount = 0;
1649+ uint32_t ubOffset = 0;
1650+ for (uint32_t curS1idx = s1StartIdx; curS1idx <= s1EndIdx; curS1idx++) {
1651+ uint32_t outOffset = info.attenOutOffset + curS1idx * constInfo.qHeadNum * constInfo.headDim + startGOffset * constInfo.headDim;
1652+ if (curS1idx != s1EndIdx) {
1653+ curDealRowCount = (curS1idx + 1) * constInfo.gSize -curStartRow;
1654+ } else {
1655+ curDealRowCount = wsMStart + dealRowCount - curStartRow;
1656+ }
1657+ ubOffset = (curStartRow - wsMStart) * columnCount;
1658+ if (unlikely((info.nextTokensPerBatch < 0) && curS1idx < (-info.nextTokensPerBatch))) {
1659+ matmul::InitOutput<OUT_T>(attentionOutGm[outOffset], constInfo.gSize * constInfo.headDim, 0);
1660+ curStartRow += curDealRowCount;
1661+ startGOffset = 0;
1662+ continue;
1663+ }
1664+ LocalTensor<OUT_T> curAttenOutUb = attenOutUb[ubOffset];
1665+ DataCopyExtParams dataCopyParams;
1666+ dataCopyParams.blockCount = curDealRowCount;
1667+ dataCopyParams.blockLen = actualColumnCount * sizeof(OUT_T);
1668+ dataCopyParams.srcStride = (columnCount - actualColumnCount) / (BYTE_BLOCK / sizeof(OUT_T));
1669+ dataCopyParams.dstStride = 0;
1670+ DataCopyPad(attentionOutGm[outOffset], curAttenOutUb, dataCopyParams);
1671+ curStartRow += curDealRowCount;
1672+ startGOffset = 0;
1673+ }
1674+}
1675+ 
1676+template <typename SFAAT>
1677+__aicore__ inline void
1678+SFAAVectorServiceGqaMsd<SFAAT>::Bmm2CastAndCopyOut(const RunInfo &info, LocalTensor<T> &bmm2ResUb,
1679+ uint32_t wsMStart, uint32_t dealRowCount, uint32_t columnCount,
1680+ uint32_t actualColumnCount)
1681+{
1682+ LocalTensor<OUT_T> tmpBmm2ResCastTensor = outputBuf1.Get<OUT_T>();
1683+ WaitFlag<AscendC::HardEvent::MTE3_V>(SYNC_OUTPUT_BUF1_FLAG);
1684+ if constexpr (IsSameType<OUT_T, bfloat16_t>::value) { // bf16 采取四舍六入五成双模式
1685+ Cast(tmpBmm2ResCastTensor, bmm2ResUb, AscendC::RoundMode::CAST_RINT, dealRowCount * columnCount);
1686+ } else {
1687+ Cast(tmpBmm2ResCastTensor, bmm2ResUb, AscendC::RoundMode::CAST_ROUND, dealRowCount * columnCount);
1688+ }
1689+ SetFlag<AscendC::HardEvent::V_MTE3>(SYNC_OUTPUT_BUF1_FLAG);
1690+ WaitFlag<AscendC::HardEvent::V_MTE3>(SYNC_OUTPUT_BUF1_FLAG);
1691+ Bmm2DataCopyOutTrans(info, tmpBmm2ResCastTensor, wsMStart, dealRowCount, columnCount, actualColumnCount);
1692+ SetFlag<AscendC::HardEvent::MTE3_V>(SYNC_OUTPUT_BUF1_FLAG);
1693+}
1694+ 
1695+template <typename SFAAT>
1696+__aicore__ inline void
1697+SFAAVectorServiceGqaMsd<SFAAT>::Bmm2ResCopyOut(const RunInfo &info, LocalTensor<T> &bmm2ResUb, uint32_t wsMStart,
1698+ uint32_t dealRowCount, uint32_t columnCount,
1699+ uint32_t actualColumnCount)
1700+{
1701+ if constexpr (IS_META) {
1702+ if (FLASH_DECODE) {
1703+ if (info.tndIsS2SplitCore) {
1704+ Bmm2FDDataCopyOut(info, bmm2ResUb, wsMStart, dealRowCount, columnCount, actualColumnCount);
1705+ } else {
1706+ Bmm2CastAndCopyOut(info, bmm2ResUb, wsMStart, dealRowCount, columnCount, actualColumnCount);
1707+ }
1708+ } else {
1709+ Bmm2CastAndCopyOut(info, bmm2ResUb, wsMStart, dealRowCount, columnCount, actualColumnCount);
1710+ }
1711+ } else {
1712+ Bmm2CastAndCopyOut(info, bmm2ResUb, wsMStart, dealRowCount, columnCount, actualColumnCount);
1713+ }
1714+}
1715+ 
1716+template <typename SFAAT>
1717+__aicore__ inline void
1718+SFAAVectorServiceGqaMsd<SFAAT>::RowDivs(LocalTensor<float> dstUb, LocalTensor<float> src0Ub, LocalTensor<float> src1Ub,
1719+ uint32_t dealRowCount, uint32_t columnCount, uint32_t actualColumnCount)
1720+{
1721+ // divs by row, 每行的元素除以相同的元素
1722+ // dstUb[i, (j * 8) : (j * 8 + 7)] = src0Ub[i, (j * 8) : (j * 8 + 7)] / src1Ub[i, 0 : 7]
1723+ // src0Ub:[dealRowCount, columnCount], src1Ub:[dealRowCount, FP32_BLOCK_ELEMENT_NUM] dstUb:[dealRowCount,
1724+ // columnCount]
1725+ uint32_t dtypeMask = FP32_REPEAT_ELEMENT_NUM;
1726+ uint32_t dLoop = actualColumnCount / dtypeMask;
1727+ uint32_t dRemain = actualColumnCount % dtypeMask;
1728+ 
1729+ BinaryRepeatParams repeatParamsDiv;
1730+ repeatParamsDiv.src0BlkStride = 1;
1731+ repeatParamsDiv.src1BlkStride = 0;
1732+ repeatParamsDiv.dstBlkStride = 1;
1733+ repeatParamsDiv.src0RepStride = columnCount / FP32_BLOCK_ELEMENT_NUM;
1734+ repeatParamsDiv.src1RepStride = 1;
1735+ repeatParamsDiv.dstRepStride = columnCount / FP32_BLOCK_ELEMENT_NUM;
1736+ uint32_t columnRepeatCount = dLoop;
1737+ if (columnRepeatCount <= dealRowCount) {
1738+ uint32_t offset = 0;
1739+ for (uint32_t i = 0; i < dLoop; i++) {
1740+ Div(dstUb[offset], src0Ub[offset], src1Ub, dtypeMask, dealRowCount, repeatParamsDiv);
1741+ offset += dtypeMask;
1742+ }
1743+ } else {
1744+ BinaryRepeatParams columnRepeatParams;
1745+ columnRepeatParams.src0BlkStride = 1;
1746+ columnRepeatParams.src1BlkStride = 0;
1747+ columnRepeatParams.dstBlkStride = 1;
1748+ columnRepeatParams.src0RepStride = 8; // 列方向上两次repeat起始地址间隔dtypeMask=64个元素，即8个block
1749+ columnRepeatParams.src1RepStride = 0;
1750+ columnRepeatParams.dstRepStride = 8; // 列方向上两次repeat起始地址间隔dtypeMask=64个元素，即8个block
1751+ uint32_t offset = 0;
1752+ for (uint32_t i = 0; i < dealRowCount; i++) {
1753+ Div(dstUb[offset], src0Ub[offset], src1Ub[i * FP32_BLOCK_ELEMENT_NUM], dtypeMask, columnRepeatCount,
1754+ columnRepeatParams);
1755+ offset += columnCount;
1756+ }
1757+ }
1758+ if (dRemain > 0) {
1759+ Div(dstUb[dLoop * dtypeMask], src0Ub[dLoop * dtypeMask], src1Ub, dRemain, dealRowCount, repeatParamsDiv);
1760+ }
1761+}
1762+ 
1763+#endif // SPARSE_FLASH_ATTENTION_ANTIQUANT_SERVICE_VECTOR_GQA_MSD_H
Aexperimental/attention/sparse_flash_attention_antiquant/op_kernel/sparse_flash_attention_antiquant_template_tiling_key.h+69-0
@@ -0,0 +1,69 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file sparse_flash_attention_antiquant_template_tiling_key.h
13+ * \brief
14+ */
15+ 
16+#ifndef SPARSE_FLASH_ATTENTION_ANTIQUANT_TEMPLATE_TILING_KEY_H
17+#define SPARSE_FLASH_ATTENTION_ANTIQUANT_TEMPLATE_TILING_KEY_H
18+ 
19+#include "ascendc/host_api/tiling/template_argument.h"
20+ 
21+#define SFAA_LAYOUT_BSND 0
22+#define SFAA_LAYOUT_TND 1
23+#define SFAA_LAYOUT_PA_BSND 2
24+#define SFAA_LAYOUT_PA_BNSD 3
25+#define SFAA_LAYOUT_PA_NZ 4
26+ 
27+#define ASCENDC_TPL_4_BW 4
28+ 
29+#define C_TEMPLATE 0
30+#define V_TEMPLATE 1
31+ 
32+#define ATTENTION_GQA_MHA 0
33+#define ATTENTION_MLA_NAIVE 1
34+#define ATTENTION_MLA_ABSORB 2
35+ 
36+// 模板参数支持的范围定义
37+ASCENDC_TPL_ARGS_DECL(SparseFlashAttentionAntiquant, // 算子OpType
38+ASCENDC_TPL_UINT_DECL(ATTENTION_MODE, ASCENDC_TPL_4_BW, ASCENDC_TPL_UI_LIST, ATTENTION_GQA_MHA,
39+ ATTENTION_MLA_NAIVE, ATTENTION_MLA_ABSORB),
40+ASCENDC_TPL_BOOL_DECL(FLASH_DECODE, 0, 1),
41+ASCENDC_TPL_UINT_DECL(LAYOUT_T, ASCENDC_TPL_4_BW, ASCENDC_TPL_UI_LIST, SFAA_LAYOUT_BSND, SFAA_LAYOUT_TND),
42+ASCENDC_TPL_UINT_DECL(KV_LAYOUT_T, ASCENDC_TPL_4_BW, ASCENDC_TPL_UI_LIST, SFAA_LAYOUT_BSND, SFAA_LAYOUT_TND,
43+ SFAA_LAYOUT_PA_BSND, SFAA_LAYOUT_PA_BNSD, SFAA_LAYOUT_PA_NZ),
44+ASCENDC_TPL_UINT_DECL(TEMPLATE_MODE, ASCENDC_TPL_4_BW, ASCENDC_TPL_UI_LIST, C_TEMPLATE, V_TEMPLATE),
45+);
46+ 
47+// 支持的模板参数组合
48+// 用于调用GET_TPL_TILING_KEY获取TilingKey时，接口内部校验TilingKey是否合法
49+ASCENDC_TPL_SEL(
50+ // 为量化当前仅支持V模板
51+ ASCENDC_TPL_ARGS_SEL(
52+ ASCENDC_TPL_UINT_DECL(ATTENTION_MODE, ASCENDC_TPL_4_BW, ASCENDC_TPL_UI_LIST, ATTENTION_GQA_MHA,
53+ ATTENTION_MLA_ABSORB),
54+ ASCENDC_TPL_BOOL_SEL(FLASH_DECODE, 0),
55+ ASCENDC_TPL_UINT_SEL(LAYOUT_T, ASCENDC_TPL_UI_LIST, SFAA_LAYOUT_BSND, SFAA_LAYOUT_TND),
56+ ASCENDC_TPL_UINT_SEL(KV_LAYOUT_T, ASCENDC_TPL_UI_LIST, SFAA_LAYOUT_PA_BSND, SFAA_LAYOUT_PA_BNSD, SFAA_LAYOUT_PA_NZ, SFAA_LAYOUT_BSND),
57+ ASCENDC_TPL_UINT_SEL(TEMPLATE_MODE, ASCENDC_TPL_UI_LIST, V_TEMPLATE), // V模板不支持非PA
58+ ),
59+ ASCENDC_TPL_ARGS_SEL(
60+ ASCENDC_TPL_UINT_DECL(ATTENTION_MODE, ASCENDC_TPL_4_BW, ASCENDC_TPL_UI_LIST, ATTENTION_GQA_MHA,
61+ ATTENTION_MLA_ABSORB),
62+ ASCENDC_TPL_BOOL_SEL(FLASH_DECODE, 1),
63+ ASCENDC_TPL_UINT_SEL(LAYOUT_T, ASCENDC_TPL_UI_LIST, SFAA_LAYOUT_BSND, SFAA_LAYOUT_TND),
64+ ASCENDC_TPL_UINT_SEL(KV_LAYOUT_T, ASCENDC_TPL_UI_LIST, SFAA_LAYOUT_PA_BSND, SFAA_LAYOUT_PA_BNSD, SFAA_LAYOUT_PA_NZ, SFAA_LAYOUT_BSND),
65+ ASCENDC_TPL_UINT_SEL(TEMPLATE_MODE, ASCENDC_TPL_UI_LIST, V_TEMPLATE), // V模板不支持非PA
66+ ),
67+);
68+ 
69+#endif // TEMPLATE_TILING_KEY
Aexperimental/attention/sparse_flash_attention_antiquant_metadata/CMakeLists.txt+0-0
The file is empty
Aexperimental/attention/sparse_flash_attention_antiquant_metadata/examples/test_aclnn_sparse_flash_attention_antiquant_metadata.cpp+236-0
@@ -0,0 +1,236 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+#include <stdio.h>
12+#include <iostream>
13+#include <tuple>
14+#include <vector>
15+#include "acl/acl.h"
16+//#include "aclnnop/aclnn_sparse_flash_attention_antiquant_metadata.h"
17+#include "../../sparse_flash_attention_antiquant/op_kernel/sparse_flash_attention_antiquant_metadata.h"
18+#include "../op_host/aclnn_sparse_flash_attention_antiquant_metadata.h"
19+ 
20+static const uint32_t batchSize = 4;
21+static const uint32_t querySeqSize = 3;
22+static const uint32_t queryHeadNum = 10;
23+static const uint32_t kvSeqSize = 10240;
24+static const uint32_t kvHeadNum = 1;
25+static const uint32_t headDim = 128;
26+static const uint32_t topKSize = 128;
27+static const uint32_t sparseBlockSize = 16;
28+static std::string layoutQuery = "BSND";
29+static std::string layoutKV = "BSND";
30+static const uint32_t sparseMode = 0;
31+static const uint32_t attentionMode = 0;
32+static const uint32_t ropeHeadDim = 64;
33+static const uint32_t sparseSharedSize = 3;
34+ 
35+static const std::vector<int32_t> actSeqLenQuery = {3, 6, 9, 12};
36+static const std::vector<int32_t> actSeqLenKV = {10240, 10240, 10240, 10240};
37+static const std::vector<int32_t> sparseSeqLenKV = {2560, 2560, 2560, 2560};
38+static const std::vector<int64_t> actSeqLenQShape = {batchSize};
39+static const std::vector<int64_t> actSeqLenKVShape = {batchSize};
40+static const std::vector<int64_t> sparseSeqLenKVShape = {batchSize};
41+static const std::vector<int64_t> actSeqLenQStride = {1};
42+static const std::vector<int64_t> actSeqLenKVStride = {1};
43+static const std::vector<int64_t> sparseSeqLenKVStride = {1};
44+static const std::vector<int64_t> metadataShape = {optiling::SFA_META_SIZE};
45+static const std::vector<int64_t> metadataStride = {1};
46+ 
47+static const bool enableActLenQuery = true;
48+static const bool enableActLenKV = true;
49+static const bool enablesparseSeqLenKV = true;
50+ 
51+std::tuple<aclTensor*, void*> CreateTensor(size_t size, // in bytes
52+ std::vector<int64_t> shape,
53+ std::vector<int64_t> stride,
54+ aclDataType dType,
55+ const void* hostData = nullptr) {
56+ void* devicePtr = nullptr;
57+ auto ret = aclrtMalloc(&devicePtr, size, ACL_MEM_MALLOC_HUGE_FIRST);
58+ if (ret != ACL_SUCCESS) {
59+ printf("aclrtMalloc %d\n", ret);
60+ return {nullptr, nullptr};
61+ }
62+ 
63+ aclTensor* tensor = aclCreateTensor(&shape[0], shape.size(), dType,
64+ &stride[0], 0, aclFormat::ACL_FORMAT_ND,
65+ &shape[0], shape.size(), devicePtr);
66+ if (tensor == nullptr) {
67+ aclrtFree(devicePtr);
68+ return {nullptr, nullptr};
69+ }
70+ 
71+ if (hostData != nullptr) {
72+ aclrtMemcpy(devicePtr, size, hostData, size, ACL_MEMCPY_HOST_TO_DEVICE);
73+ }
74+ return {tensor, devicePtr};
75+}
76+ 
77+static void DumpMeta(void* data) {
78+ optiling::detail::SfaMetaData* metaDataPtr =
79+ (optiling::detail::SfaMetaData*)data;
80+ printf("mBaseSize: %d \n", metaDataPtr->mBaseSize);
81+ printf("s2BaseSize: %d \n", metaDataPtr->s2BaseSize);
82+ printf("gS1BaseSizeOfFd: %d \n", metaDataPtr->gS1BaseSizeOfFd);
83+ printf("usedCoreNum: %d \n", metaDataPtr->usedCoreNum);
84+ printf("numOfFdHead: %d \n", metaDataPtr->numOfFdHead);
85+ printf("usedVecNumOfFd: %d \n", metaDataPtr->usedVecNumOfFd);
86+ for (uint32_t i = 0; i < metaDataPtr->usedCoreNum; i++) {
87+ printf("bN2End[%d]: %d \n", i, metaDataPtr->bN2End[i]);
88+ printf("gS1End[%d]: %d \n", i, metaDataPtr->gS1End[i]);
89+ printf("s2End[%d]: %d \n", i, metaDataPtr->s2End[i]);
90+ printf("s2SplitStartIdxOfCore[%d]: %d \n", i,
91+ metaDataPtr->fdRes.s2SplitStartIdxOfCore[i]);
92+ }
93+ 
94+ for (uint32_t i = 0; i < metaDataPtr->numOfFdHead; i++) {
95+ printf("bN2IdxOfFdHead[%d]: %d \n", i,
96+ metaDataPtr->fdRes.bN2IdxOfFdHead[i]);
97+ printf("gS1IdxOfFdHead[%d]: %d \n", i,
98+ metaDataPtr->fdRes.gS1IdxOfFdHead[i]);
99+ printf("s2SplitNumOfFdHead[%d]: %d \n", i,
100+ metaDataPtr->fdRes.s2SplitNumOfFdHead[i]);
101+ printf("gS1SplitNumOfFdHead[%d]: %d \n", i,
102+ metaDataPtr->fdRes.gS1SplitNumOfFdHead[i]);
103+ printf("gS1LastPartSizeOfFdHead[%d]: %d \n", i,
104+ metaDataPtr->fdRes.gS1LastPartSizeOfFdHead[i]);
105+ }
106+ 
107+ for (uint32_t i = 0; i < metaDataPtr->usedVecNumOfFd; i++) {
108+ printf("gS1IdxEndOfFdHead[%d]: %d \n", i,
109+ metaDataPtr->fdRes.gS1IdxEndOfFdHead[i]);
110+ printf("gS1IdxEndOfFdHeadSplit[%d]: %d \n", i,
111+ metaDataPtr->fdRes.gS1IdxEndOfFdHeadSplit[i]);
112+ }
113+}
114+ 
115+int main() {
116+ int32_t deviceId = 0;
117+ aclrtStream stream;
118+ aclError ret = 0;
119+ aclTensor* qSeqLenTensor = nullptr;
120+ void* qSeqLenDevPtr = nullptr;
121+ aclTensor* kvSeqLenTensor = nullptr;
122+ void* kvSeqLenDevPtr = nullptr;
123+ aclTensor* spSeqLenTensor = nullptr;
124+ void* spSeqLenDevPtr = nullptr;
125+ aclTensor* metadataTensor = nullptr;
126+ void* metadataDevPtr = nullptr;
127+ aclOpExecutor* executor = nullptr;
128+ uint64_t workspaceSize = 0;
129+ void* workspace = nullptr;
130+ 
131+ ret = aclInit(nullptr);
132+ if (ret != ACL_SUCCESS) {
133+ printf("aclInit %d\n", ret);
134+ return -1;
135+ }
136+ 
137+ ret = aclrtSetDevice(deviceId);
138+ if (ret != ACL_SUCCESS) {
139+ printf("aclrtSetDevice %d\n", ret);
140+ return -1;
141+ }
142+ 
143+ ret = aclrtCreateStream(&stream);
144+ if (ret != ACL_SUCCESS) {
145+ printf("aclrtCreateStream %d\n", ret);
146+ return -1;
147+ }
148+ 
149+ if (enableActLenQuery) {
150+ std::tie(qSeqLenTensor, qSeqLenDevPtr) = CreateTensor(
151+ actSeqLenQuery.size() * sizeof(actSeqLenQuery[0]), actSeqLenQShape,
152+ actSeqLenQStride, aclDataType::ACL_INT32, &actSeqLenQuery[0]);
153+ if (qSeqLenTensor == nullptr) {
154+ return -1;
155+ }
156+ }
157+ 
158+ if (enableActLenKV) {
159+ std::tie(kvSeqLenTensor, kvSeqLenDevPtr) = CreateTensor(
160+ actSeqLenKV.size() * sizeof(actSeqLenKV[0]), actSeqLenKVShape,
161+ actSeqLenKVStride, aclDataType::ACL_INT32, &actSeqLenKV[0]);
162+ if (kvSeqLenTensor == nullptr) {
163+ return -1;
164+ }
165+ }
166+ 
167+ if (enablesparseSeqLenKV) {
168+ std::tie(spSeqLenTensor, spSeqLenDevPtr) = CreateTensor(
169+ sparseSeqLenKV.size() * sizeof(sparseSeqLenKV[0]), sparseSeqLenKVShape,
170+ sparseSeqLenKVStride, aclDataType::ACL_INT32, &sparseSeqLenKV[0]);
171+ if (spSeqLenTensor == nullptr) {
172+ return -1;
173+ }
174+ }
175+ 
176+ std::tie(metadataTensor, metadataDevPtr) =
177+ CreateTensor(sizeof(int32_t) * optiling::SFA_META_SIZE, metadataShape,
178+ metadataStride, aclDataType::ACL_INT32);
179+ if (metadataTensor == nullptr) {
180+ return -1;
181+ }
182+ 
183+ ret = aclnnKvQuantSparseFlashAttentionMetadataGetWorkspaceSize(
184+ qSeqLenTensor, kvSeqLenTensor, spSeqLenTensor, batchSize, querySeqSize,
185+ queryHeadNum, kvSeqSize, kvHeadNum, headDim, topKSize, sparseBlockSize,
186+ &layoutQuery[0], &layoutKV[0], sparseMode,
187+ attentionMode, ropeHeadDim, sparseSharedSize, metadataTensor,
188+ &workspaceSize, &executor);
189+ if (ret != ACL_SUCCESS) {
190+ printf("aclnnKvQuantSparseFlashAttentionMetadataGetWorkspaceSize %d\n",
191+ ret);
192+ return -1;
193+ }
194+ 
195+ ret = aclnnKvQuantSparseFlashAttentionMetadata(workspace, workspaceSize,
196+ executor, stream);
197+ if (ret != ACL_SUCCESS) {
198+ printf("aclnnKvQuantSparseFlashAttentionMetadata %d\n", ret);
199+ return -1;
200+ }
201+ 
202+ ret = aclrtSynchronizeStream(stream);
203+ if (ret != ACL_SUCCESS) {
204+ printf("aclrtSynchronizeStream %d\n", ret);
205+ return -1;
206+ }
207+ 
208+ std::vector<int32_t> metdataHost(optiling::SFA_META_SIZE);
209+ ret = aclrtMemcpy(metdataHost.data(),
210+ metdataHost.size() * sizeof(metdataHost[0]), metadataDevPtr,
211+ optiling::SFA_META_SIZE * sizeof(int32_t),
212+ ACL_MEMCPY_DEVICE_TO_HOST);
213+ if (ret != ACL_SUCCESS) {
214+ printf("aclrtMemcpy %d\n", ret);
215+ return -1;
216+ }
217+ 
218+ DumpMeta(&metdataHost[0]);
219+ 
220+ aclDestroyTensor(qSeqLenTensor);
221+ aclDestroyTensor(kvSeqLenTensor);
222+ aclDestroyTensor(spSeqLenTensor);
223+ aclDestroyTensor(metadataTensor);
224+ 
225+ aclrtFree(qSeqLenDevPtr);
226+ aclrtFree(kvSeqLenDevPtr);
227+ aclrtFree(spSeqLenDevPtr);
228+ aclrtFree(metadataDevPtr);
229+ aclrtFree(workspace);
230+ 
231+ aclrtDestroyStream(stream);
232+ aclrtResetDevice(deviceId);
233+ aclFinalize();
234+ 
235+ return 0;
236+}
Aexperimental/attention/sparse_flash_attention_antiquant_metadata/examples/test_geir_sparse_flash_attention_antiquant_metadata.cpp+220-0
@@ -0,0 +1,220 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+#include <cstdint>
12+#include <vector>
13+#include <string>
14+#include <map>
15+ 
16+#include "graph.h"
17+#include "types.h"
18+#include "tensor.h"
19+#include "ge_error_codes.h"
20+#include "ge_api_types.h"
21+#include "ge_api.h"
22+#include "array_ops.h" // REG_OP(Data)
23+#include "ge_ir_build.h"
24+ 
25+// #include "experiment_ops.h"
26+#include "nn_other.h"
27+#include "../op_graph/sparse_flash_attention_antiquant_metadata_proto.h"
28+#include "../../sparse_flash_attention_antiquant/op_kernel/sparse_flash_attention_antiquant_metadata.h"
29+ 
30+using namespace ge;
31+ 
32+static const uint32_t batchSize = 4;
33+static const uint32_t querySeqSize = 3;
34+static const uint32_t queryHeadNum = 10;
35+static const uint32_t kvSeqSize = 10240;
36+static const uint32_t kvHeadNum = 1;
37+static const uint32_t headDim = 128;
38+static const uint32_t topKSize = 128;
39+static const uint32_t sparseBlockSize = 16;
40+static const std::string layoutQuery = "BSND";
41+static const std::string layoutKV = "BSND";
42+static const uint32_t sparseMode = 0;
43+static const uint32_t attentionMode = 0;
44+static const uint32_t ropeHeadDim = 64;
45+static const uint32_t sparseSharedSize = 3;
46+static const uint32_t aicCoreNum = 24;
47+static const uint32_t aivCoreNum = 48;
48+ 
49+static const std::vector<int32_t> actSeqLenQuery = {3, 6, 9, 12};
50+static const std::vector<int32_t> actSeqLenKV = {10240, 10240, 10240, 10240};
51+static const std::vector<int32_t> sparseSeqLenKV = {2560, 2560, 2560, 2560};
52+static const std::vector<int64_t> actSeqLenQShape = {batchSize};
53+static const std::vector<int64_t> actSeqLenKVShape = {batchSize};
54+static const std::vector<int64_t> sparseSeqLenKVShape = {batchSize};
55+static const std::vector<int64_t> metadataShape = {optiling::SFA_META_SIZE};
56+static const std::string dumpFile = "./dump";
57+ 
58+static const bool enableActLenQuery = true;
59+static const bool enableActLenKV = true;
60+static const bool enablesparseSeqLenKV = true;
61+ 
62+using namespace ge;
63+ 
64+class GeEnv {
65+public:
66+ GeEnv() {
67+ std::map<AscendString, AscendString> opt = {{"ge.exec.deviceId", "0"},
68+ {"ge.graphRunMode", "1"}};
69+ inited_ = GEInitialize(opt) == SUCCESS;
70+ }
71+ ~GeEnv() {
72+ if (inited_)
73+ GEFinalize();
74+ }
75+ 
76+ bool Ok() { return inited_; }
77+ 
78+private:
79+ bool inited_;
80+};
81+ 
82+int main(int argc, char **argv) {
83+ GeEnv geEnv;
84+ if (!geEnv.Ok()) {
85+ printf("GEInitialize fail\n");
86+ return -1;
87+ }
88+ 
89+ Graph graph("GraphKvQuantSparseFlashAttentionMetadata");
90+ 
91+ auto metaDataOp = op::KvQuantSparseFlashAttentionMetadata(
92+ "KvQuantSparseFlashAttentionMetadata-0");
93+ 
94+ // gen graph
95+ auto dataOp0 = op::Data("input0").set_attr_index(0); // Data 算子
96+ if (enableActLenQuery) {
97+ TensorDesc desc(ge::Shape(actSeqLenQShape), FORMAT_ND, DT_INT32);
98+ desc.SetPlacement(ge::kPlacementHost);
99+ desc.SetFormat(FORMAT_ND);
100+ desc.SetRealDimCnt(actSeqLenQShape.size());
101+ dataOp0.update_input_desc_x(desc);
102+ graph.AddOp(dataOp0);
103+ metaDataOp.set_input_actual_seq_lengths_query(dataOp0);
104+ }
105+ 
106+ auto dataOp1 = op::Data("input1").set_attr_index(0); // Data 算子
107+ if (enableActLenKV) {
108+ TensorDesc desc(ge::Shape(actSeqLenKVShape), FORMAT_ND, DT_INT32);
109+ desc.SetPlacement(ge::kPlacementHost);
110+ desc.SetFormat(FORMAT_ND);
111+ desc.SetRealDimCnt(actSeqLenKVShape.size());
112+ dataOp1.update_input_desc_x(desc);
113+ graph.AddOp(dataOp1);
114+ metaDataOp.set_input_actual_seq_lengths_kv(dataOp1);
115+ }
116+ 
117+ auto dataOp2 = op::Data("input2").set_attr_index(0); // Data 算子
118+ if (enablesparseSeqLenKV) {
119+ TensorDesc desc(ge::Shape(actSeqLenKVShape), FORMAT_ND, DT_INT32);
120+ desc.SetPlacement(ge::kPlacementHost);
121+ desc.SetFormat(FORMAT_ND);
122+ desc.SetRealDimCnt(sparseSeqLenKV.size());
123+ dataOp2.update_input_desc_x(desc);
124+ graph.AddOp(dataOp2);
125+ metaDataOp.set_input_sparse_seq_lengths_kv(dataOp2);
126+ }
127+ 
128+ metaDataOp.update_output_desc_metadata(TensorDesc{ge::Shape(metadataShape), FORMAT_ND, DT_INT32});
129+ metaDataOp.set_attr_batch_size(batchSize);
130+ metaDataOp.set_attr_query_seq_size(querySeqSize);
131+ metaDataOp.set_attr_query_head_num(queryHeadNum);
132+ metaDataOp.set_attr_kv_seq_size(kvSeqSize);
133+ metaDataOp.set_attr_kv_head_num(kvHeadNum);
134+ metaDataOp.set_attr_head_dim(headDim);
135+ metaDataOp.set_attr_topk_size(topKSize);
136+ metaDataOp.set_attr_sparse_block_size(sparseBlockSize);
137+ metaDataOp.set_attr_aic_core_num(aicCoreNum);
138+ metaDataOp.set_attr_aiv_core_num(aivCoreNum);
139+ metaDataOp.set_attr_layout_query(layoutQuery);
140+ metaDataOp.set_attr_layout_kv(layoutKV);
141+ metaDataOp.set_attr_sparse_mode(sparseMode);
142+ metaDataOp.set_attr_attention_mode(attentionMode);
143+ metaDataOp.set_attr_rope_head_dim(ropeHeadDim);
144+ metaDataOp.set_attr_sparse_shared_size(sparseSharedSize);
145+ graph.AddOp(metaDataOp);
146+ 
147+ // run Graph
148+ std::vector<ge::Operator> inputOps = {dataOp0, dataOp1, dataOp2};
149+ std::vector<ge::Operator> outputOps = {metaDataOp};
150+ graph.SetInputs(inputOps).SetOutputs(outputOps);
151+ 
152+ aclgrphDumpGraph(graph, dumpFile.c_str(), dumpFile.length());
153+ 
154+ std::vector<ge::Tensor> inputTensors;
155+ std::vector<ge::Tensor> outputTensors;
156+ 
157+ if (enableActLenQuery) {
158+ inputTensors.push_back(
159+ Tensor{dataOp0.get_input_desc_x(),
160+ reinterpret_cast<const uint8_t *>(&actSeqLenQuery[0]),
161+ actSeqLenQuery.size() * sizeof(actSeqLenQuery[0])});
162+ }
163+ 
164+ if (enableActLenKV) {
165+ inputTensors.push_back(
166+ Tensor{dataOp1.get_input_desc_x(),
167+ reinterpret_cast<const uint8_t *>(&actSeqLenKV[0]),
168+ actSeqLenKV.size() * sizeof(actSeqLenKV[0])});
169+ }
170+ 
171+ if (enablesparseSeqLenKV) {
172+ inputTensors.push_back(
173+ Tensor{dataOp2.get_input_desc_x(),
174+ reinterpret_cast<const uint8_t *>(&sparseSeqLenKV[0]),
175+ sparseSeqLenKV.size() * sizeof(sparseSeqLenKV[0])});
176+ }
177+ 
178+ std::map<AscendString, AscendString> build_options;
179+ auto session = std::make_shared<Session>(build_options);
180+ std::map<AscendString, AscendString> graph_options;
181+ uint32_t graph_id = 0;
182+ session->AddGraph(graph_id, graph, graph_options);
183+ if (session->RunGraph(graph_id, inputTensors, outputTensors)) {
184+ printf("RunGraph Fail\n");
185+ return -1;
186+ }
187+ 
188+ auto tensor = outputTensors[0];
189+ auto data = tensor.GetData();
190+ auto dataSize = tensor.GetTensorDesc().GetShape().GetShapeSize();
191+
192+ optiling::detail::SfaMetaData *metaDataPtr = (optiling::detail::SfaMetaData*)data;
193+ printf("mBaseSize: %d \n", metaDataPtr->mBaseSize);
194+ printf("s2BaseSize: %d \n", metaDataPtr->s2BaseSize);
195+ printf("gS1BaseSizeOfFd: %d \n", metaDataPtr->gS1BaseSizeOfFd);
196+ printf("usedCoreNum: %d \n", metaDataPtr->usedCoreNum);
197+ printf("numOfFdHead: %d \n", metaDataPtr->numOfFdHead);
198+ printf("usedVecNumOfFd: %d \n", metaDataPtr->usedVecNumOfFd);
199+ for (uint32_t i = 0; i < metaDataPtr->usedCoreNum; i ++) {
200+ printf("bN2End[%d]: %d \n", i, metaDataPtr->bN2End[i]);
201+ printf("gS1End[%d]: %d \n", i, metaDataPtr->gS1End[i]);
202+ printf("s2End[%d]: %d \n", i, metaDataPtr->s2End[i]);
203+ printf("s2SplitStartIdxOfCore[%d]: %d \n", i, metaDataPtr->fdRes.s2SplitStartIdxOfCore[i]);
204+ }
205+ 
206+ for (uint32_t i = 0; i < metaDataPtr->numOfFdHead; i ++) {
207+ printf("bN2IdxOfFdHead[%d]: %d \n", i, metaDataPtr->fdRes.bN2IdxOfFdHead[i]);
208+ printf("gS1IdxOfFdHead[%d]: %d \n", i, metaDataPtr->fdRes.gS1IdxOfFdHead[i]);
209+ printf("s2SplitNumOfFdHead[%d]: %d \n", i, metaDataPtr->fdRes.s2SplitNumOfFdHead[i]);
210+ printf("gS1SplitNumOfFdHead[%d]: %d \n", i, metaDataPtr->fdRes.gS1SplitNumOfFdHead[i]);
211+ printf("gS1LastPartSizeOfFdHead[%d]: %d \n", i, metaDataPtr->fdRes.gS1LastPartSizeOfFdHead[i]);
212+ }
213+ 
214+ for (uint32_t i = 0; i < metaDataPtr->usedVecNumOfFd; i ++) {
215+ printf("gS1IdxEndOfFdHead[%d]: %d \n", i, metaDataPtr->fdRes.gS1IdxEndOfFdHead[i]);
216+ printf("gS1IdxEndOfFdHeadSplit[%d]: %d \n", i, metaDataPtr->fdRes.gS1IdxEndOfFdHeadSplit[i]);
217+ }
218+ 
219+ return 0;
220+}
Aexperimental/attention/sparse_flash_attention_antiquant_metadata/op_graph/sparse_flash_attention_antiquant_metadata_proto.h+81-0
@@ -0,0 +1,81 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file sparse_flash_attention_antiquant_metadata_proto.h
13+ * \brief
14+ */
15+#ifndef SPARSE_FLASH_ATTENTION_ANTIQUANT_METADATA_PROTO_H
16+#define SPARSE_FLASH_ATTENTION_ANTIQUANT_METADATA_PROTO_H
17+ 
18+#include "graph/operator_reg.h"
19+#include "graph/types.h"
20+ 
21+namespace ge {
22+ 
23+/**
24+* @brief Function KvQuantSparseFlashAttentionMetadata.
25+ 
26+* @par Inputs:
27+* @li actual_seq_lengths_query: A matrix tensor. The type support int32.
28+* Efective sequence length of query in different batches.
29+* @li actual_seq_lengths_kv: A matrix tensor. The type support int32.
30+* Effective sequence length of key/value in different batches.
31+ 
32+* @par Attributes:
33+* @li batch_size: An int. batch size of QKV Tensor.
34+* @li query_seq_size: An int. Q tensor sequence len.
35+* @li query_head_num: An int. head number of query.
36+* @li kv_seq_size: An int. KV tensor sequence len.
37+* @li kv_head_num: An int. KV tensor head number.
38+* @li head_dim: An int. QKV head dim.
39+* @li topk_size: An int. topK size.
40+* @li sparse_block_size: An int. The block size in the sparse phase. Default: 1.
41+* @li layout_query: A string. Specifies the layout of `query`, the value must be one of ["BSND", "TND"]. Default: "BSND".
42+* @li layout_kv: A string. Specifies the layout of `key/value`, the value must be one of ["BSND", "TND", "PA_BSND"]. Default: "BSND".
43+* @li sparse_mode: Sparse mode. Default: 3.
44+* - 0: default mask
45+* - 3: rightDownCausal make
46+* - 4: band mask
47+* @li attention_mode: An int. Attention mode, 0: GQA/MHA; 1: MLA-naive; 2: MLA-absorb. Default: 0.
48+* @li tile_size: An int. Tile size. Default: 128.
49+* @li rope_head_dim: An int. Rope head dim. Default: 64.
50+ 
51+* @par Outputs:
52+* @li metadata: A matrix tensor. The type support int32.
53+* The output of attention structure.
54+*/
55+REG_OP(KvQuantSparseFlashAttentionMetadata)
56+ .OPTIONAL_INPUT(actual_seq_lengths_query, TensorType({DT_INT32}))
57+ .OPTIONAL_INPUT(actual_seq_lengths_kv, TensorType({DT_INT32}))
58+ .OPTIONAL_INPUT(sparse_seq_lengths_kv, TensorType({DT_INT32}))
59+ .OUTPUT(metadata, TensorType({DT_INT32}))
60+ .REQUIRED_ATTR(batch_size, Int)
61+ .REQUIRED_ATTR(query_seq_size, Int)
62+ .REQUIRED_ATTR(query_head_num, Int)
63+ .REQUIRED_ATTR(kv_seq_size, Int)
64+ .REQUIRED_ATTR(kv_head_num, Int)
65+ .REQUIRED_ATTR(head_dim, Int)
66+ .REQUIRED_ATTR(topk_size, Int)
67+ .REQUIRED_ATTR(sparse_block_size, Int)
68+ .REQUIRED_ATTR(aic_core_num, Int)
69+ .REQUIRED_ATTR(aiv_core_num, Int)
70+ .ATTR(layout_query, String, "BSND")
71+ .ATTR(layout_kv, String, "BSND")
72+ .ATTR(sparse_mode, Int, 3)
73+ .ATTR(attention_mode, Int, 0)
74+ .ATTR(rope_head_dim, Int, 64)
75+ .ATTR(sparse_shared_size, Int, 16)
76+ .ATTR(soc_version, String, "ascend910B")
77+ .OP_END_FACTORY_REG(KvQuantSparseFlashAttentionMetadata)
78+ 
79+} // namespace ge
80+ 
81+#endif // SPARSE_FLASH_ATTENTION_ANTIQUANT_METADATA_PROTO_H
Aexperimental/attention/sparse_flash_attention_antiquant_metadata/op_host/aclnn_sparse_flash_attention_antiquant_metadata.cpp+135-0
@@ -0,0 +1,135 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file aclnn_sparse_flash_attention_antiquant_metadata.cpp
13+ * \brief
14+ */
15+ 
16+#include "aclnn_sparse_flash_attention_antiquant_metadata.h"
17+#include "l0_sparse_flash_attention_antiquant_metadata.h"
18+#include "aclnn_kernels/contiguous.h"
19+#include "aclnn_kernels/reshape.h"
20+#include "aclnn/aclnn_base.h"
21+#include "aclnn_kernels/common/op_error_check.h"
22+#include "opdev/common_types.h"
23+#include "opdev/data_type_utils.h"
24+#include "opdev/format_utils.h"
25+#include "opdev/op_dfx.h"
26+#include "opdev/op_executor.h"
27+#include "opdev/op_log.h"
28+#include "opdev/tensor_view_utils.h"
29+#include "opdev/make_op_executor.h"
30+#include "opdev/platform.h"
31+ 
32+#ifdef __cplusplus
33+extern "C" {
34+#endif
35+ 
36+static aclnnStatus ParamsCheck(const aclTensor* actualSeqLengthsQueryOptional,
37+ const aclTensor* actualSeqLengthsKvOptional,
38+ const aclTensor* sparseSeqLengthsKvOptional,
39+ int64_t batchSize,
40+ int64_t querySeqSize,
41+ int64_t queryHeadNum,
42+ int64_t kvSeqSize,
43+ int64_t kvHeadNum,
44+ int64_t headDim,
45+ int64_t topkSize,
46+ int64_t sparseBlockSize,
47+ char* layoutQueryOptional,
48+ char* layoutKvOptional,
49+ int64_t sparseMode,
50+ int64_t attentionMode,
51+ int64_t ropeHeadDim,
52+ int64_t sparseSharedSize,
53+ const aclTensor* metaData) {
Y
YYuZhengzhong4月2日

ParamsCheck 函数体为空直接返回成功，未对任何输入参数进行校验。外部传入非法参数（如 batchSize <= 0、sparseBlockSize == 0）将直接透传到 AICPU kernel 导致未定义行为。quant_sals_indexer_metadata 模块也有同样的问题。

likedislike
M
mayafei
4月3日 评论：
54+ if (batchSize < 0 || querySeqSize < 0 || queryHeadNum < 0 ||
55+ kvSeqSize < 0 || kvHeadNum < 0 || headDim < 0 ||
56+ sparseBlockSize < 0 || sparseSharedSize < 0) {
57+ return ACLNN_ERR_PARAM_INVALID;
58+ }
59+ return ACLNN_SUCCESS;
60+}
61+ 
62+aclnnStatus aclnnKvQuantSparseFlashAttentionMetadataGetWorkspaceSize(
63+ const aclTensor* actualSeqLengthsQueryOptional,
64+ const aclTensor* actualSeqLengthsKvOptional,
65+ const aclTensor* sparseSeqLengthsKvOptional,
66+ int64_t batchSize,
67+ int64_t querySeqSize,
68+ int64_t queryHeadNum,
69+ int64_t kvSeqSize,
70+ int64_t kvHeadNum,
71+ int64_t headDim,
72+ int64_t topkSize,
73+ int64_t sparseBlockSize,
74+ char* layoutQueryOptional,
75+ char* layoutKvOptional,
76+ int64_t sparseMode,
77+ int64_t attentionMode,
78+ int64_t ropeHeadDim,
79+ int64_t sparseSharedSize,
80+ const aclTensor* metaData,
81+ uint64_t* workspaceSize,
82+ aclOpExecutor** executor) {
83+ L2_DFX_PHASE_1(
84+ aclnnKvQuantSparseFlashAttentionMetadata,
85+ DFX_IN(actualSeqLengthsQueryOptional, actualSeqLengthsKvOptional,
86+ sparseSeqLengthsKvOptional, batchSize, querySeqSize, queryHeadNum,
87+ kvSeqSize, kvHeadNum, headDim, topkSize, sparseBlockSize,
88+ layoutQueryOptional, layoutKvOptional,
89+ sparseMode, attentionMode, ropeHeadDim, sparseSharedSize),
90+ DFX_OUT(metaData));
91+ 
92+ auto uniqueExecutor = CREATE_EXECUTOR();
93+ CHECK_RET(uniqueExecutor.get() != nullptr, ACLNN_ERR_INNER_CREATE_EXECUTOR);
94+ 
95+ auto ret = ParamsCheck(
96+ actualSeqLengthsQueryOptional, actualSeqLengthsKvOptional,
97+ sparseSeqLengthsKvOptional, batchSize, querySeqSize, queryHeadNum,
98+ kvSeqSize, kvHeadNum, headDim, topkSize, sparseBlockSize,
99+ layoutQueryOptional, layoutKvOptional, sparseMode,
100+ attentionMode, ropeHeadDim, sparseSharedSize, metaData);
101+ CHECK_RET(ret == ACLNN_SUCCESS, ret);
102+ 
103+ const op::PlatformInfo &npuInfo = op::GetCurrentPlatformInfo();
104+ uint32_t aicCoreNum = npuInfo.GetCubeCoreNum();
105+ uint32_t aivCoreNum = npuInfo.GetVectorCoreNum();
106+ 
107+ std::string socVersionStr = npuInfo.GetSocLongVersion();
108+ const char* socVersionOptional = socVersionStr.c_str();
109+ 
110+ auto output = l0op::KvQuantSparseFlashAttentionMetadata(
111+ actualSeqLengthsQueryOptional, actualSeqLengthsKvOptional,
112+ sparseSeqLengthsKvOptional, batchSize, querySeqSize, queryHeadNum,
113+ kvSeqSize, kvHeadNum, headDim, topkSize, sparseBlockSize, aicCoreNum,
114+ aivCoreNum, layoutQueryOptional, layoutKvOptional, sparseMode,
115+ attentionMode, ropeHeadDim, sparseSharedSize, socVersionOptional, metaData,
116+ uniqueExecutor.get());
117+ CHECK_RET(output != nullptr, ACLNN_ERR_INNER_NULLPTR);
118+ 
119+ *workspaceSize = 0;
120+ uniqueExecutor.ReleaseTo(executor);
121+ return ACLNN_SUCCESS;
122+}
123+ 
124+__attribute__((visibility("default"))) aclnnStatus
125+aclnnKvQuantSparseFlashAttentionMetadata(void* workspace,
126+ uint64_t workspaceSize,
127+ aclOpExecutor* executor,
128+ aclrtStream stream) {
129+ L2_DFX_PHASE_2(aclnnKvQuantSparseFlashAttentionMetadata);
130+ return CommonOpExecutorRun(workspace, workspaceSize, executor, stream);
131+}
132+ 
133+#ifdef __cplusplus
134+}
135+#endif
Aexperimental/attention/sparse_flash_attention_antiquant_metadata/op_host/aclnn_sparse_flash_attention_antiquant_metadata.h+84-0
@@ -0,0 +1,84 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+#ifndef ACLNN_SPARSE_FLASH_ATTENTION_ANTIQUANT_METADATA_AICPU_H
12+#define ACLNN_SPARSE_FLASH_ATTENTION_ANTIQUANT_METADATA_AICPU_H
13+ 
14+#include "aclnn/aclnn_base.h"
15+ 
16+#ifdef __cplusplus
17+extern "C" {
18+#endif
19+ 
20+/* funtion: aclnnKvQuantSparseFlashAttentionMetadataGetWorkspaceSize
21+ * parameters :
22+ * actualSeqLengthsQueryOptional : optional
23+ * actualSeqLengthsKvOptional : optional
24+ * sparseSeqLengthsKvOptional : optional
25+ * batchSize : required
26+ * querySeqSize : required
27+ * queryHeadNum : required
28+ * kvSeqSize : required
29+ * kvHeadNum : required
30+ * headDim : required
31+ * topkSize : required
32+ * scaleValue : required
33+ * sparseBlockSize : required
34+ * layoutQueryOptional : optional
35+ * layoutKvOptional : optional
36+ * sparseMode : optional
37+ * attentionMode : optional
38+ * ropeHeadDim : optional
39+ * sparseShardSize : optional
40+ * out : required
41+ * workspaceSize : size of workspace(output).
42+ * executor : executor context(output).
43+ */
44+__attribute__((visibility("default"))) aclnnStatus
45+aclnnKvQuantSparseFlashAttentionMetadataGetWorkspaceSize(
46+ const aclTensor* actualSeqLengthsQueryOptional,
47+ const aclTensor* actualSeqLengthsKvOptional,
48+ const aclTensor* sparseSeqLengthsKvOptional,
49+ int64_t batchSize,
50+ int64_t querySeqSize,
51+ int64_t queryHeadNum,
52+ int64_t kvSeqSize,
53+ int64_t kvHeadNum,
54+ int64_t headDim,
55+ int64_t topkSize,
56+ int64_t sparseBlockSize,
57+ char* layoutQueryOptional,
58+ char* layoutKvOptional,
59+ int64_t sparseMode,
60+ int64_t attentionMode,
61+ int64_t ropeHeadDim,
62+ int64_t sparseShardSize,
63+ const aclTensor* metaData,
64+ uint64_t* workspaceSize,
65+ aclOpExecutor** executor);
66+ 
67+/* funtion: aclnnKvQuantSparseFlashAttentionMetadata
68+ * parameters :
69+ * workspace : workspace memory addr(input).
70+ * workspaceSize : size of workspace(input).
71+ * executor : executor context(input).
72+ * stream : acl stream.
73+ */
74+__attribute__((visibility("default"))) aclnnStatus
75+aclnnKvQuantSparseFlashAttentionMetadata(void* workspace,
76+ uint64_t workspaceSize,
77+ aclOpExecutor* executor,
78+ aclrtStream stream);
79+ 
80+#ifdef __cplusplus
81+}
82+#endif
83+ 
84+#endif // ACLNN_SPARSE_FLASH_ATTENTION_ANTIQUANT_METADATA_AICPU_H
Aexperimental/attention/sparse_flash_attention_antiquant_metadata/op_host/l0_sparse_flash_attention_antiquant_metadata.cpp+84-0
@@ -0,0 +1,84 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file l0_sparse_flash_attention_antiquant_metadata.cpp
13+ * \brief
14+ */
15+ 
16+#include "l0_sparse_flash_attention_antiquant_metadata.h"
17+#include "opdev/aicpu/aicpu_task.h"
18+#include "opdev/make_op_executor.h"
19+#include "opdev/op_def.h"
20+#include "opdev/op_dfx.h"
21+#include "opdev/op_executor.h"
22+#include "opdev/op_log.h"
23+#include "opdev/shape_utils.h"
24+ 
25+using namespace op;
26+namespace l0op {
27+OP_TYPE_REGISTER(KvQuantSparseFlashAttentionMetadata);
28+ 
29+const aclTensor* KvQuantSparseFlashAttentionMetadata(
30+ const aclTensor* actualSeqLengthsQueryOptional,
31+ const aclTensor* actualSeqLengthsKvOptional,
32+ const aclTensor* sparseSeqLengthsKvOptional,
33+ int64_t batchSize,
34+ int64_t querySeqSize,
35+ int64_t queryHeadNum,
36+ int64_t kvSeqSize,
37+ int64_t kvHeadNum,
38+ int64_t headDim,
39+ int64_t topkSize,
40+ int64_t sparseBlockSize,
41+ int64_t aicCoreNum,
42+ int64_t aivCoreNum,
43+ char* layoutQueryOptional,
44+ char* layoutKvOptional,
45+ int64_t sparseMode,
46+ int64_t attentionMode,
47+ int64_t ropeHeadDim,
48+ int64_t sparseShardSize,
49+ const char* socVersionOptional,
50+ const aclTensor* metaData,
51+ aclOpExecutor* executor) {
52+ L0_DFX(KvQuantSparseFlashAttentionMetadata, actualSeqLengthsQueryOptional,
53+ actualSeqLengthsKvOptional, sparseSeqLengthsKvOptional, batchSize,
54+ querySeqSize, queryHeadNum, kvSeqSize, kvHeadNum, headDim, topkSize,
55+ sparseBlockSize, aicCoreNum, aivCoreNum, layoutQueryOptional,
56+ layoutKvOptional, sparseMode, attentionMode, ropeHeadDim,
57+ sparseShardSize, socVersionOptional, metaData);
58+ 
59+ static internal::AicpuTaskSpace space(
60+ "KvQuantSparseFlashAttentionMetadata");
61+ 
62+ auto ret = ADD_TO_LAUNCHER_LIST_AICPU(
63+ KvQuantSparseFlashAttentionMetadata,
64+ OP_ATTR_NAMES({"batch_size", "query_seq_size", "query_head_num",
65+ "kv_seq_size", "kv_head_num", "head_dim", "topk_size",
66+ "sparse_block_size", "aic_core_num", "aiv_core_num",
67+ "layout_query", "layout_kv", "sparse_mode",
68+ "attention_mode", "rope_head_dim", "sparse_shared_size", "soc_version"}),
69+ OP_INPUT(actualSeqLengthsQueryOptional, actualSeqLengthsKvOptional,
70+ sparseSeqLengthsKvOptional),
71+ OP_OUTPUT(metaData),
72+ OP_ATTR(batchSize, querySeqSize, queryHeadNum, kvSeqSize, kvHeadNum,
73+ headDim, topkSize, sparseBlockSize, aicCoreNum, aivCoreNum,
74+ layoutQueryOptional, layoutKvOptional, sparseMode, attentionMode,
75+ ropeHeadDim, sparseShardSize, socVersionOptional));
76+ OP_CHECK(ret == ACL_SUCCESS,
77+ OP_LOGE(ACLNN_ERR_INNER_NULLPTR,
78+ "KvQuantSparseFlashAttentionMetadata"
79+ " ADD_TO_LAUNCHER_LIST_AICPU failed."),
80+ return nullptr);
81+ return metaData;
82+}
83+ 
84+} // namespace l0op
Aexperimental/attention/sparse_flash_attention_antiquant_metadata/op_host/l0_sparse_flash_attention_antiquant_metadata.h+42-0
@@ -0,0 +1,42 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+#ifndef L0_SPARSE_FLASH_ATTENTION_ANTIQUANT_METADATA_AICPU_H
12+#define L0_SPARSE_FLASH_ATTENTION_ANTIQUANT_METADATA_AICPU_H
13+ 
14+#include "opdev/op_executor.h"
15+ 
16+namespace l0op {
17+const aclTensor* KvQuantSparseFlashAttentionMetadata(
18+ const aclTensor* actualSeqLengthsQueryOptional,
19+ const aclTensor* actualSeqLengthsKvOptional,
20+ const aclTensor* sparseSeqLengthsKvOptional,
21+ int64_t batchSize,
22+ int64_t querySeqSize,
23+ int64_t queryHeadNum,
24+ int64_t kvSeqSize,
25+ int64_t kvHeadNum,
26+ int64_t headDim,
27+ int64_t topkSize,
28+ int64_t sparseBlockSize,
29+ int64_t aicCoreNum,
30+ int64_t aivCoreNum,
31+ char* layoutQueryOptional,
32+ char* layoutKvOptional,
33+ int64_t sparseMode,
34+ int64_t attentionMode,
35+ int64_t ropeHeadDim,
36+ int64_t sparseSharedSize,
37+ const char* socVersionOptional,
38+ const aclTensor* metaData,
39+ aclOpExecutor* executor);
40+} // namespace l0op
41+ 
42+#endif
Aexperimental/attention/sparse_flash_attention_antiquant_metadata/op_host/sparse_flash_attention_antiquant_metadata_infershape.cpp+42-0
@@ -0,0 +1,42 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file sparse_flash_attention_antiquant_metadata_infershape.cpp
13+ * \brief
14+ */
15+#include "register/op_impl_registry.h"
16+#include "error/ops_error.h"
17+#include "../../sparse_flash_attention_antiquant/op_kernel/sparse_flash_attention_antiquant_metadata.h"
18+#include "../op_graph/sparse_flash_attention_antiquant_metadata_proto.h"
19+ 
20+using namespace ge;
21+ 
22+namespace ops {
23+static ge::graphStatus InferShapeSparseFlashAttentionAntiquantMetaData(gert::InferShapeContext* context)
24+{
25+ gert::Shape* oShape = context->GetOutputShape(0);
26+ OPS_LOG_E_IF_NULL(context, oShape, return ge::GRAPH_FAILED);
27+ // output shape (SFA_META_SIZE, )
28+ oShape->SetDimNum(1);
29+ oShape->SetDim(0, optiling::SFA_META_SIZE);
30+ return GRAPH_SUCCESS;
31+}
32+ 
33+static ge::graphStatus InferDtypeSparseFlashAttentionAntiquantMetaData(gert::InferDataTypeContext* context)
34+{
35+ context->SetOutputDataType(0, DT_INT32);
36+ return GRAPH_SUCCESS;
37+}
38+ 
39+IMPL_OP_INFERSHAPE(KvQuantSparseFlashAttentionMetadata)
40+ .InferShape(InferShapeSparseFlashAttentionAntiquantMetaData)
41+ .InferDataType(InferDtypeSparseFlashAttentionAntiquantMetaData);
42+} // namespace ops
Aexperimental/attention/sparse_flash_attention_antiquant_metadata/op_kernel_aicpu/sparse_flash_attention_antiquant_metadata_aicpu.cpp+925-0
@@ -0,0 +1,925 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+
11+/*!
12+ * \file sparse_flash_attention_antiquant_metadata_aicpu.cpp
13+ * \brief
14+ */
15+ 
16+#include "sparse_flash_attention_antiquant_metadata_aicpu.h"
17+#include "../../sparse_flash_attention_antiquant/op_kernel/sparse_flash_attention_antiquant_metadata.h"
18+#include "../../common/cpu_context_util.h"
19+#include <cstdio>
20+#include <math.h>
21+ 
22+namespace aicpu {
23+uint32_t
24+SparseFlashAttentionAntiquantMetaDataCpuKernel::Compute(CpuKernelContext &ctx) {
25+ context_ = &ctx;
26+ bool success = Prepare(ctx) && BalanceSchedule() && GenMetaData();
27+ return success ? KERNEL_STATUS_OK : KERNEL_STATUS_PARAM_INVALID;
28+}
29+ 
30+bool SparseFlashAttentionAntiquantMetaDataCpuKernel::Prepare(
31+ CpuKernelContext &ctx) {
32+ // input
33+ actSeqLenQ_ = ctx.Input(static_cast<uint32_t>(ParamId::actSeqLenQ));
34+ actSeqLenKV_ = ctx.Input(static_cast<uint32_t>(ParamId::actSeqLenKV));
35+ sparseSeqLenKV_ = ctx.Input(static_cast<uint32_t>(ParamId::sparseSeqLenKV));
36+ // output
37+ metaData_ = ctx.Output(static_cast<uint32_t>(ParamId::metaData));
38+ 
39+ bool requiredAttrs = GetAttrValue(ctx, "batch_size", batchSize_) &&
40+ GetAttrValue(ctx, "query_seq_size", querySeqSize_) &&
41+ GetAttrValue(ctx, "query_head_num", queryHeadNum_) &&
42+ GetAttrValue(ctx, "kv_seq_size", kvSeqSize_) &&
43+ GetAttrValue(ctx, "kv_head_num", kvHeadNum_) &&
44+ GetAttrValue(ctx, "head_dim", headDim_) &&
45+ GetAttrValue(ctx, "topk_size", topKSize_) &&
46+ GetAttrValue(ctx, "sparse_block_size", sparseBlockSize_) &&
47+ GetAttrValue(ctx, "aic_core_num", aicCoreNum_) &&
48+ GetAttrValue(ctx, "aiv_core_num", aivCoreNum_);
49+ if (!requiredAttrs) {
50+ return false;
51+ }
52+ 
53+ coreNum_ = aicCoreNum_;
54+ // attributes optional
55+ GetAttrValueOpt(ctx, "layout_query", layoutQuery_);
56+ GetAttrValueOpt(ctx, "layout_kv", layoutKV_);
57+ GetAttrValueOpt(ctx, "sparse_mode", sparseMode_);
58+ GetAttrValueOpt(ctx, "attention_mode", attentionMode_);
59+ GetAttrValueOpt(ctx, "rope_head_dim", ropeHeadDim_);
60+ GetAttrValueOpt(ctx, "sparse_shared_size", sparseSharedSize_);
61+ 
62+ return (ParamsCheck() && ParamsInit());
63+}
64+ 
65+bool SparseFlashAttentionAntiquantMetaDataCpuKernel::ParamsCheck() {
66+ if (actSeqLenQ_ != nullptr) {
67+ auto shape = actSeqLenQ_->GetTensorShape();
68+ auto data = actSeqLenQ_->GetData();
69+ auto dtype = actSeqLenQ_->GetDataType();
70+ 
71+ KERNEL_CHECK_NULLPTR(shape, false,
72+ "shape of actual_seq_lengths_query is null");
73+ KERNEL_CHECK_NULLPTR(data, false,
74+ "data of actual_seq_lengths_query is null");
75+ KERNEL_CHECK_FALSE((dtype == DataType::DT_INT32), false,
76+ "dtype %u of actual_seq_lengths_query is not int32 ", dtype);
77+ 
78+ KERNEL_CHECK_FALSE(
79+ (shape->GetDims() == 1 && shape->GetDimSize(0) == batchSize_), false,
80+ "shape of actual_seq_lengths_query is not {%u,}", batchSize_);
81+ }
82+ 
83+ if (actSeqLenKV_ != nullptr) {
84+ auto shape = actSeqLenKV_->GetTensorShape();
85+ auto data = actSeqLenKV_->GetData();
86+ auto dtype = actSeqLenKV_->GetDataType();
87+ 
88+ KERNEL_CHECK_NULLPTR(shape, false,
89+ "shape of actual_seq_lengths_kv is null");
90+ KERNEL_CHECK_NULLPTR(data, false, "data of actual_seq_lengths_kv is null");
91+ KERNEL_CHECK_FALSE((dtype == DataType::DT_INT32), false,
92+ "dtype of actual_seq_lengths_kv is not int32");
93+ 
94+ KERNEL_CHECK_FALSE(
95+ (shape->GetDims() == 1 && shape->GetDimSize(0) == batchSize_), false,
96+ "shape of actual_seq_lengths_query date is not {%u,}", batchSize_);
97+ }
98+ 
99+ if (sparseSeqLenKV_ != nullptr) {
100+ auto shape = sparseSeqLenKV_->GetTensorShape();
101+ auto data = sparseSeqLenKV_->GetData();
102+ auto dtype = sparseSeqLenKV_->GetDataType();
103+ 
104+ KERNEL_CHECK_NULLPTR(shape, false,
105+ "shape of sparse_seq_lengths_kv is null");
106+ KERNEL_CHECK_NULLPTR(data, false, "data of sparse_seq_lengths_kv is null");
107+ KERNEL_CHECK_FALSE((dtype == DataType::DT_INT32), false,
108+ "dtype of sparse_seq_lengths_kv is not int32");
109+ KERNEL_CHECK_FALSE(
110+ (shape->GetDims() == 1 && shape->GetDimSize(0) == batchSize_), false,
111+ "shape of sparse_seq_lengths_kv date is not {%u,}", batchSize_);
112+ }
113+ 
114+ KERNEL_CHECK_FALSE((layoutQuery_ == "BSND" || layoutQuery_ == "TND"), false,
115+ "layout_query invalid");
116+ 
117+ KERNEL_CHECK_FALSE(
118+ (layoutKV_ == "BSND" || layoutKV_ == "TND" || layoutKV_ == "PA_BSND" || layoutKV_ == "PA_BNSD" || layoutKV_ == "PA_NZ"),
119+ false, "layout_kv invalid");
120+ KERNEL_CHECK_FALSE((sparseMode_ == 0 || sparseMode_ == 3), false,
121+ "sparse_mode invalid");
122+ KERNEL_CHECK_FALSE(
123+ (attentionMode_ == 0 || attentionMode_ == 1 || attentionMode_ == 2),
124+ false, "attention_mode invalid");
125+ 
126+ KERNEL_CHECK_NULLPTR(metaData_, false, "metadata is null");
127+ auto shape = metaData_->GetTensorShape();
128+ auto data = metaData_->GetData();
129+ auto dtype = metaData_->GetDataType();
130+ 
131+ KERNEL_CHECK_NULLPTR(shape, false, "shape of metadata is null");
132+ KERNEL_CHECK_NULLPTR(data, false, "data of metadata is null");
133+ KERNEL_CHECK_FALSE((dtype == DataType::DT_INT32), false,
134+ "dtype of metadata is not int32");
135+ KERNEL_CHECK_FALSE((shape->GetDims() == 1 &&
136+ shape->GetDimSize(0) == optiling::SFA_META_SIZE),
137+ false, "shape of sparse_seq_lengths_kv date is not {%u,}",
138+ optiling::SFA_META_SIZE);
139+ 
140+ KERNEL_CHECK_FALSE(
141+ (aicCoreNum_ != 0 && aivCoreNum_ != 0 && aivCoreNum_ % aicCoreNum_ == 0 &&
142+ aicCoreNum_ <= optiling::CORE_NUM &&
143+ aivCoreNum_ <= (2 * optiling::CORE_NUM)),
144+ false, "core num invalid aic:%u aiv:%u", aicCoreNum_,
145+ aivCoreNum_); // more limit check with platform-core
146+ 
147+ return true;
148+}
149+ 
150+bool SparseFlashAttentionAntiquantMetaDataCpuKernel::ParamsInit() {
151+ groupSize_ = queryHeadNum_ / kvHeadNum_;
152+ mBaseSize_ = groupSize_ * sparseSharedSize_;
153+ s2BaseSize_ = 1024U;
154+ gS1BaseSizeOfFd_ = 8U;
155+ return true;
156+}
157+ 
158+uint32_t SparseFlashAttentionAntiquantMetaDataCpuKernel::GetS1SeqSize(uint32_t bIdx)
159+{
160+ if (actSeqLenQ_ == nullptr) {
161+ return querySeqSize_;
162+ }
163+ const int32_t *s1Ptr = (int32_t*)actSeqLenQ_->GetData();
164+ if (layoutQuery_ == "TND") {
165+ return (bIdx == 0) ? static_cast<uint32_t>(s1Ptr[bIdx]) :
166+ static_cast<uint32_t>(s1Ptr[bIdx] - s1Ptr[bIdx - 1U]);
167+ } else {
168+ return static_cast<uint32_t>(s1Ptr[bIdx]);
169+ }
170+}
171+ 
172+uint32_t SparseFlashAttentionAntiquantMetaDataCpuKernel::GetS2SeqSize(uint32_t bIdx)
173+{
174+ uint32_t s2Size = 0;
175+ if (actSeqLenKV_ == nullptr) {
176+ s2Size = kvSeqSize_;
177+ } else {
178+ const int32_t *s2Ptr = (int32_t*)actSeqLenKV_->GetData();
179+ if (layoutKV_ == "TND") {
180+ s2Size = (bIdx == 0) ? static_cast<uint32_t>(s2Ptr[bIdx]) :
181+ static_cast<uint32_t>(s2Ptr[bIdx] - s2Ptr[bIdx - 1U]);
182+ } else {
183+ s2Size = static_cast<uint32_t>(s2Ptr[bIdx]);
184+ }
185+ }
186+ 
187+ uint32_t sparseSeqSize = GetSparseSeqSize(bIdx);
188+ if (sparseSeqSize < s2Size) {
189+ s2Size = sparseSeqSize;
190+ }
191+ return s2Size;
192+}
193+ 
194+uint32_t SparseFlashAttentionAntiquantMetaDataCpuKernel::GetSparseSeqSize(uint32_t bIdx)
195+{
196+ if (sparseSeqLenKV_ == nullptr) {
197+ return 0U;
198+ } else {
199+ const int32_t *sparseS2Ptr = (int32_t*)sparseSeqLenKV_->GetData();
200+ if (layoutKV_ == "TND") {
201+ return (bIdx == 0) ? static_cast<uint32_t>(sparseS2Ptr[bIdx]) :
202+ static_cast<uint32_t>(sparseS2Ptr[bIdx] - sparseS2Ptr[bIdx - 1U]);
203+ } else {
204+ return static_cast<uint32_t>(sparseS2Ptr[bIdx]);
205+ }
206+ }
207+}
208+ 
209+void SparseFlashAttentionAntiquantMetaDataCpuKernel::CalcSplitInfo(SplitContext &splitContext)
210+{
211+ // 计算每个batch的切分，统计是否为空batch，记录最后有效batch（每个batch的每个N2切分是一样的）
212+ SplitInfo &splitInfo = splitContext.splitInfo;
213+ for (uint32_t bIdx = 0; bIdx < batchSize_; bIdx++) {
214+ uint32_t s1Size = GetS1SeqSize(bIdx);
215+ uint32_t s2Size = GetS2SeqSize(bIdx);
216+ splitInfo.s1GBaseNum[bIdx] = (s1Size * groupSize_ + (mBaseSize_ - 1U)) / mBaseSize_;
217+ splitInfo.s1GTailSize[bIdx] = (s1Size * groupSize_) % mBaseSize_;
218+ splitInfo.s2BaseNum[bIdx] = (s2Size + s2BaseSize_ - 1U) / s2BaseSize_;
219+ splitInfo.s2TailSize[bIdx] = s2Size % s2BaseSize_;
220+ if (splitInfo.s1GBaseNum[bIdx] != 0U && splitInfo.s2BaseNum[bIdx] != 0U) {
221+ splitInfo.isKvSeqAllZero = false;
222+ }
223+ }
224+ return;
225+}
226+ 
227+int64_t SparseFlashAttentionAntiquantMetaDataCpuKernel::CalcPreTokenLeftUp(
228+ uint32_t s1Size, uint32_t s2Size)
229+{
230+ auto mode = static_cast<SparseMode>(sparseMode_);
231+ if (mode == SparseMode::BAND) {
232+ return static_cast<int64_t>(s1Size) - static_cast<int64_t>(s2Size) + preToken_;
233+ }
234+ return preToken_;
235+}
236+ 
237+int64_t SparseFlashAttentionAntiquantMetaDataCpuKernel::CalcNextTokenLeftUp(
238+ uint32_t s1Size, uint32_t s2Size)
239+{
240+ auto mode = static_cast<SparseMode>(sparseMode_);
241+ switch (mode) {
242+ case SparseMode::DEFAULT_MASK:
243+ case SparseMode::ALL_MASK:
244+ case SparseMode::LEFT_UP_CAUSAL:
245+ return nextToken_;
246+ case SparseMode::RIGHT_DOWN_CAUSAL:
247+ return static_cast<int64_t>(s2Size) - static_cast<int64_t>(s1Size);
248+ case SparseMode::BAND:
249+ return static_cast<int64_t>(s2Size) - static_cast<int64_t>(s1Size) + nextToken_;
250+ default:
251+ return nextToken_;
252+ }
253+}
254+ 
255+int64_t SparseFlashAttentionAntiquantMetaDataCpuKernel::CalcCost(
256+ uint32_t basicM, uint32_t basicS2)
257+{
258+ float a = 25.04f;
259+ float b = 2.41f;
260+ float c = 54.99f;
261+ return static_cast<int64_t>(a * basicM + b * basicS2 + c);
262+}
263+ 
264+BlockCost<int64_t> SparseFlashAttentionAntiquantMetaDataCpuKernel::CalcCostTable(uint32_t s1NormalSize,
265+ uint32_t s2NormalSize, uint32_t s1GTailSize, uint32_t s2TailSize)
266+{
267+ BlockCost<int64_t> typeCost {};
268+ typeCost[NORMAL_BLOCK][NORMAL_BLOCK] = CalcCost(s1NormalSize, s2NormalSize);
269+ typeCost[TAIL_BLOCK][NORMAL_BLOCK] = (s1GTailSize == 0U) ? 0U : CalcCost(s1GTailSize, s2NormalSize);
270+ typeCost[NORMAL_BLOCK][TAIL_BLOCK] = (s2TailSize == 0U) ? 0U : CalcCost(s1NormalSize, s2TailSize);
271+ typeCost[TAIL_BLOCK][TAIL_BLOCK] = (s1GTailSize == 0U || s2TailSize == 0U) ? 0U : CalcCost(s1GTailSize, s2TailSize);
272+ return typeCost;
273+}
274+ 
275+Range<uint32_t> SparseFlashAttentionAntiquantMetaDataCpuKernel::CalcS2Range(
276+ uint32_t s1GIdx, const BatchCache &batchCache)
277+{
278+ uint32_t s2Start = 0U;
279+ uint32_t s2End = 0U;
280+ 
281+ // actual seq == 0
282+ if (batchCache.s1Size == 0U || batchCache.s2Size == 0U) {
283+ return std::make_pair(s2Start, s2End);
284+ }
285+ 
286+ // no mask
287+ if (!attentionMode_) { //attentionMaskFlag ?
288+ s2Start = 0U;
289+ s2End = (batchCache.s2Size + s2BaseSize_ - 1U) / s2BaseSize_;
290+ return std::make_pair(s2Start, s2End);
291+ }
292+ 
293+ // 1. calc index of s2FirstToken, s2LastToken by index of s1GFirstToken, s1GLastToken
294+ int64_t s1GFirstToken = static_cast<int64_t>(s1GIdx) * static_cast<int64_t>(mBaseSize_);
295+ int64_t s1GLastToken = std::min(s1GFirstToken + static_cast<int64_t>(mBaseSize_),
296+ static_cast<int64_t>(batchCache.s1Size) * static_cast<int64_t>(groupSize_)) - 1;
297+
298+ int64_t s1FirstToken = 0;
299+ int64_t s1LastToken = 0;
300+ if (isS1G_) {
301+ s1FirstToken = s1GFirstToken / static_cast<int64_t>(groupSize_);
302+ s1LastToken = s1GLastToken / static_cast<int64_t>(groupSize_);
303+ } else {
304+ if (s1GFirstToken / batchCache.s1Size == s1GLastToken / batchCache.s1Size) {
305+ // start and end locate in one G
306+ s1FirstToken = s1GFirstToken % static_cast<int64_t>(batchCache.s1Size);
307+ s1LastToken = s1GLastToken % static_cast<int64_t>(batchCache.s1Size);
308+ } else {
309+ // start and end locate in tow or more G, but working same as crossing a complete block
310+ s1FirstToken = 0;
311+ s1LastToken = batchCache.s1Size;
312+ }
313+ }
314+ 
315+ int64_t s2FirstToken = s1FirstToken - batchCache.preTokenLeftUp;
316+ int64_t s2LastToken = s1LastToken + batchCache.nextTokenLeftUp;
317+ 
318+ // 2. trans index of token to index of block
319+ // no valid token
320+ if (s2FirstToken >= static_cast<int64_t>(batchCache.s2Size) || s2LastToken < 0 || s2LastToken < s2FirstToken) {
321+ s2Start = 0U;
322+ s2End = 0U;
323+ return std::make_pair(s2Start, s2End);
324+ }
325+ // get valid range
326+ s2FirstToken = Clip(s2FirstToken, static_cast<int64_t>(0), static_cast<int64_t>(batchCache.s2Size - 1U));
327+ s2LastToken = Clip(s2LastToken, static_cast<int64_t>(0), static_cast<int64_t>(batchCache.s2Size - 1U));
328+ 
329+ s2Start = static_cast<uint32_t>(s2FirstToken) / s2BaseSize_;
330+ s2End = static_cast<uint32_t>(s2LastToken) / s2BaseSize_ + 1U; // end of block index, Right-open interval
331+ 
332+ return std::make_pair(s2Start, s2End);
333+}
334+ 
335+void SparseFlashAttentionAntiquantMetaDataCpuKernel::CalcBatchCache(
336+ uint32_t bIdx, const SplitContext &splitContext, BatchCache &batchCache)
337+{
338+ const SplitInfo &splitInfo = splitContext.splitInfo;
339+ 
340+ batchCache.bIdx = bIdx;
341+ batchCache.s1Size = GetS1SeqSize(bIdx);
342+ batchCache.s2Size = GetS2SeqSize(bIdx);
343+ batchCache.preTokenLeftUp = CalcPreTokenLeftUp(batchCache.s1Size, batchCache.s2Size);
344+ batchCache.nextTokenLeftUp = CalcNextTokenLeftUp(batchCache.s1Size, batchCache.s2Size);
345+ batchCache.typeCost = CalcCostTable(mBaseSize_, s2BaseSize_, splitInfo.s1GTailSize[bIdx],
346+ splitInfo.s2TailSize[bIdx]);
347+}
348+ 
349+void SparseFlashAttentionAntiquantMetaDataCpuKernel::CalcS1GCache(uint32_t s1GIdx,
350+ const SplitContext &splitContext, const BatchCache &batchCache, S1GCache &s1GCache)
351+{
352+ const SplitInfo &splitInfo = splitContext.splitInfo;
353+ 
354+ s1GCache.bIdx = batchCache.bIdx;
355+ s1GCache.s1GIdx = s1GIdx;
356+ 
357+ auto s2Range = CalcS2Range(s1GIdx, batchCache);
358+ s1GCache.s2Start = s2Range.first;
359+ s1GCache.s2End = s2Range.second;
360+ 
361+ if (s1GCache.s2Start >= s1GCache.s2End || splitInfo.s1GBaseNum[batchCache.bIdx] == 0) {
362+ s1GCache.s1GBlock = 0;
363+ s1GCache.s1GCost = 0;
364+ s1GCache.s1GLastBlockCost = 0;
365+ s1GCache.s1GNormalBlockCost = 0;
366+ return;
367+ }
368+ 
369+ // 计算S2方向满块、尾块数量
370+ s1GCache.s1GBlock = s1GCache.s2End - s1GCache.s2Start;
371+ uint32_t curTailS2Num = (splitInfo.s2TailSize[batchCache.bIdx] != 0U &&
372+ s1GCache.s2End == splitInfo.s2BaseNum[batchCache.bIdx]) ? 1U : 0U;
373+ uint32_t curNormalS2Num = s1GCache.s1GBlock - curTailS2Num;
374+ if (s1GIdx == (splitInfo.s1GBaseNum[batchCache.bIdx] - 1U) && splitInfo.s1GTailSize[batchCache.bIdx] != 0U) {
375+ s1GCache.s1GCost = batchCache.typeCost[TAIL_BLOCK][NORMAL_BLOCK] * curNormalS2Num +
376+ batchCache.typeCost[TAIL_BLOCK][TAIL_BLOCK] * curTailS2Num;
377+ s1GCache.s1GLastBlockCost = curTailS2Num > 0U ? batchCache.typeCost[TAIL_BLOCK][TAIL_BLOCK] :
378+ batchCache.typeCost[TAIL_BLOCK][NORMAL_BLOCK];
379+ s1GCache.s1GNormalBlockCost = batchCache.typeCost[TAIL_BLOCK][NORMAL_BLOCK];
380+ } else {
381+ s1GCache.s1GCost = batchCache.typeCost[NORMAL_BLOCK][NORMAL_BLOCK] * curNormalS2Num +
382+ batchCache.typeCost[NORMAL_BLOCK][TAIL_BLOCK] * curTailS2Num;
383+ s1GCache.s1GLastBlockCost = curTailS2Num > 0U ? batchCache.typeCost[NORMAL_BLOCK][TAIL_BLOCK] :
384+ batchCache.typeCost[NORMAL_BLOCK][NORMAL_BLOCK];
385+ s1GCache.s1GNormalBlockCost = batchCache.typeCost[NORMAL_BLOCK][NORMAL_BLOCK];
386+ }
387+ 
388+ if (s1GCache.s1GBlock == 1) {
389+ s1GCache.s1GLastBlockCost += V0_COST;
390+ }
391+ s1GCache.s1GCost += V0_COST;
392+}
393+ 
394+void SparseFlashAttentionAntiquantMetaDataCpuKernel::CalcBatchCost(
395+ uint32_t bIdx, const SplitContext &splitContext, CostInfo &costInfo)
396+{
397+ const SplitInfo &splitInfo = splitContext.splitInfo;
398+ 
399+ costInfo.bN2CostOfEachBatch[bIdx] = 0;
400+ costInfo.bN2BlockOfEachBatch[bIdx] = 0U;
401+ costInfo.bN2LastBlockCostOfEachBatch[bIdx] = 0U;
402+ 
403+ if (GetS1SeqSize(bIdx) == 0U || GetS2SeqSize(bIdx) == 0U) {
404+ return;
405+ }
406+ 
407+ BatchCache bCache;
408+ S1GCache s1GCache;
409+ CalcBatchCache(bIdx, splitContext, bCache);
410+ for (uint32_t s1GIdx = 0; s1GIdx < splitInfo.s1GBaseNum[bIdx]; s1GIdx++) {
411+ CalcS1GCache(s1GIdx, splitContext, bCache, s1GCache);
412+ costInfo.bN2CostOfEachBatch[bIdx] += s1GCache.s1GCost;
413+ costInfo.bN2BlockOfEachBatch[bIdx] += s1GCache.s1GBlock;
414+ if(s1GCache.s1GBlock > 0){
415+ costInfo.bN2LastBlockCostOfEachBatch[bIdx] = s1GCache.s1GLastBlockCost;
416+ }
417+ }
418+}
419+ 
420+void SparseFlashAttentionAntiquantMetaDataCpuKernel::CalcCostInfo(SplitContext &splitContext)
421+{
422+ const SplitInfo &splitInfo = splitContext.splitInfo;
423+ CostInfo &costInfo = splitContext.costInfo;
424+ 
425+ if (splitInfo.isKvSeqAllZero) {
426+ costInfo.totalCost = 0;
427+ costInfo.totalBlockNum = 0U;
428+ return;
429+ }
430+ 
431+ // 计算batch的负载并记录，用于按batch分配，需要按行计算起止点，统计块数、负载
432+ for (uint32_t bIdx = 0; bIdx < batchSize_; bIdx++) {
433+ CalcBatchCost(bIdx, splitContext, costInfo);
434+ costInfo.totalCost += costInfo.bN2CostOfEachBatch[bIdx] * kvHeadNum_;
435+ costInfo.totalBlockNum += costInfo.bN2BlockOfEachBatch[bIdx] * kvHeadNum_;
436+ }
437+}
438+ 
439+void SparseFlashAttentionAntiquantMetaDataCpuKernel::UpdateCursor(const SplitContext &splitContext, AssignContext &assignContext)
440+{
441+ const SplitInfo &splitInfo = splitContext.splitInfo;
442+ const CostInfo &costInfo = splitContext.costInfo;
443+ 
444+ bool UpdateS1G = false;
445+ bool UpdateBatch = false;
446+ 
447+ // Update S2
448+ if (assignContext.curS2Idx >= assignContext.s1GCache.s2End) { // 边界assignInfo.s2End是取不到的开区间
449+ assignContext.curS2Idx = 0U;
450+ assignContext.curS1GIdx++;
451+ UpdateS1G = true;
452+ }
453+ 
454+ // Update S1G
455+ if (assignContext.curS1GIdx >= splitInfo.s1GBaseNum[assignContext.curBIdx]) {
456+ assignContext.curS1GIdx = 0U;
457+ assignContext.curBN2Idx++;
458+ }
459+ 
460+ // Update Batch
461+ if (assignContext.curBN2Idx == batchSize_ * kvHeadNum_) { // 所有负载全部分配完，设置最后一个核的右开区间，返回
462+ assignContext.curS1GIdx = 0U;
463+ assignContext.curS2Idx = 0U;
464+ assignContext.isFinished = true;
465+ return;
466+ }
467+ 
468+ if (assignContext.curBN2Idx / kvHeadNum_ != assignContext.curBIdx) {
469+ assignContext.curBIdx = assignContext.curBN2Idx / kvHeadNum_;
470+ assignContext.curS1GIdx = 0U;
471+ UpdateBatch = true;
472+ UpdateS1G = true;
473+ }
474+ 
475+ // Update Cache
476+ if (UpdateBatch) {
477+ CalcBatchCache(assignContext.curBIdx, splitContext, assignContext.batchCache);
478+ assignContext.bN2Cost = costInfo.bN2CostOfEachBatch[assignContext.curBIdx];
479+ assignContext.bN2Block = costInfo.bN2BlockOfEachBatch[assignContext.curBIdx];
480+ }
481+ if (UpdateS1G) {
482+ CalcS1GCache(assignContext.curS1GIdx, splitContext, assignContext.batchCache, assignContext.s1GCache);
483+ assignContext.curS2Idx = assignContext.s1GCache.s2Start;
484+ }
485+}
486+ 
487+void SparseFlashAttentionAntiquantMetaDataCpuKernel::AssignByBatch(const SplitContext &splitContext, AssignContext &assignContext)
488+{
489+ if (assignContext.isFinished) {
490+ return;
491+ }
492+ const CostInfo &costInfo = splitContext.costInfo;
493+ while (assignContext.bN2Cost == 0 || IsWithinTolerance(assignContext.coreCache.costLimit,
494+ costInfo.bN2LastBlockCostOfEachBatch[assignContext.curBIdx] / FA_TOLERANCE_RATIO,
495+ assignContext.coreCache.cost + assignContext.bN2Cost)) {
496+ assignContext.coreCache.cost += assignContext.bN2Cost;
497+ assignContext.coreCache.block += assignContext.bN2Block;
498+ assignContext.curBN2Idx++;
499+ 
500+ // to the end
501+ if (assignContext.curBN2Idx == batchSize_ * kvHeadNum_) {
502+ assignContext.curS1GIdx = 0U;
503+ assignContext.curS2Idx = 0U;
504+ assignContext.isFinished = true;
505+ return;
506+ }
507+ 
508+ // next batch
509+ if (assignContext.curBN2Idx / kvHeadNum_ != assignContext.curBIdx) {
510+ assignContext.curBIdx = assignContext.curBN2Idx / kvHeadNum_;
511+ CalcBatchCache(assignContext.curBIdx, splitContext, assignContext.batchCache);
512+ }
513+ 
514+ assignContext.bN2Cost = costInfo.bN2CostOfEachBatch[assignContext.curBIdx];
515+ assignContext.bN2Block = costInfo.bN2BlockOfEachBatch[assignContext.curBIdx];
516+ assignContext.curS1GIdx = 0U;
517+ CalcS1GCache(assignContext.curS1GIdx, splitContext, assignContext.batchCache, assignContext.s1GCache);
518+ assignContext.curS2Idx = assignContext.s1GCache.s2Start;
519+ }
520+}
521+ 
522+void SparseFlashAttentionAntiquantMetaDataCpuKernel::AssignByRow(const SplitContext &splitContext, AssignContext &assignContext)
523+{
524+ if (assignContext.isFinished) {
525+ return;
526+ }
527+ 
528+ while (IsWithinTolerance(assignContext.coreCache.costLimit,
529+ assignContext.s1GCache.s1GLastBlockCost / FA_TOLERANCE_RATIO,
530+ assignContext.coreCache.cost + assignContext.s1GCache.s1GCost)) {
531+ assignContext.coreCache.cost += assignContext.s1GCache.s1GCost;
532+ assignContext.coreCache.block += assignContext.s1GCache.s1GBlock;
533+ 
534+ // 当前batch被分配一行出去，更新剩余负载
535+ assignContext.bN2Cost = assignContext.bN2Cost > assignContext.s1GCache.s1GCost ?
536+ assignContext.bN2Cost - assignContext.s1GCache.s1GCost : 0;
537+ assignContext.bN2Block = assignContext.bN2Block > assignContext.s1GCache.s1GBlock ?
538+ assignContext.bN2Block - assignContext.s1GCache.s1GBlock : 0U;
539+ // 计算新一行的信息
540+ do{
541+ assignContext.curS1GIdx++;
542+ CalcS1GCache(assignContext.curS1GIdx, splitContext, assignContext.batchCache, assignContext.s1GCache);
543+ }while(assignContext.s1GCache.s1GBlock == 0);
544+ assignContext.curS2Idx = assignContext.s1GCache.s2Start;
545+ }
546+}
547+ 
548+void SparseFlashAttentionAntiquantMetaDataCpuKernel::AssignByBlock(const SplitContext &splitContext, AssignContext &assignContext)
549+{
550+ if (assignContext.isFinished) {
551+ return;
552+ }
553+ 
554+ int64_t curCost = assignContext.s1GCache.s1GNormalBlockCost;
555+ if (assignContext.curS2Idx == (assignContext.s1GCache.s2End - 1U)) {
556+ curCost = assignContext.s1GCache.s1GLastBlockCost;
557+ }
558+ int64_t realCost = curCost;
559+ if (assignContext.curS2Idx == 0 && assignContext.s1GCache.s1GBlock != 1) {
560+ realCost += V0_COST;
561+ }
562+ 
563+ while (IsWithinTolerance(assignContext.coreCache.costLimit, realCost / FA_TOLERANCE_RATIO,
564+ assignContext.coreCache.cost + realCost)) { // (costLimit - curCostOnCore) * FA_TOLERANCE_RATIO > curCost；至少分配1块
565+ assignContext.coreCache.cost += realCost;
566+ assignContext.coreCache.block++;
567+ assignContext.curS2Idx++;
568+ // 当前batch被分配一块出去，更新剩余负载
569+ assignContext.bN2Cost = assignContext.bN2Cost - realCost;
570+ // 当前行被分配一块出去，更新剩余负载
571+ assignContext.s1GCache.s1GCost = assignContext.s1GCache.s1GCost - realCost;
572+ assignContext.bN2Block--;
573+ assignContext.s1GCache.s1GBlock--;
574+ 
575+ realCost = curCost;
576+ }
577+}
578+ 
579+void SparseFlashAttentionAntiquantMetaDataCpuKernel::ForceAssign(const SplitContext &splitContext, AssignContext &assignContext)
580+{
581+ if (assignContext.isFinished) {
582+ return;
583+ }
584+ 
585+ int64_t curCost = assignContext.s1GCache.s1GNormalBlockCost;
586+ if (assignContext.curS2Idx == (assignContext.s1GCache.s2End - 1U)) {
587+ curCost = assignContext.s1GCache.s1GLastBlockCost;
588+ }
589+ int64_t realCost = curCost;
590+ if (assignContext.curS2Idx == 0 && assignContext.s1GCache.s1GBlock != 1) {
591+ realCost += V0_COST;
592+ }
593+ 
594+ assignContext.coreCache.cost += realCost;
595+ assignContext.coreCache.block++;
596+ assignContext.curS2Idx++;
597+ // 当前batch被分配一块出去，更新剩余负载
598+ assignContext.bN2Cost = assignContext.bN2Cost - realCost;
599+ assignContext.bN2Block--;
600+ // 当前行被分配一块出去，更新剩余负载
601+ assignContext.s1GCache.s1GCost = assignContext.s1GCache.s1GCost - realCost;
602+ assignContext.s1GCache.s1GBlock--;
603+ UpdateCursor(splitContext, assignContext);
604+}
605+ 
606+bool SparseFlashAttentionAntiquantMetaDataCpuKernel::IsNeedRecordFDInfo(const AssignContext &assignContext, const SplitResult &splitRes)
607+{
608+ // 切分点大概率不会刚好在行尾，因此滞后处理归约信息的统计，到下一个切分点再判断是否需要归约
609+ // 核0无需处理
610+ if (assignContext.curCoreIdx == 0U) {
611+ return false;
612+ }
613+ // 无跨核行，无需处理
614+ if (assignContext.curKvSplitPart <= 1U) {
615+ return false;
616+ }
617+ // 需要归约的行还未处理完
618+ if (assignContext.curBN2Idx == splitRes.bN2End[assignContext.curCoreIdx - 1U] &&
619+ assignContext.curS1GIdx == splitRes.gS1End[assignContext.curCoreIdx - 1U]) {
620+ return false;
621+ }
622+ return true;
623+}
624+ 
625+void SparseFlashAttentionAntiquantMetaDataCpuKernel::RecordFDInfo(const SplitContext &splitContext, const AssignContext &assignContext, SplitResult &result)
626+{
627+ const SplitInfo &splitInfo = splitContext.splitInfo;
628+ // 需要规约的行是上一个核的切分点所在位置
629+ uint32_t splitBIdx = result.bN2End[assignContext.curCoreIdx - 1U] / kvHeadNum_;
630+ uint32_t splitS1GIdx = result.gS1End[assignContext.curCoreIdx - 1U];
631+ uint32_t s1Size = GetS1SeqSize(splitBIdx);
632+ 
633+ // 计算归约数据的FD均衡划分信息
634+ uint32_t curFdS1gSize = (splitS1GIdx == splitInfo.s1GBaseNum[splitBIdx] - 1U) ?
635+ (s1Size * groupSize_ - splitS1GIdx * mBaseSize_) : mBaseSize_;
636+ uint32_t curFdS1gSplitPart = (curFdS1gSize + gS1BaseSizeOfFd_ - 1U) / gS1BaseSizeOfFd_;
637+ uint32_t curFdS1gLastPartSize = curFdS1gSize - (gS1BaseSizeOfFd_ * (curFdS1gSplitPart - 1U));
638+ // 记录
639+ result.maxS2SplitNum = std::max(result.maxS2SplitNum, assignContext.curKvSplitPart);
640+ // 若存在头归约，则切分点一定为上一个核结束的位置
641+ result.fdRes.bN2IdxOfFdHead[result.numOfFdHead] = result.bN2End[assignContext.curCoreIdx - 1U];
642+ result.fdRes.gS1IdxOfFdHead[result.numOfFdHead] = result.gS1End[assignContext.curCoreIdx - 1U];
643+ result.fdRes.s2SplitNumOfFdHead[result.numOfFdHead] = assignContext.curKvSplitPart;
644+ result.fdRes.gS1SplitNumOfFdHead[result.numOfFdHead] = curFdS1gSplitPart;
645+ result.fdRes.gS1LastPartSizeOfFdHead[result.numOfFdHead] = curFdS1gLastPartSize;
646+ result.numOfFdHead++;
647+}
648+ 
649+void SparseFlashAttentionAntiquantMetaDataCpuKernel::CalcSplitPlan(uint32_t coreNum,
650+ int64_t costLimit, const SplitContext &splitContext, SplitResult &result)
651+{
652+ const CostInfo &costInfo = splitContext.costInfo;
653+ 
654+ if (coreNum == 0U) {
655+ return;
656+ }
657+ result.maxCost = 0U;
658+ result.usedCoreNum = 0U;
659+ 
660+ AssignContext assignContext {};
661+ assignContext.curBIdx = 0U;
662+ assignContext.curS1GIdx = 0U;
663+ assignContext.unassignedCost = costInfo.totalCost;
664+ assignContext.bN2Cost = costInfo.bN2CostOfEachBatch[assignContext.curBIdx];
665+ assignContext.bN2Block = costInfo.bN2BlockOfEachBatch[assignContext.curBIdx];
666+ CalcBatchCache(assignContext.curBIdx, splitContext, assignContext.batchCache);
667+ CalcS1GCache(assignContext.curS1GIdx, splitContext, assignContext.batchCache, assignContext.s1GCache);
668+ assignContext.curS2Idx = assignContext.s1GCache.s2Start;
669+ 
670+ for (uint32_t i = 0; i < coreNum; ++i) {
671+ if (result.maxCost > costLimit) {
672+ return;
673+ }
674+ if (assignContext.isFinished || assignContext.unassignedCost <= 0) {
675+ break;
676+ }
677+ 
678+ assignContext.curCoreIdx = i;
679+ result.fdRes.s2SplitStartIdxOfCore[assignContext.curCoreIdx] = assignContext.curKvSplitPart - 1U;
680+ 
681+ assignContext.coreCache = {};
682+ assignContext.coreCache.costLimit = assignContext.unassignedCost / (coreNum - assignContext.curCoreIdx);
683+ 
684+ // 1、按整batch分配
685+ AssignByBatch(splitContext, assignContext);
686+ // 2、按行分配
687+ AssignByRow(splitContext, assignContext);
688+ // 3、按块分配
689+ AssignByBlock(splitContext, assignContext);
690+ // 4、强制分配
691+ if (assignContext.coreCache.block == 0) {
692+ ForceAssign(splitContext, assignContext);
693+ }
694+ 
695+ result.bN2End[i] = assignContext.curBN2Idx;
696+ result.gS1End[i] = assignContext.curS1GIdx;
697+ result.s2End[i] = assignContext.curS2Idx;
698+ result.maxCost = std::max(result.maxCost, assignContext.coreCache.cost);
699+ 
700+ assignContext.unassignedCost -= assignContext.coreCache.cost;
701+ 
702+ // 对之前的归约信息进行记录并清理
703+ if (IsNeedRecordFDInfo(assignContext, result)) {
704+ RecordFDInfo(splitContext, assignContext, result);
705+ assignContext.curKvSplitPart = 1U;
706+ }
707+ 
708+ // 更新S2切分信息
709+ if (assignContext.curS2Idx > assignContext.s1GCache.s2Start &&
710+ assignContext.curS2Idx <= assignContext.s1GCache.s2End) {
711+ assignContext.curKvSplitPart++;
712+ }
713+ }
714+ 
715+ result.usedCoreNum = assignContext.curCoreIdx + 1;
716+}
717+ 
718+void SparseFlashAttentionAntiquantMetaDataCpuKernel::CopyTmpResult(SplitResult &tmpRes, SplitResult &splitRes)
719+{
720+ uint64_t len = tmpRes.bN2End.size();
721+ splitRes.usedCoreNum = tmpRes.usedCoreNum;
722+ splitRes.maxCost = tmpRes.maxCost;
723+ splitRes.numOfFdHead = tmpRes.numOfFdHead;
724+ splitRes.maxS2SplitNum = tmpRes.maxS2SplitNum;
725+ 
726+ for (size_t i = 0; i < len; ++i) {
727+ splitRes.bN2End[i] = tmpRes.bN2End[i];
728+ splitRes.gS1End[i] = tmpRes.gS1End[i];
729+ splitRes.s2End[i] = tmpRes.s2End[i];
730+ 
731+ splitRes.fdRes.bN2IdxOfFdHead[i] = tmpRes.fdRes.bN2IdxOfFdHead[i];
732+ splitRes.fdRes.gS1IdxOfFdHead[i] = tmpRes.fdRes.gS1IdxOfFdHead[i];
733+ splitRes.fdRes.s2SplitNumOfFdHead[i] = tmpRes.fdRes.s2SplitNumOfFdHead[i];
734+ splitRes.fdRes.s2SplitStartIdxOfCore[i] = tmpRes.fdRes.s2SplitStartIdxOfCore[i];
735+ splitRes.fdRes.gS1SplitNumOfFdHead[i] = tmpRes.fdRes.gS1SplitNumOfFdHead[i];
736+ splitRes.fdRes.gS1LastPartSizeOfFdHead[i] = tmpRes.fdRes.gS1LastPartSizeOfFdHead[i];
737+ }
738+}
739+ 
740+void SparseFlashAttentionAntiquantMetaDataCpuKernel::ClearTmpResult(SplitResult &tmpResult)
741+{
742+ uint64_t len = tmpResult.bN2End.size();
743+ tmpResult.usedCoreNum = 0U;
744+ tmpResult.maxCost = 0;
745+ tmpResult.numOfFdHead = 0U;
746+ tmpResult.maxS2SplitNum = 0U;
747+ tmpResult.usedVecNumOfFd = 0U;
748+ 
749+ for (size_t i = 0; i < len; ++i) {
750+ tmpResult.bN2End[i] = 0U;
751+ tmpResult.gS1End[i] = 0U;
752+ tmpResult.s2End[i] = 0U;
753+ tmpResult.fdRes.bN2IdxOfFdHead[i] = 0U;
754+ tmpResult.fdRes.gS1IdxOfFdHead[i] = 0U;
755+ tmpResult.fdRes.s2SplitNumOfFdHead[i] = 0U;
756+ tmpResult.fdRes.s2SplitStartIdxOfCore[i] = 0U;
757+ tmpResult.fdRes.gS1SplitNumOfFdHead[i] = 0U;
758+ tmpResult.fdRes.gS1LastPartSizeOfFdHead[i] = 0U;
759+ }
760+}
761+ 
762+void SparseFlashAttentionAntiquantMetaDataCpuKernel::RollBackCursor(const SplitContext &splitContext,
763+ const CostInfo &costInfo, SplitResult &splitRes)
764+{
765+ for (size_t i = 0; i < splitRes.usedCoreNum; ++i) {
766+ // x, y, z
767+ if (splitRes.s2End[i] > 0U) {
768+ splitRes.s2End[i] = splitRes.s2End[i] - 1U;
769+ continue;
770+ }
771+ uint32_t bIdx = splitRes.bN2End[i] / kvHeadNum_;
772+ // x, y, 0
773+ if (splitRes.gS1End[i] > 0U) {
774+ splitRes.gS1End[i] = splitRes.gS1End[i] - 1U;
775+ splitRes.s2End[i] = splitContext.splitInfo.s2BaseNum[bIdx] - 1U;
776+ continue;
777+ }
778+ 
779+ // x, 0, 0
780+ uint32_t bN2Idx = splitRes.bN2End[i] > 0U ? splitRes.bN2End[i] - 1U : 0U;
781+ bIdx = bN2Idx / kvHeadNum_;
782+ 
783+ // last end point
784+ if (i == splitRes.usedCoreNum - 1U && costInfo.bN2BlockOfEachBatch[bIdx] == 0U) {
785+ splitRes.bN2End[i] = batchSize_ * kvHeadNum_ - 1;
786+ splitRes.gS1End[i] = splitContext.splitInfo.s1GBaseNum[bIdx] > 0 ?
787+ splitContext.splitInfo.s1GBaseNum[bIdx] - 1U : 0U;
788+ splitRes.s2End[i] = splitContext.splitInfo.s2BaseNum[bIdx] > 0 ?
789+ splitContext.splitInfo.s2BaseNum[bIdx] - 1U : 0U;
790+ continue;
791+ }
792+ 
793+ while (bN2Idx > 0U && costInfo.bN2BlockOfEachBatch[bIdx] == 0U) {
794+ bN2Idx -= 1U;
795+ bIdx = bN2Idx / kvHeadNum_;
796+ }
797+ 
798+ if (costInfo.bN2BlockOfEachBatch[bIdx] != 0U) {
799+ splitRes.bN2End[i] = bN2Idx;
800+ splitRes.gS1End[i] = splitContext.splitInfo.s1GBaseNum[bIdx] - 1U;
801+ splitRes.s2End[i] = splitContext.splitInfo.s2BaseNum[bIdx] - 1U;
802+ } else {
803+ splitRes.bN2End[i] = 0U;
804+ splitRes.gS1End[i] = 0U;
805+ splitRes.s2End[i] = 0U;
806+ }
807+ }
808+}
809+ 
810+void SparseFlashAttentionAntiquantMetaDataCpuKernel::SplitFD(SplitResult &result)
811+{
812+ uint32_t totalFDLoad = 0;
813+ uint32_t totalFDHeadSplit = 0;
814+ // 计算FD的总数据量
815+ for (uint32_t i = 0; i < result.numOfFdHead; i++) {
816+ totalFDLoad += result.fdRes.s2SplitNumOfFdHead[i] * result.fdRes.gS1SplitNumOfFdHead[i];
817+ totalFDHeadSplit += result.fdRes.gS1SplitNumOfFdHead[i];
818+ }
819+ // 基于FA开核数量，计算每个Vector需要计算的FD数据量
820+ // FD均衡的最小单位为一个归约任务的一个split，所以最多占用totalFDHeadSplit个vector
821+ uint32_t maxVectorNum = std::min(totalFDHeadSplit, result.usedCoreNum * result.vecCubeRatio);
822+ double loadThrOfVector = static_cast<double>(totalFDLoad) / static_cast<double>(maxVectorNum); // 初始化vector的负载上限
823+ int64_t loadOfCurVector = 0;
824+ uint32_t curCoreIndex = 0;
825+ uint32_t preTmpFDIndexEndOfFdHead = 0;
826+ uint32_t preTmpFDIndexEndOfFdHeadSplit = 0;
827+ for (uint32_t i = 0; i < result.numOfFdHead; i++) {
828+ uint32_t fDKVSplitNum = result.fdRes.s2SplitNumOfFdHead[i];
829+ for (uint32_t gS1SplitIdx = 0; gS1SplitIdx < result.fdRes.gS1SplitNumOfFdHead[i]; gS1SplitIdx++) {
830+ double remainSpace = loadThrOfVector - static_cast<double>(loadOfCurVector); // 计算当前vector剩余负载空间
831+ // 判断是否放在当前vector的标准是剩余空间是否能容纳一半当前归约块
832+ if (fDKVSplitNum > remainSpace * FD_TOLERANCE_RATIO) {
833+ result.fdRes.gS1IdxEndOfFdHead[curCoreIndex] = preTmpFDIndexEndOfFdHead;
834+ result.fdRes.gS1IdxEndOfFdHeadSplit[curCoreIndex] = preTmpFDIndexEndOfFdHeadSplit;
835+ curCoreIndex += 1U;
836+ totalFDLoad -= static_cast<uint32_t>(loadOfCurVector); // 当前未分配的总负载
837+ // 根据剩余负载和剩余可用vector更新负载上限，保证最后一个vector能分配所有负载
838+ loadThrOfVector = static_cast<double>(totalFDLoad) / static_cast<double>(maxVectorNum - curCoreIndex);
839+ loadOfCurVector = 0;
840+ }
841+ loadOfCurVector += fDKVSplitNum;
842+ preTmpFDIndexEndOfFdHead = i;
843+ preTmpFDIndexEndOfFdHeadSplit = gS1SplitIdx;
844+ }
845+ }
846+ result.fdRes.gS1IdxEndOfFdHead[curCoreIndex] = preTmpFDIndexEndOfFdHead;
847+ result.fdRes.gS1IdxEndOfFdHeadSplit[curCoreIndex] = preTmpFDIndexEndOfFdHeadSplit;
848+ result.usedVecNumOfFd = curCoreIndex + 1;
849+}
850+ 
851+bool SparseFlashAttentionAntiquantMetaDataCpuKernel::BalanceSchedule() {
852+ SplitContext splitContext(batchSize_);
853+ 
854+ // 1、划分基本块，统计信息
855+ CalcSplitInfo(splitContext);
856+ // 全空case
857+ if (splitContext.splitInfo.isKvSeqAllZero) {
858+ splitRes_.usedCoreNum = 1U;
859+ splitRes_.bN2End[0] = batchSize_ * kvHeadNum_ - 1;
860+ splitRes_.gS1End[0] = 0U;
861+ splitRes_.s2End[0] = 0U;
862+ return true;
863+ }
864+ CalcCostInfo(splitContext);
865+ 
866+ // 2、获取每个核的分配方案
867+ uint32_t maxCore = std::min(coreNum_, splitContext.costInfo.totalBlockNum);
868+ uint32_t minCore = static_cast<uint32_t>(
869+ std::sqrt(static_cast<float>(splitContext.costInfo.totalBlockNum) + 0.25f) + 0.5f);
870+ minCore = std::min(minCore, maxCore);
871+ 
872+ splitRes_.maxCost = INT64_MAX;
873+ splitRes_.usedCoreNum = 1U;
874+ SplitResult tmpResult {coreNum_, aivCoreNum_ / aicCoreNum_}; // C: V = 1: 2, TODO: C: V = 1: 1 ?
875+ for (uint32_t i = minCore; i <= maxCore; ++i) {
876+ CalcSplitPlan(i, splitRes_.maxCost, splitContext, tmpResult);
877+ if (tmpResult.maxCost < splitRes_.maxCost) {
878+ CopyTmpResult(tmpResult, splitRes_);
879+ }
880+ ClearTmpResult(tmpResult);
881+ }
882+ 
883+ // 3、存在FD任务，对FD进行负载均衡分配
884+ if (splitRes_.numOfFdHead > 0U) {
885+ SplitFD(splitRes_);
886+ }
887+ splitRes_.usedCoreNum = std::max(splitRes_.usedCoreNum, 1U); // 至少使用1个core
888+ RollBackCursor(splitContext, splitContext.costInfo, splitRes_);
889+ return true;
890+}
891+ 
892+bool SparseFlashAttentionAntiquantMetaDataCpuKernel::GenMetaData() {
893+ optiling::detail::SfaMetaData* metaDataPtr = (optiling::detail::SfaMetaData*)metaData_->GetData();
894+ metaDataPtr->usedCoreNum = splitRes_.usedCoreNum;
895+ metaDataPtr->numOfFdHead = splitRes_.numOfFdHead;
896+ metaDataPtr->usedVecNumOfFd = splitRes_.usedVecNumOfFd;
897+ metaDataPtr->mBaseSize = mBaseSize_;
898+ metaDataPtr->s2BaseSize = s2BaseSize_;
899+ metaDataPtr->gS1BaseSizeOfFd = gS1BaseSizeOfFd_;
900+ 
901+ for (size_t i = 0; i < coreNum_; ++i) {
902+ metaDataPtr->bN2End[i] = splitRes_.bN2End[i];
903+ metaDataPtr->gS1End[i] = splitRes_.gS1End[i];
904+ metaDataPtr->s2End[i] = splitRes_.s2End[i];
905+ metaDataPtr->fdRes.bN2IdxOfFdHead[i] = splitRes_.fdRes.bN2IdxOfFdHead[i];
906+ metaDataPtr->fdRes.gS1IdxOfFdHead[i] = splitRes_.fdRes.gS1IdxOfFdHead[i];
907+ metaDataPtr->fdRes.s2SplitNumOfFdHead[i] = splitRes_.fdRes.s2SplitNumOfFdHead[i];
908+ metaDataPtr->fdRes.s2SplitStartIdxOfCore[i] = splitRes_.fdRes.s2SplitStartIdxOfCore[i];
909+ metaDataPtr->fdRes.gS1SplitNumOfFdHead[i] = splitRes_.fdRes.gS1SplitNumOfFdHead[i];
910+ metaDataPtr->fdRes.gS1LastPartSizeOfFdHead[i] = splitRes_.fdRes.gS1LastPartSizeOfFdHead[i];
911+ }
912+ if (splitRes_.numOfFdHead > 0U) {
913+ for (size_t i = 0; i < coreNum_ * 2U; ++i) {
914+ metaDataPtr->fdRes.gS1IdxEndOfFdHead[i] = splitRes_.fdRes.gS1IdxEndOfFdHead[i];
915+ metaDataPtr->fdRes.gS1IdxEndOfFdHeadSplit[i] = splitRes_.fdRes.gS1IdxEndOfFdHeadSplit[i];
916+ }
917+ }
918+ return true;
919+}
920+ 
921+static const char *sfaKernelType = "KvQuantSparseFlashAttentionMetadata";
922+REGISTER_CPU_KERNEL(sfaKernelType, SparseFlashAttentionAntiquantMetaDataCpuKernel);
923+ 
924+}; // namespace aicpu
925+ 
Aexperimental/attention/sparse_flash_attention_antiquant_metadata/op_kernel_aicpu/sparse_flash_attention_antiquant_metadata_aicpu.h+303-0
@@ -0,0 +1,303 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file sparse_flash_attention_antiquant_metadata_aicpu.h
13+ * \brief
14+ */
15+ 
16+#ifndef SPARSE_FLASH_ATTENTION_ANTIQUANT_METADATA_AICPU_H
17+#define SPARSE_FLASH_ATTENTION_ANTIQUANT_METADATA_AICPU_H
18+ 
19+#include "cpu_context.h"
20+#include "cpu_kernel.h"
21+#include "cpu_tensor.h"
22+#include <string>
23+#include <vector>
24+#include <array>
25+ 
26+namespace aicpu {
27+constexpr int64_t FA_TOLERANCE_RATIO = 2;
28+constexpr uint32_t FD_TOLERANCE_RATIO = 2U;
29+constexpr int64_t V0_COST = 250;
30+ 
31+enum BlockType : uint32_t {
32+ NORMAL_BLOCK = 0,
33+ TAIL_BLOCK,
34+ BLOCK_MAX_TYPE
35+};
36+ 
37+enum class SparseMode : uint8_t {
38+ DEFAULT_MASK = 0,
39+ ALL_MASK,
40+ LEFT_UP_CAUSAL,
41+ RIGHT_DOWN_CAUSAL,
42+ BAND,
43+ SPARSE_BUTT,
44+};
45+ 
46+template<class T>
47+using Range = std::pair<T, T>;
48+ 
49+template<class T>
50+using BlockCost = std::array<std::array<T, static_cast<size_t>(BLOCK_MAX_TYPE)>, static_cast<size_t>(BLOCK_MAX_TYPE)>;
51+ 
52+template<typename T>
53+T Clip(T value, T minValue, T maxValue)
54+{
55+ if (value < minValue) {
56+ return minValue;
57+ }
58+ if (value > maxValue) {
59+ return maxValue;
60+ }
61+ return value;
62+}
63+ 
64+template<typename T>
65+inline bool IsWithinTolerance(T limit, T tolerance, T value)
66+{
67+ return limit + tolerance >= value;
68+}
69+ 
70+// 分核功能模块输出：FD信息，包含需要归约的数据索引及其分核信息
71+struct FlashDecodeResult {
72+ // 1、归约任务的索引信息
73+ std::vector<uint32_t> bN2IdxOfFdHead {}; // 每个归约任务的BN2索引，脚标为归约任务的序号，最大为核数-1
74+ std::vector<uint32_t> gS1IdxOfFdHead {}; // 每个归约任务的GS1索引，脚标为归约任务的序号
75+ std::vector<uint32_t> s2SplitNumOfFdHead {}; // 每个归约任务的S2核间切分份数，脚标为归约任务的序号
76+ // 2、FD负载均衡阶段，归约任务的分核（vec）信息
77+ std::vector<uint32_t> gS1SplitNumOfFdHead {}; // 每个归约任务m轴切分份数，脚标为归约任务的序号
78+ std::vector<uint32_t> gS1LastPartSizeOfFdHead {}; // 每个归约任务m轴切分的最后一份的大小，脚标为归约任务的序号
79+ std::vector<uint32_t> gS1IdxEndOfFdHead {}; // FD负载均衡阶段，每个vector的一级索引，脚标为vector ID，值为归约任务的ID
80+ std::vector<uint32_t> gS1IdxEndOfFdHeadSplit {}; // FD负载均衡阶段，每个vector的二级索引，脚标为vector ID，值为归约任务的m轴切分ID
81+ // 3、每个core处理的第1个归约任务的数据应存放的workspace位置
82+ std::vector<uint32_t> s2SplitStartIdxOfCore {};
83+ 
84+ FlashDecodeResult(uint32_t coreNum, uint32_t vecCubeRatio) :
85+ bN2IdxOfFdHead(coreNum),
86+ gS1IdxOfFdHead(coreNum),
87+ s2SplitNumOfFdHead(coreNum),
88+ gS1SplitNumOfFdHead(coreNum),
89+ gS1LastPartSizeOfFdHead(coreNum),
90+ gS1IdxEndOfFdHead(coreNum * vecCubeRatio),
91+ gS1IdxEndOfFdHeadSplit(coreNum * vecCubeRatio),
92+ s2SplitStartIdxOfCore(coreNum) {}
93+};
94+ 
95+// 分核功能模块输出：FA阶段的核间分核信息
96+struct SplitResult {
97+ uint32_t usedCoreNum { 0U }; // 使用的核数量
98+ uint32_t vecCubeRatio { 0U }; // vec 与 cube 核数比例
99+ std::vector<uint32_t> bN2End {}; // 每个核处理数据的BN2结束点
100+ std::vector<uint32_t> gS1End {}; // 每个核处理数据的GS1结束点
101+ std::vector<uint32_t> s2End {}; // 每个核处理数据的S2结束点
102+ int64_t maxCost { 0 }; // 慢核开销
103+ uint32_t numOfFdHead { 0U }; // 归约任务数量
104+ uint32_t maxS2SplitNum { 0U }; // 单个归约任务最大分核数量
105+ uint32_t usedVecNumOfFd { 0U }; // 归约过程使用的vector数量
106+ FlashDecodeResult fdRes { 0U, 0U }; // FD信息
107+ 
108+ SplitResult(uint32_t coreNum, uint32_t ratio) :
109+ bN2End(coreNum),
110+ vecCubeRatio(ratio),
111+ gS1End(coreNum),
112+ s2End(coreNum),
113+ fdRes(coreNum, ratio) {};
114+};
115+ 
116+// 分核功能模块内部使用：记录切分信息
117+struct SplitInfo {
118+ std::vector<uint32_t> s1GBaseNum {}; // S1G方向，切了多少个基本块
119+ std::vector<uint32_t> s2BaseNum {}; // S2方向，切了多少个基本块
120+ std::vector<uint32_t> s1GTailSize {}; // S1G方向，尾块size
121+ std::vector<uint32_t> s2TailSize {}; // S2方向，尾块size
122+ bool isKvSeqAllZero { true };
123+ 
124+ explicit SplitInfo(uint32_t batchSize) :
125+ s1GBaseNum(batchSize),
126+ s2BaseNum(batchSize),
127+ s1GTailSize(batchSize),
128+ s2TailSize(batchSize) {}
129+};
130+ 
131+// 分核功能模块内部使用：记录batch的开销信息
132+struct CostInfo {
133+ std::vector<int64_t> bN2CostOfEachBatch {}; // 整个batch的开销
134+ std::vector<uint32_t> bN2BlockOfEachBatch {}; // 整个batch的开销
135+ std::vector<int64_t> bN2LastBlockCostOfEachBatch {}; // batch最后一块的开销
136+ uint32_t totalBlockNum { 0U };
137+ int64_t totalCost { 0 };
138+ 
139+ explicit CostInfo(uint32_t batchSize) :
140+ bN2CostOfEachBatch(batchSize),
141+ bN2BlockOfEachBatch(batchSize),
142+ bN2LastBlockCostOfEachBatch(batchSize) {}
143+};
144+ 
145+// 分核功能模块内部使用：分核过程中，case基本信息的上下文信息，组合以减少接口传参数量
146+struct SplitContext {
147+ SplitInfo splitInfo { 0U };
148+ CostInfo costInfo { 0U };
149+ 
150+ explicit SplitContext(uint32_t batchSize) :
151+ splitInfo(batchSize),
152+ costInfo(batchSize) {}
153+};
154+ 
155+// 分核功能模块内部使用：记录batch相关的临时信息
156+struct BatchCache {
157+ uint32_t bIdx { 0U };
158+ uint32_t s1Size { 0U };
159+ uint32_t s2Size { 0U };
160+ int64_t preTokenLeftUp { 0 };
161+ int64_t nextTokenLeftUp { 0 };
162+ BlockCost<int64_t> typeCost {};
163+};
164+ 
165+// 分核功能模块内部使用：记录当前行（S1G）的临时信息
166+struct S1GCache {
167+ uint32_t bIdx { 0U };
168+ uint32_t s1GIdx { 0U };
169+ uint32_t s2Start { 0U };
170+ uint32_t s2End { 0U };
171+ int64_t s1GCost { 0 };
172+ int64_t s1GLastBlockCost { 0 };
173+ uint32_t s1GBlock { 0U };
174+ int64_t s1GNormalBlockCost { 0 };
175+};
176+ 
177+// 分核功能模块内部使用：记录分配过程中，当前核的负载信息
178+struct CoreCache {
179+ int64_t costLimit { 0 }; // 负载上限
180+ int64_t cost { 0 }; // 已分配负载
181+ uint32_t block { 0U }; // 已分配块数
182+};
183+ 
184+// 分核功能模块内部使用：记录分配过程中的上下文信息
185+struct AssignContext {
186+ uint32_t curBIdx { 0U };
187+ uint32_t curBN2Idx { 0U };
188+ uint32_t curS1GIdx { 0U };
189+ uint32_t curS2Idx { 0U };
190+ uint32_t curCoreIdx { 0U };
191+ int64_t unassignedCost { 0 };
192+ uint32_t usedCoreNum { 0U };
193+ uint32_t curKvSplitPart { 1U };
194+ 
195+ int64_t bN2Cost { 0 };
196+ uint32_t bN2Block { 0U };
197+ bool isFinished { false };
198+ BatchCache batchCache {};
199+ S1GCache s1GCache {};
200+ CoreCache coreCache {};
201+};
202+class SparseFlashAttentionAntiquantMetaDataCpuKernel : public CpuKernel {
203+public:
204+ SparseFlashAttentionAntiquantMetaDataCpuKernel() = default;
205+ ~SparseFlashAttentionAntiquantMetaDataCpuKernel() = default;
206+ uint32_t Compute(CpuKernelContext &ctx) override;
207+ 
208+private:
209+ bool Prepare(CpuKernelContext &ctx);
210+ bool ParamsCheck();
211+ bool ParamsInit();
212+ bool BalanceSchedule();
213+ bool GenMetaData();
214+ 
215+ // util
216+ uint32_t GetS1SeqSize(uint32_t bIdx);
217+ uint32_t GetS2SeqSize(uint32_t bIdx);
218+ uint32_t GetSparseSeqSize(uint32_t bIdx);
219+ int64_t CalcPreTokenLeftUp(uint32_t s1Size, uint32_t s2Size);
220+ int64_t CalcNextTokenLeftUp(uint32_t s1Size, uint32_t s2Size);
221+ Range<uint32_t> CalcS2Range(uint32_t s1GIdx,const BatchCache &batchCache);
222+ int64_t CalcCost(uint32_t basicM, uint32_t basicS2);
223+ BlockCost<int64_t> CalcCostTable(uint32_t s1NormalSize, uint32_t s2NormalSize, uint32_t s1GTailSize,
224+ uint32_t s2TailSize);
225+ 
226+ // cache calculation
227+ void CalcBatchCache(uint32_t bIdx, const SplitContext &splitContext, BatchCache &batchCache);
228+ void CalcS1GCache(uint32_t s1GIdx, const SplitContext &splitContext, const BatchCache &batchCache, S1GCache &s1GCache);
229+ void CopyTmpResult(SplitResult &tmpRes, SplitResult &splitRes);
230+ void ClearTmpResult(SplitResult &tmpRes);
231+ 
232+ // preprocess
233+ void CalcSplitInfo(SplitContext &splitContext);
234+ void CalcBatchCost(uint32_t bIdx, const SplitContext &splitContext, CostInfo &costInfo);
235+ void CalcCostInfo(SplitContext &splitContext);
236+ 
237+ // assign
238+ void UpdateCursor(const SplitContext &splitContext, AssignContext &assignContext);
239+ void AssignByBatch(const SplitContext &splitContext, AssignContext &assignContext);
240+ void AssignByRow(const SplitContext &splitContext, AssignContext &assignContext);
241+ void AssignByBlock(const SplitContext &splitContext, AssignContext &assignContext);
242+ void ForceAssign(const SplitContext &splitContext, AssignContext &assignContext);
243+ 
244+ // FD
245+ bool IsNeedRecordFDInfo(const AssignContext &assignContext, const SplitResult &splitRes);
246+ void RecordFDInfo(const SplitContext &splitContext, const AssignContext &assignContext, SplitResult &result);
247+ 
248+ // main
249+ void SplitFD(SplitResult &result);
250+ void CalcSplitPlan(uint32_t coreNum, int64_t costLimit, const SplitContext &splitContext, SplitResult &result);
251+ void SplitCore();
252+ void RollBackCursor(const SplitContext &splitContext, const CostInfo &costInfo, SplitResult &splitRes);
253+ 
254+private:
255+ CpuKernelContext* context_ = nullptr;
256+ // input
257+ Tensor *actSeqLenQ_ = nullptr;
258+ Tensor *actSeqLenKV_ = nullptr;
259+ Tensor *sparseSeqLenKV_ = nullptr;
260+ // output
261+ Tensor *metaData_ = nullptr;
262+ // attributes
263+ uint32_t aicCoreNum_ = 24U;
264+ uint32_t aivCoreNum_ = 48U;
265+ uint32_t batchSize_ = 0;
266+ uint32_t querySeqSize_ = 0;
267+ uint32_t queryHeadNum_ = 0;
268+ uint32_t kvSeqSize_ = 0;
269+ uint32_t kvHeadNum_ = 0;
270+ uint32_t headDim_ = 0;
271+ uint32_t topKSize_ = 0;
272+ uint32_t sparseBlockSize_ = 0;
273+ uint32_t sparseBlockCount_ = 0; // new
274+ std::string layoutQuery_ = "BSND";
275+ std::string layoutKV_ = "BSND";
276+ uint32_t sparseMode_ = 0;
277+ uint32_t attentionMode_ = 0;
278+ uint32_t ropeHeadDim_ = 0;
279+ uint32_t sparseSharedSize_ = 0;
280+ 
281+ // SplitParams
282+ uint32_t coreNum_ = 24U; // new
283+ int64_t preToken_ = 0; // new
284+ int64_t nextToken_ = 0; // new
285+ uint32_t groupSize_ = 0;
286+ uint32_t mBaseSize_ = 0;
287+ uint32_t s2BaseSize_ = 0;
288+ uint32_t gS1BaseSizeOfFd_ = 0;
289+ bool isS1G_ = true;
290+ SplitResult splitRes_ {24, 2};
291+ 
292+private:
293+ enum class ParamId : uint32_t {
294+ // input
295+ actSeqLenQ = 0,
296+ actSeqLenKV = 1,
297+ sparseSeqLenKV = 2,
298+ // output
299+ metaData = 0,
300+ };
301+};
302+} // namespace aicpu
303+#endif
Aexperimental/attention/sparse_flash_attention_antiquant_metadata/op_kernel_aicpu/sparse_flash_attention_antiquant_metadata_aicpu.json+15-0
@@ -0,0 +1,15 @@
1+{
2+ "KvQuantSparseFlashAttentionMetadata":{
3+ "opInfo":{
4+ "computeCost":"100",
5+ "engine":"DNN_VM_AICPU",
6+ "flagAsync":"False",
7+ "flagPartial":"False",
8+ "functionName":"RunCpuKernel",
9+ "kernelSo":"libtransformer_aicpu_kernels.so",
10+ "opKernelLib":"CUSTAICPUKernel",
11+ "userDefined":"True",
12+ "workspaceSize":"100"
13+ }
14+ }
15+}
Mscripts/ci/parse_changed_ops.py+5-1
@@ -32,7 +32,11 @@ BlackList = {
32 "moe_distribute_combine_shmem",32 "moe_distribute_combine_shmem",
33 "moe_distribute_dispatch_shmem",33 "moe_distribute_dispatch_shmem",
34 "rope_matrix",34 "rope_matrix",
35- "all_gather_matmul_v2"35+ "all_gather_matmul_v2",
36+ "quant_sals_indexer",
37+ "quant_sals_indexer_metadata",
38+ "sparse_flash_attention_antiquant",
39+ "sparse_flash_attention_antiquant_metadata"
36 }40 }
37 41 
38 42