diff --git a/test/python_test/RegisterOps.cpp b/test/python_test/RegisterOps.cpp index ab0300c..828c208 100644 --- a/test/python_test/RegisterOps.cpp +++ b/test/python_test/RegisterOps.cpp @@ -1430,6 +1430,160 @@ std::tuple moe_gating_top_k_hash_impl_npu( return std::make_tuple(y, expert_idx, out); } +std::tuple +qwen35_gdn_decode_super_op( + const at::Tensor& qkv, + const at::Tensor& z, + const at::Tensor& b, + const at::Tensor& a, + const at::Tensor& conv_weight, + at::Tensor& conv_state, + const at::Tensor& a_log, + const at::Tensor& dt_bias, + at::Tensor& ssm_state, + const at::Tensor& state_indices, + const at::Tensor& norm_weight) { + at::Tensor conv_out = at::empty_like(qkv); + at::Tensor out = at::empty_like(z); + EXEC_NPU_CMD(aclnnQwen35GdnDecodeSuperOp, + qkv, + z, + b, + a, + conv_weight, + conv_state, + a_log, + dt_bias, + ssm_state, + state_indices, + norm_weight, + conv_out, + conv_state, + ssm_state, + out); + return std::tuple( + conv_out, conv_state, ssm_state, out); +} + +std::tuple +qwen35_gdn_prefill_super_op( + const at::Tensor& mixed_qkv, + const at::Tensor& z, + const at::Tensor& b, + const at::Tensor& a, + const at::Tensor& conv_weight, + at::Tensor& conv_state, + const at::Tensor& a_log, + const at::Tensor& dt_bias, + at::Tensor& ssm_state, + const at::Tensor& norm_weight, + const at::Tensor& mask_lower, + const at::Tensor& mask_full, + const at::Tensor& minus_identity, + const at::Tensor& cu_seqlens, + int64_t conv_state_index, + int64_t ssm_state_index) { + constexpr int64_t kChunkSize = 128; + constexpr int64_t kHeadDim = 128; + constexpr int64_t kValueHeads = 24; + const int64_t total_tokens = mixed_qkv.size(0); + const int64_t num_matrices = + ((total_tokens + kChunkSize - 1) / kChunkSize) * kValueHeads; + auto opts_fp16 = mixed_qkv.options().dtype(at::kHalf); + auto opts_fp32 = mixed_qkv.options().dtype(at::kFloat); + + at::Tensor packed_qkv = at::empty({total_tokens, 5120}, opts_fp16); + at::Tensor g = at::empty({1, total_tokens, kValueHeads}, opts_fp32); + at::Tensor beta = at::empty({1, total_tokens, kValueHeads}, opts_fp16); + at::Tensor initial_state = + at::empty({1, kValueHeads, kHeadDim, kHeadDim}, opts_fp16); + at::Tensor mega_out = + at::empty({1, total_tokens, kValueHeads, kHeadDim}, opts_fp16); + at::Tensor g_sum = + at::empty({1, total_tokens, kValueHeads}, opts_fp32); + at::Tensor g_t = at::empty({kValueHeads, total_tokens}, opts_fp32); + at::Tensor beta_t = at::empty({kValueHeads, total_tokens}, opts_fp16); + at::Tensor mega_a = + at::zeros({1, total_tokens, kValueHeads, kChunkSize}, opts_fp16); + at::Tensor a_inv_f32 = + at::zeros({1, total_tokens, kValueHeads, kChunkSize}, opts_fp32); + at::Tensor a_inv = + at::zeros({1, total_tokens, kValueHeads, kChunkSize}, opts_fp16); + at::Tensor w = + at::empty({1, total_tokens, kValueHeads, kHeadDim}, opts_fp16); + at::Tensor u = at::empty_like(w); + at::Tensor h = + at::zeros({num_matrices, kHeadDim, kHeadDim}, opts_fp16); + at::Tensor v_new = at::empty_like(w); + at::Tensor final_state = + at::zeros({kValueHeads, kHeadDim, kHeadDim}, opts_fp16); + at::Tensor out = + at::empty({total_tokens, kValueHeads, kHeadDim}, mixed_qkv.options()); + + EXEC_NPU_CMD(aclnnQwen35GdnPrefillSuperOp, + mixed_qkv, + z, + b, + a, + conv_weight, + conv_state, + a_log, + dt_bias, + ssm_state, + norm_weight, + mask_lower, + mask_full, + minus_identity, + cu_seqlens, + num_matrices, + conv_state_index, + ssm_state_index, + packed_qkv, + g, + beta, + initial_state, + mega_out, + g_sum, + g_t, + beta_t, + mega_a, + a_inv_f32, + a_inv, + w, + u, + h, + v_new, + final_state, + conv_state, + ssm_state, + out); + return std::tuple(packed_qkv, + g, + beta, + initial_state, + mega_out, + final_state, + conv_state, + ssm_state, + out); +} + 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"); @@ -1452,6 +1606,12 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("recurrent_gated_delta_rule", &recurrent_gated_delta_rule, "recurrent_gated_delta_rule"); m.def("rec_constrained_topk", &rec_constrained_topk_impl_npu, "rec_constrained_topk"); m.def("mega_chunk_gdn", &mega_chunk_gdn, "mega_chunk_gdn"); + m.def("qwen35_gdn_decode_super_op", + &qwen35_gdn_decode_super_op, + "qwen35_gdn_decode_super_op"); + m.def("qwen35_gdn_prefill_super_op", + &qwen35_gdn_prefill_super_op, + "qwen35_gdn_prefill_super_op"); m.def("layer_norm_fwd", &layer_norm_fwd_impl_npu, "layer_norm_fwd"); m.def("moe_fused_add_topk", &moe_fused_add_topk_impl_npu, "moe_fused_add_topk"); m.def("moe_fused_reducesum_div", &moe_fused_reducesum_div_impl_npu, "moe_fused_reducesum_div"); diff --git a/test/python_test/test_qwen35_gdn_decode_super_op.py b/test/python_test/test_qwen35_gdn_decode_super_op.py new file mode 100644 index 0000000..4daddf6 --- /dev/null +++ b/test/python_test/test_qwen35_gdn_decode_super_op.py @@ -0,0 +1,89 @@ +# Copyright 2026 The xLLM Authors. All Rights Reserved. + +import pytest +import torch + +torch_npu = pytest.importorskip("torch_npu") +custom_ops = pytest.importorskip("custom_ops_lib") + + +HEAD_DIM = 128 +LOCAL_HEAD_SHAPES = [ + (num_k_heads, num_k_heads * ratio) + for ratio in range(1, 5) + for num_k_heads in (1, 2, 4, 8, 16) +] + + +def _has_npu(): + return hasattr(torch, "npu") and torch.npu.is_available() + + +pytestmark = pytest.mark.skipif(not _has_npu(), reason="NPU is required") + + +def _run_case(num_k_heads, num_v_heads, batch_size): + conv_dim = (2 * num_k_heads + num_v_heads) * HEAD_DIM + bf16_options = {"device": "npu:0", "dtype": torch.bfloat16} + float_options = {"device": "npu:0", "dtype": torch.float32} + + qkv = torch.ones((batch_size, conv_dim), **bf16_options) + z = torch.zeros((batch_size, num_v_heads, HEAD_DIM), **bf16_options) + b = torch.zeros((batch_size, num_v_heads), **bf16_options) + a = torch.zeros((batch_size, num_v_heads), **bf16_options) + conv_weight = torch.zeros((4, conv_dim), **bf16_options) + conv_state = torch.zeros((batch_size, 3, conv_dim), **bf16_options) + a_log = torch.zeros((num_v_heads,), **float_options) + dt_bias = torch.zeros((num_v_heads,), **float_options) + ssm_state = torch.ones( + (batch_size, num_v_heads, HEAD_DIM, HEAD_DIM), **float_options + ) + state_indices = torch.arange(batch_size, device="npu:0", dtype=torch.int32) + norm_weight = torch.ones((HEAD_DIM,), **bf16_options) + + outputs = custom_ops.qwen35_gdn_decode_super_op( + qkv, + z, + b, + a, + conv_weight, + conv_state, + a_log, + dt_bias, + ssm_state, + state_indices, + norm_weight, + ) + torch.npu.synchronize() + + assert torch.equal(outputs[3], torch.zeros_like(outputs[3])) + assert torch.equal(conv_state[:, :2], torch.zeros_like(conv_state[:, :2])) + assert torch.equal(conv_state[:, 2], qkv) + assert torch.allclose( + ssm_state, + torch.full_like(ssm_state, 0.5), + rtol=0.0, + atol=1e-6, + ) + + +@pytest.mark.parametrize( + ("num_k_heads", "num_v_heads"), + LOCAL_HEAD_SHAPES, +) +def test_all_qwen35_local_head_shapes(num_k_heads, num_v_heads): + _run_case(num_k_heads, num_v_heads, batch_size=1) + + +@pytest.mark.parametrize( + ("num_k_heads", "num_v_heads", "batch_size"), + [ + (8, 16, 2), + (8, 24, 3), + (8, 24, 4), + (8, 24, 32), + (16, 64, 2), + ], +) +def test_representative_batch_shapes(num_k_heads, num_v_heads, batch_size): + _run_case(num_k_heads, num_v_heads, batch_size) diff --git a/xllm_ops/build_aclnn.sh b/xllm_ops/build_aclnn.sh index ae48ac3..3a70e4b 100644 --- a/xllm_ops/build_aclnn.sh +++ b/xllm_ops/build_aclnn.sh @@ -133,6 +133,8 @@ elif [[ "$SOC_VERSION" =~ ^(ascend)?910b ]]; then "cache_unshared_kv" "causal_conv1d" "causal_conv1d_qkv" + "qwen35_gdn_decode_super_op" + "qwen35_gdn_prefill_super_op" "convert_kv_cache_format" "beam_search" "index_group_matmul" @@ -243,6 +245,8 @@ elif [[ "$SOC_VERSION" =~ ^ascend910_93 ]]; then "cache_unshared_kv" "causal_conv1d" "causal_conv1d_qkv" + "qwen35_gdn_decode_super_op" + "qwen35_gdn_prefill_super_op" "convert_kv_cache_format" "beam_search" "index_group_matmul" diff --git a/xllm_ops/mega_chunk_gdn/op_kernel/mega_chunk_gdn.cpp b/xllm_ops/mega_chunk_gdn/op_kernel/mega_chunk_gdn.cpp index 04f7326..84dbdeb 100644 --- a/xllm_ops/mega_chunk_gdn/op_kernel/mega_chunk_gdn.cpp +++ b/xllm_ops/mega_chunk_gdn/op_kernel/mega_chunk_gdn.cpp @@ -374,8 +374,6 @@ AICORE inline void mega_kernel_impl(GM_ADDR q_ptr, GM_ADDR k_ptr, GM_ADDR v_ptr, return; #endif - SyncAllImpl(); - #ifdef MEGA_STOP_AFTER_SYNC_BEFORE_WY return; #endif diff --git a/xllm_ops/qwen35_gdn_decode_super_op/CMakeLists.txt b/xllm_ops/qwen35_gdn_decode_super_op/CMakeLists.txt new file mode 100644 index 0000000..3e31e7f --- /dev/null +++ b/xllm_ops/qwen35_gdn_decode_super_op/CMakeLists.txt @@ -0,0 +1,7 @@ +file(GLOB CURRENT_DIRS RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*) +foreach(SUB_DIR ${CURRENT_DIRS}) + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/${SUB_DIR}/CMakeLists.txt") + add_subdirectory(${SUB_DIR}) + endif() +endforeach() +add_op_to_compiled_list() diff --git a/xllm_ops/qwen35_gdn_decode_super_op/op_host/CMakeLists.txt b/xllm_ops/qwen35_gdn_decode_super_op/op_host/CMakeLists.txt new file mode 100644 index 0000000..84648e2 --- /dev/null +++ b/xllm_ops/qwen35_gdn_decode_super_op/op_host/CMakeLists.txt @@ -0,0 +1,17 @@ +if (BUILD_OPEN_PROJECT) + target_sources(op_host_aclnn PRIVATE + qwen35_gdn_decode_super_op_def.cpp + ) +endif() + +add_ops_compile_options( + OP_NAME Qwen35GdnDecodeSuperOp + OPTIONS --cce-auto-sync=on + -Wno-deprecated-declarations + -Wno-error + -I${CMAKE_CURRENT_SOURCE_DIR}/../../../third_party/pto-isa/include +) + +if (NOT BUILD_OPS_RTY_KERNEL) + add_modules_sources(OPTYPE qwen35_gdn_decode_super_op ACLNNTYPE aclnn) +endif() diff --git a/xllm_ops/qwen35_gdn_decode_super_op/op_host/qwen35_gdn_decode_super_op_def.cpp b/xllm_ops/qwen35_gdn_decode_super_op/op_host/qwen35_gdn_decode_super_op_def.cpp new file mode 100644 index 0000000..9294c38 --- /dev/null +++ b/xllm_ops/qwen35_gdn_decode_super_op/op_host/qwen35_gdn_decode_super_op_def.cpp @@ -0,0 +1,81 @@ +#include "register/op_def_registry.h" + +namespace ops { + +class Qwen35GdnDecodeSuperOp : public OpDef { +public: + explicit Qwen35GdnDecodeSuperOp(const char* name) : OpDef(name) + { + this->Input("qkv").ParamType(REQUIRED).DataType({ge::DT_BF16}).FormatList({ge::FORMAT_ND}).AutoContiguous(); + this->Input("z").ParamType(REQUIRED).DataType({ge::DT_BF16}).FormatList({ge::FORMAT_ND}).AutoContiguous(); + this->Input("b").ParamType(REQUIRED).DataType({ge::DT_BF16}).FormatList({ge::FORMAT_ND}).AutoContiguous(); + this->Input("a").ParamType(REQUIRED).DataType({ge::DT_BF16}).FormatList({ge::FORMAT_ND}).AutoContiguous(); + this->Input("convWeight") + .ParamType(REQUIRED) + .DataType({ge::DT_BF16}) + .FormatList({ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("convState") + .ParamType(REQUIRED) + .DataType({ge::DT_BF16}) + .FormatList({ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("aLog") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT}) + .FormatList({ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("dtBias") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT}) + .FormatList({ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("ssmState") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT}) + .FormatList({ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("stateIndices") + .ParamType(REQUIRED) + .DataType({ge::DT_INT32}) + .FormatList({ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("normWeight") + .ParamType(REQUIRED) + .DataType({ge::DT_BF16}) + .FormatList({ge::FORMAT_ND}) + .AutoContiguous(); + + this->Output("convOut") + .ParamType(REQUIRED) + .DataType({ge::DT_BF16}) + .FormatList({ge::FORMAT_ND}) + .AutoContiguous(); + this->Output("convStateOut") + .ParamType(REQUIRED) + .DataType({ge::DT_BF16}) + .FormatList({ge::FORMAT_ND}) + .AutoContiguous(); + this->Output("ssmStateOut") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT}) + .FormatList({ge::FORMAT_ND}) + .AutoContiguous(); + this->Output("out").ParamType(REQUIRED).DataType({ge::DT_BF16}).FormatList({ge::FORMAT_ND}).AutoContiguous(); + + OpAICoreConfig config; + config.DynamicCompileStaticFlag(true) + .DynamicFormatFlag(false) + .DynamicRankSupportFlag(false) + .DynamicShapeSupportFlag(true) + .NeedCheckSupportFlag(false) + .PrecisionReduceFlag(false) + .ExtendCfgInfo("coreType.value", "AiCore"); + this->AICore().AddConfig("ascend910b", config); + this->AICore().AddConfig("ascend910_93", config); + } +}; + +OP_ADD(Qwen35GdnDecodeSuperOp); + +} // namespace ops diff --git a/xllm_ops/qwen35_gdn_decode_super_op/op_host/qwen35_gdn_decode_super_op_proto.cpp b/xllm_ops/qwen35_gdn_decode_super_op/op_host/qwen35_gdn_decode_super_op_proto.cpp new file mode 100644 index 0000000..0c203be --- /dev/null +++ b/xllm_ops/qwen35_gdn_decode_super_op/op_host/qwen35_gdn_decode_super_op_proto.cpp @@ -0,0 +1,26 @@ +#include "register/op_impl_registry.h" +#include "tiling_base/error_log.h" + +namespace ops { + +static ge::graphStatus InferShapeQwen35GdnDecodeSuperOp(gert::InferShapeContext* context) +{ + const gert::Shape* qkvShape = context->GetInputShape(0); + const gert::Shape* zShape = context->GetInputShape(1); + const gert::Shape* convStateShape = context->GetInputShape(5); + const gert::Shape* ssmStateShape = context->GetInputShape(8); + OP_CHECK_NULL_WITH_CONTEXT(context, qkvShape); + OP_CHECK_NULL_WITH_CONTEXT(context, zShape); + OP_CHECK_NULL_WITH_CONTEXT(context, convStateShape); + OP_CHECK_NULL_WITH_CONTEXT(context, ssmStateShape); + + *context->GetOutputShape(0) = *qkvShape; + *context->GetOutputShape(1) = *convStateShape; + *context->GetOutputShape(2) = *ssmStateShape; + *context->GetOutputShape(3) = *zShape; + return ge::GRAPH_SUCCESS; +} + +IMPL_OP_INFERSHAPE(Qwen35GdnDecodeSuperOp).InferShape(InferShapeQwen35GdnDecodeSuperOp); + +} // namespace ops diff --git a/xllm_ops/qwen35_gdn_decode_super_op/op_host/qwen35_gdn_decode_super_op_tiling.cpp b/xllm_ops/qwen35_gdn_decode_super_op/op_host/qwen35_gdn_decode_super_op_tiling.cpp new file mode 100644 index 0000000..1b61e17 --- /dev/null +++ b/xllm_ops/qwen35_gdn_decode_super_op/op_host/qwen35_gdn_decode_super_op_tiling.cpp @@ -0,0 +1,132 @@ +#include + +#include "register/op_def_registry.h" +#include "tiling/platform/platform_ascendc.h" + +#include "qwen35_gdn_decode_super_op_tiling.h" + +namespace optiling { + +namespace { + +constexpr int64_t kHeadDim = 128; +constexpr int64_t kConvStateLen = 3; +constexpr int64_t kMaxNumCacheSlots = 1024; +constexpr int64_t kMaxBatchSize = 32; + +bool IsSupportedShape(int64_t numKHeads, int64_t numVHeads) +{ + if (numKHeads < 1 || numKHeads > 16 || + (numKHeads & (numKHeads - 1)) != 0 || + numVHeads % numKHeads != 0) { + return false; + } + const int64_t vHeadsPerK = numVHeads / numKHeads; + return vHeadsPerK >= 1 && vHeadsPerK <= 4; +} + +ge::graphStatus GetShapeInfo(gert::TilingContext* context, + int64_t& batchSize, + int64_t& convTileCount, + int64_t& numVHeads, + int64_t& numKHeads) +{ + const auto* qkvShape = context->GetInputShape(0); + const auto* zShape = context->GetInputShape(1); + const auto* convWeightShape = context->GetInputShape(4); + const auto* convStateShape = context->GetInputShape(5); + const auto* ssmStateShape = context->GetInputShape(8); + if (qkvShape == nullptr || zShape == nullptr || convWeightShape == nullptr || + convStateShape == nullptr || ssmStateShape == nullptr) { + return ge::GRAPH_FAILED; + } + + const gert::Shape& qkv = qkvShape->GetStorageShape(); + const gert::Shape& z = zShape->GetStorageShape(); + const gert::Shape& convWeight = convWeightShape->GetStorageShape(); + const gert::Shape& convState = convStateShape->GetStorageShape(); + const gert::Shape& ssmState = ssmStateShape->GetStorageShape(); + if (qkv.GetDimNum() != 2 || z.GetDimNum() != 3 || + convWeight.GetDimNum() != 2 || convState.GetDimNum() != 3 || + ssmState.GetDimNum() != 4) { + return ge::GRAPH_FAILED; + } + + batchSize = qkv.GetDim(0); + const int64_t convDim = qkv.GetDim(1); + numVHeads = z.GetDim(1); + const int64_t qkWidth = convDim - numVHeads * kHeadDim; + if (batchSize < 1 || batchSize > kMaxBatchSize || convDim <= 0 || + convDim % kHeadDim != 0 || qkWidth <= 0 || + qkWidth % (2 * kHeadDim) != 0 || z.GetDim(0) != batchSize || + z.GetDim(2) != kHeadDim || convWeight.GetDim(0) != 4 || + convWeight.GetDim(1) != convDim || convState.GetDim(0) < 1 || + convState.GetDim(0) > kMaxNumCacheSlots || + convState.GetDim(1) != kConvStateLen || + convState.GetDim(2) != convDim || + ssmState.GetDim(0) != convState.GetDim(0) || + ssmState.GetDim(1) != numVHeads || + ssmState.GetDim(2) != kHeadDim || + ssmState.GetDim(3) != kHeadDim) { + return ge::GRAPH_FAILED; + } + + numKHeads = qkWidth / (2 * kHeadDim); + if (!IsSupportedShape(numKHeads, numVHeads)) { + return ge::GRAPH_FAILED; + } + convTileCount = convDim / kHeadDim; + return ge::GRAPH_SUCCESS; +} + +} // namespace + +static ge::graphStatus Qwen35GdnDecodeSuperOpTiling(gert::TilingContext* context) +{ + int64_t batchSize = 0; + int64_t convTileCount = 0; + int64_t numVHeads = 0; + int64_t numKHeads = 0; + if (GetShapeInfo(context, batchSize, convTileCount, numVHeads, numKHeads) != + ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } + + Qwen35GdnDecodeSuperOpTilingData tiling; + tiling.set_batch_size(batchSize); + tiling.set_num_k_heads(numKHeads); + tiling.set_num_v_heads(numVHeads); + tiling.SaveToBuffer(context->GetRawTilingData()->GetData(), + context->GetRawTilingData()->GetCapacity()); + context->GetRawTilingData()->SetDataSize(tiling.GetDataSize()); + context->SetTilingKey(batchSize == 1 ? 2 : 1); + + auto platform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo()); + const uint32_t aicCoreNum = platform.GetCoreNumAic(); + const uint32_t aivCoreNum = platform.GetCoreNumAiv(); + const uint32_t headCount = + static_cast(batchSize * numVHeads); + const uint32_t convTaskCount = static_cast(convTileCount); + const uint32_t taskCount = batchSize == 1 + ? std::max((convTaskCount + 1) / 2, headCount) + : std::max(convTaskCount, headCount); + const uint32_t usedAivCoreNum = std::min(taskCount, aivCoreNum); + const uint32_t blockDim = + platform.CalcTschBlockDim(usedAivCoreNum, aicCoreNum, aivCoreNum); + if (blockDim == 0) { + return ge::GRAPH_FAILED; + } + + context->SetBlockDim(blockDim); + if (context->SetScheduleMode(1) != ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } + + size_t* workspaceSizes = context->GetWorkspaceSizes(1); + workspaceSizes[0] = platform.GetLibApiWorkSpaceSize(); + return ge::GRAPH_SUCCESS; +} + +IMPL_OP_OPTILING(Qwen35GdnDecodeSuperOp).Tiling(Qwen35GdnDecodeSuperOpTiling); + +} // namespace optiling diff --git a/xllm_ops/qwen35_gdn_decode_super_op/op_host/qwen35_gdn_decode_super_op_tiling.h b/xllm_ops/qwen35_gdn_decode_super_op/op_host/qwen35_gdn_decode_super_op_tiling.h new file mode 100644 index 0000000..b728869 --- /dev/null +++ b/xllm_ops/qwen35_gdn_decode_super_op/op_host/qwen35_gdn_decode_super_op_tiling.h @@ -0,0 +1,16 @@ +/* Copyright 2026 The xLLM Authors. All Rights Reserved. */ + +#pragma once + +#include "register/tilingdata_base.h" + +namespace optiling { +BEGIN_TILING_DATA_DEF(Qwen35GdnDecodeSuperOpTilingData) +TILING_DATA_FIELD_DEF(int64_t, batch_size); +TILING_DATA_FIELD_DEF(int64_t, num_k_heads); +TILING_DATA_FIELD_DEF(int64_t, num_v_heads); +END_TILING_DATA_DEF; + +REGISTER_TILING_DATA_CLASS(Qwen35GdnDecodeSuperOp, + Qwen35GdnDecodeSuperOpTilingData) +} // namespace optiling diff --git a/xllm_ops/qwen35_gdn_decode_super_op/op_kernel/qwen35_gdn_decode_pto_kernel.h b/xllm_ops/qwen35_gdn_decode_super_op/op_kernel/qwen35_gdn_decode_pto_kernel.h new file mode 100644 index 0000000..5c782ca --- /dev/null +++ b/xllm_ops/qwen35_gdn_decode_super_op/op_kernel/qwen35_gdn_decode_pto_kernel.h @@ -0,0 +1,698 @@ +/* Copyright 2026 The xLLM Authors. All Rights Reserved. */ + +#pragma once + +#include "kernel_operator.h" +#include + +#include +#include + +namespace qwen35_decode_pto { + +using namespace pto; + +template +using TileUbDataND = + pto::Tile; + +template +using TileUbDataDN = + pto::Tile; + +template +AICORE PTO_INLINE void CopyGmToUb(__gm__ T1 *handle, int32_t ubAddress, + int32_t ubOffset, int32_t validRows, + int32_t validCols) +{ + static_assert(std::is_same_v); + pto::Shape shape; + shape.shape[3] = validRows; + shape.shape[4] = validCols; + pto::GlobalTensor< + T1, pto::Shape, + pto::Stride> + tensor(handle, shape); + + TileUbDataND + tile(validRows, validCols); + pto::TASSIGN(tile, ubAddress + ubOffset * sizeof(T2)); + pto::TLOAD(tile, tensor); + + if constexpr (PadVal != pto::PadValue::Null) { + if (validRows != static_cast(UbRows) || + validCols != static_cast(UbCols)) { + TileUbDataND padded; + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + pto::TASSIGN(padded, ubAddress + ubOffset * sizeof(T2)); + pto::TFILLPAD_INPLACE(padded, tile); + set_flag(PIPE_V, PIPE_MTE2, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID0); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + pipe_barrier(PIPE_V); + } + } +} + +template +AICORE PTO_INLINE void CopyUbToGm(__gm__ T1 *handle, int32_t ubAddress, + int32_t ubOffset, int32_t validRows, + int32_t validCols) +{ + static_assert(std::is_same_v); + pto::Shape shape; + shape.shape[3] = validRows; + shape.shape[4] = validCols; + pto::GlobalTensor< + T1, pto::Shape, + pto::Stride> + tensor(handle, shape); + + constexpr bool kUseNd = static_cast(UbCols) * sizeof(T2) >= 32; + if constexpr (kUseNd) { + TileUbDataND tile( + validRows, validCols); + pto::TASSIGN(tile, ubAddress + ubOffset * sizeof(T2)); + pto::TSTORE(tensor, tile); + } else { + TileUbDataDN tile( + validRows, validCols); + pto::TASSIGN(tile, ubAddress + ubOffset * sizeof(T2)); + pto::TSTORE(tensor, tile); + } +} + +template +AICORE PTO_INLINE void Sigmoid(TileUbDataND &dst, + TileUbDataND &src) +{ + TMULS(src, src, -1); + pipe_barrier(PIPE_V); + TEXP(src, src); + pipe_barrier(PIPE_V); + TADDS(src, src, 1); + pipe_barrier(PIPE_V); + TRECIP(dst, src); +} + +template +AICORE PTO_INLINE void Silu(TileUbDataND &dst, + TileUbDataND &src, + TileUbDataND &tmp) +{ + TMOV(tmp, src); + pipe_barrier(PIPE_V); + Sigmoid(dst, src); + pipe_barrier(PIPE_V); + TMUL(dst, tmp, dst); +} + +template +AICORE PTO_INLINE void MulAddDst(TileUbDataND &dst, + TileUbDataND &src0, + TileUbDataND &src1, + TileUbDataND &tmp) +{ + TMUL(tmp, src0, src1); + pipe_barrier(PIPE_V); + TADD(dst, dst, tmp); +} + +template +__tf__ PTO_INTERNAL void OuterProduct128Impl( + typename TileDst::TileDType __out__ dst, + typename TileRow::TileDType __in__ row, + typename TileCol::TileDType __in__ col) +{ + __ubuf__ float *dstPtr = + reinterpret_cast<__ubuf__ float *>(__cce_get_tile_ptr(dst)); + __ubuf__ float *rowPtr = + reinterpret_cast<__ubuf__ float *>(__cce_get_tile_ptr(row)); + __ubuf__ uint32_t *colPtr = + reinterpret_cast<__ubuf__ uint32_t *>(__cce_get_tile_ptr(col)); + __ubuf__ uint32_t *broadcast = + reinterpret_cast<__ubuf__ uint32_t *>(TMP_UB_OFFSET); + __ubuf__ float *colValues = reinterpret_cast<__ubuf__ float *>(broadcast); + + vbrcb(broadcast, colPtr, 1, 8, 16); + pipe_barrier(PIPE_V); + vmul(dstPtr, rowPtr, colValues, 128, 1, 1, 0, 16, 0, 1); + vmul(dstPtr + 64, rowPtr + 64, colValues, 128, 1, 1, 0, 16, 0, 1); +} + +template +AICORE PTO_INLINE void OuterProduct128(TileDst &dst, TileRow &row, TileCol &col) +{ + OuterProduct128Impl(dst.data(), row.data(), + col.data()); +} + +template +__tf__ PTO_INTERNAL void ColSum128Impl(typename TileDst::TileDType __out__ dst, + typename TileSrc::TileDType __in__ src) +{ + __ubuf__ float *dstPtr = + reinterpret_cast<__ubuf__ float *>(__cce_get_tile_ptr(dst)); + __ubuf__ float *srcPtr = + reinterpret_cast<__ubuf__ float *>(__cce_get_tile_ptr(src)); + + set_mask_count(); + set_vector_mask(0, 128); +#define QWEN35_COLSUM_STAGE(rows) \ + for (uint32_t i = 0; i < (rows) / 2; ++i) { \ + vadd(srcPtr + i * 128, srcPtr + 2 * i * 128, \ + srcPtr + (2 * i + 1) * 128, 0, 1, 1, 1, 8, 8, 8); \ + } \ + pipe_barrier(PIPE_V) + QWEN35_COLSUM_STAGE(128); + QWEN35_COLSUM_STAGE(64); + QWEN35_COLSUM_STAGE(32); + QWEN35_COLSUM_STAGE(16); + QWEN35_COLSUM_STAGE(8); + QWEN35_COLSUM_STAGE(4); + QWEN35_COLSUM_STAGE(2); +#undef QWEN35_COLSUM_STAGE + set_mask_norm(); + set_vector_mask(-1, -1); + copy_ubuf_to_ubuf(dstPtr, srcPtr, 0, 1, 16, 0, 0); + pipe_barrier(PIPE_V); +} + +template +AICORE PTO_INLINE void ColSum128(TileDst &dst, TileSrc &src) +{ + ColSum128Impl(dst.data(), src.data()); +} + +constexpr uint16_t kSyncAivOnlyFlag = 14; +constexpr uint16_t kSyncModeShift = 4; +constexpr uint16_t kSyncFlagShift = 8; + +AICORE PTO_INLINE uint16_t GetFftsMessage(uint16_t mode, uint16_t flag) +{ + return 0x1 + ((mode & 0x3) << kSyncModeShift) + + ((flag & 0xf) << kSyncFlagShift); +} + +AICORE PTO_INLINE void SyncAllAiv() +{ + pipe_barrier(PIPE_ALL); + ffts_cross_core_sync(PIPE_MTE3, GetFftsMessage(0, kSyncAivOnlyFlag)); + wait_flag_dev(kSyncAivOnlyFlag); +} + +// The generated PTO kernel body follows. The tile shape stays fixed while +// model-dependent tensor strides and loop bounds come from host tiling data. + +template +AICORE PTO_INLINE void Run( + __gm__ bfloat16_t *qkv_handle, __gm__ bfloat16_t *z_handle, + __gm__ bfloat16_t *b_handle, __gm__ bfloat16_t *a_handle, + __gm__ bfloat16_t *conv_weight_handle, + __gm__ bfloat16_t *conv_state_handle, __gm__ float *a_log_handle, + __gm__ float *dt_bias_handle, __gm__ float *ssm_state_handle, + __gm__ int *state_indices_handle, __gm__ bfloat16_t *norm_weight_handle, + __gm__ bfloat16_t *conv_out_handle, + __gm__ bfloat16_t *conv_state_out_handle, + __gm__ float *ssm_state_out_handle, __gm__ bfloat16_t *out_handle, + int32_t num_k_heads, int32_t num_v_heads, int32_t runtime_batch_size) +{ + constexpr int32_t kCompiledNumCacheSlots = 1024; + constexpr int32_t kHeadDim = 128; + constexpr int32_t kConvStateLen = 3; + constexpr int32_t kSsmHeadElements = kHeadDim * kHeadDim; + constexpr int32_t kMaxNumKHeads = 16; + constexpr int32_t kMaxNumVHeads = 64; + constexpr int32_t kMaxBatchSize = 32; + constexpr int32_t kMaxConvDim = + (2 * kMaxNumKHeads + kMaxNumVHeads) * kHeadDim; + constexpr int32_t kMaxConvWeightElements = 4 * kMaxConvDim; + constexpr int32_t kMaxConvStateElements = + kCompiledNumCacheSlots * kConvStateLen * kMaxConvDim; + constexpr int32_t kMaxSsmStateElements = + kCompiledNumCacheSlots * kMaxNumVHeads * kSsmHeadElements; + const int32_t batch_size = IsBatchOne ? 1 : runtime_batch_size; + const int32_t conv_dim = + (2 * num_k_heads + num_v_heads) * kHeadDim; + const int32_t conv_tile_count = conv_dim / kHeadDim; + const int32_t conv_state_stride = kConvStateLen * conv_dim; + const int32_t v_heads_per_k = num_v_heads / num_k_heads; + const int32_t v_width = num_v_heads * kHeadDim; + const int32_t ssm_state_stride = num_v_heads * kSsmHeadElements; + auto cid = get_block_idx(); + + qwen35_decode_pto::TileUbDataND w_half0; + TASSIGN(w_half0, 0); + qwen35_decode_pto::TileUbDataND w_half1; + TASSIGN(w_half1, 256); + qwen35_decode_pto::TileUbDataND w_half2; + TASSIGN(w_half2, 512); + qwen35_decode_pto::TileUbDataND w_half3; + TASSIGN(w_half3, 768); + qwen35_decode_pto::TileUbDataND w0; + TASSIGN(w0, 1024); + qwen35_decode_pto::TileUbDataND w1; + TASSIGN(w1, 1536); + qwen35_decode_pto::TileUbDataND w2; + TASSIGN(w2, 2048); + qwen35_decode_pto::TileUbDataND w3; + TASSIGN(w3, 2560); + qwen35_decode_pto::TileUbDataND hist_half0; + TASSIGN(hist_half0, 3072); + qwen35_decode_pto::TileUbDataND hist_half1; + TASSIGN(hist_half1, 3328); + qwen35_decode_pto::TileUbDataND hist_half2; + TASSIGN(hist_half2, 3584); + qwen35_decode_pto::TileUbDataND x_half; + TASSIGN(x_half, 3840); + qwen35_decode_pto::TileUbDataND hist0; + TASSIGN(hist0, 4096); + qwen35_decode_pto::TileUbDataND hist1; + TASSIGN(hist1, 4608); + qwen35_decode_pto::TileUbDataND hist2; + TASSIGN(hist2, 5120); + qwen35_decode_pto::TileUbDataND x_fp32; + TASSIGN(x_fp32, 5632); + qwen35_decode_pto::TileUbDataND conv_acc; + TASSIGN(conv_acc, 6144); + qwen35_decode_pto::TileUbDataND conv_tmp; + TASSIGN(conv_tmp, 6656); + qwen35_decode_pto::TileUbDataND conv_y; + TASSIGN(conv_y, 7168); + qwen35_decode_pto::TileUbDataND y_half; + TASSIGN(y_half, 7680); + qwen35_decode_pto::TileUbDataND save_half0; + TASSIGN(save_half0, 7936); + qwen35_decode_pto::TileUbDataND save_half1; + TASSIGN(save_half1, 8192); + qwen35_decode_pto::TileUbDataND save_half2; + TASSIGN(save_half2, 8448); + qwen35_decode_pto::TileUbDataND q_half; + TASSIGN(q_half, 8704); + qwen35_decode_pto::TileUbDataND k_half; + TASSIGN(k_half, 8960); + qwen35_decode_pto::TileUbDataND a_half; + TASSIGN(a_half, 9216); + qwen35_decode_pto::TileUbDataND b_half; + TASSIGN(b_half, 9248); + qwen35_decode_pto::TileUbDataND q_fp32; + TASSIGN(q_fp32, 9280); + qwen35_decode_pto::TileUbDataND k_fp32; + TASSIGN(k_fp32, 9792); + qwen35_decode_pto::TileUbDataND scalar; + TASSIGN(scalar, 10304); + qwen35_decode_pto::TileUbDataND scalar2; + TASSIGN(scalar2, 10336); + qwen35_decode_pto::TileUbDataND norm_sq; + TASSIGN(norm_sq, 10368); + qwen35_decode_pto::TileUbDataND norm_val; + TASSIGN(norm_val, 10880); + qwen35_decode_pto::TileUbDataND tmp_ub; + TASSIGN(tmp_ub, 10912); + qwen35_decode_pto::TileUbDataND scalar_tmp; + TASSIGN(scalar_tmp, 19104); + qwen35_decode_pto::TileUbDataND exp_a_buf; + TASSIGN(exp_a_buf, 19136); + qwen35_decode_pto::TileUbDataND scalar_work; + TASSIGN(scalar_work, 19168); + qwen35_decode_pto::TileUbDataND norm_half; + TASSIGN(norm_half, 152320); + qwen35_decode_pto::TileUbDataND z_half; + TASSIGN(z_half, 152576); + qwen35_decode_pto::TileUbDataND weight_half; + TASSIGN(weight_half, 152832); + qwen35_decode_pto::TileUbDataND norm_fp32; + TASSIGN(norm_fp32, 153088); + qwen35_decode_pto::TileUbDataND z_fp32; + TASSIGN(z_fp32, 153600); + qwen35_decode_pto::TileUbDataND weight_fp32; + TASSIGN(weight_fp32, 154112); + qwen35_decode_pto::TileUbDataND square_fp32; + TASSIGN(square_fp32, 154624); + qwen35_decode_pto::TileUbDataND rms; + TASSIGN(rms, 155136); + qwen35_decode_pto::TileUbDataND gate_fp32; + TASSIGN(gate_fp32, 155168); + qwen35_decode_pto::TileUbDataND final_half; + TASSIGN(final_half, 155680); + qwen35_decode_pto::TileUbDataND v_half; + TASSIGN(v_half, 19200); + qwen35_decode_pto::TileUbDataND h_vec; + TASSIGN(h_vec, 19456); + qwen35_decode_pto::TileUbDataND v_fp32; + TASSIGN(v_fp32, 84992); + qwen35_decode_pto::TileUbDataND compute_buf; + TASSIGN(compute_buf, 85504); + qwen35_decode_pto::TileUbDataND pred; + TASSIGN(pred, 151040); + qwen35_decode_pto::TileUbDataND delta; + TASSIGN(delta, 151552); + qwen35_decode_pto::TileUbDataND out_half; + TASSIGN(out_half, 152064); + auto vid = get_subblockid(); +#if defined(__DAV_C220_VEC__) + set_mask_norm(); + set_vector_mask(-1, -1); + const int32_t vector_core_idx = cid * get_subblockdim() + vid; + const int32_t vector_core_count = get_block_num() * get_subblockdim(); + const int32_t batch_one_state_idx = + IsBatchOne ? *state_indices_handle : 0; + + for (int32_t conv_tile = vector_core_idx; conv_tile < conv_tile_count; + conv_tile += vector_core_count) { + const int32_t channel_offset = conv_tile * kHeadDim; + qwen35_decode_pto::CopyGmToUb(conv_weight_handle + channel_offset, 0, 0, 1, 128); + qwen35_decode_pto::CopyGmToUb(conv_weight_handle + channel_offset + conv_dim, 256, 0, 1, 128); + qwen35_decode_pto::CopyGmToUb(conv_weight_handle + channel_offset + 2 * conv_dim, 512, 0, 1, 128); + qwen35_decode_pto::CopyGmToUb(conv_weight_handle + channel_offset + 3 * conv_dim, 768, 0, 1, 128); + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID1); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID1); + TCVT(w0, w_half0, RoundMode::CAST_NONE); + TCVT(w1, w_half1, RoundMode::CAST_NONE); + TCVT(w2, w_half2, RoundMode::CAST_NONE); + TCVT(w3, w_half3, RoundMode::CAST_NONE); + + for (int32_t batch_idx = 0; batch_idx < batch_size; ++batch_idx) { + if constexpr (!IsBatchOne) { + pipe_barrier(PIPE_ALL); + pipe_barrier(PIPE_ALL); + } + const int32_t state_idx = IsBatchOne + ? batch_one_state_idx + : *(state_indices_handle + batch_idx); + qwen35_decode_pto::CopyGmToUb(conv_state_handle + state_idx * conv_state_stride + channel_offset, 3072, 0, 1, 128); + qwen35_decode_pto::CopyGmToUb(conv_state_handle + state_idx * conv_state_stride + channel_offset + conv_dim, 3328, 0, 1, 128); + qwen35_decode_pto::CopyGmToUb(conv_state_handle + state_idx * conv_state_stride + channel_offset + 2 * conv_dim, 3584, 0, 1, 128); + qwen35_decode_pto::CopyGmToUb(qkv_handle + batch_idx * conv_dim + channel_offset, 3840, 0, 1, 128); + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID2); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID2); + TCVT(hist0, hist_half0, RoundMode::CAST_NONE); + TCVT(hist1, hist_half1, RoundMode::CAST_NONE); + TCVT(hist2, hist_half2, RoundMode::CAST_NONE); + TCVT(x_fp32, x_half, RoundMode::CAST_NONE); + pipe_barrier(PIPE_V); + TMUL(conv_acc, w0, hist0); + TMUL(conv_tmp, w1, hist1); + pipe_barrier(PIPE_V); + TADD(conv_acc, conv_acc, conv_tmp); + pipe_barrier(PIPE_V); + TMUL(conv_tmp, w2, hist2); + pipe_barrier(PIPE_V); + TADD(conv_acc, conv_acc, conv_tmp); + pipe_barrier(PIPE_V); + qwen35_decode_pto::TileUbDataND conv_acc_temp_0_muladddst_tmp; + TASSIGN(conv_acc_temp_0_muladddst_tmp, 155936); + qwen35_decode_pto::MulAddDst(conv_acc, x_fp32, w3, conv_acc_temp_0_muladddst_tmp); + pipe_barrier(PIPE_V); + qwen35_decode_pto::TileUbDataND conv_y_temp_0_silu_tmp; + TASSIGN(conv_y_temp_0_silu_tmp, 155936); + qwen35_decode_pto::Silu(conv_y, conv_acc, conv_y_temp_0_silu_tmp); + pipe_barrier(PIPE_V); + TCVT(y_half, conv_y, RoundMode::CAST_RINT); + TCVT(save_half0, hist1, RoundMode::CAST_RINT); + TCVT(save_half1, hist2, RoundMode::CAST_RINT); + TCVT(save_half2, x_fp32, RoundMode::CAST_RINT); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID3); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID3); + qwen35_decode_pto::CopyUbToGm(conv_out_handle + batch_idx * conv_dim + channel_offset, 7680, 0, 1, 128); + qwen35_decode_pto::CopyUbToGm(conv_state_out_handle + state_idx * conv_state_stride + channel_offset, 7936, 0, 1, 128); + pipe_barrier(PIPE_MTE3); + qwen35_decode_pto::CopyUbToGm(conv_state_out_handle + state_idx * conv_state_stride + channel_offset + conv_dim, 8192, 0, 1, 128); + pipe_barrier(PIPE_MTE3); + qwen35_decode_pto::CopyUbToGm(conv_state_out_handle + state_idx * conv_state_stride + channel_offset + 2 * conv_dim, 8448, 0, 1, 128); + if constexpr (!IsBatchOne) { + pipe_barrier(PIPE_ALL); + pipe_barrier(PIPE_ALL); + } + } + if constexpr (IsBatchOne) { + if (conv_tile + vector_core_count < conv_tile_count) { + set_flag(PIPE_V, PIPE_MTE2, EVENT_ID3); + wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID3); + set_flag(PIPE_MTE3, PIPE_V, EVENT_ID3); + wait_flag(PIPE_MTE3, PIPE_V, EVENT_ID3); + } + } + } + qwen35_decode_pto::SyncAllAiv(); + + for (int32_t head_index = vector_core_idx; + head_index < batch_size * num_v_heads; + head_index += vector_core_count) { + if constexpr (!IsBatchOne) { + pipe_barrier(PIPE_ALL); + pipe_barrier(PIPE_ALL); + } + const int32_t batch_idx = + IsBatchOne ? 0 : head_index / num_v_heads; + const int32_t head_idx = + IsBatchOne ? head_index : head_index % num_v_heads; + const int32_t qk_head_idx = head_idx / v_heads_per_k; + const int32_t state_idx_1 = IsBatchOne + ? batch_one_state_idx + : *(state_indices_handle + batch_idx); + qwen35_decode_pto::CopyGmToUb(conv_out_handle + batch_idx * conv_dim + qk_head_idx * kHeadDim, 8704, 0, 1, 128); + qwen35_decode_pto::CopyGmToUb(conv_out_handle + batch_idx * conv_dim + num_k_heads * kHeadDim + qk_head_idx * kHeadDim, 8960, 0, 1, 128); + qwen35_decode_pto::CopyGmToUb(a_handle + head_index, 9216, 0, 1, 1); + qwen35_decode_pto::CopyGmToUb(b_handle + head_index, 9248, 0, 1, 1); + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID6); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID6); + TCVT(q_fp32, q_half, RoundMode::CAST_NONE); + TCVT(k_fp32, k_half, RoundMode::CAST_NONE); + qwen35_decode_pto::TileUbDataND a_half_temp_0; + TASSIGN(a_half_temp_0, 9216 + 0 * 2); + qwen35_decode_pto::TileUbDataND scalar_temp_0; + TASSIGN(scalar_temp_0, 10304 + 0 * 4); + TCVT(scalar_temp_0, a_half_temp_0, RoundMode::CAST_NONE); + qwen35_decode_pto::TileUbDataND b_half_temp_0; + TASSIGN(b_half_temp_0, 9248 + 0 * 2); + qwen35_decode_pto::TileUbDataND scalar2_temp_0; + TASSIGN(scalar2_temp_0, 10336 + 0 * 4); + TCVT(scalar2_temp_0, b_half_temp_0, RoundMode::CAST_NONE); + pipe_barrier(PIPE_V); + TMUL(norm_sq, q_fp32, q_fp32); + pipe_barrier(PIPE_V); + qwen35_decode_pto::TileUbDataDN norm_val_temp_0; + TASSIGN(norm_val_temp_0, 10880 + 0 * 4); + qwen35_decode_pto::TileUbDataND tmp_ub_temp_0; + TASSIGN(tmp_ub_temp_0, 10912 + 0 * 4); + TROWSUM(norm_val_temp_0, norm_sq, tmp_ub_temp_0); + pipe_barrier(PIPE_V); + TADDS(norm_val, norm_val, 1.000000e-06f); + pipe_barrier(PIPE_V); + qwen35_decode_pto::TileUbDataND norm_val_temp_1; + TASSIGN(norm_val_temp_1, 10880 + 0 * 4); + qwen35_decode_pto::TileUbDataND scalar_tmp_temp_0; + TASSIGN(scalar_tmp_temp_0, 19104 + 0 * 4); + TSQRT(scalar_tmp_temp_0, norm_val_temp_1); + set_flag(PIPE_V, PIPE_S, EVENT_ID0); + wait_flag(PIPE_V, PIPE_S, EVENT_ID0); + float q_norm_scalar = scalar_tmp.GetValue(0); + TMULS(q_fp32, q_fp32, 1.0f / q_norm_scalar); + pipe_barrier(PIPE_V); + TMULS(q_fp32, q_fp32, 8.838835e-02f); + TMUL(norm_sq, k_fp32, k_fp32); + pipe_barrier(PIPE_V); + qwen35_decode_pto::TileUbDataDN norm_val_temp_2; + TASSIGN(norm_val_temp_2, 10880 + 0 * 4); + qwen35_decode_pto::TileUbDataND tmp_ub_temp_1; + TASSIGN(tmp_ub_temp_1, 10912 + 0 * 4); + TROWSUM(norm_val_temp_2, norm_sq, tmp_ub_temp_1); + pipe_barrier(PIPE_V); + TADDS(norm_val, norm_val, 1.000000e-06f); + pipe_barrier(PIPE_V); + qwen35_decode_pto::TileUbDataND norm_val_temp_3; + TASSIGN(norm_val_temp_3, 10880 + 0 * 4); + qwen35_decode_pto::TileUbDataND scalar_tmp_temp_1; + TASSIGN(scalar_tmp_temp_1, 19104 + 0 * 4); + TSQRT(scalar_tmp_temp_1, norm_val_temp_3); + set_flag(PIPE_V, PIPE_S, EVENT_ID0); + wait_flag(PIPE_V, PIPE_S, EVENT_ID0); + float k_norm_scalar = scalar_tmp.GetValue(0); + TMULS(k_fp32, k_fp32, 1.0f / k_norm_scalar); + set_flag(PIPE_V, PIPE_MTE2, EVENT_ID7); + wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID7); + qwen35_decode_pto::CopyGmToUb(a_log_handle + head_idx, 19104, 0, 1, 1); + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + qwen35_decode_pto::TileUbDataND scalar_tmp_temp_2; + TASSIGN(scalar_tmp_temp_2, 19104 + 0 * 4); + qwen35_decode_pto::TileUbDataND exp_a_buf_temp_0; + TASSIGN(exp_a_buf_temp_0, 19136 + 0 * 4); + TEXP(exp_a_buf_temp_0, scalar_tmp_temp_2); + qwen35_decode_pto::CopyGmToUb(dt_bias_handle + head_idx, 10880, 0, 1, 1); + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID1); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID1); + pipe_barrier(PIPE_V); + qwen35_decode_pto::TileUbDataND scalar_temp_1; + TASSIGN(scalar_temp_1, 10304 + 0 * 4); + qwen35_decode_pto::TileUbDataND norm_val_temp_4; + TASSIGN(norm_val_temp_4, 10880 + 0 * 4); + qwen35_decode_pto::TileUbDataND scalar_tmp_temp_3; + TASSIGN(scalar_tmp_temp_3, 19104 + 0 * 4); + TADD(scalar_tmp_temp_3, scalar_temp_1, norm_val_temp_4); + set_flag(PIPE_V, PIPE_S, EVENT_ID0); + wait_flag(PIPE_V, PIPE_S, EVENT_ID0); + float x_gate = scalar_tmp.GetValue(0); + if (2.000000e+01f < x_gate) { + TADDS(scalar_work, scalar_tmp, 0.000000e+00f); + } else { + qwen35_decode_pto::TileUbDataND scalar_tmp_temp_4; + TASSIGN(scalar_tmp_temp_4, 19104 + 0 * 4); + qwen35_decode_pto::TileUbDataND scalar_work_temp_0; + TASSIGN(scalar_work_temp_0, 19168 + 0 * 4); + TEXP(scalar_work_temp_0, scalar_tmp_temp_4); + pipe_barrier(PIPE_V); + TADDS(norm_val, scalar_work, 1.000000e+00f); + pipe_barrier(PIPE_V); + qwen35_decode_pto::TileUbDataND norm_val_temp_5; + TASSIGN(norm_val_temp_5, 10880 + 0 * 4); + qwen35_decode_pto::TileUbDataND scalar_work_temp_1; + TASSIGN(scalar_work_temp_1, 19168 + 0 * 4); + TLOG(scalar_work_temp_1, norm_val_temp_5); + } + pipe_barrier(PIPE_ALL); + qwen35_decode_pto::TileUbDataND exp_a_buf_temp_1; + TASSIGN(exp_a_buf_temp_1, 19136 + 0 * 4); + qwen35_decode_pto::TileUbDataND scalar_work_temp_2; + TASSIGN(scalar_work_temp_2, 19168 + 0 * 4); + qwen35_decode_pto::TileUbDataND scalar_tmp_temp_5; + TASSIGN(scalar_tmp_temp_5, 19104 + 0 * 4); + TMUL(scalar_tmp_temp_5, exp_a_buf_temp_1, scalar_work_temp_2); + pipe_barrier(PIPE_V); + TMULS(scalar_tmp, scalar_tmp, -1.000000e+00f); + pipe_barrier(PIPE_V); + qwen35_decode_pto::TileUbDataND scalar_tmp_temp_6; + TASSIGN(scalar_tmp_temp_6, 19104 + 0 * 4); + qwen35_decode_pto::TileUbDataND scalar_work_temp_3; + TASSIGN(scalar_work_temp_3, 19168 + 0 * 4); + TEXP(scalar_work_temp_3, scalar_tmp_temp_6); + set_flag(PIPE_V, PIPE_S, EVENT_ID0); + wait_flag(PIPE_V, PIPE_S, EVENT_ID0); + float decay = scalar_work.GetValue(0); + pipe_barrier(PIPE_V); + TMULS(scalar_tmp, scalar2, -1.000000e+00f); + pipe_barrier(PIPE_V); + qwen35_decode_pto::TileUbDataND scalar_tmp_temp_7; + TASSIGN(scalar_tmp_temp_7, 19104 + 0 * 4); + qwen35_decode_pto::TileUbDataND scalar_work_temp_4; + TASSIGN(scalar_work_temp_4, 19168 + 0 * 4); + TEXP(scalar_work_temp_4, scalar_tmp_temp_7); + pipe_barrier(PIPE_V); + TADDS(scalar_tmp, scalar_work, 1.000000e+00f); + pipe_barrier(PIPE_V); + qwen35_decode_pto::TileUbDataND scalar_tmp_temp_8; + TASSIGN(scalar_tmp_temp_8, 19104 + 0 * 4); + qwen35_decode_pto::TileUbDataND scalar_work_temp_5; + TASSIGN(scalar_work_temp_5, 19168 + 0 * 4); + TRECIP(scalar_work_temp_5, scalar_tmp_temp_8); + set_flag(PIPE_V, PIPE_S, EVENT_ID0); + wait_flag(PIPE_V, PIPE_S, EVENT_ID0); + float beta_gate = scalar_work.GetValue(0); + qwen35_decode_pto::CopyGmToUb(conv_out_handle + batch_idx * conv_dim + 2 * num_k_heads * kHeadDim + head_idx * kHeadDim, 19200, 0, 1, 128); + qwen35_decode_pto::CopyGmToUb(ssm_state_handle + state_idx_1 * ssm_state_stride + head_idx * kSsmHeadElements, 19456, 0, 128, 128); + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID2); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID2); + TCVT(v_fp32, v_half, RoundMode::CAST_NONE); + TMULS(h_vec, h_vec, decay); + pipe_barrier(PIPE_V); + { + qwen35_decode_pto::TileUbDataDN k_row; + TASSIGN(k_row, reinterpret_cast(k_fp32.data())); + TROWEXPANDMUL(compute_buf, h_vec, k_row); + } + pipe_barrier(PIPE_V); + qwen35_decode_pto::ColSum128( + pred, compute_buf); + TSUB(delta, v_fp32, pred); + pipe_barrier(PIPE_V); + TMULS(delta, delta, beta_gate); + pipe_barrier(PIPE_V); + { + qwen35_decode_pto::TileUbDataDN k_row; + TASSIGN(k_row, reinterpret_cast(k_fp32.data())); + qwen35_decode_pto::OuterProduct128( + compute_buf, delta, k_row); + } + pipe_barrier(PIPE_V); + TADD(h_vec, h_vec, compute_buf); + pipe_barrier(PIPE_V); + { + qwen35_decode_pto::TileUbDataDN q_row; + TASSIGN(q_row, reinterpret_cast(q_fp32.data())); + TROWEXPANDMUL(compute_buf, h_vec, q_row); + } + pipe_barrier(PIPE_V); + qwen35_decode_pto::ColSum128( + pred, compute_buf); + TCVT(out_half, pred, RoundMode::CAST_RINT); + pipe_barrier(PIPE_V); + TMOV(norm_half, out_half); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID3); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID3); + qwen35_decode_pto::CopyUbToGm(ssm_state_out_handle + state_idx_1 * ssm_state_stride + head_idx * kSsmHeadElements, 19456, 0, 128, 128); + qwen35_decode_pto::CopyGmToUb(z_handle + batch_idx * v_width + head_idx * kHeadDim, 152576, 0, 1, 128); + qwen35_decode_pto::CopyGmToUb(norm_weight_handle + 0, 152832, 0, 1, 128); + pipe_barrier(PIPE_V); + TCVT(norm_fp32, norm_half, RoundMode::CAST_NONE); + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID4); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID4); + TCVT(z_fp32, z_half, RoundMode::CAST_NONE); + TCVT(weight_fp32, weight_half, RoundMode::CAST_NONE); + pipe_barrier(PIPE_V); + TMUL(square_fp32, norm_fp32, norm_fp32); + pipe_barrier(PIPE_V); + qwen35_decode_pto::TileUbDataDN rms_temp_0; + TASSIGN(rms_temp_0, 155136 + 0 * 4); + qwen35_decode_pto::TileUbDataND tmp_ub_temp_4; + TASSIGN(tmp_ub_temp_4, 10912 + 0 * 4); + TROWSUM(rms_temp_0, square_fp32, tmp_ub_temp_4); + pipe_barrier(PIPE_V); + TMULS(rms, rms, 1.0f / 1.280000e+02f); + pipe_barrier(PIPE_V); + TADDS(rms, rms, 1.000000e-06f); + pipe_barrier(PIPE_V); + qwen35_decode_pto::TileUbDataND rms_temp_1; + TASSIGN(rms_temp_1, 155136 + 0 * 4); + qwen35_decode_pto::TileUbDataND scalar_tmp_temp_9; + TASSIGN(scalar_tmp_temp_9, 19104 + 0 * 4); + TSQRT(scalar_tmp_temp_9, rms_temp_1); + pipe_barrier(PIPE_V); + set_flag(PIPE_V, PIPE_S, EVENT_ID0); + wait_flag(PIPE_V, PIPE_S, EVENT_ID0); + auto scalar_tmp_scalar_temp_0 = scalar_tmp.GetValue(0); + TMULS(norm_fp32, norm_fp32, 1.0f / scalar_tmp_scalar_temp_0); + pipe_barrier(PIPE_V); + TMUL(norm_fp32, norm_fp32, weight_fp32); + qwen35_decode_pto::TileUbDataND gate_fp32_temp_0_silu_tmp; + TASSIGN(gate_fp32_temp_0_silu_tmp, 155936); + qwen35_decode_pto::Silu(gate_fp32, z_fp32, gate_fp32_temp_0_silu_tmp); + pipe_barrier(PIPE_V); + TMUL(norm_fp32, norm_fp32, gate_fp32); + pipe_barrier(PIPE_V); + TCVT(final_half, norm_fp32, RoundMode::CAST_RINT); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID5); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID5); + qwen35_decode_pto::CopyUbToGm(out_handle + batch_idx * v_width + head_idx * kHeadDim, 155680, 0, 1, 128); + pipe_barrier(PIPE_ALL); + if constexpr (!IsBatchOne) { + pipe_barrier(PIPE_ALL); + } + } +#endif +} + +} // namespace qwen35_decode_pto diff --git a/xllm_ops/qwen35_gdn_decode_super_op/op_kernel/qwen35_gdn_decode_super_op.cpp b/xllm_ops/qwen35_gdn_decode_super_op/op_kernel/qwen35_gdn_decode_super_op.cpp new file mode 100644 index 0000000..a5e77ca --- /dev/null +++ b/xllm_ops/qwen35_gdn_decode_super_op/op_kernel/qwen35_gdn_decode_super_op.cpp @@ -0,0 +1,69 @@ +/* Copyright 2026 The xLLM Authors. All Rights Reserved. */ + +#include "qwen35_gdn_decode_pto_kernel.h" + +struct Qwen35GdnDecodeSuperOpTilingData { + int64_t batch_size; + int64_t num_k_heads; + int64_t num_v_heads; +}; + +template +AICORE PTO_INLINE void RunQwen35Decode( + GM_ADDR qkv, GM_ADDR z, GM_ADDR b, GM_ADDR a, GM_ADDR convWeight, + GM_ADDR convState, GM_ADDR aLog, GM_ADDR dtBias, GM_ADDR ssmState, + GM_ADDR stateIndices, GM_ADDR normWeight, GM_ADDR convOut, + GM_ADDR convStateOut, GM_ADDR ssmStateOut, GM_ADDR out, + int32_t numKHeads, int32_t numVHeads, int32_t batchSize) +{ + qwen35_decode_pto::Run( + reinterpret_cast<__gm__ bfloat16_t *>(qkv), + reinterpret_cast<__gm__ bfloat16_t *>(z), + reinterpret_cast<__gm__ bfloat16_t *>(b), + reinterpret_cast<__gm__ bfloat16_t *>(a), + reinterpret_cast<__gm__ bfloat16_t *>(convWeight), + reinterpret_cast<__gm__ bfloat16_t *>(convState), + reinterpret_cast<__gm__ float *>(aLog), + reinterpret_cast<__gm__ float *>(dtBias), + reinterpret_cast<__gm__ float *>(ssmState), + reinterpret_cast<__gm__ int *>(stateIndices), + reinterpret_cast<__gm__ bfloat16_t *>(normWeight), + reinterpret_cast<__gm__ bfloat16_t *>(convOut), + reinterpret_cast<__gm__ bfloat16_t *>(convStateOut), + reinterpret_cast<__gm__ float *>(ssmStateOut), + reinterpret_cast<__gm__ bfloat16_t *>(out), + numKHeads, numVHeads, batchSize); +} + +extern "C" __global__ __aicore__ void qwen35_gdn_decode_super_op( + GM_ADDR qkv, GM_ADDR z, GM_ADDR b, GM_ADDR a, GM_ADDR convWeight, + GM_ADDR convState, GM_ADDR aLog, GM_ADDR dtBias, GM_ADDR ssmState, + GM_ADDR stateIndices, GM_ADDR normWeight, GM_ADDR convOut, + GM_ADDR convStateOut, GM_ADDR ssmStateOut, GM_ADDR out, + GM_ADDR workspace, GM_ADDR tiling) +{ + KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_MIX_AIC_1_2); + REGISTER_TILING_DEFAULT(Qwen35GdnDecodeSuperOpTilingData); + GET_TILING_DATA_WITH_STRUCT(Qwen35GdnDecodeSuperOpTilingData, + tilingData, tiling); + (void)workspace; + + if constexpr (TILING_KEY_IS(2)) { + RunQwen35Decode( + qkv, z, b, a, convWeight, convState, aLog, dtBias, ssmState, + stateIndices, normWeight, convOut, convStateOut, ssmStateOut, + out, static_cast(tilingData.num_k_heads), + static_cast(tilingData.num_v_heads), 1); + } else if constexpr (TILING_KEY_IS(1)) { + const int32_t batchSize = static_cast(tilingData.batch_size); + RunQwen35Decode( + qkv, z, b, a, convWeight, convState, aLog, dtBias, ssmState, + stateIndices, normWeight, convOut, convStateOut, ssmStateOut, + out, static_cast(tilingData.num_k_heads), + static_cast(tilingData.num_v_heads), batchSize); + } +} + +// The generated mixed-kernel wrapper calls matmul::clearWorkspace. Keep this +// include after PTO code to avoid the CANN DYNAMIC/pto::DYNAMIC name collision. +#include "lib/matmul_intf.h" diff --git a/xllm_ops/qwen35_gdn_prefill_super_op/CMakeLists.txt b/xllm_ops/qwen35_gdn_prefill_super_op/CMakeLists.txt new file mode 100644 index 0000000..a385854 --- /dev/null +++ b/xllm_ops/qwen35_gdn_prefill_super_op/CMakeLists.txt @@ -0,0 +1,11 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved. + +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() diff --git a/xllm_ops/qwen35_gdn_prefill_super_op/op_host/CMakeLists.txt b/xllm_ops/qwen35_gdn_prefill_super_op/op_host/CMakeLists.txt new file mode 100644 index 0000000..0e49a66 --- /dev/null +++ b/xllm_ops/qwen35_gdn_prefill_super_op/op_host/CMakeLists.txt @@ -0,0 +1,27 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved. + +add_op_to_compiled_list() + +if (BUILD_OPEN_PROJECT) + target_sources(op_host_aclnn PRIVATE + qwen35_gdn_prefill_super_op_def.cpp + ) +endif() + +add_ops_compile_options( + OP_NAME Qwen35GdnPrefillSuperOp + OPTIONS --cce-auto-sync=on + -Wno-deprecated-declarations + -Wno-error + -DGDN_KERNEL_NAME=qwen35_gdn_prefill_super_op + -DGDN_D=128 + -DGDN_C=128 + -I${CMAKE_CURRENT_SOURCE_DIR} + -I${CMAKE_CURRENT_SOURCE_DIR}/../op_kernel + -I${CMAKE_CURRENT_SOURCE_DIR}/../../causal_conv1d/op_kernel + -I${CMAKE_CURRENT_SOURCE_DIR}/../../../third_party/pto-isa/include +) + +if (NOT BUILD_OPS_RTY_KERNEL) + add_modules_sources(OPTYPE qwen35_gdn_prefill_super_op ACLNNTYPE aclnn) +endif() diff --git a/xllm_ops/qwen35_gdn_prefill_super_op/op_host/qwen35_gdn_prefill_super_op_def.cpp b/xllm_ops/qwen35_gdn_prefill_super_op/op_host/qwen35_gdn_prefill_super_op_def.cpp new file mode 100644 index 0000000..267ec57 --- /dev/null +++ b/xllm_ops/qwen35_gdn_prefill_super_op/op_host/qwen35_gdn_prefill_super_op_def.cpp @@ -0,0 +1,75 @@ +/* Copyright 2026 The xLLM Authors. All Rights Reserved. */ + +#include "register/op_def_registry.h" + +namespace ops { +class Qwen35GdnPrefillSuperOp : public OpDef { +public: + explicit Qwen35GdnPrefillSuperOp(const char *name) : OpDef(name) + { + this->Input("mixed_qkv").ParamType(REQUIRED).DataType({ge::DT_BF16}).Format({ge::FORMAT_ND}).AutoContiguous(); + this->Input("z").ParamType(REQUIRED).DataType({ge::DT_BF16}).Format({ge::FORMAT_ND}).AutoContiguous(); + this->Input("b").ParamType(REQUIRED).DataType({ge::DT_BF16}).Format({ge::FORMAT_ND}).AutoContiguous(); + this->Input("a").ParamType(REQUIRED).DataType({ge::DT_BF16}).Format({ge::FORMAT_ND}).AutoContiguous(); + this->Input("conv_weight") + .ParamType(REQUIRED) + .DataType({ge::DT_BF16}) + .Format({ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("conv_state") + .ParamType(REQUIRED) + .DataType({ge::DT_BF16}) + .Format({ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("a_log").ParamType(REQUIRED).DataType({ge::DT_FLOAT}).Format({ge::FORMAT_ND}).AutoContiguous(); + this->Input("dt_bias").ParamType(REQUIRED).DataType({ge::DT_FLOAT}).Format({ge::FORMAT_ND}).AutoContiguous(); + this->Input("ssm_state") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT}) + .Format({ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("norm_weight") + .ParamType(REQUIRED) + .DataType({ge::DT_BF16}) + .Format({ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("mask_lower").ParamType(REQUIRED).DataType({ge::DT_FLOAT}).Format({ge::FORMAT_ND}); + this->Input("mask_full").ParamType(REQUIRED).DataType({ge::DT_FLOAT}).Format({ge::FORMAT_ND}); + this->Input("minus_identity").ParamType(REQUIRED).DataType({ge::DT_FLOAT16}).Format({ge::FORMAT_ND}); + this->Input("cu_seqlens") + .ParamType(REQUIRED) + .DataType({ge::DT_INT32}) + .Format({ge::FORMAT_ND}) + .AutoContiguous(); + + this->Output("packed_qkv").ParamType(REQUIRED).DataType({ge::DT_FLOAT16}).Format({ge::FORMAT_ND}); + this->Output("g").ParamType(REQUIRED).DataType({ge::DT_FLOAT}).Format({ge::FORMAT_ND}); + this->Output("beta").ParamType(REQUIRED).DataType({ge::DT_FLOAT16}).Format({ge::FORMAT_ND}); + this->Output("initial_state").ParamType(REQUIRED).DataType({ge::DT_FLOAT16}).Format({ge::FORMAT_ND}); + this->Output("mega_out").ParamType(REQUIRED).DataType({ge::DT_FLOAT16}).Format({ge::FORMAT_ND}); + this->Output("g_sum").ParamType(REQUIRED).DataType({ge::DT_FLOAT}).Format({ge::FORMAT_ND}); + this->Output("g_t").ParamType(REQUIRED).DataType({ge::DT_FLOAT}).Format({ge::FORMAT_ND}); + this->Output("beta_t").ParamType(REQUIRED).DataType({ge::DT_FLOAT16}).Format({ge::FORMAT_ND}); + this->Output("mega_a").ParamType(REQUIRED).DataType({ge::DT_FLOAT16}).Format({ge::FORMAT_ND}); + this->Output("a_inv_f32").ParamType(REQUIRED).DataType({ge::DT_FLOAT}).Format({ge::FORMAT_ND}); + this->Output("a_inv").ParamType(REQUIRED).DataType({ge::DT_FLOAT16}).Format({ge::FORMAT_ND}); + this->Output("w").ParamType(REQUIRED).DataType({ge::DT_FLOAT16}).Format({ge::FORMAT_ND}); + this->Output("u").ParamType(REQUIRED).DataType({ge::DT_FLOAT16}).Format({ge::FORMAT_ND}); + this->Output("h").ParamType(REQUIRED).DataType({ge::DT_FLOAT16}).Format({ge::FORMAT_ND}); + this->Output("v_new").ParamType(REQUIRED).DataType({ge::DT_FLOAT16}).Format({ge::FORMAT_ND}); + this->Output("final_state").ParamType(REQUIRED).DataType({ge::DT_FLOAT16}).Format({ge::FORMAT_ND}); + this->Output("conv_state_out").ParamType(REQUIRED).DataType({ge::DT_BF16}).Format({ge::FORMAT_ND}); + this->Output("ssm_state_out").ParamType(REQUIRED).DataType({ge::DT_FLOAT}).Format({ge::FORMAT_ND}); + this->Output("out").ParamType(REQUIRED).DataType({ge::DT_BF16}).Format({ge::FORMAT_ND}); + + this->Attr("num_matrices").AttrType(REQUIRED).Int(0); + this->Attr("conv_state_index").AttrType(REQUIRED).Int(0); + this->Attr("ssm_state_index").AttrType(REQUIRED).Int(0); + + this->AICore().AddConfig("ascend910b"); + this->AICore().AddConfig("ascend910_93"); + } +}; + +OP_ADD(Qwen35GdnPrefillSuperOp); +} // namespace ops diff --git a/xllm_ops/qwen35_gdn_prefill_super_op/op_host/qwen35_gdn_prefill_super_op_infershape.cpp b/xllm_ops/qwen35_gdn_prefill_super_op/op_host/qwen35_gdn_prefill_super_op_infershape.cpp new file mode 100644 index 0000000..5bd5e12 --- /dev/null +++ b/xllm_ops/qwen35_gdn_prefill_super_op/op_host/qwen35_gdn_prefill_super_op_infershape.cpp @@ -0,0 +1,98 @@ +/* Copyright 2026 The xLLM Authors. All Rights Reserved. */ + +#include + +#include "register/op_def_registry.h" + +namespace { +constexpr int64_t kHeadDim = 128; +constexpr int64_t kValueHeads = 24; +constexpr int64_t kChunkSize = 128; + +enum InputIndex { + MIXED_QKV_INDEX = 0, + Z_INDEX, + B_INDEX, + A_INDEX, + CONV_WEIGHT_INDEX, + CONV_STATE_INDEX, + A_LOG_INDEX, + DT_BIAS_INDEX, + SSM_STATE_INDEX, + NORM_WEIGHT_INDEX, + MASK_LOWER_INDEX, + MASK_FULL_INDEX, + MINUS_IDENTITY_INDEX, + CU_SEQLENS_INDEX, +}; + +enum OutputIndex { + PACKED_QKV_INDEX = 0, + G_INDEX, + BETA_INDEX, + INITIAL_STATE_INDEX, + MEGA_OUT_INDEX, + G_SUM_INDEX, + G_T_INDEX, + BETA_T_INDEX, + MEGA_A_INDEX, + A_INV_F32_INDEX, + A_INV_INDEX, + W_INDEX, + U_INDEX, + H_INDEX, + V_NEW_INDEX, + FINAL_STATE_INDEX, + CONV_STATE_OUT_INDEX, + SSM_STATE_OUT_INDEX, + OUT_INDEX, +}; + +void SetShape(gert::Shape *shape, std::initializer_list dims) +{ + shape->SetDimNum(dims.size()); + size_t index = 0; + for (auto dim : dims) { + shape->SetDim(index++, dim); + } +} +} // namespace + +namespace ge { +static ge::graphStatus InferShape(gert::InferShapeContext *context) +{ + const gert::Shape *mixedShape = context->GetInputShape(MIXED_QKV_INDEX); + const gert::Shape *convStateShape = context->GetInputShape(CONV_STATE_INDEX); + const gert::Shape *ssmStateShape = context->GetInputShape(SSM_STATE_INDEX); + if (mixedShape == nullptr || convStateShape == nullptr || ssmStateShape == nullptr || + mixedShape->GetDimNum() != 2) { + return GRAPH_FAILED; + } + + const int64_t totalTokens = mixedShape->GetDim(0); + const int64_t numMatrices = ((totalTokens + kChunkSize - 1) / kChunkSize) * kValueHeads; + + SetShape(context->GetOutputShape(PACKED_QKV_INDEX), {totalTokens, 5120}); + SetShape(context->GetOutputShape(G_INDEX), {1, totalTokens, kValueHeads}); + SetShape(context->GetOutputShape(BETA_INDEX), {1, totalTokens, kValueHeads}); + SetShape(context->GetOutputShape(INITIAL_STATE_INDEX), {1, kValueHeads, kHeadDim, kHeadDim}); + SetShape(context->GetOutputShape(MEGA_OUT_INDEX), {1, totalTokens, kValueHeads, kHeadDim}); + SetShape(context->GetOutputShape(G_SUM_INDEX), {1, totalTokens, kValueHeads}); + SetShape(context->GetOutputShape(G_T_INDEX), {kValueHeads, totalTokens}); + SetShape(context->GetOutputShape(BETA_T_INDEX), {kValueHeads, totalTokens}); + SetShape(context->GetOutputShape(MEGA_A_INDEX), {1, totalTokens, kValueHeads, kChunkSize}); + SetShape(context->GetOutputShape(A_INV_F32_INDEX), {1, totalTokens, kValueHeads, kChunkSize}); + SetShape(context->GetOutputShape(A_INV_INDEX), {1, totalTokens, kValueHeads, kChunkSize}); + SetShape(context->GetOutputShape(W_INDEX), {1, totalTokens, kValueHeads, kHeadDim}); + SetShape(context->GetOutputShape(U_INDEX), {1, totalTokens, kValueHeads, kHeadDim}); + SetShape(context->GetOutputShape(H_INDEX), {numMatrices, kHeadDim, kHeadDim}); + SetShape(context->GetOutputShape(V_NEW_INDEX), {1, totalTokens, kValueHeads, kHeadDim}); + SetShape(context->GetOutputShape(FINAL_STATE_INDEX), {kValueHeads, kHeadDim, kHeadDim}); + *context->GetOutputShape(CONV_STATE_OUT_INDEX) = *convStateShape; + *context->GetOutputShape(SSM_STATE_OUT_INDEX) = *ssmStateShape; + SetShape(context->GetOutputShape(OUT_INDEX), {totalTokens, kValueHeads, kHeadDim}); + return GRAPH_SUCCESS; +} + +IMPL_OP_INFERSHAPE(Qwen35GdnPrefillSuperOp).InferShape(InferShape); +} // namespace ge diff --git a/xllm_ops/qwen35_gdn_prefill_super_op/op_host/qwen35_gdn_prefill_super_op_tiling.cpp b/xllm_ops/qwen35_gdn_prefill_super_op/op_host/qwen35_gdn_prefill_super_op_tiling.cpp new file mode 100644 index 0000000..af925b8 --- /dev/null +++ b/xllm_ops/qwen35_gdn_prefill_super_op/op_host/qwen35_gdn_prefill_super_op_tiling.cpp @@ -0,0 +1,114 @@ +/* Copyright 2026 The xLLM Authors. All Rights Reserved. */ + +#include "qwen35_gdn_prefill_super_op_tiling.h" + +#include +#include +#include + +#include "register/op_def_registry.h" +#include "tiling/platform/platform_ascendc.h" + +namespace { +constexpr uint32_t kChunkSize = 128; +constexpr uint32_t kValueHeads = 24; +constexpr uint32_t kHalfBytes = 2; +constexpr uint32_t kVectorCoreCount = 40; +constexpr uint32_t kConvDimTiles = 2; + +enum InputIndex { + MIXED_QKV_INDEX = 0, +}; + +enum AttrIndex { + NUM_MATRICES_ATTR = 0, + CONV_STATE_INDEX_ATTR, + SSM_STATE_INDEX_ATTR, +}; + +uint32_t CeilDiv(uint32_t value, uint32_t divisor) +{ + return divisor == 0 ? 0 : (value + divisor - 1) / divisor; +} + +uint64_t CalcUserWorkspaceBytes(uint32_t blockDim) +{ + const uint64_t tileBytes = static_cast(kChunkSize) * kChunkSize * kHalfBytes; + constexpr uint64_t kWorkspaceTileCount = 11; + return static_cast(blockDim) * kWorkspaceTileCount * tileBytes; +} +} // namespace + +namespace optiling { +static ge::graphStatus TilingFunc(gert::TilingContext *context) +{ + auto platform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo()); + uint32_t blockDim = std::max(platform.GetCoreNumAic(), 1); + const gert::StorageShape *mixedShape = context->GetInputShape(MIXED_QKV_INDEX); + if (mixedShape == nullptr || mixedShape->GetOriginShape().GetDimNum() != 2 || + mixedShape->GetOriginShape().GetDim(1) != 5120) { + return ge::GRAPH_FAILED; + } + + const int64_t totalTokens64 = mixedShape->GetOriginShape().GetDim(0); + if (totalTokens64 <= 0 || totalTokens64 > std::numeric_limits::max()) { + return ge::GRAPH_FAILED; + } + const uint32_t totalTokens = static_cast(totalTokens64); + const uint32_t tokenCoreBudget = kVectorCoreCount / kConvDimTiles; + const uint32_t tokenBlockSize = CeilDiv(totalTokens, tokenCoreBudget); + const uint32_t tokenBlockCount = CeilDiv(totalTokens, tokenBlockSize); + + auto attrs = context->GetAttrs(); + if (attrs == nullptr || attrs->GetAttrPointer(NUM_MATRICES_ATTR) == nullptr || + attrs->GetAttrPointer(CONV_STATE_INDEX_ATTR) == nullptr || + attrs->GetAttrPointer(SSM_STATE_INDEX_ATTR) == nullptr) { + return ge::GRAPH_FAILED; + } + const int64_t numMatricesAttr = *attrs->GetAttrPointer(NUM_MATRICES_ATTR); + const int64_t convStateIndex = *attrs->GetAttrPointer(CONV_STATE_INDEX_ATTR); + const int64_t ssmStateIndex = *attrs->GetAttrPointer(SSM_STATE_INDEX_ATTR); + if (numMatricesAttr <= 0 || convStateIndex < 0 || ssmStateIndex < 0) { + return ge::GRAPH_FAILED; + } + + Qwen35GdnPrefillSuperOpTilingData tiling; + tiling.set_block_dim(blockDim); + tiling.set_num_matrices(static_cast(numMatricesAttr)); + tiling.set_num_heads(kValueHeads); + tiling.set_num_key_heads(8); + tiling.set_token_block_size(tokenBlockSize); + tiling.set_token_block_count(tokenBlockCount); + tiling.set_conv_state_index(convStateIndex); + tiling.set_ssm_state_index(ssmStateIndex); + tiling.set_batch_size(1); + tiling.set_seq_len(totalTokens); + tiling.set_total_tokens(totalTokens); + + context->SetBlockDim(blockDim); + if (context->SetScheduleMode(1) != ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } + tiling.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity()); + context->GetRawTilingData()->SetDataSize(tiling.GetDataSize()); + + size_t *workspaceSizes = context->GetWorkspaceSizes(1); + if (workspaceSizes == nullptr) { + return ge::GRAPH_FAILED; + } + workspaceSizes[0] = CalcUserWorkspaceBytes(blockDim) + platform.GetLibApiWorkSpaceSize(); + return ge::GRAPH_SUCCESS; +} + +struct Qwen35GdnPrefillSuperOpCompileInfo {}; + +static ge::graphStatus TilingParse(gert::TilingParseContext *context) +{ + (void)context; + return ge::GRAPH_SUCCESS; +} + +IMPL_OP_OPTILING(Qwen35GdnPrefillSuperOp) + .Tiling(TilingFunc) + .TilingParse(TilingParse); +} // namespace optiling diff --git a/xllm_ops/qwen35_gdn_prefill_super_op/op_host/qwen35_gdn_prefill_super_op_tiling.h b/xllm_ops/qwen35_gdn_prefill_super_op/op_host/qwen35_gdn_prefill_super_op_tiling.h new file mode 100644 index 0000000..075bf19 --- /dev/null +++ b/xllm_ops/qwen35_gdn_prefill_super_op/op_host/qwen35_gdn_prefill_super_op_tiling.h @@ -0,0 +1,23 @@ +/* Copyright 2026 The xLLM Authors. All Rights Reserved. */ + +#pragma once + +#include "register/tilingdata_base.h" + +namespace optiling { +BEGIN_TILING_DATA_DEF(Qwen35GdnPrefillSuperOpTilingData) +TILING_DATA_FIELD_DEF(uint32_t, block_dim); +TILING_DATA_FIELD_DEF(uint32_t, num_matrices); +TILING_DATA_FIELD_DEF(uint32_t, num_heads); +TILING_DATA_FIELD_DEF(uint32_t, num_key_heads); +TILING_DATA_FIELD_DEF(uint32_t, token_block_size); +TILING_DATA_FIELD_DEF(uint32_t, token_block_count); +TILING_DATA_FIELD_DEF(int64_t, conv_state_index); +TILING_DATA_FIELD_DEF(int64_t, ssm_state_index); +TILING_DATA_FIELD_DEF(int64_t, batch_size); +TILING_DATA_FIELD_DEF(int64_t, seq_len); +TILING_DATA_FIELD_DEF(int64_t, total_tokens); +END_TILING_DATA_DEF; + +REGISTER_TILING_DATA_CLASS(Qwen35GdnPrefillSuperOp, Qwen35GdnPrefillSuperOpTilingData) +} // namespace optiling diff --git a/xllm_ops/qwen35_gdn_prefill_super_op/op_kernel/chunk_cumsum.cpp b/xllm_ops/qwen35_gdn_prefill_super_op/op_kernel/chunk_cumsum.cpp new file mode 100644 index 0000000..cae3b94 --- /dev/null +++ b/xllm_ops/qwen35_gdn_prefill_super_op/op_kernel/chunk_cumsum.cpp @@ -0,0 +1,371 @@ +// ============================================================================ +// chunk_cumsum_kernel.cpp — Prefix sum of gate values G along time dimension +// +// Mathematical operation (per chunk of C tokens, independently per head h): +// g_sum[t, h] = Σ_{i=0}^{t} g[i, h] for t = 0 .. valid-1 +// +// Input: g [total_tokens, H] float, BSND layout — raw gate values +// Output: g_sum [total_tokens, H] float — cumulative sums +// +// The prefix sum enables downstream kernels to compute exponential decay +// coefficients: exp(g_sum[i] - g_sum[j]) gives the cumulative gate +// from token j to token i within a chunk. +// +// Architecture: Vec-only kernel (no Cube/GEMM). Single Vec sub-block. +// Pipeline: MTE2(load) → Vec(compute) → MTE3(store), serialized per chunk. +// +// NPU memory hierarchy used: +// GM (Global Memory) → UB (Unified Buffer, on-chip SRAM, Vec-accessible) +// +// ─── PTO / NPU Primer for This Kernel ────────────────────────────────────── +// +// AI Core: The basic processing unit of an NPU, analogous to a Streaming +// Multiprocessor (SM) on a GPU. A single chip has many AI cores, and each +// core runs the same kernel code on different data (SPMD model). +// +// Memory hierarchy (outer → inner): +// GM (Global Memory) — Off-chip DRAM, like GPU HBM. Large (several GB) +// but high latency. All AI cores share GM. +// UB (Unified Buffer) — On-chip SRAM, ~256 KB per AI core. Like GPU +// shared memory. Very fast, but small. The Vec engine can only operate +// on data that lives in UB, so every tensor must be DMA'd in first. +// +// Hardware pipes (execute in parallel, like independent GPU warps): +// Vec — SIMD vector processor. Performs element-wise math (add, mul, etc.) +// on data already in UB. Think of it as a wide SIMD ALU. +// MTE2 — DMA engine for loads: copies data from GM → UB. +// MTE3 — DMA engine for stores: copies data from UB → GM. +// Cube — Matrix engine for GEMMs (not used in this kernel). +// +// Synchronization (set_flag / wait_flag): +// Because Vec, MTE2, and MTE3 run in parallel on separate hardware, you +// must explicitly synchronize them to ensure data is ready: +// set_flag(SRC_PIPE, DST_PIPE, event): SRC signals that it is done. +// wait_flag(SRC_PIPE, DST_PIPE, event): DST blocks until the signal. +// Example: After MTE2 loads data into UB, Vec must wait_flag before reading +// it. This is like a fine-grained torch.cuda.synchronize() between pipes. +// Events (EVENT_ID0 .. EVENT_ID7) are semaphore indices. +// +// ============================================================================ + +#include +#include "acl/acl.h" +using namespace pto; + +// NumHeads and ChunkSize are template parameters so the compiler can optimize +// tile sizes, unroll loops, and compute UB addresses at compile time. +#ifndef GDN_C +#define GDN_C 128 +#endif + +// ── PTO type aliases (device-only, guarded by __CCE_AICORE__) ─────────────── +// UB tile in row-major (ND) layout, used by Vec engine. +// T=dtype, R×C=static shape, RV×CV=valid region, P=pad value for TLOAD. +// +// Think of UbND as: torch.empty((R, C), dtype=T) allocated in on-chip SRAM (UB). +// - TileType::Vec = this tile lives in UB, operated on by the Vec (SIMD) engine +// - BLayout::RowMajor = row-major storage, like C arrays or numpy default +// - RV, CV = "valid" region within the R×C buffer (for handling partial/tail chunks) +// - PadValue = what to fill outside the valid region during TLOAD (Zero or Null) +// - 512 = alignment in bytes (hardware requirement for efficient DMA) +#ifdef __CCE_AICORE__ +template +using UbND = pto::Tile; +#endif + +template +AICORE void cumsum_kernel(__gm__ float *g_ptr, __gm__ float *g_sum_ptr, __gm__ int32_t *cu_seqlens, int64_t batch_size, + int64_t seq_len) +{ + // get_block_idx(): Returns this AI core's index (0..block_num-1). + // Like blockIdx.x in CUDA — identifies which core this code runs on. + // get_block_num(): Total number of AI cores launched (like gridDim.x in CUDA). + // get_subblockid(): Returns 0 or 1 — selects which Vec sub-block within the core. + // Each AI core has 2 Vec sub-blocks that can run in parallel. + auto cid = get_block_idx(); + auto block_num = get_block_num(); + auto vid = get_subblockid(); + +// #if defined(__DAV_C220_VEC__): This block only compiles for the Vec core pass. +// The bisheng compiler makes 3 passes over the same source file: +// Pass 1: __DAV_C220_VEC__ defined → compiles Vec (SIMD) code +// Pass 2: __DAV_C220_CUBE__ defined → compiles Cube (matrix) code +// Pass 3: neither defined → compiles host (CPU) launcher code +// Using these guards lets us put Vec, Cube, and host code in one file. +#if defined(__DAV_C220_VEC__) + if (vid != 0) return; + + // set_mask_norm(): Reset Vec mask to normal mode (all lanes active). + // set_vector_mask(-1, -1): Enable all SIMD lanes (128 lanes for fp32). + // The -1 sets all 64 bits to 1 in each of the two 64-bit mask registers. + // This is like setting torch's computation to operate on all elements. + set_mask_norm(); + set_vector_mask(-1, -1); + + // HeadTileCols: NumHeads rounded up to 8-element alignment (32B for float) + // HTC = NumHeads rounded up to nearest multiple of 8. + // Why? The Vec engine processes data in 32-byte granularity. + // For float (4 bytes), that's 8 elements per SIMD "word". + // Rounding up ensures every row is a whole number of SIMD words, + // avoiding partial-lane issues. The extra columns are zero-padded. + // Example: NumHeads=16 → HTC=16 (already aligned), NumHeads=13 → HTC=16. + constexpr int32_t HTC = ((NumHeads + 7) / 8) * 8; + constexpr int32_t BlockBytes = ChunkSize * HTC * static_cast(sizeof(float)); + constexpr int32_t RowBytes = HTC * static_cast(sizeof(float)); + + // ── UB memory layout ────────────────────────────────────────────────── + // [0 .. BlockBytes) = g input (ChunkSize × HTC floats) + // [BlockBytes .. 2*BlockBytes) = g_sum output + // [2*BlockBytes .. 2*BlockBytes+RowBytes) = row accumulator (1 × HTC) + constexpr int32_t GUbAddr = 0; + constexpr int32_t SUbAddr = BlockBytes; + constexpr int32_t AccUbAddr = BlockBytes * 2; + + // GlobalTensor types for g/g_sum in [total_tokens, NumHeads] layout. + // 5D shape with last two dims dynamic; stride encodes row pitch. + // + // GlobalTensor is a "view" into GM (Global Memory), like torch.as_strided(). + // GlobalTensor(base_ptr, shape) + // Shape<1,1,1,DYNAMIC,DYNAMIC> = 5D shape where first 3 dims are 1 (unused), + // last 2 dims are set at runtime (valid rows × NumHeads). + // Stride<1,1,1,NumHeads,1> = stride between elements. The 4th stride = NumHeads + // means consecutive rows in GM are NumHeads elements apart (BSND layout: + // token[t] at offset t*NumHeads, head[h] at offset h within that token). + // This is equivalent to: + // g_gm = torch.as_strided(g_ptr, size=[valid, NumHeads], stride=[NumHeads, 1]) + using GmShape = Shape<1, 1, 1, DYNAMIC, DYNAMIC>; + using GmStride = Stride<1, 1, 1, NumHeads, 1>; + using GmFloat = GlobalTensor; + + // Pre-assign row accumulator at fixed UB address + // TASSIGN(tile, address): Binds a tile descriptor to a fixed byte address in UB. + // Think of it as: tile = ub_memory[address:address+sizeof(tile)] + // This does NOT allocate or move data — it just tells the hardware where the tile lives. + // We manually manage UB memory layout (like a memory pool) via compile-time addresses. + UbND acc_ub; + TASSIGN(acc_ub, AccUbAddr); + + int64_t num_seqs = batch_size; + + // ── Fixed-length sequence path (cu_seqlens == nullptr) ──────────────── + if (cu_seqlens == nullptr) { + int64_t chunks_per_seq = (seq_len + ChunkSize - 1) / ChunkSize; + int64_t total_chunks = num_seqs * chunks_per_seq; + + // Work distribution: Each AI core processes chunks in a round-robin pattern. + // Core `cid` handles chunks cid, cid+block_num, cid+2*block_num, ... + // This is the NPU equivalent of CUDA's grid-stride loop: + // for (int i = blockIdx.x; i < total; i += gridDim.x) + for (int64_t gi = static_cast(cid); gi < total_chunks; gi += static_cast(block_num)) { + int64_t seq_idx = gi / chunks_per_seq; + int64_t local_chunk = gi % chunks_per_seq; + int64_t bos = seq_idx * seq_len; + int64_t chunk_start = bos + local_chunk * ChunkSize; + int64_t remaining = seq_len - local_chunk * ChunkSize; + int32_t valid = static_cast(remaining < ChunkSize ? remaining : ChunkSize); + + // ── DMA: load g[chunk_start .. +valid] from GM → UB (MTE2 pipe) ── + // Constructs a GlobalTensor view over the g array, loads into UB, + // then zero-pads the tail region (rows beyond `valid`, cols beyond + // NumHeads up to the 8-aligned HTC) so downstream Vec ops see zeros. + { + GmShape gs; + gs.shape[3] = valid; + gs.shape[4] = NumHeads; + GmFloat g_gm(g_ptr + chunk_start * NumHeads, gs); + UbND g_load(valid, NumHeads); + TASSIGN(g_load, GUbAddr); + // TLOAD(ub_tile, gm_tensor): DMA transfer from GM → UB. + // Equivalent to: ub_tile[:valid, :NumHeads] = gm_tensor[:valid, :NumHeads] + // This is an ASYNC operation on the MTE2 pipe — the CPU/Vec engine can do + // other work while DMA is in progress. You must call set_flag/wait_flag + // before reading the loaded data. + TLOAD(g_load, g_gm); + if (valid != ChunkSize || NumHeads != HTC) { + UbND g_pad; + TASSIGN(g_pad, GUbAddr); + // TFILLPAD_INPLACE(full_tile, partial_tile): Zero-fills the region outside + // the valid area of partial_tile. + // Equivalent to: + // full_tile[valid:ChunkSize, :] = 0 # zero rows beyond valid + // full_tile[:, NumHeads:HTC] = 0 # zero cols beyond NumHeads (alignment padding) + // This ensures downstream Vec operations see clean zeros in padded regions. + TFILLPAD_INPLACE(g_pad, g_load); + } + } + // ── Synchronization: MTE2 → Vec ──────────────────────────────────── + // set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0): Signal from MTE2 (DMA load + // engine) to Vec (SIMD engine) that the DMA transfer is complete. + // wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0): Vec waits here until MTE2 + // has set the flag. After this, UB data from TLOAD is safe to read. + // Think of it as: torch.cuda.synchronize() but fine-grained per pipe. + // EVENT_ID0 is a semaphore index (0-7 available). + // MTE2 → Vec sync: wait for DMA load to finish before Vec reads UB + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + + // ── Vec compute: prefix sum over rows (all H heads in parallel) ─── + // Row 0: acc[h] = g[0,h]; g_sum[0,h] = acc[h] + UbND g_row_0; + TASSIGN(g_row_0, GUbAddr); + // TMOV(dst, src): Element-wise copy, like dst = src.clone() in UB. + TMOV(acc_ub, g_row_0); + // pipe_barrier(PIPE_V): Ensures all pending Vec (SIMD) operations complete + // before the next Vec instruction begins. Needed because Vec ops are pipelined + // and may not finish in order. Think of it as a local __syncthreads() for the + // Vec engine only. Much lighter than set_flag/wait_flag (which sync across + // different hardware units). + pipe_barrier(PIPE_V); + + UbND s_row_0; + TASSIGN(s_row_0, SUbAddr); + TMOV(s_row_0, acc_ub); + pipe_barrier(PIPE_V); + + // Rows 1..valid-1: acc[h] += g[i,h]; g_sum[i,h] = acc[h] + for (int32_t i = 1; i < valid; ++i) { + UbND g_row_i; + TASSIGN(g_row_i, GUbAddr + i * RowBytes); + // TADD(dst, a, b): Element-wise add, like dst = a + b. All in UB. + // Operates on all HTC elements in parallel (SIMD). + TADD(acc_ub, acc_ub, g_row_i); + pipe_barrier(PIPE_V); + + UbND s_row_i; + TASSIGN(s_row_i, SUbAddr + i * RowBytes); + TMOV(s_row_i, acc_ub); + pipe_barrier(PIPE_V); + } + + // Zero-fill rows beyond valid (tail padding for downstream kernels) + // TEXPANDS(tile, scalar): Fill entire tile with a scalar value. + // Equivalent to: tile[:] = scalar (like torch.full_like(tile, scalar)) + TEXPANDS(acc_ub, 0.0f); + pipe_barrier(PIPE_V); + for (int32_t i = valid; i < ChunkSize; ++i) { + UbND s_row_i; + TASSIGN(s_row_i, SUbAddr + i * RowBytes); + TMOV(s_row_i, acc_ub); + pipe_barrier(PIPE_V); + } + + // ── DMA: store g_sum from UB → GM (MTE3 pipe) ──────────────────── + // ── Synchronization: Vec → MTE3 ─────────────────────────────────── + // Vec signals MTE3 that computation is done and UB data is ready to store. + // MTE3 (DMA store engine) waits for this before reading UB for TSTORE. + // Without this sync, MTE3 might read stale/partial data from UB. + // Vec → MTE3 sync: ensure Vec writes to UB are visible before DMA + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + + { + GmShape ss; + ss.shape[3] = valid; + ss.shape[4] = NumHeads; + GmFloat gs_gm(g_sum_ptr + chunk_start * NumHeads, ss); + UbND s_store(valid, NumHeads); + TASSIGN(s_store, SUbAddr); + // TSTORE(gm_tensor, ub_tile): DMA transfer from UB → GM. + // Equivalent to: gm_tensor[:valid, :NumHeads] = ub_tile[:valid, :NumHeads] + // Async on MTE3 pipe. Must sync (Vec→MTE3) before calling, and sync + // (MTE3→Vec) after if reusing the same UB region. + TSTORE(gs_gm, s_store); + } + // ── Synchronization: MTE3 → Vec ─────────────────────────────────── + // MTE3 signals Vec that the DMA store is complete and UB can be reused. + // Vec waits before starting the next iteration's TLOAD into the same UB region. + // Without this, the next TLOAD could overwrite data still being stored. + // MTE3 → Vec sync: wait for DMA store before reusing UB next iter + set_flag(PIPE_MTE3, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE3, PIPE_V, EVENT_ID0); + } + } + // ── Variable-length sequence path (cu_seqlens != nullptr) ───────────── + else { + int64_t gi = 0; + for (int64_t si = 0; si < num_seqs; ++si) { + int64_t bos = static_cast(cu_seqlens[si]); + int64_t eos = static_cast(cu_seqlens[si + 1]); + int64_t slen = eos - bos; + int64_t nc = (slen + ChunkSize - 1) / ChunkSize; + + for (int64_t c = 0; c < nc; ++c) { + if (gi % static_cast(block_num) == static_cast(cid)) { + int64_t chunk_start = bos + c * ChunkSize; + int64_t remaining = slen - c * ChunkSize; + int32_t valid = static_cast(remaining < ChunkSize ? remaining : ChunkSize); + + // Load g chunk from GM → UB, zero-padded + { + GmShape gs; + gs.shape[3] = valid; + gs.shape[4] = NumHeads; + GmFloat g_gm(g_ptr + chunk_start * NumHeads, gs); + UbND g_load(valid, NumHeads); + TASSIGN(g_load, GUbAddr); + TLOAD(g_load, g_gm); + if (valid != ChunkSize || NumHeads != HTC) { + UbND g_pad; + TASSIGN(g_pad, GUbAddr); + TFILLPAD_INPLACE(g_pad, g_load); + } + } + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + + // Prefix sum: acc = g[0]; g_sum[0] = acc + UbND g_row_0; + TASSIGN(g_row_0, GUbAddr); + TMOV(acc_ub, g_row_0); + pipe_barrier(PIPE_V); + + UbND s_row_0; + TASSIGN(s_row_0, SUbAddr); + TMOV(s_row_0, acc_ub); + pipe_barrier(PIPE_V); + + // acc += g[i]; g_sum[i] = acc + for (int32_t i = 1; i < valid; ++i) { + UbND g_row_i; + TASSIGN(g_row_i, GUbAddr + i * RowBytes); + TADD(acc_ub, acc_ub, g_row_i); + pipe_barrier(PIPE_V); + + UbND s_row_i; + TASSIGN(s_row_i, SUbAddr + i * RowBytes); + TMOV(s_row_i, acc_ub); + pipe_barrier(PIPE_V); + } + + // Zero-fill padding rows + TEXPANDS(acc_ub, 0.0f); + pipe_barrier(PIPE_V); + for (int32_t i = valid; i < ChunkSize; ++i) { + UbND s_row_i; + TASSIGN(s_row_i, SUbAddr + i * RowBytes); + TMOV(s_row_i, acc_ub); + pipe_barrier(PIPE_V); + } + + // Store g_sum to GM + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + + { + GmShape ss; + ss.shape[3] = valid; + ss.shape[4] = NumHeads; + GmFloat gs_gm(g_sum_ptr + chunk_start * NumHeads, ss); + UbND s_store(valid, NumHeads); + TASSIGN(s_store, SUbAddr); + TSTORE(gs_gm, s_store); + } + set_flag(PIPE_MTE3, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE3, PIPE_V, EVENT_ID0); + } + gi++; + } + } + } +#endif +} diff --git a/xllm_ops/qwen35_gdn_prefill_super_op/op_kernel/chunk_h.cpp b/xllm_ops/qwen35_gdn_prefill_super_op/op_kernel/chunk_h.cpp new file mode 100644 index 0000000..2f15acc --- /dev/null +++ b/xllm_ops/qwen35_gdn_prefill_super_op/op_kernel/chunk_h.cpp @@ -0,0 +1,853 @@ +// ============================================================================ +// chunk_h_kernel.cpp — Recurrent hidden state update for GatedDeltaNet +// +// Mathematical recurrence per chunk c: +// S_{c+1} = exp(g_last) * S_c + K^T @ V +// +// where g_last = exp(g[valid-1]) is the chunk's final gate value, S is the +// D×D hidden state, K ∈ ℝ^{C×D}, V ∈ ℝ^{C×D}, and g ∈ ℝ^C is the per-token +// gate. +// +// ── Cube phase (two GEMMs per chunk, sequentially): ────────────────────── +// 1. WS = W @ S project current state through W (wy_fast output) +// W ∈ ℝ^{C×D}, S ∈ ℝ^{D×D} → WS ∈ ℝ^{C×D} +// 2. KV = K^T @ V outer product of keys and values (transpose_A!) +// K stored as D×C, V ∈ ℝ^{C×D} → KV ∈ ℝ^{D×D} +// +// ── Vec phase (two sub-blocks handle upper/lower C/2 rows): ───────────── +// For each chunk: +// 1. Load K, G (pre-transposed), U (from wy_fast) +// 2. Compute coeff[i] = exp(g[i] - g[valid-1]) — time-decay scaling +// Uses TROWEXPAND to broadcast coefficients across D columns +// 3. Scale K: K_scaled[i,:] = K[i,:] * coeff[i] +// 4. Load WS from Cube workspace, compute V_new = U - WS (residual) +// 5. Store V_new and K_scaled to workspace for Cube's next iteration +// 6. Update state: S = exp(g_last) * S + KV (from Cube workspace) +// 7. Store final state FS after last chunk +// +// Cross-core sync: Cube→Vec flags for WS/KV ready, Vec→Cube flags for +// K/S ready. +// +// Inputs: +// K [total_tokens, Hg, D] half — keys (BSND layout; GQA/MQA group heads) +// W [total_tokens, H, D] half — wy_fast output (BSND layout) +// U [total_tokens, H, D] half — values pre-residual (BSND layout) +// G [H, total_tokens] float — pre-transposed cumulative gates +// S [total_chunks, H, D, D] half — per-chunk state snapshots (output) +// V [total_tokens, H, D] half — residual-corrected values (output) +// FS [batch, H, D, D] half — final state per sequence (output) +// H0 [batch, H, D, D] half — optional initial state per sequence +// workspace [per-core scratch] — Cube↔Vec communication buffer +// +// NPU memory hierarchy: +// GM → L1 (Cube-accessible) → L0A/L0B/L0C (Cube GEMM registers) +// GM → UB (Vec-accessible, on-chip SRAM) +// Cross-core sync via FFTS (Fast Fine-grained Task Synchronization) +// +// ── PTO / NPU Primer ────────────────────────────────────────────────── +// This is the most complex kernel in the GDN suite. It implements the +// recurrent state update, requiring sequential chunk processing (chunks +// within a sequence CANNOT be parallelized — each depends on the previous). +// +// Key PTO APIs (numpy/torch equivalents): +// TLOAD(dst, gm) — dst = gm_data (DMA: GM→L1 or GM→UB) +// TSTORE(gm, src) — gm_data = src (DMA: UB/L0C→GM) +// TASSIGN(tile, addr) — tile = memory[addr] (bind tile to buffer address) +// TCVT(dst, src, mode) — dst = src.float()/.half() +// TMOV(dst, src) — dst = src.clone() +// TADD(d, a, b) — d = a + b +// TSUB(d, a, b) — d = a - b +// TMUL(d, a, b) — d = a * b +// TMULS(d, s, scalar) — d = s * scalar (scalar multiply) +// TADDS(d, s, scalar) — d = s + scalar (scalar add) +// TEXP(d, s) — d = torch.exp(s) +// TEXPANDS(tile, scalar) — tile[:] = scalar (fill with constant) +// TROWEXPAND(2d, col) — 2d[i,j] = col[i] (broadcast col across row dim) +// TFILLPAD(dst, src) — zero-fill L1 tile padding (for tail chunks) +// TEXTRACT(l0, l1, r, c) — L1 sub-tile → L0A/L0B +// TRESHAPE(zn, nz) — reinterpret layout NZ↔ZN (logical transpose, free) +// TMATMUL(C, A, B) — C = A @ B (Cube GEMM, fp16 inputs → fp32 accum) +// set_flag/wait_flag — pipe sync within same core +// ffts_cross_core_sync — cross-core signal Cube↔Vec +// wait_flag_dev(flag) — wait for cross-core signal +// GetValue(idx) — read a single scalar from a UB tile (slow, use sparingly) +// +// ── Workspace memory layout (shared between Cube and Vec via GM) ────── +// Each AI core has its own workspace region to avoid contention: +// WS_WS [C×D]: Cube writes WS = W @ S here → Vec reads it +// WS_K [D×C]: Vec writes K_scaled here → Cube reads it for KV = K^T @ V +// WS_S [D×D]: Vec writes current state S here → Cube reads it for GEMM 1 +// WS_KV [D×D]: Cube writes KV = K^T @ V here → Vec reads it to update S +// +// Data flow per chunk (think of it as a ping-pong between Cube and Vec): +// Vec: write S₀ to WS_S → signal Cube (flag 3) +// Cube: read S from WS_S, load W → compute WS = W@S → write WS_WS → signal Vec (flag 0) +// Vec: read WS, compute V_new = U - WS, compute K_scaled → write WS_K → signal Cube (flag 1) +// Cube: read K from WS_K, load V → compute KV = K^T@V → write WS_KV → signal Vec (flag 2) +// Vec: read KV, update S = exp(g_last)*S + KV → write S to WS_S → signal Cube (flag 3) +// ... repeat for next chunk ... +// ============================================================================ + +#include +#include +#include "acl/acl.h" +using namespace pto; + +#ifdef __CCE_AICORE__ + +namespace { + +using GmShape2D = pto::Shape<1, 1, 1, pto::DYNAMIC, pto::DYNAMIC>; +using GmStride2D = pto::Stride<1, 1, 1, pto::DYNAMIC, 1>; + +template +using GmTensor2D = pto::GlobalTensor; + +template +using DynMatL1 = pto::Tile; + +template +using DynVecTile = pto::Tile; + +template +using DynAccTile = pto::TileAcc; + +template +using TileMatL1 = pto::Tile; + +template +using TileMatL1ZN = pto::Tile; + +template +using TileMatL0A = pto::Tile; + +template +using TileMatL0B = pto::Tile; + +template +using TileUbDataND = pto::Tile; + +template +using TileUbDataDN = pto::Tile; + +// PTO cheat sheet for the recurrent kernel: +// - `GlobalTensor` is a GM tensor view with explicit runtime shape/stride. +// - `Tile<..., Mat, ...>` lives in L1 and feeds Cube matmul instructions. +// - `Tile<..., Vec, ...>` lives in UB for elementwise vector work. +// - `TileAcc` is a Cube accumulator tile. +// - `TLOAD` / `TSTORE` are DMA copies between GM and on-chip memory. +// - `TROWEXPAND` broadcasts a column vector across the feature dimension. +// - `TFILLPAD(_INPLACE)` zero-pads tail rows so full-tile code can still run. + +template +AICORE PTO_INLINE void +gemm_v0(std::conditional_t, TileMatL1> &A, + std::conditional_t, TileMatL1> &B, + pto::TileAcc &C, bool clear) +{ + // Local K-sliced matmul helper: + // C = A @ B + // PTO exposes the L1/L0 staging explicitly, so this stays as a tiny file- + // local helper instead of a shared wrapper. + // + // PyTorch mental model: + // C = 0 + // for k0 in range(0, K, kL0Size): + // C += A[:, k0:k1] @ B[k0:k1, :] + constexpr uint32_t kL0Size = 128; + const uint32_t kL0split = (K + kL0Size - 1) / kL0Size; + + auto war_event_id = (event_t)(((int)EVENT_ID0 + 1) % 8); + set_flag(PIPE_MTE2, PIPE_MTE1, war_event_id); + wait_flag(PIPE_MTE2, PIPE_MTE1, war_event_id); + + for (uint32_t kL0Idx = 0; kL0Idx < kL0split; ++kL0Idx) { + const bool initflag = clear && (kL0Idx == 0); + const bool is_tail_block = (kL0Idx == kL0split - 1); + + if (is_tail_block) { + TileMatL0A l0a; + TileMatL0B l0b; + pto::TASSIGN(l0a, 0x0); + pto::TASSIGN(l0b, 0x0); + + set_flag(PIPE_M, PIPE_MTE1, war_event_id); + wait_flag(PIPE_M, PIPE_MTE1, war_event_id); + + if constexpr (!transpose_A) { + pto::TEXTRACT(l0a, A, 0, kL0Idx * K_tail); + } else { + TileMatL1ZN A_t; + pto::TRESHAPE(A_t, A); + pto::TEXTRACT(l0a, A_t, 0, kL0Idx * K_tail); + } + + if constexpr (!transpose_B) { + pto::TEXTRACT(l0b, B, kL0Idx * K_tail, 0); + } else { + TileMatL1ZN B_t; + pto::TRESHAPE(B_t, B); + pto::TEXTRACT(l0b, B_t, kL0Idx * K_tail, 0); + } + + set_flag(PIPE_MTE1, PIPE_M, war_event_id); + wait_flag(PIPE_MTE1, PIPE_M, war_event_id); + + if (initflag) { + pto::TMATMUL(C, l0a, l0b); + } else { + pto::TMATMUL_ACC(C, C, l0a, l0b); + } + } else { + TileMatL0A l0a; + TileMatL0B l0b; + pto::TASSIGN(l0a, 0x0); + pto::TASSIGN(l0b, 0x0); + + set_flag(PIPE_M, PIPE_MTE1, war_event_id); + wait_flag(PIPE_M, PIPE_MTE1, war_event_id); + + set_flag(PIPE_FIX, PIPE_M, war_event_id); + wait_flag(PIPE_FIX, PIPE_M, war_event_id); + + if constexpr (!transpose_A) { + pto::TEXTRACT(l0a, A, 0, kL0Idx * kL0Size); + } else { + TileMatL1ZN A_t; + pto::TRESHAPE(A_t, A); + pto::TEXTRACT(l0a, A_t, 0, kL0Idx * kL0Size); + } + + if constexpr (!transpose_B) { + pto::TEXTRACT(l0b, B, kL0Idx * kL0Size, 0); + } else { + TileMatL1ZN B_t; + pto::TRESHAPE(B_t, B); + pto::TEXTRACT(l0b, B_t, kL0Idx * kL0Size, 0); + } + + set_flag(PIPE_MTE1, PIPE_M, war_event_id); + wait_flag(PIPE_MTE1, PIPE_M, war_event_id); + + if (initflag) { + pto::TMATMUL(C, l0a, l0b); + } else { + pto::TMATMUL_ACC(C, C, l0a, l0b); + } + + set_flag(PIPE_MTE1, PIPE_MTE2, war_event_id); + wait_flag(PIPE_MTE1, PIPE_MTE2, war_event_id); + } + } + + set_flag(PIPE_MTE1, PIPE_MTE2, war_event_id); + wait_flag(PIPE_MTE1, PIPE_MTE2, war_event_id); + + set_flag(PIPE_M, PIPE_FIX, war_event_id); + wait_flag(PIPE_M, PIPE_FIX, war_event_id); +} + +} // namespace + +#endif + +template +AICORE void chunk_h_kernel(__gm__ half *K_handle, __gm__ half *W_handle, __gm__ half *U_handle, __gm__ float *G_handle, + __gm__ half *S_handle, __gm__ half *V_handle, __gm__ half *FS_handle, __gm__ half *H0_handle, + int64_t has_initial_state, __gm__ half *workspace_handle, __gm__ int32_t *cu_seqlens, + int64_t batch_size, int64_t seq_len, int64_t total_tokens, uint32_t num_key_heads, + __gm__ float *FS_cache_handle, int64_t fs_cache_index) +{ + // chunk_h advances the recurrent hidden state chunk by chunk: + // ws_i = W_i @ S_i + // v_i_new = U_i - ws_i + // k_i_tilde = exp(g_last - g_i) * K_i + // S_{i+1} = exp(g_last) * S_i + k_i_tilde^T @ v_i_new. + // + // Shapes for one (sequence, head, chunk): + // W_i, U_i, K_i, V_i_new : [valid, D] + // S_i, S_{i+1} : [D, D] + // + // PyTorch / NumPy sketch: + // ws = W_i @ S_i + // v_new = U_i - ws + // decay = exp(g_last - g_i)[:, None] + // k_tilde = decay * K_i + // kv = k_tilde.T @ v_new + // S = exp(g_last) * S + kv + // + // PTO split: + // Cube forms the two matmuls (`W_i @ S_i` and `K_i^T @ V_i_new`). + // Vec does the elementwise gating/decay and carries the running state. + auto cid = get_block_idx(); + auto block_num = get_block_num(); + + constexpr int32_t D = HiddenSize; + constexpr int32_t C = ChunkSize; + constexpr int32_t H = NumHeads; + const int32_t Hg = static_cast(num_key_heads); + if (Hg <= 0 || (H % Hg) != 0) return; + const int32_t GROUP = H / Hg; + constexpr int32_t HalfC = C / 2; + constexpr int32_t BSND_QKV_STRIDE = H * D; + const int32_t BSND_K_STRIDE = Hg * D; + constexpr int32_t DD = D * D; + + constexpr int32_t WS_WS = 0; + constexpr int32_t WS_K = DD; + constexpr int32_t WS_S = DD * 2; + constexpr int32_t WS_KV = DD * 3; + constexpr int32_t WS_PER_CORE = DD * 4; + + TileMatL1 s_l1; + TASSIGN(s_l1, 0); + TileMatL1 w_l1; + TASSIGN(w_l1, D * D * sizeof(half)); + TileAcc ws_l0; + TASSIGN(ws_l0, 0); + TileMatL1 k_l1; + TASSIGN(k_l1, (DD + C * D) * sizeof(half)); + TileMatL1 v_l1; + TASSIGN(v_l1, (DD + C * D + D * C) * sizeof(half)); + TileAcc kv_l0; + TASSIGN(kv_l0, C * D * sizeof(float)); + + constexpr int32_t G_BLOCK_UB = 0; + // Leading UB scratch: legacy kernels used ``C * NumHeads * sizeof(float)``, which overflows UB when + // ``NumHeads`` is 32/48/64. Keep the same slack as the historical ``GDN_H=16`` build (8192 bytes). + constexpr int32_t ZERO_UB = ChunkSize * 16 * static_cast(sizeof(float)); + constexpr int32_t S_UB = ZERO_UB + 64 * sizeof(float); + constexpr int32_t K_UB_HALF = S_UB + HalfC * D * sizeof(float); + constexpr int32_t G_UB = K_UB_HALF + HalfC * D * sizeof(half); + constexpr int32_t U_UB_HALF = G_UB + C * sizeof(float); + constexpr int32_t K_UB = U_UB_HALF + HalfC * D * sizeof(half); + constexpr int32_t G_V_UB = K_UB + HalfC * D * sizeof(float); + constexpr int32_t COEFF_UB = G_V_UB + 64 * sizeof(float); + constexpr int32_t U_UB = COEFF_UB + 64 * sizeof(float); + constexpr int32_t WS_UB = U_UB + HalfC * D * sizeof(float); + constexpr int32_t KV_UB = U_UB_HALF; + constexpr int32_t S_UB_HALF = WS_UB + HalfC * D * sizeof(float); + + TileUbDataND zero_ub; + TASSIGN(zero_ub, ZERO_UB); + TileUbDataND s_ub; + TASSIGN(s_ub, S_UB); + TileUbDataND k_ub_half; + TASSIGN(k_ub_half, K_UB_HALF); + TileUbDataND g_ub; + TASSIGN(g_ub, G_UB); + TileUbDataND s_ub_half; + TASSIGN(s_ub_half, S_UB_HALF); + TileUbDataND u_ub_half; + TASSIGN(u_ub_half, U_UB_HALF); + TileUbDataND k_ub; + TASSIGN(k_ub, K_UB); + TileUbDataND g_v_ub; + TASSIGN(g_v_ub, G_V_UB); + TileUbDataND coeff_ub; + TASSIGN(coeff_ub, COEFF_UB); + TileUbDataND u_ub; + TASSIGN(u_ub, U_UB); + TileUbDataND ws_ub; + TASSIGN(ws_ub, WS_UB); + TileUbDataND kv_ub; + TASSIGN(kv_ub, KV_UB); + + auto vid = get_subblockid(); + + int64_t num_seqs = batch_size; + int64_t total_work = num_seqs * H; + +#if defined(__DAV_C220_CUBE__) + for (int64_t wi = 0; wi < (total_work + block_num - 1) / block_num; ++wi) { + int64_t pid = wi * block_num + cid; + if (pid >= total_work) break; + + int64_t head = pid % H; + int64_t seq_idx = pid / H; + + int64_t bos, slen; + int64_t chunk_offset = 0; + if (cu_seqlens != nullptr) { + bos = static_cast(cu_seqlens[seq_idx]); + int64_t eos = static_cast(cu_seqlens[seq_idx + 1]); + slen = eos - bos; + for (int64_t si = 0; si < seq_idx; ++si) { + int64_t sb = static_cast(cu_seqlens[si]); + int64_t se = static_cast(cu_seqlens[si + 1]); + chunk_offset += (se - sb + C - 1) / C; + } + } else { + bos = seq_idx * seq_len; + slen = seq_len; + chunk_offset = seq_idx * ((seq_len + C - 1) / C); + } + int64_t num_chunks = (slen + C - 1) / C; + int64_t ws_base = static_cast(cid) * WS_PER_CORE; + // One per-core scratch region stores: + // WS_WS : ws = W_i @ S_i + // WS_K : k_tilde + // WS_S : running state S_i + // WS_KV : k_tilde^T @ v_i_new + + for (int32_t ci = 0; ci < num_chunks; ++ci) { + wait_flag_dev(3); + + int64_t chunk_start = bos + static_cast(ci) * C; + int64_t valid = slen - static_cast(ci) * C; + if (valid > C) valid = C; + + { + GmShape2D s_shape(D, D); + GmStride2D s_stride(D); + GmTensor2D s_global(workspace_handle + ws_base + WS_S, s_shape, s_stride); + DynMatL1 s_l1_load(D, D); + TASSIGN(s_l1_load, 0); + // Load the previous recurrent state S_i from per-core workspace. + TLOAD(s_l1_load, s_global); + } + + int64_t w_offset = ((chunk_start)*H + head) * D; + { + GmShape2D w_shape(static_cast(valid), D); + GmStride2D w_stride(BSND_QKV_STRIDE); + GmTensor2D w_global(W_handle + w_offset, w_shape, w_stride); + DynMatL1 w_l1_load(static_cast(valid), D); + TASSIGN(w_l1_load, D * D * static_cast(sizeof(half))); + TLOAD(w_l1_load, w_global); + if (valid != C) { + TFILLPAD(w_l1_load, w_l1_load); + } + } + + set_flag(PIPE_FIX, PIPE_M, EVENT_ID0); + wait_flag(PIPE_FIX, PIPE_M, EVENT_ID0); + // Apply the carried recurrent state to every token in this chunk. + gemm_v0(w_l1, s_l1, ws_l0, (bool)1); + + { + GmShape2D ws_shape(C, D); + GmStride2D ws_stride(D); + GmTensor2D ws_global(workspace_handle + ws_base + WS_WS, ws_shape, ws_stride); + DynAccTile ws_store(C, D); + TASSIGN(ws_store, 0); + // Save ws_i so the Vec phase can do `v_new = U_i - ws_i`. + TSTORE(ws_global, ws_store); + } + ffts_cross_core_sync(PIPE_FIX, 1 | (2 << 4) | (0 << 8)); + + wait_flag_dev(1); + + { + GmShape2D k_shape(D, C); + GmStride2D k_stride(C); + GmTensor2D k_global(workspace_handle + ws_base + WS_K, k_shape, k_stride); + DynMatL1 k_l1_load(D, C); + TASSIGN(k_l1_load, (DD + C * D) * static_cast(sizeof(half))); + TLOAD(k_l1_load, k_global); + } + + int64_t v_offset = ((chunk_start)*H + head) * D; + { + GmShape2D v_shape(static_cast(valid), D); + GmStride2D v_stride(BSND_QKV_STRIDE); + GmTensor2D v_global(V_handle + v_offset, v_shape, v_stride); + DynMatL1 v_l1_load(static_cast(valid), D); + TASSIGN(v_l1_load, (DD + C * D + D * C) * static_cast(sizeof(half))); + TLOAD(v_l1_load, v_global); + if (valid != C) { + TFILLPAD(v_l1_load, v_l1_load); + } + } + + set_flag(PIPE_FIX, PIPE_M, EVENT_ID0); + wait_flag(PIPE_FIX, PIPE_M, EVENT_ID0); + // This chunk contributes the additive update K_i^T V_i to the state recurrence. + gemm_v0(k_l1, v_l1, kv_l0, (bool)1); + + { + GmShape2D kv_shape(D, D); + GmStride2D kv_stride(D); + GmTensor2D kv_global(workspace_handle + ws_base + WS_KV, kv_shape, kv_stride); + DynAccTile kv_store(D, D); + TASSIGN(kv_store, C * D * static_cast(sizeof(float))); + // Save kv = k_tilde^T @ v_i_new so Vec can finish the state update. + TSTORE(kv_global, kv_store); + } + ffts_cross_core_sync(PIPE_FIX, 1 | (2 << 4) | (2 << 8)); + } + } +#endif +#if defined(__DAV_C220_VEC__) + set_mask_norm(); + set_vector_mask(-1, -1); + + // Vec owns the running recurrent state S_i and updates it after every chunk. + for (int64_t wi = 0; wi < (total_work + block_num - 1) / block_num; ++wi) { + int64_t pid = wi * block_num + cid; + if (pid >= total_work) break; + + int64_t head = pid % H; + int64_t head_g = head / GROUP; + int64_t seq_idx = pid / H; + + int64_t bos, slen; + int64_t chunk_offset = 0; + if (cu_seqlens != nullptr) { + bos = static_cast(cu_seqlens[seq_idx]); + int64_t eos = static_cast(cu_seqlens[seq_idx + 1]); + slen = eos - bos; + for (int64_t si = 0; si < seq_idx; ++si) { + int64_t sb = static_cast(cu_seqlens[si]); + int64_t se = static_cast(cu_seqlens[si + 1]); + chunk_offset += (se - sb + C - 1) / C; + } + } else { + bos = seq_idx * seq_len; + slen = seq_len; + chunk_offset = seq_idx * ((seq_len + C - 1) / C); + } + int64_t num_chunks = (slen + C - 1) / C; + int64_t ws_base = static_cast(cid) * WS_PER_CORE; + + set_flag(PIPE_V, PIPE_S, EVENT_ID0); + wait_flag(PIPE_V, PIPE_S, EVENT_ID0); + TEXPANDS(zero_ub, 0.0f); + set_flag(PIPE_V, PIPE_S, EVENT_ID0); + wait_flag(PIPE_V, PIPE_S, EVENT_ID0); + if (has_initial_state != 0) { + int64_t h0_offset = (seq_idx * H + head) * DD + vid * HalfC * D; + GmShape2D h0_shape(HalfC, D); + GmStride2D h0_stride(D); + GmTensor2D h0_global(H0_handle + h0_offset, h0_shape, h0_stride); + DynVecTile h0_load(HalfC, D); + TASSIGN(h0_load, S_UB_HALF); + TLOAD(h0_load, h0_global); + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + TCVT(s_ub, s_ub_half, pto::RoundMode::CAST_NONE); + } else { + // Start each sequence/head recurrence from S_0 = 0. + TEXPANDS(s_ub, 0.0f); + TCVT(s_ub_half, s_ub, pto::RoundMode::CAST_NONE); + } + // Preserve MegaChunkGdn's FP16 final-state rounding point before the + // persistent cache is materialized as FP32. + pipe_barrier(PIPE_V); + TCVT(s_ub, s_ub_half, pto::RoundMode::CAST_NONE); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + { + // `workspace_handle` is a `half*`, so all offsets here are in half elements. + GmShape2D s_shape(HalfC, D); + GmStride2D s_stride(D); + GmTensor2D s_global(workspace_handle + ws_base + WS_S + vid * HalfC * D, s_shape, s_stride); + DynVecTile s_store(HalfC, D); + TASSIGN(s_store, S_UB_HALF); + TSTORE(s_global, s_store); + } + { + int64_t s_out_offset = (chunk_offset * H + head) * DD; + GmShape2D s_out_shape(HalfC, D); + GmStride2D s_out_stride(D); + GmTensor2D s_out_global(S_handle + s_out_offset + vid * HalfC * D, s_out_shape, s_out_stride); + DynVecTile s_out_store(HalfC, D); + TASSIGN(s_out_store, S_UB_HALF); + TSTORE(s_out_global, s_out_store); + } + ffts_cross_core_sync(PIPE_MTE3, 1 | (2 << 4) | (3 << 8)); + + int64_t chunk_start_0 = bos; + int64_t valid0 = slen; + if (valid0 > C) valid0 = C; + // Vec work is split by row stripe, not by individual token. For the first + // chunk we compute exactly how many live rows belong to this sub-block's + // HalfC stripe so short tails do not overrun the packed BSND input. + int32_t valid_rows_0 = static_cast(valid0 - static_cast(vid) * HalfC); + if (valid_rows_0 < 0) valid_rows_0 = 0; + if (valid_rows_0 > HalfC) valid_rows_0 = HalfC; + + int64_t k_offset_0 = (chunk_start_0 * Hg + head_g) * D + vid * HalfC * BSND_K_STRIDE; + if (valid_rows_0 > 0) { + GmShape2D k_shape(valid_rows_0, D); + GmStride2D k_stride(BSND_K_STRIDE); + GmTensor2D k_global(K_handle + k_offset_0, k_shape, k_stride); + DynVecTile k_load(valid_rows_0, D); + TASSIGN(k_load, K_UB_HALF); + TLOAD(k_load, k_global); + if (valid_rows_0 != HalfC) { + TFILLPAD_INPLACE(k_ub_half, k_load); + } + } else { + // Empty stripe (typically vid=1 on a very short tail chunk): synthesize + // a zero tile so later full-width vector math and workspace stores still + // observe proper padding semantics. + TEXPANDS(k_ub, 0.0f); + TCVT(k_ub_half, k_ub, pto::RoundMode::CAST_NONE); + } + + { + GmShape2D g_shape(1, static_cast(valid0)); + GmStride2D g_stride(1); + GmTensor2D g_global(G_handle + head * total_tokens + chunk_start_0, g_shape, g_stride); + DynVecTile g_load(1, static_cast(valid0)); + TASSIGN(g_load, G_UB); + TLOAD(g_load, g_global); + if (valid0 != C) { + TFILLPAD_INPLACE(g_ub, g_load); + } + } + + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + + for (int32_t ci = 0; ci < static_cast(num_chunks); ++ci) { + int64_t chunk_start = bos + static_cast(ci) * C; + int64_t valid = slen - static_cast(ci) * C; + if (valid > C) valid = C; + int32_t valid_rows = static_cast(valid - static_cast(vid) * HalfC); + if (valid_rows < 0) valid_rows = 0; + if (valid_rows > HalfC) valid_rows = HalfC; + // Each Vec subblock owns one contiguous HalfC-row stripe of the chunk. + // For short tail chunks, `valid_rows` may be smaller or even zero. This + // is the key fix that keeps ragged tails and dense varlen boundary mixes + // from reading or writing beyond the live rows in this stripe. + + int64_t u_offset = (chunk_start * H + head) * D + vid * HalfC * BSND_QKV_STRIDE; + if (valid_rows > 0) { + GmShape2D u_shape(valid_rows, D); + GmStride2D u_stride(BSND_QKV_STRIDE); + GmTensor2D u_global(U_handle + u_offset, u_shape, u_stride); + DynVecTile u_load(valid_rows, D); + TASSIGN(u_load, U_UB_HALF); + TLOAD(u_load, u_global); + if (valid_rows != HalfC) { + TFILLPAD_INPLACE(u_ub_half, u_load); + } + } else { + // No live rows for this stripe in the current chunk; keep the tile + // explicitly zero-padded so the remainder of the recurrence logic can + // run in full-tile form without special-casing every later step. + TEXPANDS(u_ub, 0.0f); + TCVT(u_ub_half, u_ub, pto::RoundMode::CAST_NONE); + } + + TCVT(k_ub, k_ub_half, pto::RoundMode::CAST_NONE); + + TileUbDataND g_ub_temp; + TASSIGN(g_ub_temp, G_UB + vid * 64 * sizeof(float)); + TMOV(g_v_ub, g_ub_temp); + + set_flag(PIPE_V, PIPE_S, EVENT_ID0); + wait_flag(PIPE_V, PIPE_S, EVENT_ID0); + float g_last = g_ub.GetValue(static_cast(valid) - 1); + // Rebase the chunk gate around g_last so the intra-chunk decay stays numerically local. + // Torch-like: + // coeff = exp(g_last - g_rows_owned_by_this_subblock) + TADDS(coeff_ub, g_v_ub, -g_last); + pipe_barrier(PIPE_V); + TSUB(coeff_ub, zero_ub, coeff_ub); + pipe_barrier(PIPE_V); + TEXP(coeff_ub, coeff_ub); + + TEXP(g_ub, g_ub); + + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + TCVT(u_ub, u_ub_half, pto::RoundMode::CAST_NONE); + + TileUbDataDN coeff_col_ub; + TASSIGN(coeff_col_ub, COEFF_UB); + TileUbDataND coeff_2d_ub; + TASSIGN(coeff_2d_ub, WS_UB); + // Broadcast one decay scalar per token row across the D feature columns: + // coeff_2d[row, :] = coeff[row] + TROWEXPAND(coeff_2d_ub, coeff_col_ub); + pipe_barrier(PIPE_V); + // `k_ub` now holds k_tilde = exp(g_last - g_i) * K_i. + TMUL(k_ub, k_ub, coeff_2d_ub); + pipe_barrier(PIPE_V); + + wait_flag_dev(0); + { + GmShape2D ws_shape(HalfC, D); + GmStride2D ws_stride(D); + GmTensor2D ws_global(workspace_handle + ws_base + WS_WS + vid * HalfC * D, ws_shape, ws_stride); + DynVecTile ws_load(HalfC, D); + TASSIGN(ws_load, U_UB_HALF); + TLOAD(ws_load, ws_global); + } + + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + TCVT(ws_ub, u_ub_half, pto::RoundMode::CAST_NONE); + // v_i_new = U_i - W_i @ S_i. + // In PyTorch notation: + // u_ub = u_ub - ws_ub + TSUB(u_ub, u_ub, ws_ub); + TCVT(u_ub_half, u_ub, pto::RoundMode::CAST_NONE); + TCVT(k_ub_half, k_ub, pto::RoundMode::CAST_NONE); + + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + + int64_t v_offset = (chunk_start * H + head) * D + vid * HalfC * BSND_QKV_STRIDE; + if (valid_rows > 0) { + GmShape2D v_shape(valid_rows, D); + GmStride2D v_stride(BSND_QKV_STRIDE); + GmTensor2D v_global(V_handle + v_offset, v_shape, v_stride); + DynVecTile v_store(valid_rows, D); + TASSIGN(v_store, U_UB_HALF); + TSTORE(v_global, v_store); + } + + // Spill both V_i_new and k_i_tilde so the Cube stage can form + // k_i_tilde^T @ V_i_new for this chunk. + { + GmShape2D k_shape(HalfC, D); + GmStride2D k_stride(D); + GmTensor2D k_global(workspace_handle + ws_base + WS_K + vid * HalfC * D, k_shape, k_stride); + DynVecTile k_store(HalfC, D); + TASSIGN(k_store, K_UB_HALF); + TSTORE(k_global, k_store); + } + + ffts_cross_core_sync(PIPE_MTE3, 1 | (2 << 4) | (1 << 8)); + + set_flag(PIPE_MTE3, PIPE_S, EVENT_ID0); + wait_flag(PIPE_MTE3, PIPE_S, EVENT_ID0); + float exp_g_last = g_ub.GetValue(static_cast(valid) - 1); + // Carry the recurrence across chunks: S_{i+1} = exp(g_last) * S_i + K_i^T V_i. + TMULS(s_ub, s_ub, exp_g_last); + + set_flag(PIPE_V, PIPE_MTE2, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID0); + if (ci + 1 < static_cast(num_chunks)) { + int64_t next_start = bos + static_cast(ci + 1) * C; + int64_t next_valid = slen - static_cast(ci + 1) * C; + if (next_valid > C) next_valid = C; + int32_t next_valid_rows = static_cast(next_valid - static_cast(vid) * HalfC); + if (next_valid_rows < 0) next_valid_rows = 0; + if (next_valid_rows > HalfC) next_valid_rows = HalfC; + + int64_t nk_off = (next_start * Hg + head_g) * D + vid * HalfC * BSND_K_STRIDE; + if (next_valid_rows > 0) { + GmShape2D k_shape(next_valid_rows, D); + GmStride2D k_stride(BSND_K_STRIDE); + GmTensor2D k_global(K_handle + nk_off, k_shape, k_stride); + DynVecTile k_load(next_valid_rows, D); + TASSIGN(k_load, K_UB_HALF); + TLOAD(k_load, k_global); + if (next_valid_rows != HalfC) { + TFILLPAD_INPLACE(k_ub_half, k_load); + } + } else { + // Same tail-safe zero materialization for the prefetch path: the next + // chunk may have no rows in this stripe even though the other stripe + // is still active. + TEXPANDS(k_ub, 0.0f); + TCVT(k_ub_half, k_ub, pto::RoundMode::CAST_NONE); + } + + { + GmShape2D g_shape(1, static_cast(next_valid)); + GmStride2D g_stride(1); + GmTensor2D g_global(G_handle + head * total_tokens + next_start, g_shape, g_stride); + DynVecTile g_load(1, static_cast(next_valid)); + TASSIGN(g_load, G_UB); + TLOAD(g_load, g_global); + if (next_valid != C) { + TFILLPAD_INPLACE(g_ub, g_load); + } + } + } + + wait_flag_dev(2); + { + GmShape2D kv_shape(HalfC, D); + GmStride2D kv_stride(D); + GmTensor2D kv_global(workspace_handle + ws_base + WS_KV + vid * HalfC * D, kv_shape, kv_stride); + DynVecTile kv_load(HalfC, D); + TASSIGN(kv_load, S_UB_HALF); + TLOAD(kv_load, kv_global); + } + + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + TCVT(kv_ub, s_ub_half, pto::RoundMode::CAST_NONE); + pipe_barrier(PIPE_ALL); + // Finish S_{i+1} = exp(g_last) * S_i + k_i_tilde^T @ v_i_new. + // Torch-like: + // s_ub = s_ub + kv_ub + TADD(s_ub, s_ub, kv_ub); + TCVT(s_ub_half, s_ub, pto::RoundMode::CAST_NONE); + + if (ci + 1 < static_cast(num_chunks)) { + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + { + GmShape2D s_shape(HalfC, D); + GmStride2D s_stride(D); + GmTensor2D s_global(workspace_handle + ws_base + WS_S + vid * HalfC * D, s_shape, s_stride); + DynVecTile s_store(HalfC, D); + TASSIGN(s_store, S_UB_HALF); + TSTORE(s_global, s_store); + } + + // Expose the post-chunk state so the next chunk (and debug/verification + // outputs) can see S_{i+1}. Conceptually: + // S_handle[chunk_idx + 1, head] = S_{i+1} + int64_t s_out_offset = ((chunk_offset + ci + 1) * H + head) * DD; + { + GmShape2D s_out_shape(HalfC, D); + GmStride2D s_out_stride(D); + GmTensor2D s_out_global(S_handle + s_out_offset + vid * HalfC * D, s_out_shape, s_out_stride); + DynVecTile s_out_store(HalfC, D); + TASSIGN(s_out_store, S_UB_HALF); + TSTORE(s_out_global, s_out_store); + } + ffts_cross_core_sync(PIPE_MTE3, 1 | (2 << 4) | (3 << 8)); + } + + if (ci + 1 < static_cast(num_chunks)) { + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + } + } + + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + int64_t fs_offset = (seq_idx * H + head) * DD; + { + GmShape2D fs_shape(HalfC, D); + GmStride2D fs_stride(D); + GmTensor2D fs_global(FS_handle + fs_offset + vid * HalfC * D, fs_shape, fs_stride); + DynVecTile fs_store(HalfC, D); + TASSIGN(fs_store, S_UB_HALF); + TSTORE(fs_global, fs_store); + } + { + int64_t cache_offset = (fs_cache_index * H + head) * DD + vid * HalfC * D; + GmShape2D cache_shape(HalfC, D); + GmStride2D cache_stride(D); + GmTensor2D cache_global(FS_cache_handle + cache_offset, cache_shape, cache_stride); + DynVecTile cache_store(HalfC, D); + TASSIGN(cache_store, S_UB); + TSTORE(cache_global, cache_store); + } + set_flag(PIPE_MTE3, PIPE_S, EVENT_ID0); + wait_flag(PIPE_MTE3, PIPE_S, EVENT_ID0); + } +#endif +} diff --git a/xllm_ops/qwen35_gdn_prefill_super_op/op_kernel/chunk_o.cpp b/xllm_ops/qwen35_gdn_prefill_super_op/op_kernel/chunk_o.cpp new file mode 100644 index 0000000..e05a81a --- /dev/null +++ b/xllm_ops/qwen35_gdn_prefill_super_op/op_kernel/chunk_o.cpp @@ -0,0 +1,1197 @@ +// ============================================================================ +// chunk_o_kernel.cpp — Output computation for GatedDeltaNet (chunk-wise) +// +// Mathematical operation (per chunk of C tokens, per head h): +// +// O = (QK_gated @ V) + exp(g) * (Q @ S) +// = intra_chunk_attention + inter_chunk_state_contribution +// +// where: +// Q, K, V ∈ ℝ^{C×D} — query/key/value projections for this chunk +// S ∈ ℝ^{D×D} — accumulated hidden state entering this chunk +// G ∈ ℝ^{C} — cumulative gate values (pre-transposed [H,T]) +// Msk ∈ ℝ^{C×C} — lower-triangular causal mask +// +// Cube phase (3 GEMMs per chunk): +// 1. QK = Q @ K^T — intra-chunk attention scores +// 2. QS = Q @ S — query applied to accumulated state +// 3. QKV = QK_gated @ V — gated attention applied to values +// +// Vec phase (two sub-blocks process upper/lower C/2 rows): +// a. Load G → compute gating coefficients: +// coeff[i,j] = exp(min(g[i] - g[j], 0)) * mask[i,j] +// b. Apply gating to QK: QK_gated = QK * coeff +// c. Scale QS by exp(g): QS_gated = QS * exp(g_row) +// d. Combine: O = QS_gated + QKV +// e. Store O to GM in BSND layout +// +// Cross-core sync protocol (Cube ↔ Vec via FFTS): +// flag 0: Cube→Vec — QK and QS results ready in workspace +// flag 1: Vec→Cube — QK_gated written back, Cube can proceed to GEMM 3 +// flag 2: Cube→Vec — QKV result ready in workspace +// flag 3: Vec→Cube — Vec done with this chunk, Cube can reuse workspace +// +// NPU memory hierarchy used: +// GM → L1 (Cube-accessible) → L0A/L0B (matrix engines) → L0C (accumulator) +// GM → UB (Vec-accessible, on-chip SRAM) +// +// ── PTO / NPU Primer ────────────────────────────────────────────────── +// This kernel combines matrix multiplication (Cube) with element-wise gating +// (Vec) in a tightly coordinated 3-GEMM + gating pipeline per chunk. +// +// Execution timeline for one chunk: +// Cube: GEMM1(Q@K^T) → GEMM2(Q@S) → store QK,QS → signal Vec ──────┐ +// Vec: (meanwhile) load G, compute gating coefficients │ +// Vec: ←── wait for Cube signal ──── apply gating to QK → QK_gated │ +// Vec: store QK_gated → signal Cube ────────────────────────────────┐│ +// Cube: ←── wait for Vec signal ──── GEMM3(QK_gated@V) → store QKV ─┘│ +// Vec: ←── wait for Cube signal ──── scale QS, combine O=QKV+QS_g │ +// Vec: store O → signal Cube "done" ─────────────────────────────────┘ +// +// numpy pseudocode for the entire chunk computation: +// QK = Q @ K.T # GEMM 1 +// QS = Q @ S # GEMM 2 +// coeff = exp(min(g_row - g_col, 0)) * mask # gating (dynamic PTO) +// (``static_baseline/run_chunk_o_static.py`` uses exp(g_row-g_col) without min.) +// QK_gated = QK * coeff # apply gating +// QKV = QK_gated @ V # GEMM 3 +// O = QKV + QS * np.exp(g_row).reshape(-1, 1) # final output +// +// Key PTO APIs (with numpy/torch equivalents): +// TLOAD(dst, gm) — dst = gm_data (DMA: GM→UB/L1, async) +// TSTORE(gm, src) — gm = src (DMA: UB/L0C→GM, async) +// TASSIGN(tile, addr) — bind tile descriptor to buffer address +// TCVT(dst, src, mode) — type cast: dst = src.float() or .half() +// TMOV(dst, src) — copy: dst = src.clone() +// TADD(d, a, b) — d = a + b +// TSUB(d, a, b) — d = a - b +// TMUL(d, a, b) — d = a * b +// TMINS(d, s, val) — d = torch.clamp(s, max=val) +// TEXP(d, s) — d = torch.exp(s) +// TROWEXPAND(2d, col) — 2d[i,j] = col[i] (broadcast column→rows) +// TCOLEXPAND(2d, row) — 2d[i,j] = row[j] (broadcast row→columns) +// TEXTRACT(l0, l1, r, c) — copy L1 sub-tile → L0A/L0B (Cube input regs) +// TRESHAPE(zn, nz) — reinterpret L1 fractal layout (transpose, free) +// TMATMUL(C, A, B) — C = A @ B (Cube engine, fp16→fp32 accum) +// set_flag / wait_flag — synchronize pipes within same AI core +// ffts_cross_core_sync — signal across Cube↔Vec cores +// wait_flag_dev(flag) — wait for cross-core signal +// ============================================================================ + +#include +#include "acl/acl.h" +using namespace pto; + +#ifndef GDN_D +#define GDN_D 128 +#endif + +#ifndef GDN_C +#define GDN_C 128 +#endif + +// ── PTO type aliases (device-only, guarded for host pass safety) ──────────── +// The bisheng compiler performs 3 passes: vec core, cube core (__CCE_AICORE__ +// defined), and host (__CCE_AICORE__ NOT defined). Type aliases using PTO +// tile types must be guarded so the host pass never sees them. +#ifdef __CCE_AICORE__ + +// UbND = Unified Buffer tile, row-major (ND) layout, for Vec SIMD ops. +// Like torch.empty((R, C), dtype=T) in fast on-chip SRAM (~256KB). +// RV, CV = valid region (handles dynamic shapes, partial chunks). +// PadValue::Zero = fill with 0 outside valid region during TLOAD. +// T=dtype, R×C=static shape, RV×CV=valid region, P=pad fill for TLOAD. +template +using UbND = pto::Tile; + +// UbDN = UB tile in column-major (DN) layout. +// Needed as source for TROWEXPAND which requires column-format input. +// TROWEXPAND takes a column vector and broadcasts it across all columns +// of a destination ND tile: dst[i,j] = col[i] for all j. +template +using UbDN = pto::Tile; + +// L1Mat = L1 cache tile in NZ fractal format — standard Cube GEMM input. +// Data is loaded here from GM via TLOAD, then fed to L0A/L0B via TEXTRACT. +template +using L1Mat = pto::Tile; + +// L1MatZN = ZN fractal format — used for transposed GEMM operands. +// TRESHAPE(l1_zn, l1_nz) converts NZ→ZN = logical matrix transpose (free, no data movement). +template +using L1MatZN = pto::Tile; + +using GmShape2D = pto::Shape<1, 1, 1, pto::DYNAMIC, pto::DYNAMIC>; +using GmStride2D = pto::Stride<1, 1, 1, pto::DYNAMIC, 1>; + +template +using GmTensor2D = pto::GlobalTensor; + +#endif // __CCE_AICORE__ + +template +static inline AICORE void chunk_o_kernel(__gm__ half *Q_handle, __gm__ half *K_handle, __gm__ half *V_handle, + __gm__ half *S_handle, __gm__ float *G_handle, __gm__ float *Msk_handle, + __gm__ half *workspace_qk_handle, __gm__ half *workspace_qs_qkv_handle, + __gm__ half *workspace_qk_gated_handle, __gm__ half *O_handle, + __gm__ int32_t *cu_seqlens, int64_t batch_size, int64_t seq_len, + int64_t total_tokens, uint32_t num_key_heads) +{ + // Half the chunk — each Vec sub-block handles C/2 rows independently. + constexpr int32_t HalfChunk = ChunkSize / 2; + // KTail / CTail: the number of valid elements in the last 128-element tile + // when D or C isn't a multiple of 128. Used internally by PTO for partial tiles. + constexpr uint32_t KTail = (HiddenSize % 128 == 0) ? 128 : (HiddenSize % 128); + constexpr uint32_t CTail = (ChunkSize % 128 == 0) ? 128 : (ChunkSize % 128); + + constexpr int32_t H = NumHeads; + const int32_t Hg = static_cast(num_key_heads); + if (Hg <= 0 || (H % Hg) != 0) return; + const int32_t GROUP = H / Hg; + constexpr int32_t BSND_V_STRIDE = H * HiddenSize; + const int32_t BSND_QK_STRIDE = Hg * HiddenSize; + + // Workspace sizes (in elements) shared between Cube and Vec via GM + constexpr int32_t WsQKSize = ChunkSize * ChunkSize; + constexpr int32_t WsQSSize = ChunkSize * HiddenSize; + constexpr int32_t WsGatedSize = ChunkSize * ChunkSize; + + // ── UB memory map (byte addresses within Unified Buffer) ───────────── + constexpr int32_t GUbAddr = 0; + constexpr int32_t MskUbAddr = 512; + constexpr int32_t QKUbAddr = 33280; + constexpr int32_t GvUbAddr = 66048; + constexpr int32_t CoeffUbAddr = 66304; + constexpr int32_t QKHalfUbAddr = 99072; + constexpr int32_t QSHalfUbAddr = 115456; + constexpr int32_t QSUbAddr = 131840; + constexpr int32_t OHalfUbAddr = 164608; + constexpr int32_t OUbAddr = QKUbAddr; + + // cid = which AI core am I? (0..block_num-1). Used to partition work items. + auto cid = get_block_idx(); + // block_num = total number of AI cores running this kernel in parallel. + auto block_num = get_block_num(); + // vid = Vec sub-block ID (0 or 1). Each Vec core has 2 sub-blocks that + // process the upper (vid=0) and lower (vid=1) halves of C/2 rows. + auto vid = get_subblockid(); + + int64_t num_seqs = batch_size; + + // ── L1 tiles for Cube GEMM operands ────────────────────────────────── + // L1 holds matrices in NZ (col-major fractal) format for the matrix engine. + // Each tile is assigned a fixed L1 byte address to avoid runtime allocation. + // + // ── L1 tile layout for Cube GEMMs ──────────────────────────────────── + // L1 cache (~1MB) is manually partitioned for the 3 GEMMs: + // q_l1 at 0: Q [C×D] — shared by GEMM 1 and GEMM 2 + // k_l1 at 32768: K [C×D] — used in GEMM 1 (transposed via TRESHAPE) + // s_l1 at 65536: S [D×D] — accumulated state, used in GEMM 2 + // qk_gated at 98304: QK_gated [C×C] — from Vec, used in GEMM 3 + // v_l1 at 131072: V [C×D] — values, used in GEMM 3 + L1Mat q_l1; + TASSIGN(q_l1, 0); + L1Mat k_l1; + TASSIGN(k_l1, 32768); + TileAcc qk_l0; + TASSIGN(qk_l0, 0); + L1Mat s_l1; + TASSIGN(s_l1, 65536); + TileAcc qs_l0; + TASSIGN(qs_l0, 65536); + L1Mat qk_gated_l1; + TASSIGN(qk_gated_l1, 98304); + L1Mat v_l1; + TASSIGN(v_l1, 131072); + TileAcc qkv_l0; + TASSIGN(qkv_l0, 0); + + // ── UB tiles for Vec element-wise operations ───────────────────────── + // UB (Unified Buffer) is on-chip SRAM accessible by the Vec engine. + // Tiles here are row-major (ND) for standard element-wise ops. + // + // ── UB tile layout for Vec element-wise ops ────────────────────────── + // Each Vec sub-block (vid=0 or vid=1) processes C/2 rows of the C×C or C×D + // matrices. The UB layout (byte addresses) is designed so all needed tiles + // fit simultaneously in the ~256KB UB without overlapping: + // g_ub: gate values [1, C] float @ 0 + // msk_ub: causal mask [C/2, C] float @ 512 (loaded once, reused) + // qk_ub: QK scores in float [C/2, C] @ 33280 (after cast from half) + // g_v_ub: this sub-block's gate slice [1, C/2] @ 66048 + // coeff_ub: gating coefficients [C/2, C] float @ 66304 + // qk_ub_half: QK in half [C/2, C] @ 99072 + // qs_ub_half: QS in half [C/2, D] @ 115456 + // qs_ub: QS in float [C/2, D] @ 131840 + // o_ub_half: output O in half [C/2, D] @ 164608 + // o_ub: output O in float [C/2, D] @ QKUbAddr (reuses qk_ub space) + UbND g_ub; + TASSIGN(g_ub, GUbAddr); + UbND msk_ub; + TASSIGN(msk_ub, MskUbAddr); + UbND qk_ub; + TASSIGN(qk_ub, QKUbAddr); + UbND g_v_ub; + TASSIGN(g_v_ub, GvUbAddr); + UbND coeff_ub; + TASSIGN(coeff_ub, CoeffUbAddr); + UbND qk_ub_half; + TASSIGN(qk_ub_half, QKHalfUbAddr); + UbND qs_ub_half; + TASSIGN(qs_ub_half, QSHalfUbAddr); + UbND qs_ub; + TASSIGN(qs_ub, QSUbAddr); + UbND o_ub_half; + TASSIGN(o_ub_half, OHalfUbAddr); + UbND o_ub; + TASSIGN(o_ub, OUbAddr); + + // Total work items = (batches * chunks_per_sequence * heads). + // Each AI core (cid) picks every block_num-th work item (round-robin). + int64_t total_work = 0; + if (cu_seqlens == nullptr) { + int64_t chunks_per_seq = (seq_len + ChunkSize - 1) / ChunkSize; + total_work = num_seqs * chunks_per_seq * NumHeads; + } + +// ===================================================================== +// CUBE CORE — Three GEMMs per chunk: QK, QS, QKV +// Each AI core processes a different (chunk, head) pair. The Cube engine +// performs the heavy matmuls, then writes results to GM workspace for +// the Vec engine to apply gating and produce the final output. +// ===================================================================== +#if defined(__DAV_C220_CUBE__) + if (cu_seqlens == nullptr) { + // ── Fixed-length sequence path ────────────────────────────────────── + int64_t chunks_per_seq = (seq_len + ChunkSize - 1) / ChunkSize; + int64_t global_chunk_base = 0; + bool first_cube_iter = true; + + for (int64_t work_idx = static_cast(cid); work_idx < total_work; + work_idx += static_cast(block_num)) { + // Wait for Vec to finish with previous chunk's workspace (flag 3) + if (!first_cube_iter) wait_flag_dev(3); + set_flag(PIPE_FIX, PIPE_M, EVENT_ID0); + wait_flag(PIPE_FIX, PIPE_M, EVENT_ID0); + + int32_t head_idx = static_cast(work_idx % NumHeads); + int32_t head_g = head_idx / GROUP; + int64_t chunk_head_idx = work_idx / NumHeads; + int64_t seq_idx = chunk_head_idx / chunks_per_seq; + int64_t ci = chunk_head_idx % chunks_per_seq; + + int64_t bos = seq_idx * seq_len; + int64_t slen = seq_len; + int64_t chunk_start = ci * ChunkSize; + int64_t remaining = slen - chunk_start; + int32_t valid_rows = static_cast(remaining < ChunkSize ? remaining : ChunkSize); + int64_t chunk_token_start = bos + chunk_start; + int32_t row_offset = static_cast(vid) * HalfChunk; + int32_t local_rows = valid_rows - row_offset; + if (local_rows < 0) local_rows = 0; + if (local_rows > HalfChunk) local_rows = HalfChunk; + + int64_t qk_off = (chunk_token_start * static_cast(Hg) + static_cast(head_g)) * + static_cast(HiddenSize); + int64_t v_off = (chunk_token_start * static_cast(H) + static_cast(head_idx)) * + static_cast(HiddenSize); + + int64_t chunk_global_idx = seq_idx * chunks_per_seq + ci; + int64_t s_offset = (chunk_global_idx * NumHeads + head_idx) * static_cast(HiddenSize) * + static_cast(HiddenSize); + + // ── Load Q [valid_rows × D] from GM → L1 ──────────────────────── + // GlobalTensor describes the GM layout with BSND strides. + // TLOAD performs DMA (MTE2 pipe). TFILLPAD zero-pads tail rows so + // downstream GEMMs see a clean C×D matrix. + { + L1Mat _l1(valid_rows, HiddenSize); + TASSIGN(_l1, 0); + GmShape2D _gs(valid_rows, HiddenSize); + GmStride2D _stride(BSND_QK_STRIDE); + GmTensor2D _gm(Q_handle + qk_off, _gs, _stride); + TLOAD(_l1, _gm); + if (valid_rows != ChunkSize) TFILLPAD(_l1, _l1); + } + // ── Load K [valid_rows × D] from GM → L1 ──────────────────────── + { + L1Mat _l1(valid_rows, HiddenSize); + TASSIGN(_l1, 32768); + GmShape2D _gs(valid_rows, HiddenSize); + GmStride2D _stride(BSND_QK_STRIDE); + GmTensor2D _gm(K_handle + qk_off, _gs, _stride); + TLOAD(_l1, _gm); + if (valid_rows != ChunkSize) TFILLPAD(_l1, _l1); + } + + // ── GEMM 1: QK = Q @ K^T (intra-chunk attention scores) ──────── + // ── GEMM 1: QK = Q @ K^T ───────────────────────────────────────── + // numpy: QK = Q @ K.T → [C×D] @ [D×C] = [C×C] + // + // How transpose works on NPU: + // K is loaded into L1 in NZ (col-major fractal) format. + // TRESHAPE(l1_zn, k_l1) reinterprets it as ZN (row-major fractal) = K^T. + // This is a ZERO-COST operation — no data movement, just metadata change. + // TEXTRACT then loads the transposed view into L0B. + // + // Cube GEMM pipeline: + // TEXTRACT(l0a, q_l1, 0, 0) — Q → L0A (left operand) + // TEXTRACT(l0b, k_zn, 0, 0) — K^T → L0B (right operand) + // TMATMUL(qk_l0, l0a, l0b) — QK = L0A × L0B → L0C accumulator + // + // transpose_B: TRESHAPE converts k_l1 from NZ → ZN fractal layout, + // effectively transposing K before TEXTRACT loads it into L0B. + { + TileLeft _l0a; + TileRight _l0b; + TASSIGN(_l0a, 0x0); + TASSIGN(_l0b, 0x0); + auto _we = EVENT_ID1; + set_flag(PIPE_MTE2, PIPE_MTE1, _we); + wait_flag(PIPE_MTE2, PIPE_MTE1, _we); + set_flag(PIPE_M, PIPE_MTE1, _we); + wait_flag(PIPE_M, PIPE_MTE1, _we); + TEXTRACT(_l0a, q_l1, 0, 0); + L1MatZN _bzn; + TRESHAPE(_bzn, k_l1); + TEXTRACT(_l0b, _bzn, 0, 0); + set_flag(PIPE_MTE1, PIPE_M, _we); + wait_flag(PIPE_MTE1, PIPE_M, _we); + TMATMUL(qk_l0, _l0a, _l0b); + set_flag(PIPE_MTE1, PIPE_MTE2, _we); + wait_flag(PIPE_MTE1, PIPE_MTE2, _we); + set_flag(PIPE_M, PIPE_FIX, _we); + wait_flag(PIPE_M, PIPE_FIX, _we); + } + + // ── Load S [D × D] from GM → L1 (accumulated hidden state) ───── + { + L1Mat _l1(HiddenSize, HiddenSize); + TASSIGN(_l1, 65536); + Shape<1, 1, 1, DYNAMIC, DYNAMIC> _gs; + _gs.shape[3] = HiddenSize; + _gs.shape[4] = HiddenSize; + GlobalTensor> _gm(S_handle + s_offset, _gs); + TLOAD(_l1, _gm); + } + + // ── GEMM 2: QS = Q @ S (query applied to accumulated state) ──── + { + TileLeft _l0a; + TileRight _l0b; + TASSIGN(_l0a, 0x0); + TASSIGN(_l0b, 0x0); + auto _we = EVENT_ID1; + set_flag(PIPE_MTE2, PIPE_MTE1, _we); + wait_flag(PIPE_MTE2, PIPE_MTE1, _we); + set_flag(PIPE_M, PIPE_MTE1, _we); + wait_flag(PIPE_M, PIPE_MTE1, _we); + TEXTRACT(_l0a, q_l1, 0, 0); + TEXTRACT(_l0b, s_l1, 0, 0); + set_flag(PIPE_MTE1, PIPE_M, _we); + wait_flag(PIPE_MTE1, PIPE_M, _we); + TMATMUL(qs_l0, _l0a, _l0b); + set_flag(PIPE_MTE1, PIPE_MTE2, _we); + wait_flag(PIPE_MTE1, PIPE_MTE2, _we); + set_flag(PIPE_M, PIPE_FIX, _we); + wait_flag(PIPE_M, PIPE_FIX, _we); + } + + // ── Store QK [C × C] from L0C → GM workspace (fp32→fp16 cast) ─── + // TSTORE on TileAcc triggers MTE3 DMA with implicit type conversion. + { + TileAcc _l0(ChunkSize, ChunkSize); + TASSIGN(_l0, 0); + Shape<1, 1, 1, DYNAMIC, DYNAMIC> _gs; + _gs.shape[3] = ChunkSize; + _gs.shape[4] = ChunkSize; + GlobalTensor> _gm( + workspace_qk_handle + static_cast(cid) * WsQKSize, _gs); + TSTORE(_gm, _l0); + } + + // ── Store QS [C × D] from L0C → GM workspace ──────────────────── + { + TileAcc _l0(ChunkSize, HiddenSize); + TASSIGN(_l0, 65536); + Shape<1, 1, 1, DYNAMIC, DYNAMIC> _gs; + _gs.shape[3] = ChunkSize; + _gs.shape[4] = HiddenSize; + GlobalTensor> _gm( + workspace_qs_qkv_handle + static_cast(cid) * WsQSSize, _gs); + TSTORE(_gm, _l0); + } + + // Signal Vec: QK and QS are ready (flag 0, Cube→Vec) + // ── Cross-core sync protocol ────────────────────────────────────── + // Cube and Vec are SEPARATE physical cores. They exchange data through GM + // and coordinate via FFTS flags. Think of it as two processes communicating + // through shared memory with semaphores. + // + // ffts_cross_core_sync(PIPE_FIX, config): + // config = 1 | (mode << 4) | (flag_id << 8) + // mode=2: broadcast signal to all cores in this block + // flag_id: identifies which signal (0, 1, 2, 3) + // + // Protocol for this kernel: + // flag 0: Cube→Vec "QK and QS are ready in workspace" + // flag 1: Vec→Cube "QK_gated is ready for GEMM 3" + // flag 2: Cube→Vec "QKV (GEMM 3 result) is ready" + // flag 3: Vec→Cube "I'm done with this chunk, you can reuse workspace" + ffts_cross_core_sync(PIPE_FIX, 1 | (2 << 4) | (0 << 8)); + + // Wait for Vec to write QK_gated back (flag 1, Vec→Cube) + wait_flag_dev(1); + + set_flag(PIPE_FIX, PIPE_M, EVENT_ID0); + wait_flag(PIPE_FIX, PIPE_M, EVENT_ID0); + + // ── Load QK_gated [C × C] from GM workspace → L1 ──────────────── + { + L1Mat _l1(ChunkSize, ChunkSize); + TASSIGN(_l1, 98304); + Shape<1, 1, 1, DYNAMIC, DYNAMIC> _gs; + _gs.shape[3] = ChunkSize; + _gs.shape[4] = ChunkSize; + GlobalTensor> _gm( + workspace_qk_gated_handle + static_cast(cid) * WsGatedSize, _gs); + TLOAD(_l1, _gm); + } + // ── Load V [valid_rows × D] from GM → L1 ──────────────────────── + { + L1Mat _l1(valid_rows, HiddenSize); + TASSIGN(_l1, 131072); + Shape<1, 1, 1, DYNAMIC, DYNAMIC> _gs; + _gs.shape[3] = valid_rows; + _gs.shape[4] = HiddenSize; + GlobalTensor> _gm(V_handle + v_off, _gs); + TLOAD(_l1, _gm); + if (valid_rows != ChunkSize) TFILLPAD(_l1, _l1); + } + + // ── GEMM 3: QKV = QK_gated @ V (gated attention → values) ────── + { + TileLeft _l0a; + TileRight _l0b; + TASSIGN(_l0a, 0x0); + TASSIGN(_l0b, 0x0); + auto _we = EVENT_ID1; + set_flag(PIPE_MTE2, PIPE_MTE1, _we); + wait_flag(PIPE_MTE2, PIPE_MTE1, _we); + set_flag(PIPE_M, PIPE_MTE1, _we); + wait_flag(PIPE_M, PIPE_MTE1, _we); + TEXTRACT(_l0a, qk_gated_l1, 0, 0); + TEXTRACT(_l0b, v_l1, 0, 0); + set_flag(PIPE_MTE1, PIPE_M, _we); + wait_flag(PIPE_MTE1, PIPE_M, _we); + TMATMUL(qkv_l0, _l0a, _l0b); + set_flag(PIPE_MTE1, PIPE_MTE2, _we); + wait_flag(PIPE_MTE1, PIPE_MTE2, _we); + set_flag(PIPE_M, PIPE_FIX, _we); + wait_flag(PIPE_M, PIPE_FIX, _we); + } + + // ── Store QKV [C × D] from L0C → GM workspace ─────────────────── + // ── Workspace buffer reuse ──────────────────────────────────────── + // workspace_qs_qkv_handle is shared between QS (GEMM 2 output) and QKV + // (GEMM 3 output). This is safe because: + // 1. Vec reads QS BEFORE Cube writes QKV to the same buffer + // 2. The cross-core flags ensure proper ordering: + // - flag 0: QS ready (Vec reads QS) + // - flag 1: QK_gated ready (Vec done reading QS, Cube can write QKV) + // - flag 2: QKV ready (Vec reads QKV from same buffer) + { + TileAcc _l0(ChunkSize, HiddenSize); + TASSIGN(_l0, 0); + Shape<1, 1, 1, DYNAMIC, DYNAMIC> _gs; + _gs.shape[3] = ChunkSize; + _gs.shape[4] = HiddenSize; + GlobalTensor> _gm( + workspace_qs_qkv_handle + static_cast(cid) * WsQSSize, _gs); + TSTORE(_gm, _l0); + } + + // Signal Vec: QKV is ready (flag 2, Cube→Vec) + ffts_cross_core_sync(PIPE_FIX, 1 | (2 << 4) | (2 << 8)); + first_cube_iter = false; + } + } else { + // ── Variable-length sequence path (cu_seqlens != nullptr) ────────── + int64_t gi = 0; + int64_t chunk_global_idx = 0; + bool first_cube_iter_v = true; + for (int64_t si = 0; si < num_seqs; ++si) { + int64_t bos = static_cast(cu_seqlens[si]); + int64_t eos = static_cast(cu_seqlens[si + 1]); + int64_t slen = eos - bos; + int64_t nc = (slen + ChunkSize - 1) / ChunkSize; + + for (int64_t ci = 0; ci < nc; ++ci) { + for (int32_t h = 0; h < NumHeads; ++h) { + if (gi % static_cast(block_num) == static_cast(cid)) { + if (!first_cube_iter_v) wait_flag_dev(3); + set_flag(PIPE_FIX, PIPE_M, EVENT_ID0); + wait_flag(PIPE_FIX, PIPE_M, EVENT_ID0); + + int64_t chunk_start = ci * ChunkSize; + int64_t remaining = slen - chunk_start; + int32_t valid_rows = static_cast(remaining < ChunkSize ? remaining : ChunkSize); + int64_t chunk_token_start = bos + chunk_start; + int32_t head_idx = h; + int32_t head_g = head_idx / GROUP; + + int64_t qk_off = (chunk_token_start * static_cast(Hg) + static_cast(head_g)) * + static_cast(HiddenSize); + int64_t v_off = (chunk_token_start * static_cast(H) + static_cast(head_idx)) * + static_cast(HiddenSize); + int64_t s_offset = (chunk_global_idx * NumHeads + head_idx) * static_cast(HiddenSize) * + static_cast(HiddenSize); + + // Load Q + { + L1Mat _l1(valid_rows, HiddenSize); + TASSIGN(_l1, 0); + GmShape2D _gs(valid_rows, HiddenSize); + GmStride2D _stride(BSND_QK_STRIDE); + GmTensor2D _gm(Q_handle + qk_off, _gs, _stride); + TLOAD(_l1, _gm); + if (valid_rows != ChunkSize) TFILLPAD(_l1, _l1); + } + // Load K + { + L1Mat _l1(valid_rows, HiddenSize); + TASSIGN(_l1, 32768); + GmShape2D _gs(valid_rows, HiddenSize); + GmStride2D _stride(BSND_QK_STRIDE); + GmTensor2D _gm(K_handle + qk_off, _gs, _stride); + TLOAD(_l1, _gm); + if (valid_rows != ChunkSize) TFILLPAD(_l1, _l1); + } + + // GEMM 1: QK = Q @ K^T (transpose_B via TRESHAPE NZ→ZN) + { + TileLeft _l0a; + TileRight _l0b; + TASSIGN(_l0a, 0x0); + TASSIGN(_l0b, 0x0); + auto _we = EVENT_ID1; + set_flag(PIPE_MTE2, PIPE_MTE1, _we); + wait_flag(PIPE_MTE2, PIPE_MTE1, _we); + set_flag(PIPE_M, PIPE_MTE1, _we); + wait_flag(PIPE_M, PIPE_MTE1, _we); + TEXTRACT(_l0a, q_l1, 0, 0); + L1MatZN _bzn; + TRESHAPE(_bzn, k_l1); + TEXTRACT(_l0b, _bzn, 0, 0); + set_flag(PIPE_MTE1, PIPE_M, _we); + wait_flag(PIPE_MTE1, PIPE_M, _we); + TMATMUL(qk_l0, _l0a, _l0b); + set_flag(PIPE_MTE1, PIPE_MTE2, _we); + wait_flag(PIPE_MTE1, PIPE_MTE2, _we); + set_flag(PIPE_M, PIPE_FIX, _we); + wait_flag(PIPE_M, PIPE_FIX, _we); + } + + // Load S + { + L1Mat _l1(HiddenSize, HiddenSize); + TASSIGN(_l1, 65536); + Shape<1, 1, 1, DYNAMIC, DYNAMIC> _gs; + _gs.shape[3] = HiddenSize; + _gs.shape[4] = HiddenSize; + GlobalTensor> _gm(S_handle + s_offset, + _gs); + TLOAD(_l1, _gm); + } + + // GEMM 2: QS = Q @ S + { + TileLeft _l0a; + TileRight _l0b; + TASSIGN(_l0a, 0x0); + TASSIGN(_l0b, 0x0); + auto _we = EVENT_ID1; + set_flag(PIPE_MTE2, PIPE_MTE1, _we); + wait_flag(PIPE_MTE2, PIPE_MTE1, _we); + set_flag(PIPE_M, PIPE_MTE1, _we); + wait_flag(PIPE_M, PIPE_MTE1, _we); + TEXTRACT(_l0a, q_l1, 0, 0); + TEXTRACT(_l0b, s_l1, 0, 0); + set_flag(PIPE_MTE1, PIPE_M, _we); + wait_flag(PIPE_MTE1, PIPE_M, _we); + TMATMUL(qs_l0, _l0a, _l0b); + set_flag(PIPE_MTE1, PIPE_MTE2, _we); + wait_flag(PIPE_MTE1, PIPE_MTE2, _we); + set_flag(PIPE_M, PIPE_FIX, _we); + wait_flag(PIPE_M, PIPE_FIX, _we); + } + + // Store QK → workspace + { + TileAcc _l0(ChunkSize, ChunkSize); + TASSIGN(_l0, 0); + Shape<1, 1, 1, DYNAMIC, DYNAMIC> _gs; + _gs.shape[3] = ChunkSize; + _gs.shape[4] = ChunkSize; + GlobalTensor> _gm( + workspace_qk_handle + static_cast(cid) * WsQKSize, _gs); + TSTORE(_gm, _l0); + } + + // Store QS → workspace + { + TileAcc _l0(ChunkSize, HiddenSize); + TASSIGN(_l0, 65536); + Shape<1, 1, 1, DYNAMIC, DYNAMIC> _gs; + _gs.shape[3] = ChunkSize; + _gs.shape[4] = HiddenSize; + GlobalTensor> _gm( + workspace_qs_qkv_handle + static_cast(cid) * WsQSSize, _gs); + TSTORE(_gm, _l0); + } + + // Cube→Vec: QK & QS ready (flag 0) + ffts_cross_core_sync(PIPE_FIX, 1 | (2 << 4) | (0 << 8)); + + // Wait Vec→Cube: QK_gated ready (flag 1) + wait_flag_dev(1); + + set_flag(PIPE_FIX, PIPE_M, EVENT_ID0); + wait_flag(PIPE_FIX, PIPE_M, EVENT_ID0); + + // Load QK_gated + { + L1Mat _l1(ChunkSize, ChunkSize); + TASSIGN(_l1, 98304); + Shape<1, 1, 1, DYNAMIC, DYNAMIC> _gs; + _gs.shape[3] = ChunkSize; + _gs.shape[4] = ChunkSize; + GlobalTensor> _gm( + workspace_qk_gated_handle + static_cast(cid) * WsGatedSize, _gs); + TLOAD(_l1, _gm); + } + // Load V + { + L1Mat _l1(valid_rows, HiddenSize); + TASSIGN(_l1, 131072); + Shape<1, 1, 1, DYNAMIC, DYNAMIC> _gs; + _gs.shape[3] = valid_rows; + _gs.shape[4] = HiddenSize; + GlobalTensor> _gm(V_handle + v_off, + _gs); + TLOAD(_l1, _gm); + if (valid_rows != ChunkSize) TFILLPAD(_l1, _l1); + } + + // GEMM 3: QKV = QK_gated @ V + { + TileLeft _l0a; + TileRight _l0b; + TASSIGN(_l0a, 0x0); + TASSIGN(_l0b, 0x0); + auto _we = EVENT_ID1; + set_flag(PIPE_MTE2, PIPE_MTE1, _we); + wait_flag(PIPE_MTE2, PIPE_MTE1, _we); + set_flag(PIPE_M, PIPE_MTE1, _we); + wait_flag(PIPE_M, PIPE_MTE1, _we); + TEXTRACT(_l0a, qk_gated_l1, 0, 0); + TEXTRACT(_l0b, v_l1, 0, 0); + set_flag(PIPE_MTE1, PIPE_M, _we); + wait_flag(PIPE_MTE1, PIPE_M, _we); + TMATMUL(qkv_l0, _l0a, _l0b); + set_flag(PIPE_MTE1, PIPE_MTE2, _we); + wait_flag(PIPE_MTE1, PIPE_MTE2, _we); + set_flag(PIPE_M, PIPE_FIX, _we); + wait_flag(PIPE_M, PIPE_FIX, _we); + } + + { + TileAcc _l0(ChunkSize, HiddenSize); + TASSIGN(_l0, 0); + Shape<1, 1, 1, DYNAMIC, DYNAMIC> _gs; + _gs.shape[3] = ChunkSize; + _gs.shape[4] = HiddenSize; + GlobalTensor> _gm( + workspace_qs_qkv_handle + static_cast(cid) * WsQSSize, _gs); + TSTORE(_gm, _l0); + } + + ffts_cross_core_sync(PIPE_FIX, 1 | (2 << 4) | (2 << 8)); + first_cube_iter_v = false; + } + gi++; + } + chunk_global_idx++; + } + } + } +#endif + +// ===================================================================== +// VEC CORE — Gating, element-wise ops, output assembly +// Two Vec sub-blocks (vid=0,1) process upper/lower C/2 rows in parallel. +// Each sub-block independently: +// 1. Computes gating coefficients from G and the causal mask +// 2. Applies gating to the Cube's QK result → QK_gated +// 3. Scales the Cube's QS result by exp(g) +// 4. Combines QKV + scaled QS → final output O +// ===================================================================== +#if defined(__DAV_C220_VEC__) + // Vec engine initialization: set_mask_norm selects "normal" masking mode, + // and set_vector_mask(-1, -1) enables ALL SIMD lanes (no masking). + set_mask_norm(); + set_vector_mask(-1, -1); + + // ── Load causal mask once (reused across all chunks) ───────────────── + // ── Causal mask (loaded once, reused) ───────────────────────────────── + // The causal mask is a C×C lower-triangular matrix of 0s and 1s: + // mask[i,j] = 1 if i >= j else 0 + // Each sub-block loads its C/2 rows. Applied via TMUL to zero out + // non-causal (future) attention scores. + // + // Each sub-block (vid=0,1) loads its C/2 rows of the C×C lower-tri mask. + { + Shape<1, 1, 1, DYNAMIC, DYNAMIC> _gs; + _gs.shape[3] = HalfChunk; + _gs.shape[4] = ChunkSize; + GlobalTensor> _gm( + Msk_handle + static_cast(vid) * HalfChunk * ChunkSize, _gs); + UbND _ld(HalfChunk, ChunkSize); + TASSIGN(_ld, MskUbAddr); + TLOAD(_ld, _gm); + } + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + + if (cu_seqlens == nullptr) { + // ── Fixed-length sequence path ────────────────────────────────────── + int64_t chunks_per_seq = (seq_len + ChunkSize - 1) / ChunkSize; + + for (int64_t work_idx = static_cast(cid); work_idx < total_work; + work_idx += static_cast(block_num)) { + int32_t head_idx = static_cast(work_idx % NumHeads); + int64_t chunk_head_idx = work_idx / NumHeads; + int64_t seq_idx = chunk_head_idx / chunks_per_seq; + int64_t ci = chunk_head_idx % chunks_per_seq; + + int64_t bos = seq_idx * seq_len; + int64_t slen = seq_len; + int64_t chunk_start = ci * ChunkSize; + int64_t remaining = slen - chunk_start; + int32_t valid_rows = static_cast(remaining < ChunkSize ? remaining : ChunkSize); + int64_t chunk_token_start = bos + chunk_start; + int32_t row_offset = static_cast(vid) * HalfChunk; + int32_t local_rows = valid_rows - row_offset; + if (local_rows < 0) local_rows = 0; + if (local_rows > HalfChunk) local_rows = HalfChunk; + + if (local_rows > 0) { + // ── Load G [1 × valid_rows] — gate values for this chunk ──────── + // G is pre-transposed to [H, total_tokens], contiguous per head. + { + Shape<1, 1, 1, DYNAMIC, DYNAMIC> _gs; + _gs.shape[3] = 1; + _gs.shape[4] = valid_rows; + GlobalTensor> _gm( + G_handle + static_cast(head_idx) * total_tokens + chunk_token_start, _gs); + UbND _ld(1, valid_rows); + TASSIGN(_ld, GUbAddr); + TLOAD(_ld, _gm); + if (valid_rows != ChunkSize) { + UbND _pd; + TASSIGN(_pd, GUbAddr); + TFILLPAD_INPLACE(_pd, _ld); + } + } + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + + // ── Compute gating coefficients ────────────────────────────────── + // ── Gating coefficient computation (numpy pseudocode) ───────────── + // For this sub-block's rows (vid=0: rows 0..C/2-1, vid=1: rows C/2..C-1): + // + // g_row = g[my_start:my_start+C/2] # my gates (shape [C/2]) + // g_col = g[0:C] # full chunk gates (shape [C]) + // + // # Broadcast to 2D matrices: + // g_r_2d = g_row[:, None] * np.ones((1, C)) # TROWEXPAND: [C/2, C] + // g_c_2d = np.ones((C/2, 1)) * g_col[None, :] # TCOLEXPAND: [C/2, C] + // coeff = exp(min(g_r_2d - g_c_2d, 0)) * mask + // + // # Also compute exp(g_row) for QS scaling: + // exp_g_row = np.exp(g_row) # TEXP + UbND g_ub_temp_0; + TASSIGN(g_ub_temp_0, + GUbAddr + static_cast(vid) * HalfChunk * static_cast(sizeof(float))); + TMOV(g_v_ub, g_ub_temp_0); + + // Broadcast g_row into [C/2 × C] and g_col into [C/2 × C] + UbND g_r_2d; + TASSIGN(g_r_2d, QSUbAddr); + UbDN g_v_col; + TASSIGN(g_v_col, GvUbAddr); + TROWEXPAND(g_r_2d, g_v_col); // g_r_2d[i,j] = g_row[i] + TCOLEXPAND(coeff_ub, g_ub); // coeff[i,j] = g_col[j] + TSUB(coeff_ub, g_r_2d, coeff_ub); // d = g_row - g_col + pipe_barrier(PIPE_V); + TMINS(coeff_ub, coeff_ub, 0.0f); + pipe_barrier(PIPE_V); + TEXP(coeff_ub, coeff_ub); + pipe_barrier(PIPE_V); + TMUL(coeff_ub, coeff_ub, msk_ub); + pipe_barrier(PIPE_V); + TEXP(g_v_ub, g_v_ub); // exp(g_row) for QS scaling + } + + // ── Wait for Cube→Vec flag 0: QK & QS ready ───────────────────── + wait_flag_dev(0); + if (local_rows == 0) { + ffts_cross_core_sync(PIPE_MTE3, 1 | (2 << 4) | (1 << 8)); + wait_flag_dev(2); + ffts_cross_core_sync(PIPE_MTE3, 1 | (2 << 4) | (3 << 8)); + continue; + } + + // ── Load QK [C/2 × C] from workspace → UB ─────────────────────── + { + Shape<1, 1, 1, DYNAMIC, DYNAMIC> _gs; + _gs.shape[3] = local_rows; + _gs.shape[4] = ChunkSize; + GlobalTensor> _gm( + workspace_qk_handle + static_cast(cid) * WsQKSize + + static_cast(vid) * HalfChunk * ChunkSize, + _gs); + UbND _ld(local_rows, ChunkSize); + TASSIGN(_ld, QKHalfUbAddr); + TLOAD(_ld, _gm); + if (local_rows != HalfChunk) { + TFILLPAD_INPLACE(qk_ub_half, _ld); + } + } + + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + TCVT(qk_ub, qk_ub_half, pto::RoundMode::CAST_NONE); + + set_flag(PIPE_V, PIPE_MTE2, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID0); + + // ── Load QS [C/2 × D] from workspace → UB ─────────────────────── + { + Shape<1, 1, 1, DYNAMIC, DYNAMIC> _gs; + _gs.shape[3] = local_rows; + _gs.shape[4] = HiddenSize; + GlobalTensor> _gm( + workspace_qs_qkv_handle + static_cast(cid) * WsQSSize + + static_cast(vid) * HalfChunk * HiddenSize, + _gs); + UbND _ld(local_rows, HiddenSize); + TASSIGN(_ld, QSHalfUbAddr); + TLOAD(_ld, _gm); + if (local_rows != HalfChunk) { + TFILLPAD_INPLACE(qs_ub_half, _ld); + } + } + + // ── Apply gating: QK_gated = QK * exp(d*mask)*mask + TMUL(qk_ub, qk_ub, coeff_ub); + TCVT(qk_ub_half, qk_ub, pto::RoundMode::CAST_NONE); + + // ── Store QK_gated [C/2 × C] → workspace for Cube's GEMM 3 ───── + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + { + Shape<1, 1, 1, DYNAMIC, DYNAMIC> _gs; + _gs.shape[3] = local_rows; + _gs.shape[4] = ChunkSize; + GlobalTensor> _gm( + workspace_qk_gated_handle + static_cast(cid) * WsGatedSize + + static_cast(vid) * HalfChunk * ChunkSize, + _gs); + UbND _st(local_rows, ChunkSize); + TASSIGN(_st, QKHalfUbAddr); + TSTORE(_gm, _st); + } + // Vec→Cube: QK_gated ready (flag 1) + ffts_cross_core_sync(PIPE_MTE3, 1 | (2 << 4) | (1 << 8)); + + // ── Scale QS by exp(g): QS_gated = QS * exp(g_row) ────────────── + // ── Scale QS by exp(g): inter-chunk state contribution ──────────── + // numpy: QS_scaled = QS * np.exp(g_row)[:, None] (broadcast across D columns) + // TROWEXPAND broadcasts the scalar exp(g[i]) for each row i across all D columns, + // then TMUL applies it element-wise. This gates how much the accumulated state + // contributes to each token's output. + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + TCVT(qs_ub, qs_ub_half, pto::RoundMode::CAST_NONE); + UbND g_exp_2d; + TASSIGN(g_exp_2d, CoeffUbAddr); + UbDN g_v_col2; + TASSIGN(g_v_col2, GvUbAddr); + TROWEXPAND(g_exp_2d, g_v_col2); // broadcast exp(g_row) across columns + pipe_barrier(PIPE_V); + TMUL(qs_ub, qs_ub, g_exp_2d); // QS_gated = QS * exp(g_row) + + // ── Wait for Cube→Vec flag 2: QKV ready ───────────────────────── + wait_flag_dev(2); + + // ── Load QKV [C/2 × D] from workspace → UB ────────────────────── + { + Shape<1, 1, 1, DYNAMIC, DYNAMIC> _gs; + _gs.shape[3] = local_rows; + _gs.shape[4] = HiddenSize; + GlobalTensor> _gm( + workspace_qs_qkv_handle + static_cast(cid) * WsQSSize + + static_cast(vid) * HalfChunk * HiddenSize, + _gs); + UbND _ld(local_rows, HiddenSize); + TASSIGN(_ld, OHalfUbAddr); + TLOAD(_ld, _gm); + if (local_rows != HalfChunk) { + TFILLPAD_INPLACE(o_ub_half, _ld); + } + } + + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + + // ── Combine: O = QS_gated + QKV ───────────────────────────────── + // ── Final output: O = QKV + QS_scaled ───────────────────────────── + // numpy: O = (QK_gated @ V) + (Q @ S) * exp(g)[:, None] + // = intra_chunk_attention + inter_chunk_state_contribution + // TCVT half→float for QKV, then TADD, then TCVT float→half for output. + TCVT(o_ub, o_ub_half, pto::RoundMode::CAST_NONE); + TADD(o_ub, qs_ub, o_ub); + TCVT(o_ub_half, o_ub, pto::RoundMode::CAST_NONE); + + // ── Store O [C/2 × D] → GM in BSND layout ─────────────────────── + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + + int64_t o_offset = (chunk_token_start * static_cast(H) + static_cast(head_idx)) * + static_cast(HiddenSize) + + static_cast(vid) * HalfChunk * static_cast(BSND_V_STRIDE); + + { + Shape<1, 1, 1, DYNAMIC, DYNAMIC> _gs; + _gs.shape[3] = local_rows; + _gs.shape[4] = HiddenSize; + GlobalTensor> _gm(O_handle + o_offset, _gs); + UbND _st(local_rows, HiddenSize); + TASSIGN(_st, OHalfUbAddr); + TSTORE(_gm, _st); + } + + // Vec→Cube: done with this chunk (flag 3) + ffts_cross_core_sync(PIPE_MTE3, 1 | (2 << 4) | (3 << 8)); + } + } else { + // ── Variable-length sequence path (cu_seqlens != nullptr) ────────── + int64_t gi = 0; + for (int64_t si = 0; si < num_seqs; ++si) { + int64_t bos = static_cast(cu_seqlens[si]); + int64_t eos = static_cast(cu_seqlens[si + 1]); + int64_t slen = eos - bos; + int64_t nc = (slen + ChunkSize - 1) / ChunkSize; + + for (int64_t ci = 0; ci < nc; ++ci) { + for (int32_t h = 0; h < NumHeads; ++h) { + if (gi % static_cast(block_num) == static_cast(cid)) { + int64_t chunk_start = ci * ChunkSize; + int64_t remaining = slen - chunk_start; + int32_t valid_rows = static_cast(remaining < ChunkSize ? remaining : ChunkSize); + int64_t chunk_token_start = bos + chunk_start; + int32_t head_idx = h; + int32_t row_offset = static_cast(vid) * HalfChunk; + int32_t local_rows = valid_rows - row_offset; + if (local_rows < 0) local_rows = 0; + if (local_rows > HalfChunk) local_rows = HalfChunk; + + if (local_rows > 0) { + // Load G + { + Shape<1, 1, 1, DYNAMIC, DYNAMIC> _gs; + _gs.shape[3] = 1; + _gs.shape[4] = valid_rows; + GlobalTensor> _gm( + G_handle + static_cast(head_idx) * total_tokens + chunk_token_start, _gs); + UbND _ld(1, valid_rows); + TASSIGN(_ld, GUbAddr); + TLOAD(_ld, _gm); + if (valid_rows != ChunkSize) { + UbND _pd; + TASSIGN(_pd, GUbAddr); + TFILLPAD_INPLACE(_pd, _ld); + } + } + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + + // Compute gating coefficients (same math as fixed-length path — see detailed pseudocode + // above) + UbND g_ub_temp_v; + TASSIGN(g_ub_temp_v, GUbAddr + static_cast(vid) * HalfChunk * + static_cast(sizeof(float))); + TMOV(g_v_ub, g_ub_temp_v); + + UbND g_r_2d_v; + TASSIGN(g_r_2d_v, QSUbAddr); + UbDN g_v_col_v; + TASSIGN(g_v_col_v, GvUbAddr); + TROWEXPAND(g_r_2d_v, g_v_col_v); + TCOLEXPAND(coeff_ub, g_ub); + TSUB(coeff_ub, g_r_2d_v, coeff_ub); // d = g_row - g_col + pipe_barrier(PIPE_V); + TMINS(coeff_ub, coeff_ub, 0.0f); + pipe_barrier(PIPE_V); + TEXP(coeff_ub, coeff_ub); + pipe_barrier(PIPE_V); + TMUL(coeff_ub, coeff_ub, msk_ub); + pipe_barrier(PIPE_V); + TEXP(g_v_ub, g_v_ub); + } + + wait_flag_dev(0); + if (local_rows == 0) { + ffts_cross_core_sync(PIPE_MTE3, 1 | (2 << 4) | (1 << 8)); + wait_flag_dev(2); + ffts_cross_core_sync(PIPE_MTE3, 1 | (2 << 4) | (3 << 8)); + } else { + // Load QK from workspace + { + Shape<1, 1, 1, DYNAMIC, DYNAMIC> _gs; + _gs.shape[3] = local_rows; + _gs.shape[4] = ChunkSize; + GlobalTensor> _gm( + workspace_qk_handle + static_cast(cid) * WsQKSize + + static_cast(vid) * HalfChunk * ChunkSize, + _gs); + UbND _ld(local_rows, + ChunkSize); + TASSIGN(_ld, QKHalfUbAddr); + TLOAD(_ld, _gm); + if (local_rows != HalfChunk) { + TFILLPAD_INPLACE(qk_ub_half, _ld); + } + } + + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + TCVT(qk_ub, qk_ub_half, pto::RoundMode::CAST_NONE); + + set_flag(PIPE_V, PIPE_MTE2, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID0); + + // Load QS from workspace + { + Shape<1, 1, 1, DYNAMIC, DYNAMIC> _gs; + _gs.shape[3] = local_rows; + _gs.shape[4] = HiddenSize; + GlobalTensor> _gm( + workspace_qs_qkv_handle + static_cast(cid) * WsQSSize + + static_cast(vid) * HalfChunk * HiddenSize, + _gs); + UbND _ld(local_rows, + HiddenSize); + TASSIGN(_ld, QSHalfUbAddr); + TLOAD(_ld, _gm); + if (local_rows != HalfChunk) { + TFILLPAD_INPLACE(qs_ub_half, _ld); + } + } + + TMUL(qk_ub, qk_ub, coeff_ub); + TCVT(qk_ub_half, qk_ub, pto::RoundMode::CAST_NONE); // float→half for GM store + + // Store QK_gated → workspace + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + { + Shape<1, 1, 1, DYNAMIC, DYNAMIC> _gs; + _gs.shape[3] = local_rows; + _gs.shape[4] = ChunkSize; + GlobalTensor> _gm( + workspace_qk_gated_handle + static_cast(cid) * WsGatedSize + + static_cast(vid) * HalfChunk * ChunkSize, + _gs); + UbND _st(local_rows, ChunkSize); + TASSIGN(_st, QKHalfUbAddr); + TSTORE(_gm, _st); + } + // Vec→Cube: QK_gated ready (flag 1) + ffts_cross_core_sync(PIPE_MTE3, 1 | (2 << 4) | (1 << 8)); + + // Scale QS by exp(g): QS_scaled = QS * exp(g_row)[:, None] + // (same inter-chunk state scaling as fixed-length path) + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + TCVT(qs_ub, qs_ub_half, pto::RoundMode::CAST_NONE); // half→float for Vec math + + UbND g_exp_2d_v; + TASSIGN(g_exp_2d_v, CoeffUbAddr); + UbDN g_v_col2_v; + TASSIGN(g_v_col2_v, GvUbAddr); + TROWEXPAND(g_exp_2d_v, g_v_col2_v); + pipe_barrier(PIPE_V); + TMUL(qs_ub, qs_ub, g_exp_2d_v); + + wait_flag_dev(2); + + // Load QKV from workspace + { + Shape<1, 1, 1, DYNAMIC, DYNAMIC> _gs; + _gs.shape[3] = local_rows; + _gs.shape[4] = HiddenSize; + GlobalTensor> _gm( + workspace_qs_qkv_handle + static_cast(cid) * WsQSSize + + static_cast(vid) * HalfChunk * HiddenSize, + _gs); + UbND _ld(local_rows, + HiddenSize); + TASSIGN(_ld, OHalfUbAddr); + TLOAD(_ld, _gm); + if (local_rows != HalfChunk) { + TFILLPAD_INPLACE(o_ub_half, _ld); + } + } + + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + + // O = QS_gated + QKV (final output: intra-chunk attention + inter-chunk state) + TCVT(o_ub, o_ub_half, pto::RoundMode::CAST_NONE); // half→float + TADD(o_ub, qs_ub, o_ub); // O = QS_scaled + QKV + TCVT(o_ub_half, o_ub, pto::RoundMode::CAST_NONE); // float→half for GM store + + // Store O → GM + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + + int64_t o_offset = + (chunk_token_start * static_cast(H) + static_cast(head_idx)) * + static_cast(HiddenSize) + + static_cast(vid) * HalfChunk * static_cast(BSND_V_STRIDE); + + { + Shape<1, 1, 1, DYNAMIC, DYNAMIC> _gs; + _gs.shape[3] = local_rows; + _gs.shape[4] = HiddenSize; + GlobalTensor> _gm( + O_handle + o_offset, _gs); + UbND _st(local_rows, HiddenSize); + TASSIGN(_st, OHalfUbAddr); + TSTORE(_gm, _st); + } + + // Vec→Cube: done with this chunk (flag 3) + ffts_cross_core_sync(PIPE_MTE3, 1 | (2 << 4) | (3 << 8)); + } + } + gi++; + } + } + } + } +#endif +} diff --git a/xllm_ops/qwen35_gdn_prefill_super_op/op_kernel/qwen35_gdn_prefill_super_op.cpp b/xllm_ops/qwen35_gdn_prefill_super_op/op_kernel/qwen35_gdn_prefill_super_op.cpp new file mode 100644 index 0000000..9a4d21f --- /dev/null +++ b/xllm_ops/qwen35_gdn_prefill_super_op/op_kernel/qwen35_gdn_prefill_super_op.cpp @@ -0,0 +1,914 @@ +// Qwen3.5 TP2 prefill state super-kernel. +// +// Same pipeline as pto_mega_kernel, with value heads (H) and key heads (Hg) +// passed at runtime. H dispatches to finite compile-time specializations. +// +// Stages: +// 1. cumsum (Vec) +// 2. transpose (Vec) +// 3. kkt (Cube+Vec) — K has Hg heads; β,g,A use H value heads +// 4. solve_tril (Cube) +// 5. wy_fast (Vec+Cube) +// 6. chunk_h (Cube+Vec) +// 7. chunk_o (Cube+Vec) + +#ifndef GDN_D +#define GDN_D 128 +#endif +#ifndef GDN_C +#define GDN_C 128 +#endif +#ifndef MEMORY_BASE +#define MEMORY_BASE +#endif +#ifndef GDN_KERNEL_NAME +#define GDN_KERNEL_NAME launch_qwen35_gdn_prefill_super_op +#endif +// Note the codegen parser does not support arguments of form "type *name", only "type* name" +// clang-format off +#ifndef GM_ADDR +#define GM_ADDR __gm__ uint8_t* +#endif +// clang-format off + +#include "acl/acl.h" +#include "kernel_operator.h" +#define CAUSAL_CONV1D_SKIP_TPL_REGISTRATION +#include "../../causal_conv1d/op_kernel/causal_conv1d_fn.h" +#undef CAUSAL_CONV1D_SKIP_TPL_REGISTRATION +#include +#include +using namespace pto; + +namespace qwen35_mem { +template +__aicore__ inline void CopyGmToUb(AscendC::LocalTensor dst, AscendC::GlobalTensor src, + uint32_t realSrcN = 1, uint32_t maskShapeM = dstM, + uint32_t maskShapeN = dstN, T padValue = T(0)) +{ + if (maskShapeM != dstM || maskShapeN != dstN) { + AscendC::Duplicate(dst, padValue, dstM * dstN); + AscendC::PipeBarrier(); + } + const bool isPad = maskShapeN != dstN; + AscendC::DataCopyExtParams params(maskShapeM, maskShapeN * sizeof(T), + (realSrcN - maskShapeN) * sizeof(T), + (dstN - maskShapeN) * sizeof(T) / 32, 0); + AscendC::DataCopyPadExtParams padParams(isPad, 0, isPad ? 1 : 0, padValue); + AscendC::DataCopyPad(dst, src, params, padParams); +} + +template +__aicore__ inline void CopyUbToGm(AscendC::GlobalTensor dst, AscendC::LocalTensor src, + uint32_t realDstN = 1, uint32_t maskShapeM = srcM, + uint32_t maskShapeN = srcN) +{ + AscendC::DataCopyExtParams params(maskShapeM, maskShapeN * sizeof(T), + (srcN - maskShapeN) * sizeof(T) / 32, + (realDstN - maskShapeN) * sizeof(T), 0); + AscendC::DataCopyPad(dst, src, params); +} + +template +__aicore__ inline void CopyUbToUb(AscendC::LocalTensor dst, AscendC::LocalTensor src) +{ + AscendC::DataCopy(dst, src, length); +} + +template +__aicore__ inline void ReduceRows(AscendC::LocalTensor dst, AscendC::LocalTensor src, + AscendC::LocalTensor tmp) +{ + uint32_t shape[] = {M, N}; + AscendC::ReduceSum(dst, src, tmp, shape, true); +} + +template +__aicore__ inline void Broadcast(AscendC::LocalTensor dst, AscendC::LocalTensor src, + AscendC::LocalTensor &tmp, + const uint32_t dstShape[dim], const uint32_t srcShape[dim]) +{ + AscendC::Broadcast(dst, src, dstShape, srcShape, tmp); +} +} // namespace qwen35_mem + +struct Qwen35GdnPrefillSuperOpTilingData { + uint32_t block_dim; + uint32_t num_matrices; + uint32_t num_heads; + uint32_t num_key_heads; + uint32_t token_block_size; + uint32_t token_block_count; + int64_t conv_state_index; + int64_t ssm_state_index; + int64_t batch_size; + int64_t seq_len; + int64_t total_tokens; +}; + +// =================================================================== +// Device-only helpers (shared with standard mega-kernel) +// =================================================================== +#ifdef __CCE_AICORE__ + +constexpr uint16_t SYNC_AIV_FLAG = 12; +constexpr uint16_t SYNC_AIC_FLAG = 11; +constexpr uint16_t SYNC_AIC_AIV_FLAG = 13; +constexpr uint16_t SYNC_AIV_ONLY_ALL = 14; +constexpr uint16_t SYNC_MODE_SHIFT_VALUE = 4; +constexpr uint16_t SYNC_FLAG_SHIFT_VALUE = 8; + +AICORE inline uint16_t GetffstMsg(uint16_t mode, uint16_t flagId) +{ + return (0x1 + ((mode & 0x3) << SYNC_MODE_SHIFT_VALUE) + ((flagId & 0xf) << SYNC_FLAG_SHIFT_VALUE)); +} + +template +AICORE inline void SyncAllImpl() +{ + pipe_barrier(PIPE_ALL); + if constexpr (isAIVOnly) { + ffts_cross_core_sync(PIPE_MTE3, GetffstMsg(0x0, SYNC_AIV_ONLY_ALL)); + wait_flag_dev(SYNC_AIV_ONLY_ALL); + return; + } +#if defined(__DAV_C220_CUBE__) + wait_flag_dev(SYNC_AIV_FLAG); + ffts_cross_core_sync(PIPE_FIX, GetffstMsg(0x0, SYNC_AIC_FLAG)); + wait_flag_dev(SYNC_AIC_FLAG); + ffts_cross_core_sync(PIPE_MTE3, GetffstMsg(0x02, SYNC_AIC_AIV_FLAG)); +#elif defined(__DAV_C220_VEC__) + ffts_cross_core_sync(PIPE_MTE3, GetffstMsg(0x02, SYNC_AIV_FLAG)); + wait_flag_dev(SYNC_AIC_AIV_FLAG); +#endif +} + +namespace NsCausalConv1d { +template +class Qwen35PrefillConv : public CausalConv1d { +public: + __aicore__ inline void Init(GM_ADDR x, GM_ADDR weight, GM_ADDR convState, GM_ADDR packedQkv, + const CausalConv1dTilingData *tilingData) + { + this->ResetRuntimeState(tilingData); + this->xGm.SetGlobalBuffer(reinterpret_cast<__gm__ T *>(x)); + this->weightGm.SetGlobalBuffer(reinterpret_cast<__gm__ T *>(weight)); + this->convStatesGm.SetGlobalBuffer(reinterpret_cast<__gm__ T *>(convState)); + this->packedQkvYGm.SetGlobalBuffer(reinterpret_cast<__gm__ half *>(packedQkv)); + this->InitSharedBuffersAndEvents(); + } + + __aicore__ inline void Process(int32_t cacheIndex) + { + const int32_t blockIdx = static_cast(GetBlockIdx()); + const int32_t tokenCount = static_cast(this->tilingData_->cuSeqlen); + const auto task = ResolveFnPackedQkvBlockTask( + blockIdx, static_cast(this->tilingData_->tokenBlockCnt), + static_cast(this->tilingData_->tokenBlockSize), tokenCount, + static_cast(this->tilingData_->baseDimCnt), + static_cast(this->tilingData_->baseDim), static_cast(this->tilingData_->dim)); + if (task.valid) { + this->ProcessFnChunk(cacheIndex, false, 0, tokenCount, task.tokenStart, + task.tokenEnd - task.tokenStart, task.channelStart, + task.baseDimSize, static_cast(this->tilingData_->dim)); + } + this->ReleaseEvents(); + } +}; +} // namespace NsCausalConv1d + +AICORE inline void RunQwen35PrefillConv(GM_ADDR mixedQkv, GM_ADDR convWeight, GM_ADDR convState, + GM_ADDR packedQkv, int64_t totalTokens, uint32_t tokenBlockSize, + uint32_t tokenBlockCount, int64_t cacheIndex) +{ +#if defined(__DAV_C220_VEC__) + CausalConv1dTilingData tiling{}; + tiling.dim = 5120; + tiling.cuSeqlen = totalTokens; + tiling.seqLen = totalTokens; + tiling.inputMode = 0; + tiling.width = 4; + tiling.stateLen = 3; + tiling.numCacheLines = cacheIndex + 1; + tiling.batch = 1; + tiling.activationMode = CAUSAL_CONV1D_ACTIVATION_SILU_PACKED_QKV; + tiling.padSlotId = -1; + tiling.hasBias = 0; + tiling.baseDim = 3072; + tiling.baseDimCnt = 2; + tiling.hasNumAcceptedTokens = 0; + tiling.hasCacheIndices = 0; + tiling.hasInitialStateMode = 0; + tiling.hasInitStateWorkspace = 0; + tiling.tokenBlockSize = tokenBlockSize; + tiling.tokenBlockCnt = tokenBlockCount; + tiling.hasExplicitTokenSeqRanges = 0; + + NsCausalConv1d::Qwen35PrefillConv + op; + op.Init(mixedQkv, convWeight, convState, packedQkv, &tiling); + op.Process(static_cast(cacheIndex)); +#endif +} + +AICORE inline int32_t QwenVectorTaskId() +{ +#if defined(__DAV_C220_VEC__) + return static_cast(AscendC::GetBlockIdx() / 2 * 2 + AscendC::GetSubBlockIdx()); +#else + return 0; +#endif +} + +AICORE inline void PrepareQwen35Gate(GM_ADDR aPtr, GM_ADDR bPtr, GM_ADDR aLogPtr, GM_ADDR dtBiasPtr, + GM_ADDR gPtr, GM_ADDR betaPtr, int64_t totalTokens) +{ +#if defined(__DAV_C220_VEC__) + AscendC::TPipe pipe; + AscendC::TBuf ub; + pipe.InitBuffer(ub, 32768); + + auto aLog = ub.GetWithOffset(64, 0); + auto dtBias = ub.GetWithOffset(64, 256); + auto aBf16 = ub.GetWithOffset(64, 512); + auto bBf16 = ub.GetWithOffset(64, 640); + auto x = ub.GetWithOffset(64, 768); + auto betaX = ub.GetWithOffset(64, 1024); + auto absX = ub.GetWithOffset(64, 1280); + auto tmp = ub.GetWithOffset(64, 1536); + auto betaFp32 = ub.GetWithOffset(64, 1792); + auto betaRounded = ub.GetWithOffset(64, 2048); + auto betaHalf = ub.GetWithOffset(64, 2176); + auto cmpMask = ub.GetWithOffset(64, 2304); + auto sigmoidTmp = ub.GetWithOffset(2048, 2368); + + AscendC::GlobalTensor aLogGm; + AscendC::GlobalTensor dtBiasGm; + AscendC::GlobalTensor aGm; + AscendC::GlobalTensor bGm; + AscendC::GlobalTensor gGm; + AscendC::GlobalTensor betaGm; + aLogGm.SetGlobalBuffer(reinterpret_cast<__gm__ float *>(aLogPtr)); + dtBiasGm.SetGlobalBuffer(reinterpret_cast<__gm__ float *>(dtBiasPtr)); + aGm.SetGlobalBuffer(reinterpret_cast<__gm__ bfloat16_t *>(aPtr)); + bGm.SetGlobalBuffer(reinterpret_cast<__gm__ bfloat16_t *>(bPtr)); + gGm.SetGlobalBuffer(reinterpret_cast<__gm__ float *>(gPtr)); + betaGm.SetGlobalBuffer(reinterpret_cast<__gm__ half *>(betaPtr)); + + qwen35_mem::CopyGmToUb(aLog[0], aLogGm[0], 24, 1, 24, 0.0f); + qwen35_mem::CopyGmToUb(dtBias[0], dtBiasGm[0], 24, 1, 24, 0.0f); + AscendC::SetFlag(0); + AscendC::WaitFlag(0); + AscendC::Exp(aLog[0], aLog[0], 64); + AscendC::PipeBarrier(); + AscendC::Muls(aLog[0], aLog[0], -1.0f, 64); + + const int32_t taskId = QwenVectorTaskId(); + constexpr int32_t kVectorTasks = 40; + for (int64_t token = taskId; token < totalTokens; token += kVectorTasks) { + AscendC::PipeBarrier(); + qwen35_mem::CopyGmToUb( + aBf16[0], aGm[token * 24], 24, 1, 24, bfloat16_t(0.0f)); + qwen35_mem::CopyGmToUb( + bBf16[0], bGm[token * 24], 24, 1, 24, bfloat16_t(0.0f)); + AscendC::SetFlag(1); + AscendC::WaitFlag(1); + AscendC::Cast(x[0], aBf16[0], AscendC::RoundMode::CAST_NONE, 64); + AscendC::PipeBarrier(); + AscendC::Add(x[0], x[0], dtBias[0], 64); + AscendC::PipeBarrier(); + AscendC::Abs(absX[0], x[0], 64); + AscendC::Muls(tmp[0], absX[0], -1.0f, 64); + AscendC::PipeBarrier(); + AscendC::Exp(betaFp32[0], tmp[0], 64); + AscendC::PipeBarrier(); + AscendC::Adds(betaFp32[0], betaFp32[0], 1.0f, 64); + AscendC::PipeBarrier(); + AscendC::Ln(tmp[0], betaFp32[0], 64); + AscendC::CompareScalar(cmpMask[0], x[0], 20.0f, AscendC::CMPMODE::GT, 64); + AscendC::Add(betaX[0], x[0], absX[0], 64); + AscendC::PipeBarrier(); + AscendC::Muls(betaX[0], betaX[0], 0.5f, 64); + AscendC::Add(betaX[0], betaX[0], tmp[0], 64); + AscendC::PipeBarrier(); + AscendC::Select(betaX[0], cmpMask[0], x[0], betaX[0], + AscendC::SELMODE::VSEL_TENSOR_TENSOR_MODE, 64); + AscendC::Mul(x[0], aLog[0], betaX[0], 64); + + AscendC::Cast(betaFp32[0], bBf16[0], AscendC::RoundMode::CAST_NONE, 64); + AscendC::PipeBarrier(); + AscendC::Sigmoid(betaFp32[0], betaFp32[0], sigmoidTmp[0], 64); + AscendC::PipeBarrier(); + AscendC::Cast(betaRounded[0], betaFp32[0], AscendC::RoundMode::CAST_RINT, 64); + AscendC::PipeBarrier(); + AscendC::Cast(betaFp32[0], betaRounded[0], AscendC::RoundMode::CAST_NONE, 64); + AscendC::Cast(betaHalf[0], betaFp32[0], AscendC::RoundMode::CAST_RINT, 64); + AscendC::SetFlag(2); + AscendC::WaitFlag(2); + qwen35_mem::CopyUbToGm(gGm[token * 24], x[0], 24, 1, 24); + qwen35_mem::CopyUbToGm(betaGm[token * 24], betaHalf[0], 24, 1, 24); + } + AscendC::PipeBarrier(); +#endif +} + +AICORE inline void ZeroInitialState(GM_ADDR initialStatePtr) +{ +#if defined(__DAV_C220_VEC__) + AscendC::TPipe pipe; + AscendC::TBuf ub; + pipe.InitBuffer(ub, 512); + auto zero = ub.Get(); + AscendC::Duplicate(zero, half(0.0f), 256); + AscendC::GlobalTensor state; + state.SetGlobalBuffer(reinterpret_cast<__gm__ half *>(initialStatePtr)); + constexpr int64_t kElements = 24 * 128 * 128; + constexpr int64_t kChunk = 256; + const int32_t taskId = QwenVectorTaskId(); + for (int64_t offset = static_cast(taskId) * kChunk; offset < kElements; + offset += 40 * kChunk) { + const int64_t count = (offset + kChunk <= kElements) ? kChunk : (kElements - offset); + AscendC::DataCopy(state[offset], zero, count); + } + AscendC::PipeBarrier(); +#endif +} + +AICORE inline void ZeroMegaScratch(GM_ADDR megaAPtr, GM_ADDR aInvF32Ptr, + GM_ADDR aInvPtr, GM_ADDR hPtr, + GM_ADDR finalStatePtr, + int64_t totalTokens, + uint32_t numMatrices) +{ +#if defined(__DAV_C220_VEC__) + constexpr int64_t kChunk = 256; + constexpr int64_t kTasks = 40; + constexpr int64_t kHeads = 24; + constexpr int64_t kHeadDim = 128; + const int32_t taskId = QwenVectorTaskId(); + const int64_t matrixElements = totalTokens * kHeads * kHeadDim; + const int64_t hElements = static_cast(numMatrices) * kHeadDim * kHeadDim; + const int64_t finalStateElements = kHeads * kHeadDim * kHeadDim; + + AscendC::TPipe pipe; + AscendC::TBuf ub; + pipe.InitBuffer(ub, 1536); + auto zeroHalf = ub.GetWithOffset(kChunk, 0); + auto zeroFloat = ub.GetWithOffset(kChunk, 512); + AscendC::Duplicate(zeroHalf, half(0.0f), kChunk); + AscendC::Duplicate(zeroFloat, 0.0f, kChunk); + AscendC::SetFlag(5); + AscendC::WaitFlag(5); + + AscendC::GlobalTensor megaA; + AscendC::GlobalTensor aInvF32; + AscendC::GlobalTensor aInv; + AscendC::GlobalTensor h; + AscendC::GlobalTensor finalState; + megaA.SetGlobalBuffer(reinterpret_cast<__gm__ half *>(megaAPtr)); + aInvF32.SetGlobalBuffer(reinterpret_cast<__gm__ float *>(aInvF32Ptr)); + aInv.SetGlobalBuffer(reinterpret_cast<__gm__ half *>(aInvPtr)); + h.SetGlobalBuffer(reinterpret_cast<__gm__ half *>(hPtr)); + finalState.SetGlobalBuffer(reinterpret_cast<__gm__ half *>(finalStatePtr)); + + for (int64_t offset = static_cast(taskId) * kChunk; + offset < matrixElements; offset += kTasks * kChunk) { + AscendC::DataCopy(megaA[offset], zeroHalf, kChunk); + AscendC::DataCopy(aInvF32[offset], zeroFloat, kChunk); + AscendC::DataCopy(aInv[offset], zeroHalf, kChunk); + } + for (int64_t offset = static_cast(taskId) * kChunk; + offset < hElements; offset += kTasks * kChunk) { + AscendC::DataCopy(h[offset], zeroHalf, kChunk); + } + for (int64_t offset = static_cast(taskId) * kChunk; + offset < finalStateElements; offset += kTasks * kChunk) { + AscendC::DataCopy(finalState[offset], zeroHalf, kChunk); + } + AscendC::PipeBarrier(); +#endif +} + +AICORE inline void RunQwen35OutputNorm(GM_ADDR megaOutPtr, GM_ADDR zPtr, GM_ADDR normWeightPtr, + GM_ADDR outPtr, int64_t totalTokens) +{ +#if defined(__DAV_C220_VEC__) + constexpr int32_t kRows = 8; + constexpr int32_t kHeadDim = 128; + constexpr int32_t kElements = kRows * kHeadDim; + constexpr float kScale = 0.08838834764831845f; + + AscendC::TPipe pipe; + AscendC::TBuf ub; + pipe.InitBuffer(ub, 49152); + auto weightBf16 = ub.GetWithOffset(128, 0); + auto weightBase = ub.GetWithOffset(128, 256); + auto weight = ub.GetWithOffset(kElements, 768); + auto xHalf = ub.GetWithOffset(kElements, 4864); + auto zBf16 = ub.GetWithOffset(kElements, 6912); + auto x = ub.GetWithOffset(kElements, 8960); + auto z = ub.GetWithOffset(kElements, 13056); + auto silu = ub.GetWithOffset(kElements, 17152); + auto square = ub.GetWithOffset(kElements, 21248); + auto rounded = ub.GetWithOffset(kElements, 25344); + auto rms = ub.GetWithOffset(kRows, 27392); + auto rms2d = ub.GetWithOffset(kElements, 27424); + auto tmp = ub.GetWithOffset(4096, 31520); + + AscendC::GlobalTensor megaOut; + AscendC::GlobalTensor gate; + AscendC::GlobalTensor normWeight; + AscendC::GlobalTensor out; + megaOut.SetGlobalBuffer(reinterpret_cast<__gm__ half *>(megaOutPtr)); + gate.SetGlobalBuffer(reinterpret_cast<__gm__ bfloat16_t *>(zPtr)); + normWeight.SetGlobalBuffer(reinterpret_cast<__gm__ bfloat16_t *>(normWeightPtr)); + out.SetGlobalBuffer(reinterpret_cast<__gm__ bfloat16_t *>(outPtr)); + + qwen35_mem::CopyGmToUb( + weightBf16[0], normWeight[0], 128, 1, 128, bfloat16_t(0.0f)); + AscendC::SetFlag(2); + AscendC::WaitFlag(2); + AscendC::Cast(weightBase[0], weightBf16[0], AscendC::RoundMode::CAST_NONE, 128); + for (int32_t row = 0; row < kRows; ++row) { + AscendC::PipeBarrier(); + qwen35_mem::CopyUbToUb(weight[row * 128], weightBase[0]); + } + + const int64_t numRows = totalTokens * 24; + const int64_t numChunks = (numRows + kRows - 1) / kRows; + const int32_t taskId = QwenVectorTaskId(); + for (int64_t chunk = taskId; chunk < numChunks; chunk += 40) { + const int64_t rowStart = chunk * kRows; + const int32_t validRows = + static_cast((rowStart + kRows <= numRows) ? kRows : (numRows - rowStart)); + for (int32_t row = 0; row < validRows; ++row) { + qwen35_mem::CopyGmToUb( + xHalf[row * 128], megaOut[(rowStart + row) * 128], 128, 1, 128, half(0.0f)); + qwen35_mem::CopyGmToUb( + zBf16[row * 128], gate[(rowStart + row) * 128], 128, 1, 128, bfloat16_t(0.0f)); + } + for (int32_t row = validRows; row < kRows; ++row) { + AscendC::Duplicate(xHalf[row * 128], half(0.0f), 128); + AscendC::Duplicate(zBf16[row * 128], bfloat16_t(0.0f), 128); + } + AscendC::SetFlag(3); + AscendC::WaitFlag(3); + AscendC::Cast(x[0], xHalf[0], AscendC::RoundMode::CAST_NONE, kElements); + AscendC::Muls(x[0], x[0], kScale, kElements); + AscendC::PipeBarrier(); + AscendC::Cast(rounded[0], x[0], AscendC::RoundMode::CAST_RINT, kElements); + AscendC::PipeBarrier(); + AscendC::Cast(x[0], rounded[0], AscendC::RoundMode::CAST_NONE, kElements); + AscendC::Cast(z[0], zBf16[0], AscendC::RoundMode::CAST_NONE, kElements); + AscendC::PipeBarrier(); + AscendC::Silu(silu[0], z[0], kElements); + AscendC::Mul(square[0], x[0], x[0], kElements); + AscendC::PipeBarrier(); + qwen35_mem::ReduceRows(rms[0], square[0], tmp[0]); + AscendC::PipeBarrier(); + AscendC::Muls(rms[0], rms[0], 1.0f / 128.0f, kRows); + AscendC::Adds(rms[0], rms[0], 1.0e-6f, kRows); + AscendC::PipeBarrier(); + AscendC::Sqrt(rms[0], rms[0], kRows); + AscendC::PipeBarrier(); + qwen35_mem::Broadcast( + rms2d[0], rms[0], tmp, (uint32_t[]){8, 128}, (uint32_t[]){8, 1}); + AscendC::PipeBarrier(); + AscendC::Div(x[0], x[0], rms2d[0], kElements); + AscendC::Mul(x[0], x[0], weight[0], kElements); + AscendC::PipeBarrier(); + AscendC::Mul(x[0], x[0], silu[0], kElements); + AscendC::PipeBarrier(); + AscendC::Cast(rounded[0], x[0], AscendC::RoundMode::CAST_RINT, kElements); + AscendC::SetFlag(4); + AscendC::WaitFlag(4); + for (int32_t row = 0; row < validRows; ++row) { + qwen35_mem::CopyUbToGm( + out[(rowStart + row) * 128], rounded[row * 128], 128, 1, 128); + } + AscendC::PipeBarrier(); + } +#endif +} + +template +AICORE inline void mega_transpose_TH_to_HT(__gm__ T *src, __gm__ T *dst, int64_t T_len) +{ +#if defined(__DAV_C220_VEC__) + if (get_subblockid() != 0) return; + set_mask_norm(); + set_vector_mask(-1, -1); + + auto cid = get_block_idx(); + auto block_num = get_block_num(); + + constexpr int32_t BLOCK = 128; + constexpr int32_t H = static_cast(H_val); + constexpr int32_t ES = static_cast(sizeof(T)); + constexpr int32_t AlignBytes = 32; + constexpr int32_t AlignRows = AlignBytes / ES; + constexpr int32_t MinTransposeCols = 16; + constexpr int32_t AlignElems = (AlignRows > MinTransposeCols) ? AlignRows : MinTransposeCols; + constexpr int32_t HP = ((H + AlignElems - 1) / AlignElems) * AlignElems; + constexpr int32_t SRC_UB = 0; + constexpr int32_t DST_UB = SRC_UB + BLOCK * HP * ES; + constexpr int32_t TMP_UB = DST_UB + HP * BLOCK * ES; + + using UBSrcFull = + Tile; + using UBSrcDyn = + Tile; + using UBDst = Tile; + using UBDstDyn = Tile; + using UBTmp = Tile; + + using UBRow = Tile; + using UBRowDyn = Tile; + using UBHeadDyn = Tile; + + using Gm2D = Shape<1, 1, 1, DYNAMIC, DYNAMIC>; + using Gm1D = Shape<1, 1, 1, 1, DYNAMIC>; + using GmSrcS = Stride<1, 1, 1, H, 1>; + using GmHeadS = Stride<1, 1, 1, 1, H>; + using GmS1 = Stride<1, 1, 1, 1, 1>; + + if constexpr (H < MinTransposeCols) { + int64_t num_tok_blocks = (T_len + BLOCK - 1) / BLOCK; + for (int64_t bi = static_cast(cid); bi < num_tok_blocks; bi += static_cast(block_num)) { + int64_t t0 = bi * BLOCK; + int32_t valid = (t0 + BLOCK <= T_len) ? BLOCK : static_cast(T_len - t0); + + for (int32_t h = 0; h < H; ++h) { + Gm1D gs; + gs.shape[4] = valid; + UBHeadDyn row(valid); + TASSIGN(row, SRC_UB); + { + GlobalTensor gm(src + t0 * H + h, gs); + TLOAD(row, gm); + } + set_flag(PIPE_MTE2, PIPE_MTE3, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_MTE3, EVENT_ID0); + { + GlobalTensor gm(dst + h * T_len + t0, gs); + TSTORE(gm, row); + } + set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0); + wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0); + } + } + return; + } + + UBSrcFull ub_src; + TASSIGN(ub_src, SRC_UB); + UBDst ub_dst; + TASSIGN(ub_dst, DST_UB); + UBTmp ub_tmp; + TASSIGN(ub_tmp, TMP_UB); + + int64_t num_tok_blocks = (T_len + BLOCK - 1) / BLOCK; + + for (int64_t bi = static_cast(cid); bi < num_tok_blocks; bi += static_cast(block_num)) { + int64_t t0 = bi * BLOCK; + int32_t valid = (t0 + BLOCK <= T_len) ? BLOCK : static_cast(T_len - t0); + + { + Gm2D gs; + gs.shape[3] = valid; + gs.shape[4] = H; + GlobalTensor gm(src + t0 * H, gs); + UBSrcDyn ld(valid, H); + TASSIGN(ld, SRC_UB); + TLOAD(ld, gm); + if (valid != BLOCK || H != HP) TFILLPAD_INPLACE(ub_src, ld); + } + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + + TTRANS(ub_dst, ub_src, ub_tmp); + + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + + for (int32_t h = 0; h < H; ++h) { + Gm1D gs; + gs.shape[4] = valid; + GlobalTensor gm(dst + h * T_len + t0, gs); + UBRowDyn st(1, valid); + TASSIGN(st, DST_UB + h * BLOCK * ES); + TSTORE(gm, st); + } + set_flag(PIPE_MTE3, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE3, PIPE_V, EVENT_ID0); + } +#endif +} + +template +AICORE inline void mega_cast_fp32_to_fp16_bsnd(__gm__ float *src, __gm__ half *dst, uint32_t num_matrices, + int64_t total_tokens) +{ +#if defined(__DAV_C220_VEC__) + if (get_subblockid() != 0) return; + set_mask_norm(); + set_vector_mask(-1, -1); + + auto cid = get_block_idx(); + auto block_num = get_block_num(); + + constexpr int32_t F32_UB = 0; + constexpr int32_t F16_UB = C * static_cast(sizeof(float)); + + using SrcUB = Tile; + using DynSrcUB = + Tile; + using DstUB = Tile; + using DynDstUB = Tile; + using Gm1D = Shape<1, 1, 1, 1, DYNAMIC>; + using GmS1 = Stride<1, 1, 1, 1, 1>; + + SrcUB src_ub; + TASSIGN(src_ub, F32_UB); + DstUB dst_ub; + TASSIGN(dst_ub, F16_UB); + + for (uint32_t m = cid; m < num_matrices; m += block_num) { + uint32_t h = m % static_cast(H); + uint32_t chunk_idx = m / static_cast(H); + + for (int64_t t = 0; t < total_tokens; ++t) { + int64_t off = t * static_cast(H * C) + static_cast(h * C); + + { + Gm1D gs; + gs.shape[4] = C; + GlobalTensor gm(src + off, gs); + SrcUB ld; + TASSIGN(ld, F32_UB); + TLOAD(ld, gm); + } + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + + TCVT(dst_ub, src_ub, RoundMode::CAST_NONE); + + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + { + Gm1D gs; + gs.shape[4] = C; + GlobalTensor gm(dst + off, gs); + DstUB st; + TASSIGN(st, F16_UB); + TSTORE(gm, st); + } + set_flag(PIPE_MTE3, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE3, PIPE_V, EVENT_ID0); + } + } +#endif +} + +#endif // __CCE_AICORE__ + +// =================================================================== +// Include original kernel implementations in separate namespaces. +// =================================================================== + +namespace mk_cumsum { +#include "chunk_cumsum.cpp" +} + +namespace mk_kkt { +#include "scaled_dot_kkt.cpp" +} + +namespace mk_solve { +#include "tri_inverse_impl.cpp" +} + +namespace mk_wy { +#include "wy_fast.cpp" +} + +namespace mk_h { +#include "chunk_h.cpp" +} + +namespace mk_o { +#include "chunk_o.cpp" +} + +AICORE inline void mega_solve_tril(__gm__ half *out, __gm__ half *in, __gm__ half *minus_id, uint32_t matrix_size, + uint32_t num_matrices, uint32_t num_bsnd_heads, __gm__ int32_t *cu_seqlens, + uint32_t is_lower) +{ + if (num_matrices <= get_block_num()) + mk_solve::runKernelTriInvRecUnroll(out, in, minus_id, num_matrices, + num_bsnd_heads, cu_seqlens, is_lower); + else if (num_matrices <= 2u * get_block_num()) + mk_solve::runKernelTriInvRecUnroll(out, in, minus_id, num_matrices, + num_bsnd_heads, cu_seqlens, is_lower); + else + mk_solve::runKernelTriInvRecUnroll(out, in, minus_id, num_matrices, + num_bsnd_heads, cu_seqlens, is_lower); +} + +template +AICORE inline void mega_kernel_impl(GM_ADDR q_ptr, GM_ADDR k_ptr, GM_ADDR v_ptr, GM_ADDR g_in_ptr, GM_ADDR beta_ptr, + GM_ADDR msk_lower_ptr, GM_ADDR msk_full_ptr, GM_ADDR minus_id_ptr, + GM_ADDR cu_seqlens_ptr, GM_ADDR o_ptr, GM_ADDR g_sum_ptr, GM_ADDR g_t_ptr, + GM_ADDR beta_t_ptr, GM_ADDR A_ptr, GM_ADDR A_inv_f32_ptr, GM_ADDR A_inv_ptr, + GM_ADDR w_ptr, GM_ADDR u_ptr, GM_ADDR s_ptr, GM_ADDR v_new_ptr, GM_ADDR fs_ptr, + GM_ADDR h0_ptr, int64_t has_initial_state, GM_ADDR kkt_ws_ptr, + GM_ADDR wy_ws_a1_ptr, GM_ADDR wy_ws_a2_ptr, GM_ADDR h_ws_ptr, + GM_ADDR o_ws_qk_ptr, GM_ADDR o_ws_qs_ptr, GM_ADDR o_ws_gated_ptr, + uint32_t num_key_heads, int64_t batch_size, int64_t seq_len, + int64_t total_tokens, uint32_t num_matrices, + GM_ADDR ssm_state_ptr, int64_t ssm_state_index) +{ + constexpr int32_t D = GDN_D; + constexpr int32_t C = GDN_C; + + if (num_key_heads == 0 || (static_cast(H) % num_key_heads) != 0) { + return; + } + + mk_cumsum::cumsum_kernel(reinterpret_cast<__gm__ float *>(g_in_ptr), + reinterpret_cast<__gm__ float *>(g_sum_ptr), + reinterpret_cast<__gm__ int32_t *>(cu_seqlens_ptr), batch_size, seq_len); + +#ifdef MEGA_STOP_AFTER_CUMSUM + pipe_barrier(PIPE_ALL); + return; +#endif + + SyncAllImpl(); + +#ifdef MEGA_STOP_AFTER_SYNC1 + return; +#endif + + mega_transpose_TH_to_HT(reinterpret_cast<__gm__ float *>(g_sum_ptr), + reinterpret_cast<__gm__ float *>(g_t_ptr), total_tokens); + mega_transpose_TH_to_HT(reinterpret_cast<__gm__ half *>(beta_ptr), + reinterpret_cast<__gm__ half *>(beta_t_ptr), total_tokens); + +#ifdef MEGA_STOP_AFTER_TRANSPOSE + pipe_barrier(PIPE_ALL); + return; +#endif + + SyncAllImpl(); + + mk_kkt::kkt_kernel( + reinterpret_cast<__gm__ half *>(k_ptr), reinterpret_cast<__gm__ half *>(beta_t_ptr), + reinterpret_cast<__gm__ float *>(g_t_ptr), reinterpret_cast<__gm__ float *>(msk_lower_ptr), + reinterpret_cast<__gm__ half *>(kkt_ws_ptr), reinterpret_cast<__gm__ half *>(A_ptr), + reinterpret_cast<__gm__ int32_t *>(cu_seqlens_ptr), batch_size, seq_len, total_tokens, num_key_heads); + +#if defined(__DAV_C220_CUBE__) + pipe_barrier(PIPE_ALL); + wait_flag_dev(2); + wait_flag_dev(3); +#endif + +#ifdef MEGA_STOP_AFTER_KKT + pipe_barrier(PIPE_ALL); + return; +#endif + + SyncAllImpl(); + + mega_solve_tril(reinterpret_cast<__gm__ half *>(A_inv_ptr), reinterpret_cast<__gm__ half *>(A_ptr), + reinterpret_cast<__gm__ half *>(minus_id_ptr), C, num_matrices, H, + reinterpret_cast<__gm__ int32_t *>(cu_seqlens_ptr), 1); + +#ifdef MEGA_STOP_AFTER_SOLVE + pipe_barrier(PIPE_ALL); + return; +#endif + + SyncAllImpl(); + +#ifdef MEGA_STOP_AFTER_CAST + pipe_barrier(PIPE_ALL); + return; +#endif + +#ifdef MEGA_STOP_AFTER_SYNC_BEFORE_WY + return; +#endif + + mk_wy::wy_fast_kernel( + reinterpret_cast<__gm__ half *>(k_ptr), reinterpret_cast<__gm__ half *>(v_ptr), + reinterpret_cast<__gm__ half *>(beta_t_ptr), reinterpret_cast<__gm__ float *>(g_t_ptr), + reinterpret_cast<__gm__ half *>(A_inv_ptr), reinterpret_cast<__gm__ half *>(wy_ws_a1_ptr), + reinterpret_cast<__gm__ half *>(wy_ws_a2_ptr), reinterpret_cast<__gm__ half *>(w_ptr), + reinterpret_cast<__gm__ half *>(u_ptr), reinterpret_cast<__gm__ int32_t *>(cu_seqlens_ptr), batch_size, seq_len, + total_tokens, num_key_heads); + +#if defined(__DAV_C220_VEC__) + if (get_block_idx() < num_matrices) { + pipe_barrier(PIPE_ALL); + wait_flag_dev(3); + wait_flag_dev(4); + } +#endif + +#ifdef MEGA_STOP_AFTER_WY + pipe_barrier(PIPE_ALL); + return; +#endif + + SyncAllImpl(); + + mk_h::chunk_h_kernel( + reinterpret_cast<__gm__ half *>(k_ptr), reinterpret_cast<__gm__ half *>(w_ptr), + reinterpret_cast<__gm__ half *>(u_ptr), reinterpret_cast<__gm__ float *>(g_t_ptr), + reinterpret_cast<__gm__ half *>(s_ptr), reinterpret_cast<__gm__ half *>(v_new_ptr), + reinterpret_cast<__gm__ half *>(fs_ptr), reinterpret_cast<__gm__ half *>(h0_ptr), has_initial_state, + reinterpret_cast<__gm__ half *>(h_ws_ptr), reinterpret_cast<__gm__ int32_t *>(cu_seqlens_ptr), batch_size, + seq_len, total_tokens, num_key_heads, reinterpret_cast<__gm__ float *>(ssm_state_ptr), ssm_state_index); + +#ifdef MEGA_STOP_AFTER_H + pipe_barrier(PIPE_ALL); + return; +#endif + + SyncAllImpl(); + + mk_o::chunk_o_kernel( + reinterpret_cast<__gm__ half *>(q_ptr), reinterpret_cast<__gm__ half *>(k_ptr), + reinterpret_cast<__gm__ half *>(v_new_ptr), reinterpret_cast<__gm__ half *>(s_ptr), + reinterpret_cast<__gm__ float *>(g_t_ptr), reinterpret_cast<__gm__ float *>(msk_full_ptr), + reinterpret_cast<__gm__ half *>(o_ws_qk_ptr), reinterpret_cast<__gm__ half *>(o_ws_qs_ptr), + reinterpret_cast<__gm__ half *>(o_ws_gated_ptr), reinterpret_cast<__gm__ half *>(o_ptr), + reinterpret_cast<__gm__ int32_t *>(cu_seqlens_ptr), batch_size, seq_len, total_tokens, num_key_heads); + +#if defined(__DAV_C220_CUBE__) + if (get_block_idx() < num_matrices) { + pipe_barrier(PIPE_ALL); + wait_flag_dev(3); + } +#endif +} + +extern "C" __global__ __aicore__ void +GDN_KERNEL_NAME(GM_ADDR mixed_qkv_ptr, GM_ADDR z_ptr, GM_ADDR b_ptr, GM_ADDR a_input_ptr, + GM_ADDR conv_weight_ptr, GM_ADDR conv_state_ptr, GM_ADDR a_log_ptr, GM_ADDR dt_bias_ptr, + GM_ADDR ssm_state_ptr, GM_ADDR norm_weight_ptr, GM_ADDR msk_lower_ptr, GM_ADDR msk_full_ptr, + GM_ADDR minus_id_ptr, GM_ADDR cu_seqlens_ptr, GM_ADDR packed_qkv_ptr, GM_ADDR g_in_ptr, + GM_ADDR beta_ptr, GM_ADDR initial_state_ptr, GM_ADDR o_ptr, GM_ADDR g_sum_ptr, GM_ADDR g_t_ptr, + GM_ADDR beta_t_ptr, GM_ADDR A_ptr, GM_ADDR A_inv_f32_ptr, GM_ADDR A_inv_ptr, GM_ADDR w_ptr, + GM_ADDR u_ptr, GM_ADDR s_ptr, GM_ADDR v_new_ptr, GM_ADDR fs_ptr, GM_ADDR conv_state_out_ptr, + GM_ADDR ssm_state_out_ptr, GM_ADDR out_ptr, GM_ADDR workspace, GM_ADDR tiling) +{ + KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_MIX_AIC_1_2); + REGISTER_TILING_DEFAULT(Qwen35GdnPrefillSuperOpTilingData); + GET_TILING_DATA_WITH_STRUCT(Qwen35GdnPrefillSuperOpTilingData, tiling_data, tiling); + (void)conv_state_out_ptr; + + RunQwen35PrefillConv(mixed_qkv_ptr, conv_weight_ptr, conv_state_ptr, packed_qkv_ptr, + tiling_data.total_tokens, tiling_data.token_block_size, + tiling_data.token_block_count, tiling_data.conv_state_index); + PrepareQwen35Gate(a_input_ptr, b_ptr, a_log_ptr, dt_bias_ptr, g_in_ptr, beta_ptr, + tiling_data.total_tokens); + ZeroInitialState(initial_state_ptr); + ZeroMegaScratch(A_ptr, A_inv_f32_ptr, A_inv_ptr, s_ptr, fs_ptr, + tiling_data.total_tokens, tiling_data.num_matrices); + SyncAllImpl(); + + GM_ADDR q_ptr = packed_qkv_ptr; + GM_ADDR k_ptr = packed_qkv_ptr + static_cast(tiling_data.total_tokens) * 1024 * sizeof(half); + GM_ADDR v_ptr = packed_qkv_ptr + static_cast(tiling_data.total_tokens) * 2048 * sizeof(half); + GM_ADDR user_ws = AscendC::GetUserWorkspace(workspace); + const uint64_t tile_bytes = static_cast(GDN_C) * GDN_C * sizeof(half); + + GM_ADDR kkt_ws_ptr = user_ws; + GM_ADDR wy_ws_a1_ptr = kkt_ws_ptr + static_cast(tiling_data.block_dim) * 2 * tile_bytes; + GM_ADDR wy_ws_a2_ptr = wy_ws_a1_ptr + static_cast(tiling_data.block_dim) * tile_bytes; + GM_ADDR h_ws_ptr = wy_ws_a2_ptr + static_cast(tiling_data.block_dim) * tile_bytes; + GM_ADDR o_ws_qk_ptr = h_ws_ptr + static_cast(tiling_data.block_dim) * 4 * tile_bytes; + GM_ADDR o_ws_qs_ptr = o_ws_qk_ptr + static_cast(tiling_data.block_dim) * tile_bytes; + GM_ADDR o_ws_gated_ptr = o_ws_qs_ptr + static_cast(tiling_data.block_dim) * tile_bytes; + + mega_kernel_impl<24>( + q_ptr, k_ptr, v_ptr, g_in_ptr, beta_ptr, msk_lower_ptr, msk_full_ptr, minus_id_ptr, + cu_seqlens_ptr, o_ptr, g_sum_ptr, g_t_ptr, beta_t_ptr, A_ptr, A_inv_f32_ptr, + A_inv_ptr, w_ptr, u_ptr, s_ptr, v_new_ptr, fs_ptr, initial_state_ptr, 0, + kkt_ws_ptr, wy_ws_a1_ptr, wy_ws_a2_ptr, h_ws_ptr, o_ws_qk_ptr, o_ws_qs_ptr, + o_ws_gated_ptr, tiling_data.num_key_heads, tiling_data.batch_size, + tiling_data.seq_len, tiling_data.total_tokens, tiling_data.num_matrices, + ssm_state_out_ptr, tiling_data.ssm_state_index); + + SyncAllImpl(); + RunQwen35OutputNorm(o_ptr, z_ptr, norm_weight_ptr, out_ptr, tiling_data.total_tokens); +} + +// The CANN wrapper generated for mixed AIC/AIV kernels calls matmul::clearWorkspace +// after including this source. Keep this include after PTO code so CANN's DYNAMIC +// enum does not collide with pto::DYNAMIC in the kernel templates above. +#include "lib/matmul_intf.h" diff --git a/xllm_ops/qwen35_gdn_prefill_super_op/op_kernel/scaled_dot_kkt.cpp b/xllm_ops/qwen35_gdn_prefill_super_op/op_kernel/scaled_dot_kkt.cpp new file mode 100644 index 0000000..91a82f9 --- /dev/null +++ b/xllm_ops/qwen35_gdn_prefill_super_op/op_kernel/scaled_dot_kkt.cpp @@ -0,0 +1,602 @@ +// ============================================================================ +// scaled_dot_kkt_kernel.cpp — Intra-chunk attention matrix for GatedDeltaNet +// +// Computes A = mask(KK^T · gating_coeff) per chunk, where: +// KK^T ∈ ℝ^{C×C} = K @ K^T (Cube engine, GEMM) +// coeff[i,j] = exp(clamp(g[i]+log(β[i]) - g[j], max=0)) (Vec engine) +// A[i,j] = KK^T[i,j] · coeff[i,j] · causal_mask[i,j] +// +// Inputs: +// K [total_tokens, Hg, D] half — key vectors (BSND along seq; stride Hg * D) +// Beta [H, total_tokens] half — gate bias per **value** head (pre-transposed) +// G [H, total_tokens] float — cumulative gate sum per **value** head +// Msk [C, C] float — lower-triangular causal mask +// +// Output: +// A [total_tokens, H, C] half — gated attention matrix in BSND +// +// Architecture: Cube + Vec cross-core kernel. +// Cube phase: K→L1, GEMM K@K^T→L0C, store to workspace (GM) +// Vec phase: load workspace KK^T, compute gating coefficients, apply mask +// +// Cross-core sync: Cube signals Vec via FFTS flag after each chunk's KK^T +// is written to workspace. Vec signals back when workspace buffer is free. +// Two workspace slots alternate (double-buffering via slot = ci & 1). +// +// Vec sub-blocks: Two sub-blocks (vid=0,1) process upper/lower halves of +// the C×C attention matrix in parallel (HalfChunk rows each). +// +// NPU memory hierarchy: +// GM → L1 (Cube-accessible) → L0A/L0B (GEMM operands) → L0C (accumulator) +// GM → UB (Vec-accessible SRAM) +// +// ── PTO / NPU Primer for This Kernel ────────────────────────────────── +// NPU Architecture (simplified): +// Each "AI Core" (like a GPU SM) has: +// - Cube engine: matrix multiply unit (like GPU Tensor Cores), works on L0A/L0B/L0C +// - Vec engine: SIMD vector unit (like GPU CUDA cores), works on UB (Unified Buffer) +// - MTE2: DMA engine for loading data: GM → L1 or GM → UB +// - MTE3: DMA engine for storing data: UB → GM or L0C → GM +// - MTE1: DMA engine for L1 → L0A/L0B transfers (internal to Cube pipeline) +// Memory hierarchy (fast→slow): L0 registers > L1 cache > UB (SRAM) > GM (HBM) +// Cube and Vec run on SEPARATE cores — they communicate via GM + cross-core flags. +// +// Key PTO APIs used in this kernel (with numpy/torch equivalents): +// TASSIGN(tile, addr) — Bind tile to UB/L1/L0 address (tile = memory[addr]) +// TLOAD(dst, gm_tensor) — DMA load: dst = gm_tensor (async, MTE2 pipe) +// TSTORE(gm, src) — DMA store: gm = src (async, MTE3 pipe) +// TFILLPAD(dst, src) — Zero-fill padding: dst[outside valid] = 0 +// TFILLPAD_INPLACE(d, s) — Same but in-place for UB tiles +// TEXTRACT(l0, l1, r, c) — Copy L1 sub-block → L0A or L0B (MTE1 pipe) +// TRESHAPE(dst, src) — Reinterpret L1 tile layout (NZ↔ZN for transpose) +// TMATMUL(C, A, B) — Matrix multiply: C = A @ B in Cube engine +// TCVT(dst, src, mode) — Type conversion: like dst = src.float() or src.half() +// TMOV(dst, src) — Copy: dst = src.clone() +// TADD(d, a, b) — Element-wise add: d = a + b +// TSUB(d, a, b) — Element-wise subtract: d = a - b +// TMUL(d, a, b) — Element-wise multiply: d = a * b +// TMINS(d, s, val) — Clamp max: d = torch.clamp(s, max=val) +// TEXP(d, s) — Element-wise exp: d = torch.exp(s) +// TLOG(d, s) — Element-wise log: d = torch.log(s) +// TROWEXPAND(2d, col) — Broadcast column → rows: 2d[i,j] = col[i] +// TCOLEXPAND(2d, row) — Broadcast row → cols: 2d[i,j] = row[j] +// set_flag(P1, P2, EVT) — Signal from pipe P1 to pipe P2 (like a semaphore post) +// wait_flag(P1, P2, EVT) — Wait for signal from P1 (like a semaphore wait) +// pipe_barrier(PIPE_V) — Local Vec barrier (ensure all Vec ops complete) +// pipe_barrier(PIPE_ALL) — Barrier for all local pipes +// ffts_cross_core_sync() — Cross-core signal (Cube↔Vec, different physical cores) +// wait_flag_dev(flag) — Wait for cross-core signal +// ============================================================================ + +#include // PTO (Performance Tile Operator): NPU kernel API +#include "acl/acl.h" // ACL (Ascend Computing Language): runtime API +using namespace pto; + +// NumHeads, HiddenSize, and ChunkSize are template parameters. The mega kernel +// dispatches NumHeads from a runtime value to a finite set of specializations. +#ifndef GDN_D +#define GDN_D 128 // D = hidden dimension per head +#endif + +#ifndef GDN_C +#define GDN_C 128 // C = chunk size (tokens processed per chunk) +#endif + +// ── PTO type aliases (device-only, guarded by __CCE_AICORE__) ─────────────── +// These are only compiled for the NPU device compiler (__CCE_AICORE__ is defined +// when compiling for AI Core hardware, similar to __CUDA_ARCH__ in CUDA). +#ifdef __CCE_AICORE__ +// UbND = UB tile in row-major (ND) layout for Vec engine. +// Think of it as: torch.empty((R, C), dtype=T) in on-chip SRAM. +// RV, CV = valid region (for dynamic shapes, like a[:valid_rows, :valid_cols]) +// The Vec engine (SIMD unit) reads/writes these tiles for element-wise ops. +template +using UbND = pto::Tile; + +// UbDN = UB tile in column-major (DN) layout — needed for TROWEXPAND source. +// TROWEXPAND requires its source vector in column-major (transposed) format. +// Same physical memory (UB SRAM), just different indexing convention. +template +using UbDN = pto::Tile; + +// L1Mat = L1 cache tile in NZ fractal format (col-major blocks, row-major within). +// This is the standard input format for the Cube matrix engine. +// Think of it as a matrix in L1 cache ready for GEMM. +// NZ = "Normal-Z": the default fractal layout that Cube expects for left/right operands. +template +using L1Mat = pto::Tile; + +// L1MatZN = L1 tile in ZN fractal format (row-major blocks, col-major within). +// Used when you need to transpose a matrix before GEMM: +// TRESHAPE(l1_zn, l1_nz) reinterprets NZ→ZN layout = logical transpose. +// This is FREE (no data movement) — it just changes how the Cube reads the bits. +template +using L1MatZN = pto::Tile; + +using GmShape2D = pto::Shape<1, 1, 1, pto::DYNAMIC, pto::DYNAMIC>; +using GmStride2D = pto::Stride<1, 1, 1, pto::DYNAMIC, 1>; + +template +using GmTensor2D = pto::GlobalTensor; +#endif + +// ── Main kernel function (runs on each AI core) ────────────────────── +// Template parameters: NumHeads (H value), HiddenSize, ChunkSize. +// GROUP = H/Hg; Cube loads K at head_g = head_idx / GROUP. +// +// __gm__: Marks pointers as Global Memory (HBM) — the NPU equivalent of +// CUDA's device memory. All input/output tensors live in GM. +template +AICORE void kkt_kernel(__gm__ half *K_handle, __gm__ half *Beta_handle, __gm__ float *G_handle, + __gm__ float *Msk_handle, __gm__ half *workspace_handle, __gm__ half *A_handle, + __gm__ int32_t *cu_seqlens, int64_t batch_size, int64_t seq_len, int64_t total_tokens, + uint32_t num_key_heads) +{ + constexpr int32_t HalfChunk = ChunkSize / 2; + constexpr int32_t ChunkSquare = ChunkSize * ChunkSize; + const int32_t key_heads = static_cast(num_key_heads); + if (key_heads <= 0 || (NumHeads % key_heads) != 0) return; + const int32_t group = NumHeads / key_heads; + const int32_t bsnd_qk_stride = key_heads * HiddenSize; + // KTail: number of valid columns in the last 128-wide fractal block of K. + // If HiddenSize is a multiple of 128, the last block is fully used (128). + // Otherwise it's the remainder. Used internally by TLOAD for partial blocks. + constexpr uint32_t KTail = (HiddenSize % 128 == 0) ? 128 : (HiddenSize % 128); + + // ── UB address map (manual memory planning) ───────────────────────── + // The UB is a flat SRAM; we manually assign byte offsets for each tile. + // This is like malloc'ing fixed regions — no dynamic allocator on NPU. + constexpr int32_t GUbAddr = 0; // g_ub: cumulative gates [1×C] + constexpr int32_t BetaHalfUbAddr = 512; // beta_ub_half: gate bias fp16 [1×C/2] + constexpr int32_t BetaUbAddr = 640; // beta_ub: gate bias fp32 [1×C/2] + constexpr int32_t GvUbAddr = 896; // g_v_ub: combined gate+bias [1×C/2] + constexpr int32_t AUbAddr = 1152; // a_ub: attention sub-block fp32 [C/2×C] + constexpr int32_t GRUbAddr = 33920; // g_r_ub: row gates [1×C/2] + constexpr int32_t GCUbAddr = 34176; // g_c_ub: column gates [1×C] + constexpr int32_t MskUbAddr = 34688; // msk_ub: causal mask [C/2×C] + constexpr int32_t GR2dUbAddr = 67456; // g_r_2d_ub: broadcast row gates [C/2×C] + constexpr int32_t GC2dUbAddr = 124800; // g_c_2d_ub: broadcast col gates [C/2×C] + constexpr int32_t CoeffUbAddr = 157568; // coeff_ub: gating coefficient [C/2×C] + // a_ub_half overlaps g_r_2d — safe because they're never live simultaneously + constexpr int32_t AUbHalfAddr = GR2dUbAddr; + + auto cid = get_block_idx(); // Which AI core am I? (like CUDA blockIdx.x) + auto block_num = get_block_num(); // Total AI cores launched (like CUDA gridDim.x) + // ── Vec sub-block parallelism ───────────────────────────────────────── + // Each AI core has 2 Vec sub-blocks (vid=0 and vid=1). + // They share the same UB memory but run independently in parallel. + // Here, vid=0 processes rows [0, C/2) and vid=1 processes rows [C/2, C). + // This halves the per-sub-block work and doubles Vec throughput. + auto vid = get_subblockid(); // 0 or 1: which Vec sub-block am I? + + // Work distribution: each (sequence, head) pair is one "work item". + // AI cores split work round-robin, just like CUDA blocks split a grid. + int64_t num_seqs = batch_size; + int64_t total_work = num_seqs * NumHeads; + + // ── Cube-side tile declarations ───────────────────────────────────── + // Cube-side tiles: K in L1 (NZ format), accumulator in L0C + L1Mat k_l1; + TASSIGN(k_l1, 0); + // TileAcc: L0C accumulator tile for GEMM results. + // The Cube engine always accumulates in float32 for precision, even when + // inputs are fp16. Think of it as: result = torch.matmul(a.half(), b.half()).float() + // When stored to GM via TSTORE with a half GlobalTensor, automatic fp32→fp16 cast occurs. + TileAcc a_l0; + TASSIGN(a_l0, 0); + + // ── Vec-side UB tile declarations ──────────────────────────────────── + // These tiles live in UB (Unified Buffer, the Vec engine's SRAM scratchpad). + // Each TASSIGN binds a tile handle to a fixed UB byte offset (our manual alloc). + // Vec-side UB tiles for gating computation + UbND g_ub; + TASSIGN(g_ub, GUbAddr); + UbND beta_ub_half; + TASSIGN(beta_ub_half, BetaHalfUbAddr); + UbND beta_ub; + TASSIGN(beta_ub, BetaUbAddr); + UbND g_v_ub; + TASSIGN(g_v_ub, GvUbAddr); + UbND a_ub; + TASSIGN(a_ub, AUbAddr); + UbND g_r_ub; + TASSIGN(g_r_ub, GRUbAddr); + UbND g_c_ub; + TASSIGN(g_c_ub, GCUbAddr); + UbND msk_ub; + TASSIGN(msk_ub, MskUbAddr); + UbND g_r_2d_ub; + TASSIGN(g_r_2d_ub, GR2dUbAddr); + UbND g_c_2d_ub; + TASSIGN(g_c_2d_ub, GC2dUbAddr); + UbND coeff_ub; + TASSIGN(coeff_ub, CoeffUbAddr); + UbND a_ub_half; + TASSIGN(a_ub_half, AUbHalfAddr); + + // ======================================================================== + // CUBE PHASE: Compute KK^T = K @ K^T for each chunk via GEMM + // + // ── How GEMM works on NPU (the "Cube pipeline") ────────────────────── + // The matrix multiply pipeline has 3 stages: + // Step 1: TLOAD loads data from GM → L1 (MTE2 pipe) + // Step 2: TEXTRACT copies sub-blocks from L1 → L0A/L0B (MTE1 pipe) + // L0A holds the left operand, L0B holds the right operand + // Step 3: TMATMUL multiplies L0A × L0B → L0C accumulator (M pipe) + // + // For K @ K^T: (numpy: KK_T = K @ K.T) + // Left operand: K [C×D] loaded into L1 in NZ format + // Right operand: K^T — same data, but we TRESHAPE to ZN format + // (TRESHAPE is FREE — it just reinterprets the fractal layout as transposed) + // Result: KK^T [C×C] in L0C (float32 accumulator, even though inputs are fp16) + // ======================================================================== + // __DAV_C220_CUBE__: This code only compiles for the Cube core. + // On NPU, Cube and Vec are separate compilation targets (like two different GPUs). +#if defined(__DAV_C220_CUBE__) + // Outer loop: iterate over all (sequence, head) work items assigned to this core + for (int64_t work_idx = 0; work_idx < (total_work + block_num - 1) / block_num; ++work_idx) { + int64_t pid = work_idx * static_cast(block_num) + static_cast(cid); + if (pid >= total_work) continue; + + // Map linear work index → (sequence, head) pair + int32_t head_idx = static_cast(pid % NumHeads); + int64_t seq_idx = pid / NumHeads; + + // Resolve sequence boundaries: cu_seqlens for variable-length, else fixed stride + int64_t bos, slen; + if (cu_seqlens != nullptr) { + // Variable-length sequences (packed tensor): cu_seqlens = [0, len0, len0+len1, ...] + bos = static_cast(cu_seqlens[seq_idx]); + slen = static_cast(cu_seqlens[seq_idx + 1]) - bos; + } else { + // Fixed-length sequences: each is seq_len tokens starting at seq_idx*seq_len + bos = seq_idx * seq_len; + slen = seq_len; + } + // Ceiling division: how many ChunkSize-sized chunks cover this sequence + int64_t num_chunks = (slen + ChunkSize - 1) / ChunkSize; + + // ── Double-buffering via workspace slots ────────────────────────── + // slot = ci & 1: alternates between 0 and 1 each chunk iteration. + // Cube writes KK^T to workspace[slot], then signals Vec. + // While Vec processes slot[0], Cube can write slot[1] (next chunk). + // This overlaps Cube computation with Vec computation for pipelining. + for (int64_t ci = 0; ci < num_chunks; ++ci) { + int32_t slot = static_cast(ci & 1); + // Wait for Vec to finish reading the previous KK^T from this slot + wait_flag_dev(2 + slot); + pipe_barrier(PIPE_ALL); + + int64_t chunk_start = ci * ChunkSize; + int64_t remaining = slen - chunk_start; + int32_t valid_rows = static_cast(remaining < ChunkSize ? remaining : ChunkSize); + + // BSND key layout [Seq, Hg, D]: token stride Hg * D (see BSND_QK_STRIDE). + // Value head head_idx maps to head_g = head_idx / GROUP for shared K rows. + int32_t head_g = head_idx / group; + int64_t k_offset = ((bos + chunk_start) * static_cast(key_heads) + static_cast(head_g)) * + static_cast(HiddenSize); + + // ── Load K chunk from GM → L1 (MTE2 pipe) ────────────────────── + // DYNAMIC shape: valid_rows may be < ChunkSize for the last chunk. + // GlobalTensor describes the GM layout with strides (BSND interleaved). + // TLOAD triggers the MTE2 DMA engine to copy from GM (HBM) → L1 (on-chip cache). + // If the chunk is partial, TFILLPAD zero-fills the padding region + // so the GEMM doesn't produce garbage from uninitialized memory. + { + L1Mat _l1(valid_rows, HiddenSize); + TASSIGN(_l1, 0); + GmShape2D _gs(valid_rows, HiddenSize); + GmStride2D _stride(bsnd_qk_stride); + GmTensor2D _gm(K_handle + k_offset, _gs, _stride); + TLOAD(_l1, _gm); + if (valid_rows != ChunkSize) TFILLPAD(_l1, _l1); + } + + // ── GEMM: KK^T = K @ K^T (L1→L0A/L0B→L0C) ──────────────────── + // K is [C×D] in L1 NZ; K^T obtained via ZN reshape of same tile. + // + // ── WAR (Write-After-Read) synchronization ──────────────────────── + // Before TEXTRACT (MTE1) writes new data to L0A/L0B, we must ensure: + // 1. MTE2 has finished loading L1 (MTE2→MTE1 sync) + // 2. Cube M pipe has finished reading previous L0A/L0B data (M→MTE1 sync) + // After TEXTRACT, before TMATMUL: + // 3. MTE1→M sync ensures L0A/L0B data is ready for the matrix engine + // After TMATMUL completes: + // 4. M→FIX sync ensures the L0C accumulator can be read + // This is like ensuring a producer-consumer chain is properly ordered. + // WAR sync: MTE2→MTE1, M→MTE1 before extract; MTE1→M before matmul. + { + TileLeft _l0a; + TileRight _l0b; + TASSIGN(_l0a, 0x0); + TASSIGN(_l0b, 0x0); + auto _we = EVENT_ID1; + set_flag(PIPE_MTE2, PIPE_MTE1, _we); + wait_flag(PIPE_MTE2, PIPE_MTE1, _we); + set_flag(PIPE_M, PIPE_MTE1, _we); + wait_flag(PIPE_M, PIPE_MTE1, _we); + // Left operand: K in NZ format, extract directly to L0A + TEXTRACT(_l0a, k_l1, 0, 0); + // Right operand: K^T via ZN reshape of same L1 tile, extract to L0B + L1MatZN _bzn; + TRESHAPE(_bzn, k_l1); + TEXTRACT(_l0b, _bzn, 0, 0); + set_flag(PIPE_MTE1, PIPE_M, _we); + wait_flag(PIPE_MTE1, PIPE_M, _we); + TMATMUL(a_l0, _l0a, _l0b); + set_flag(PIPE_MTE1, PIPE_MTE2, _we); + wait_flag(PIPE_MTE1, PIPE_MTE2, _we); + set_flag(PIPE_M, PIPE_FIX, _we); + wait_flag(PIPE_M, PIPE_FIX, _we); + } + + // ── Store KK^T from L0C → workspace GM (with fp32→fp16 cast) ─── + { + TileAcc _l0(ChunkSize, ChunkSize); + TASSIGN(_l0, 0); + Shape<1, 1, 1, DYNAMIC, DYNAMIC> _gs; + _gs.shape[3] = ChunkSize; + _gs.shape[4] = ChunkSize; + GlobalTensor> _gm( + workspace_handle + (static_cast(cid) * 2 + slot) * ChunkSquare, _gs); + TSTORE(_gm, _l0); + } + + // ── Cross-core synchronization (Cube → Vec) ────────────────────── + // ffts_cross_core_sync(pipe, config): Signal across physical cores. + // Unlike set_flag/wait_flag (which sync pipes within ONE core), this syncs + // between the Cube core and Vec core (they are separate hardware units). + // + // Config encoding: 1 | (mode << 4) | (flag_id << 8) + // mode=2: broadcast to all cores on same block + // flag_id: which flag to set (0,1,2,3...) + // + // The receiving side calls wait_flag_dev(flag_id) to wait for this signal. + // + // In this kernel: + // Cube sets flag 0/1 → Vec waits on wait_flag_dev(0/1) (KK^T ready) + // Vec sets flag 2/3 → Cube waits on wait_flag_dev(2/3) (workspace free) + // + // Signal Vec that this slot's KK^T is ready + ffts_cross_core_sync(PIPE_FIX, 1 | (2 << 4) | (slot << 8)); + } + } +#endif + + // ======================================================================== + // VEC PHASE: Apply gating and causal mask to KK^T + // coeff[i,j] = exp(min(g[i]+log(β[i]) - g[j], 0)) + // A[i,j] = KK^T[i,j] · coeff[i,j] · mask[i,j] + // Each sub-block (vid=0,1) handles HalfChunk rows of the C×C matrix. + // + // ── Gating computation (numpy pseudocode) ───────────────────────────── + // # For each sub-block's C/2 rows (vid selects upper or lower half): + // g_row = g_sum[row_offset:row_offset+C/2] # this sub-block's gates + // g_v = g_row + np.log(beta[row_offset:row_offset+C/2]) # combined gate+bias + // g_col = g_sum[0:C] # full chunk gates + // + // # Broadcast to 2D matrices for element-wise ops: + // g_r_2d = np.tile(g_v.reshape(-1, 1), (1, C)) # TROWEXPAND + // g_c_2d = np.tile(g_col.reshape(1, -1), (C/2, 1)) # TCOLEXPAND + // + // # Gating coefficient: exponential decay, clamped to ≤ 1 + // coeff = np.exp(np.minimum(g_r_2d - g_c_2d, 0)) # TSUB → TMINS → TEXP + // + // # Final: A = KK_T * coeff * causal_mask + // A = KK_T[my_rows] * coeff * mask[my_rows] # TMUL × 2 + // ======================================================================== + // __DAV_C220_VEC__: This code only compiles for the Vec core. +#if defined(__DAV_C220_VEC__) + // set_mask_norm / set_vector_mask: configure the SIMD mask for Vec ops. + // (-1, -1) means "all lanes active" — process every element. + // (Like CUDA's __activemask() returning all 1s for a full warp.) + set_mask_norm(); + set_vector_mask(-1, -1); + + // ── Load causal mask (lower triangular) once, reused across all chunks ── + // vid=0 loads the top half (rows 0..C/2-1), vid=1 loads the bottom half. + // The mask is [C×C] in GM; each sub-block loads its [C/2×C] portion. + { + Shape<1, 1, 1, DYNAMIC, DYNAMIC> _gs; + _gs.shape[3] = HalfChunk; + _gs.shape[4] = ChunkSize; + GlobalTensor> _gm( + Msk_handle + static_cast(vid) * HalfChunk * ChunkSize, _gs); + UbND _ld(HalfChunk, ChunkSize); + TASSIGN(_ld, MskUbAddr); + TLOAD(_ld, _gm); + } + // MTE2→V sync: ensure mask DMA is complete before Vec reads it + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + + // Initial cross-core sync: release both workspace slots so Cube can start. + // Vec tells Cube "slots 0 and 1 are free" by setting flags 2 and 3. + // Without this, Cube would hang on wait_flag_dev(2/3) at the first iteration. + ffts_cross_core_sync(PIPE_MTE3, 1 | (2 << 4) | (2 << 8)); + ffts_cross_core_sync(PIPE_MTE3, 1 | (2 << 4) | (3 << 8)); + + for (int64_t work_idx = 0; work_idx < (total_work + block_num - 1) / block_num; ++work_idx) { + int64_t pid = work_idx * static_cast(block_num) + static_cast(cid); + if (pid >= total_work) continue; + + int32_t head_idx = static_cast(pid % NumHeads); + int64_t seq_idx = pid / NumHeads; + + int64_t bos, slen; + if (cu_seqlens != nullptr) { + bos = static_cast(cu_seqlens[seq_idx]); + slen = static_cast(cu_seqlens[seq_idx + 1]) - bos; + } else { + bos = seq_idx * seq_len; + slen = seq_len; + } + int64_t num_chunks = (slen + ChunkSize - 1) / ChunkSize; + + for (int64_t ci = 0; ci < num_chunks; ++ci) { + int32_t slot = static_cast(ci & 1); + + int64_t chunk_start = ci * ChunkSize; + int64_t remaining = slen - chunk_start; + int32_t valid_rows = static_cast(remaining < ChunkSize ? remaining : ChunkSize); + // row_offset: which half of the C×C matrix this sub-block handles + // vid=0 → rows [0, C/2), vid=1 → rows [C/2, C) + int32_t row_offset = static_cast(vid) * HalfChunk; + // local_valid: how many rows in this sub-block are real (not padding) + // Handles the case where the last chunk has fewer than C valid rows + int32_t local_valid = valid_rows > row_offset + ? (valid_rows - row_offset < HalfChunk ? valid_rows - row_offset : HalfChunk) + : 0; + + if (local_valid > 0) { + // ── Load G (full chunk, 1×C) and Beta (sub-block rows, 1×HalfC) ── + // G is [H, total_tokens] float — contiguous per head + { + Shape<1, 1, 1, DYNAMIC, DYNAMIC> _gs; + _gs.shape[3] = 1; + _gs.shape[4] = valid_rows; + GlobalTensor> _gm( + G_handle + static_cast(head_idx) * total_tokens + (bos + chunk_start), _gs); + UbND _ld(1, valid_rows); + TASSIGN(_ld, GUbAddr); + TLOAD(_ld, _gm); + if (valid_rows != ChunkSize) { + UbND _pd; + TASSIGN(_pd, GUbAddr); + TFILLPAD_INPLACE(_pd, _ld); + } + } + + // Beta is [H, total_tokens] half — contiguous per head + { + Shape<1, 1, 1, DYNAMIC, DYNAMIC> _gs; + _gs.shape[3] = 1; + _gs.shape[4] = local_valid; + GlobalTensor> _gm( + Beta_handle + static_cast(head_idx) * total_tokens + (bos + chunk_start + row_offset), + _gs); + UbND _ld(1, local_valid); + TASSIGN(_ld, BetaHalfUbAddr); + TLOAD(_ld, _gm); + if (local_valid != HalfChunk) { + UbND _pd; + TASSIGN(_pd, BetaHalfUbAddr); + TFILLPAD_INPLACE(_pd, _ld); + } + } + } + + // Wait for Cube to finish writing KK^T for this slot + wait_flag_dev(slot); + pipe_barrier(PIPE_ALL); + + if (local_valid > 0) { + // ── Compute gating coefficient ──────────────────────────────── + // Step 1: Convert beta from fp16→fp32 for precision + // Step 2: g_v[i] = g[row_offset+i] + log(β[i]) — combined row gate + // Step 3: Broadcast g_v (rows) and g (cols) to 2D matrices + // Step 4: coeff = exp(min(g_v_2d - g_2d, 0)) — clamped exponential gating + // g_v[i] = g[row_offset+i] + log(β[i]) — combined row gate + TCVT(beta_ub, beta_ub_half, pto::RoundMode::CAST_NONE); + // g_ub_temp points to the sub-block's portion of g within the full g_ub. + // row_offset * sizeof(float) is the byte offset into the g_ub tile. + UbND g_ub_temp; + TASSIGN(g_ub_temp, GUbAddr + row_offset * static_cast(sizeof(float))); + TMOV(g_v_ub, g_ub_temp); // g_v = g[row_offset:row_offset+C/2] + pipe_barrier(PIPE_V); // Wait for TMOV to complete + + TLOG(beta_ub, beta_ub); // beta_ub = log(beta) in-place + pipe_barrier(PIPE_V); + TADD(g_v_ub, g_v_ub, beta_ub); // g_v = g_sub + log(beta) — the combined gate + pipe_barrier(PIPE_V); + TMOV(g_r_ub, g_v_ub); // Copy to g_r for row-broadcast + TMOV(g_c_ub, g_ub); // Copy full g to g_c for col-broadcast + pipe_barrier(PIPE_V); + + // Broadcast g_v to rows, g to columns → 2D gating matrix + // coeff[i,j] = exp(min(g_v[i] - g[j], 0)) + // + // g_r_ub_temp is a column-major (DN) alias of g_r_ub, required because + // TROWEXPAND expects its source in column-major layout. + UbDN g_r_ub_temp; + TASSIGN(g_r_ub_temp, GRUbAddr); + TROWEXPAND(g_r_2d_ub, g_r_ub_temp); // g_r_2d[i,j] = g_v[i] for all j + TCOLEXPAND(g_c_2d_ub, g_c_ub); // g_c_2d[i,j] = g[j] for all i + pipe_barrier(PIPE_V); + TSUB(coeff_ub, g_r_2d_ub, g_c_2d_ub); // coeff[i,j] = g_v[i] - g[j] + pipe_barrier(PIPE_V); + TMINS(coeff_ub, coeff_ub, 0.0f); // clamp to ≤ 0 (coeff will be ≤ 1 after exp) + pipe_barrier(PIPE_V); + TEXP(coeff_ub, coeff_ub); // coeff = exp(clamped_diff) ∈ (0, 1] + + // V→MTE2 sync: ensure gating computation is done before we start + // loading KK^T from workspace (we need coeff ready for the multiply later, + // and we want to overlap the DMA load with the preceding Vec work). + set_flag(PIPE_V, PIPE_MTE2, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID0); + + // ── Load KK^T sub-block from workspace (fp16) ──────────────── + // workspace layout: [core_id * 2 + slot][C×C], we load our sub-block's + // [C/2×C] portion (offset by vid * HalfChunk * ChunkSize elements). + { + Shape<1, 1, 1, DYNAMIC, DYNAMIC> _gs; + _gs.shape[3] = HalfChunk; + _gs.shape[4] = ChunkSize; + GlobalTensor> _gm( + workspace_handle + (static_cast(cid) * 2 + slot) * ChunkSquare + + static_cast(vid) * HalfChunk * ChunkSize, + _gs); + UbND _ld(HalfChunk, ChunkSize); + TASSIGN(_ld, AUbHalfAddr); + TLOAD(_ld, _gm); + } + + // MTE2→V sync: KK^T data is now in UB, safe for Vec to read + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + + // ── Apply gating and mask: A = KK^T · coeff · mask ─────────── + // 1. Convert KK^T from fp16 → fp32 (Cube stored it as fp16 to save GM bandwidth) + TCVT(a_ub, a_ub_half, pto::RoundMode::CAST_NONE); + // 2. Element-wise multiply by gating coefficient + TMUL(a_ub, a_ub, coeff_ub); + // 3. Element-wise multiply by causal mask (lower triangular, zeros above diagonal) + TMUL(a_ub, a_ub, msk_ub); + // 4. Convert result back to fp16 for output + TCVT(a_ub_half, a_ub, pto::RoundMode::CAST_NONE); + + // V→MTE3 sync: Vec computation done, safe for DMA store to begin + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + + // ── Store A sub-block to output GM ──────────────────────────── + // Output A is in BSND layout: [total_tokens, NumHeads, ChunkSize] + // Each row of A corresponds to one token's attention weights for this head. + // Stride between consecutive tokens = NumHeads * ChunkSize (BSND interleaved). + int64_t a_gm_offset = + ((bos + chunk_start + row_offset) * NumHeads + head_idx) * static_cast(ChunkSize); + + { + Shape<1, 1, 1, DYNAMIC, DYNAMIC> _gs; + _gs.shape[3] = local_valid; + _gs.shape[4] = ChunkSize; + GlobalTensor> _gm( + A_handle + a_gm_offset, _gs); + UbND _st(local_valid, ChunkSize); + TASSIGN(_st, AUbHalfAddr); + TSTORE(_gm, _st); + } + } + + pipe_barrier(PIPE_ALL); + // Signal Cube that this workspace slot is free for reuse. + // Flag (2+slot): slot 0 → flag 2, slot 1 → flag 3. + // Cube is waiting on wait_flag_dev(2+slot) before writing the next chunk. + ffts_cross_core_sync(PIPE_MTE3, 1 | (2 << 4) | ((2 + slot) << 8)); + } + } +#endif +} diff --git a/xllm_ops/qwen35_gdn_prefill_super_op/op_kernel/tri_inverse_impl.cpp b/xllm_ops/qwen35_gdn_prefill_super_op/op_kernel/tri_inverse_impl.cpp new file mode 100644 index 0000000..6ff58fa --- /dev/null +++ b/xllm_ops/qwen35_gdn_prefill_super_op/op_kernel/tri_inverse_impl.cpp @@ -0,0 +1,690 @@ +/** +Copyright (c) 2026 Huawei Technologies Co., Ltd. +All rights reserved. + +See LICENSE in the root of the software repository: +https://github.com/huawei-csl/pto-kernels/ +for the full License text. +*/ + +#ifndef MEMORY_BASE +#define MEMORY_BASE +#endif +#include + +using namespace pto; + +AICORE inline uint32_t CeilDiv(uint32_t value, uint32_t divisor) +{ + return (value + divisor - 1) / divisor; +} + +#define BSND_OFFSET(tile_id, N, S, D) (((tile_id) / (N)) * (S) * (N) * (D) + ((tile_id) % (N)) * (D)) + +/* + * For aligned BSND, tile_id enumerates chunk-major then head-major and maps to + * a fixed-stride address inside the dense BSND tensor. + */ +AICORE inline uint32_t GetBSNDFixedTileOffset(uint32_t tile_id, uint32_t num_bsnd_heads, uint32_t matrix_size) +{ + return BSND_OFFSET(tile_id, num_bsnd_heads, matrix_size, matrix_size); +} + +/** + * @brief Struct containing starting address and size of a single tile + */ +struct BSNDVarlenTileInfo { + uint32_t bsnd_offset; /**< Contains the starting index in the global tensor */ + uint32_t valid_size; /**< This is the size (num_rows/cols) of the tile */ +}; + +/* + * For cu_seqlens-based varlen BSND, tile_id still enumerates chunk-major then + * head-major. We recover the owning sequence by scanning cu_seqlens and + * counting chunks per sequence. + */ +AICORE inline BSNDVarlenTileInfo GetBSNDVarlenTileInfoFromCuSeqlens(uint32_t tile_id, uint32_t num_bsnd_heads, + uint32_t matrix_size, __gm__ int32_t *cu_seqlens) +{ + const uint32_t head_idx = tile_id % num_bsnd_heads; + const uint32_t chunk_idx = tile_id / num_bsnd_heads; + + uint32_t seq_start = static_cast(cu_seqlens[0]); + uint32_t accumulated_chunks = 0; + for (uint32_t seq_idx = 0;; ++seq_idx) { + const uint32_t seq_end = static_cast(cu_seqlens[seq_idx + 1]); + const uint32_t seq_len = seq_end - seq_start; + const uint32_t seq_num_chunks = CeilDiv(seq_len, matrix_size); + if (chunk_idx < accumulated_chunks + seq_num_chunks) { + const uint32_t local_chunk_idx = chunk_idx - accumulated_chunks; + const uint32_t row_start = seq_start + local_chunk_idx * matrix_size; + const uint32_t valid_size = min(static_cast(seq_end - row_start), matrix_size); + return {row_start * num_bsnd_heads * matrix_size + head_idx * matrix_size, valid_size}; + } + accumulated_chunks += seq_num_chunks; + seq_start = seq_end; + } +} + +/* + * @brief: Takes as input two matrices of size MatrixSize * MatrixSize each. + * The src matrix lies in L1, while the dst matrix lies either in L0A or L0B. + * This kernel copies only the diagonal blocks (fractals) of size FractalSize * + * FractalSize from the src matrix to the dst matrix. + * + * @tparam InputT Input data type (fp16). + * @tparam FractalSize Size of each fractal matrix (diagonal block). + * @tparam MatrixSize Size of the entire input/output matrices. + * @tparam SrcL1TileT The actual tile type of the src matrix. + * @tparam DstL0TileT The actual tile type of the dst matrix. + * + * @param src Tile in L1 memory. + * @param dst Tile in L0A or L0B memory. + */ +template +AICORE inline void CopyDiagonalFractalsL1ToL0(SrcL1TileT src, DstL0TileT dst) +{ + constexpr uint32_t NumFractals = MatrixSize / FractalSize; + constexpr bool is_left = std::is_same_v>; + constexpr TileType LeftOrRight = is_left ? TileType::Left : TileType::Right; + constexpr SLayout InnerLayout = is_left ? SLayout::RowMajor : SLayout::ColMajor; + + Tile + fractals[NumFractals]; + const std::uintptr_t starting_address = reinterpret_cast(dst.data()); + for (uint32_t i = 0; i < NumFractals; ++i) { + TASSIGN(fractals[i], starting_address + i * FractalSize * (MatrixSize + FractalSize) * sizeof(InputT)); + TEXTRACT(fractals[i], src, i * FractalSize, i * FractalSize); + } +} + +/* + * @brief: Takes as input two matrices of size MatrixSize * MatrixSize each, + * and an integer block_size. The src matrix lies in L1, while the dst matrix + * either in L0A or L0B. This method copies some of the diagonal blocks from the + * input to the output as follows: + * - If dst is in L0A (left): copy even diagonal blocks 0, 2, 4, ... + * - If dst is in L0B (right): copy odd blocks 1, 3, 5, ... + * Important note: the dst matrix should be initialized to all-zeros before + * calling this method + * + * @tparam InputT Input data type (fp16). + * @tparam FractalSize Size of each fractal matrix (diagonal block). + * @tparam MatrixSize Size of the entire input/output matrices. + * @tparam SrcL1TileT The actual tile type of the src matrix. + * @tparam DstL0TileT The actual tile type of the dst matrix. + * + * @param src Tile in L1 memory. + * @param dst Tile in L0A or L0B memory. + * @param block_size Size of diagonal blocks. Needs: block_size >= FractalSize. + */ +template +AICORE inline void CopyOddOrEvenBlocksL1ToL0(SrcL1TileT src, DstL0TileT dst, uint32_t block_size, + bool swap_parity = false) +{ + constexpr bool is_left = std::is_same_v>; + constexpr TileType LeftOrRight = is_left ? TileType::Left : TileType::Right; + constexpr SLayout InnerLayout = is_left ? SLayout::RowMajor : SLayout::ColMajor; + + // Default: left→even(0), right→odd(1). swap_parity flips this. + const uint32_t starting_block_index = (is_left ? 0u : 1u) ^ (swap_parity ? 1u : 0u); + + const uint32_t num_blocks = MatrixSize / block_size; + const uint32_t num_fractals_per_block = block_size / FractalSize; + + // might need fewer fractals if block_size < FractalSize + Tile + fractals[MatrixSize / FractalSize]; + + const std::uintptr_t starting_address = reinterpret_cast(dst.data()); + for (uint32_t i = 0; i < num_fractals_per_block; ++i) { + for (uint32_t j = 0; j < num_fractals_per_block; ++j) { + for (uint32_t b = starting_block_index; b < num_blocks; b += 2) { + const uint32_t offset = b * (MatrixSize + FractalSize) * block_size /* block_offset */ + + i * MatrixSize * FractalSize /* col_fractal_offset */ + + j * FractalSize * FractalSize /* row_fractal_offset */; + TASSIGN(fractals[b], starting_address + offset * sizeof(InputT)); + TEXTRACT(fractals[b], src, b * block_size + i * FractalSize, b * block_size + j * FractalSize); + } + } + } +} + +/* + * @brief: Prepares Identity and Zeros matrix. + * + * @tparam TileL1AB The type of the input tiles in L1. + * @tparam TileL0A The type of the input tiles in L0A. + * @tparam TileL0B The type of the input tiles in L0B. + * @tparam TileL0C The type of the input tiles in L0C. + * + * @param I_neg_l1_tile Tile containing the -I (negative identity) matrix. + * @param Zero_l1_tile Tile to store the all-zero matrix. + * @param I_l1_tile Tile to store the identity matrix. + * @param a_l0_tile Tile in L0A for matmuls. + * @param b_l0_tile Tile in L0B for matmuls. + * @param c_l0_tile Tile in L0C for matmuls. + */ +template +AICORE inline void PrepareAuxiliaryMatrices(TileL1AB I_neg_l1_tile, TileL1AB Zero_l1_tile, TileL1AB I_l1_tile, + TileL0A a_l0_tile, TileL0B b_l0_tile, TileL0C c_l0_tile) +{ + TMOV(a_l0_tile, I_neg_l1_tile); // a_l0 initialized with I_neg + TMOV(b_l0_tile, I_neg_l1_tile); // b_l0 initialized with I_neg + set_flag(PIPE_MTE1, PIPE_M, static_cast(0)); + wait_flag(PIPE_MTE1, PIPE_M, static_cast(0)); + + TMATMUL(c_l0_tile, a_l0_tile, b_l0_tile); // c_l0 contains I + set_flag(PIPE_M, PIPE_FIX, static_cast(0)); + wait_flag(PIPE_M, PIPE_FIX, static_cast(0)); + + TMOV(I_l1_tile, c_l0_tile); // I_l1 now contains I + set_flag(PIPE_FIX, PIPE_MTE1, static_cast(0)); + wait_flag(PIPE_FIX, PIPE_MTE1, static_cast(0)); + + TMOV(b_l0_tile, I_l1_tile); // b_l0 contains I + set_flag(PIPE_MTE1, PIPE_M, static_cast(0)); + wait_flag(PIPE_MTE1, PIPE_M, static_cast(0)); + + TMATMUL_ACC(c_l0_tile, c_l0_tile, a_l0_tile, + b_l0_tile); // c_l0 contains zeros + set_flag(PIPE_M, PIPE_FIX, static_cast(0)); + wait_flag(PIPE_M, PIPE_FIX, static_cast(0)); + + TMOV(Zero_l1_tile, c_l0_tile); // Zeros_l1 now contains zeros + set_flag(PIPE_FIX, PIPE_MTE1, static_cast(0)); + wait_flag(PIPE_FIX, PIPE_MTE1, static_cast(0)); +} + +/* + * @brief: Inverts a single matrix / tile of the global tensor. + * The first part of the algorithm inverts the FractalSize * FractalSize + * diagonal blocks of the input matrix (inv_trick part). The second phase + * assembles the partial inverses using the cube unig (recursive part). + * + * @tparam InputT The type of the input elements. + * @tparam TileL1AB The type of the input tiles in L1. + * @tparam TileL0A The type of the input tiles in L0A. + * @tparam TileL0B The type of the input tiles in L0B. + * @tparam TileL0C The type of the input tiles in L0C. + * @tparam MatrixSize Size of the entire input/output matrices. + * @tparam FractalSize Size of matrix fractals. + * @tparam NumTilesPerCubeIter How many matrices to load and invert in a single + * cube iteration. + * + * @param X_l1_tile Tile in L1 used for intermediate computations. + * @param I_l1_tile Tile containing the identity matrix. + * @param I_neg_l1_tile Tile containing the negative identity matrix. + * @param M_neg_l1_tile Tile containing the negative input matrix. + * @param Zero_l1_tile Tile containing the all-zero matrix. + * @param Y_l1_tile Tile in L1 used for intermediate computations. + * @param a_l0_tile* Array of two tiles in L0A (for double-buffering). + * @param b_l0_tile* Array of two tiles in L0B (for double-buffering). + * @param c_l0_tile* Tile in L0C for matmuls. + * @param tile_id Index of the current tile (used for sync). + */ +template +AICORE inline void InvertSingleTile(TileL1AB X_l1_tile, TileL1AB I_l1_tile, TileL1AB I_neg_l1_tile, + TileL1AB M_neg_l1_tile, TileL1AB Zero_l1_tile, TileL1AB Y_l1_tile, + TileL0A *a_l0_tile, TileL0B *b_l0_tile, TileL0C *c_l0_tile, const uint32_t tile_id, + const bool swap_parity = false) +{ + const event_t event_0 = static_cast(tile_id); + const event_t event_1 = static_cast(tile_id + NumTilesPerCubeIter); + + TMOV(b_l0_tile[0], Y_l1_tile); // b_l0[0] contains M + TMOV(a_l0_tile[0], I_neg_l1_tile); // a_l0[0] contains I_neg + set_flag(PIPE_MTE1, PIPE_M, event_0); + TMOV(a_l0_tile[1], Zero_l1_tile); + TMOV(b_l0_tile[1], Zero_l1_tile); + set_flag(PIPE_MTE1, PIPE_M, event_1); + wait_flag(PIPE_MTE1, PIPE_M, event_1); + set_flag(PIPE_M, PIPE_MTE1, event_1); + wait_flag(PIPE_M, PIPE_MTE1, event_1); + CopyDiagonalFractalsL1ToL0(Y_l1_tile, a_l0_tile[1]); // a_l0[1] = diag_fractals(M) + CopyDiagonalFractalsL1ToL0(Y_l1_tile, b_l0_tile[1]); // b_l0[1] = diag_fractals(M) + set_flag(PIPE_MTE1, PIPE_M, event_1); + + /* First Matmul: event_0 */ + wait_flag(PIPE_MTE1, PIPE_M, event_0); + TMATMUL(c_l0_tile[0], a_l0_tile[0], b_l0_tile[0]); // c_l0[0] contains M_neg + set_flag(PIPE_M, PIPE_FIX, event_0); + set_flag(PIPE_M, PIPE_MTE1, event_0); + + wait_flag(PIPE_M, PIPE_FIX, event_0); + TMOV(M_neg_l1_tile, c_l0_tile[0]); // M_neg_l1 now contains M_neg + set_flag(PIPE_FIX, PIPE_M, event_0); + + /* Second Matmul: event_1 */ + wait_flag(PIPE_MTE1, PIPE_M, event_1); + set_flag(PIPE_MTE1, PIPE_M, event_1); + TMATMUL(c_l0_tile[1], a_l0_tile[1], + b_l0_tile[1]); // c_l0[1] contains diag_fractals(M)^2 + set_flag(PIPE_M, PIPE_FIX, event_1); + wait_flag(PIPE_M, PIPE_FIX, event_1); + TMOV(Y_l1_tile, + c_l0_tile[1]); // Y_l1 now contains diag_fractals(M)^2 + set_flag(PIPE_FIX, PIPE_M, event_1); + wait_flag(PIPE_FIX, PIPE_M, event_1); + + /* Third Matmul: event_0*/ + wait_flag(PIPE_M, PIPE_MTE1, event_0); + TMOV(b_l0_tile[0], I_neg_l1_tile); // b_l0[0] contains I_neg + TMOV(a_l0_tile[0], I_neg_l1_tile); // a_l0[0] contains I_neg + set_flag(PIPE_MTE1, PIPE_M, event_0); + + wait_flag(PIPE_MTE1, PIPE_M, event_0); + wait_flag(PIPE_FIX, PIPE_M, event_0); + wait_flag(PIPE_MTE1, PIPE_M, event_1); + TMATMUL(c_l0_tile[0], a_l0_tile[1], + b_l0_tile[0]); // c_l0[0] = diag_fractals(M_neg) + set_flag(PIPE_M, PIPE_FIX, event_0); + wait_flag(PIPE_M, PIPE_FIX, event_0); + set_flag(PIPE_FIX, PIPE_M, event_0); + wait_flag(PIPE_FIX, PIPE_M, event_0); + + TMATMUL_ACC(c_l0_tile[0], c_l0_tile[0], a_l0_tile[0], + b_l0_tile[0]); // c_l0[0] has I-diag_fractals(M) + set_flag(PIPE_M, PIPE_FIX, event_1); + wait_flag(PIPE_M, PIPE_FIX, event_1); + TMOV(X_l1_tile, c_l0_tile[0]); // X_l1 now contains I-diag_fractals(M) + + /* + * Inv Trick part: + * X = I - M + * Y = M + * block_size = 1 + * while block_size < FractalSize / 2: + * Y = Y @ Y + * X = X + X @ Y + * block_size *= 2 + */ + set_flag(PIPE_FIX, PIPE_M, event_0); // store c + set_flag(PIPE_M, PIPE_MTE1, event_0); // load matrices for matmuls + set_flag(PIPE_FIX, PIPE_MTE1, event_0); + set_flag(PIPE_FIX, PIPE_M, event_1); // only for update Y + set_flag(PIPE_M, PIPE_MTE1, event_1); // only for update Y + set_flag(PIPE_FIX, PIPE_MTE1, event_1); // only for update Y + for (uint32_t block_size = 1; block_size < FractalSize / 2; block_size *= 2) { + wait_flag(PIPE_M, PIPE_MTE1, event_0); + TMOV(b_l0_tile[0], I_l1_tile); + wait_flag(PIPE_FIX, PIPE_MTE1, event_0); + TMOV(a_l0_tile[0], X_l1_tile); + set_flag(PIPE_MTE1, PIPE_M, event_0); + + wait_flag(PIPE_FIX, PIPE_MTE1, event_1); + TMOV(b_l0_tile[1], Y_l1_tile); + set_flag(PIPE_MTE1, PIPE_M, event_1); + + wait_flag(PIPE_FIX, PIPE_M, event_0); // from previous iter + wait_flag(PIPE_MTE1, PIPE_M, event_0); // from loading a_l0[0], b_l0[0] + TMATMUL(c_l0_tile[0], a_l0_tile[0], b_l0_tile[0]); // c_l0[0] contains X + set_flag(PIPE_M, PIPE_FIX, event_0); + wait_flag(PIPE_M, PIPE_FIX, event_0); + set_flag(PIPE_FIX, PIPE_M, event_0); + wait_flag(PIPE_FIX, PIPE_M, event_0); + + if (block_size < FractalSize / 4) { // Update Y except in last iteration + wait_flag(PIPE_M, PIPE_MTE1, event_1); // from previous iter + TMOV(a_l0_tile[1], Y_l1_tile); + wait_flag(PIPE_MTE1, PIPE_M, event_1); + set_flag(PIPE_MTE1, PIPE_M, event_1); + + wait_flag(PIPE_MTE1, PIPE_M, event_1); + wait_flag(PIPE_FIX, PIPE_M, event_1); // from previous iter + TMATMUL(c_l0_tile[1], a_l0_tile[1], b_l0_tile[1]); + set_flag(PIPE_M, PIPE_MTE1, event_1); // for next iter + set_flag(PIPE_M, PIPE_FIX, event_1); + set_flag(PIPE_MTE1, PIPE_M, event_1); + + wait_flag(PIPE_M, PIPE_FIX, event_1); + TMOV(Y_l1_tile, c_l0_tile[1]); + set_flag(PIPE_FIX, PIPE_M, event_1); // for next iter + } + set_flag(PIPE_FIX, PIPE_MTE1, event_1); // for next iter + + wait_flag(PIPE_MTE1, PIPE_M, event_1); + TMATMUL_ACC(c_l0_tile[0], c_l0_tile[0], a_l0_tile[0], + b_l0_tile[1]); // c_l0[0] has X + X @ Y + set_flag(PIPE_M, PIPE_MTE1, event_0); + set_flag(PIPE_M, PIPE_FIX, event_0); + + wait_flag(PIPE_M, PIPE_FIX, event_0); + TMOV(X_l1_tile, c_l0_tile[0]); + set_flag(PIPE_FIX, PIPE_M, event_0); // for next iter + set_flag(PIPE_FIX, PIPE_MTE1, event_0); // for next iter + } + wait_flag(PIPE_FIX, PIPE_MTE1, event_1); // only for update Y + wait_flag(PIPE_M, PIPE_MTE1, event_1); // only for update Y + wait_flag(PIPE_FIX, PIPE_M, event_1); // only for update Y + wait_flag(PIPE_FIX, PIPE_MTE1, event_0); + wait_flag(PIPE_M, PIPE_MTE1, event_0); + wait_flag(PIPE_FIX, PIPE_M, event_0); + + /* + * Unrolled recursion part: + * Upper-tri (swap_parity=false): + * LX = even_blocks(X), RX = odd_blocks(X) + * Y = LX @ (-M) + I, X = Y @ RX + LX + * Lower-tri (swap_parity=true): + * RX = even→L0A(odd via swap), LX = odd→L0B(even via swap) + * Y = RX @ (-M) + I, X = Y @ LX + RX + */ + TMOV(b_l0_tile[1], M_neg_l1_tile); // b_l0[1] contains M_neg + TMOV(a_l0_tile[0], I_l1_tile); // a_l0[0] contains I + + if constexpr (MatrixSize > FractalSize) { + set_flag(PIPE_FIX, PIPE_M, event_1); + } + set_flag(PIPE_M, PIPE_MTE1, event_1); + set_flag(PIPE_M, PIPE_MTE1, event_0); + set_flag(PIPE_FIX, PIPE_MTE1, event_1); + set_flag(PIPE_FIX, PIPE_M, event_0); + for (uint32_t block_size = FractalSize; block_size < MatrixSize; block_size *= 2) { + wait_flag(PIPE_M, PIPE_MTE1, event_0); // Wait for last iter a_l0[1] + TMOV(a_l0_tile[1], Zero_l1_tile); + + wait_flag(PIPE_M, PIPE_MTE1, event_1); + TMOV(b_l0_tile[0], I_l1_tile); + set_flag(PIPE_MTE1, PIPE_M, event_0); + + wait_flag(PIPE_FIX, PIPE_MTE1, event_1); // Wait to write last X + CopyOddOrEvenBlocksL1ToL0(X_l1_tile, a_l0_tile[1], block_size, + swap_parity); // a_l0[1]: even(LX) or odd(RX) + set_flag(PIPE_MTE1, PIPE_M, event_1); + + wait_flag(PIPE_MTE1, PIPE_M, event_0); + wait_flag(PIPE_FIX, PIPE_M, event_0); // Wait c_l0[0] from previous iter + TMATMUL(c_l0_tile[0], a_l0_tile[0], b_l0_tile[0]); // c_l0[0] has I + + wait_flag(PIPE_MTE1, PIPE_M, event_1); + wait_flag(PIPE_FIX, PIPE_M, event_1); // Wait c_l0[1] from previous iter + TMATMUL(c_l0_tile[1], a_l0_tile[1], b_l0_tile[0]); // c_l0[1] contains LX + set_flag(PIPE_M, PIPE_MTE1, event_1); // allow to load RX on b_l0[0] + + TMATMUL_ACC(c_l0_tile[0], c_l0_tile[0], a_l0_tile[1], + b_l0_tile[1]); // c_l0[0] <- LX * M_neg + I + set_flag(PIPE_M, PIPE_FIX, event_0); + set_flag(PIPE_M, PIPE_MTE1, event_0); + + wait_flag(PIPE_M, PIPE_FIX, event_0); + TMOV(Y_l1_tile, c_l0_tile[0]); // Y_l1 contains LX * M_neg + I + set_flag(PIPE_FIX, PIPE_MTE1, event_0); + set_flag(PIPE_FIX, PIPE_M, event_0); + + /* Load complementary blocks of X in L0B */ + wait_flag(PIPE_M, PIPE_MTE1, event_1); + TMOV(b_l0_tile[0], Zero_l1_tile); + CopyOddOrEvenBlocksL1ToL0(X_l1_tile, b_l0_tile[0], block_size, + swap_parity); // b_l0[0]: odd(RX) or even(LX) + + wait_flag(PIPE_M, PIPE_MTE1, event_0); // Wait for previous use of a_l0[1] + wait_flag(PIPE_FIX, PIPE_MTE1, event_0); // Wait for Y_l1 + TMOV(a_l0_tile[1], Y_l1_tile); // a_l0[1] contains LX * M_neg + I + set_flag(PIPE_MTE1, PIPE_M, event_0); + + wait_flag(PIPE_MTE1, PIPE_M, event_0); + TMATMUL_ACC(c_l0_tile[1], c_l0_tile[1], a_l0_tile[1], b_l0_tile[0]); + set_flag(PIPE_M, PIPE_MTE1, event_0); // next iter can read on a_l0[1] + set_flag(PIPE_M, PIPE_MTE1, event_1); // next iter can read on b_l0[0] + set_flag(PIPE_M, PIPE_FIX, event_0); + wait_flag(PIPE_M, PIPE_FIX, event_0); + + if (block_size < MatrixSize / 2) { // Update X_l1 except in last iteration + TMOV(X_l1_tile, c_l0_tile[1]); + set_flag(PIPE_FIX, PIPE_M, event_1); // release c_l0[1] for next iter + } + set_flag(PIPE_FIX, PIPE_MTE1, event_1); + } + wait_flag(PIPE_M, PIPE_MTE1, event_0); + wait_flag(PIPE_M, PIPE_MTE1, event_1); + wait_flag(PIPE_FIX, PIPE_M, event_0); + wait_flag(PIPE_FIX, PIPE_MTE1, event_1); // Write c_l0[1] to X_l1 +} + +/* + * @brief: Runs the main kernel (inverts all matrices in the tensor) + * + * @tparam InputT The type of the input elements. + * @tparam OutputT The type of the output elements. + * @tparam MatrixSize Size of the entire input/output matrices. + * @tparam NumTilesPerCubeIter How many matrices to load and invert in a single + * cube iteration. + * @tparam IsBSND If IsBSND is false, then the last two dimensions represent a + * 2D triangular matrix in row-major format, while the other dimensions are + * batch dimensions. If IsBSND is true, then the dimensions represent in order: + * B batch size, S sequence length (which is chunked in tiles of size D), N + * number of heads (equivalent to a second batch dimension for this kernel), and + * D chunk size. The inverse is over the dimensions S (chunked) and D, row-major + * within each tile. + * + * @param M_inv pointer to the global memory to store the final inverse. + * @param M Pointer to the global tensor matrix in global memory. + * @param I_neg Pointer to global memory that contains the negative identity. + * @param total_tiles The total number of matrices to invert. + * @param num_bsnd_heads The number of heads, only for BSND format. + */ +template +AICORE inline void TriInvRecUnrollKernel(__gm__ StoreT *M_inv, __gm__ InputT *M, __gm__ InputT *I_neg, + uint32_t total_tiles, uint32_t num_bsnd_heads = 0, + __gm__ int32_t *cu_seqlens = nullptr, uint32_t is_lower = 0) +{ + /* Initializations */ + constexpr uint32_t TileLen = MatrixSize * MatrixSize; + constexpr uint32_t FractalSize = 16; // fractal size for half + constexpr uint32_t NumFractalsRowWise = MatrixSize / FractalSize; + constexpr uint32_t NumL0Buffers = 2; + + if (get_block_idx() * NumTilesPerCubeIter >= total_tiles) { + return; + } + + using GlobalTileShapeIn = TileShape2D; + using GlobalTileStridesIn = + typename std::conditional, + Stride<1, 1, 1, -1, 1>>::type; + using GlobalTileIn = GlobalTensor; + using GlobalTileDynamicShape = Shape<1, 1, 1, DYNAMIC, DYNAMIC>; + using GlobalTileDynamicStride = Stride<1, 1, 1, DYNAMIC, 1>; + using GlobalTileDynamicIn = GlobalTensor; + using GlobalTileStridesINeg = BaseShape2D; + using GlobalTileINeg = GlobalTensor; + + using GlobalTileShapeOut = TileShape2D; + using GlobalTileStridesOut = + typename std::conditional, + Stride<1, 1, 1, -1, 1>>::type; + using GlobalTileOut = GlobalTensor; + using GlobalTileDynamicOut = GlobalTensor; + using TileL1AB = Tile; + using TileL1ABDynamic = Tile; + + // L0 Memory + using TileL0A = TileLeft; + using TileL0B = TileRight; + using TileL0C = TileAcc; + using TileL0CDynamic = TileAcc; + + GlobalTileINeg I_neg_global_in(I_neg); + + TileL1AB X_l1_tile; + TileL1AB I_l1_tile; + TileL1AB I_neg_l1_tile; + TileL1AB M_neg_l1_tile; + TileL1AB Zero_l1_tile; + TileL1AB Y_l1_tile[NumTilesPerCubeIter]; + + TileL0A a_l0_tile[NumL0Buffers]; + TileL0B b_l0_tile[NumL0Buffers]; + TileL0C c_l0_tile[NumL0Buffers]; + + TASSIGN(I_l1_tile, 0x0); + TASSIGN(I_neg_l1_tile, 0x0 + TileLen * sizeof(InputT)); + TASSIGN(Zero_l1_tile, 0x0 + 2 * TileLen * sizeof(InputT)); + TASSIGN(M_neg_l1_tile, 0x0 + 3 * TileLen * sizeof(InputT)); + TASSIGN(X_l1_tile, 0x0 + 4 * TileLen * sizeof(InputT)); + for (uint32_t tile_id = 0; tile_id < NumTilesPerCubeIter; ++tile_id) { + TASSIGN(Y_l1_tile[tile_id], 0x0 + (5 + tile_id) * TileLen * sizeof(InputT)); + } + + for (uint32_t buffer_num = 0; buffer_num < NumL0Buffers; ++buffer_num) { + TASSIGN(a_l0_tile[buffer_num], 0x0 + buffer_num * TileLen * sizeof(InputT)); + TASSIGN(b_l0_tile[buffer_num], 0x0 + buffer_num * TileLen * sizeof(InputT)); + TASSIGN(c_l0_tile[buffer_num], 0x0 + buffer_num * TileLen * sizeof(OutputT)); + } + TLOAD(I_neg_l1_tile, I_neg_global_in); + set_flag(PIPE_MTE2, PIPE_MTE1, static_cast(0)); + wait_flag(PIPE_MTE2, PIPE_MTE1, static_cast(0)); + + PrepareAuxiliaryMatrices(I_neg_l1_tile, Zero_l1_tile, I_l1_tile, a_l0_tile[0], + b_l0_tile[0], c_l0_tile[0]); + + const uint32_t max_iters_per_aic = CeilDiv(total_tiles, (uint32_t)(NumTilesPerCubeIter * get_block_num())); + + /* Main iteration - Compute all tiles */ + uint32_t bsnd_tile_offsets[NumTilesPerCubeIter] = {0}; + uint32_t bsnd_tile_valid_sizes[NumTilesPerCubeIter] = {0}; + uint32_t next_tile_id_that_waits_for_pipe_fix_pipe_m = 0; + set_flag(PIPE_FIX, PIPE_M, static_cast(next_tile_id_that_waits_for_pipe_fix_pipe_m)); + for (uint32_t tile_id = 0; tile_id < NumTilesPerCubeIter; ++tile_id) { + set_flag(PIPE_M, PIPE_MTE2, static_cast(tile_id)); + } + for (uint32_t cube_iter = 0; cube_iter < max_iters_per_aic; ++cube_iter) { + const uint32_t global_index = (cube_iter * get_block_num() + get_block_idx()) * NumTilesPerCubeIter; + if (global_index >= total_tiles) { + break; + } + for (uint32_t tile_id = 0; (tile_id < NumTilesPerCubeIter) && (global_index + tile_id < total_tiles); + ++tile_id) { + if constexpr (IsBSND) { + const uint32_t global_tile_id = global_index + tile_id; + if (cu_seqlens != nullptr) { + const BSNDVarlenTileInfo tile_info = + GetBSNDVarlenTileInfoFromCuSeqlens(global_tile_id, num_bsnd_heads, MatrixSize, cu_seqlens); + bsnd_tile_offsets[tile_id] = tile_info.bsnd_offset; + bsnd_tile_valid_sizes[tile_id] = tile_info.valid_size; + } else { + bsnd_tile_offsets[tile_id] = GetBSNDFixedTileOffset(global_tile_id, num_bsnd_heads, MatrixSize); + bsnd_tile_valid_sizes[tile_id] = MatrixSize; + } + const uint32_t bsnd_offset = bsnd_tile_offsets[tile_id]; + const uint32_t valid_size = bsnd_tile_valid_sizes[tile_id]; + const int row_stride = static_cast(MatrixSize * num_bsnd_heads); + wait_flag(PIPE_M, PIPE_MTE2, static_cast(tile_id)); + if (valid_size < MatrixSize) { + TileL1ABDynamic Y_dyn_l1_tile(valid_size, valid_size); + TASSIGN(Y_dyn_l1_tile, 0x0 + (5 + tile_id) * TileLen * sizeof(InputT)); + GlobalTileDynamicIn M_global_in_dyn( + M + bsnd_offset, {1, 1, 1, static_cast(valid_size), static_cast(valid_size)}, + {1, 1, 1, row_stride, 1}); + TLOAD(Y_dyn_l1_tile, M_global_in_dyn); + set_flag(PIPE_MTE2, PIPE_MTE1, static_cast(tile_id)); + wait_flag(PIPE_MTE2, PIPE_MTE1, static_cast(tile_id)); + TFILLPAD(Y_dyn_l1_tile, Y_dyn_l1_tile); + } else { + GlobalTileIn M_global_in(M + bsnd_offset, {}, {row_stride}); + TLOAD(Y_l1_tile[tile_id], M_global_in); + } + } else { + GlobalTileIn M_global_in(M + (global_index + tile_id) * TileLen); + wait_flag(PIPE_M, PIPE_MTE2, static_cast(tile_id)); + TLOAD(Y_l1_tile[tile_id], + M_global_in); // Copies NumTilesPerCubeIter tiles at once + } + set_flag(PIPE_MTE2, PIPE_MTE1, static_cast(tile_id)); + } + + constexpr uint32_t final_c_buffer_index = MatrixSize > FractalSize ? 1 : 0; + for (uint32_t tile_id = 0; (tile_id < NumTilesPerCubeIter) && (global_index + tile_id < total_tiles); + ++tile_id) { + // Wait for previous cube iter to write result + wait_flag(PIPE_FIX, PIPE_M, static_cast(tile_id)); + // Wait for loading new matrices from GM + wait_flag(PIPE_MTE2, PIPE_MTE1, static_cast(tile_id)); + + InvertSingleTile( + X_l1_tile, I_l1_tile, I_neg_l1_tile, M_neg_l1_tile, Zero_l1_tile, Y_l1_tile[tile_id], a_l0_tile, + b_l0_tile, c_l0_tile, tile_id, is_lower != 0); + + // Allow next cube_iter to proceed for this tile_id + set_flag(PIPE_M, PIPE_MTE2, static_cast(tile_id)); + + /* Store result */ + if constexpr (IsBSND) { + const uint32_t bsnd_offset = bsnd_tile_offsets[tile_id]; + const uint32_t valid_size = bsnd_tile_valid_sizes[tile_id]; + const int row_stride = static_cast(MatrixSize * num_bsnd_heads); + if (valid_size < MatrixSize) { + TileL0CDynamic c_l0_tail_tile(valid_size, valid_size); + TASSIGN(c_l0_tail_tile, 0x0 + final_c_buffer_index * TileLen * sizeof(OutputT)); + GlobalTileDynamicOut M_inv_global_out_dyn( + M_inv + bsnd_offset, {1, 1, 1, static_cast(valid_size), static_cast(valid_size)}, + {1, 1, 1, row_stride, 1}); + TSTORE(M_inv_global_out_dyn, c_l0_tail_tile); + } else { + GlobalTileOut M_inv_global_out(M_inv + bsnd_offset, {}, {row_stride}); + TSTORE(M_inv_global_out, c_l0_tile[final_c_buffer_index]); + } + } else { + GlobalTileOut M_inv_global_out(M_inv + (global_index + tile_id) * TileLen); + TSTORE(M_inv_global_out, c_l0_tile[final_c_buffer_index]); + } + next_tile_id_that_waits_for_pipe_fix_pipe_m = (tile_id + 1) % NumTilesPerCubeIter; + set_flag(PIPE_FIX, PIPE_M, static_cast(next_tile_id_that_waits_for_pipe_fix_pipe_m)); + } + } + for (uint32_t tile_id = 0; tile_id < NumTilesPerCubeIter; ++tile_id) { + wait_flag(PIPE_M, PIPE_MTE2, static_cast(tile_id)); + } + wait_flag(PIPE_FIX, PIPE_M, static_cast(next_tile_id_that_waits_for_pipe_fix_pipe_m)); +} + +/* + * @brief: Computes the inverses of the blocks of tensor M + */ +template +AICORE void runKernelTriInvRecUnroll(__gm__ StoreT *M_inv, __gm__ InputT *M, __gm__ InputT *I_neg, uint32_t total_tiles, + uint32_t num_bsnd_heads = 0, __gm__ int32_t *cu_seqlens = nullptr, + uint32_t is_lower = 0) +{ +#if (__CHECK_FEATURE_AT_PRECOMPILE) || (__CCE_AICORE__ == 220 && defined(__DAV_C220_CUBE__)) // Cube compilation + + TriInvRecUnrollKernel( + M_inv, M, I_neg, total_tiles, num_bsnd_heads, cu_seqlens, is_lower); +#else +// Nothing to do on AIV +#endif +} + +template +AICORE void run_tri_inv_rec_unroll(__gm__ float *tensor_out, __gm__ InputT *tensor_in, __gm__ InputT *minus_identity_in, + uint32_t matrix_size, uint32_t num_matrices, uint32_t num_bsnd_heads, + __gm__ int32_t *cu_seqlens = nullptr, uint32_t is_lower = 0) +{ + static_assert(std::is_same_v, "tri_inv_rec_unroll supports only fp16."); + switch (matrix_size) { + case 16: + runKernelTriInvRecUnroll( + tensor_out, tensor_in, minus_identity_in, num_matrices, num_bsnd_heads, cu_seqlens, is_lower); + break; + case 32: + runKernelTriInvRecUnroll( + tensor_out, tensor_in, minus_identity_in, num_matrices, num_bsnd_heads, cu_seqlens, is_lower); + break; + case 64: + runKernelTriInvRecUnroll( + tensor_out, tensor_in, minus_identity_in, num_matrices, num_bsnd_heads, cu_seqlens, is_lower); + break; + case 128: + runKernelTriInvRecUnroll( + tensor_out, tensor_in, minus_identity_in, num_matrices, num_bsnd_heads, cu_seqlens, is_lower); + break; + } +} diff --git a/xllm_ops/qwen35_gdn_prefill_super_op/op_kernel/wy_fast.cpp b/xllm_ops/qwen35_gdn_prefill_super_op/op_kernel/wy_fast.cpp new file mode 100644 index 0000000..530f85a --- /dev/null +++ b/xllm_ops/qwen35_gdn_prefill_super_op/op_kernel/wy_fast.cpp @@ -0,0 +1,857 @@ +// ============================================================================ +// wy_fast_kernel.cpp — WY representation for GatedDeltaNet chunk recurrence +// +// Computes the WY update matrices U and W for each chunk of C tokens: +// U = A2 @ V where A2 = A * beta_2d (beta-scaled attention) +// W = A1 @ K where A1 = A * (exp(g)*beta)_2d (gate+beta-scaled attention) +// +// beta is the decay factor, g is the gate value, A is the triangular attention +// matrix (from the kkt kernel). The column-broadcast notation x_2d means +// expanding a 1xC vector into a C/2 x C matrix by replicating across rows. +// +// Architecture: Vec+Cube cooperative kernel using cross-core synchronization. +// +// Vec core (two sub-blocks for upper/lower C/2 rows): +// For each chunk: +// 1. Load beta [H,T] and A [B,S,H,C], compute A2 = A * beta_2d -> ws +// 2. Load G [H,T], compute A1 = A * (exp(g)*beta)_2d -> ws +// 3. Signal Cube via cross-core flags when workspaces are ready +// +// Cube core (waits for Vec signals): +// For each chunk: +// 1. Load K, V from BSND layout into L1 +// 2. Load A2 from workspace -> GEMM: U = A2 @ V +// 3. Load A1 from workspace -> GEMM: W = A1 @ K +// 4. Store U, W back to BSND layout +// +// NPU memory hierarchy used: +// GM -> UB (Vec), GM -> L1 -> L0A/L0B -> L0C -> GM (Cube) +// +// ── PTO / NPU Primer ────────────────────────────────────────────────── +// This kernel uses BOTH the Cube engine (matrix multiply) and Vec engine +// (SIMD element-wise ops), running on SEPARATE physical cores that +// communicate via Global Memory (GM) + cross-core flags (FFTS). +// +// Execution flow: +// Vec core: load A,beta,G → compute A2,A1 → store to GM workspace +// Cube core: wait for workspace → load A2/A1 + K/V → GEMM → store U,W +// +// Key PTO APIs (with numpy/torch equivalents): +// TLOAD(ub_tile, gm) — ub_tile = gm[...] (DMA: GM→UB, async MTE2) +// TSTORE(gm, ub_tile) — gm[...] = ub_tile (DMA: UB→GM, async MTE3) +// TCVT(dst, src, mode) — dst = src.float() or .half() (type conversion) +// TMOV(dst, src) — dst = src.clone() +// TMUL(d, a, b) — d = a * b (element-wise) +// TEXP(d, s) — d = torch.exp(s) +// TCOLEXPAND(2d, row) — 2d[i,j] = row[j] (broadcast row across all rows) +// TEXTRACT(l0, l1, r, c) — L1 sub-block → L0A/L0B (MTE1 for Cube GEMM) +// TMATMUL(C, A, B) — C = A @ B in Cube engine (fp16→fp32 accumulate) +// set_flag / wait_flag — sync between pipes on SAME core +// ffts_cross_core_sync — signal ACROSS Cube↔Vec cores +// wait_flag_dev(flag) — wait for cross-core signal +// ============================================================================ + +#include +#include "acl/acl.h" +#include +using namespace pto; + +#ifndef GDN_D +#define GDN_D 128 +#endif + +#ifndef GDN_C +#define GDN_C 128 +#endif + +#ifdef __CCE_AICORE__ + +namespace { + +template +using TileMatL1 = pto::Tile; + +template +using TileMatL1ZN = pto::Tile; + +template +using TileMatL0A = pto::Tile; + +template +using TileMatL0B = pto::Tile; + +template +using TileUbDataND = pto::Tile; + +template +using TileUbDataDN = pto::Tile; + +using GmShape2D = pto::Shape<1, 1, 1, pto::DYNAMIC, pto::DYNAMIC>; +using GmStride2D = pto::Stride<1, 1, 1, pto::DYNAMIC, 1>; + +template +using GmTensor2D = pto::GlobalTensor; + +template +using DynMatL1 = pto::Tile; + +template +using DynVecTile = pto::Tile; + +template +using DynAccTile = pto::TileAcc; + +// PTO cheat sheet for readers coming from PyTorch / NumPy: +// - `GlobalTensor` is a GM tensor view with explicit shape/stride metadata. +// - `Tile<..., Mat, ...>` is an on-chip matrix tile used by Cube kernels. +// - `Tile<..., Vec, ...>` is an on-chip UB tile used by SIMD vector kernels. +// - `TileAcc` is the matmul accumulator tile. +// - `TLOAD` / `TSTORE` are DMA copies between GM and local memory. +// - `TCOLEXPAND` is broadcast like `x[None, :].expand(rows, -1)`. +// - `TMUL`, `TEXP`, `TCVT` are vector ops on UB tiles. + +template +AICORE PTO_INLINE void +gemm_v0(std::conditional_t, TileMatL1> &A, + std::conditional_t, TileMatL1> &B, + pto::TileAcc &C, bool clear) +{ + // Local K-sliced matmul helper: + // C = A @ B + // PTO exposes the L1 -> L0 -> Cube movement explicitly, so keeping this tiny + // helper local lets readers see the schedule without hiding it in a repo-wide + // wrapper layer. + // + // PyTorch mental model: + // C = 0 + // for k0 in range(0, K, kL0Size): + // C += A[:, k0:k1] @ B[k0:k1, :] + constexpr uint32_t kL0Size = 128; + const uint32_t kL0split = (K + kL0Size - 1) / kL0Size; + + auto war_event_id = (event_t)(((int)EVENT_ID0 + 1) % 8); + set_flag(PIPE_MTE2, PIPE_MTE1, war_event_id); + wait_flag(PIPE_MTE2, PIPE_MTE1, war_event_id); + + for (uint32_t kL0Idx = 0; kL0Idx < kL0split; ++kL0Idx) { + const bool initflag = clear && (kL0Idx == 0); + const bool is_tail_block = (kL0Idx == kL0split - 1); + + if (is_tail_block) { + TileMatL0A l0a; + TileMatL0B l0b; + pto::TASSIGN(l0a, 0x0); + pto::TASSIGN(l0b, 0x0); + + set_flag(PIPE_M, PIPE_MTE1, war_event_id); + wait_flag(PIPE_M, PIPE_MTE1, war_event_id); + + if constexpr (!transpose_A) { + pto::TEXTRACT(l0a, A, 0, kL0Idx * K_tail); + } else { + TileMatL1ZN A_t; + pto::TRESHAPE(A_t, A); + pto::TEXTRACT(l0a, A_t, 0, kL0Idx * K_tail); + } + + if constexpr (!transpose_B) { + pto::TEXTRACT(l0b, B, kL0Idx * K_tail, 0); + } else { + TileMatL1ZN B_t; + pto::TRESHAPE(B_t, B); + pto::TEXTRACT(l0b, B_t, kL0Idx * K_tail, 0); + } + + set_flag(PIPE_MTE1, PIPE_M, war_event_id); + wait_flag(PIPE_MTE1, PIPE_M, war_event_id); + + if (initflag) { + pto::TMATMUL(C, l0a, l0b); + } else { + pto::TMATMUL_ACC(C, C, l0a, l0b); + } + } else { + TileMatL0A l0a; + TileMatL0B l0b; + pto::TASSIGN(l0a, 0x0); + pto::TASSIGN(l0b, 0x0); + + set_flag(PIPE_M, PIPE_MTE1, war_event_id); + wait_flag(PIPE_M, PIPE_MTE1, war_event_id); + + set_flag(PIPE_FIX, PIPE_M, war_event_id); + wait_flag(PIPE_FIX, PIPE_M, war_event_id); + + if constexpr (!transpose_A) { + pto::TEXTRACT(l0a, A, 0, kL0Idx * kL0Size); + } else { + TileMatL1ZN A_t; + pto::TRESHAPE(A_t, A); + pto::TEXTRACT(l0a, A_t, 0, kL0Idx * kL0Size); + } + + if constexpr (!transpose_B) { + pto::TEXTRACT(l0b, B, kL0Idx * kL0Size, 0); + } else { + TileMatL1ZN B_t; + pto::TRESHAPE(B_t, B); + pto::TEXTRACT(l0b, B_t, kL0Idx * kL0Size, 0); + } + + set_flag(PIPE_MTE1, PIPE_M, war_event_id); + wait_flag(PIPE_MTE1, PIPE_M, war_event_id); + + if (initflag) { + pto::TMATMUL(C, l0a, l0b); + } else { + pto::TMATMUL_ACC(C, C, l0a, l0b); + } + + set_flag(PIPE_MTE1, PIPE_MTE2, war_event_id); + wait_flag(PIPE_MTE1, PIPE_MTE2, war_event_id); + } + } + + set_flag(PIPE_MTE1, PIPE_MTE2, war_event_id); + wait_flag(PIPE_MTE1, PIPE_MTE2, war_event_id); + + set_flag(PIPE_M, PIPE_FIX, war_event_id); + wait_flag(PIPE_M, PIPE_FIX, war_event_id); +} + +} // namespace + +#endif + +template +AICORE void wy_fast_kernel(__gm__ half *K_handle, __gm__ half *V_handle, __gm__ half *Beta_handle, + __gm__ float *G_handle, __gm__ half *A_handle, __gm__ half *workspace_a1_handle, + __gm__ half *workspace_a2_handle, __gm__ half *W_handle, __gm__ half *U_handle, + __gm__ int32_t *cu_seqlens, int64_t batch_size, int64_t seq_len, int64_t total_tokens, + uint32_t num_key_heads) +{ + // WY recompute materializes two diagonal reweightings of the same A tile: + // A2[:, j] = A[:, j] * beta_j + // A1[:, j] = A[:, j] * exp(g_j) * beta_j + // and then forms the two branch outputs + // U = A2 @ V, W = A1 @ K. + // + // Shapes for one (sequence, head, chunk): + // A_chunk : [valid, valid] + // beta : [valid] + // g : [valid] + // K, V : [valid, D] + // + // PyTorch / NumPy sketch: + // A2 = A_chunk * beta[None, :] + // A1 = A_chunk * (exp(g) * beta)[None, :] + // U = A2 @ V_chunk + // W = A1 @ K_chunk + // + // PTO split: + // Vec builds the two reweighted A tiles in workspace. + // Cube later consumes those workspaces in two GEMMs. + constexpr int32_t HalfChunk = ChunkSize / 2; + constexpr uint32_t KTail = (HiddenSize % 128 == 0) ? 128 : (HiddenSize % 128); + + constexpr int32_t H = NumHeads; + const int32_t Hg = static_cast(num_key_heads); + if (Hg <= 0 || (H % Hg) != 0) return; + const int32_t GROUP = H / Hg; + constexpr int32_t BSND_V_STRIDE = H * HiddenSize; + const int32_t BSND_QK_STRIDE = Hg * HiddenSize; + + constexpr int32_t GHeadTileCols = ((NumHeads + 7) / 8) * 8; + constexpr int32_t BetaHeadTileCols = ((NumHeads + 15) / 16) * 16; + + constexpr int32_t BetaHalfUbAddr = 0; + constexpr int32_t A1HalfUbAddr = 256; + constexpr int32_t BetaUbAddr = 16640; + constexpr int32_t BetaRUbAddr = 17152; + constexpr int32_t Beta2dUbAddr = 17664; + constexpr int32_t TmpUbAddr = 50432; + constexpr int32_t A1UbAddr = 75008; + constexpr int32_t A2UbAddr = 107776; + constexpr int32_t A2HalfUbAddr = 140544; + constexpr int32_t GUbAddr = 156928; + constexpr int32_t GRUbAddr = 157440; + constexpr int32_t G2dUbAddr = 157952; + + constexpr int32_t GBlockUbAddr = TmpUbAddr; + constexpr int32_t BetaBlockUbAddr = TmpUbAddr; + + constexpr int32_t WsA1Size = ChunkSize * ChunkSize; + constexpr int32_t WsA2Size = ChunkSize * ChunkSize; + + auto cid = get_block_idx(); + auto block_num = get_block_num(); + auto vid = get_subblockid(); + + int64_t num_seqs = batch_size; + + TileUbDataND beta_ub_half; + TASSIGN(beta_ub_half, BetaHalfUbAddr); + TileUbDataND a1_ub_half; + TASSIGN(a1_ub_half, A1HalfUbAddr); + TileUbDataND beta_ub; + TASSIGN(beta_ub, BetaUbAddr); + TileUbDataND beta_r_ub; + TASSIGN(beta_r_ub, BetaRUbAddr); + TileUbDataND beta_2d_ub; + TASSIGN(beta_2d_ub, Beta2dUbAddr); + TileUbDataND tmp_ub; + TASSIGN(tmp_ub, TmpUbAddr); + TileUbDataND a1_ub; + TASSIGN(a1_ub, A1UbAddr); + TileUbDataND a2_ub; + TASSIGN(a2_ub, A2UbAddr); + TileUbDataND a2_ub_half; + TASSIGN(a2_ub_half, A2HalfUbAddr); + TileUbDataND g_ub; + TASSIGN(g_ub, GUbAddr); + TileUbDataND g_r_ub; + TASSIGN(g_r_ub, GRUbAddr); + TileUbDataND g_2d_ub; + TASSIGN(g_2d_ub, G2dUbAddr); + + TileMatL1 k_l1; + TASSIGN(k_l1, 0); + TileMatL1 v_l1; + TASSIGN(v_l1, 32768); + TileMatL1 a2_l1; + TASSIGN(a2_l1, 65536); + TileAcc u_l0; + TASSIGN(u_l0, 0); + TileMatL1 a1_l1; + TASSIGN(a1_l1, 98304); + TileAcc w_l0; + TASSIGN(w_l0, 65536); + + int64_t total_work = 0; + if (cu_seqlens == nullptr) { + int64_t chunks_per_seq = (seq_len + ChunkSize - 1) / ChunkSize; + total_work = num_seqs * chunks_per_seq * NumHeads; + } + +#if defined(__DAV_C220_VEC__) + set_mask_norm(); + set_vector_mask(-1, -1); + + // Vec prepares the two reweighted A workspaces (`A2` and `A1`) that the + // Cube phase consumes later. + if (cu_seqlens == nullptr) { + bool first_iter = true; + int64_t gi = 0; + for (int64_t seq_idx = 0; seq_idx < num_seqs; ++seq_idx) { + int64_t bos = seq_idx * seq_len; + int64_t slen = seq_len; + int64_t nc = (slen + ChunkSize - 1) / ChunkSize; + + for (int64_t ci = 0; ci < nc; ++ci) { + for (int32_t head_idx = 0; head_idx < NumHeads; ++head_idx) { + if (gi % static_cast(block_num) == static_cast(cid)) { + int64_t chunk_start = ci * ChunkSize; + int64_t remaining = slen - chunk_start; + int32_t valid_rows = static_cast(remaining < ChunkSize ? remaining : ChunkSize); + int64_t chunk_token_start = bos + chunk_start; + // Each Vec sub-block owns one HalfChunk-row stripe of the chunk. + // For a tail chunk, the upper stripe (vid=0) may hold fewer than + // 64 rows, and the lower stripe (vid=1) may hold only a suffix or + // no rows at all. `local_rows` is the exact number of live rows in + // THIS sub-block's stripe. + int32_t local_rows = valid_rows - static_cast(vid) * HalfChunk; + if (local_rows < 0) local_rows = 0; + if (local_rows > HalfChunk) local_rows = HalfChunk; + + // Beta is pre-transposed to [H, total_tokens] for contiguous loads. + { + GmShape2D beta_shape(1, valid_rows); + GmStride2D beta_stride(1); + GmTensor2D beta_global( + Beta_handle + static_cast(head_idx) * total_tokens + chunk_token_start, + beta_shape, beta_stride); + DynVecTile beta_load(1, valid_rows); + TASSIGN(beta_load, BetaHalfUbAddr); + TLOAD(beta_load, beta_global); + if (valid_rows != ChunkSize) { + TFILLPAD_INPLACE(beta_ub_half, beta_load); + } + } + + // Load only the live rows for this sub-block, then zero-pad the + // remainder of the HalfChunk tile. The Cube phase always consumes + // a full [HalfChunk, ChunkSize] workspace tile, so stale rows here + // would leak garbage into ragged tails and cross-sequence boundaries. + if (local_rows > 0) { + int64_t a_gm_offset = + ((chunk_token_start + static_cast(vid) * HalfChunk) * NumHeads + head_idx) * + static_cast(ChunkSize); + GmShape2D a_shape(local_rows, ChunkSize); + GmStride2D a_stride(NumHeads * ChunkSize); + GmTensor2D a_global(A_handle + a_gm_offset, a_shape, a_stride); + DynVecTile a_load(local_rows, ChunkSize); + TASSIGN(a_load, A1HalfUbAddr); + TLOAD(a_load, a_global); + if (local_rows != HalfChunk) { + TFILLPAD_INPLACE(a1_ub_half, a_load); + } + } else { + // Fully empty lower-half tail: materialize an all-zero tile so the + // workspace still looks like a correctly padded HalfChunk block. + TEXPANDS(a1_ub, 0.0f); + pipe_barrier(PIPE_V); + TCVT(a1_ub_half, a1_ub, pto::RoundMode::CAST_NONE); + } + + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + + TCVT(beta_ub, beta_ub_half, pto::RoundMode::CAST_NONE); + pipe_barrier(PIPE_V); + TMOV(beta_r_ub, beta_ub); + pipe_barrier(PIPE_V); + // Replicate beta_j across rows so every column j of A gets the same beta. + // PyTorch-like: + // beta_2d = beta[None, :].expand(HalfChunk, ChunkSize) + TCOLEXPAND(beta_2d_ub, beta_r_ub); + + TCVT(a1_ub, a1_ub_half, pto::RoundMode::CAST_NONE); + // Form the beta-scaled tile that the later U = A2 * V matmul consumes. + // a2_ub = a1_ub * beta_2d_ub + TMUL(a2_ub, a1_ub, beta_2d_ub); + TCVT(a2_ub_half, a2_ub, pto::RoundMode::CAST_NONE); + + if (!first_iter) wait_flag_dev(3); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + { + GmShape2D a2_shape(HalfChunk, ChunkSize); + GmStride2D a2_stride(ChunkSize); + GmTensor2D workspace_a2_global(workspace_a2_handle + + static_cast(cid) * WsA2Size + + static_cast(vid) * HalfChunk * ChunkSize, + a2_shape, a2_stride); + TSTORE(workspace_a2_global, a2_ub_half); + } + pipe_barrier(PIPE_ALL); + ffts_cross_core_sync(PIPE_MTE3, 1 | (2 << 4) | (2 << 8)); + + // G is pre-transposed to [H, total_tokens] for contiguous loads. + { + GmShape2D g_shape(1, valid_rows); + GmStride2D g_stride(1); + GmTensor2D g_global( + G_handle + static_cast(head_idx) * total_tokens + chunk_token_start, g_shape, + g_stride); + DynVecTile g_load(1, valid_rows); + TASSIGN(g_load, GUbAddr); + TLOAD(g_load, g_global); + if (valid_rows != ChunkSize) { + TFILLPAD_INPLACE(g_ub, g_load); + } + } + + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + + // Build the g-based column weights before forming the W = A1 * K branch. + // Torch-like: + // g_weight = exp(g) * beta + TEXP(g_ub, g_ub); + pipe_barrier(PIPE_V); + TMUL(g_ub, g_ub, beta_ub); + pipe_barrier(PIPE_V); + TMOV(g_r_ub, g_ub); + pipe_barrier(PIPE_V); + TCOLEXPAND(g_2d_ub, g_r_ub); + // A1 keeps the same A columns but multiplies each one by exp(g_j) * beta_j. + // a1_ub = a1_ub * g_weight[None, :] + TMUL(a1_ub, a1_ub, g_2d_ub); + TCVT(a1_ub_half, a1_ub, pto::RoundMode::CAST_NONE); + + if (!first_iter) wait_flag_dev(4); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + { + GmShape2D a1_shape(HalfChunk, ChunkSize); + GmStride2D a1_stride(ChunkSize); + GmTensor2D workspace_a1_global(workspace_a1_handle + + static_cast(cid) * WsA1Size + + static_cast(vid) * HalfChunk * ChunkSize, + a1_shape, a1_stride); + TSTORE(workspace_a1_global, a1_ub_half); + } + pipe_barrier(PIPE_ALL); + ffts_cross_core_sync(PIPE_MTE3, 1 | (2 << 4) | (1 << 8)); + first_iter = false; + } + gi++; + } + } + } + } else { + // Same WY math as above; only the work enumeration changes for varlen input. + int64_t gi = 0; + bool first_iter_v = true; + for (int64_t si = 0; si < num_seqs; ++si) { + int64_t bos = static_cast(cu_seqlens[si]); + int64_t eos = static_cast(cu_seqlens[si + 1]); + int64_t slen = eos - bos; + int64_t nc = (slen + ChunkSize - 1) / ChunkSize; + + for (int64_t ci = 0; ci < nc; ++ci) { + for (int32_t h = 0; h < NumHeads; ++h) { + if (gi % static_cast(block_num) == static_cast(cid)) { + int64_t chunk_start = ci * ChunkSize; + int64_t remaining = slen - chunk_start; + int32_t valid_rows = static_cast(remaining < ChunkSize ? remaining : ChunkSize); + int64_t chunk_token_start = bos + chunk_start; + // Same HalfChunk ownership rule as the fixed-length path above: + // each Vec sub-block handles one 64-row stripe, and ragged varlen + // tails may leave that stripe partially full or fully empty. + int32_t local_rows = valid_rows - static_cast(vid) * HalfChunk; + if (local_rows < 0) local_rows = 0; + if (local_rows > HalfChunk) local_rows = HalfChunk; + int32_t head_idx = h; + + // Beta is pre-transposed to [H, total_tokens] for contiguous loads. + { + GmShape2D beta_shape(1, valid_rows); + GmStride2D beta_stride(1); + GmTensor2D beta_global( + Beta_handle + static_cast(head_idx) * total_tokens + chunk_token_start, + beta_shape, beta_stride); + DynVecTile beta_load(1, valid_rows); + TASSIGN(beta_load, BetaHalfUbAddr); + TLOAD(beta_load, beta_global); + if (valid_rows != ChunkSize) { + TFILLPAD_INPLACE(beta_ub_half, beta_load); + } + } + + // Tail-safe A loading is especially important in varlen mode because + // the final chunk of one sequence may be immediately followed by the + // first chunk of the next sequence in packed storage. + if (local_rows > 0) { + int64_t a_gm_offset = + ((chunk_token_start + static_cast(vid) * HalfChunk) * NumHeads + head_idx) * + static_cast(ChunkSize); + GmShape2D a_shape(local_rows, ChunkSize); + GmStride2D a_stride(NumHeads * ChunkSize); + GmTensor2D a_global(A_handle + a_gm_offset, a_shape, a_stride); + DynVecTile a_load(local_rows, ChunkSize); + TASSIGN(a_load, A1HalfUbAddr); + TLOAD(a_load, a_global); + if (local_rows != HalfChunk) { + TFILLPAD_INPLACE(a1_ub_half, a_load); + } + } else { + // Empty stripe for this sub-block: write zeros so the downstream + // full-tile Cube GEMM sees valid padding rather than old workspace. + TEXPANDS(a1_ub, 0.0f); + pipe_barrier(PIPE_V); + TCVT(a1_ub_half, a1_ub, pto::RoundMode::CAST_NONE); + } + + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + + TCVT(beta_ub, beta_ub_half, pto::RoundMode::CAST_NONE); + pipe_barrier(PIPE_V); + TMOV(beta_r_ub, beta_ub); + pipe_barrier(PIPE_V); + TCOLEXPAND(beta_2d_ub, beta_r_ub); + + TCVT(a1_ub, a1_ub_half, pto::RoundMode::CAST_NONE); + // Form the beta-scaled tile that the later U = A2 * V matmul consumes. + TMUL(a2_ub, a1_ub, beta_2d_ub); + TCVT(a2_ub_half, a2_ub, pto::RoundMode::CAST_NONE); + + if (!first_iter_v) wait_flag_dev(3); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + { + GmShape2D a2_shape(HalfChunk, ChunkSize); + GmStride2D a2_stride(ChunkSize); + GmTensor2D workspace_a2_global(workspace_a2_handle + + static_cast(cid) * WsA2Size + + static_cast(vid) * HalfChunk * ChunkSize, + a2_shape, a2_stride); + TSTORE(workspace_a2_global, a2_ub_half); + } + pipe_barrier(PIPE_ALL); + ffts_cross_core_sync(PIPE_MTE3, 1 | (2 << 4) | (2 << 8)); + + // G is pre-transposed to [H, total_tokens] for contiguous loads. + { + GmShape2D g_shape(1, valid_rows); + GmStride2D g_stride(1); + GmTensor2D g_global( + G_handle + static_cast(head_idx) * total_tokens + chunk_token_start, g_shape, + g_stride); + DynVecTile g_load(1, valid_rows); + TASSIGN(g_load, GUbAddr); + TLOAD(g_load, g_global); + if (valid_rows != ChunkSize) { + TFILLPAD_INPLACE(g_ub, g_load); + } + } + + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + + // Build the g-based column weights before forming the W = A1 * K branch. + TEXP(g_ub, g_ub); + pipe_barrier(PIPE_V); + TMUL(g_ub, g_ub, beta_ub); + pipe_barrier(PIPE_V); + TMOV(g_r_ub, g_ub); + pipe_barrier(PIPE_V); + TCOLEXPAND(g_2d_ub, g_r_ub); + TMUL(a1_ub, a1_ub, g_2d_ub); + TCVT(a1_ub_half, a1_ub, pto::RoundMode::CAST_NONE); + + if (!first_iter_v) wait_flag_dev(4); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + { + GmShape2D a1_shape(HalfChunk, ChunkSize); + GmStride2D a1_stride(ChunkSize); + GmTensor2D workspace_a1_global(workspace_a1_handle + + static_cast(cid) * WsA1Size + + static_cast(vid) * HalfChunk * ChunkSize, + a1_shape, a1_stride); + TSTORE(workspace_a1_global, a1_ub_half); + } + pipe_barrier(PIPE_ALL); + ffts_cross_core_sync(PIPE_MTE3, 1 | (2 << 4) | (1 << 8)); + first_iter_v = false; + } + gi++; + } + } + } + } +#endif + +#if defined(__DAV_C220_CUBE__) + // Cube consumes the two Vec-generated workspaces and turns them into the + // branch outputs U and W. + if (cu_seqlens == nullptr) { + int64_t gi = 0; + for (int64_t seq_idx = 0; seq_idx < num_seqs; ++seq_idx) { + int64_t bos = seq_idx * seq_len; + int64_t slen = seq_len; + int64_t nc = (slen + ChunkSize - 1) / ChunkSize; + + for (int64_t ci = 0; ci < nc; ++ci) { + for (int32_t head_idx = 0; head_idx < NumHeads; ++head_idx) { + if (gi % static_cast(block_num) == static_cast(cid)) { + int64_t chunk_start = ci * ChunkSize; + int64_t remaining = slen - chunk_start; + int32_t valid_rows = static_cast(remaining < ChunkSize ? remaining : ChunkSize); + int64_t chunk_token_start = bos + chunk_start; + + int32_t head_g = head_idx / GROUP; + int64_t k_off = (chunk_token_start * static_cast(Hg) + static_cast(head_g)) * + static_cast(HiddenSize); + int64_t v_off = (chunk_token_start * static_cast(H) + static_cast(head_idx)) * + static_cast(HiddenSize); + + { + GmShape2D k_shape(valid_rows, HiddenSize); + GmStride2D k_stride(BSND_QK_STRIDE); + GmTensor2D k_global(K_handle + k_off, k_shape, k_stride); + DynMatL1 k_l1_load(valid_rows, HiddenSize); + TASSIGN(k_l1_load, 0); + TLOAD(k_l1_load, k_global); + if (valid_rows != ChunkSize) { + TFILLPAD(k_l1_load, k_l1_load); + } + } + { + GmShape2D v_shape(valid_rows, HiddenSize); + GmStride2D v_stride(BSND_V_STRIDE); + GmTensor2D v_global(V_handle + v_off, v_shape, v_stride); + DynMatL1 v_l1_load(valid_rows, HiddenSize); + TASSIGN(v_l1_load, 32768); + TLOAD(v_l1_load, v_global); + if (valid_rows != ChunkSize) { + TFILLPAD(v_l1_load, v_l1_load); + } + } + + wait_flag_dev(2); + { + GmShape2D a2_shape(ChunkSize, ChunkSize); + GmStride2D a2_stride(ChunkSize); + GmTensor2D workspace_a2_global( + workspace_a2_handle + static_cast(cid) * WsA2Size, a2_shape, a2_stride); + // Load the Vec-prepared A2 tile: + // A2 = A * beta[None, :] + TLOAD(a2_l1, workspace_a2_global); + } + + set_flag(PIPE_FIX, PIPE_M, EVENT_ID0); + wait_flag(PIPE_FIX, PIPE_M, EVENT_ID0); + // U = A2 * V keeps the beta-scaled path separate from the K-side update. + gemm_v0(a2_l1, v_l1, u_l0, true); + + { + GmShape2D u_shape(valid_rows, HiddenSize); + GmStride2D u_stride(BSND_V_STRIDE); + GmTensor2D u_global(U_handle + v_off, u_shape, u_stride); + DynAccTile u_store(valid_rows, HiddenSize); + TASSIGN(u_store, 0); + // Store only the valid token rows even though the accumulator tile is + // physically ChunkSize x HiddenSize. + TSTORE(u_global, u_store); + } + ffts_cross_core_sync(PIPE_FIX, 1 | (2 << 4) | (3 << 8)); + + wait_flag_dev(1); + { + GmShape2D a1_shape(ChunkSize, ChunkSize); + GmStride2D a1_stride(ChunkSize); + GmTensor2D workspace_a1_global( + workspace_a1_handle + static_cast(cid) * WsA1Size, a1_shape, a1_stride); + // Load the Vec-prepared A1 tile: + // A1 = A * (exp(g) * beta)[None, :] + TLOAD(a1_l1, workspace_a1_global); + } + + set_flag(PIPE_FIX, PIPE_M, EVENT_ID0); + wait_flag(PIPE_FIX, PIPE_M, EVENT_ID0); + // W = A1 * K uses the g-reweighted path for the complementary WY factor. + gemm_v0(a1_l1, k_l1, w_l0, true); + + { + GmShape2D w_shape(valid_rows, HiddenSize); + GmStride2D w_stride(BSND_V_STRIDE); + GmTensor2D w_global(W_handle + v_off, w_shape, w_stride); + DynAccTile w_store(valid_rows, HiddenSize); + TASSIGN(w_store, 65536); + TSTORE(w_global, w_store); + } + ffts_cross_core_sync(PIPE_FIX, 1 | (2 << 4) | (4 << 8)); + } + gi++; + } + } + } + } else { + int64_t gi = 0; + for (int64_t si = 0; si < num_seqs; ++si) { + int64_t bos = static_cast(cu_seqlens[si]); + int64_t eos = static_cast(cu_seqlens[si + 1]); + int64_t slen = eos - bos; + int64_t nc = (slen + ChunkSize - 1) / ChunkSize; + + for (int64_t ci = 0; ci < nc; ++ci) { + for (int32_t h = 0; h < NumHeads; ++h) { + if (gi % static_cast(block_num) == static_cast(cid)) { + int64_t chunk_start = ci * ChunkSize; + int64_t remaining = slen - chunk_start; + int32_t valid_rows = static_cast(remaining < ChunkSize ? remaining : ChunkSize); + int64_t chunk_token_start = bos + chunk_start; + int32_t head_idx = h; + + int32_t head_g = head_idx / GROUP; + int64_t k_off = (chunk_token_start * static_cast(Hg) + static_cast(head_g)) * + static_cast(HiddenSize); + int64_t v_off = (chunk_token_start * static_cast(H) + static_cast(head_idx)) * + static_cast(HiddenSize); + + { + GmShape2D k_shape(valid_rows, HiddenSize); + GmStride2D k_stride(BSND_QK_STRIDE); + GmTensor2D k_global(K_handle + k_off, k_shape, k_stride); + DynMatL1 k_l1_load(valid_rows, HiddenSize); + TASSIGN(k_l1_load, 0); + TLOAD(k_l1_load, k_global); + if (valid_rows != ChunkSize) { + TFILLPAD(k_l1_load, k_l1_load); + } + } + { + GmShape2D v_shape(valid_rows, HiddenSize); + GmStride2D v_stride(BSND_V_STRIDE); + GmTensor2D v_global(V_handle + v_off, v_shape, v_stride); + DynMatL1 v_l1_load(valid_rows, HiddenSize); + TASSIGN(v_l1_load, 32768); + TLOAD(v_l1_load, v_global); + if (valid_rows != ChunkSize) { + TFILLPAD(v_l1_load, v_l1_load); + } + } + + wait_flag_dev(2); + { + GmShape2D a2_shape(ChunkSize, ChunkSize); + GmStride2D a2_stride(ChunkSize); + GmTensor2D workspace_a2_global( + workspace_a2_handle + static_cast(cid) * WsA2Size, a2_shape, a2_stride); + TLOAD(a2_l1, workspace_a2_global); + } + + set_flag(PIPE_FIX, PIPE_M, EVENT_ID0); + wait_flag(PIPE_FIX, PIPE_M, EVENT_ID0); + // U = A2 * V keeps the beta-scaled path separate from the K-side update. + gemm_v0(a2_l1, v_l1, u_l0, true); + + { + GmShape2D u_shape(valid_rows, HiddenSize); + GmStride2D u_stride(BSND_V_STRIDE); + GmTensor2D u_global(U_handle + v_off, u_shape, u_stride); + DynAccTile u_store(valid_rows, HiddenSize); + TASSIGN(u_store, 0); + TSTORE(u_global, u_store); + } + ffts_cross_core_sync(PIPE_FIX, 1 | (2 << 4) | (3 << 8)); + + wait_flag_dev(1); + { + GmShape2D a1_shape(ChunkSize, ChunkSize); + GmStride2D a1_stride(ChunkSize); + GmTensor2D workspace_a1_global( + workspace_a1_handle + static_cast(cid) * WsA1Size, a1_shape, a1_stride); + TLOAD(a1_l1, workspace_a1_global); + } + + set_flag(PIPE_FIX, PIPE_M, EVENT_ID0); + wait_flag(PIPE_FIX, PIPE_M, EVENT_ID0); + // W = A1 * K uses the g-reweighted path for the complementary WY factor. + gemm_v0(a1_l1, k_l1, w_l0, true); + + { + GmShape2D w_shape(valid_rows, HiddenSize); + GmStride2D w_stride(BSND_V_STRIDE); + GmTensor2D w_global(W_handle + v_off, w_shape, w_stride); + DynAccTile w_store(valid_rows, HiddenSize); + TASSIGN(w_store, 65536); + TSTORE(w_global, w_store); + } + ffts_cross_core_sync(PIPE_FIX, 1 | (2 << 4) | (4 << 8)); + } + gi++; + } + } + } + } +#endif +}