diff --git a/test/python_test/RegisterOps.cpp b/test/python_test/RegisterOps.cpp index a9717b9..f61ff5e 100644 --- a/test/python_test/RegisterOps.cpp +++ b/test/python_test/RegisterOps.cpp @@ -89,6 +89,34 @@ at::Tensor x_attention_impl_npu(const at::Tensor& query, return attnOut; } +at::Tensor x_attention_v2_impl_npu(const at::Tensor& query, + const at::Tensor& key_cache, + const at::Tensor& value_cache, + const at::Tensor& unshared_key, + const at::Tensor& unshared_value, + const c10::optional& shared_block_tables, + const c10::optional& unshared_block_tables, + const at::Tensor& actual_shared_kvlen, + const at::Tensor& decode_step, + double scale_value = 0.0) { + at::Tensor attnOut = at::empty_like(query); + + EXEC_NPU_CMD(aclnnXAttentionV2, + query, + key_cache, + value_cache, + unshared_key, + unshared_value, + unshared_block_tables, + actual_shared_kvlen, + decode_step, + shared_block_tables, + scale_value, + attnOut); + + return attnOut; +} + std::tuple beam_search_impl_npu( const at::Tensor& log_probs, const at::Tensor& top_tokens, @@ -1987,6 +2015,7 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("select_unshared_kv", &select_unshared_kv_impl_npu, "select_unshared_kv"); m.def("cache_unshared_kv", &cache_unshared_kv_impl_npu, "cache_unshared_kv"); m.def("x_attention", &x_attention_impl_npu, "x_attention"); + m.def("x_attention_v2", &x_attention_v2_impl_npu, "x_attention_v2"); m.def("beam_search", &beam_search_impl_npu, "beam_search"); m.def("beam_search_group", &beam_search_group_impl_npu, "beam_search_group"); m.def("beam_search_rec_final_select", diff --git a/test/python_test/custom_ops.py b/test/python_test/custom_ops.py index 16e100c..0a14521 100644 --- a/test/python_test/custom_ops.py +++ b/test/python_test/custom_ops.py @@ -15,6 +15,17 @@ def x_attention_npu(query, key_cache, value_cache, unshared_key, unshared_value, return custom_ops_lib.x_attention(query, key_cache, value_cache, unshared_key, unshared_value, shared_block_tables, unshared_block_tables, actual_shared_kvlen, decode_step, scale_value) + +# x_attention_v2 +def x_attention_v2_npu(query, key_cache, value_cache, unshared_key, unshared_value, actual_shared_kvlen, decode_step, + shared_block_tables = None, + unshared_block_tables = None, + scale_value = None): + if scale_value is None: + scale_value = 0.0 + return custom_ops_lib.x_attention_v2(query, key_cache, value_cache, unshared_key, unshared_value, + shared_block_tables, unshared_block_tables, actual_shared_kvlen, decode_step, scale_value) + # reshape cache kv def cache_unshared_kv_npu(x_key_block, x_value_block, curr_key, curr_value, block_table, decode_step): return custom_ops_lib.cache_unshared_kv(x_key_block, x_value_block, curr_key, curr_value, block_table, decode_step) diff --git a/test/python_test/test_x_attention_v2.py b/test/python_test/test_x_attention_v2.py new file mode 100644 index 0000000..bb78581 --- /dev/null +++ b/test/python_test/test_x_attention_v2.py @@ -0,0 +1,575 @@ +#!/usr/bin/env python3 +# Copyright 2025 The xLLM Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== + +import os +import random +import copy +import pytest +import torch + +from dataclasses import dataclass +from ml_dtypes import bfloat16 + + +torch_npu = pytest.importorskip("torch_npu") +custom_ops = pytest.importorskip("custom_ops") + +WORKSPACE = os.path.dirname(os.path.abspath(__file__)) +torch.manual_seed(1) + + +def gen_seqlen(max_q_seqlen: int, max_kv_seqlen: int, is_varied_len: int, batch: int): + q_seqlen_list = [] + kv_seqlen_list = [] + if is_varied_len == 0: + q_seqlen_list = [max_q_seqlen] * batch + kv_seqlen_list = [max_kv_seqlen] * batch + else: + for _ in range(batch): + q_seq = random.randint(1, max_q_seqlen) + kv_seq = random.randint(1, max_kv_seqlen) + q_seqlen_list.append(q_seq) + kv_seqlen_list.append(kv_seq) + return q_seqlen_list, kv_seqlen_list + + +class TestFlashAttentionInfer: + @dataclass + class AttentionInputs: + # [bs, 1, headnum, headdim] + query: torch.Tensor + # shared_kv_type=1: [num_blocks, block_size, kv_head, headdim] + # shared_kv_type=0: [num_shared_kv, kv_head, headdim] + key_cache: torch.Tensor + value_cache: torch.Tensor + # unshared_kv_type=0: [bs (request_num * beam_size), kv_head, max_decode_step, headdim] + # unshared_kv_type=1: [max_request_num, beam_size, kv_head, max_decode_step, headdim] + unshared_k: torch.Tensor + unshared_v: torch.Tensor + # shared_kv_type=1: [request_num, max_blocks_per_batch] max_blocks_per_batch = ceil(max_shared_kvlen / block_size) + # shared_kv_type=0: None + block_tables: list + # unshared_kv_type=1: [request_num] list of block indices for each request + # unshared_kv_type=0: None + unshared_block_tables: list + # [request_num] (1, 1, 1) + q_seqlen_list: list + # [request_num] (share_len1, share_len2) + k_seqlen_list: list + global_mask: any + mask_type: int + shape_param: any + + @dataclass + class GenDataParams: + q_seqlen_list: list + k_seqlen_list: list + beam_size: int + unshared_kvlen: int + num_heads: int + kv_heads: int + head_size: int + num_blocks: int + block_size: int + mask_type: int + dtype: any + shared_kv_type: int + unshared_kv_type: int + + @classmethod + def group_matmul(cls, head, kv_head, left, right, right_row=None, right_col=None): + group_num = head // kv_head + score = None + for i in range(kv_head): + if right_row is None: + current_right = right[i : (i + 1), :, :] + else: + current_right = right[i : (i + 1), :right_row, :right_col] + left_group = left[i * group_num : (i + 1) * group_num, :, :] + group_score = torch.matmul(left_group.to(torch.float32), current_right.to(torch.float32)) + score = group_score if score is None else torch.cat((score, group_score), dim=0) + return score + + @classmethod + def softmax_numpy(cls, sim): + row_max = torch.max(sim, dim=-1, keepdim=True).values + sim_sub = sim - row_max + sim_sub = torch.exp(sim_sub) + row_sum = torch.sum(sim_sub, dim=-1, keepdim=True) + soft_res = sim_sub # no div rowsum + return soft_res, row_max, row_sum + + def ref_masked_attention(self, + query, # (q_seqlen, num_heads, head_size) + key, # (k_seqlen, kv_heads, head_size) + value, + scale: float, + mask # (q_seqlen, k_seqlen) + ): + # Q * K.T + query = query.permute(1, 0, 2) + key = key.permute(1, 2, 0) + sim_high = self.group_matmul(query.shape[0], key.shape[0], query, key) # (head_num, q_seqlen, k_seqlen) + sim_high = sim_high * scale + # softmax + p_high, gm, gl = self.softmax_numpy(sim_high) + p = p_high.to(query.dtype) + p_high = p_high.to(torch.float32) + value = value.permute(1, 0, 2) + out_high = self.group_matmul(query.shape[0], key.shape[0], p_high, value) + out = self.group_matmul(query.shape[0], key.shape[0], p, value) + out_high = out_high.permute(1, 0, 2) + out = out.permute(1, 0, 2) + out = out.to(query.dtype) + return out, out_high, gm, gl + + def ref_single_query_unshared_kv_attention(self, + attention_inputs: "TestFlashAttentionInfer.AttentionInputs", + output: torch.Tensor, + true_out: torch.Tensor, + unshared_gl: torch.Tensor, + unshared_gm: torch.Tensor, + ) -> None: + num_heads = attention_inputs.shape_param.num_heads + kv_heads = attention_inputs.shape_param.kv_heads + head_size = attention_inputs.shape_param.head_size + beam_size = attention_inputs.shape_param.beam_size + request_num = len(attention_inputs.q_seqlen_list) + batch = beam_size * request_num + decode_step = attention_inputs.shape_param.unshared_kvlen + unshared_kv_type = attention_inputs.shape_param.unshared_kv_type + max_decode_step = attention_inputs.unshared_k.shape[-2] if len(attention_inputs.unshared_k.shape) == 5 else attention_inputs.unshared_k.shape[2] + + scale = 1.0 / (head_size ** 0.5) + + if unshared_kv_type == 0: + # Continuous format: [batch, kv_heads, max_decode_step, head_size] + assert attention_inputs.query.shape == (batch, num_heads, head_size) + assert attention_inputs.unshared_k.shape == (batch, kv_heads, max_decode_step, head_size) + assert attention_inputs.unshared_v.shape == (batch, kv_heads, max_decode_step, head_size) + + for i in range(batch): + q = attention_inputs.query[i : i + 1, :, :] + k = attention_inputs.unshared_k[i, :, :, :] + v = attention_inputs.unshared_v[i, :, :, :] + # Transpose for group_matmul + q_t = q.permute(1, 0, 2) + k_t = k.permute(0, 2, 1) + + sim = self.group_matmul(num_heads, kv_heads, q_t, k_t, head_size, decode_step) # [num_heads, 1, unshared_kvlen] + sim = sim * scale + + # Softmax with stats + p, gm, gl = self.softmax_numpy(sim) + gm = gm.permute(1, 0, 2) # (q_seqlen, num_heads, 1) + gl = gl.permute(1, 0, 2) + p_high = p.to(torch.float32) + out_high = self.group_matmul(num_heads, kv_heads, p_high, v, decode_step, head_size) + out_high = out_high.permute(1, 0, 2) + p_low = p.to(attention_inputs.query.dtype) + out_low = self.group_matmul(num_heads, kv_heads, p_low, v, decode_step, head_size) + out_low = out_low.permute(1, 0, 2) + out_low = out_low.to(attention_inputs.query.dtype) + + # Write outputs + output[i : i + 1, :, :] = out_low + true_out[i : i + 1, :, :] = out_high + + unshared_gm[i, :, :] = gm[:, :, :] + unshared_gl[i, :, :] = gl[:, :, :] + else: + # Paged format: [max_request_num, beam_size, kv_heads, max_decode_step, head_size] + assert attention_inputs.query.shape == (batch, num_heads, head_size) + assert len(attention_inputs.unshared_k.shape) == 5 + assert attention_inputs.unshared_k.shape == attention_inputs.unshared_v.shape + max_request_num = attention_inputs.unshared_k.shape[0] + + # Use unshared_block_tables to map request to cache index + for req_idx in range(request_num): + # Get the cache index for this request from unshared_block_tables + if attention_inputs.unshared_block_tables is not None and len(attention_inputs.unshared_block_tables) > req_idx: + cache_idx = attention_inputs.unshared_block_tables[req_idx][0] # First block index for this request + else: + # Fallback: use request index directly + cache_idx = req_idx + + for beam_idx in range(beam_size): + i = req_idx * beam_size + beam_idx + q = attention_inputs.query[i : i + 1, :, :] + + # Get unshared_k and unshared_v from paged format using cache_idx + k = attention_inputs.unshared_k[cache_idx, beam_idx, :, :, :] # [kv_heads, max_decode_step, head_size] + v = attention_inputs.unshared_v[cache_idx, beam_idx, :, :, :] # [kv_heads, max_decode_step, head_size] + + # Transpose for group_matmul + q_t = q.permute(1, 0, 2) + k_t = k.permute(0, 2, 1) # [kv_heads, head_size, max_decode_step] + + sim = self.group_matmul(num_heads, kv_heads, q_t, k_t, head_size, decode_step) # [num_heads, 1, unshared_kvlen] + sim = sim * scale + + # Softmax with stats + p, gm, gl = self.softmax_numpy(sim) + gm = gm.permute(1, 0, 2) # (q_seqlen, num_heads, 1) + gl = gl.permute(1, 0, 2) + p_high = p.to(torch.float32) + out_high = self.group_matmul(num_heads, kv_heads, p_high, v, decode_step, head_size) + out_high = out_high.permute(1, 0, 2) + p_low = p.to(attention_inputs.query.dtype) + out_low = self.group_matmul(num_heads, kv_heads, p_low, v, decode_step, head_size) + out_low = out_low.permute(1, 0, 2) + out_low = out_low.to(attention_inputs.query.dtype) + + # Write outputs + output[i : i + 1, :, :] = out_low + true_out[i : i + 1, :, :] = out_high + + unshared_gm[i, :, :] = gm[:, :, :] + unshared_gl[i, :, :] = gl[:, :, :] + + def ref_single_query_shared_kv_attention( + self, + attention_inputs: "TestFlashAttentionInfer.AttentionInputs", + output, + true_out, + shared_gl, + shared_gm, + ) -> None: + num_heads = attention_inputs.shape_param.num_heads + kv_heads = attention_inputs.shape_param.kv_heads + head_size_qk = attention_inputs.shape_param.head_size + head_size_vo = attention_inputs.shape_param.head_size + block_size = attention_inputs.shape_param.block_size + beam_size = attention_inputs.shape_param.beam_size + request_num = len(attention_inputs.shape_param.q_seqlen_list) + shared_kv_type = attention_inputs.shape_param.shared_kv_type + cu_seqlen = 0 + kv_seqlen_now = 0 + layout = "TND" + + for i in range(request_num): + q_seqlen = int(beam_size) + k_seqlen = int(attention_inputs.k_seqlen_list[i]) + if layout == "TND": + q = attention_inputs.query[cu_seqlen : (cu_seqlen + q_seqlen), :, :] + elif layout == "BSND": + q = attention_inputs.query[i, :, :, :] + keys = [] + values = [] + if shared_kv_type == 1: + block_table = attention_inputs.block_tables[i] + for j in range(k_seqlen): + block_number = int(block_table[j // block_size]) + block_offset = j % block_size + + k = attention_inputs.key_cache[block_number, block_offset, :, :] + k = k.reshape(kv_heads, head_size_qk) + keys.append(k) + + v = attention_inputs.value_cache[block_number, block_offset, :, :] + v = v.reshape(kv_heads, head_size_vo) + values.append(v) + else: + for j in range(k_seqlen): + k = attention_inputs.key_cache[kv_seqlen_now + j, :, :] + # k shape is already (kv_heads, head_size_qk) + keys.append(k) + + v = attention_inputs.value_cache[kv_seqlen_now + j, :, :] + # v shape is already (kv_heads, head_size_vo) + values.append(v) + + keys = torch.stack(keys, axis=0) + values = torch.stack(values, axis=0) + scale = 1.0 / (head_size_qk ** 0.5) + mask = None + out, out_high, gm, gl = self.ref_masked_attention(q, keys, values, scale, mask) + out = out.reshape(-1, num_heads, head_size_vo) + out_high = out_high.reshape(-1, num_heads, head_size_vo) + gm = gm.permute(1, 0, 2) # (q_seqlen, num_heads, 1) + gl = gl.permute(1, 0, 2) + output[cu_seqlen : cu_seqlen + q_seqlen, :, :] = out + true_out[cu_seqlen : cu_seqlen + q_seqlen, :, :] = out_high + shared_gl[cu_seqlen : cu_seqlen + q_seqlen, :, :] = gl + shared_gm[cu_seqlen : cu_seqlen + q_seqlen, :, :] = gm + cu_seqlen += q_seqlen + kv_seqlen_now += k_seqlen + + def call_device_op(self, attention_inputs: "TestFlashAttentionInfer.AttentionInputs", + q, k, v, unshared_k, unshared_v, + block_tables, unshared_block_tables, + actual_shared_kvlen, decode_step): + + shared_kv_type = attention_inputs.shape_param.shared_kv_type + unshared_kv_type = attention_inputs.shape_param.unshared_kv_type + + q = q.npu() + k = k.npu() + v = v.npu() + unshared_k = unshared_k.npu() + unshared_v = unshared_v.npu() + + if shared_kv_type == 1: + block_tables = torch.tensor(copy.deepcopy(block_tables), dtype=torch.int32).npu() + else: + block_tables = None + + if unshared_kv_type == 1: + unshared_block_tables = torch.tensor(copy.deepcopy(unshared_block_tables), dtype=torch.int32).npu() + else: + unshared_block_tables = None + + actual_shared_kvlen = torch.tensor(actual_shared_kvlen, dtype=torch.int32).npu() + decode_step_tensor = torch.tensor([decode_step], dtype=torch.int32).npu() + + # ========== DEBUG ========== + print("="*80) + print(f" q shape: {q.shape}") + print(f" k shape: {k.shape if k is not None else None}") + print(f" v shape: {v.shape if v is not None else None}") + print(f" unshared_k shape: {unshared_k.shape}") + print(f" unshared_v shape: {unshared_v.shape}") + print(f" block_tables shape: {block_tables.shape if block_tables is not None else None}") + print(f" unshared_block_tables shape: {unshared_block_tables.shape if unshared_block_tables is not None else None}") + print(f" actual_shared_kvlen shape: {actual_shared_kvlen.shape}") + print(f" decode_step_tensor shape: {decode_step_tensor.shape}") + print("="*80 + "\n") + + attn_out = custom_ops.x_attention_v2_npu( + q, k, v, unshared_k, unshared_v, actual_shared_kvlen, decode_step_tensor, + block_tables, unshared_block_tables + ) + return attn_out + + def calc_data(self, gen_data_params: "TestFlashAttentionInfer.GenDataParams"): + head_size_qk = gen_data_params.head_size + head_size_vo = gen_data_params.head_size + q_min_range = -1.0 + q_max_range = 1.0 + kv_min_range = -1.0 + kv_max_range = 1.0 + beam_size = gen_data_params.beam_size + request_num = len(gen_data_params.k_seqlen_list) + decode_step = gen_data_params.unshared_kvlen + max_decode_step = 3 + + num_tokens = sum(gen_data_params.q_seqlen_list) * beam_size + num_shared_kv = sum(gen_data_params.k_seqlen_list) + + batch_size = request_num * beam_size + torch_dtype = gen_data_params.dtype + query = (torch.empty((num_tokens, gen_data_params.num_heads, head_size_qk), dtype=torch_dtype) + .uniform_(q_min_range, q_max_range)) + max_k_seqlen = max(gen_data_params.k_seqlen_list) + block_tables = [] # (request_num, max_num_blocks_per_seq) + key_cache = None + value_cache = None + + # Generate shared KV cache based on shared_kv_type + if gen_data_params.shared_kv_type == 1: + # Paged format: [num_blocks, block_size, kv_heads, head_dim] + key_cache = (torch.empty( + (gen_data_params.num_blocks, gen_data_params.block_size, gen_data_params.kv_heads, head_size_qk), + dtype=torch_dtype + ).uniform_(kv_min_range, kv_max_range)) + value_cache = (torch.empty( + (gen_data_params.num_blocks, gen_data_params.block_size, gen_data_params.kv_heads, head_size_vo), + dtype=torch_dtype + ).uniform_(kv_min_range, kv_max_range)) + max_num_blocks_per_seq = (max_k_seqlen + gen_data_params.block_size - 1) // gen_data_params.block_size + for i in range(request_num): + block_table = [max_num_blocks_per_seq * i + j for j in range(max_num_blocks_per_seq)] + block_tables.append(block_table) + else: + # Continuous format: [num_shared_kv, kv_heads, head_dim] + key_cache = (torch.empty( + (num_shared_kv, gen_data_params.kv_heads, head_size_qk), + dtype=torch_dtype + ).uniform_(kv_min_range, kv_max_range)) + + value_cache = (torch.empty( + (num_shared_kv, gen_data_params.kv_heads, head_size_vo), + dtype=torch_dtype + ).uniform_(kv_min_range, kv_max_range)) + block_tables = None + + # Generate unshared KV cache based on unshared_kv_type + unshared_key = None + unshared_value = None + unshared_block_tables = None + + if gen_data_params.unshared_kv_type == 0: + # Continuous format: [request_num * beam_size, kv_heads, max_decode_step, head_dim] + unshared_key = (torch.empty( + (batch_size, gen_data_params.kv_heads, max_decode_step, head_size_qk), dtype=torch_dtype + ).uniform_(kv_min_range, kv_max_range)) + unshared_value = (torch.empty( + (batch_size, gen_data_params.kv_heads, max_decode_step, head_size_vo), dtype=torch_dtype + ).uniform_(kv_min_range, kv_max_range)) + else: + # Paged format: [max_request_num, beam_size, kv_heads, max_decode_step, head_dim] + # Use request_num as max_request_num for simplicity (can be larger in real scenarios) + max_request_num = request_num + unshared_key = (torch.empty( + (max_request_num, beam_size, gen_data_params.kv_heads, max_decode_step, head_size_qk), dtype=torch_dtype + ).uniform_(kv_min_range, kv_max_range)) + unshared_value = (torch.empty( + (max_request_num, beam_size, gen_data_params.kv_heads, max_decode_step, head_size_vo), dtype=torch_dtype + ).uniform_(kv_min_range, kv_max_range)) + + # Generate unshared_block_tables + # Each request has a block table mapping to its unshared KV cache + # For simplicity, each request maps to its own index in the paged cache + unshared_block_tables = [] + for i in range(request_num): + # Each request maps to its own index (i) in the paged cache + unshared_block_tables.append([request_num - 1 - i]) + + shape_out = (num_tokens, gen_data_params.num_heads, head_size_vo) + sum_max_shape_out = (num_tokens, gen_data_params.num_heads, 1) + shared_ref_out = torch.zeros(shape_out, dtype=torch_dtype) + shared_true_out = torch.zeros(shape_out, dtype=torch.float32) + shared_gl = torch.zeros(sum_max_shape_out, dtype=torch.float32) + shared_gm = torch.zeros(sum_max_shape_out, dtype=torch.float32) + + unshared_ref_out = torch.zeros(shape_out, dtype=torch_dtype) + unshared_true_out = torch.zeros(shape_out, dtype=torch.float32) + unshared_gl = torch.zeros(sum_max_shape_out, dtype=torch.float32) + unshared_gm = torch.zeros(sum_max_shape_out, dtype=torch.float32) + + attention_inputs = self.AttentionInputs( + query, + key_cache, + value_cache, + unshared_key, + unshared_value, + block_tables, + unshared_block_tables, + gen_data_params.q_seqlen_list, + gen_data_params.k_seqlen_list, + None, + gen_data_params.mask_type, + gen_data_params, + ) + + self.ref_single_query_shared_kv_attention( + attention_inputs, shared_ref_out, shared_true_out, shared_gl, shared_gm + ) + + self.ref_single_query_unshared_kv_attention( + attention_inputs, unshared_ref_out, unshared_true_out, unshared_gl, unshared_gm + ) + + gm = torch.maximum(shared_gm, unshared_gm) + update_shared_expgm = torch.exp(shared_gm - gm) + update_unshared_expgm = torch.exp(unshared_gm - gm) + gl = shared_gl * update_shared_expgm + unshared_gl * update_unshared_expgm + tmp_shared_true = shared_true_out * update_shared_expgm + tmp_unshared_true = unshared_true_out * update_unshared_expgm + tmp_add = tmp_shared_true + tmp_unshared_true + final_true_out = tmp_add / gl + + # Prepare actual_shared_kvlen for device op + actual_shared_kvlen = gen_data_params.k_seqlen_list + + npu_res = self.call_device_op( + attention_inputs, query, key_cache, value_cache, unshared_key, unshared_value, + block_tables, unshared_block_tables, actual_shared_kvlen, decode_step + ) + golden_res = final_true_out + npu_res = npu_res.cpu().float() + assert torch.allclose(npu_res, golden_res, atol=0.001, rtol=0.001) + +@pytest.mark.parametrize("dtype,request_num,beam_size,q_seqlen,kv_seqlen,unshared_seqlen,num_head,kv_heads,embedding_size,block_size,is_varied_len,mask_type,shared_kv_type,unshared_kv_type", [ + (torch.bfloat16, 1, 128, 1, 1024, 2, 32, 8, 128, 128, 0, 0, 0, 1), + (torch.bfloat16, 1, 256, 1, 1024, 2, 32, 8, 128, 128, 0, 0, 0, 1), + (torch.bfloat16, 1, 512, 1, 1024, 2, 32, 8, 128, 128, 0, 0, 0, 1), + (torch.bfloat16, 1, 1024, 1, 1024, 2, 32, 8, 128, 128, 0, 0, 0, 1), + (torch.bfloat16, 1, 2048, 1, 1024, 2, 32, 8, 128, 128, 0, 0, 0, 1), + (torch.bfloat16, 1, 4096, 1, 1024, 2, 32, 8, 128, 128, 0, 0, 0, 1), + (torch.bfloat16, 2, 128, 1, 1024, 2, 32, 8, 128, 128, 0, 0, 0, 1), + (torch.bfloat16, 2, 256, 1, 1024, 2, 32, 8, 128, 128, 0, 0, 0, 1), + (torch.bfloat16, 2, 512, 1, 1024, 2, 32, 8, 128, 128, 0, 0, 0, 1), + (torch.bfloat16, 2, 1024, 1, 1024, 2, 32, 8, 128, 128, 0, 0, 0, 1), + (torch.bfloat16, 2, 2048, 1, 1024, 2, 32, 8, 128, 128, 0, 0, 0, 1), + (torch.bfloat16, 2, 4096, 1, 1024, 2, 32, 8, 128, 128, 0, 0, 0, 1), + (torch.bfloat16, 2, 128, 1, 2048, 2, 32, 8, 128, 128, 0, 0, 0, 1), + (torch.bfloat16, 2, 256, 1, 2048, 2, 32, 8, 128, 128, 0, 0, 0, 1), + (torch.bfloat16, 2, 512, 1, 2048, 2, 32, 8, 128, 128, 0, 0, 0, 1), + (torch.bfloat16, 2, 1024, 1, 2048, 2, 32, 8, 128, 128, 0, 0, 0, 1), + (torch.bfloat16, 2, 2048, 1, 2048, 2, 32, 8, 128, 128, 0, 0, 0, 1), + (torch.bfloat16, 2, 4096, 1, 2048, 2, 32, 8, 128, 128, 0, 0, 0, 1), + # (torch.bfloat16, 8, 128, 1, 2048, 2, 32, 8, 128, 128, 0, 0, 0, 1), + # (torch.bfloat16, 8, 256, 1, 2048, 2, 32, 8, 128, 128, 0, 0, 0, 1), + # (torch.bfloat16, 8, 512, 1, 2048, 2, 32, 8, 128, 128, 0, 0, 0, 1), + # (torch.bfloat16, 8, 1024, 1, 2048, 2, 32, 8, 128, 128, 0, 0, 0, 1), + # (torch.bfloat16, 8, 2048, 1, 2048, 2, 32, 8, 128, 128, 0, 0, 0, 1), + # (torch.bfloat16, 8, 4096, 1, 2048, 2, 32, 8, 128, 128, 0, 0, 0, 1), + # (torch.bfloat16, 16, 128, 1, 2048, 2, 32, 8, 128, 128, 0, 0, 0, 1), + # (torch.bfloat16, 16, 256, 1, 2048, 2, 32, 8, 128, 128, 0, 0, 0, 1), + # (torch.bfloat16, 16, 512, 1, 2048, 2, 32, 8, 128, 128, 0, 0, 0, 1), + # (torch.bfloat16, 16, 1024, 1, 2048, 2, 32, 8, 128, 128, 0, 0, 0, 1), + # (torch.bfloat16, 16, 2048, 1, 2048, 2, 32, 8, 128, 128, 0, 0, 0, 1), + # (torch.bfloat16, 16, 4096, 1, 2048, 2, 32, 8, 128, 128, 0, 0, 0, 1), + # (torch.bfloat16, 32, 128, 1, 2048, 2, 32, 8, 128, 128, 0, 0, 0, 1), + # (torch.bfloat16, 32, 256, 1, 2048, 2, 32, 8, 128, 128, 0, 0, 0, 1), + # (torch.bfloat16, 32, 512, 1, 2048, 2, 32, 8, 128, 128, 0, 0, 0, 1), + # (torch.bfloat16, 32, 1024, 1, 2048, 2, 32, 8, 128, 128, 0, 0, 0, 1), + # (torch.bfloat16, 32, 2048, 1, 2048, 2, 32, 8, 128, 128, 0, 0, 0, 1), + # (torch.bfloat16, 32, 4096, 1, 2048, 2, 32, 8, 128, 128, 0, 0, 0, 1), +]) +def test_x_attention_v2_npu(dtype,request_num,beam_size,q_seqlen,kv_seqlen,unshared_seqlen,num_head,kv_heads,embedding_size,block_size,is_varied_len,mask_type,shared_kv_type,unshared_kv_type): + # Device selection (skip if no NPU available) + try: + torch_npu.npu.set_device(0) + except Exception as e: + pytest.skip(f"NPU device not available: {e}") + + # request = 5 + # beam_size = 512 # must >= 128 + # q_seqlen = 1 # must be 1 + # kv_seqlen = 4090 # shared_kv_len + # unshared_seqlen = 2 + # num_head = 8 + # kv_heads = 8 + # embedding_size = 128 + # block_size = 128 + # is_varied_len = 0 + # mask_type = 0 + # shared_kv_type = 1 + # unshared_kv_type = 0 + + q_seqlen_list, kv_seqlen_list = gen_seqlen(q_seqlen, kv_seqlen, is_varied_len, request_num) + max_kv_seqlen = max(kv_seqlen_list) + num_blocks = request_num * ((max_kv_seqlen + block_size - 1) // block_size) + + test_obj = TestFlashAttentionInfer() + gen_data_params = test_obj.GenDataParams( + q_seqlen_list, + kv_seqlen_list, + beam_size, + unshared_seqlen, + num_head, + kv_heads, + embedding_size, + num_blocks, + block_size, + mask_type, + dtype, + shared_kv_type, + unshared_kv_type + ) + test_obj.calc_data(gen_data_params) diff --git a/xllm_ops/build_aclnn.sh b/xllm_ops/build_aclnn.sh index 8eb12f7..4918d27 100644 --- a/xllm_ops/build_aclnn.sh +++ b/xllm_ops/build_aclnn.sh @@ -335,6 +335,7 @@ elif [[ "$SOC_VERSION" =~ ^ascend950 ]]; then # ### JD's in-house operators #### "beam_search_group" "x_attention" + "x_attention_v2" "cache_unshared_kv" "causal_conv1d" "causal_conv1d_qkv" diff --git a/xllm_ops/x_attention_v2/CMakeLists.txt b/xllm_ops/x_attention_v2/CMakeLists.txt new file mode 100644 index 0000000..86b3082 --- /dev/null +++ b/xllm_ops/x_attention_v2/CMakeLists.txt @@ -0,0 +1,19 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +file(GLOB CURRENT_DIRS RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*) +if(NOT ENABLE_TEST AND NOT BENCHMARK) + list(REMOVE_ITEM CURRENT_DIRS tests) +endif() +foreach(SUB_DIR ${CURRENT_DIRS}) + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/${SUB_DIR}/CMakeLists.txt") + add_subdirectory(${SUB_DIR}) + endif() +endforeach() \ No newline at end of file diff --git a/xllm_ops/x_attention_v2/op_host/CMakeLists.txt b/xllm_ops/x_attention_v2/op_host/CMakeLists.txt new file mode 100644 index 0000000..5c39b6e --- /dev/null +++ b/xllm_ops/x_attention_v2/op_host/CMakeLists.txt @@ -0,0 +1,29 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- +add_op_to_compiled_list() + +if (BUILD_OPEN_PROJECT) + target_sources(op_host_aclnn PRIVATE + x_attention_v2_def.cpp + ) +endif() + +add_ops_compile_options( + OP_NAME XAttentionV2 + OPTIONS --cce-auto-sync=on + -Wno-deprecated-declarations + -Werror + -I${CANN_3RD_LIB_PATH}/catlass/include + -I${CMAKE_CURRENT_LIST_DIR}/../../../ +) + +if (NOT BUILD_OPS_RTY_KERNEL) + add_modules_sources(OPTYPE x_attention_v2 ACLNNTYPE aclnn) +endif() \ No newline at end of file diff --git a/xllm_ops/x_attention_v2/op_host/x_attention_v2_def.cpp b/xllm_ops/x_attention_v2/op_host/x_attention_v2_def.cpp new file mode 100644 index 0000000..e678a0d --- /dev/null +++ b/xllm_ops/x_attention_v2/op_host/x_attention_v2_def.cpp @@ -0,0 +1,81 @@ +/* Copyright 2025 The xLLM Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://gitcode.com/xLLM-AI/xllm_ops/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "register/op_def_registry.h" + +namespace ops { +class XAttentionV2 : public OpDef { +public: + explicit XAttentionV2(const char* name) : OpDef(name) + { + this->Input("query") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT16, ge::DT_BF16}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND}); + this->Input("shared_key_block") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT16, ge::DT_BF16}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND}); + this->Input("shared_value_block") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT16, ge::DT_BF16}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND}); + this->Input("unshared_key_block") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT16, ge::DT_BF16}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND}); + this->Input("unshared_value_block") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT16, ge::DT_BF16}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND}); + this->Input("unshared_block_table") + .ParamType(OPTIONAL) + .DataType({ge::DT_INT32, ge::DT_INT32}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND}); + this->Input("shared_kv_lens") + .ParamType(REQUIRED) + .DataType({ge::DT_INT32, ge::DT_INT32}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND}); + this->Input("decode_step") + .ParamType(REQUIRED) + .DataType({ge::DT_INT32, ge::DT_INT32}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND}); + this->Input("shared_block_table") + .ParamType(OPTIONAL) + .DataType({ge::DT_INT32, ge::DT_INT32}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND}); + this->Output("attn_out") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT16, ge::DT_BF16}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND}); + this->Attr("scale_value").AttrType(OPTIONAL).Float(0.0); + + this->AICore().AddConfig("ascend950"); + + } +}; + +OP_ADD(XAttentionV2); +} diff --git a/xllm_ops/x_attention_v2/op_host/x_attention_v2_proto.cpp b/xllm_ops/x_attention_v2/op_host/x_attention_v2_proto.cpp new file mode 100644 index 0000000..bf920d1 --- /dev/null +++ b/xllm_ops/x_attention_v2/op_host/x_attention_v2_proto.cpp @@ -0,0 +1,39 @@ +/* Copyright 2025 The xLLM Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://gitcode.com/xLLM-AI/xllm_ops/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "x_attention_v2_tiling.h" +#include "register/op_def_registry.h" +#include "tiling/platform/platform_ascendc.h" + + +namespace ge { +static ge::graphStatus InferShape(gert::InferShapeContext* context) +{ + const gert::Shape* x1_shape = context->GetInputShape(0); + gert::Shape* y_shape = context->GetOutputShape(0); + *y_shape = *x1_shape; + return GRAPH_SUCCESS; +} +static ge::graphStatus InferDataType(gert::InferDataTypeContext *context) +{ + const auto inputDataType = context->GetInputDataType(0); + context->SetOutputDataType(0, inputDataType); + return ge::GRAPH_SUCCESS; +} + +IMPL_OP_INFERSHAPE(XAttentionV2) + .InferShape(InferShape) + .InferDataType(InferDataType); +} diff --git a/xllm_ops/x_attention_v2/op_host/x_attention_v2_tiling.cpp b/xllm_ops/x_attention_v2/op_host/x_attention_v2_tiling.cpp new file mode 100644 index 0000000..7be4d02 --- /dev/null +++ b/xllm_ops/x_attention_v2/op_host/x_attention_v2_tiling.cpp @@ -0,0 +1,375 @@ +/* Copyright 2025 The xLLM Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://gitcode.com/xLLM-AI/xllm_ops/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "x_attention_v2_tiling.h" +#include "register/op_def_registry.h" +#include "tiling/platform/platform_ascendc.h" +#include + +#define OP_LOGE(nodeName, fmt, ...) \ + printf(fmt, ##__VA_ARGS__); \ + printf("\n") + + +namespace optiling { + +enum InputIndex { + QUERY = 0, + SHARED_KEY_BLOCK, + SHARED_VALUE_BLOCK, + UNSHARED_KEY_BLOCK, + UNSHARED_VALUE_BLOCK, + UNSHARED_BLOCK_TABLE, + SHARED_KV_LENS, + DECODE_STEP, + SHARED_BLOCK_TABLE, +}; + +constexpr int32_t NUM2 = 2; +constexpr int32_t NUM3 = 3; +constexpr int32_t NUM4 = 4; +constexpr int32_t UNSHARED_Q_TILE = 128; +constexpr int32_t UNSHARED_KV_TILE = 256; +constexpr uint32_t Q_S_BLOCK_TILE = 128; +constexpr uint32_t BLOCK_SIZE = 128; +constexpr int32_t WORKSPACE_BLOCK_SIZE_DB = 128 * 128 * 4; // row * col * blockStackNum +constexpr int32_t UNSHARED_WORKSPACE_BLOCK_SIZE_DB = 128 * 256; // unshared no pinpong +constexpr int32_t FLOAT_BLOCK_SIZE = 8; +constexpr int32_t SCALE_VALUE_ATTR_INDEX = 0; + +class TilingXAttentionV2Func { + public: + explicit TilingXAttentionV2Func(gert::TilingContext* tiling_context) + : tiling_context_(tiling_context) {} + ge::graphStatus RunTiling(); + private: + uint64_t GetTilingKey() const; + private: + XAttentionV2TilingData tiling_data_; + gert::TilingContext* tiling_context_ = nullptr; + uint32_t sharedBlockDim = 0; + uint32_t unsharedBlockDim = 0; + uint32_t cubeCoreNum; + uint32_t vecCoreNum; + bool isSharedPaged{false}; + bool isUnsharedPaged{false}; + uint32_t inputDtype{0}; + ge::graphStatus ParseInputShapeAndAttrs(); + ge::graphStatus FillBasicTilingData(); + ge::graphStatus FillBasicTilingData4NewKind(); + void FillSharedSplitCoreTilingData(); + void FillUnsharedSplitCoreTilingData(); + void FillCombineScaleTilingData(); + void BalanceAicore(); + void SetWorkspaces(); + uint32_t GetQNBlockTile(int64_t qSeqlen, uint32_t groupSize); + +}; + + +ge::graphStatus TilingXAttentionV2Func::FillBasicTilingData4NewKind() +{ + auto queryShape = tiling_context_->GetInputShape(QUERY)->GetStorageShape(); + // shared [total_num_tokens, kv_head, head_dim] + auto sharedKeyBlockShape = tiling_context_->GetInputShape(SHARED_KEY_BLOCK)->GetStorageShape(); + // unshared [max_request_num, beamsize, kv_head, max_decode_step, head_dim] + // unshared_blk_tb [bs, request_idx] + auto unsharedKeyBlockShape = tiling_context_->GetInputShape(UNSHARED_KEY_BLOCK)->GetStorageShape(); + auto unsharedBlockTableShape = tiling_context_->GetOptionalInputShape(UNSHARED_BLOCK_TABLE)->GetStorageShape(); + + int32_t numTokens = queryShape.GetDim(0); + int32_t qHeadNum = queryShape.GetDim(1); + int32_t embeddingSize = queryShape.GetDim(2); + int32_t batch = unsharedBlockTableShape.GetDim(0); + int32_t kvHeadNum = sharedKeyBlockShape.GetDim(1); + int32_t maxDecodeStep = unsharedKeyBlockShape.GetDim(NUM3); + int32_t beamSize = numTokens / batch; + + float scaleValue = static_cast(1.0 / std::sqrt(1.0 * embeddingSize)); + auto attrs = tiling_context_->GetAttrs(); + if (attrs != nullptr) { + const auto* attr_scale_value = attrs->GetAttrPointer(SCALE_VALUE_ATTR_INDEX); + if (attr_scale_value != nullptr && *attr_scale_value > 0.0f) { + scaleValue = *attr_scale_value; + } + } + + // set tiling data + tiling_data_.set_batch(batch); + tiling_data_.set_numHeads(qHeadNum); + tiling_data_.set_kvHeads(kvHeadNum); + tiling_data_.set_embeddingSize(embeddingSize); + tiling_data_.set_beamSize(beamSize); + tiling_data_.set_scaleValue(scaleValue); + tiling_data_.set_maskType(0); + tiling_data_.set_blockSize(BLOCK_SIZE); + tiling_data_.set_numTokens(numTokens); + tiling_data_.set_maxDecodeStep(maxDecodeStep); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus TilingXAttentionV2Func::ParseInputShapeAndAttrs() +{ + auto dType = tiling_context_->GetInputTensor(InputIndex::QUERY)->GetDataType(); + if (dType == ge::DT_FLOAT16) { + inputDtype = 0; + } else if (dType == ge::DT_BF16) { + inputDtype = 1; + } + auto sharedBlockTableShapePtr = tiling_context_->GetOptionalInputShape(InputIndex::SHARED_BLOCK_TABLE); + auto unsharedBlockTableShapePtr = tiling_context_->GetOptionalInputShape(InputIndex::UNSHARED_BLOCK_TABLE); + isSharedPaged = (sharedBlockTableShapePtr != nullptr); + isUnsharedPaged = (unsharedBlockTableShapePtr != nullptr); + if (isSharedPaged && !isUnsharedPaged) { + return FillBasicTilingData(); + } else if (!isSharedPaged && isUnsharedPaged) { + return FillBasicTilingData4NewKind(); + } + OP_LOGE(tiling_context_->GetNodeName(), "unexpected input combination between shared_block_table and unshared_block_table."); + return ge::GRAPH_FAILED; +} + +void TilingXAttentionV2Func::BalanceAicore() +{ + // Support dynamic calculation based on the amount of computation in the future. + sharedBlockDim = 12; + unsharedBlockDim = cubeCoreNum - sharedBlockDim; + return; +} + +uint32_t TilingXAttentionV2Func::GetQNBlockTile(int64_t qSeqlen, uint32_t groupSize) +{ + uint32_t qRowNumCeil = 128; + // A trick is used to ensure the qN tile is a even number, + // thus most tasks have balanced workload between two vec cores, + // and each vec core possess no more than 64 rows when all-rounded row num is no larger than 128, + // aiding the coding of rescale block + uint32_t qNBlockTile = (qRowNumCeil / qSeqlen) / 2 * 2; + qNBlockTile = std::min(qNBlockTile, groupSize); + qNBlockTile = std::max(qNBlockTile, static_cast(1)); + // The current shared x_attention kernel has accuracy error on qNBlockTile != 1. + // Limit the qNBlockTile parameter to 1 to ensure correct kernel results until the issue is fixed. + qNBlockTile = 1; + return qNBlockTile; +} + + +void TilingXAttentionV2Func::FillUnsharedSplitCoreTilingData() +{ + tiling_data_.set_unsharedCoreNum(unsharedBlockDim); + tiling_data_.set_groupSize(tiling_data_.get_numHeads() / tiling_data_.get_kvHeads()); + uint32_t totalGroupCount = tiling_data_.get_beamSize() * tiling_data_.get_kvHeads(); + // for no PA scenario, calculation can cross batch + if (!isUnsharedPaged) { + totalGroupCount *= tiling_data_.get_batch(); + } + uint32_t maxGroupCountPerLoop = std::min(UNSHARED_Q_TILE / tiling_data_.get_groupSize(), + UNSHARED_KV_TILE / tiling_data_.get_maxDecodeStep()); + // ensure each task handles same group count + while (maxGroupCountPerLoop > 1 && + (totalGroupCount % maxGroupCountPerLoop != 0 || maxGroupCountPerLoop % FLOAT_BLOCK_SIZE != 0)) + --maxGroupCountPerLoop; + tiling_data_.set_unshareGroupCountPerLoop(maxGroupCountPerLoop); + uint32_t totalTaskNum = totalGroupCount / maxGroupCountPerLoop; + if (isUnsharedPaged) { + tiling_data_.set_unsharedLoopCountPerBatch(totalTaskNum); + totalTaskNum *= tiling_data_.get_batch(); + } + uint32_t unsharedFullCoreNum = unsharedBlockDim; + uint32_t unsharedTaskNumHead = totalTaskNum / unsharedBlockDim; + uint32_t unsharedTaskNumTail = unsharedTaskNumHead; + uint32_t remainTask = totalTaskNum % unsharedBlockDim; + if (remainTask != 0) { + unsharedFullCoreNum = remainTask; + unsharedTaskNumHead += 1; + } + tiling_data_.set_unsharedFullCoreNum(unsharedFullCoreNum); + tiling_data_.set_unsharedTaskNumHead(unsharedTaskNumHead); + tiling_data_.set_unsharedTaskNumTail(unsharedTaskNumTail); + +} + + +void TilingXAttentionV2Func::SetWorkspaces() +{ + auto platform_info = + platform_ascendc::PlatformAscendC(tiling_context_->GetPlatformInfo()); + size_t systemWorkspaceSize = static_cast(platform_info.GetLibApiWorkSpaceSize()); + size_t userWorkspaceSize = 0; + + uint64_t qoSize = tiling_data_.get_numTokens() + * tiling_data_.get_numHeads() + * tiling_data_.get_embeddingSize() + * sizeof(int16_t); + // Attention occupied space + // TODO: Only apply for one temporary space, affecting preload function, long sequence scenario needs extra processing + uint64_t mm1OutSize = (sharedBlockDim * WORKSPACE_BLOCK_SIZE_DB + + unsharedBlockDim * UNSHARED_WORKSPACE_BLOCK_SIZE_DB) * NUM3 * sizeof(float);; + uint64_t smOnlineOutSize = (sharedBlockDim * WORKSPACE_BLOCK_SIZE_DB + + unsharedBlockDim * UNSHARED_WORKSPACE_BLOCK_SIZE_DB) * NUM3 * sizeof(int16_t); + uint64_t mm2OutSize = sharedBlockDim * WORKSPACE_BLOCK_SIZE_DB * NUM3 * sizeof(float); + // oUpdate currently not used + uint64_t updateSize = 0; // sharedBlockDim * WORKSPACE_BLOCK_SIZE_DB * NUM3 * sizeof(float); + tiling_data_.set_mm1OutSize(mm1OutSize); + tiling_data_.set_smOnlineOutSize(smOnlineOutSize); + tiling_data_.set_mm2OutSize(mm2OutSize); + tiling_data_.set_updateSize(updateSize); + + // combine required output occupied space + uint64_t sumMaxSize = tiling_data_.get_numTokens() * tiling_data_.get_numHeads() * sizeof(float) * NUM2; + uint64_t attnOutSize = qoSize * 2; + uint64_t combineWorkspaceSize = sumMaxSize * FLOAT_BLOCK_SIZE + attnOutSize; + uint64_t unsharedcombineWorkspaceSize = sumMaxSize + attnOutSize; + tiling_data_.set_sharedWorkspaceSize(combineWorkspaceSize); // new line + + userWorkspaceSize = mm1OutSize + smOnlineOutSize + mm2OutSize + updateSize + combineWorkspaceSize + + unsharedcombineWorkspaceSize; + size_t* workspace = tiling_context_->GetWorkspaceSizes(1); + workspace[0] = systemWorkspaceSize + userWorkspaceSize; +} + +void TilingXAttentionV2Func::FillCombineScaleTilingData() +{ + uint32_t rowNum = tiling_data_.get_batch() * + tiling_data_.get_beamSize() * + tiling_data_.get_numHeads(); + uint32_t columnSize = tiling_data_.get_embeddingSize(); + + uint32_t rowNumPerCore = rowNum / cubeCoreNum; // number of rows per core + uint32_t rowNumTailPerCore = rowNum % cubeCoreNum; // remaining rows, need to be allocated to the first few cores + tiling_data_.set_combineFormerCoreNum(rowNumTailPerCore); + tiling_data_.set_combineFormerRowNum(rowNumPerCore + 1); + tiling_data_.set_combineTailRowNum(rowNumPerCore); + tiling_data_.set_combineCoreNum(cubeCoreNum); +} + +void TilingXAttentionV2Func::FillSharedSplitCoreTilingData() +{ + uint32_t totalTaskNum = 0; + uint32_t groupSize = tiling_data_.get_numHeads() / tiling_data_.get_kvHeads(); + int64_t qSeqlen = tiling_data_.get_beamSize(); + uint32_t curQNBlockTile = GetQNBlockTile(qSeqlen, groupSize); + uint32_t qNBlockNumPerGroup = (groupSize + curQNBlockTile - 1) / curQNBlockTile; + uint32_t curQNBlockNum = qNBlockNumPerGroup * tiling_data_.get_kvHeads(); + uint32_t curQSBlockTile = Q_S_BLOCK_TILE; + uint32_t curQSBlockNum = (qSeqlen + curQSBlockTile - 1) / curQSBlockTile; + uint32_t curTaskNum = curQNBlockNum * curQSBlockNum; + uint32_t firstSharedBatchTaskNum = curTaskNum; + totalTaskNum = curTaskNum * tiling_data_.get_batch(); + tiling_data_.set_firstSharedBatchTaskNum(firstSharedBatchTaskNum); + tiling_data_.set_sharedTotalTaskNum(totalTaskNum); + tiling_data_.set_sharedCoreNum(sharedBlockDim); +} + + +ge::graphStatus TilingXAttentionV2Func::FillBasicTilingData() +{ + auto queryShape = tiling_context_->GetInputShape(QUERY)->GetStorageShape(); + auto sharedKeyBlockShape = tiling_context_->GetInputShape(SHARED_KEY_BLOCK)->GetStorageShape(); + auto unsharedKeyBlockShape = tiling_context_->GetInputShape(UNSHARED_KEY_BLOCK)->GetStorageShape(); + auto sharedBlockTableShape = tiling_context_->GetOptionalInputShape(SHARED_BLOCK_TABLE)->GetStorageShape(); + + int32_t numTokens = queryShape.GetDim(0); + int32_t qHeadNum = queryShape.GetDim(1); + int32_t embeddingSize = queryShape.GetDim(2); + int32_t batch = sharedBlockTableShape.GetDim(0); + int32_t maxNumBlocksPerBatch = sharedBlockTableShape.GetDim(1); + int32_t blockNum = sharedKeyBlockShape.GetDim(0); + int32_t blockSize = sharedKeyBlockShape.GetDim(1); + int32_t kvHeadNum = sharedKeyBlockShape.GetDim(2); + int32_t maxDecodeStep = unsharedKeyBlockShape.GetDim(2); + int32_t beamSize = numTokens / batch; + + float scaleValue = static_cast(1.0 / std::sqrt(1.0 * embeddingSize)); + auto attrs = tiling_context_->GetAttrs(); + if (attrs != nullptr) { + const auto* attr_scale_value = attrs->GetAttrPointer(SCALE_VALUE_ATTR_INDEX); + if (attr_scale_value != nullptr && *attr_scale_value > 0.0f) { + scaleValue = *attr_scale_value; + } + } + + // 设置tiling信息 + tiling_data_.set_batch(batch); + tiling_data_.set_numHeads(qHeadNum); + tiling_data_.set_kvHeads(kvHeadNum); + tiling_data_.set_embeddingSize(embeddingSize); + tiling_data_.set_beamSize(beamSize); + tiling_data_.set_scaleValue(scaleValue); + tiling_data_.set_maskType(0); + tiling_data_.set_numTokens(numTokens); + tiling_data_.set_numBlocks(blockNum); + tiling_data_.set_blockSize(blockSize); + tiling_data_.set_maxNumBlocksPerBatch(maxNumBlocksPerBatch); + tiling_data_.set_maxDecodeStep(maxDecodeStep); + return ge::GRAPH_SUCCESS; +} + +uint64_t TilingXAttentionV2Func::GetTilingKey() const { + uint64_t resKey = 0; + resKey = uint32_t(isSharedPaged << NUM3) + uint32_t(isUnsharedPaged << NUM2) + (inputDtype << 1); + // shared continous unshared paged key: bf16(6) fp16(4) 0 1 (0/1) + // shared paged unshared continous key: bf16(10) fp16(8) 1 0 (0/1) + return resKey; +} + +ge::graphStatus TilingXAttentionV2Func::RunTiling() +{ + // Get platform hardware information + auto platform_info = + platform_ascendc::PlatformAscendC(tiling_context_->GetPlatformInfo()); + cubeCoreNum = platform_info.GetCoreNumAic(); + vecCoreNum = platform_info.GetCoreNumAiv(); + + BalanceAicore(); + auto ret = ParseInputShapeAndAttrs(); + if (ret != ge::GRAPH_SUCCESS) { + OP_LOGE(tiling_context_->GetNodeName(), "fill basic tiling failed."); + return ge::GRAPH_FAILED; + } + + FillSharedSplitCoreTilingData(); + FillUnsharedSplitCoreTilingData(); + FillCombineScaleTilingData(); + SetWorkspaces(); + + // Save tilingData + tiling_data_.SaveToBuffer(tiling_context_->GetRawTilingData()->GetData(), + tiling_context_->GetRawTilingData()->GetCapacity()); + tiling_context_->GetRawTilingData()->SetDataSize(tiling_data_.GetDataSize()); + tiling_context_->SetBlockDim(cubeCoreNum); + + tiling_context_->SetTilingKey(GetTilingKey()); + + return ge::GRAPH_SUCCESS; +} + + +static ge::graphStatus TilingFunc(gert::TilingContext* context) +{ + TilingXAttentionV2Func tilingObject(context); + auto ret = tilingObject.RunTiling(); + + if (ret != ge::GRAPH_SUCCESS) { + OP_LOGE(context->GetNodeName(), "xAttention tiling failed."); + return ge::GRAPH_FAILED; + } + + return ge::GRAPH_SUCCESS; +} +// --------------------------Tiling函数及TilingPrepare函数注册-------- +IMPL_OP_OPTILING(XAttentionV2) + .Tiling(TilingFunc); +} diff --git a/xllm_ops/x_attention_v2/op_host/x_attention_v2_tiling.h b/xllm_ops/x_attention_v2/op_host/x_attention_v2_tiling.h new file mode 100644 index 0000000..1ed1261 --- /dev/null +++ b/xllm_ops/x_attention_v2/op_host/x_attention_v2_tiling.h @@ -0,0 +1,56 @@ +/* Copyright 2025 The xLLM Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://gitcode.com/xLLM-AI/xllm_ops/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "register/tilingdata_base.h" +#include "tiling/tiling_api.h" + +namespace optiling { +BEGIN_TILING_DATA_DEF(XAttentionV2TilingData) + TILING_DATA_FIELD_DEF(uint32_t, numHeads); + TILING_DATA_FIELD_DEF(uint32_t, kvHeads); + TILING_DATA_FIELD_DEF(uint32_t, embeddingSize); + TILING_DATA_FIELD_DEF(uint32_t, batch); + TILING_DATA_FIELD_DEF(uint32_t, beamSize); + TILING_DATA_FIELD_DEF(float, scaleValue); + TILING_DATA_FIELD_DEF(uint32_t, maskType); + TILING_DATA_FIELD_DEF(uint32_t, numTokens); + TILING_DATA_FIELD_DEF(uint32_t, numBlocks); + TILING_DATA_FIELD_DEF(uint32_t, blockSize); + TILING_DATA_FIELD_DEF(uint32_t, sharedCoreNum); + TILING_DATA_FIELD_DEF(uint32_t, maxNumBlocksPerBatch); + TILING_DATA_FIELD_DEF(uint32_t, firstSharedBatchTaskNum); + TILING_DATA_FIELD_DEF(uint32_t, sharedTotalTaskNum); + TILING_DATA_FIELD_DEF(uint64_t, mm1OutSize); + TILING_DATA_FIELD_DEF(uint64_t, smOnlineOutSize); + TILING_DATA_FIELD_DEF(uint64_t, mm2OutSize); + TILING_DATA_FIELD_DEF(uint64_t, updateSize); + TILING_DATA_FIELD_DEF(uint32_t, rowSumMaxSize); + TILING_DATA_FIELD_DEF(uint64_t, sharedWorkspaceSize); + TILING_DATA_FIELD_DEF(uint32_t, groupSize); + TILING_DATA_FIELD_DEF(uint32_t, maxDecodeStep); + TILING_DATA_FIELD_DEF(uint32_t, unsharedCoreNum); + TILING_DATA_FIELD_DEF(uint32_t, unshareGroupCountPerLoop); + TILING_DATA_FIELD_DEF(uint32_t, unsharedFullCoreNum); + TILING_DATA_FIELD_DEF(uint32_t, unsharedTaskNumHead); + TILING_DATA_FIELD_DEF(uint32_t, unsharedTaskNumTail); + TILING_DATA_FIELD_DEF(uint32_t, unsharedLoopCountPerBatch); + TILING_DATA_FIELD_DEF(uint32_t, combineFormerCoreNum); + TILING_DATA_FIELD_DEF(uint32_t, combineFormerRowNum); + TILING_DATA_FIELD_DEF(uint32_t, combineTailRowNum); + TILING_DATA_FIELD_DEF(uint32_t, combineCoreNum); +END_TILING_DATA_DEF; + +REGISTER_TILING_DATA_CLASS(XAttentionV2, XAttentionV2TilingData) +} diff --git a/xllm_ops/x_attention_v2/op_kernel/x_attention_v2.cpp b/xllm_ops/x_attention_v2/op_kernel/x_attention_v2.cpp new file mode 100644 index 0000000..a2baf71 --- /dev/null +++ b/xllm_ops/x_attention_v2/op_kernel/x_attention_v2.cpp @@ -0,0 +1,60 @@ +/* Copyright 2025 The xLLM Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://gitcode.com/xLLM-AI/xllm_ops/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "kernel_operator.h" +#include "x_attention_v2_catlass_helper.h" +#include "lib/matmul_intf.h" + +#define CALL_XATTN_V2_KERNEL(INPUT_TYPE, SHARED_PAGED_FLAG, UNSHARED_PAGED_FLAG) \ + do { \ + if (coreIdx < tiling_data.sharedCoreNum) { \ + CallSharedInferKernelShort(params, &tiling_data); \ + } else { \ + CallUnsharedInferKernel(params, &tiling_data); \ + } \ + AscendC::SyncAll(); \ + CallCombineScale(params, &tiling_data); \ + } while (0) + +using namespace AscendC; + +extern "C" __global__ __aicore__ void x_attention_v2(GM_ADDR query, GM_ADDR shared_key_block, GM_ADDR shared_value_block, + GM_ADDR unshared_key_block, GM_ADDR unshared_value_block, GM_ADDR unshared_block_table, + GM_ADDR shared_kv_lens, GM_ADDR decode_step, GM_ADDR shared_block_table, GM_ADDR attn_out, GM_ADDR workspace, GM_ADDR tiling) { + // workspace use; [s,p,oTemp,oUpdate,shared_workspace,unshared_workspace] + KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_MIX_AIC_1_2); + GET_TILING_DATA(tiling_data, tiling); + + GM_ADDR s = workspace; + GM_ADDR p = workspace + tiling_data.mm1OutSize; + GM_ADDR oTemp = p + tiling_data.smOnlineOutSize; + GM_ADDR oUpdate = oTemp + tiling_data.mm2OutSize; + GM_ADDR shared_workspace = oUpdate + tiling_data.updateSize; + GM_ADDR unshared_workspace = shared_workspace + tiling_data.sharedWorkspaceSize; + int64_t coreIdx = AscendC::GetBlockIdx() / AscendC::GetSubBlockNum(); + + XAttnKernelParams params{query, shared_key_block, shared_value_block, unshared_key_block, unshared_value_block, + shared_block_table, unshared_block_table, shared_kv_lens, decode_step, s, p, oTemp, oUpdate, shared_workspace, + unshared_workspace, attn_out, tiling}; + if (TILING_KEY_IS(4)) { // 0b0100 + CALL_XATTN_V2_KERNEL(half, false, true); + } else if (TILING_KEY_IS(6)) { // 0b0110 + CALL_XATTN_V2_KERNEL(bfloat16_t, false, true); + } else if (TILING_KEY_IS(8)) { // 0b1000 + CALL_XATTN_V2_KERNEL(half, true, false); + } else if (TILING_KEY_IS(10)) { // 0b1010 + CALL_XATTN_V2_KERNEL(bfloat16_t, true, false); + } +} diff --git a/xllm_ops/x_attention_v2/op_kernel/x_attention_v2_catlass_helper.h b/xllm_ops/x_attention_v2/op_kernel/x_attention_v2_catlass_helper.h new file mode 100644 index 0000000..e46310b --- /dev/null +++ b/xllm_ops/x_attention_v2/op_kernel/x_attention_v2_catlass_helper.h @@ -0,0 +1,146 @@ +/* Copyright 2025 The xLLM Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://gitcode.com/xLLM-AI/xllm_ops/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#ifndef X_ATTN_V2_CATLASS_HELPER_H +#define X_ATTN_V2_CATLASS_HELPER_H +#include "x_attention_v2_catlass_kernel.h" + +template +CATLASS_DEVICE void CallSharedInferKernelShort(const XAttnKernelParams& params, XAttentionV2TilingData* tilingData) { + using ArchTag = Arch::AtlasA2; + using ElementQ = INPUT_T; + using LayoutQ = layout::RowMajor; + using ElementK = INPUT_T; + using LayoutK = layout::ColumnMajor; + using ElementV = INPUT_T; + using LayoutV = layout::RowMajor; + using ElementS = float; + using LayoutS = layout::RowMajor; + using ElementP = INPUT_T; + using LayoutP = layout::RowMajor; + using ElementO = INPUT_T; + using LayoutO = layout::RowMajor; + using ElementMask = INPUT_T; + using LayoutMask = layout::RowMajor; + using ElementOTmp = float; + using LayoutOTmp = layout::RowMajor; + using ElementUpdate = float; + using LayoutUpdate = layout::RowMajor; + // L1TileShape::K must be embdding + using L1TileShape = GemmShape<128, 128, 128>; + using L0TileShape = L1TileShape; + // GEMM Block, implement Q @ K^T of Flash Attention Infer + // using DispatchPolicyQK = Gemm::MmadAtlasA2FAIQK; + using DispatchPolicyQK = Gemm::MmadAtlasA2FAIQKSplitRow; + + using QType = Gemm::GemmType; + using KType = Gemm::GemmType; + using SType = Gemm::GemmType; + using BlockMmadQK = Gemm::Block::BlockMmad; + + // Shared Epilogue Block, update rowsum rowmax and copyOut on lastStackTile + using DispatchPolicyOnlineSoftmax = Epilogue::EpilogueAtlasA2OnlineSoftmaxCopySumMax; + using PType = Gemm::GemmType; + using maskType = Gemm::GemmType; + using EpilogueOnlineSoftmax = Epilogue::Block::BlockEpilogue; + + // GEMM Block, implement P @ V of Flash Attention Infer + // using DispatchPolicyPV = Gemm::MmadAtlasA2FAIPV; + using DispatchPolicyPV = Gemm::MmadAtlasA2FAIPVSplitRow; + + using VType = Gemm::GemmType; + using OTmpType = Gemm::GemmType; + using BlockMmadPV = Gemm::Block::BlockMmad; + + // Shared Epilogue RescaleO，do not div rowSum or cast on lastStackTile + using DispatchPolicyRescaleO = Epilogue::EpilogueAtlasA2RescaleOWithoutDivSum; + using OType = Gemm::GemmType; + using OUpdateType = Gemm::GemmType; + using EpilogueRescaleO = Epilogue::Block::BlockEpilogue; + + using SharedFAInferKernel = SharedFAInferKernelShort< + BlockMmadQK, BlockMmadPV, EpilogueOnlineSoftmax, EpilogueRescaleO, isPAEnabled>; + + SharedFAInferKernel sharedInferKernel(tilingData); + sharedInferKernel(params); +} + +template +CATLASS_DEVICE void CallUnsharedInferKernel(const XAttnKernelParams& params, XAttentionV2TilingData* tilingData) { +using ArchTag = Arch::AtlasA2; +using ElementQ = INPUT_T; +using LayoutQ = layout::RowMajor; +using ElementK = INPUT_T; +using LayoutK = layout::ColumnMajor; +using ElementV = INPUT_T; +using LayoutV = layout::RowMajor; +using ElementS = float; +using LayoutS = layout::RowMajor; +using ElementP = INPUT_T; +using LayoutP = layout::RowMajor; +using ElementO = INPUT_T; +using LayoutO = layout::RowMajor; +using ElementMask = INPUT_T; +using LayoutMask = layout::RowMajor; +using ElementOTmp = float; +using LayoutOTmp = layout::RowMajor; +using QType = Gemm::GemmType; +using KType = Gemm::GemmType; +using SType = Gemm::GemmType; +using PType = Gemm::GemmType; +using maskType = Gemm::GemmType; +using VType = Gemm::GemmType; +using OTmpType = Gemm::GemmType; + +using QKL1TileShape = GemmShape<128, 256, 128>; +using QKL0TileShape = QKL1TileShape; +using MmadDispatchPolicyQK = Gemm::MmadAtlasA2UnsharedFAQK; +using BlockMmadQK = Gemm::Block::BlockMmad; + +using DispatchPolicyFAUnsharedSoftmax = Epilogue::EpilogueAtlasA2FAUnsharedSoftmax; +using PType = Gemm::GemmType; +using maskType = Gemm::GemmType; +using EpilogueFAUnsharedSoftmax = Epilogue::Block::BlockEpilogue; + +using PVL1TileShape = GemmShape<128, 128, 256>; +using PVL0TileShape = PVL1TileShape; +using MmadDispatchPolicyPV = Gemm::MmadAtlasA2UnsharedFAPV; +using BlockMmadPV = Gemm::Block::BlockMmad; +using UnsharedFAInferKernel = UnsharedFAInferKernel; + +UnsharedFAInferKernel unsharedInferKernel(tilingData); +unsharedInferKernel(params); +} + + +template +CATLASS_DEVICE void CallCombineScale(const XAttnKernelParams& params, XAttentionV2TilingData* tilingData) { +using ArchTag = Arch::AtlasA2; +using ElementInput = float; +using LayoutInput = layout::RowMajor; +// MatrixShape +using ElementOutput = INPUT_T; +using LayoutOutput = layout::RowMajor; +using InputType = Gemm::GemmType; +using OutputType = Gemm::GemmType; +using DispatchPolicyCombineScale = Epilogue::EpilogueAtlasA2CombineScale; +using BlockEpilogueCombineScale = Epilogue::Block::BlockEpilogue; +using CombineScaleKernel = CombineScaleKernel; +CombineScaleKernel combineScaleKernel(tilingData); +combineScaleKernel(params); +} + + +#endif \ No newline at end of file diff --git a/xllm_ops/x_attention_v2/op_kernel/x_attention_v2_catlass_kernel.h b/xllm_ops/x_attention_v2/op_kernel/x_attention_v2_catlass_kernel.h new file mode 100644 index 0000000..8523422 --- /dev/null +++ b/xllm_ops/x_attention_v2/op_kernel/x_attention_v2_catlass_kernel.h @@ -0,0 +1,1629 @@ +/* Copyright 2025 The xLLM Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://gitcode.com/xLLM-AI/xllm_ops/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#ifndef X_ATTN_V2_CATLASS_KERNEL_H +#define X_ATTN_V2_CATLASS_KERNEL_H + +#include "catlass/arch/arch.hpp" +#include "catlass/arch/cross_core_sync.hpp" +#include "catlass/arch/resource.hpp" +#include "catlass/catlass.hpp" +#include "catlass/debug.hpp" +#include "catlass/epilogue/block/block_epilogue.hpp" +#include "catlass/epilogue/dispatch_policy.hpp" +#include "catlass/gemm/block/block_mmad.hpp" +#include "catlass/gemm/dispatch_policy.hpp" +#include "catlass/gemm/gemm_type.hpp" +#include "catlass/layout/layout.hpp" +#include "common/kernel_common.hpp" +#include "kernel_operator.h" +using namespace Catlass; + +template < + class BlockMmadQK, + class BlockMmadPV, + class EpilogueFAUnsharedSoftmax, + bool PAGED_CACHE_FLAG + > +class UnsharedFAInferKernel { + public: + using ArchTag = typename BlockMmadQK::ArchTag; + using L1TileShape = typename BlockMmadQK::L1TileShape; + using ElementQ = typename BlockMmadQK::ElementA; + using LayoutQ = typename BlockMmadQK::LayoutA; + using ElementK = typename BlockMmadQK::ElementB; + using LayoutK = typename BlockMmadQK::LayoutB; + using ElementS = typename BlockMmadQK::ElementC; + using LayoutS = typename BlockMmadQK::LayoutC; + + using ElementP = typename BlockMmadPV::ElementA; + using LayoutP = typename BlockMmadPV::LayoutA; + using ElementV = typename BlockMmadPV::ElementB; + using LayoutV = typename BlockMmadPV::LayoutB; + using ElementO = typename BlockMmadPV::ElementC; + using LayoutO = typename BlockMmadPV::LayoutC; + + CATLASS_DEVICE + UnsharedFAInferKernel(XAttentionV2TilingData* tilingDataPtr): faTilingData(tilingDataPtr) { + } + + template + CATLASS_DEVICE void operator()(XAttnKernelParams const ¶ms); + + template <> + CATLASS_DEVICE void operator()(XAttnKernelParams const ¶ms) { + static constexpr uint32_t L1_QK_SIZE = BlockMmadQK::L1TileShape::M * BlockMmadQK::L1TileShape::K + * sizeof(ElementQ) + + BlockMmadQK::L1TileShape::N * BlockMmadQK::L1TileShape::K + * sizeof(ElementK) * 2; + + BlockMmadQK blockMmadQK(resource); + BlockMmadPV blockMmadPV(resource, L1_QK_SIZE); + + // __gm__ XATilingData *faTilingData = reinterpret_cast<__gm__ XATilingData *>(params.tiling); + + AscendC::GlobalTensor gQ; + gQ.SetGlobalBuffer((__gm__ ElementQ *)params.q); + AscendC::GlobalTensor gUnsharedK; + gUnsharedK.SetGlobalBuffer((__gm__ ElementK *)params.unshared_k); + AscendC::GlobalTensor gUnsharedV; + gUnsharedV.SetGlobalBuffer((__gm__ ElementV *)params.unshared_v); + AscendC::GlobalTensor gS; + gS.SetGlobalBuffer((__gm__ ElementS *)params.s); + AscendC::GlobalTensor gP; + gP.SetGlobalBuffer((__gm__ ElementP *)params.p); + AscendC::GlobalTensor gUnsharedO; + gUnsharedO.SetGlobalBuffer((__gm__ ElementO *)params.unshared_workspace); + AscendC::GlobalTensor gUnsharedBlockTable; + gUnsharedBlockTable.SetGlobalBuffer((__gm__ int32_t *)(params.unsharedBlockTable)); + + uint32_t maxDecodeStep = faTilingData->maxDecodeStep; + uint32_t embeddingSize = faTilingData->embeddingSize; + uint32_t groupSize = faTilingData->groupSize; + uint32_t unsharedCoreNum = faTilingData->unsharedCoreNum; + uint32_t unshareGroupCountPerLoop = faTilingData->unshareGroupCountPerLoop; + uint32_t unsharedTaskNumHead = faTilingData->unsharedTaskNumHead; + uint32_t unsharedTaskNumTail = faTilingData->unsharedTaskNumTail; + uint32_t unsharedFullCoreNum = faTilingData->unsharedFullCoreNum; + uint32_t unsharedLoopCountPerBatch = faTilingData->unsharedLoopCountPerBatch; + uint32_t coreNumShared = faTilingData->sharedCoreNum; + + uint32_t relativeCoreIdx = AscendC::GetBlockIdx() - coreNumShared; + uint32_t taskStartIdx = relativeCoreIdx <= unsharedFullCoreNum ? relativeCoreIdx * unsharedTaskNumHead: + (unsharedFullCoreNum * unsharedTaskNumHead + (relativeCoreIdx - unsharedFullCoreNum) * unsharedTaskNumTail) ; + uint32_t taskEndIdx = taskStartIdx + (relativeCoreIdx >= unsharedFullCoreNum ? unsharedTaskNumTail : unsharedTaskNumHead); + //cce::printf("aic blockIdx:%d, sub block num:%d, loopCount:%d\n", coreIdx, AscendC::GetSubBlockNum(), taskEndIdx - taskStartIdx); + + uint64_t qOBaseBlockSize = groupSize * unshareGroupCountPerLoop * embeddingSize; + uint64_t kvBaseBlockSize = maxDecodeStep * unshareGroupCountPerLoop * embeddingSize; + uint64_t gmQOffset = taskStartIdx * qOBaseBlockSize; + uint64_t gmOOffset = gmQOffset; + uint64_t gmKOffset = 0; + GetKVOffset(taskStartIdx, unsharedLoopCountPerBatch, kvBaseBlockSize, gUnsharedBlockTable, gmKOffset); + uint64_t gmVOffset = gmKOffset; + uint64_t gmSOffset = (coreNumShared * WORKSPACE_BLOCK_SIZE_DB + + relativeCoreIdx * UNSHARED_WORKSPACE_BLOCK_SIZE_DB) * NUM3; + uint64_t gmPOffset = gmSOffset; + uint64_t sRelativeOffset = 0; + uint64_t pRelativeOffset = 0; + uint64_t actualPOffset = 0; + uint64_t actualSOffset = 0; + auto kvLen = unshareGroupCountPerLoop * maxDecodeStep; + auto qLen = unshareGroupCountPerLoop * groupSize; + auto kvLenAlign = (kvLen + 7) / 8 * 8; + LayoutQ layoutQTemp(qLen, embeddingSize); + LayoutK layoutKTemp(embeddingSize, kvLen); + LayoutV layoutVTemp(kvLen, embeddingSize); + LayoutO layoutOTemp(qLen, embeddingSize); + LayoutS layoutSTemp(qLen, kvLen, kvLenAlign); + LayoutP layoutPTemp(qLen, kvLen, kvLenAlign); + + GemmCoord actualBlockShapeQK{qLen, kvLen, embeddingSize}; + GemmCoord actualBlockShapePV{qLen, embeddingSize, kvLen}; + for (uint32_t taskIdx = taskStartIdx, nextTaskIdx; taskIdx < taskEndIdx + PRE_LAUNCH; ++taskIdx) { + if (taskIdx < taskEndIdx) { + sRelativeOffset = (taskIdx % (PRE_LAUNCH + 1)) * UNSHARED_WORKSPACE_BLOCK_SIZE_DB; + actualSOffset = gmSOffset + sRelativeOffset; + blockMmadQK( + gQ[gmQOffset], gUnsharedK[gmKOffset], gS[actualSOffset], + layoutQTemp, layoutKTemp, layoutSTemp, actualBlockShapeQK); + Arch::CrossCoreSetFlag<0x2, PIPE_FIX>(qkReady); + gmQOffset += qOBaseBlockSize; + nextTaskIdx = taskIdx + 1; + if (nextTaskIdx % unsharedLoopCountPerBatch == 0) { + GetKVOffset(nextTaskIdx, unsharedLoopCountPerBatch, kvBaseBlockSize, gUnsharedBlockTable, + gmKOffset); + } else { + gmKOffset += kvBaseBlockSize; + } + } + if (taskIdx >= taskStartIdx + PRE_LAUNCH) { + pRelativeOffset = ((taskIdx - PRE_LAUNCH) % (PRE_LAUNCH + 1)) * UNSHARED_WORKSPACE_BLOCK_SIZE_DB; + actualPOffset = gmPOffset + pRelativeOffset; + blockMmadPV( + gP[actualPOffset], + gUnsharedV[gmVOffset], gUnsharedO[gmOOffset], + layoutPTemp, layoutVTemp, layoutOTemp, + actualBlockShapePV, softmaxReady); + gmOOffset += qOBaseBlockSize; + nextTaskIdx = taskIdx + 1 - PRE_LAUNCH; + if (nextTaskIdx % unsharedLoopCountPerBatch == 0) { + GetKVOffset(nextTaskIdx, unsharedLoopCountPerBatch, kvBaseBlockSize, gUnsharedBlockTable, + gmVOffset); + } else { + gmVOffset += kvBaseBlockSize; + } + } + } + } + template <> + CATLASS_DEVICE void operator()(XAttnKernelParams const ¶ms) { + AscendC::SetFlag(EVENT_ID0); + AscendC::SetFlag(EVENT_ID0); + + AscendC::GlobalTensor gS; + gS.SetGlobalBuffer((__gm__ ElementS *)params.s); + AscendC::GlobalTensor gP; + gP.SetGlobalBuffer((__gm__ ElementP *)params.p); + AscendC::GlobalTensor gDecodeStep; + gDecodeStep.SetGlobalBuffer((__gm__ uint32_t *)params.decodeStep); + + uint32_t batch = faTilingData->batch; + uint32_t beamSize = faTilingData->beamSize; + uint32_t numHeads = faTilingData->numHeads; + uint32_t unsharedKvSeqLen = gDecodeStep.GetValue(0); + uint32_t maxDecodeStep = faTilingData->maxDecodeStep; + uint32_t embeddingSize = faTilingData->embeddingSize; + uint32_t groupSize = faTilingData->groupSize; + uint32_t unsharedCoreNum = faTilingData->unsharedCoreNum; + uint32_t unshareGroupCountPerLoop = faTilingData->unshareGroupCountPerLoop; + uint32_t unsharedTaskNumHead = faTilingData->unsharedTaskNumHead; + uint32_t unsharedTaskNumTail = faTilingData->unsharedTaskNumTail; + uint32_t unsharedFullCoreNum = faTilingData->unsharedFullCoreNum; + float scaleValue = faTilingData->scaleValue; + uint32_t coreNumShared = faTilingData->sharedCoreNum; + + uint32_t gUnsharedOffset = batch * beamSize * numHeads * embeddingSize; + AscendC::GlobalTensor gUnsharedGm; + gUnsharedGm.SetGlobalBuffer(((__gm__ ElementO *)params.unshared_workspace) + gUnsharedOffset); + gUnsharedOffset += batch * beamSize * numHeads; + AscendC::GlobalTensor gUnsharedGl; + gUnsharedGl.SetGlobalBuffer(((__gm__ ElementO *)params.unshared_workspace) + gUnsharedOffset); + + uint32_t relativeCoreIdx = AscendC::GetBlockIdx() / AscendC::GetSubBlockNum() - coreNumShared; + uint32_t taskStartIdx = relativeCoreIdx <= unsharedFullCoreNum ? relativeCoreIdx * unsharedTaskNumHead: + (unsharedFullCoreNum * unsharedTaskNumHead + (relativeCoreIdx - unsharedFullCoreNum) * unsharedTaskNumTail) ; + uint32_t taskEndIdx = taskStartIdx + (relativeCoreIdx >= unsharedFullCoreNum ? unsharedTaskNumTail : unsharedTaskNumHead); + + EpilogueFAUnsharedSoftmax epilogueFAUnsharedSoftmax(resource, scaleValue, unsharedKvSeqLen, maxDecodeStep, unshareGroupCountPerLoop, + groupSize); + + uint64_t gmSPOffset = (coreNumShared * WORKSPACE_BLOCK_SIZE_DB + + relativeCoreIdx * UNSHARED_WORKSPACE_BLOCK_SIZE_DB) * NUM3; + uint64_t spRelativeOffset = 0; + uint64_t actualSPOffset = 0; + int64_t gmGlBaseBlockSize = groupSize * unshareGroupCountPerLoop; + uint64_t gmUnsharedGmGlOffset = taskStartIdx * gmGlBaseBlockSize; + auto kvLen = unshareGroupCountPerLoop * maxDecodeStep; + auto qLen = unshareGroupCountPerLoop * groupSize; + auto kvLenAlign = (kvLen + 7) / 8 * 8; + LayoutS layoutSTemp(qLen, kvLen, kvLenAlign); + LayoutP layoutPTemp(qLen, kvLen, kvLenAlign); + GemmCoord actualBlockShapeQK{unshareGroupCountPerLoop * groupSize, unshareGroupCountPerLoop * maxDecodeStep, + embeddingSize}; + for (uint32_t taskIdx = taskStartIdx; taskIdx < taskEndIdx; ++taskIdx) { + spRelativeOffset = (taskIdx % (PRE_LAUNCH + 1)) * UNSHARED_WORKSPACE_BLOCK_SIZE_DB; + actualSPOffset = gmSPOffset + spRelativeOffset; + Arch::CrossCoreWaitFlag(qkReady); + // FA unshared softmax + epilogueFAUnsharedSoftmax( + gP[actualSPOffset], gS[actualSPOffset], gUnsharedGm[gmUnsharedGmGlOffset], + gUnsharedGl[gmUnsharedGmGlOffset], layoutPTemp, layoutSTemp, + actualBlockShapeQK, unshareGroupCountPerLoop + ); + Arch::CrossCoreSetFlag<0x2, PIPE_MTE3>(softmaxReady); + gmUnsharedGmGlOffset += gmGlBaseBlockSize; + } + AscendC::WaitFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + } + private: + CATLASS_DEVICE void GetKVOffset(uint32_t taskIdx, uint32_t unsharedLoopCountPerBatch, uint64_t kvBaseBlockSize, + const AscendC::GlobalTensor &gUnsharedBlockTable, uint64_t &offset) { + if constexpr(PAGED_CACHE_FLAG) { + uint32_t batchId = gUnsharedBlockTable.GetValue(taskIdx / unsharedLoopCountPerBatch); + uint32_t remainCnt = taskIdx % unsharedLoopCountPerBatch; + offset = (batchId * unsharedLoopCountPerBatch + remainCnt) * kvBaseBlockSize; + } else { + offset = taskIdx * kvBaseBlockSize; + } + } + private: + Arch::Resource resource; + Arch::CrossCoreFlag qkReady{QK_READY_ID}; + Arch::CrossCoreFlag softmaxReady{SOFTMAX_READY_ID}; + XAttentionV2TilingData* faTilingData; +}; + + /* + FASharedInferKernel + Compute Stream + 1. BlockMmadQK + 2. OnlineSoftmax + 3. BlockMmadPV + 4. EpilogueRescaleO 最后rescaleO 不需要div rowsum + */ + template < + class BlockMmadQK, + class BlockMmadPV, + class BlockMmadQKTail, + class BlockMmadPVTail, + class EpilogueOnlineSoftmax, + class EpilogueRescaleO, + bool PAGED_CACHE_FLAG = true + > + class SharedFAInferKernel { + public: + using ArchTag = typename BlockMmadQK::ArchTag; + using L1TileShape = typename BlockMmadQK::L1TileShape; + using ElementQ = typename BlockMmadQK::ElementA; + using LayoutQ = typename BlockMmadQK::LayoutA; + using ElementK = typename BlockMmadQK::ElementB; + using LayoutK = typename BlockMmadQK::LayoutB; + using ElementS = typename BlockMmadQK::ElementC; + using LayoutS = typename BlockMmadQK::LayoutC; + + using ElementP = typename BlockMmadPV::ElementA; + using LayoutP = typename BlockMmadPV::LayoutA; + using ElementV = typename BlockMmadPV::ElementB; + using LayoutV = typename BlockMmadPV::LayoutB; + + using ElementMask = typename EpilogueOnlineSoftmax::ElementMask; + using LayoutMask = typename EpilogueOnlineSoftmax::LayoutMask; + + using ElementO = typename EpilogueRescaleO::ElementOutput; + using LayoutO = typename EpilogueRescaleO::LayoutOutput; + + using ElementOTmp = typename EpilogueRescaleO::ElementInput; + using LayoutOTmp = typename EpilogueRescaleO::LayoutInput; + + // Methods + CATLASS_DEVICE + SharedFAInferKernel(XAttentionV2TilingData* tilingDataPtr): faTilingData(tilingDataPtr) { + } + + template + CATLASS_DEVICE void operator()(XAttnKernelParams const ¶ms); + + template <> + CATLASS_DEVICE void operator()(XAttnKernelParams const ¶ms) { + AscendC::SetFlag(EVENT_ID0); + AscendC::SetFlag(EVENT_ID1); + AscendC::SetFlag(EVENT_ID2); + AscendC::SetFlag(EVENT_ID3); + AscendC::SetFlag(EVENT_ID4); + AscendC::SetFlag(EVENT_ID5); + AscendC::SetFlag(EVENT_ID6); + AscendC::SetFlag(EVENT_ID7); + AscendC::SetFlag(EVENT_ID0); + AscendC::SetFlag(EVENT_ID1); + AscendC::SetFlag(EVENT_ID0); + AscendC::SetFlag(EVENT_ID1); + AscendC::SetFlag(EVENT_ID2); + AscendC::SetFlag(EVENT_ID3); + AscendC::SetFlag(EVENT_ID4); + AscendC::SetFlag(EVENT_ID5); + AscendC::SetFlag(EVENT_ID6); + AscendC::SetFlag(EVENT_ID7); + static constexpr uint32_t L1_QK_SIZE = BlockMmadQK::L1TileShape::M * BlockMmadQK::L1TileShape::K + * sizeof(ElementQ) + + BlockMmadQK::L1TileShape::N * BlockMmadQK::L1TileShape::K + * sizeof(ElementK) * 2; + BlockMmadQK blockMmadQK(resource); + BlockMmadPV blockMmadPV(resource, L1_QK_SIZE); + + BlockMmadQKTail blockMmadQKTail(resource); + BlockMmadPVTail blockMmadPVTail(resource, L1_QK_SIZE); + // __gm__ XATilingData *faTilingData = reinterpret_cast<__gm__ XATilingData *>(params.tiling); + uint64_t mm1OutSize = faTilingData->mm1OutSize; + uint64_t smOnlineOutSize = faTilingData->smOnlineOutSize; + uint32_t batch = faTilingData->batch; // requestNum + uint32_t beamSize = faTilingData->beamSize; + uint32_t qHeads = faTilingData->numHeads; + uint32_t kvHeads = faTilingData->kvHeads; + uint32_t embed = faTilingData->embeddingSize; + uint32_t pagedBlockSize = faTilingData->blockSize; + uint32_t sharedCoreNum = faTilingData->sharedCoreNum; + uint32_t maxNumBlocksPerBatch = faTilingData->maxNumBlocksPerBatch; + uint32_t curTotalTaskNum = faTilingData->firstSharedBatchTaskNum; + uint32_t totalTaskNum = faTilingData->sharedTotalTaskNum; + uint32_t blockSize = faTilingData->blockSize; + uint32_t maskType = faTilingData->maskType; + float scaleValue = faTilingData->scaleValue; + + AscendC::GlobalTensor gQ; + gQ.SetGlobalBuffer((__gm__ ElementQ *)params.q); + AscendC::GlobalTensor gK; + gK.SetGlobalBuffer((__gm__ ElementK *)params.k_cache); + AscendC::GlobalTensor gV; + gV.SetGlobalBuffer((__gm__ ElementK *)params.v_cache); + AscendC::GlobalTensor gBlockTable; + gBlockTable.SetGlobalBuffer((__gm__ int32_t *)(params.sharedBlockTable)); + AscendC::GlobalTensor gActualKvseqlen; + gActualKvseqlen.SetGlobalBuffer((__gm__ int32_t *)params.actualKvseqlen); + AscendC::GlobalTensor gS; + gS.SetGlobalBuffer((__gm__ ElementS *)params.s); + AscendC::GlobalTensor gP; + gP.SetGlobalBuffer((__gm__ ElementP *)params.p); + AscendC::GlobalTensor gOTmp; + gOTmp.SetGlobalBuffer((__gm__ ElementOTmp *)params.oTemp); + + uint64_t strideQO = qHeads * embed; + uint64_t strideKV = kvHeads * embed; + uint32_t embedRound = RoundUp(embed); + uint32_t groupSize = qHeads / kvHeads; + + uint32_t coreIdx = AscendC::GetBlockIdx(); + uint32_t coreNum = sharedCoreNum; // TODO coreNum Need To be modified + curTotalTaskNum = 0; + uint32_t preTotalTaskNum = 0; + uint32_t curBatch = 0; + uint64_t qBOffset = 0; + uint64_t kBOffset = 0; + uint64_t vBOffset = 0; + uint64_t blockBOffset = 0; + int64_t qSeqlen = 0; + int64_t kvSeqlen = 0; + uint32_t curQNBlockTile; + uint32_t qNBlockNumPerGroup; + uint32_t curQNBlockNum; + int64_t curQSBlockTile; + uint32_t curQSBlockNum; + + preTotalTaskNum = curTotalTaskNum; + qSeqlen = beamSize; + kvSeqlen = gActualKvseqlen.GetValue(curBatch); + curQSBlockTile = GetQSBlockTile(kvSeqlen); + curQNBlockTile = GetQNBlockTile(qSeqlen, groupSize); + qNBlockNumPerGroup = CeilDiv(groupSize, curQNBlockTile); + curQNBlockNum = qNBlockNumPerGroup * kvHeads; + curQSBlockNum = CeilDiv(qSeqlen, curQSBlockTile); + curTotalTaskNum += curQNBlockNum * curQSBlockNum; + for (uint32_t taskIdx = coreIdx; taskIdx < totalTaskNum; taskIdx += uint32_t(coreNum)) { + while (taskIdx >= curTotalTaskNum) { + ++curBatch; + preTotalTaskNum = curTotalTaskNum; + qBOffset += qSeqlen * strideQO; + if constexpr (!PAGED_CACHE_FLAG) { + kBOffset += kvSeqlen * strideKV; + vBOffset += kvSeqlen * strideKV; + } else { + blockBOffset += maxNumBlocksPerBatch; + } + qSeqlen = beamSize; + kvSeqlen = gActualKvseqlen.GetValue(curBatch); + curQSBlockTile = GetQSBlockTile(kvSeqlen); + curQNBlockTile = GetQNBlockTile(qSeqlen, groupSize); + qNBlockNumPerGroup = CeilDiv(groupSize, curQNBlockTile); + curQNBlockNum = qNBlockNumPerGroup * kvHeads; + curQSBlockNum = CeilDiv(qSeqlen, curQSBlockTile); + curTotalTaskNum += curQNBlockNum * curQSBlockNum; + } + uint32_t taskIdxCurBatch = taskIdx - preTotalTaskNum; + uint32_t qSBlockIdx = taskIdxCurBatch / curQNBlockNum; + uint32_t qNBlockIdx = taskIdxCurBatch - qSBlockIdx * curQNBlockNum; + uint32_t qNBlockIdxCurGroup = qNBlockIdx % qNBlockNumPerGroup; + uint32_t kvHeadIdx = qNBlockIdx / qNBlockNumPerGroup; + uint32_t qHeadIdx = kvHeadIdx * groupSize + qNBlockIdxCurGroup * curQNBlockTile; + uint64_t gmQOffset = qBOffset + qSBlockIdx * curQSBlockTile * strideQO + qHeadIdx * embed; + uint64_t gmKOffset = kBOffset + kvHeadIdx * embed; + uint64_t gmVOffset = vBOffset + kvHeadIdx * embed; + uint32_t qSBlockSize = (qSBlockIdx == (curQSBlockNum - 1)) ? (qSeqlen - qSBlockIdx * curQSBlockTile) + : curQSBlockTile; + uint32_t qNBlockSize = (qNBlockIdxCurGroup == (qNBlockNumPerGroup - 1)) + ? (groupSize - qNBlockIdxCurGroup * curQNBlockTile) + : curQNBlockTile; + uint32_t rowNum = qSBlockSize * qNBlockSize; + uint32_t rowNumRound = AlignUp(rowNum, BLOCK_SIZE); + uint32_t noSkipKvS = kvSeqlen; + uint32_t noMaskKvS = kvSeqlen; + uint32_t noMaskTailS = 0; + // if (maskType != 0) { + // uint32_t diffS = kvSeqlen - qSeqlen; + // noSkipKvS = (qSBlockIdx + 1) * curQSBlockTile + diffS; + // noSkipKvS = Min((uint32_t)kvSeqlen, noSkipKvS); + // noMaskKvS = noSkipKvS - qSBlockSize; + // noMaskTailS = noMaskKvS % pagedBlockSize; + // } + uint32_t maskedKvS = qSBlockSize; + uint32_t kvSLoopNumNoMask = CeilDiv(noMaskKvS, pagedBlockSize); + uint32_t kvSLoopNumTotal = CeilDiv(noSkipKvS, pagedBlockSize); + uint32_t blockStackNum = 4; + uint32_t stackSeqTile; + uint32_t stackSeqTileRound = blockStackNum * 128; + int32_t preLaunch = 2; + int32_t totalStackSeqNum = (maskType != 0) ? (CeilDiv(noMaskKvS, blockStackNum * pagedBlockSize) + 1) + : CeilDiv(noMaskKvS, blockStackNum * pagedBlockSize); + int32_t stackSeqCount = 0; + + LayoutQ layoutQTemp(rowNum, embed); + LayoutK layoutKTemp(strideKV, blockStackNum * pagedBlockSize); + LayoutV layoutVTemp(blockStackNum * pagedBlockSize, strideKV); + blockMmadQK.loadQGM(gQ[gmQOffset], layoutQTemp, rowNum, qNBlockSize, qHeads); + for (uint32_t kvSIdx = 0; kvSIdx < kvSLoopNumNoMask; kvSIdx += blockStackNum) { + if (kvSIdx < kvSLoopNumNoMask) { + if (kvSIdx + blockStackNum > kvSLoopNumNoMask - 1) { + stackSeqTile = noMaskKvS - kvSIdx * pagedBlockSize; + } else { + stackSeqTile = pagedBlockSize * blockStackNum; + } + uint32_t SWorkSpacePingPongFlag = stackSeqCount % (preLaunch + 1); + uint64_t gmSOffset = coreIdx * WORKSPACE_BLOCK_SIZE_DB * (preLaunch + 1) + + SWorkSpacePingPongFlag * WORKSPACE_BLOCK_SIZE_DB; + GemmCoord actualBlockShapeQK{rowNum, stackSeqTile, embed}; + if constexpr (!PAGED_CACHE_FLAG) { + blockMmadQK( + gQ[gmQOffset], gK[gmKOffset], gS[gmSOffset], gBlockTable, layoutQTemp, layoutKTemp, + actualBlockShapeQK, kvSIdx, kvSLoopNumNoMask, pagedBlockSize, noMaskKvS, strideKV + ); + } else { + blockMmadQK( + gQ[gmQOffset], gK[gmKOffset], gS[gmSOffset], gBlockTable[blockBOffset], layoutQTemp, + layoutKTemp, actualBlockShapeQK, kvSIdx, kvSLoopNumNoMask, pagedBlockSize, noMaskKvS, + strideKV + ); + } + Arch::CrossCoreSetFlag<0x2, PIPE_FIX>(qkReady); + } + if (kvSIdx >= preLaunch * blockStackNum) { + uint32_t nowkvSIdx = kvSIdx - preLaunch * blockStackNum; + if (nowkvSIdx + blockStackNum > kvSLoopNumNoMask - 1) { + stackSeqTile = noMaskKvS - nowkvSIdx * pagedBlockSize; + } else { + stackSeqTile = pagedBlockSize * blockStackNum; + } + uint32_t PVWorkSpacePingPongFlag = (stackSeqCount - preLaunch) % (preLaunch + 1); + uint64_t gmPOffset = coreIdx * WORKSPACE_BLOCK_SIZE_DB * (preLaunch + 1) + + PVWorkSpacePingPongFlag * WORKSPACE_BLOCK_SIZE_DB; + uint64_t gmOTmpOffset = gmPOffset; + LayoutP layoutPTemp(rowNum, stackSeqTileRound); + GemmCoord actualBlockShapePV{rowNum, embed, stackSeqTile}; + if constexpr (!PAGED_CACHE_FLAG) { + blockMmadPV( + gP[gmPOffset], gV[gmVOffset], gOTmp[gmOTmpOffset], gBlockTable, layoutPTemp, layoutVTemp, + actualBlockShapePV, nowkvSIdx, kvSLoopNumNoMask, pagedBlockSize, noMaskKvS, strideKV, + softmaxReady + ); + } else { + blockMmadPV( + gP[gmPOffset], gV[gmVOffset], gOTmp[gmOTmpOffset], gBlockTable[blockBOffset], layoutPTemp, + layoutVTemp, actualBlockShapePV, nowkvSIdx, kvSLoopNumNoMask, pagedBlockSize, noMaskKvS, + strideKV, softmaxReady + ); + } + Arch::CrossCoreSetFlag<0x2, PIPE_FIX>(pvReady); + } + stackSeqCount++; + } + + /* + * for the secondary loop + * while masked, it deals the CV stage1(Qk^t/SMOnline) of the final base block(typical shape [128, 512]), + * and the CV stage2(PV/rescaleO) of the last (prelaunch+1) base blocks while not masked, it deals only the + * CV stage2(PV/rescaleO) of the last (prelaunch) base blocks + */ + + // deal secondary loop conditions + uint32_t maskedStartIdx = (maskType != 0) ? ((noMaskTailS != 0) ? (kvSLoopNumNoMask - 1) : kvSLoopNumNoMask) + : AlignUp(kvSLoopNumNoMask, blockStackNum); + uint32_t noMaskTailInteStackNum = (noMaskKvS / pagedBlockSize) % blockStackNum; + noMaskTailInteStackNum = (noMaskTailInteStackNum != 0) ? noMaskTailInteStackNum + : ((noMaskTailS != 0) ? 0 : blockStackNum); + uint32_t preLaunchStackNum = (maskType != 0) ? ((preLaunch - 1) * blockStackNum + noMaskTailInteStackNum) + : (preLaunch * blockStackNum); + + // masked kvSeqlen loop + + for (uint32_t kvSIdx = maskedStartIdx; kvSIdx < kvSLoopNumTotal + preLaunchStackNum;) { + if ((kvSIdx < kvSLoopNumTotal) && (stackSeqCount <= totalStackSeqNum - 1)) { + stackSeqTile = maskedKvS; + uint32_t SWorkSpacePingPongFlag = stackSeqCount % (preLaunch + 1); + uint64_t gmSOffset = coreIdx * WORKSPACE_BLOCK_SIZE_DB * (preLaunch + 1) + + SWorkSpacePingPongFlag * WORKSPACE_BLOCK_SIZE_DB; + GemmCoord actualBlockShapeQK{rowNum, stackSeqTile, embed}; + if constexpr (!PAGED_CACHE_FLAG) { + blockMmadQKTail( + gQ[gmQOffset], gK[gmKOffset], gS[gmSOffset], gBlockTable, layoutQTemp, layoutKTemp, + actualBlockShapeQK, kvSIdx, kvSLoopNumTotal, pagedBlockSize, noSkipKvS, strideKV, + noMaskTailS, 1 + ); + } else { + blockMmadQKTail( + gQ[gmQOffset], gK[gmKOffset], gS[gmSOffset], gBlockTable[blockBOffset], layoutQTemp, + layoutKTemp, actualBlockShapeQK, kvSIdx, kvSLoopNumTotal, pagedBlockSize, noSkipKvS, + strideKV, noMaskTailS, 1 + ); + } + Arch::CrossCoreSetFlag<0x2, PIPE_FIX>(qkReady); + } + + if (kvSIdx >= preLaunchStackNum) { + uint32_t delayedKvSIdx = kvSIdx - preLaunchStackNum; + + if (delayedKvSIdx + blockStackNum > kvSLoopNumTotal - 1 && (maskType != 0)) { + stackSeqTile = maskedKvS; + } else if (delayedKvSIdx + blockStackNum > kvSLoopNumNoMask - 1) { + stackSeqTile = noMaskKvS - delayedKvSIdx * pagedBlockSize; + } else { + stackSeqTile = pagedBlockSize * blockStackNum; + } + uint32_t PVWorkSpacePingPongFlag = (stackSeqCount - preLaunch) % (preLaunch + 1); + uint64_t gmPOffset = coreIdx * WORKSPACE_BLOCK_SIZE_DB * (preLaunch + 1) + + PVWorkSpacePingPongFlag * WORKSPACE_BLOCK_SIZE_DB; + uint64_t gmOTmpOffset = gmPOffset; + LayoutP layoutPTemp(rowNum, stackSeqTileRound); + GemmCoord actualBlockShapePV{rowNum, embed, stackSeqTile}; + + if ((stackSeqCount - preLaunch == totalStackSeqNum - 1) && (maskType != 0)) { // 加mask + if constexpr (!PAGED_CACHE_FLAG) { + blockMmadPVTail( + gP[gmPOffset], gV[gmVOffset], gOTmp[gmOTmpOffset], gBlockTable, layoutPTemp, + layoutVTemp, actualBlockShapePV, delayedKvSIdx, kvSLoopNumTotal, pagedBlockSize, + noSkipKvS, strideKV, softmaxReady, noMaskTailS, 1 + ); + } else { + blockMmadPVTail( + gP[gmPOffset], gV[gmVOffset], gOTmp[gmOTmpOffset], gBlockTable[blockBOffset], + layoutPTemp, layoutVTemp, actualBlockShapePV, delayedKvSIdx, kvSLoopNumTotal, + pagedBlockSize, noSkipKvS, strideKV, softmaxReady, noMaskTailS, 1 + ); + } + } else { // 不加mask + if constexpr (!PAGED_CACHE_FLAG) { + blockMmadPV( + gP[gmPOffset], gV[gmVOffset], gOTmp[gmOTmpOffset], gBlockTable, layoutPTemp, + layoutVTemp, actualBlockShapePV, delayedKvSIdx, kvSLoopNumNoMask, pagedBlockSize, + noMaskKvS, strideKV, softmaxReady + ); + } else { + blockMmadPV( + gP[gmPOffset], gV[gmVOffset], gOTmp[gmOTmpOffset], gBlockTable[blockBOffset], + layoutPTemp, layoutVTemp, actualBlockShapePV, delayedKvSIdx, kvSLoopNumNoMask, + pagedBlockSize, noMaskKvS, strideKV, softmaxReady + ); + } + } + Arch::CrossCoreSetFlag<0x2, PIPE_FIX>(pvReady); + } + if ((maskType != 0) && (stackSeqCount - preLaunch == totalStackSeqNum - 2)) { + kvSIdx += noMaskTailInteStackNum; + } else { + kvSIdx += blockStackNum; + } + stackSeqCount++; + } + } + + AscendC::WaitFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID1); + AscendC::WaitFlag(EVENT_ID2); + AscendC::WaitFlag(EVENT_ID3); + AscendC::WaitFlag(EVENT_ID4); + AscendC::WaitFlag(EVENT_ID5); + AscendC::WaitFlag(EVENT_ID6); + AscendC::WaitFlag(EVENT_ID7); + + AscendC::WaitFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID1); + + AscendC::WaitFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID1); + AscendC::WaitFlag(EVENT_ID2); + AscendC::WaitFlag(EVENT_ID3); + AscendC::WaitFlag(EVENT_ID4); + AscendC::WaitFlag(EVENT_ID5); + AscendC::WaitFlag(EVENT_ID6); + AscendC::WaitFlag(EVENT_ID7); + } + + + template <> + CATLASS_DEVICE void operator()(XAttnKernelParams const ¶ms) { + AscendC::SetFlag(EVENT_ID0); + AscendC::SetFlag(EVENT_ID1); + AscendC::SetFlag(EVENT_ID2); + AscendC::SetFlag(EVENT_ID3); + AscendC::SetFlag(EVENT_ID4); + AscendC::SetFlag(EVENT_ID5); + + AscendC::SetFlag(EVENT_ID0); + AscendC::SetFlag(EVENT_ID1); + AscendC::SetFlag(EVENT_ID3); + AscendC::SetFlag(EVENT_ID2); + AscendC::SetFlag(EVENT_ID2); + + // Get tiling parameters + // __gm__ XATilingData *faTilingData = reinterpret_cast<__gm__ XATilingData *>(params.tiling); + uint64_t mm1OutSize = faTilingData->mm1OutSize; + uint64_t smOnlineOutSize = faTilingData->smOnlineOutSize; + uint64_t mm2OutSize = faTilingData->mm2OutSize; + uint32_t batch = faTilingData->batch; + uint32_t beamSize = faTilingData->beamSize; + uint32_t qHeads = faTilingData->numHeads; + uint32_t kvHeads = faTilingData->kvHeads; + uint32_t embed = faTilingData->embeddingSize; + uint32_t sharedCoreNum = faTilingData->sharedCoreNum; + uint32_t pagedBlockSize = faTilingData->blockSize; + uint32_t maxNumBlocksPerBatch = faTilingData->maxNumBlocksPerBatch; + uint32_t firstBatchTaskNum = faTilingData->firstSharedBatchTaskNum; + uint32_t totalTaskNum = faTilingData->sharedTotalTaskNum; + uint32_t maskType = faTilingData->maskType; + float scaleValue = faTilingData->scaleValue; + + uint64_t gOffsetTempO = batch * beamSize * qHeads * embed; + uint64_t gMaxOffset = batch * beamSize * qHeads * SOFTMAX_BROAD_SIZE;; + // Get the memory offset address of the input on Global Memory + AscendC::GlobalTensor gActualKvseqlen; + gActualKvseqlen.SetGlobalBuffer((__gm__ int32_t *)params.actualKvseqlen); + AscendC::GlobalTensor gO; + gO.SetGlobalBuffer((__gm__ ElementO *)params.o); + AscendC::GlobalTensor gS; + gS.SetGlobalBuffer((__gm__ ElementS *)params.s); + AscendC::GlobalTensor gP; + gP.SetGlobalBuffer((__gm__ ElementP *)params.p); + AscendC::GlobalTensor gOTmp; + gOTmp.SetGlobalBuffer((__gm__ ElementOTmp *)params.oTemp); + AscendC::GlobalTensor gOUpdate; + gOUpdate.SetGlobalBuffer((__gm__ ElementOTmp *)params.oUpdate); + + // shared Gm and Gl output + AscendC::GlobalTensor gSharedO; + AscendC::GlobalTensor gSharedSum; + AscendC::GlobalTensor gSharedMax; + gSharedO.SetGlobalBuffer((__gm__ float *)params.shared_workspace); + gSharedMax.SetGlobalBuffer((__gm__ ElementS *)params.shared_workspace + gOffsetTempO); + gSharedSum.SetGlobalBuffer((__gm__ ElementS *)params.shared_workspace + gOffsetTempO + gMaxOffset); + + + uint32_t groupSize = qHeads / kvHeads; + uint32_t embedRound = RoundUp(embed, BLOCK_SIZE); + + EpilogueOnlineSoftmax epilogueOnlineSoftmax(resource, scaleValue); + EpilogueRescaleO epilogueRescaleO(resource); + + // uint32_t curTotalTaskNum = 0; + uint32_t preTotalTaskNum = 0; + uint32_t curBatch = 0; + uint32_t oBatchOffset = 0; + uint32_t qSeqlen = static_cast(beamSize); + uint32_t kvSeqlen = static_cast(gActualKvseqlen.GetValue(curBatch)); + uint32_t curQNBlockTile = GetQNBlockTile(qSeqlen, groupSize); + uint32_t qNBlockNumPerGroup = CeilDiv(groupSize, curQNBlockTile); + uint32_t curQNBlockNum = qNBlockNumPerGroup * kvHeads; + uint32_t curQSBlockTile = GetQSBlockTile(kvSeqlen); + uint32_t curQSBlockNum = CeilDiv(qSeqlen, curQSBlockTile); + uint32_t curTotalTaskNum = firstBatchTaskNum; + + uint32_t coreIdx = AscendC::GetBlockIdx() / AscendC::GetSubBlockNum(); + // coreNum Need to be changed + uint32_t coreNum = sharedCoreNum; + + // Go through each task. + for (uint32_t taskIdx = coreIdx; taskIdx < totalTaskNum; taskIdx += uint32_t(coreNum)) { + // Get the offset of each core on the GM. + while (taskIdx >= curTotalTaskNum) { + curBatch++; + oBatchOffset += qSeqlen * qHeads * embed; + preTotalTaskNum = curTotalTaskNum; + qSeqlen = static_cast(beamSize); + kvSeqlen = static_cast(gActualKvseqlen.GetValue(curBatch)); + curQNBlockTile = GetQNBlockTile(qSeqlen, groupSize); + qNBlockNumPerGroup = CeilDiv(groupSize, curQNBlockTile); + curQNBlockNum = qNBlockNumPerGroup * kvHeads; + curQSBlockTile = GetQSBlockTile(kvSeqlen); + curQSBlockNum = CeilDiv(qSeqlen, curQSBlockTile); + curTotalTaskNum += curQNBlockNum * curQSBlockNum; + } + uint32_t taskIdxCurBatch = taskIdx - preTotalTaskNum; + uint32_t qSBlockIdx = taskIdxCurBatch / curQNBlockNum; + uint32_t qNBlockIdx = taskIdxCurBatch % curQNBlockNum; + uint32_t qNBlockIdxCurGroup = qNBlockIdx % qNBlockNumPerGroup; + + uint32_t oSOffset = qSBlockIdx * curQSBlockTile * qHeads * embed; + uint32_t kvNIdx = qNBlockIdx / qNBlockNumPerGroup; + uint32_t qStartNIdx = kvNIdx * groupSize + qNBlockIdxCurGroup * curQNBlockTile; + uint32_t oNOffset = qStartNIdx * embed; + uint32_t gmOffsetO = oBatchOffset + oSOffset + oNOffset; + + // shared sum max workspace offset + // Calculate shared workspace offset for this core's task + // gsharedout size: [batch*beamSize, numHeads, 8] + // Each core handles: [qSBlockSize, qNBlockSize] elements + // Offset = curBatch * beamSize * qHeads * 8 + qSBlockIdx * curQSBlockTile * qHeads * 8 + qStartNIdx * 8; + uint32_t gSharedOffset = curBatch * qSeqlen * qHeads * SOFTMAX_BROAD_SIZE + + qSBlockIdx * curQSBlockTile * qHeads * SOFTMAX_BROAD_SIZE + qStartNIdx * SOFTMAX_BROAD_SIZE; + // cce::printf("gSharedOffset:%d\n", gSharedOffset); + uint32_t qSBlockSize = (qSBlockIdx == (curQSBlockNum - 1)) ? (qSeqlen - qSBlockIdx * curQSBlockTile) + : curQSBlockTile; + uint32_t qNBlockSize = (qNBlockIdxCurGroup == (qNBlockNumPerGroup - 1)) + ? (groupSize - qNBlockIdxCurGroup * curQNBlockTile) + : curQNBlockTile; + uint32_t rowNum = qSBlockSize * qNBlockSize; + uint32_t rowNumRound = RoundUp(rowNum, BLOCK_SIZE); + + uint32_t noSkipKvS = kvSeqlen; + uint32_t noMaskKvS = kvSeqlen; + uint32_t noMaskTailS = 0; + if (maskType != 0) { + uint32_t diffS = kvSeqlen - qSeqlen; + noSkipKvS = (qSBlockIdx + 1) * curQSBlockTile + diffS; + noSkipKvS = Min(kvSeqlen, noSkipKvS); + noMaskKvS = noSkipKvS - qSBlockSize; + noMaskTailS = noMaskKvS % pagedBlockSize; + } + uint32_t maskedKvS = qSBlockSize; + uint32_t kvSLoopNumTotal = CeilDiv(noSkipKvS, pagedBlockSize); + uint32_t kvSLoopNumNoMask = CeilDiv(noMaskKvS, pagedBlockSize); + uint32_t blockStackNum = 4; + uint32_t stackSeqTilePad = blockStackNum * pagedBlockSize; + uint32_t stackSeqTile; + int32_t preLaunch = 2; + // totalStackSeqNum = 1 + int32_t totalStackSeqNum = (maskType != 0) ? (CeilDiv(noMaskKvS, blockStackNum * pagedBlockSize) + 1) + : CeilDiv(noMaskKvS, blockStackNum * pagedBlockSize); + int32_t stackSeqCount = 0; + + // no mask kvSeqlen loop + for (uint32_t kvSIdx = 0; kvSIdx < kvSLoopNumNoMask; kvSIdx += blockStackNum) { + + if (kvSIdx + blockStackNum > kvSLoopNumNoMask - 1) { + stackSeqTile = noMaskKvS - kvSIdx * pagedBlockSize; + } else { + stackSeqTile = pagedBlockSize * blockStackNum; + } + uint32_t isLastStackTile = (kvSIdx + blockStackNum > kvSLoopNumNoMask - 1) ? 1 : 0; + uint32_t stackSeqTileRound = RoundUp(stackSeqTile, BLOCK_SIZE); + LayoutS layOutS(rowNum, stackSeqTile, stackSeqTilePad); + LayoutP layOutP(rowNum, stackSeqTile, stackSeqTilePad); + GemmCoord actualBlockShapeQK{rowNum, stackSeqTile, embed}; + uint32_t curStackTileMod = stackSeqCount % (preLaunch + 1); + uint32_t gmOffsetS = coreIdx * WORKSPACE_BLOCK_SIZE_DB * (preLaunch + 1) + // cube core offset + curStackTileMod * WORKSPACE_BLOCK_SIZE_DB; // single cube core db offset + // vec core offset will be processed within epilogue block + uint32_t gmOffsetP = gmOffsetS; + // AscendC::printf("stackSeqCount:%d\n", stackSeqCount); + Arch::CrossCoreWaitFlag(qkReady); + + // online softmax + epilogueOnlineSoftmax( + gP[gmOffsetP], gS[gmOffsetS], gSharedMax[gSharedOffset], gSharedSum[gSharedOffset], layOutP, layOutS, actualBlockShapeQK, (stackSeqCount == 0), + isLastStackTile, qSBlockSize, qNBlockSize, curStackTileMod, qHeads + ); + Arch::CrossCoreSetFlag<0x2, PIPE_MTE3>(softmaxReady); + + if (kvSIdx >= preLaunch * blockStackNum) { + uint32_t delayedKvSIdx = kvSIdx - preLaunch * blockStackNum; + if (delayedKvSIdx + blockStackNum > kvSLoopNumNoMask - 1) { + stackSeqTile = noMaskKvS - kvSIdx * pagedBlockSize; + } else { + stackSeqTile = pagedBlockSize * blockStackNum; + } + LayoutO layoutO(qSeqlen, embed * qHeads); + LayoutOTmp layoutOTmp(rowNum, embed, embedRound); + GemmCoord actualBlockShapePV{rowNum, embed, stackSeqTile}; + uint32_t curStackTileMod = (stackSeqCount - preLaunch) % (preLaunch + 1); + uint32_t gmOffsetOTmp = coreIdx * WORKSPACE_BLOCK_SIZE_DB * (preLaunch + 1) + + curStackTileMod * WORKSPACE_BLOCK_SIZE_DB; + Arch::CrossCoreWaitFlag(pvReady); + // rescale O + epilogueRescaleO( + gO[gmOffsetO], gOTmp[gmOffsetOTmp], gSharedO[gmOffsetO], layoutO, layoutOTmp, actualBlockShapePV, qSBlockSize, + qNBlockSize, (stackSeqCount - preLaunch == 0), 0, curStackTileMod + ); + } + stackSeqCount++; + } + /* + * for the secondary loop + * while masked, it deals the CV stage1(Qk^t/SMOnline) of the final base block(typical shape [128, 512]), + * and the CV stage2(PV/rescaleO) of the last (prelaunch+1) base blocks while unmasked, it deals the CV + * stage1(Qk^t/SMOnline) of the last (prelaunch+1) base blocks + */ + // deal secondary loop conditions + uint32_t maskedStartIdx = (maskType != 0) ? ((noMaskTailS != 0) ? (kvSLoopNumNoMask - 1) : kvSLoopNumNoMask) + : AlignUp(kvSLoopNumNoMask, blockStackNum); + uint32_t noMaskTailInteStackNum = (noMaskKvS / pagedBlockSize) % blockStackNum; + noMaskTailInteStackNum = (noMaskTailInteStackNum != 0) ? noMaskTailInteStackNum + : ((noMaskTailS != 0) ? 0 : blockStackNum); + uint32_t preLaunchStackNum = (maskType != 0) ? ((preLaunch - 1) * blockStackNum + noMaskTailInteStackNum) + : (preLaunch * blockStackNum); + // masked kvSeqlen loop + // maskedStartIdx = AlignUp(kvSLoopNumNoMask, blockStackNum) + // kvSLoopNumTotal = kvSLoopNumNoMask so maskedStartIdx means kvSIdx >= kvSLoopNumTotal + for (uint32_t kvSIdx = maskedStartIdx; kvSIdx < kvSLoopNumTotal + preLaunchStackNum;) { + // In no mask scenario, kvSIdx will not be less than kvSLoopNumTotal, so it will not enter this if + if ((kvSIdx < kvSLoopNumTotal) && (stackSeqCount <= totalStackSeqNum - 1)) { + // stackSeqTile = maskedKvS; + // uint32_t stackSeqTileRound = RoundUp(stackSeqTile, BLOCK_SIZE); + // LayoutS layOutS(rowNum, stackSeqTile, stackSeqTilePad); + // LayoutP layOutP(rowNum, stackSeqTile, stackSeqTilePad); + // LayoutMask layOutMask(1024, 1024, 1024); + // GemmCoord actualBlockShapeQK{rowNum, stackSeqTile, embed}; + // uint32_t curStackTileMod = stackSeqCount % (preLaunch + 1); + // uint32_t gmOffsetS = coreIdx * WORKSPACE_BLOCK_SIZE_DB * (preLaunch + 1) + // cube core offset + // curStackTileMod * WORKSPACE_BLOCK_SIZE_DB; // single cube core db offset + // // vec core offset will be processed within epilogue block + // uint32_t gmOffsetP = gmOffsetS; + // // online softmax + // epilogueOnlineSoftmax( + // gP[gmOffsetP], gS[gmOffsetS], gMask, layOutP, layOutS, layOutMask, actualBlockShapeQK, + // (stackSeqCount == 0), qSBlockSize, qNBlockSize, curStackTileMod, qkReady + // ); + Arch::CrossCoreSetFlag<0x2, PIPE_MTE3>(softmaxReady); + } + if (kvSIdx >= preLaunchStackNum) { + uint32_t delayedKvSIdx = kvSIdx - preLaunchStackNum; + if (delayedKvSIdx + blockStackNum > kvSLoopNumTotal - 1 && (maskType != 0)) { + stackSeqTile = maskedKvS; + } else if (delayedKvSIdx + blockStackNum > kvSLoopNumNoMask - 1) { + stackSeqTile = noMaskKvS - delayedKvSIdx * pagedBlockSize; + } else { + stackSeqTile = pagedBlockSize * blockStackNum; + } + LayoutO layoutO(qSBlockSize, embed * qHeads); + LayoutOTmp layoutOTmp(rowNum, embed, embedRound); + GemmCoord actualBlockShapePV{rowNum, embed, stackSeqTile}; + uint32_t curStackTileMod = (stackSeqCount - preLaunch) % (preLaunch + 1); + uint32_t gmOffsetOTmp = coreIdx * WORKSPACE_BLOCK_SIZE_DB * (preLaunch + 1) + + curStackTileMod * WORKSPACE_BLOCK_SIZE_DB; + Arch::CrossCoreWaitFlag(pvReady); + // rescale O + epilogueRescaleO( + gO[gmOffsetO], gOTmp[gmOffsetOTmp], gSharedO[gmOffsetO], layoutO, layoutOTmp, actualBlockShapePV, qSBlockSize, + qNBlockSize, (stackSeqCount - preLaunch == 0), + (stackSeqCount - preLaunch == totalStackSeqNum - 1), curStackTileMod + ); + } + if ((maskType != 0) && (stackSeqCount - preLaunch == totalStackSeqNum - 2)) { + kvSIdx += noMaskTailInteStackNum; + } else { + kvSIdx += blockStackNum; + } + stackSeqCount++; + } + } + + AscendC::WaitFlag(EVENT_ID2); + AscendC::WaitFlag(EVENT_ID3); + AscendC::WaitFlag(EVENT_ID4); + AscendC::WaitFlag(EVENT_ID5); + + AscendC::WaitFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID1); + AscendC::WaitFlag(EVENT_ID2); + AscendC::WaitFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID1); + AscendC::WaitFlag(EVENT_ID2); + AscendC::WaitFlag(EVENT_ID3); + } + + private: + Arch::Resource resource; + Arch::CrossCoreFlag qkReady{QK_READY_ID}; + Arch::CrossCoreFlag softmaxReady{SOFTMAX_READY_ID}; + Arch::CrossCoreFlag pvReady{PV_READY_ID}; + XAttentionV2TilingData* faTilingData; + }; + + + /* + FASharedInferKernelShort + Compute Stream 短序列优化场景 + 1. BlockMmadQK + 2. OnlineSoftmax + 3. BlockMmadPV + 4. EpilogueRescaleO + */ + template < + class BlockMmadQK, + class BlockMmadPV, + class EpilogueOnlineSoftmax, + class EpilogueRescaleO, + bool PAGED_CACHE_FLAG = true + > + class SharedFAInferKernelShort { + public: + using ArchTag = typename BlockMmadQK::ArchTag; + using L1TileShape = typename BlockMmadQK::L1TileShape; + using ElementQ = typename BlockMmadQK::ElementA; + using LayoutQ = typename BlockMmadQK::LayoutA; + using ElementK = typename BlockMmadQK::ElementB; + using LayoutK = typename BlockMmadQK::LayoutB; + using ElementS = typename BlockMmadQK::ElementC; + using LayoutS = typename BlockMmadQK::LayoutC; + + using ElementP = typename BlockMmadPV::ElementA; + using LayoutP = typename BlockMmadPV::LayoutA; + using ElementV = typename BlockMmadPV::ElementB; + using LayoutV = typename BlockMmadPV::LayoutB; + + using ElementMask = typename EpilogueOnlineSoftmax::ElementMask; + using LayoutMask = typename EpilogueOnlineSoftmax::LayoutMask; + + using ElementO = typename EpilogueRescaleO::ElementOutput; + using LayoutO = typename EpilogueRescaleO::LayoutOutput; + + using ElementOTmp = typename EpilogueRescaleO::ElementInput; + using LayoutOTmp = typename EpilogueRescaleO::LayoutInput; + + using ElementUpdate = typename EpilogueRescaleO::ElementUpdate; + using LayoutUpdate = typename EpilogueRescaleO::LayoutUpdate; + // Methods + CATLASS_DEVICE + SharedFAInferKernelShort(XAttentionV2TilingData* tilingDataPtr): faTilingData(tilingDataPtr) { + } + + + struct TaskQue { + uint32_t taskIdx; + uint64_t gmOffsetV; + uint64_t gmOffsetO; + uint32_t rowNum; + uint32_t stackSeqTile; + uint32_t kvSIdx; + uint32_t blockBOffset; + uint32_t qSeqlen; + uint32_t qSBlockSize; + uint32_t qNBlockSize; + uint32_t kvSLoopNumTotal; + + CATLASS_DEVICE + TaskQue() + {} + + CATLASS_DEVICE + void SetValue(uint32_t taskIdx_, uint64_t gmOffsetV_, uint64_t gmOffsetO_, uint32_t rowNum_, + uint32_t stackSeqTile_, uint32_t kvSIdx_, uint32_t blockBOffset_, uint32_t qSeqlen_, uint32_t qSBlockSize_, + uint32_t qNBlockSize_, uint32_t kvSLoopNumTotal_) + { + taskIdx = taskIdx_; + gmOffsetV = gmOffsetV_; + gmOffsetO = gmOffsetO_; + rowNum = rowNum_; + stackSeqTile = stackSeqTile_; + kvSIdx = kvSIdx_; + blockBOffset = blockBOffset_; + qSeqlen = qSeqlen_; + qSBlockSize = qSBlockSize_; + qNBlockSize = qNBlockSize_; + kvSLoopNumTotal = kvSLoopNumTotal_; + } + }; + + template + CATLASS_DEVICE void operator()(XAttnKernelParams const ¶ms); + + template <> + CATLASS_DEVICE void operator()(XAttnKernelParams const ¶ms) { + AscendC::SetFlag(EVENT_ID0); + AscendC::SetFlag(EVENT_ID1); + + AscendC::SetFlag(EVENT_ID0); + AscendC::SetFlag(EVENT_ID1); + AscendC::SetFlag(EVENT_ID0); + AscendC::SetFlag(EVENT_ID1); + AscendC::SetFlag(EVENT_ID2); + // TODO L1TileShape可能会有区别 + static constexpr uint32_t L1_QK_SIZE = + BlockMmadQK::L1TileShape::M * BlockMmadQK::L1TileShape::K * sizeof(ElementQ); + + BlockMmadQK blockMmadQK(resource); + BlockMmadPV blockMmadPV(resource, L1_QK_SIZE); + + //__gm__ XATilingData *faTilingData = reinterpret_cast<__gm__ XATilingData *>(params.tiling); + uint64_t mm1OutSize = faTilingData->mm1OutSize; + uint64_t smOnlineOutSize = faTilingData->smOnlineOutSize; + uint32_t batch = faTilingData->batch; // requestNum + uint32_t beamSize = faTilingData->beamSize; + uint32_t qHeads = faTilingData->numHeads; + uint32_t kvHeads = faTilingData->kvHeads; + uint32_t embed = faTilingData->embeddingSize; + uint32_t pagedBlockSize = faTilingData->blockSize; + uint32_t sharedCoreNum = faTilingData->sharedCoreNum; + uint32_t maxNumBlocksPerBatch = faTilingData->maxNumBlocksPerBatch; + uint32_t curTotalTaskNum = faTilingData->firstSharedBatchTaskNum; + uint32_t totalTaskNum = faTilingData->sharedTotalTaskNum; + uint32_t blockSize = faTilingData->blockSize; + uint32_t maskType = faTilingData->maskType; + float scaleValue = faTilingData->scaleValue; + + AscendC::GlobalTensor gQ; + gQ.SetGlobalBuffer((__gm__ ElementQ *)params.q); + AscendC::GlobalTensor gK; + gK.SetGlobalBuffer((__gm__ ElementK *)params.k_cache); + AscendC::GlobalTensor gV; + gV.SetGlobalBuffer((__gm__ ElementK *)params.v_cache); + AscendC::GlobalTensor gBlockTable; + gBlockTable.SetGlobalBuffer((__gm__ int32_t *)(params.sharedBlockTable)); + //AscendC::GlobalTensor gActualQseqlen; + //gActualQseqlen.SetGlobalBuffer((__gm__ int64_t *)params.actualQseqlen); + AscendC::GlobalTensor gActualKvseqlen; + gActualKvseqlen.SetGlobalBuffer((__gm__ int32_t *)params.actualKvseqlen); + AscendC::GlobalTensor gS; + gS.SetGlobalBuffer((__gm__ ElementS *)params.s); + AscendC::GlobalTensor gP; + gP.SetGlobalBuffer((__gm__ ElementP *)params.p); + AscendC::GlobalTensor gOTmp; + gOTmp.SetGlobalBuffer((__gm__ ElementOTmp *)params.oTemp); + AscendC::GlobalTensor gSharedO; + gSharedO.SetGlobalBuffer((__gm__ ElementS *)params.shared_workspace); + + uint64_t strideQO = qHeads * embed; + uint64_t strideKV = kvHeads * embed; + uint32_t embedRound = RoundUp(embed); + uint32_t groupSize = qHeads / kvHeads; + + uint32_t coreIdx = AscendC::GetBlockIdx(); + uint32_t coreNum = sharedCoreNum; // TODO coreNum Need To be modified + uint32_t preTotalTaskNum = 0; + uint32_t curBatch = 0; + uint64_t qBOffset = 0; + uint64_t kBOffset = 0; + uint64_t vBOffset = 0; + uint64_t blockBOffset = 0; + uint64_t oBOffset = 0; + + int64_t qSeqlen = 0; + int64_t kvSeqlen = 0; + uint32_t curQNBlockTile; + uint32_t qNBlockNumPerGroup; + uint32_t curQNBlockNum; + int64_t curQSBlockTile; + uint32_t curQSBlockNum; + uint32_t blockStackNum = 4; + uint32_t stackSeqTilePad = blockStackNum * pagedBlockSize; + + qSeqlen = beamSize; + kvSeqlen = gActualKvseqlen.GetValue(curBatch); + curQSBlockTile = GetQSBlockTile(kvSeqlen); + curQNBlockTile = GetQNBlockTile(qSeqlen, groupSize); + qNBlockNumPerGroup = CeilDiv(groupSize, curQNBlockTile); + curQNBlockNum = qNBlockNumPerGroup * kvHeads; + curQSBlockNum = CeilDiv(qSeqlen, curQSBlockTile); + + uint32_t l1KPPingPongFlag = 0; + uint32_t l0ABPingPongFlag = 0; + uint32_t l0CPingPongFlag = 0; + uint32_t workspaceStagesFlag = 0; + uint32_t loopCnt = 0; + TaskQue taskQue[PRE_LAUNCH + 1]; + LayoutK layoutKTemp(strideKV, stackSeqTilePad); + LayoutV layoutVTemp(stackSeqTilePad, strideKV); + + for (uint32_t taskIdx = coreIdx; taskIdx < totalTaskNum + PRE_LAUNCH * coreNum; taskIdx += uint32_t(coreNum)) { + while(taskIdx >= curTotalTaskNum && taskIdx < totalTaskNum) { + ++curBatch; + preTotalTaskNum = curTotalTaskNum; + qBOffset += qSeqlen * strideQO; + if (!PAGED_CACHE_FLAG) { + kBOffset += kvSeqlen * strideKV; + vBOffset += kvSeqlen * strideKV; + } else { + blockBOffset += maxNumBlocksPerBatch; + } + oBOffset += qSeqlen * strideQO; + + kvSeqlen = gActualKvseqlen.GetValue(curBatch); + curQNBlockTile = GetQNBlockTile(qSeqlen, groupSize); + qNBlockNumPerGroup = CeilDiv(groupSize, curQNBlockTile); + curQNBlockNum = qNBlockNumPerGroup * kvHeads; + curQSBlockTile = GetQSBlockTile(kvSeqlen); + curQSBlockNum = CeilDiv(qSeqlen, curQSBlockTile); + curTotalTaskNum += curQNBlockNum * curQSBlockNum; + } + uint32_t taskIdxCurBatch = taskIdx - preTotalTaskNum; + uint32_t qSBlockIdx = taskIdxCurBatch / curQNBlockNum; + uint32_t qNBlockIdx = taskIdxCurBatch - qSBlockIdx * curQNBlockNum; + uint32_t qNBlockIdxCurGroup = qNBlockIdx % qNBlockNumPerGroup; + + uint32_t kvHeadIdx = qNBlockIdx / qNBlockNumPerGroup; + uint32_t qHeadIdx = kvHeadIdx * groupSize + qNBlockIdxCurGroup * curQNBlockTile; + + uint64_t gmOffsetQ = qBOffset + qSBlockIdx * curQSBlockTile * strideQO + qHeadIdx * embed; + uint64_t gmOffsetK = kBOffset + kvHeadIdx * embed; + uint64_t gmOffsetV = vBOffset + kvHeadIdx * embed; + uint64_t gmOffsetO = oBOffset + qSBlockIdx * curQSBlockTile * strideQO + qHeadIdx * embed; + + uint32_t qSBlockSize = + (qSBlockIdx == (curQSBlockNum - 1)) ? (qSeqlen - qSBlockIdx * curQSBlockTile) : curQSBlockTile; + uint32_t qNBlockSize = (qNBlockIdxCurGroup == (qNBlockNumPerGroup - 1)) + ? (groupSize - qNBlockIdxCurGroup * curQNBlockTile) + : curQNBlockTile; + + uint32_t rowNum = qSBlockSize * qNBlockSize; + uint32_t noSkipKvS = kvSeqlen; + uint32_t kvSLoopNumTotal = CeilDiv(noSkipKvS, pagedBlockSize); + if (taskIdx >= totalTaskNum) { + kvSLoopNumTotal = 1; + } + + uint32_t stackSeqTile; + + LayoutQ layoutQTemp(rowNum, embed); + if (taskIdx < totalTaskNum) { + blockMmadQK.loadQGM(gQ[gmOffsetQ], layoutQTemp, rowNum, qNBlockSize, qHeads); + } + + for (uint32_t kvSIdx = 0; kvSIdx < kvSLoopNumTotal; kvSIdx += blockStackNum) { + if (taskIdx < totalTaskNum) { + if (kvSIdx + blockStackNum > kvSLoopNumTotal - 1) { + stackSeqTile = noSkipKvS - kvSIdx * pagedBlockSize; + } else { + stackSeqTile = stackSeqTilePad; + } + uint64_t gmOffsetS = coreIdx * WORKSPACE_BLOCK_SIZE_DB * (PRE_LAUNCH + 1) + + workspaceStagesFlag * WORKSPACE_BLOCK_SIZE_DB; + uint32_t taskStagesFlag = (taskIdx / coreNum) % (PRE_LAUNCH + 1); + GemmCoord actualBlockShapeQK{rowNum, stackSeqTile, embed}; + LayoutS layOutS(rowNum, stackSeqTile, stackSeqTilePad); + + blockMmadQK(gQ[gmOffsetQ], + gK[gmOffsetK], + gS[gmOffsetS], + gBlockTable[blockBOffset], + layoutQTemp, + layoutKTemp, + layOutS, + actualBlockShapeQK, + kvSIdx, + pagedBlockSize, + strideKV, + l1KPPingPongFlag, + l0ABPingPongFlag, + l0CPingPongFlag); + Arch::CrossCoreSetFlag<0x2, PIPE_FIX>(qkReady); + + taskQue[workspaceStagesFlag].SetValue(taskIdx, + gmOffsetV, + gmOffsetO, + rowNum, + stackSeqTile, + kvSIdx, + blockBOffset, + qSeqlen, + qSBlockSize, + qNBlockSize, + kvSLoopNumTotal); + } + if (loopCnt >= PRE_LAUNCH) { + uint32_t nowWorkspaceStagesFlag = + (workspaceStagesFlag + (PRE_LAUNCH + 1) - PRE_LAUNCH) % (PRE_LAUNCH + 1); + uint32_t nowTaskStagesFlag = (taskQue[nowWorkspaceStagesFlag].taskIdx / coreNum) % (PRE_LAUNCH + 1); + uint32_t nowkvSIdx = taskQue[nowWorkspaceStagesFlag].kvSIdx; + uint32_t nowRowNum = taskQue[nowWorkspaceStagesFlag].rowNum; + stackSeqTile = taskQue[nowWorkspaceStagesFlag].stackSeqTile; + uint64_t gmOffsetOTmp = coreIdx * WORKSPACE_BLOCK_SIZE_DB * (PRE_LAUNCH + 1) + + nowWorkspaceStagesFlag * WORKSPACE_BLOCK_SIZE_DB; + GemmCoord actualBlockShapePV{nowRowNum, embed, stackSeqTile}; + LayoutOTmp layoutOTmp(nowRowNum, embed, embedRound); + LayoutP layoutPTemp(nowRowNum, stackSeqTile, stackSeqTilePad); + uint64_t gmOffsetP = coreIdx * WORKSPACE_BLOCK_SIZE_DB * (PRE_LAUNCH + 1) + + nowWorkspaceStagesFlag * WORKSPACE_BLOCK_SIZE_DB; + blockMmadPV(gP[gmOffsetP], + gV[taskQue[nowWorkspaceStagesFlag].gmOffsetV], + gOTmp[gmOffsetOTmp], + gBlockTable[taskQue[nowWorkspaceStagesFlag].blockBOffset], + layoutPTemp, + layoutVTemp, + layoutOTmp, + actualBlockShapePV, + nowkvSIdx, + pagedBlockSize, + strideKV, + softmaxReady, + l1KPPingPongFlag, + l0ABPingPongFlag, + l0CPingPongFlag); + Arch::CrossCoreSetFlag<0x2, PIPE_FIX>(pvReady); + } + loopCnt++; + workspaceStagesFlag = (workspaceStagesFlag + 1) % (PRE_LAUNCH + 1); + } + + } + + AscendC::WaitFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID1); + + AscendC::WaitFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID1); + + AscendC::WaitFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID1); + AscendC::WaitFlag(EVENT_ID2); + } + + + template <> + CATLASS_DEVICE void operator()(XAttnKernelParams const ¶ms) { + AscendC::SetFlag(EVENT_ID0); + AscendC::SetFlag(EVENT_ID1); + AscendC::SetFlag(EVENT_ID2); + AscendC::SetFlag(EVENT_ID3); + AscendC::SetFlag(EVENT_ID4); + AscendC::SetFlag(EVENT_ID5); + AscendC::SetFlag(EVENT_ID6); + + AscendC::SetFlag(EVENT_ID0); + AscendC::SetFlag(EVENT_ID1); + AscendC::SetFlag(EVENT_ID3); + AscendC::SetFlag(EVENT_ID2); + AscendC::SetFlag(EVENT_ID2); + + // Get tiling parameters + //__gm__ XATilingData *faTilingData = reinterpret_cast<__gm__ XATilingData *>(params.tiling); + uint64_t mm1OutSize = faTilingData->mm1OutSize; + uint64_t smOnlineOutSize = faTilingData->smOnlineOutSize; + uint64_t mm2OutSize = faTilingData->mm2OutSize; + uint32_t batch = faTilingData->batch; + uint32_t beamSize = faTilingData->beamSize; + uint32_t qHeads = faTilingData->numHeads; + uint32_t kvHeads = faTilingData->kvHeads; + uint32_t embed = faTilingData->embeddingSize; + uint32_t sharedCoreNum = faTilingData->sharedCoreNum; + uint32_t pagedBlockSize = faTilingData->blockSize; + uint32_t maxNumBlocksPerBatch = faTilingData->maxNumBlocksPerBatch; + uint32_t firstBatchTaskNum = faTilingData->firstSharedBatchTaskNum; + uint32_t totalTaskNum = faTilingData->sharedTotalTaskNum; + uint32_t maskType = faTilingData->maskType; + float scaleValue = faTilingData->scaleValue; + + uint64_t gOffsetTempO = batch * beamSize * qHeads * embed; + uint64_t gMaxOffset = batch * beamSize * qHeads * SOFTMAX_BROAD_SIZE; + // Get the memory offset address of the input on Global Memory + //AscendC::GlobalTensor gMask; + //gMask.SetGlobalBuffer((__gm__ ElementMask *)params.mask); + //AscendC::GlobalTensor gActualQseqlen; + //gActualQseqlen.SetGlobalBuffer((__gm__ int64_t *)params.actualQseqlen); + AscendC::GlobalTensor gActualKvseqlen; + gActualKvseqlen.SetGlobalBuffer((__gm__ int32_t *)params.actualKvseqlen); + AscendC::GlobalTensor gO; + gO.SetGlobalBuffer((__gm__ ElementO *)params.o); + AscendC::GlobalTensor gS; + gS.SetGlobalBuffer((__gm__ ElementS *)params.s); + AscendC::GlobalTensor gP; + gP.SetGlobalBuffer((__gm__ ElementP *)params.p); + AscendC::GlobalTensor gOTmp; + gOTmp.SetGlobalBuffer((__gm__ ElementOTmp *)params.oTemp); + AscendC::GlobalTensor gOUpdate; + gOUpdate.SetGlobalBuffer((__gm__ ElementOTmp *)params.oUpdate); + + // shared Gm and Gl output + AscendC::GlobalTensor gSharedO; + AscendC::GlobalTensor gSharedSum; + AscendC::GlobalTensor gSharedMax; + gSharedO.SetGlobalBuffer((__gm__ float *)params.shared_workspace); + gSharedMax.SetGlobalBuffer((__gm__ ElementS *)params.shared_workspace + gOffsetTempO); + gSharedSum.SetGlobalBuffer((__gm__ ElementS *)params.shared_workspace + gOffsetTempO + gMaxOffset); + + + uint32_t groupSize = qHeads / kvHeads; + uint32_t embedRound = RoundUp(embed, BLOCK_SIZE); + + EpilogueOnlineSoftmax epilogueOnlineSoftmax(resource, scaleValue); + EpilogueRescaleO epilogueRescaleO(resource); + + // uint32_t curTotalTaskNum = 0; + uint32_t preTotalTaskNum = 0; + uint32_t curBatch = 0; + uint32_t oBOffset = 0; + uint32_t qSeqlen = static_cast(beamSize); + uint32_t kvSeqlen = static_cast(gActualKvseqlen.GetValue(curBatch)); + uint32_t curQNBlockTile = GetQNBlockTile(qSeqlen, groupSize); + uint32_t qNBlockNumPerGroup = CeilDiv(groupSize, curQNBlockTile); + uint32_t curQNBlockNum = qNBlockNumPerGroup * kvHeads; + uint32_t curQSBlockTile = GetQSBlockTile(kvSeqlen); + uint32_t curQSBlockNum = CeilDiv(qSeqlen, curQSBlockTile); + uint32_t curTotalTaskNum = firstBatchTaskNum; + uint32_t blockStackNum = 4; + uint32_t stackSeqTilePad = blockStackNum * pagedBlockSize; + uint32_t coreIdx = AscendC::GetBlockIdx() / AscendC::GetSubBlockNum(); + // coreNum Need to be changed TODO + uint32_t coreNum = sharedCoreNum; + + TaskQue taskQue[PRE_LAUNCH + 1]; + uint32_t loopCnt = 0; + uint32_t workspaceStagesFlag = 0; + // Go through each task. + for (uint32_t taskIdx = coreIdx; taskIdx < totalTaskNum + PRE_LAUNCH * coreNum; taskIdx += uint32_t(coreNum)) { + // Get the offset of each core on the GM. + while (taskIdx >= curTotalTaskNum && taskIdx < totalTaskNum) { + curBatch++; + oBOffset += qSeqlen * qHeads * embed; + preTotalTaskNum = curTotalTaskNum; + qSeqlen = static_cast(beamSize); + kvSeqlen = static_cast(gActualKvseqlen.GetValue(curBatch)); + curQNBlockTile = GetQNBlockTile(qSeqlen, groupSize); + qNBlockNumPerGroup = CeilDiv(groupSize, curQNBlockTile); + curQNBlockNum = qNBlockNumPerGroup * kvHeads; + curQSBlockTile = GetQSBlockTile(kvSeqlen); + curQSBlockNum = CeilDiv(qSeqlen, curQSBlockTile); + curTotalTaskNum += curQNBlockNum * curQSBlockNum; + } + uint32_t taskIdxCurBatch = taskIdx - preTotalTaskNum; + uint32_t qSBlockIdx = taskIdxCurBatch / curQNBlockNum; + uint32_t qNBlockIdx = taskIdxCurBatch % curQNBlockNum; + uint32_t qNBlockIdxCurGroup = qNBlockIdx % qNBlockNumPerGroup; + + uint32_t oSOffset = qSBlockIdx * curQSBlockTile * qHeads * embed; + uint32_t kvNIdx = qNBlockIdx / qNBlockNumPerGroup; + uint32_t qStartNIdx = kvNIdx * groupSize + qNBlockIdxCurGroup * curQNBlockTile; + uint32_t oNOffset = qStartNIdx * embed; + uint32_t gmOffsetO = oBOffset + oSOffset + oNOffset; + + // shared sum max workspace offset + // Calculate shared workspace offset for this core's task + // gsharedout size: [batch*beamSize, numHeads, 8] + // Each core handles: [qSBlockSize, qNBlockSize] elements + // Offset = curBatch * beamSize * qHeads * 8 + qSBlockIdx * curQSBlockTile * qHeads * 8 + qStartNIdx * 8; + uint32_t gSharedOffset = curBatch * qSeqlen * qHeads * SOFTMAX_BROAD_SIZE + + qSBlockIdx * curQSBlockTile * qHeads * SOFTMAX_BROAD_SIZE + qStartNIdx * SOFTMAX_BROAD_SIZE; + + + uint32_t qSBlockSize = (qSBlockIdx == (curQSBlockNum - 1)) ? (qSeqlen - qSBlockIdx * curQSBlockTile) + : curQSBlockTile; + uint32_t qNBlockSize = (qNBlockIdxCurGroup == (qNBlockNumPerGroup - 1)) + ? (groupSize - qNBlockIdxCurGroup * curQNBlockTile) + : curQNBlockTile; + uint32_t rowNum = qSBlockSize * qNBlockSize; + uint32_t rowNumRound = RoundUp(rowNum, BLOCK_SIZE); + + uint32_t noSkipKvS = kvSeqlen; + uint32_t noMaskKvS = kvSeqlen; + uint32_t noMaskTailS = 0; + + uint32_t stackSeqTile = 0; + uint32_t kvSLoopNumTotal = CeilDiv(noSkipKvS, pagedBlockSize); + if (taskIdx >= totalTaskNum) { + kvSLoopNumTotal = 1; + } + uint32_t isLastStackTile = 0; + // no mask kvSeqlen loop + for (uint32_t kvSIdx = 0; kvSIdx < kvSLoopNumTotal; kvSIdx += blockStackNum) { + if (taskIdx < totalTaskNum) { + if (kvSIdx + blockStackNum > kvSLoopNumTotal - 1) { + stackSeqTile = noSkipKvS - kvSIdx * pagedBlockSize; + isLastStackTile = 1; + } else { + stackSeqTile = stackSeqTilePad; + } + uint64_t gmOffsetS = coreIdx * WORKSPACE_BLOCK_SIZE_DB * (PRE_LAUNCH + 1) + + workspaceStagesFlag * WORKSPACE_BLOCK_SIZE_DB; + uint32_t taskStagesFlag = (taskIdx / coreNum) % (PRE_LAUNCH + 1); + GemmCoord actualBlockShapeQK{rowNum, stackSeqTile, embed}; + LayoutS layOutS(rowNum, stackSeqTile, stackSeqTilePad); + LayoutP layOutP(rowNum, stackSeqTile, stackSeqTilePad); + uint64_t gmOffsetP = gmOffsetS; + uint32_t kvSStartIdx = kvSIdx * pagedBlockSize; + uint32_t kvSEndIdx = kvSStartIdx + stackSeqTile; + Arch::CrossCoreWaitFlag(qkReady); + epilogueOnlineSoftmax( + gP[gmOffsetP], gS[gmOffsetS], gSharedMax[gSharedOffset], gSharedSum[gSharedOffset], + layOutP, layOutS, actualBlockShapeQK, (kvSIdx == 0), + isLastStackTile, qSBlockSize, qNBlockSize, workspaceStagesFlag, qHeads + ); + Arch::CrossCoreSetFlag<0x2, PIPE_MTE3>(softmaxReady); + taskQue[workspaceStagesFlag].SetValue(taskIdx, + 0, + gmOffsetO, + rowNum, + stackSeqTile, + kvSIdx, + 0, + qSeqlen, + qSBlockSize, + qNBlockSize, + kvSLoopNumTotal); + + } + if (loopCnt >= PRE_LAUNCH) { + uint32_t nowWorkspaceStagesFlag = + (workspaceStagesFlag + (PRE_LAUNCH + 1) - PRE_LAUNCH) % (PRE_LAUNCH + 1); + uint32_t nowTaskStagesFlag = (taskQue[nowWorkspaceStagesFlag].taskIdx / coreNum) % (PRE_LAUNCH + 1); + uint32_t nowkvSIdx = taskQue[nowWorkspaceStagesFlag].kvSIdx; + uint32_t nowRowNum = taskQue[nowWorkspaceStagesFlag].rowNum; + stackSeqTile = taskQue[nowWorkspaceStagesFlag].stackSeqTile; + // uint32_t curStackTileMod = (stackSeqCount - PRE_LAUNCH) % (PRE_LAUNCH + 1); + uint64_t gmOffsetOTmp = coreIdx * WORKSPACE_BLOCK_SIZE_DB * (PRE_LAUNCH + 1) + + nowWorkspaceStagesFlag * WORKSPACE_BLOCK_SIZE_DB; + GemmCoord actualBlockShapePV{nowRowNum, embed, stackSeqTile}; + LayoutOTmp layoutOTmp(nowRowNum, embed, embedRound); + LayoutO layoutO(taskQue[nowWorkspaceStagesFlag].qSeqlen, embed * qHeads); + LayoutUpdate layoutUpdate(nowRowNum, embed, embedRound); + uint64_t gmOffsetUpdate = (uint64_t)(coreIdx * WORKSPACE_BLOCK_SIZE_DB); + + Arch::CrossCoreWaitFlag(pvReady); + // rescale O + epilogueRescaleO( + gO[taskQue[nowWorkspaceStagesFlag].gmOffsetO], + gOTmp[gmOffsetOTmp], + gSharedO[taskQue[nowWorkspaceStagesFlag].gmOffsetO], + layoutO, layoutOTmp, actualBlockShapePV, + taskQue[nowWorkspaceStagesFlag].qSBlockSize, + taskQue[nowWorkspaceStagesFlag].qNBlockSize, + (nowkvSIdx == 0), + nowkvSIdx + blockStackNum >= taskQue[nowWorkspaceStagesFlag].kvSLoopNumTotal, + nowWorkspaceStagesFlag + ); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + } + loopCnt++; + workspaceStagesFlag = (workspaceStagesFlag + 1) % (PRE_LAUNCH + 1); + + } + } + + AscendC::WaitFlag(EVENT_ID2); + AscendC::WaitFlag(EVENT_ID3); + AscendC::WaitFlag(EVENT_ID4); + AscendC::WaitFlag(EVENT_ID5); + AscendC::WaitFlag(EVENT_ID6); + + AscendC::WaitFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID1); + AscendC::WaitFlag(EVENT_ID2); + AscendC::WaitFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID1); + AscendC::WaitFlag(EVENT_ID2); + AscendC::WaitFlag(EVENT_ID3); + } + + private: + Arch::Resource resource; + Arch::CrossCoreFlag qkReady{QK_READY_ID}; + Arch::CrossCoreFlag softmaxReady{SOFTMAX_READY_ID}; + Arch::CrossCoreFlag pvReady{PV_READY_ID}; + XAttentionV2TilingData* faTilingData; + }; + + +template +class CombineScaleKernel { +public: + using ArchTag = typename EpilogueCombineScale::ArchTag; + using ElementOutput = typename EpilogueCombineScale::ElementOutput; + using ElementInput = typename EpilogueCombineScale::ElementInput; + + CATLASS_DEVICE + CombineScaleKernel(XAttentionV2TilingData* tilingDataPtr): faTilingData(tilingDataPtr) {} + + template + CATLASS_DEVICE void operator()(XAttnKernelParams const ¶ms); + + template <> + CATLASS_DEVICE void operator()(XAttnKernelParams const ¶ms) { + return; + } + + template <> + CATLASS_DEVICE void operator()(XAttnKernelParams const ¶ms) { + AscendC::SetFlag(EVENT_ID0); + AscendC::SetFlag(EVENT_ID1); + AscendC::SetFlag(EVENT_ID2); + AscendC::SetFlag(EVENT_ID3); + AscendC::SetFlag(EVENT_ID4); + AscendC::SetFlag(EVENT_ID5); + + AscendC::SetFlag(EVENT_ID0); + AscendC::SetFlag(EVENT_ID1); + AscendC::SetFlag(EVENT_ID3); + AscendC::SetFlag(EVENT_ID2); + AscendC::SetFlag(EVENT_ID2); + + EpilogueCombineScale epilogueCombineScale(resource); + + // get tiling params + // __gm__ XATilingData *faTilingData = reinterpret_cast<__gm__ XATilingData *>(params.tiling); + uint32_t combineFormerCoreNum = faTilingData->combineFormerCoreNum; + uint32_t combineFormerRowNum = faTilingData->combineFormerRowNum; + uint32_t combineTailRowNum = faTilingData->combineTailRowNum; + uint32_t numTokens = faTilingData->numTokens; + uint32_t qHeads = faTilingData->numHeads; + uint32_t coreIdx = AscendC::GetBlockIdx() / AscendC::GetSubBlockNum(); + uint32_t coreNum = AscendC::GetBlockNum(); + uint32_t currentCoreRowNum = 0; + uint32_t embed = faTilingData->embeddingSize; + uint32_t gmOffsetO = 0; + uint32_t gmOffsetSharedGl = 0; + uint32_t gmOffsetSharedGm = 0; + uint32_t gmOffsetUnSharedGl = 0; + uint32_t gmOffsetUnSharedGm = 0; + uint32_t sharedSumMaxGmOffset = numTokens * qHeads * SOFTMAX_BROAD_SIZE; + uint32_t unsharedSumMaxGmOffset = numTokens * qHeads; + uint32_t attnOutputOffset = numTokens * qHeads * embed; + if (coreIdx >= faTilingData->combineCoreNum) { + return; + } + + if (coreIdx < combineFormerCoreNum) { + currentCoreRowNum = combineFormerRowNum; + gmOffsetO = coreIdx * combineFormerRowNum * embed; + gmOffsetUnSharedGl = coreIdx *combineFormerRowNum; + gmOffsetUnSharedGm = gmOffsetUnSharedGl; + gmOffsetSharedGl = coreIdx *combineFormerRowNum * SOFTMAX_BROAD_SIZE; + gmOffsetSharedGm = gmOffsetSharedGl; + } else { + currentCoreRowNum = combineTailRowNum; + gmOffsetO = combineFormerCoreNum * combineFormerRowNum * embed + + (coreIdx - combineFormerCoreNum) * combineTailRowNum * embed; + gmOffsetUnSharedGl = combineFormerCoreNum * combineFormerRowNum + + (coreIdx - combineFormerCoreNum) * combineTailRowNum; + gmOffsetUnSharedGm = gmOffsetUnSharedGl; + gmOffsetSharedGl = gmOffsetUnSharedGl * SOFTMAX_BROAD_SIZE; + gmOffsetSharedGm = gmOffsetSharedGl; + } + // shared_workspace [attnout:[attnOut * 4Bytes], + // gm:[attnOutputOffset * 4Bytes], + // gl:[attnOutputOffset * 4Bytes]] + AscendC::GlobalTensor gSharedGm; + gSharedGm.SetGlobalBuffer((__gm__ ElementInput *)params.shared_workspace + attnOutputOffset); + AscendC::GlobalTensor gSharedGl; + gSharedGl.SetGlobalBuffer((__gm__ ElementInput *)params.shared_workspace + attnOutputOffset + + sharedSumMaxGmOffset); + AscendC::GlobalTensor gUnsharedGm; + gUnsharedGm.SetGlobalBuffer((__gm__ ElementInput *)params.unshared_workspace + attnOutputOffset); + AscendC::GlobalTensor gUnsharedGl; + gUnsharedGl.SetGlobalBuffer((__gm__ ElementInput *)params.unshared_workspace + attnOutputOffset + + unsharedSumMaxGmOffset); + AscendC::GlobalTensor gSharedOut; + gSharedOut.SetGlobalBuffer((__gm__ ElementInput *)params.shared_workspace); + AscendC::GlobalTensor gUnsharedOut; + gUnsharedOut.SetGlobalBuffer((__gm__ ElementInput *)params.unshared_workspace); + AscendC::GlobalTensor gFinalOut; + gFinalOut.SetGlobalBuffer((__gm__ ElementOutput *)params.o); + + MatrixCoord actualBlockShape(currentCoreRowNum, embed); + //cce::printf("coreIdx:%d, gmOffsetO:%d, currentCoreRowNum:%d, gmOffsetGl:%d, gmOffsetGm:%d\n", AscendC::GetBlockIdx(), + //gmOffsetO, currentCoreRowNum, gmOffsetGl, gmOffsetGm); + epilogueCombineScale( + gSharedGm[gmOffsetSharedGm], gUnsharedGm[gmOffsetUnSharedGm], + gSharedGl[gmOffsetSharedGl], gUnsharedGl[gmOffsetUnSharedGl], + gSharedOut[gmOffsetO], gUnsharedOut[gmOffsetO], + gFinalOut[gmOffsetO], actualBlockShape + ); + + // AscendC::PipeBarrier(); + + + AscendC::WaitFlag(EVENT_ID2); + AscendC::WaitFlag(EVENT_ID3); + AscendC::WaitFlag(EVENT_ID4); + AscendC::WaitFlag(EVENT_ID5); + + AscendC::WaitFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID1); + AscendC::WaitFlag(EVENT_ID2); + AscendC::WaitFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID1); + AscendC::WaitFlag(EVENT_ID2); + AscendC::WaitFlag(EVENT_ID3); + } +private: + Arch::Resource resource; + XAttentionV2TilingData* faTilingData; +}; + +#endif \ No newline at end of file